From d5609ac1fc7787d924d37b8c7bb25869e609ac5e Mon Sep 17 00:00:00 2001 From: clairesonglee Date: Wed, 1 Jul 2026 17:19:24 -0700 Subject: [PATCH 001/127] ci: revert to fla 0.4.x (#836) --- .../hooks/train/pretrain/megatron/requirements-megatron.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/runner/helpers/hooks/train/pretrain/megatron/requirements-megatron.txt b/runner/helpers/hooks/train/pretrain/megatron/requirements-megatron.txt index 7fbaa85a1..09d5a624c 100644 --- a/runner/helpers/hooks/train/pretrain/megatron/requirements-megatron.txt +++ b/runner/helpers/hooks/train/pretrain/megatron/requirements-megatron.txt @@ -1,4 +1,4 @@ # Add Megatron pretrain hook private Python deps here when needed. # Qwen3.5 / gated-delta-net attention (Megatron core GatedDeltaNet) needs FLA. -flash-linear-attention==0.5.1 +flash-linear-attention~=0.4.0 causal-conv1d~=1.5 From 77eaeea85fcb870bb3d9f152e2dec1901648af8f Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Mon, 6 Jul 2026 05:07:44 +0300 Subject: [PATCH 002/127] ci: bump Primus-Turbo/AITER pins for Flux diffusion (mxfp4) + skip CI on draft PRs (#806) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. This PR targets `main` and is the **first to merge**. The content branches are cut from this branch (`feat/flux/ci-env`), not raw `main`, so they carry the bumped pins + draft guard on their own heads — which lets the rest of the stack open as drafts immediately, before this PR merges. ## What this changes Updates the public-CI dependency pins to the versions the Flux diffusion feature needs: - Bumps `PRIMUS_TURBO_COMMIT` to a recent Primus-Turbo `main` build that exposes the mxfp4 `gemm_fp4_impl(..., preshuffled=...)` fast path, and `PRIMUS_TURBO_AITER_COMMIT` to AITER `v0.1.14.post1`. Mirrors the same turbo commit into `benchmark.yaml` (the docker build picks both up through `ci.yaml`'s build-args). - Adds a **draft-skip guard** to the expensive jobs (`build-docker`, `run-unittest-torch`, `run-unittest-jax`) so the rest of the stack can be opened as Draft PRs without firing the full docker-build + GPU pipeline on every open/sync; `code-lint` stays ungated. Two parts: add `ready_for_review` to the `pull_request` trigger `types` (the default `opened/synchronize/reopened` set omits it, so without this the draft→ready flip triggers nothing) **and** add `if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }}` to the three heavy jobs. ## Why it merges first (but need not block opening) On the current `main` pins the mxfp4 path fails (`gemm_fp4_impl` rejects the `preshuffled` arg), so content PRs can only go green if their head branch carries the new pins — which stacking on `feat/flux/ci-env` provides without a merge. Merging this first matters at *merge* time: once it's on `main`, children auto-retargeted to `main` inherit the fix and the guard. No `MEGATRON_PATH` change is needed: the diffusion unit tests run in-process and `tests/conftest.py` already puts the `submodules: recursive`-checked-out `third_party/Megatron-LM` on `sys.path` when `MEGATRON_PATH` is unset. ## Dependencies None — this is the root prerequisite. ## Test plan No unit tests of its own. After the runner rebuilds the image, confirm a smoke run is green: `pytest tests/unit_tests/backends/megatron/diffusion -k mxfp4`. The guard can be sanity-checked by opening a throwaway draft PR (expect only `code-lint`) and marking it ready (expect the full pipeline). ## Files 2 (`.github/workflows/ci.yaml` — pins + draft guard; `.github/workflows/benchmark.yaml` — mirrored turbo pin). --------- Co-authored-by: Flux Split Co-authored-by: WangLingxun Co-authored-by: Xiaoming-AMD Co-authored-by: luiza-amd --- .github/workflows/ci.yaml | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index f3916a33d..c19a82543 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -8,6 +8,9 @@ on: tags: - "v*" pull_request: + # Re-list the implicit defaults (opened/synchronize/reopened) and add + # ready_for_review so flipping a Draft PR to Ready triggers its first run. + types: [opened, synchronize, reopened, ready_for_review] concurrency: group: ${{ github.workflow }}-${{ github.head_ref || github.event.merge_group.head_ref || github.ref }} @@ -21,7 +24,7 @@ permissions: env: PRIMUS_TURBO_COMMIT: a04a233cbfb468dbe21600cbf9db70953428b25c # feat: force use nt layout gemm in bwd (#386) - PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # v0.14.0.post1 + PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # AITER v0.1.14.post1 (tag commit) — required by Primus-Turbo main aiter_utils.py ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355) UCCL_COMMIT: 5afb4117893c58cc0c8557d9286336141a301053 # [EP]: fix fp8 error of internode_ll on amd gfx950 arch. (#710) TRITON_COMMIT: 88b227e23f0445f3f695bad05bbf1a363b4f50e0 @@ -69,6 +72,9 @@ jobs: build-docker: needs: [code-lint] + # Skip on Draft PRs (keep code-lint ungated); the event_name guard preserves + # push/tag and workflow_dispatch runs where github.event.pull_request is absent. + if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }} runs-on: build-docker strategy: matrix: @@ -221,6 +227,7 @@ jobs: # PRIMUS_WORKDIR: /wekafs/primus-data/primus_safe_ci/torch PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32: 1 needs: [code-lint] + if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }} # runs-on: [primus-lm-cicd-torch-j8knc] runs-on: [primus-lm-cicd-v26.3-tas8n-a16-40] steps: @@ -499,6 +506,7 @@ jobs: # PRIMUS_WORKDIR: /wekafs/primus-data/primus_safe_ci/jax PRIMUS_WORKDIR: /mnt/apps_proxy/tas/0_public/primus_docker_jax_ci/actions-runner needs: [code-lint] + if: ${{ github.event_name != 'pull_request' || !github.event.pull_request.draft }} runs-on: [primus-jax-tas-runner] # docker container primus_jax_github_runner on tas a16-31 steps: - run: echo "🎉 Begin Primus-Turbo Checkout." From 73cd33da5d9f4b3130390e1bc6a0d3824ca68a5e Mon Sep 17 00:00:00 2001 From: Fuyuan Jing <167437074+amd-fuyuajin@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:47:56 -0400 Subject: [PATCH 003/127] add mixtral-8x22B config files for maxtext backend (#841) --- .../MI300X/mixtral_8x22B-pretrain.yaml | 41 +++++++++++++++++++ .../MI355X/mixtral_8x22B-pretrain.yaml | 41 +++++++++++++++++++ .../configs/models/maxtext/mixtral_8x22B.yaml | 5 +++ 3 files changed, 87 insertions(+) create mode 100644 examples/maxtext/configs/MI300X/mixtral_8x22B-pretrain.yaml create mode 100644 examples/maxtext/configs/MI355X/mixtral_8x22B-pretrain.yaml create mode 100644 primus/configs/models/maxtext/mixtral_8x22B.yaml diff --git a/examples/maxtext/configs/MI300X/mixtral_8x22B-pretrain.yaml b/examples/maxtext/configs/MI300X/mixtral_8x22B-pretrain.yaml new file mode 100644 index 000000000..3db838b20 --- /dev/null +++ b/examples/maxtext/configs/MI300X/mixtral_8x22B-pretrain.yaml @@ -0,0 +1,41 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:mixtral_8x22B-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: maxtext + config: pre_trainer.yaml + + # model to run + model: mixtral_8x22B.yaml + overrides: + run_name: "mixtral_8x22b_training" + base_output_directory: "./output" + steps: 50 + profiler: "" + + # data + dataset_type: "synthetic" + hf_access_token: ${HF_TOKEN:""} + + # checkpoint + enable_checkpointing: false + async_checkpointing: false + + # inter-node parallelism strategy + dcn_data_parallelism: -1 + dcn_fsdp_parallelism: 1 + + # intra-node parallelism strategy + ici_fsdp_parallelism: 1 + ici_data_parallelism: 1 + ici_expert_parallelism: -1 + + sparse_matmul: false + megablox: false + capacity_factor: 1 + max_target_length: 4096 + per_device_batch_size: 4 + remat_policy: "save_dot_with_context_except_mlp" diff --git a/examples/maxtext/configs/MI355X/mixtral_8x22B-pretrain.yaml b/examples/maxtext/configs/MI355X/mixtral_8x22B-pretrain.yaml new file mode 100644 index 000000000..d4ea7fd6b --- /dev/null +++ b/examples/maxtext/configs/MI355X/mixtral_8x22B-pretrain.yaml @@ -0,0 +1,41 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:mixtral_8x22B-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: maxtext + config: pre_trainer.yaml + + # model to run + model: mixtral_8x22B.yaml + overrides: + run_name: "mixtral_8x22b_training" + base_output_directory: "./output" + steps: 50 + profiler: "" + + # data + dataset_type: "synthetic" + hf_access_token: ${HF_TOKEN:""} + + # checkpoint + enable_checkpointing: false + async_checkpointing: false + + # inter-node parallelism strategy + dcn_data_parallelism: -1 + dcn_fsdp_parallelism: 1 + + # intra-node parallelism strategy + ici_fsdp_parallelism: 1 + ici_data_parallelism: 1 + ici_expert_parallelism: -1 + + sparse_matmul: false + megablox: false + capacity_factor: 1 + max_target_length: 4096 + per_device_batch_size: 8 + remat_policy: "save_dot_with_context_except_mlp" diff --git a/primus/configs/models/maxtext/mixtral_8x22B.yaml b/primus/configs/models/maxtext/mixtral_8x22B.yaml new file mode 100644 index 000000000..cccbceab3 --- /dev/null +++ b/primus/configs/models/maxtext/mixtral_8x22B.yaml @@ -0,0 +1,5 @@ +extends: + - model_base.yaml + +model_name: "mixtral-8x22b" +tokenizer_path: "mistralai/Mixtral-8x22B-v0.1" From 868d32ef158917659e42c78891e8f9274508af1b Mon Sep 17 00:00:00 2001 From: Fuyuan Jing <167437074+amd-fuyuajin@users.noreply.github.com> Date: Mon, 6 Jul 2026 21:48:10 -0400 Subject: [PATCH 004/127] improve the logging format for the megatron backend (#842) ## made some minor changes to the logging format of Megatron backend for training steps 1. Currently, we label TFLOP/s/GPU as throughput. Per customer's feedback, throughput is ambiguous and not the best to describe TFLOP/s/GPU. Change it to `Compute per GPU`. 2. For the tokens/s/GPU, we report two numbers in the format of `###/###`. It's not clear what they are. The number before `/` is perf of the step, the number after `/` is arithmetic mean of many steps. Since we measure tokens/s/GPU, which is a rate, it's better to use harmonic mean. So, change this to harmonic mean. **Before change** image **After change** image --- .../training_log/print_rank_last_patches.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/primus/backends/megatron/patches/training_log/print_rank_last_patches.py b/primus/backends/megatron/patches/training_log/print_rank_last_patches.py index 9485b3116..a43a83b26 100644 --- a/primus/backends/megatron/patches/training_log/print_rank_last_patches.py +++ b/primus/backends/megatron/patches/training_log/print_rank_last_patches.py @@ -425,7 +425,7 @@ def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> s idx = parsed.throughput_index if idx is not None and 0 <= idx < len(parsed.segments): parsed.segments[idx] = ( - f"throughput per GPU (TFLOP/s/GPU): " f"{tflops_value:.1f}/{avg_tflops:.1f}" + f"compute per GPU (TFLOP/s/GPU): {tflops_value:.1f} (avg {avg_tflops:.1f})" ) # ---------------- Tokens/s ---------------- @@ -466,7 +466,16 @@ def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> s warning_rank_0(f"[Patch:megatron.training_log] No token throughput") return log_string - avg_tokens = sum(self._recent_token_throughputs) / len(self._recent_token_throughputs) + # Use the harmonic mean for the token throughput average. The harmonic + # mean is the correct way to average rates (tokens/s) over iterations of + # equal token count, since it weights slow iterations more heavily. + positive_token_throughputs = [t for t in self._recent_token_throughputs if t > 0] + if positive_token_throughputs: + avg_tokens = len(positive_token_throughputs) / sum( + 1.0 / t for t in positive_token_throughputs + ) + else: + avg_tokens = 0.0 # Append token throughput directly after the TFLOP throughput within # the same segment. We do not create a new segment to keep the log @@ -475,7 +484,7 @@ def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> s if idx is not None and 0 <= idx < len(parsed.segments): parsed.segments[idx] = ( f"{parsed.segments[idx]} " - f" | tokens per GPU (tokens/s/GPU): {token_value:.1f}/{avg_tokens:.1f}" + f" | tokens/s/GPU inst/harmonic mean: {token_value:.1f}/{avg_tokens:.1f}" ) # String result is ignored by the main patch when parsed is provided. From e2f7ee6760286b3e51b4ea8cac6a08ba27f2c9c2 Mon Sep 17 00:00:00 2001 From: luiza-amd Date: Tue, 7 Jul 2026 11:27:05 +0300 Subject: [PATCH 005/127] merge to main of feat(flux): core runtime + Megatron adapter scaffolding (#856) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **Important**: The content of this PR was approved and merged https://github.com/AMD-AGI/Primus/pull/807, but not to main (auto-target to main was not triggered and the lesson was taken into account for further PRs). **The purpose of this PR is to finalize the merge to main and no new changes were introduced.** Base of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/ci-env` and auto-retargets to `main` once that merges. Every other content PR stacks on this one. ## What this changes The shared runtime scaffolding the rest of the feature builds on: core runtime state + train-runtime wiring, the Megatron adapter and base/pretrain trainers, and a patch auto-loader (`patches/__init__.py`) that imports every `*_patches.py` in the package on import, so each later layer only drops in its own patch file with no registry edit. Also carries the shared test root (`tests/conftest.py`, `tests/utils.py`) and a one-line `.gitignore` change (the bare `data` ignore → root-anchored `/data/*`, and nothing else) that stops Git from ignoring the in-repo `data/` source and config directories the later layers add. Kept deliberately small so it can land first. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); no functional dependency on the turbo bump. ## Test plan `pytest tests/unit_tests/core tests/unit_tests/backends/megatron`; lint/pre-commit clean. Validated locally on an AMD GPU container: 58 passed. ## Files 39 (core runtime, Megatron adapter/trainers, base patch loader, shared test root, `.gitignore`). --------- Co-authored-by: Flux Split Co-authored-by: WangLingxun Co-authored-by: Xiaoming-AMD Co-authored-by: Flux Split Trial --- .gitignore | 2 +- .../hummingbirdxt_posttrain_trainer.py | 5 +- .../maxtext/maxtext_pretrain_trainer.py | 7 +- primus/backends/megatron/__init__.py | 21 +- primus/backends/megatron/megatron_adapter.py | 98 ++++- .../megatron/megatron_base_trainer.py | 87 ++++- .../megatron/megatron_pretrain_trainer.py | 123 ++++-- .../backends/megatron/megatron_sft_trainer.py | 5 +- primus/backends/megatron/patches/__init__.py | 17 + .../backends/megatron/patches/_patch_guard.py | 45 +++ .../megatron/patches/args/__init__.py | 4 +- .../patches/args/checkpoint_path_patches.py | 43 ++- .../patches/args/hsdp_args_patches.py | 60 +++ .../megatron/patches/build_model_patches.py | 65 +++- .../patches/distributed_init_patches.py | 76 ++++ .../backends/megatron/patches/env_patches.py | 7 +- .../megatron/patches/mp_sync_skip_patches.py | 63 ++++ .../backends/megatron/training/global_vars.py | 2 +- .../megatron_bridge_base_trainer.py | 5 +- .../megatron_bridge_posttrain_trainer.py | 5 +- .../megatron_bridge_pretrain_trainer.py | 5 +- .../torchtitan/torchtitan_pretrain_trainer.py | 7 +- primus/core/backend/backend_adapter.py | 12 +- primus/core/backend/backend_registry.py | 2 +- primus/core/config/primus_config.py | 21 +- primus/core/launcher/parser.py | 17 + primus/core/runtime/runtime_state.py | 34 ++ primus/core/runtime/train_runtime.py | 147 +++++++- primus/core/trainer/base_trainer.py | 45 ++- primus/modules/base_module.py | 23 +- primus/modules/module_utils.py | 20 +- .../modules/trainer/megatron/pre_trainer.py | 10 +- primus/modules/trainer/megatron/trainer.py | 26 +- primus/modules/trainer/megatron/utils.py | 27 ++ tests/conftest.py | 137 +++++++ .../unit_tests/backends/megatron/conftest.py | 153 +++++--- ...on_pretrain_trainer_overridable_methods.py | 115 ++++++ .../megatron/patches/test_patch_guard.py | 75 ++++ .../megatron/test_megatron_adapter.py | 189 ++++++---- .../megatron/test_megatron_base_trainer.py | 213 +++++++++++ .../core/backend/test_backend_adapter.py | 19 +- .../core/backend/test_backend_registry.py | 210 ++++------- .../core/config/test_primus_config.py | 61 +++ .../core/runtime/test_train_runtime.py | 350 +++++++++++++++++- .../core/trainer/test_base_trainer.py | 184 +++++---- tests/utils.py | 72 ++-- 46 files changed, 2410 insertions(+), 504 deletions(-) create mode 100644 primus/backends/megatron/patches/_patch_guard.py create mode 100644 primus/backends/megatron/patches/args/hsdp_args_patches.py create mode 100644 primus/backends/megatron/patches/distributed_init_patches.py create mode 100644 primus/backends/megatron/patches/mp_sync_skip_patches.py create mode 100644 primus/core/runtime/runtime_state.py create mode 100644 tests/conftest.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_megatron_pretrain_trainer_overridable_methods.py create mode 100644 tests/unit_tests/backends/megatron/patches/test_patch_guard.py create mode 100644 tests/unit_tests/backends/megatron/test_megatron_base_trainer.py create mode 100644 tests/unit_tests/core/config/test_primus_config.py diff --git a/.gitignore b/.gitignore index 97c0687f4..3353a8bdb 100644 --- a/.gitignore +++ b/.gitignore @@ -11,5 +11,5 @@ local/ .gitmodules output experiment -data +/data/* pp_simulation_result diff --git a/primus/backends/hummingbirdxt/hummingbirdxt_posttrain_trainer.py b/primus/backends/hummingbirdxt/hummingbirdxt_posttrain_trainer.py index 2356fe114..aed8315b8 100644 --- a/primus/backends/hummingbirdxt/hummingbirdxt_posttrain_trainer.py +++ b/primus/backends/hummingbirdxt/hummingbirdxt_posttrain_trainer.py @@ -11,8 +11,9 @@ class HummingbirdXTPosttrainTrainer(BaseTrainer): - def __init__(self, backend_args: Any): - super().__init__(backend_args=backend_args) + def __init__(self, backend_args: Any = None, **kwargs): + # Accept and forward runtime context kwargs so BaseTrainer can filter them. + super().__init__(backend_args=backend_args, **kwargs) def setup(self): pass diff --git a/primus/backends/maxtext/maxtext_pretrain_trainer.py b/primus/backends/maxtext/maxtext_pretrain_trainer.py index e521041a2..8009230ac 100644 --- a/primus/backends/maxtext/maxtext_pretrain_trainer.py +++ b/primus/backends/maxtext/maxtext_pretrain_trainer.py @@ -41,8 +41,11 @@ class MaxTextPretrainTrainer(BaseTrainer): Trainer class for MaxText pre-training. """ - def __init__(self, backend_args: Any): - super().__init__(backend_args=backend_args) + def __init__(self, backend_args: Any = None, **kwargs): + # The core runtime instantiates every trainer with BaseModule-style + # context kwargs (module_name, primus_config, module_rank, ...). MaxText + # does not need them; accept and forward so BaseTrainer can filter them. + super().__init__(backend_args=backend_args, **kwargs) # Training state (populated in init()) self.train_config: Optional[Any] = None diff --git a/primus/backends/megatron/__init__.py b/primus/backends/megatron/__init__.py index 8cf9e639f..939a1c4ba 100644 --- a/primus/backends/megatron/__init__.py +++ b/primus/backends/megatron/__init__.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -12,3 +12,22 @@ BackendRegistry.register_adapter("megatron", MegatronAdapter) BackendRegistry.register_trainer_class(MegatronPretrainTrainer, "megatron") BackendRegistry.register_trainer_class(MegatronSFTTrainer, "megatron", "sft") + +# Export trainers for convenience +# Use lazy import for FluxPretrainTrainer to avoid Megatron dependency +# when importing data pipeline components +__all__ = [ + "MegatronAdapter", + "FluxPretrainTrainer", + "MegatronPretrainTrainer", + "MegatronSFTTrainer", +] + + +def __getattr__(name): + """Lazy import for trainer classes to avoid Megatron dependency on module import.""" + if name == "FluxPretrainTrainer": + from primus.backends.megatron.flux_pretrain_trainer import FluxPretrainTrainer + + return FluxPretrainTrainer + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/primus/backends/megatron/megatron_adapter.py b/primus/backends/megatron/megatron_adapter.py index 6f1e58321..66574ec00 100644 --- a/primus/backends/megatron/megatron_adapter.py +++ b/primus/backends/megatron/megatron_adapter.py @@ -1,9 +1,12 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### +import importlib +from typing import Optional + import primus.backends.megatron.patches # noqa: F401 # Register patches from primus.backends.megatron.argument_builder import MegatronArgBuilder from primus.core.backend.backend_adapter import BackendAdapter @@ -18,8 +21,27 @@ def __init__(self, framework="megatron"): super().__init__(framework) self.third_party_dir_name = "Megatron-LM" - def load_trainer_class(self, stage: str = "pretrain"): - """Return the trainer class for the specified training stage.""" + def load_trainer_class(self, stage: str = "pretrain", trainer_class: Optional[str] = None): + """ + Return the Trainer class for the specified training stage or trainer class name. + + Args: + stage: Training stage (e.g., "pretrain", "sft"). Defaults to "pretrain". + trainer_class: Optional specific trainer class name for dynamic loading. + If provided, this takes precedence over stage-based selection. + + Returns: + Trainer class + + Raises: + ImportError: If trainer class cannot be imported + ValueError: If stage is invalid or trainer class is not found + """ + # If trainer_class is specified, attempt dynamic loading + if trainer_class: + return self._load_trainer_class_by_name(trainer_class) + + # Fallback to stage-based selection via the backend registry try: trainer_cls = BackendRegistry.get_trainer_class(self.framework, stage=stage) except (ValueError, AssertionError) as exc: @@ -31,6 +53,76 @@ def load_trainer_class(self, stage: str = "pretrain"): log_rank_0(f"[Primus:MegatronAdapter] Loaded trainer class: {trainer_cls.__name__}") return trainer_cls + def _load_trainer_class_by_name(self, trainer_class: str): + """ + Dynamically load trainer class by name. + + Args: + trainer_class: Trainer class name (e.g., "FluxPretrainTrainer", "MegatronPretrainTrainer") + + Returns: + Trainer class + + Raises: + ImportError: If trainer class cannot be imported + ValueError: If trainer class is not found + """ + # Define trainer registry for Megatron backend + # This maps trainer class names to their module paths + MEGATRON_TRAINERS = { + "MegatronPretrainTrainer": "primus.backends.megatron.megatron_pretrain_trainer.MegatronPretrainTrainer", + "FluxPretrainTrainer": "primus.backends.megatron.flux_pretrain_trainer.FluxPretrainTrainer", + } + + if trainer_class not in MEGATRON_TRAINERS: + # Try to load from common locations as fallback + log_rank_0( + f"[Primus:MegatronAdapter] Trainer '{trainer_class}' not in registry, attempting dynamic import..." + ) + + possible_paths = [ + f"primus.backends.megatron.{trainer_class.lower()}.{trainer_class}", + f"primus.modules.trainer.megatron.{trainer_class.lower()}.{trainer_class}", + f"primus.backends.megatron.{trainer_class}", + ] + + for module_path in possible_paths: + try: + module_name, class_name = module_path.rsplit(".", 1) + module = importlib.import_module(module_name) + trainer_cls = getattr(module, class_name) + log_rank_0( + f"[Primus:MegatronAdapter] Successfully loaded trainer: {trainer_class} from {module_name}" + ) + return trainer_cls + except (ImportError, AttributeError): + continue + + # If all attempts failed, provide helpful error message + available = list(MEGATRON_TRAINERS.keys()) + raise ValueError( + f"Trainer class '{trainer_class}' not found.\n" + f"Available trainers: {', '.join(available) if available else 'none'}\n" + f"Hint: Set 'trainer_class' in your experiment YAML to a registered trainer " + f"(e.g. trainer_class: FluxPretrainTrainer), register it in MEGATRON_TRAINERS, " + f"or ensure it is importable from standard locations." + ) + + # Load from registry + trainer_path = MEGATRON_TRAINERS[trainer_class] + module_path, class_name = trainer_path.rsplit(".", 1) + + try: + module = importlib.import_module(module_path) + trainer_cls = getattr(module, class_name) + log_rank_0(f"[Primus:MegatronAdapter] Loaded trainer: {trainer_class} from {module_path}") + return trainer_cls + except (ImportError, AttributeError) as e: + raise ImportError( + f"Failed to load trainer class '{trainer_class}' from '{module_path}': {e}\n" + f"Hint: Check that the module exists and the class name is correct." + ) from e + def detect_backend_version(self) -> str: """Detect Megatron-LM version via AST parsing (avoids __init__.py execution).""" import ast diff --git a/primus/backends/megatron/megatron_base_trainer.py b/primus/backends/megatron/megatron_base_trainer.py index 9f6a84692..89cbb3813 100644 --- a/primus/backends/megatron/megatron_base_trainer.py +++ b/primus/backends/megatron/megatron_base_trainer.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -14,7 +14,6 @@ ) from primus.backends.megatron.training.mlflow_setup import upload_mlflow_artifacts from primus.core.trainer.base_trainer import BaseTrainer -from primus.core.utils.env import flush_before_hard_exit from primus.modules.module_utils import log_rank_0, warning_rank_0 @@ -23,9 +22,84 @@ class MegatronBaseTrainer(BaseTrainer): def setup(self): """Setup Megatron runtime: set global vars and patch parse_args.""" + self._ensure_megatron_path() set_primus_global_variables(self.backend_args) self._patch_parse_args() + def _ensure_megatron_path(self): + """Ensure Megatron-LM path is in sys.path before any megatron imports.""" + import os + import sys + from pathlib import Path + + # First, check if megatron is already importable + try: + import megatron # type: ignore + + return # Path already set correctly + except ImportError: + pass + + # Try multiple methods to find the Megatron-LM path + megatron_paths = [] + + # Method 1: From PRIMUS_PATH environment variable (most reliable) + primus_path = os.getenv("PRIMUS_PATH") + if primus_path: + megatron_path = Path(primus_path) / "third_party" / "Megatron-LM" + if megatron_path.exists(): + megatron_paths.append(str(megatron_path)) + + # Method 2: From current working directory (works in container) + try: + cwd = Path.cwd() + # Check if we're in /workspace/Primus or a subdirectory + if "Primus" in str(cwd): + # Find the Primus root + primus_root = cwd + while primus_root.name != "Primus" and primus_root != primus_root.parent: + primus_root = primus_root.parent + if primus_root.name == "Primus": + megatron_path = primus_root / "third_party" / "Megatron-LM" + if megatron_path.exists(): + megatron_paths.append(str(megatron_path)) + except Exception: + pass + + # Method 3: From current file location + try: + repo_root = Path(__file__).resolve().parents[3] + megatron_path = repo_root / "third_party" / "Megatron-LM" + if megatron_path.exists(): + megatron_paths.append(str(megatron_path)) + except Exception: + pass + + # Method 4: Check if already in sys.path + for path in sys.path: + path_obj = Path(path) + if path_obj.exists(): + megatron_pkg = path_obj / "megatron" + if megatron_pkg.exists() and megatron_pkg.is_dir(): + return # Path already set correctly + + # Add paths to sys.path if not already present + for path in megatron_paths: + if path not in sys.path: + sys.path.insert(0, path) + log_rank_0(f"[Primus:MegatronBaseTrainer] Added Megatron-LM to sys.path: {path}") + + # Verify the path was set correctly by trying to import + try: + import megatron # type: ignore + + log_rank_0("[Primus:MegatronBaseTrainer] Successfully verified megatron import") + except ImportError as e: + log_rank_0( + f"[Primus:MegatronBaseTrainer] WARNING: Failed to import megatron after path setup: {e}" + ) + log_rank_0(f"[Primus:MegatronBaseTrainer] sys.path: {sys.path[:5]}") + def init(self): """Initialize Megatron training components.""" log_rank_0("Initializing Megatron training...") @@ -53,7 +127,14 @@ def cleanup(self, on_error: bool = False): if exit_fast and not on_error: log_rank_0("[MegatronBaseTrainer] PRIMUS_EXIT_FAST=1 -> os._exit(0)") - flush_before_hard_exit() + # Flush stdout/stderr so the final log lines are not lost. + try: + import sys + + sys.stdout.flush() + sys.stderr.flush() + except Exception: # pragma: no cover + pass os._exit(0) def _finalize_mlflow_artifacts(self): diff --git a/primus/backends/megatron/megatron_pretrain_trainer.py b/primus/backends/megatron/megatron_pretrain_trainer.py index 5a2e795c4..4955491d8 100644 --- a/primus/backends/megatron/megatron_pretrain_trainer.py +++ b/primus/backends/megatron/megatron_pretrain_trainer.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -11,6 +11,70 @@ class MegatronPretrainTrainer(MegatronBaseTrainer): """Trainer for Megatron-LM pre-training.""" + def get_forward_step(self): + """ + Return forward step function for training loop. + + Override this method in subclasses to provide custom forward step functions. + Default implementation returns GPT forward step. + + Note: This method assumes Megatron-LM path is already set up by + MegatronBaseTrainer._ensure_megatron_path() during setup(). + + Returns: + Callable: Forward step function compatible with Megatron's training loop. + """ + # Path should already be set by _ensure_megatron_path() during setup() + try: + from pretrain_gpt import forward_step # type: ignore + except ImportError as e: + log_rank_0(f"[MegatronPretrainTrainer] Failed to import forward_step from pretrain_gpt: {e}") + log_rank_0( + "[MegatronPretrainTrainer] This indicates a configuration issue. " + "Ensure Megatron-LM is properly installed at third_party/Megatron-LM" + ) + raise ImportError( + "Could not import forward_step from pretrain_gpt. " + "This indicates that Megatron-LM path setup failed during setup(). " + "Ensure Megatron-LM is properly installed and _ensure_megatron_path() succeeded." + ) from e + return forward_step + + def get_datasets_provider(self): + """ + Return dataset provider function for training. + + Override this method in subclasses to provide custom dataset providers. + Default implementation returns GPT dataset provider. + + Note: This method assumes Megatron-LM path is already set up by + MegatronBaseTrainer._ensure_megatron_path() during setup(). + + Returns: + Callable: Dataset provider function compatible with Megatron's training loop. + """ + # Path should already be set by _ensure_megatron_path() during setup() + try: + from pretrain_gpt import train_valid_test_datasets_provider # type: ignore + except ImportError as e: + log_rank_0( + f"[MegatronPretrainTrainer] Failed to import train_valid_test_datasets_provider " + f"from pretrain_gpt: {e}" + ) + log_rank_0( + "[MegatronPretrainTrainer] This indicates a configuration issue. " + "Ensure Megatron-LM is properly installed at third_party/Megatron-LM" + ) + raise ImportError( + "Could not import train_valid_test_datasets_provider from pretrain_gpt. " + "This indicates that Megatron-LM path setup failed during setup(). " + "Ensure Megatron-LM is properly installed and _ensure_megatron_path() succeeded." + ) from e + + provider = train_valid_test_datasets_provider + provider.is_distributed = True # Always True to match Megatron's behavior + return provider + def train(self): """Execute Megatron pre-training.""" log_rank_0("Executing Megatron pretrain...") @@ -22,11 +86,10 @@ def train(self): from primus.core.utils.import_utils import get_model_provider - # Determine model type (gpt or mamba) from backend_args + # Determine model type (gpt, mamba, or diffusion) from backend_args model_type = getattr(self.backend_args, "model_type", "gpt") log_rank_0(f"-detected model_type: {model_type}") - # Import the appropriate training components based on model_type if model_type == "mamba": from pretrain_mamba import ( # type: ignore forward_step, @@ -34,16 +97,16 @@ def train(self): ) log_rank_0("Using Mamba model provider and training components") + # Upstream pretrain entrypoints set this in their __main__ blocks, but Primus imports the + # provider directly and calls pretrain() programmatically. Without restoring this flag, + # only TP rank 0 enters dataset construction while the core dataset builder still issues + # distributed barriers, which deadlocks for TP>1. + train_valid_test_datasets_provider.is_distributed = True else: - from pretrain_gpt import ( # type: ignore - forward_step, - train_valid_test_datasets_provider, - ) - - log_rank_0("Using GPT model provider and training components") - - # Configure training components - train_valid_test_datasets_provider.is_distributed = True + # Use overridable methods so subclasses (e.g. diffusion/Flux) can plug in their own + # forward_step / dataset_provider. Defaults pull from pretrain_gpt. + forward_step = self.get_forward_step() + train_valid_test_datasets_provider = self.get_datasets_provider() # Handle Megatron version differences (v0.12.0 vs newer with inprocess_restart) wrapped_pretrain = pretrain @@ -65,12 +128,14 @@ def train(self): if "store" in sig.parameters: kwargs["store"] = store - # Get model provider with correct model_type - # Only pass model_type if it's not the default to maintain compatibility - if model_type != "gpt": - model_provider = get_model_provider(model_type=model_type) - else: - model_provider = get_model_provider() + # Resolve model_provider: prefer subclass-provided attribute (e.g. diffusion/Flux), + # else fall back to registry-based lookup with model_type for mamba/gpt. + model_provider = getattr(self, "model_provider", None) + if model_provider is None: + if model_type != "gpt": + model_provider = get_model_provider(model_type=model_type) + else: + model_provider = get_model_provider() log_rank_0(f"-model_provider: {model_provider}") # Patch Megatron's get_forward_backward_func to support dump_pp_data @@ -101,13 +166,21 @@ def patched_get_forward_backward_func(*args, **kwargs): if hasattr(mt_training, "get_forward_backward_func"): mt_training.get_forward_backward_func = patched_get_forward_backward_func - wrapped_pretrain( - train_valid_test_datasets_provider, - model_provider, - ModelType.encoder_or_decoder, - forward_step, - **kwargs, - ) + try: + wrapped_pretrain( + train_valid_test_datasets_provider, + model_provider, + ModelType.encoder_or_decoder, + forward_step, + **kwargs, + ) + log_rank_0("[MegatronPretrainTrainer] pretrain() completed successfully") + except Exception as e: + log_rank_0(f"[MegatronPretrainTrainer] ERROR in pretrain(): {type(e).__name__}: {e}") + import traceback + + log_rank_0(f"[MegatronPretrainTrainer] Traceback: {traceback.format_exc()}") + raise # Dump PP visualization data if enabled try: diff --git a/primus/backends/megatron/megatron_sft_trainer.py b/primus/backends/megatron/megatron_sft_trainer.py index 7a14350b6..0f95eebc9 100644 --- a/primus/backends/megatron/megatron_sft_trainer.py +++ b/primus/backends/megatron/megatron_sft_trainer.py @@ -30,14 +30,15 @@ class MegatronSFTTrainer(MegatronBaseTrainer): - Common Megatron initialization patterns """ - def __init__(self, backend_args: Any): + def __init__(self, backend_args: Any = None, **kwargs): """ Initialize Megatron SFT trainer. Args: backend_args: Megatron-LM argument namespace (from MegatronArgBuilder) + **kwargs: Runtime context kwargs forwarded to BaseTrainer for filtering. """ - super().__init__(backend_args=backend_args) + super().__init__(backend_args=backend_args, **kwargs) # Initialize LoRA if enabled self.peft = None diff --git a/primus/backends/megatron/patches/__init__.py b/primus/backends/megatron/patches/__init__.py index 7b294fd70..6ae718c8b 100644 --- a/primus/backends/megatron/patches/__init__.py +++ b/primus/backends/megatron/patches/__init__.py @@ -8,6 +8,23 @@ Megatron Patch Collection This module defines the public entrypoint for applying Megatron-specific patches. + +Registration vs. application (important for non-Flux / LLM users): + Importing this package eagerly imports every ``*_patches.py`` module, which + runs each module's ``@register_patch`` side effect. This registration is + GLOBAL: it happens for *every* Megatron job (LLM or diffusion), because + ``megatron_adapter`` imports this package unconditionally. + + Registration only adds a patch to the registry; it does NOT apply it. Each + patch carries a ``condition=...`` predicate (typically gated on a config + flag such as ``torch_compile.enable``, ``use_fsdp2_fp8_all_gather``, or a + diffusion-specific arg) that is evaluated at ``run_patches`` time. A patch + whose condition is False is a no-op for that job. + + Consequence: diffusion/Flux-specific patches are registered for LLM jobs but + should not take effect there. When adding a patch, make its ``condition`` + precise so it cannot alter unrelated (e.g. non-Flux) training. An opt-in + "patch profile" mechanism could make this stricter in the future. """ import importlib diff --git a/primus/backends/megatron/patches/_patch_guard.py b/primus/backends/megatron/patches/_patch_guard.py new file mode 100644 index 000000000..c86dfc0b8 --- /dev/null +++ b/primus/backends/megatron/patches/_patch_guard.py @@ -0,0 +1,45 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Idempotency guard for monkeypatches that *wrap* (compose) a target callable. + +The core patch runner (``primus.core.patches``) has no built-in re-apply guard: +running the same phase twice re-invokes every patch handler. For patches that +replace a class method or module attribute outright that is harmless (the +assignment is idempotent), but patches that *wrap* an existing callable -- e.g. +the stacked ``train_step`` wrappers (FP8 cache refresh, delayed-scaling +preamble, wall-clock timer) -- would wrap again on a second run and silently +double their side effects. + +This helper records applied patch keys in a sentinel set attached to the +*object that owns the wrapped attribute* (typically a module). Keying on the +owner object -- which is stable across re-runs -- rather than on the wrapped +callable makes the guard robust even when several wrappers compose on the same +attribute. + +This lives under ``primus.backends.megatron.patches`` (feature-owned) on +purpose: the shared ``primus.core.patches`` framework is inherited and must not +grow feature-specific behavior. +""" + +from typing import Any + +_SENTINEL_ATTR = "_primus_applied_patch_keys" + + +def is_patched(target: Any, key: str) -> bool: + """Return True if ``key`` has already been applied to ``target``.""" + applied = getattr(target, _SENTINEL_ATTR, None) + return applied is not None and key in applied + + +def mark_patched(target: Any, key: str) -> None: + """Record that ``key`` has been applied to ``target``.""" + applied = getattr(target, _SENTINEL_ATTR, None) + if applied is None: + applied = set() + setattr(target, _SENTINEL_ATTR, applied) + applied.add(key) diff --git a/primus/backends/megatron/patches/args/__init__.py b/primus/backends/megatron/patches/args/__init__.py index a994f95b8..647847380 100644 --- a/primus/backends/megatron/patches/args/__init__.py +++ b/primus/backends/megatron/patches/args/__init__.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. +# Copyright (c) 2026, Advanced Micro Devices, Inc. # # See LICENSE for license information. ############################################################################### @@ -19,6 +19,7 @@ from . import ( # noqa: F401 checkpoint_path_patches, data_path_split_patches, + hsdp_args_patches, iterations_to_skip_default_patches, logging_level_patches, mock_data_patches, @@ -35,6 +36,7 @@ "wandb_config_patches", "logging_level_patches", "data_path_split_patches", + "hsdp_args_patches", "mock_data_patches", "sequence_parallel_tp1_patches", "iterations_to_skip_default_patches", diff --git a/primus/backends/megatron/patches/args/checkpoint_path_patches.py b/primus/backends/megatron/patches/args/checkpoint_path_patches.py index c67a77b42..b0736551f 100644 --- a/primus/backends/megatron/patches/args/checkpoint_path_patches.py +++ b/primus/backends/megatron/patches/args/checkpoint_path_patches.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -20,20 +20,39 @@ def patch_checkpoint_path(ctx: PatchContext): """ Configure checkpoint save path. - Sets args.save to /checkpoints and warns if user - provided a different path. + Behavior: + - args.save is None (YAML explicitly set `save: null`, or unset): + user has opted out of saving — leave args.save = None so that + Megatron's `checkpoint_and_decide_exit` (training.py:2340) skips + save_checkpoint_and_time(...). This is what allows runs configured + with `save: null` + `save_interval: ` to avoid hitting the + checkpoint save path entirely. + - args.save is non-None: + user has opted in to saving — route the destination to + /checkpoints, warning if it differs from what the user + specified. """ args = ctx.extra.get("backend_args", {}) primus_config = ctx.extra.get("primus_config", {}) - if args and primus_config.exp_root_path: - ckpt_path = os.path.abspath(os.path.join(primus_config.exp_root_path, "checkpoints")) + if not args or not primus_config.exp_root_path: + return - if hasattr(args, "save") and args.save is not None and args.save != ckpt_path: - log_rank_0( - f"[Patch:megatron.args.checkpoint_path][WARN] " - f"args.save is deprecated; overriding to: {ckpt_path}" - ) + # Respect explicit opt-out via `save: null` in YAML. + if not hasattr(args, "save") or args.save is None: + log_rank_0( + "[Patch:megatron.args.checkpoint_path] " + "args.save is None (opt-out); skipping save-path override." + ) + return - args.save = ckpt_path - log_rank_0(f"[Patch:megatron.args.checkpoint_path] save → {ckpt_path}") + ckpt_path = os.path.abspath(os.path.join(primus_config.exp_root_path, "checkpoints")) + + if args.save != ckpt_path: + log_rank_0( + f"[Patch:megatron.args.checkpoint_path][WARN] " + f"args.save is deprecated; overriding to: {ckpt_path}" + ) + + args.save = ckpt_path + log_rank_0(f"[Patch:megatron.args.checkpoint_path] save → {ckpt_path}") diff --git a/primus/backends/megatron/patches/args/hsdp_args_patches.py b/primus/backends/megatron/patches/args/hsdp_args_patches.py new file mode 100644 index 000000000..da1d76bc5 --- /dev/null +++ b/primus/backends/megatron/patches/args/hsdp_args_patches.py @@ -0,0 +1,60 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +from primus.core.patches import PatchContext, register_patch +from primus.modules.module_utils import log_kv_rank_0, log_rank_0 + + +@register_patch( + "megatron.args.hsdp", + backend="megatron", + phase="build_args", + description=( + "Map data_parallel_replicate_degree to Megatron's " + "num_distributed_optimizer_instances so that initialize_model_parallel " + "creates the HSDP process groups (intra-partial shard + inter-instance replicate)." + ), +) +def patch_hsdp_args(ctx: PatchContext): + """ + Propagate data_parallel_replicate_degree through to Megatron. + + data_parallel_replicate_degree is a Primus-only config key that is not + recognized by MegatronArgBuilder and would be silently dropped. This + patch reads it from module_config.params, injects it into backend_args, + and -- when > 1 -- maps it to num_distributed_optimizer_instances so that + Megatron's parallel_state.initialize_model_parallel creates the required + intra-partial (shard) and inter-instance (replicate) process groups. + """ + args = ctx.extra.get("backend_args") + module_config = ctx.extra.get("module_config") + + if not args or not module_config: + return + + replicate_degree = getattr(module_config.params, "data_parallel_replicate_degree", 1) + args.data_parallel_replicate_degree = replicate_degree + + if replicate_degree > 1: + if not getattr(args, "use_torch_fsdp2", False): + raise ValueError("data_parallel_replicate_degree > 1 requires use_torch_fsdp2=True") + + ckpt_fmt = getattr(args, "ckpt_format", "torch_dist") + if ckpt_fmt != "torch_dcp": + raise ValueError( + f"HSDP (data_parallel_replicate_degree={replicate_degree}) requires " + f"ckpt_format='torch_dcp', got '{ckpt_fmt}'. The torch_dist format's " + f"checkpoint offset logic assumes flat DP and is incompatible with a 2D mesh." + ) + + args.num_distributed_optimizer_instances = replicate_degree + log_rank_0( + f"[Patch:megatron.args.hsdp] HSDP enabled: " + f"data_parallel_replicate_degree={replicate_degree} " + f"-> num_distributed_optimizer_instances={replicate_degree}" + ) + + log_kv_rank_0("[Patch:megatron.args.hsdp] data_parallel_replicate_degree", str(replicate_degree)) diff --git a/primus/backends/megatron/patches/build_model_patches.py b/primus/backends/megatron/patches/build_model_patches.py index a96403184..601524c8b 100644 --- a/primus/backends/megatron/patches/build_model_patches.py +++ b/primus/backends/megatron/patches/build_model_patches.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. +# Copyright (c) 2026, Advanced Micro Devices, Inc. # # See LICENSE for license information. ############################################################################### @@ -10,14 +10,16 @@ This module contains patches that modify Megatron's model construction behavior to better integrate with Primus. -Current patch: +Current patches: - Disable the second DDP construction inside ``torch.cuda.stream()`` in ``megatron.training.training.get_model`` by temporarily replacing ``torch.cuda.stream`` with a no-op context manager while calling the original ``get_model``. + - Skip Float16Module wrapping when using FSDP2 FP32 param optimizer (model + parameters must stay in FP32 for FSDP2's MixedPrecisionPolicy). """ -from primus.core.patches import PatchContext, register_patch +from primus.core.patches import PatchContext, get_args, register_patch from primus.modules.module_utils import log_rank_0 @@ -29,6 +31,7 @@ "Monkey patch megatron.training.training.get_model to disable the " "second DDP construction inside torch.cuda.stream()." ), + condition=lambda ctx: not getattr(get_args(ctx), "disable_build_model_patches", False), ) def patch_megatron_get_model_disable_second_ddp(ctx: PatchContext) -> None: """ @@ -66,3 +69,59 @@ def _noop_stream(*_a, **_k): setattr(_patched_get_model, "_primus_disable_second_ddp", True) training.get_model = _patched_get_model log_rank_0("[Patch:megatron.get_model] Disabled second DDP via torch.cuda.stream no-op wrapper") + + +@register_patch( + "megatron.training.training.skip_float16_module", + backend="megatron", + phase="before_train", + description=( + "Skip Float16Module wrapping when using FSDP2 FP32 param optimizer. " + "Model parameters must stay FP32 for FSDP2's MixedPrecisionPolicy to handle casting." + ), + priority=40, + condition=lambda ctx: ( + getattr(get_args(ctx), "use_fsdp2_fp32_param_optimizer", False) + and getattr(get_args(ctx), "use_torch_fsdp2", False) + and getattr(get_args(ctx), "bf16", False) + ), +) +def patch_skip_float16_module(ctx: PatchContext) -> None: + """Replace Float16Module with a passthrough wrapper that preserves + the attributes Megatron's training loop expects (module, config, + vp_size, vp_stage, pg_collection) but does NOT convert parameters + to BF16/FP16. + + This is necessary because FSDP2's MixedPrecisionPolicy handles + the FP32->BF16 casting during forward/backward, and Float16Module + would convert parameters to BF16 before FSDP2 wrapping, defeating + the purpose of FP32 parameter storage. + """ + from megatron.core.transformer.module import MegatronModule + + class _FSDP2PassthroughModule(MegatronModule): + """Drop-in replacement for Float16Module that keeps params in FP32.""" + + def __init__(self, config, module): + super().__init__(config) + self.config = config + self.add_module("module", module) + self.vp_size = config.virtual_pipeline_model_parallel_size + self.vp_stage = getattr(module, "vp_stage", None) + self.pg_collection = getattr(module, "pg_collection", None) + + def set_input_tensor(self, input_tensor): + return self.module.set_input_tensor(input_tensor) + + def forward(self, *inputs, **kwargs): + return self.module(*inputs, **kwargs) + + import megatron.core.transformer.module as module_mod + import megatron.training.training as training_mod + + module_mod.Float16Module = _FSDP2PassthroughModule + training_mod.Float16Module = _FSDP2PassthroughModule + log_rank_0( + "[Patch:skip_float16_module] Replaced Float16Module with FP32 passthrough " + "(FSDP2 MixedPrecisionPolicy handles BF16 casting)" + ) diff --git a/primus/backends/megatron/patches/distributed_init_patches.py b/primus/backends/megatron/patches/distributed_init_patches.py new file mode 100644 index 000000000..7d7da8b92 --- /dev/null +++ b/primus/backends/megatron/patches/distributed_init_patches.py @@ -0,0 +1,76 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Distributed initialization patches (FSDP2 only). + +Patches Megatron's _initialize_distributed to pass device_id to +torch.distributed.init_process_group. Without device_id, RCCL guesses +the GPU-to-rank mapping, which causes deadlocks on MI355X after the +first FSDP2 iteration. + +This patch is gated on use_torch_fsdp2 because device_id triggers eager +RCCL communicator creation for the world PG and all ~26 Megatron sub-groups, +consuming ~768 MiB of additional GPU memory. Under DDP this extra allocation +pushes large models over the GPU memory limit. +""" + +import os + +import torch + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.modules.module_utils import log_rank_0 + + +@register_patch( + "megatron.distributed.init_process_group_device_id", + backend="megatron", + phase="before_train", + description=( + "Inject device_id into torch.distributed.init_process_group to " + "prevent RCCL device mapping deadlocks on MI355X (FSDP2 only)." + ), + priority=10, + condition=lambda ctx: getattr(get_args(ctx), "use_torch_fsdp2", False), +) +def patch_init_process_group_device_id(ctx: PatchContext): + """ + Wrap _initialize_distributed so that torch.distributed.init_process_group + receives an explicit device_id. + + Megatron computes device_id = torch.device(f'cuda:{args.local_rank}') but + never passes it to init_process_group. On MI355X with RCCL, the missing + device_id causes PyTorch to guess the GPU-to-rank mapping, leading to + deadlocks on the second FSDP2 iteration. + """ + import megatron.training.initialize as init_module + import torch.distributed as dist + + _orig_init_distributed = init_module._initialize_distributed + + def _patched_initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, store): + _orig_init_pg = dist.init_process_group + + def _init_pg_with_device_id(*args, **kwargs): + if "device_id" not in kwargs and torch.cuda.is_available(): + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + kwargs["device_id"] = torch.device(f"cuda:{local_rank}") + log_rank_0( + f"[Patch:device_id] Injected device_id=cuda:{local_rank} " f"into init_process_group" + ) + return _orig_init_pg(*args, **kwargs) + + dist.init_process_group = _init_pg_with_device_id + try: + _orig_init_distributed(get_embedding_ranks, get_position_embedding_ranks, store) + finally: + dist.init_process_group = _orig_init_pg + + init_module._initialize_distributed = _patched_initialize_distributed + log_rank_0( + "[Patch:device_id] Patched _initialize_distributed to inject " "device_id into init_process_group" + ) diff --git a/primus/backends/megatron/patches/env_patches.py b/primus/backends/megatron/patches/env_patches.py index 79df41f9b..c8c9ae4c0 100644 --- a/primus/backends/megatron/patches/env_patches.py +++ b/primus/backends/megatron/patches/env_patches.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -25,6 +25,11 @@ backend="megatron", phase="setup", description="Set CUDA_DEVICE_MAX_CONNECTIONS based on FSDP configuration", + condition=lambda ctx: not getattr( + getattr(ctx.extra.get("module_config"), "params", None), + "disable_env_patches", + False, + ), ) def set_cuda_device_max_connections(ctx: PatchContext): """ diff --git a/primus/backends/megatron/patches/mp_sync_skip_patches.py b/primus/backends/megatron/patches/mp_sync_skip_patches.py new file mode 100644 index 000000000..b7b9170c2 --- /dev/null +++ b/primus/backends/megatron/patches/mp_sync_skip_patches.py @@ -0,0 +1,63 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Skip redundant model-parallel synchronization when TP=1 and PP=1. + +When running pure data-parallel (TP=1, PP=1), the model-parallel process +group has size 1. The upstream Megatron ``train_step`` and ``training_log`` +unconditionally call ``logical_and_across_model_parallel_group`` and +``reduce_max_stat_across_model_parallel_group`` which each issue an +``all_reduce`` + ``.item()`` pair. With a size-1 group these are no-ops +that still force a GPU-to-CPU synchronization barrier, draining the +asynchronous GPU pipeline. + +This patch replaces those two module-level names inside +``megatron.training.training`` with lightweight pass-through functions +that preserve the return-type contract (``bool`` and ``float | None``) +without issuing any collective or ``.item()`` call. +""" + +import torch + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.modules.module_utils import log_rank_0 + + +def _is_pure_dp(ctx: PatchContext) -> bool: + """True when TP=1 and PP=1 -- the MP group has size 1.""" + args = get_args(ctx) + if args is None: + return False + tp = getattr(args, "tensor_model_parallel_size", 1) + pp = getattr(args, "pipeline_model_parallel_size", 1) + return tp == 1 and pp == 1 + + +@register_patch( + "megatron.training.skip_redundant_mp_sync", + backend="megatron", + phase="before_train", + description="Skip redundant model-parallel all-reduces when TP=1 and PP=1", + condition=_is_pure_dp, + priority=35, +) +def patch_skip_redundant_mp_sync(ctx: PatchContext): + import megatron.training.training as megatron_training + + def _passthrough_logical_and(val): + return val + + def _passthrough_reduce_max(val): + if val is None: + return None + if isinstance(val, torch.Tensor): + return val.item() + return val + + megatron_training.logical_and_across_model_parallel_group = _passthrough_logical_and + megatron_training.reduce_max_stat_across_model_parallel_group = _passthrough_reduce_max + log_rank_0("[Patch:skip_redundant_mp_sync] " "Replaced MP sync functions with passthrough (TP=1, PP=1)") diff --git a/primus/backends/megatron/training/global_vars.py b/primus/backends/megatron/training/global_vars.py index 5b2ae4825..4e0a52a86 100644 --- a/primus/backends/megatron/training/global_vars.py +++ b/primus/backends/megatron/training/global_vars.py @@ -152,7 +152,7 @@ def _ensure_var_is_initialized(var, name): def _ensure_var_is_not_initialized(var, name): - """Make sure the input variable is not None.""" + """Make sure the input variable is None (not yet initialized).""" assert var is None, "{} is already initialized.".format(name) diff --git a/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py index bf594e88e..2f46dfe7b 100644 --- a/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py +++ b/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py @@ -36,20 +36,21 @@ class MegatronBridgeBaseTrainer(BaseTrainer): - Handle Megatron-Bridge specific initialization and setup """ - def __init__(self, backend_args: Any): + def __init__(self, backend_args: Any = None, **kwargs): """ Initialize Megatron-Bridge base trainer. Args: backend_args: Megatron-Bridge configuration as SimpleNamespace (from MegatronBridgeArgBuilder) + **kwargs: Runtime context kwargs forwarded to BaseTrainer for filtering. """ log_rank_0("=" * 80) log_rank_0("Initializing MegatronBridgeBaseTrainer...") log_rank_0("=" * 80) # Initialize BaseTrainer - super().__init__(backend_args=backend_args) + super().__init__(backend_args=backend_args, **kwargs) set_primus_global_variables(self.backend_args) import primus.backends.megatron.patches # noqa: F401 diff --git a/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py index 1e95e0c3d..4ed170ec4 100644 --- a/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py +++ b/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py @@ -47,15 +47,16 @@ class MegatronBridgePosttrainTrainer(MegatronBridgeBaseTrainer): # Task type identifier for logging TASK_TYPE = "Post-training (SFT/Instruction Tuning)" - def __init__(self, backend_args: Any): + def __init__(self, backend_args: Any = None, **kwargs): """ Initialize Megatron-Bridge posttrain trainer. Args: backend_args: Megatron-Bridge argument namespace (from MegatronBridgeArgBuilder) + **kwargs: Runtime context kwargs forwarded to BaseTrainer for filtering. """ # Initialize MegatronBridgeBaseTrainer (which initializes BaseTrainer) - super().__init__(backend_args=backend_args) + super().__init__(backend_args=backend_args, **kwargs) def setup(self): """ diff --git a/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py index 9ee771f13..ac73d7362 100644 --- a/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py +++ b/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py @@ -39,14 +39,15 @@ class MegatronBridgePretrainTrainer(MegatronBridgeBaseTrainer): - Unified training workflow and patch management """ - def __init__(self, backend_args: Any): + def __init__(self, backend_args: Any = None, **kwargs): """ Initialize Megatron-Bridge pretrain trainer. Args: backend_args: Megatron-Bridge argument namespace (from MegatronBridgeArgBuilder) + **kwargs: Runtime context kwargs forwarded to BaseTrainer for filtering. """ - super().__init__(backend_args=backend_args) + super().__init__(backend_args=backend_args, **kwargs) def setup(self): """ diff --git a/primus/backends/torchtitan/torchtitan_pretrain_trainer.py b/primus/backends/torchtitan/torchtitan_pretrain_trainer.py index d2e88d095..8cc6d1c98 100644 --- a/primus/backends/torchtitan/torchtitan_pretrain_trainer.py +++ b/primus/backends/torchtitan/torchtitan_pretrain_trainer.py @@ -17,11 +17,14 @@ class TorchTitanPretrainTrainer(BaseTrainer): """Trainer class for TorchTitan pre-training.""" - def __init__(self, backend_args: Any): + def __init__(self, backend_args: Any = None, **kwargs): # Patch TorchTitan logger before any other initialization self._patch_torchtitan_logger() - super().__init__(backend_args=backend_args) + # The core runtime instantiates every trainer with BaseModule-style + # context kwargs (module_name, primus_config, module_rank, ...). Accept + # and forward them so BaseTrainer can filter them cooperatively. + super().__init__(backend_args=backend_args, **kwargs) self._trainer: Optional["Trainer"] = None # type: ignore[name-defined] def _patch_torchtitan_logger(self): diff --git a/primus/core/backend/backend_adapter.py b/primus/core/backend/backend_adapter.py index ce58505bc..738191b5f 100644 --- a/primus/core/backend/backend_adapter.py +++ b/primus/core/backend/backend_adapter.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -117,12 +117,18 @@ def _use_path(path: str, error_msg: str) -> str: # ============================================================================ @abstractmethod - def load_trainer_class(self, stage: str = "pretrain"): + def load_trainer_class(self, stage: str = "pretrain", trainer_class: Optional[str] = None): """ Return backend Trainer class registered in `BackendRegistry`. + Args: + stage: Training stage (e.g., "pretrain", "sft"). Defaults to "pretrain". + trainer_class: Optional specific trainer class name for dynamic loading. + If provided, this takes precedence over stage-based selection. + Default behavior: - - Lookup trainer via `BackendRegistry.get_trainer_class(self.framework, stage=stage)` + - If `trainer_class` is provided, attempt to dynamically load it + - Otherwise, lookup trainer via `BackendRegistry.get_trainer_class(self.framework, stage=stage)` Backends can override this method if they need special resolution rules. """ diff --git a/primus/core/backend/backend_registry.py b/primus/core/backend/backend_registry.py index 1f78d3f5b..e11001877 100644 --- a/primus/core/backend/backend_registry.py +++ b/primus/core/backend/backend_registry.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### diff --git a/primus/core/config/primus_config.py b/primus/core/config/primus_config.py index a9886fa1a..720a3316e 100644 --- a/primus/core/config/primus_config.py +++ b/primus/core/config/primus_config.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. +# Copyright (c) 2026, Advanced Micro Devices, Inc. # # See LICENSE for license information. ############################################################################### @@ -53,7 +53,7 @@ def _normalize_module_for_runtime(module_cfg: SimpleNamespace, module_name: str) if not getattr(normalized, "name", None): setattr(normalized, "name", module_name) - reserved_keys = {"name", "framework", "config", "model", "params"} + reserved_keys = {"name", "framework", "config", "model", "params", "trainer_class"} # Start from any existing params dict/namespace if provided. existing_params = getattr(normalized, "params", {}) params = _to_plain_dict(existing_params) @@ -121,11 +121,26 @@ def load_primus_config(config_path: Path, cli_args: Any | None = None) -> Simple cfg.platform = platform_config # Build modules list from legacy PrimusConfig.module_keys/get_module_config. - cfg.modules = [ + # + # `_normalize_module_for_runtime` deepcopies each module namespace before + # reshaping it, so `legacy_cfg` itself is left pristine. This is what makes + # it safe to expose `legacy_cfg` via `cfg._legacy` below: the legacy + # `PrimusConfig` that downstream consumers (e.g. `BaseModule`) read through + # `get_module_config(...)` keeps its original module shape. + modules: list[SimpleNamespace] = [ _normalize_module_for_runtime(legacy_cfg.get_module_config(module_name), module_name) for module_name in getattr(legacy_cfg, "module_keys", []) ] + cfg.modules = modules + + # Expose the underlying legacy PrimusConfig (additive, non-breaking) so the + # core runtime can reuse it instead of re-parsing the YAML a second time. + # BaseModule (inherited) still needs the PrimusConfig class interface + # (get_module_config method, set_global_variables, exp_* properties) that the + # SimpleNamespace above does not provide. `legacy_cfg` is pristine (see above). + cfg._legacy = legacy_cfg + return cfg diff --git a/primus/core/launcher/parser.py b/primus/core/launcher/parser.py index 753977f9c..c4a8d28dd 100644 --- a/primus/core/launcher/parser.py +++ b/primus/core/launcher/parser.py @@ -1,3 +1,6 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + from __future__ import annotations import argparse @@ -275,6 +278,15 @@ def parse_trainer_module(self, module_name: str): for key in ("config", "model"): yaml_utils.check_key_in_namespace(module, key) + # ---- Preserve module-level attributes before loading presets ---- + # Attributes like 'trainer_class' are at the module level in YAML + # and should be preserved even after loading config/model presets + preserved_attrs = {} + reserved_module_keys = {"name", "framework", "config", "model", "overrides", "params"} + for key, value in vars(module).items(): + if key not in reserved_module_keys and not key.startswith("_"): + preserved_attrs[key] = value + # ---- Load module config ---- model_format = self.get_model_format(framework) @@ -283,6 +295,11 @@ def parse_trainer_module(self, module_name: str): module_config.name = f"exp.modules.{module_name}.config" module_config.framework = framework + # Restore preserved module-level attributes + for key, value in preserved_attrs.items(): + if not hasattr(module_config, key): # Don't override if preset already has it + setattr(module_config, key, value) + # ---- Load model config ---- model_config_dict = PresetLoader.load(module.model, model_format, config_type="models") model_config = yaml_utils.dict_to_nested_namespace(model_config_dict) diff --git a/primus/core/runtime/runtime_state.py b/primus/core/runtime/runtime_state.py new file mode 100644 index 000000000..ffecb310c --- /dev/null +++ b/primus/core/runtime/runtime_state.py @@ -0,0 +1,34 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Runtime State Management for Training. + +This module provides a dedicated RuntimeState object for storing dynamic +runtime metrics that change during training (e.g., image dimensions, timesteps). +This is separate from backend_args (which is for configuration parameters). +""" + +from dataclasses import dataclass, field +from typing import Any, Dict + + +@dataclass +class RuntimeState: + """ + Dedicated runtime state object for per-iteration training metrics. + + This is separate from backend_args (which is for configuration) and + provides a clean place to store dynamic runtime metrics that change + during training (e.g., image dimensions, timesteps, etc.). + """ + + # Per-iteration metrics (updated each forward step) + last_metrics: Dict[str, Any] = field(default_factory=dict) + + def update_metrics(self, metrics: Dict[str, Any]) -> None: + """Update last_metrics with new metrics from forward step.""" + self.last_metrics.update(metrics) diff --git a/primus/core/runtime/train_runtime.py b/primus/core/runtime/train_runtime.py index fbcedd991..67d72f364 100644 --- a/primus/core/runtime/train_runtime.py +++ b/primus/core/runtime/train_runtime.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. +# Copyright (c) 2026, Advanced Micro Devices, Inc. # # Primus Runtime Orchestrator for Training # @@ -23,6 +23,7 @@ ) from primus.core.patches import run_patches from primus.core.runtime.logging import init_worker_logger +from primus.core.runtime.runtime_state import RuntimeState from primus.core.utils.arg_utils import parse_cli_overrides from primus.core.utils.env_setup import setup_training_env from primus.core.utils.yaml_utils import ( @@ -56,6 +57,7 @@ class TrainContext: trainer: Any = None backend_args: Any = None backend_version: Optional[str] = None + runtime_state: Optional[RuntimeState] = None # Distributed context rank: int = 0 @@ -135,7 +137,9 @@ def _get_backend_version(self) -> Optional[str]: self.ctx.backend_version = None return self.ctx.backend_version - def _run_phase_patches(self, phase: str, backend_args: Any = None) -> None: + def _run_phase_patches( + self, phase: str, backend_args: Any = None, runtime_state: Optional[RuntimeState] = None + ) -> None: """ Apply a patch phase in a single, runtime-owned place. @@ -156,6 +160,7 @@ def _run_phase_patches(self, phase: str, backend_args: Any = None) -> None: "backend_args": backend_args, "primus_config": self.ctx.primus_config, "module_config": self.ctx.module_config, + "runtime_state": runtime_state, }, ) @@ -174,6 +179,12 @@ def _initialize_configuration(self, module_name: str, overrides: Optional[List[s primus_cfg = load_primus_config(cfg_path, self.args) + # Reuse the legacy PrimusConfig that load_primus_config already parsed + # (exposed as `_legacy`) instead of re-parsing the YAML a second time. + # BaseModule needs the PrimusConfig interface (get_module_config / global + # vars); the SimpleNamespace `primus_cfg` drives the new core runtime. + primus_config_obj = primus_cfg._legacy + # Resolve module configuration via core helper. module_cfg = get_module_config(primus_cfg, module_name) available_modules = get_module_names(primus_cfg) or ["none"] @@ -188,11 +199,12 @@ def _initialize_configuration(self, module_name: str, overrides: Optional[List[s raise ValueError(f"[Primus:TrainRuntime] Module '{module_name}' missing 'framework'.") # Initialize TrainContext based on raw configuration (before CLI overrides). + # Use primus_config_obj (PrimusConfig) for BaseModule compatibility self.ctx = TrainContext( config_path=cfg_path, data_path=Path(getattr(self.args, "data_path", "./data")), module_name=module_name, - primus_config=primus_cfg, + primus_config=primus_config_obj, # Use PrimusConfig object, not SimpleNamespace module_config=module_cfg, framework=framework, ) @@ -247,6 +259,10 @@ def _initialize_adapter(self) -> None: assert self.ctx is not None, "TrainContext must be initialized before backend adapter." backend_path = getattr(self.args, "backend_path", None) + # CRITICAL: Set up backend path BEFORE importing backend module + # This ensures megatron is importable when patches are loaded + self._setup_backend_path_early(backend=self.ctx.framework, backend_path=backend_path) + adapter = BackendRegistry.get_adapter(backend=self.ctx.framework, backend_path=backend_path) assert ( @@ -258,6 +274,49 @@ def _initialize_adapter(self) -> None: self.ctx.adapter = adapter + def _setup_backend_path_early(self, backend: str, backend_path=None) -> None: + """Set up backend path before backend module is imported.""" + import os + import sys + from pathlib import Path + + # For Megatron backend, set up the path early + if backend == "megatron": + megatron_paths = [] + + # Method 1: From backend_path argument + if backend_path: + megatron_path = Path(backend_path) + if megatron_path.exists(): + megatron_paths.append(str(megatron_path)) + + # Method 2: From PRIMUS_PATH environment variable + primus_path = os.getenv("PRIMUS_PATH") + if primus_path: + megatron_path = Path(primus_path) / "third_party" / "Megatron-LM" + if megatron_path.exists(): + megatron_paths.append(str(megatron_path)) + + # Method 3: From current working directory + try: + cwd = Path.cwd() + if "Primus" in str(cwd): + primus_root = cwd + while primus_root.name != "Primus" and primus_root != primus_root.parent: + primus_root = primus_root.parent + if primus_root.name == "Primus": + megatron_path = primus_root / "third_party" / "Megatron-LM" + if megatron_path.exists(): + megatron_paths.append(str(megatron_path)) + except Exception: + pass + + # Add paths to sys.path if not already present + for path in megatron_paths: + if path not in sys.path: + sys.path.insert(0, path) + log_rank_0(f"[Primus:Runtime] Early setup: Added Megatron-LM to sys.path: {path}") + def _initialize_trainer(self) -> None: assert ( self.ctx is not None and self.ctx.adapter is not None @@ -273,8 +332,14 @@ def _initialize_trainer(self) -> None: backend_args = adapter.convert_config(module_config.params) self.ctx.backend_args = backend_args + # Create runtime state object (for dynamic per-iteration metrics) + # Initialize early so it's available for all patch phases + self.ctx.runtime_state = RuntimeState() + # Phase: build_args (after args creation, before trainer instantiation) - self._run_phase_patches(phase="build_args", backend_args=backend_args) + self._run_phase_patches( + phase="build_args", backend_args=backend_args, runtime_state=self.ctx.runtime_state + ) # Log final args after patches, then merge module_config.params into backend_args log_dict_aligned("Final backend args (after patches)", backend_args) @@ -290,18 +355,64 @@ def _initialize_trainer(self) -> None: primus_only_params = {key: params_dict[key] for key in sorted(primus_only_keys)} log_dict_aligned("Primus-specific parameters", primus_only_params) + # Extract trainer_class from module_config BEFORE merging params into backend_args + # (After merge, module_config.params becomes backend_args, and trainer_class might be lost) + trainer_class = None + + # First, check module_config directly (top-level attribute) + if hasattr(module_config, "trainer_class"): + trainer_class = module_config.trainer_class + # Second, check module_config.params (nested in params, BEFORE merge) + elif hasattr(module_config.params, "trainer_class"): + trainer_class = module_config.params.trainer_class + # Merge backend_args into params (backend_args overrides params) merge_namespace(backend_args, module_config.params, allow_override=False, excepts=[]) module_config.params = backend_args - # Load trainer class and instantiate + # Third, check backend_args (after merge, in case it was preserved) + if not trainer_class and hasattr(backend_args, "trainer_class"): + trainer_class = backend_args.trainer_class + + # Log summary of trainer_class extraction (keep one summary log for troubleshooting) + if trainer_class: + log_rank_0(f"[TrainRuntime] Using trainer_class: {trainer_class}") + else: + log_rank_0(f"[TrainRuntime] WARNING: trainer_class not found, will use default for stage") + + # Extract stage for fallback stage = getattr(module_config.params, "stage", "pretrain") or "pretrain" - TrainerClass = adapter.load_trainer_class(stage=stage) - trainer = TrainerClass(backend_args=backend_args) + + log_rank_0(f"[TrainRuntime] Loading trainer: stage={stage}, trainer_class={trainer_class}") + # Only forward trainer_class to adapters that support it (Megatron). Other + # backend adapters keep their stage-only signature, so passing trainer_class + # unconditionally would raise TypeError on the shared core-runtime path. + trainer_cls = adapter.load_trainer_class( + stage=stage, + **({"trainer_class": trainer_class} if trainer_class else {}), + ) + log_rank_0(f"[TrainRuntime] Loaded trainer class: {trainer_cls.__name__}") + + # Initialize trainer with both BaseModule and BaseTrainer arguments + # BaseModule needs: module_name, primus_config, module_rank, module_world_size, module_master_addr, module_master_port + # BaseTrainer needs: backend_args + # Note: BaseModule.__init__ expects keyword arguments, not positional + trainer = trainer_cls( + backend_args=backend_args, + module_name=self.ctx.module_name, + primus_config=self.ctx.primus_config, + module_rank=self.ctx.rank, + module_world_size=self.ctx.world_size, + module_master_addr=self.ctx.master_addr, + module_master_port=self.ctx.master_port, + ) assert trainer is not None, f"Failed to create trainer instance for framework '{self.ctx.framework}'." self.ctx.trainer = trainer + # Attach runtime_state to trainer instance for easy access in forward_step() + trainer.runtime_state = self.ctx.runtime_state + def _run_trainer_lifecycle(self) -> None: assert ( self.ctx is not None and self.ctx.trainer is not None @@ -320,18 +431,32 @@ def _log_step(step_name: str, func): trainer = self.ctx.trainer # 1) Optional setup phase - self._run_phase_patches(phase="setup", backend_args=self.ctx.backend_args) + self._run_phase_patches( + phase="setup", backend_args=self.ctx.backend_args, runtime_state=self.ctx.runtime_state + ) _log_step("Setup", trainer.setup) # 2) Initialize training components _log_step("Init", trainer.init) # 3) Execute training - self._run_phase_patches(phase="before_train", backend_args=self.ctx.backend_args) - _log_step("Training", trainer.train) + self._run_phase_patches( + phase="before_train", backend_args=self.ctx.backend_args, runtime_state=self.ctx.runtime_state + ) + import gc + + gc.disable() + log_rank_0("Python GC disabled for training (matching NeMo memory management)") + try: + _log_step("Training", trainer.train) + finally: + gc.enable() + log_rank_0("Python GC re-enabled after training") # 4) Cleanup and finalize - self._run_phase_patches(phase="after_train", backend_args=self.ctx.backend_args) + self._run_phase_patches( + phase="after_train", backend_args=self.ctx.backend_args, runtime_state=self.ctx.runtime_state + ) _log_step("Cleanup", trainer.cleanup) # --------------------------- Cleanup ---------------------------------- # diff --git a/primus/core/trainer/base_trainer.py b/primus/core/trainer/base_trainer.py index 107c44438..86e447166 100644 --- a/primus/core/trainer/base_trainer.py +++ b/primus/core/trainer/base_trainer.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -45,26 +45,49 @@ class BaseTrainer(ABC): MegatronPretrainTrainer, TorchtitanPretrainTrainer, etc. """ - def __init__(self, backend_args: Any): + def __init__(self, backend_args: Any = None, *args, **kwargs): """ Initialize base trainer. Args: backend_args: Backend-specific arguments (e.g., from MegatronArgBuilder) + *args, **kwargs: Additional arguments (filtered to prevent reaching object.__init__()) Note: Distributed environment and logging should be initialized globally before creating trainer instances. """ - self.backend_args = backend_args - - # Resolve distributed environment directly from torchrun-style env vars - dist_env = get_torchrun_env() - self.rank = dist_env["rank"] - self.world_size = dist_env["world_size"] - self.local_rank = dist_env["local_rank"] - self.master_addr = dist_env["master_addr"] - self.master_port = dist_env["master_port"] + # Filter backend_args from kwargs if not provided as positional/keyword argument + if backend_args is None: + backend_args = kwargs.pop("backend_args", None) + + try: + from abc import ABC + + ABC.__init__(self) + + self.backend_args = backend_args + + dist_env = get_torchrun_env() + self.rank = dist_env["rank"] + self.world_size = dist_env["world_size"] + self.local_rank = dist_env["local_rank"] + self.master_addr = dist_env["master_addr"] + self.master_port = dist_env["master_port"] + + # Cooperative multiple inheritance: pass kwargs to BaseModule if present in MRO, + # otherwise call super().__init__() with no args to avoid object.__init__() error. + from primus.modules.base_module import BaseModule + + if BaseModule in type(self).__mro__: + super().__init__(**kwargs) + else: + super().__init__() + except Exception: + import traceback + + traceback.print_exc() + raise @abstractmethod def setup(self): diff --git a/primus/modules/base_module.py b/primus/modules/base_module.py index 615ba8627..71fa19f02 100644 --- a/primus/modules/base_module.py +++ b/primus/modules/base_module.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -28,8 +28,8 @@ def __init__( module_world_size: int, module_master_addr: str = None, module_master_port: int = None, + **kwargs, # Accept extra kwargs for cooperative multiple inheritance ): - # module will be initialized by multiple worker processes self.module_name = module_name assert primus_config is not None @@ -71,6 +71,25 @@ def __init__( # setup logger for worker self.setup_worker_logger(module_rank, module_world_size) + # Call super() for cooperative multiple inheritance + # MRO for MegatronTrainer: MegatronTrainer -> BaseTrainer -> BaseModule -> ABC -> object + # + # Flow: + # 1. MegatronTrainer.__init__() calls super().__init__() -> BaseTrainer.__init__() + # 2. BaseTrainer.__init__() consumes backend_args, then calls super().__init__(**kwargs) -> BaseModule.__init__() + # Note: BaseTrainer passes kwargs containing module_name, primus_config, etc. to BaseModule + # 3. BaseModule.__init__() receives module_name, primus_config, etc. as positional/keyword args + # The kwargs dict may still contain these keys, but BaseModule has already consumed them as positional args + # 4. BaseModule.__init__() calls super().__init__() -> ABC.__init__() -> object.__init__() + # + # The issue: If kwargs still contains keys, passing them to super() will eventually reach object.__init__() + # which doesn't accept arguments. We must ensure kwargs is empty or only contains args for the next class. + # Since the next class after BaseModule is ABC (which doesn't accept kwargs), we should not pass kwargs. + + # Don't pass kwargs to super() - the next class in MRO is ABC, which doesn't need them + # All required args for BaseModule have been consumed as positional/keyword arguments above + super().__init__() + @abstractmethod def init(self, *args, **kwargs): raise NotImplementedError diff --git a/primus/modules/module_utils.py b/primus/modules/module_utils.py index 52503d985..11c9621d9 100644 --- a/primus/modules/module_utils.py +++ b/primus/modules/module_utils.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -73,7 +73,14 @@ def log_kv_rank_0(key, value): log_func(key, value, module_name, function_name, line) -def debug_rank_0(msg, *args, **kwargs): +def debug_rank_0(*args, **kwargs): + if not args: + return + sep = kwargs.pop("sep", " ") + kwargs.pop("end", None) + kwargs.pop("file", None) + kwargs.pop("flush", None) + msg = sep.join(str(a) for a in args) log_func = logger.debug_with_caller caller = inspect.stack()[1] @@ -86,7 +93,14 @@ def debug_rank_0(msg, *args, **kwargs): log_func(msg, module_name, function_name, line) -def debug_rank_all(msg, *args, **kwargs): +def debug_rank_all(*args, **kwargs): + if not args: + return + sep = kwargs.pop("sep", " ") + kwargs.pop("end", None) + kwargs.pop("file", None) + kwargs.pop("flush", None) + msg = sep.join(str(a) for a in args) log_func = logger.debug_with_caller caller = inspect.stack()[1] diff --git a/primus/modules/trainer/megatron/pre_trainer.py b/primus/modules/trainer/megatron/pre_trainer.py index d34788b04..4609f9c1a 100644 --- a/primus/modules/trainer/megatron/pre_trainer.py +++ b/primus/modules/trainer/megatron/pre_trainer.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -132,7 +132,13 @@ def __init__(self, *args, **kwargs): f"Megatron backend does not support unregistered config keys." ) - super().__init__(*args, **kwargs) + try: + super().__init__(*args, **kwargs) + except Exception: + import traceback + + traceback.print_exc() + raise def get_batch(self, data_iterator, vp_stage=None): """Generate a batch.""" diff --git a/primus/modules/trainer/megatron/trainer.py b/primus/modules/trainer/megatron/trainer.py index 429facb0c..130d0edca 100644 --- a/primus/modules/trainer/megatron/trainer.py +++ b/primus/modules/trainer/megatron/trainer.py @@ -1,6 +1,6 @@ ############################################################################### # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. -# Modification Copyright© 2025 Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -1072,6 +1072,14 @@ def setup_model_and_optimizer( checkpointing_context=self.checkpointing_context, skip_load_to_model_and_opt=HAVE_FSDP2 and args.use_torch_fsdp2, ) + if ( + HAVE_FSDP2 + and args.use_torch_fsdp2 + and optimizer is not None + and hasattr(optimizer, "finalize_dist_ckpt_load") + and args.iteration > 0 + ): + optimizer.finalize_dist_ckpt_load(args.iteration) timers("load-checkpoint").stop(barrier=True) timers.log(["load-checkpoint"]) one_logger and one_logger.log_metrics( @@ -1806,6 +1814,22 @@ def run_forward_backward_func(optimizer=None): timers("optimizer").stop() + if getattr(args, "use_fsdp2_fp8_all_gather", False): + from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + precompute_fp8_scales_for_fsdp, + ) + + precompute_fp8_scales_for_fsdp( + model[0], + stochastic_rounding=getattr(args, "fp8_all_gather_stochastic_rounding", False), + ) + + # FSDP2FP32Optimizer returns grad_norm as a GPU tensor to avoid a + # torch.compile graph break inside the compiled optimizer.step(). + # Materialize to float here, outside the compiled region. + if isinstance(grad_norm, torch.Tensor): + grad_norm = grad_norm.item() + # when freezing sub-models we may have a mixture of successful and unsucessful ranks, # so we must gather across mp ranks update_successful = logical_and_across_model_parallel_group(update_successful) diff --git a/primus/modules/trainer/megatron/utils.py b/primus/modules/trainer/megatron/utils.py index 75a30885a..d97842547 100644 --- a/primus/modules/trainer/megatron/utils.py +++ b/primus/modules/trainer/megatron/utils.py @@ -487,6 +487,30 @@ def _get_sync_free_moe_options(args) -> dict: return sync_free_moe[stage] +# FSDP2 custom optimizer selection flags. Each one monkeypatches +# get_megatron_optimizer at the same priority (50), so at most one may be set. +_FSDP2_OPTIMIZER_FLAGS = ( + "use_fsdp2_fp32_param_optimizer", + "use_fsdp2_bf16_master_weight_optimizer", +) + + +def validate_fsdp2_optimizer_exclusivity(args) -> None: + """Ensure at most one FSDP2 custom optimizer flag is enabled. + + Enabling more than one would silently let whichever optimizer patch applies + last win the monkeypatch, so raise ValueError to fail loudly at + arg-validation time (before training starts). + """ + enabled = [flag for flag in _FSDP2_OPTIMIZER_FLAGS if getattr(args, flag, False)] + if len(enabled) > 1: + raise ValueError( + "Conflicting FSDP2 optimizer selection: at most one of " + f"{list(_FSDP2_OPTIMIZER_FLAGS)} may be enabled, but got {enabled}. " + "Enable exactly one." + ) + + def validate_args_on_rocm(args): # Deterministic mode if args.deterministic_mode: @@ -509,6 +533,9 @@ def validate_args_on_rocm(args): assert not getattr( args, "use_turbo_parallel_linear", False ), "use_turbo_parallel_linear has been removed; please use use_turbo_gemm instead." + + validate_fsdp2_optimizer_exclusivity(args) + use_turbo_gemm = getattr(args, "use_turbo_gemm", False) # Turbo FP8 linear check if args.fp8 and use_turbo_gemm: diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..b4048ee90 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,137 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +import os +import sys +from pathlib import Path + + +def _warmup_aiter_nondeterministic_mha_bwd() -> None: + """Load aiter's nondeterministic flash-attention backward kernel before TE. + + aiter and TransformerEngine both statically bundle composable_kernel + (``ck_tile``). If ``transformer_engine`` is imported before aiter's + nondeterministic ``mha_bwd`` JIT module is first loaded, that kernel's + ``.so`` resolves its ck_tile host launch path against TE's (different) CK + copy and launches with an invalid grid/block config -> at runtime the + backward dies with:: + + HIP Function Failed (.../ck_tile/host/kernel_launch.hpp,110) + invalid configuration argument + + Triggering one tiny nondeterministic backward here (in the root conftest's + ``pytest_configure``, before any test module imports TE) loads that kernel + against aiter's own CK first, after which it stays correct for the rest of + the session. The deterministic kernel is unaffected, so we only need to warm + the nondeterministic variant. Best-effort: never let warmup break the suite. + """ + try: + import torch + + if not torch.cuda.is_available(): + return + + import math + + import primus_turbo.pytorch as pt + + q, k, v = ( + torch.randn(1, 8, 1, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + for _ in range(3) + ) + out = pt.ops.flash_attn_func( + q, + k, + v, + dropout_p=0.0, + softmax_scale=1.0 / math.sqrt(64), + causal=False, + window_size=(-1, -1), + deterministic=False, + return_lse=False, + ) + out.float().sum().backward() + torch.cuda.synchronize() + except Exception: + # Warmup is a best-effort mitigation; if aiter/CUDA is unavailable or + # the API shifts, fall through silently and let the affected tests + # surface their own errors. + pass + + +def _selection_includes_megatron(config) -> bool: + """True if the pytest selection could include a backends/megatron suite. + + Covers both the unit (``tests/unit_tests/backends/megatron``) and integration + (``tests/integration_tests/backends/megatron``) trees: the diffusion + integration tests run the same aiter attention backward and need the same + crash mitigations (deepbind hook + warmup). + + Path-level heuristic only (does not inspect -k / -m); the mitigations are + best-effort, so a conservative path check is sufficient. Fails OPEN: any + detection error returns True, because a missed mitigation degrades to a hard + mha_bwd crash, not just a slowdown. + + Uses ``config.args`` (populated post-parse, includes the default rootdir on + a bare run) rather than ``config.getoption("file_or_dir")`` (empty on a bare + run) -- do not "simplify" to the latter, it flips the empty-case semantics. + """ + try: + base = Path(__file__).resolve().parent + megatron_dirs = ( + base / "unit_tests" / "backends" / "megatron", + base / "integration_tests" / "backends" / "megatron", + ) + args = list(getattr(config, "args", []) or []) + if not args: + return True # whole-suite run (defensive; config.args is never empty) + for arg in args: + path = Path(str(arg).split("::", 1)[0]).resolve() + for megatron_dir in megatron_dirs: + # arg is the megatron dir, an ancestor of it, or a path inside it + if ( + path == megatron_dir + or megatron_dir.is_relative_to(path) + or path.is_relative_to(megatron_dir) + ): + return True + return False + except Exception: + # Fail open: never let a detection error suppress the crash-prevention mitigations. + return True + + +def pytest_configure(config): + project_root = Path(__file__).resolve().parent.parent + if str(project_root) not in sys.path: + sys.path.insert(0, str(project_root)) + + megatron_path = os.environ.get("MEGATRON_PATH") + if megatron_path is None or not os.path.exists(megatron_path): + megatron_path = project_root / "third_party" / "Megatron-LM" + if str(megatron_path) not in sys.path: + sys.path.append(str(megatron_path)) + + # Only needed for the megatron suite, so skip for unrelated selections + # (fail-open: detection errors still run it). + if _selection_includes_megatron(config): + # ORDER MATTERS: install the aiter RTLD_DEEPBIND import hook BEFORE the + # warmup below (or any test) first imports aiter's mha backward + # extension. The hook wraps importlib.import_module so the pinned + # aiter::mha_bwd binds its own ck_tile instead of TE's stale vendored + # libmha (ROCm/aiter#1332); once ``module_fmha_v3_bwd`` is already in + # sys.modules the hook is a no-op, so the warmup MUST NOT run first or + # the hd128 backward launches with an invalid grid config and aborts the + # whole process. Mirrors the production megatron.turbo.aiter_deepbind + # before_train patch (which has no such ordering hazard). + from tests.utils import install_aiter_deepbind_hook + + install_aiter_deepbind_hook() + # Must run before any test module imports transformer_engine. See the + # helper docstring for the aiter<->TE composable_kernel load-order issue. + _warmup_aiter_nondeterministic_mha_bwd() + else: + print("[conftest] skipping aiter mha_bwd mitigations (no megatron tests selected)") diff --git a/tests/unit_tests/backends/megatron/conftest.py b/tests/unit_tests/backends/megatron/conftest.py index 4c6599b92..8c807ac95 100644 --- a/tests/unit_tests/backends/megatron/conftest.py +++ b/tests/unit_tests/backends/megatron/conftest.py @@ -1,34 +1,63 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + """ -Pytest fixtures for Megatron backend tests. +Shared pytest fixtures for Megatron backend tests. -This conftest provides fixtures for tests that need Megatron parallel state -initialization, such as tests for Primus Turbo layers and other Megatron-Core -components that require distributed setup. +Provides reusable fixtures for parallel state initialization used across +optimizer, diffusion, and other Megatron test suites. """ import os -import socket +from types import SimpleNamespace import pytest import torch -import torch.distributed as dist + +from primus.core.utils import logger -def _find_free_port() -> int: - """Ask the kernel for a free TCP port on localhost. +@pytest.fixture(autouse=True, scope="session") +def setup_logger(): + """Initialize the Primus logger for megatron tests that use log_rank_0. - Binding to port 0 lets the OS pick a port that is guaranteed free at this - instant, which avoids the EADDRINUSE failures seen when guessing a random - port out of a fixed range that overlaps the ephemeral port range. + The Primus ``_logger`` global is ``None`` until configured, so tests that + log (e.g. test_warmup_convergence.py) raise ``AttributeError`` when run + without a prior logger-init. This session-scoped autouse fixture removes + that ordering dependency. Subdirectories may override it with an identically + named fixture (see diffusion/training/conftest.py). """ - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return s.getsockname()[1] + logger_cfg = logger.LoggerConfig( + exp_root_path=os.environ.get("UT_LOG_PATH", "ut_out"), + work_group="develop", + user_name="root", + exp_name="unittest", + module_name="UT-training", + file_sink_level="DEBUG", + stderr_sink_level="INFO", + node_ip="localhost", + rank=os.environ.get("RANK", 0), + world_size=os.environ.get("WORLD_SIZE", 1), + ) + logger.setup_logger(logger_cfg, is_head=False) + + +def _is_mxfp4_supported(): + if not torch.cuda.is_available(): + return False + try: + from primus_turbo.pytorch.core.low_precision import check_mxfp4_support + + supported, _ = check_mxfp4_support() + return supported + except ImportError: + return False + + +requires_mxfp4 = pytest.mark.skipif( + not _is_mxfp4_supported(), + reason="Requires gfx950+ (MI355X) for MXFP4 support", +) @pytest.fixture(scope="function") @@ -48,31 +77,37 @@ def init_parallel_state(): Example: @pytest.fixture(autouse=True) - def setup_parallel(self, init_parallel_state, monkeypatch): + def setup_parallel(self, init_parallel_state): '''Auto-use parallel state for this test class.''' pass """ - # Only import after sys.path is set up (which happens in pytest_configure) from megatron.core import parallel_state as ps # Initialize torch.distributed if not already initialized - if not dist.is_initialized(): - # Use a kernel-assigned free port to avoid conflicts with other running - # tests. Set MASTER_PORT explicitly (not setdefault) so a stale value - # left in the environment cannot diverge from the port we actually use. - port = _find_free_port() - os.environ["MASTER_ADDR"] = "127.0.0.1" - os.environ["MASTER_PORT"] = str(port) - - dist.init_process_group( - backend="nccl" if torch.cuda.is_available() else "gloo", - init_method=f"tcp://127.0.0.1:{port}", - world_size=1, - rank=0, - ) + if not torch.distributed.is_initialized(): + import os + import socket + + # Use OS-assigned ephemeral port to avoid TIME_WAIT conflicts + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", str(port)) - # Check if model parallel already initialized and destroy if so - # This ensures clean state for each test + try: + torch.distributed.init_process_group( + backend="nccl" if torch.cuda.is_available() else "gloo", + init_method=f"tcp://127.0.0.1:{port}", + world_size=1, + rank=0, + ) + except Exception as e: + pytest.skip(f"Could not initialize distributed: {e}") + + # Check if model parallel already initialized if ps.model_parallel_is_initialized(): ps.destroy_model_parallel() @@ -88,22 +123,29 @@ def setup_parallel(self, init_parallel_state, monkeypatch): # This is required for: # 1. ColumnParallelLinear and other tensor parallel layers that use get_cuda_rng_tracker().fork() # 2. CUDA graph support in layers (enable_cuda_graph=True) + from megatron.core.tensor_parallel import random as tp_random + if torch.cuda.is_available(): + # Initialize RNG tracker with CUDA graph support BEFORE calling model_parallel_cuda_manual_seed + # Try Transformer Engine RNG tracker first (best for CUDA graphs, used by Megatron's own tests) try: - from megatron.core.tensor_parallel import random as tp_random - - # Initialize RNG tracker with CUDA graph support BEFORE calling model_parallel_cuda_manual_seed - # Try Transformer Engine RNG tracker first (best for CUDA graphs, used by Megatron's own tests) - try: - tp_random.initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) - except (ImportError, AssertionError): - # Fallback to native PyTorch CUDA graph RNG support if TE not available - tp_random.initialize_rng_tracker(use_cudagraphable_rng=True, force_reset=True) - - tp_random.model_parallel_cuda_manual_seed(42) - except ImportError: - # RNG tracker initialization is optional - skip if not available - pass + tp_random.initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True) + except (ImportError, AssertionError): + # Fallback to native PyTorch CUDA graph RNG support if TE not available + tp_random.initialize_rng_tracker(use_cudagraphable_rng=True, force_reset=True) + + tp_random.model_parallel_cuda_manual_seed(42) + + # Initialize Megatron global args with minimal defaults required by + # PrimusTurboLocalAttention and PrimusTorchFullyShardedDataParallel + from megatron.training.global_vars import set_args + + set_args( + SimpleNamespace( + enable_turbo_attention_float8=False, + data_parallel_replicate_degree=1, + ) + ) yield @@ -111,7 +153,12 @@ def setup_parallel(self, init_parallel_state, monkeypatch): if ps.model_parallel_is_initialized(): ps.destroy_model_parallel() + # Reset Megatron global args + import megatron.training.global_vars as gvars + + gvars._GLOBAL_ARGS = None + # Cleanup torch.distributed for single-process tests # (In multi-process torchrun tests, the process group persists across tests) - if dist.is_initialized(): - dist.destroy_process_group() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_megatron_pretrain_trainer_overridable_methods.py b/tests/unit_tests/backends/megatron/diffusion/training/test_megatron_pretrain_trainer_overridable_methods.py new file mode 100644 index 000000000..5c2c45e64 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_megatron_pretrain_trainer_overridable_methods.py @@ -0,0 +1,115 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for overridable methods in MegatronPretrainTrainer. + +Tests verify that get_forward_step() and get_datasets_provider() methods +can be overridden by subclasses while maintaining backward compatibility. +""" + +import sys +import types +from types import SimpleNamespace +from typing import Any, List, Tuple + +import pytest +import torch + +from primus.backends.megatron.megatron_pretrain_trainer import MegatronPretrainTrainer + +# train() imports the real megatron.core/training stack (transitively primus_turbo/TE), which +# initializes CUDA at import time. Gate the two tests that call train() so they run only on GPU. +requires_gpu = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="train() imports real megatron.core/training which initialize CUDA at import", +) + + +def _build_trainer(monkeypatch: pytest.MonkeyPatch) -> MegatronPretrainTrainer: + """Helper to build MegatronPretrainTrainer with a stubbed MegatronBaseTrainer.""" + + # Stub out MegatronBaseTrainer.__init__ to avoid real Megatron imports/patching. + def dummy_init(self, backend_args: Any = None): + self.backend_args = backend_args + + monkeypatch.setattr( + "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.__init__", + dummy_init, + ) + + # Silence logging from the trainer module. + monkeypatch.setattr( + "primus.backends.megatron.megatron_pretrain_trainer.log_rank_0", + lambda *args, **kwargs: None, + ) + + backend_args = SimpleNamespace() + + return MegatronPretrainTrainer(backend_args=backend_args) + + +class TestMegatronPretrainTrainerOverridableMethods: + """Tests for overridable methods in MegatronPretrainTrainer.""" + + @requires_gpu + def test_backward_compatibility_gpt_training_still_works(self, monkeypatch: pytest.MonkeyPatch): + """Test that existing GPT training still works (backward compatibility).""" + trainer = _build_trainer(monkeypatch) + + calls: List[Tuple[tuple, dict]] = [] + + # Use the real megatron modules that train() imports; patch only the integration seams. + import megatron.core.pipeline_parallel as mpp + import megatron.training + import megatron.training.inprocess_restart as inprocess_restart + import megatron.training.training as mt_training + from megatron.core.enums import ModelType + + def fake_pretrain(*args, **kwargs): + raise AssertionError("fake_pretrain was called directly; expected wrapped_pretrain to be used") + + def wrapped_pretrain(*args, store=None, **kwargs): + calls.append((args, {"store": store, **kwargs})) + + def fake_maybe_wrap(fn): + assert fn is fake_pretrain + return wrapped_pretrain, "STORE" + + monkeypatch.setattr(megatron.training, "pretrain", fake_pretrain) + monkeypatch.setattr(inprocess_restart, "maybe_wrap_for_inprocess_restart", fake_maybe_wrap) + + # train() reassigns get_forward_backward_func on both real modules; snapshot both through + # monkeypatch so they are restored and the mutation does not leak into later tests. + monkeypatch.setattr(mpp, "get_forward_backward_func", mpp.get_forward_backward_func) + monkeypatch.setattr(mt_training, "get_forward_backward_func", mt_training.get_forward_backward_func) + + model_provider = object() + monkeypatch.setattr( + "primus.core.utils.import_utils.get_model_provider", + lambda *args, **kwargs: model_provider, + ) + + # pretrain_gpt is a Megatron example script, not importable here -> mock it via sys.modules. + train_valid_test_datasets_provider = SimpleNamespace(is_distributed=False) + pretrain_gpt_mod = types.SimpleNamespace( + forward_step="FORWARD_STEP", + train_valid_test_datasets_provider=train_valid_test_datasets_provider, + ) + monkeypatch.setitem(sys.modules, "pretrain_gpt", pretrain_gpt_mod) + + # Execute training + trainer.train() + + # Verify backward compatibility: should work exactly as before + assert train_valid_test_datasets_provider.is_distributed is True + assert len(calls) == 1 + (args, kwargs) = calls[0] + + # Positional arguments should match expected GPT training pattern + assert args[0] is train_valid_test_datasets_provider + assert args[1] is model_provider + assert args[2] is ModelType.encoder_or_decoder + assert args[3] == "FORWARD_STEP" + + assert kwargs == {"store": "STORE"} diff --git a/tests/unit_tests/backends/megatron/patches/test_patch_guard.py b/tests/unit_tests/backends/megatron/patches/test_patch_guard.py new file mode 100644 index 000000000..97737f16c --- /dev/null +++ b/tests/unit_tests/backends/megatron/patches/test_patch_guard.py @@ -0,0 +1,75 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for the feature-owned patch idempotency guard. + +These validate that: + - a guarded "wrapping" patch applies exactly once even if its handler runs + twice (e.g. before_train re-invoked), and + - distinct patches still compose (each wraps once), which is required for the + stacked train_step wrappers (FP8 cache, delayed scaling, wall-clock timer). +""" + +from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched + + +class _FakeModule: + """Stand-in for ``megatron.training.training`` (a module object).""" + + +def _apply_wrapping_patch(module, key): + """Mimic the guarded train_step wrap pattern used by the FP8/timer patches.""" + if is_patched(module, key): + return + original = module.train_step + + def wrapped(*args, **kwargs): + module.call_log.append(key) + return original(*args, **kwargs) + + module.train_step = wrapped + mark_patched(module, key) + + +def test_is_patched_false_before_mark(): + module = _FakeModule() + assert not is_patched(module, "k") + + +def test_mark_then_is_patched(): + module = _FakeModule() + mark_patched(module, "k") + assert is_patched(module, "k") + assert not is_patched(module, "other") + + +def test_guarded_wrap_applies_once_on_reapply(): + module = _FakeModule() + module.call_log = [] + module.train_step = lambda: "base" + + # Apply the same patch twice (simulating before_train running twice). + _apply_wrapping_patch(module, "megatron.train_step.demo") + _apply_wrapping_patch(module, "megatron.train_step.demo") + + module.train_step() + # Wrapped exactly once -> the side effect is recorded once per call. + assert module.call_log == ["megatron.train_step.demo"] + + +def test_distinct_patches_still_compose(): + module = _FakeModule() + module.call_log = [] + module.train_step = lambda: "base" + + _apply_wrapping_patch(module, "patch.a") + _apply_wrapping_patch(module, "patch.b") + # Re-run both: idempotent, so still only one layer each. + _apply_wrapping_patch(module, "patch.a") + _apply_wrapping_patch(module, "patch.b") + + module.train_step() + assert sorted(module.call_log) == ["patch.a", "patch.b"] diff --git a/tests/unit_tests/backends/megatron/test_megatron_adapter.py b/tests/unit_tests/backends/megatron/test_megatron_adapter.py index 66d610f52..74e4a098c 100644 --- a/tests/unit_tests/backends/megatron/test_megatron_adapter.py +++ b/tests/unit_tests/backends/megatron/test_megatron_adapter.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -123,63 +123,6 @@ def test_detect_version_not_found_raises_error(self, tmp_path, monkeypatch): assert "Cannot locate" in str(exc_info.value) -class TestMegatronAdapterConfigConversion: - """Test configuration conversion from Primus to Megatron.""" - - @patch("primus.backends.megatron.megatron_adapter.MegatronArgBuilder") - @patch("primus.backends.megatron.megatron_adapter.log_rank_0") - def test_convert_config_basic(self, mock_log, mock_builder_class): - """Test basic config conversion workflow.""" - # Setup mock builder - mock_builder = Mock() - mock_builder_class.return_value = mock_builder - - # Mock finalize to return SimpleNamespace with args - mock_args = SimpleNamespace( - micro_batch_size=4, - global_batch_size=32, - seq_length=2048, - ) - mock_builder.finalize.return_value = mock_args - - # Create mock params - mock_params = { - "micro_batch_size": 4, - "global_batch_size": 32, - "seq_length": 2048, - } - - # Test conversion - adapter = MegatronAdapter() - result = adapter.convert_config(mock_params) - - # Verify builder was called correctly - mock_builder.update.assert_called_once_with(mock_params) - mock_builder.finalize.assert_called_once() - - # Verify result - assert result == mock_args - assert result.micro_batch_size == 4 - assert result.global_batch_size == 32 - assert result.seq_length == 2048 - - @patch("primus.backends.megatron.megatron_adapter.MegatronArgBuilder") - @patch("primus.backends.megatron.megatron_adapter.log_rank_0") - def test_convert_config_empty_params(self, mock_log, mock_builder_class): - """Test config conversion with empty params dict.""" - mock_builder = Mock() - mock_builder_class.return_value = mock_builder - mock_builder.finalize.return_value = SimpleNamespace() - - mock_params = {} - - adapter = MegatronAdapter() - result = adapter.convert_config(mock_params) - - mock_builder.update.assert_called_once_with(mock_params) - assert isinstance(result, SimpleNamespace) - - class TestMegatronAdapterTrainerLoading: """Test trainer class loading.""" @@ -214,31 +157,127 @@ def test_load_trainer_class_invalid_stage(self): assert "backend trainer not registered" in str(exc_info.value) -class TestMegatronAdapterBackendPreparation: - """Test backend preparation workflow.""" +class TestMegatronAdapterDynamicTrainerLoading: + """Test dynamic trainer class loading by name.""" - @patch("primus.modules.module_utils.log_rank_0") - def test_prepare_backend_success(self, mock_log): - """Test successful backend preparation.""" + @patch("primus.backends.megatron.megatron_adapter.log_rank_0") + def test_load_trainer_class_by_name_takes_precedence_over_stage(self, mock_log): + """Test that trainer_class parameter takes precedence over stage.""" adapter = MegatronAdapter() - mock_config = Mock() - # Should not raise - prepare_backend is inherited from BackendAdapter - adapter.prepare_backend(mock_config) + # Even with invalid stage, trainer_class should work + result = adapter.load_trainer_class(stage="invalid_stage", trainer_class="MegatronPretrainTrainer") + + from primus.backends.megatron.megatron_pretrain_trainer import ( + MegatronPretrainTrainer, + ) + + assert result == MegatronPretrainTrainer + + @patch("primus.backends.megatron.megatron_adapter.log_rank_0") + def test_load_trainer_class_fallback_tries_multiple_paths(self, mock_log): + """Test fallback tries multiple import paths in order for unregistered trainers.""" + import importlib + + # Use a trainer name NOT in MEGATRON_TRAINERS to exercise the fallback + mock_trainer_class = type("CustomExperimentalTrainer", (), {}) + mock_module = Mock() + mock_module.CustomExperimentalTrainer = mock_trainer_class + + import_calls = [] + + def side_effect(module_name, *args, **kwargs): + import_calls.append(module_name) + if len(import_calls) <= 2: + raise ImportError(f"Path {len(import_calls)} failed") + return mock_module + + with patch.object(importlib, "import_module", side_effect=side_effect): + adapter = MegatronAdapter() + result = adapter.load_trainer_class(trainer_class="CustomExperimentalTrainer") + + assert result == mock_trainer_class + assert len(import_calls) == 3 + assert "primus.backends.megatron.customexperimentaltrainer" in import_calls[0].lower() + @patch("primus.backends.megatron.megatron_adapter.log_rank_0") + def test_load_trainer_class_fallback_all_paths_fail(self, mock_log): + """Test error when all fallback paths fail.""" + import importlib + + # All import attempts fail + def side_effect(module_name, *args, **kwargs): + raise ImportError("Module not found") + + with patch.object(importlib, "import_module", side_effect=side_effect): + adapter = MegatronAdapter() + + with pytest.raises(ValueError) as exc_info: + adapter.load_trainer_class(trainer_class="NonExistentTrainer") -class TestMegatronAdapterInitialization: - """Test MegatronAdapter initialization.""" + error_msg = str(exc_info.value) + # Verify exact error message format from implementation + assert "Trainer class 'NonExistentTrainer' not found" in error_msg + assert "Available trainers:" in error_msg + assert "Hint:" in error_msg + assert "MEGATRON_TRAINERS" in error_msg or "standard locations" in error_msg - def test_initialization_default(self): - """Test adapter initialization with default framework name.""" + @patch("primus.backends.megatron.megatron_adapter.log_rank_0") + def test_load_trainer_class_falsy_values_fall_back_to_stage(self, mock_log): + """Test that falsy trainer_class values (None, empty string) fall back to stage-based selection.""" adapter = MegatronAdapter() - assert adapter.framework == "megatron" - def test_initialization_custom_framework(self): - """Test adapter initialization with custom framework name.""" - adapter = MegatronAdapter(framework="custom_megatron") - assert adapter.framework == "custom_megatron" + from primus.backends.megatron.megatron_pretrain_trainer import ( + MegatronPretrainTrainer, + ) + + # Test None falls back to stage + result = adapter.load_trainer_class(stage="pretrain", trainer_class=None) + assert result == MegatronPretrainTrainer + + # Test empty string falls back to stage (falsy) + result = adapter.load_trainer_class(stage="pretrain", trainer_class="") + assert result == MegatronPretrainTrainer + + def test_load_trainer_class_registry_import_error_provides_context(self): + """Test that import errors from registry provide helpful context.""" + import importlib + + def side_effect(module_name, *args, **kwargs): + raise ImportError("Module not found") + + with patch.object(importlib, "import_module", side_effect=side_effect): + adapter = MegatronAdapter() + + with pytest.raises(ImportError) as exc_info: + adapter.load_trainer_class(trainer_class="MegatronPretrainTrainer") + + error_msg = str(exc_info.value) + # Verify exact error message format from implementation + assert "Failed to load trainer class 'MegatronPretrainTrainer'" in error_msg + assert "Hint:" in error_msg + assert "Check that the module exists" in error_msg + + def test_load_trainer_class_registry_attribute_error_provides_context(self): + """Test that attribute errors from registry provide helpful context.""" + import importlib + + # Module imports but class doesn't exist + mock_module = Mock() + del mock_module.MegatronPretrainTrainer # Ensure attribute doesn't exist + + def side_effect(module_name, *args, **kwargs): + return mock_module + + with patch.object(importlib, "import_module", side_effect=side_effect): + adapter = MegatronAdapter() + + with pytest.raises(ImportError) as exc_info: + adapter.load_trainer_class(trainer_class="MegatronPretrainTrainer") + + error_msg = str(exc_info.value) + assert "Failed to load trainer class" in error_msg + assert "Hint:" in error_msg class TestMegatronAdapterIntegration: diff --git a/tests/unit_tests/backends/megatron/test_megatron_base_trainer.py b/tests/unit_tests/backends/megatron/test_megatron_base_trainer.py new file mode 100644 index 000000000..b1d097d2b --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_megatron_base_trainer.py @@ -0,0 +1,213 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for MegatronBaseTrainer. + +Tests path resolution, parse_args patching, and setup orchestration. +""" + +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from primus.backends.megatron.megatron_base_trainer import MegatronBaseTrainer + + +def _build_trainer(monkeypatch: pytest.MonkeyPatch, backend_args=None): + """Helper to build MegatronBaseTrainer with stubbed dependencies.""" + + # Create a concrete subclass for testing (MegatronBaseTrainer is abstract) + class ConcreteMegatronBaseTrainer(MegatronBaseTrainer): + def train(self): + pass + + # Stub out BaseTrainer.__init__ to avoid real distributed env setup + def dummy_base_init(self, backend_args=None, *args, **kwargs): + self.backend_args = backend_args + self.rank = 0 + self.world_size = 1 + self.local_rank = 0 + self.master_addr = "localhost" + self.master_port = 12345 + + monkeypatch.setattr( + "primus.core.trainer.base_trainer.BaseTrainer.__init__", + dummy_base_init, + ) + + # Silence logging + monkeypatch.setattr( + "primus.backends.megatron.megatron_base_trainer.log_rank_0", + lambda *args, **kwargs: None, + ) + + if backend_args is None: + backend_args = SimpleNamespace() + + return ConcreteMegatronBaseTrainer(backend_args=backend_args) + + +class TestMegatronBaseTrainer: + """Tests for MegatronBaseTrainer setup and path resolution.""" + + def test_ensure_megatron_path_from_primus_path_env(self, monkeypatch: pytest.MonkeyPatch): + """Test path resolution from PRIMUS_PATH environment variable.""" + trainer = _build_trainer(monkeypatch) + + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + primus_path = Path(tmpdir) / "Primus" + megatron_path = primus_path / "third_party" / "Megatron-LM" + megatron_path.mkdir(parents=True) + + monkeypatch.setenv("PRIMUS_PATH", str(primus_path)) + + # Force `import megatron` to raise ImportError by caching None in + # sys.modules. Unlike mocking builtins.__import__, this leaves + # Python's import machinery intact for all other modules. + saved_megatron = sys.modules.get("megatron", _SENTINEL := object()) + sys.modules["megatron"] = None + # Strip paths that either contain "megatron" in the name or have a + # megatron/ subdirectory, so Method 4 doesn't short-circuit. + original_path = sys.path[:] + sys.path[:] = [ + p for p in sys.path if "megatron" not in p.lower() and not (Path(p) / "megatron").is_dir() + ] + try: + trainer._ensure_megatron_path() + assert str(megatron_path) in sys.path + finally: + sys.path[:] = original_path + if saved_megatron is _SENTINEL: + sys.modules.pop("megatron", None) + else: + sys.modules["megatron"] = saved_megatron + + def test_ensure_megatron_path_from_cwd(self, monkeypatch: pytest.MonkeyPatch): + """Test path resolution from current working directory.""" + trainer = _build_trainer(monkeypatch) + + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + primus_path = Path(tmpdir) / "Primus" + megatron_path = primus_path / "third_party" / "Megatron-LM" + megatron_path.mkdir(parents=True) + + monkeypatch.delenv("PRIMUS_PATH", raising=False) + + def mock_cwd(): + return primus_path / "some" / "subdirectory" + + monkeypatch.setattr(Path, "cwd", staticmethod(mock_cwd)) + + saved_megatron = sys.modules.get("megatron", _SENTINEL := object()) + sys.modules["megatron"] = None + original_path = sys.path[:] + sys.path[:] = [ + p for p in sys.path if "megatron" not in p.lower() and not (Path(p) / "megatron").is_dir() + ] + try: + trainer._ensure_megatron_path() + assert str(megatron_path) in sys.path + finally: + sys.path[:] = original_path + if saved_megatron is _SENTINEL: + sys.modules.pop("megatron", None) + else: + sys.modules["megatron"] = saved_megatron + + def test_ensure_megatron_path_from_file_location(self, monkeypatch: pytest.MonkeyPatch): + """Test path resolution from current file location.""" + trainer = _build_trainer(monkeypatch) + + # Unset PRIMUS_PATH and mock cwd to not contain Primus + monkeypatch.delenv("PRIMUS_PATH", raising=False) + monkeypatch.setattr(Path, "cwd", staticmethod(lambda: Path("/some/other/path"))) + + # Mock __file__ to point to a Primus subdirectory + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + primus_path = Path(tmpdir) / "Primus" + megatron_path = primus_path / "third_party" / "Megatron-LM" + megatron_path.mkdir(parents=True) + + # Mock __file__ to be in primus/backends/megatron/ + fake_file = primus_path / "primus" / "backends" / "megatron" / "megatron_base_trainer.py" + fake_file.parent.mkdir(parents=True) + fake_file.touch() + + def mock_resolve(self): + return fake_file + + monkeypatch.setattr(Path, "resolve", mock_resolve) + + saved_megatron = sys.modules.get("megatron", _SENTINEL := object()) + sys.modules["megatron"] = None + original_path = sys.path[:] + # Isolate Method 3 (file-location) by neutralizing Method 4, which + # returns early if any sys.path entry already contains a real + # ``megatron`` package dir. A name-only filter misses a pip-installed + # megatron-core under site-packages (no "megatron" in the path name), + # so also drop entries that actually contain a ``megatron`` dir. + sys.path[:] = [ + p for p in sys.path if "megatron" not in p.lower() and not (Path(p) / "megatron").is_dir() + ] + try: + trainer._ensure_megatron_path() + assert str(megatron_path) in sys.path + finally: + sys.path[:] = original_path + if saved_megatron is _SENTINEL: + sys.modules.pop("megatron", None) + else: + sys.modules["megatron"] = saved_megatron + + def test_ensure_megatron_path_already_importable(self, monkeypatch: pytest.MonkeyPatch): + """Test that path setup is skipped if megatron is already importable.""" + trainer = _build_trainer(monkeypatch) + + megatron_mod = types.ModuleType("megatron") + monkeypatch.setitem(sys.modules, "megatron", megatron_mod) + + path_before = sys.path[:] + trainer._ensure_megatron_path() + + assert sys.path == path_before + + def test_patch_parse_args(self, monkeypatch: pytest.MonkeyPatch): + """Test that parse_args is patched in both locations.""" + trainer = _build_trainer(monkeypatch) + + import megatron.training.arguments as megatron_args_mod + import megatron.training.initialize as megatron_init_mod + + original_parse_args_args = megatron_args_mod.parse_args + original_parse_args_init = megatron_init_mod.parse_args + + backend_args = SimpleNamespace(test_param="value") + trainer.backend_args = backend_args + + try: + trainer._patch_parse_args() + + assert megatron_args_mod.parse_args is not original_parse_args_args + assert megatron_init_mod.parse_args is not original_parse_args_init + + result_args = megatron_args_mod.parse_args() + result_init = megatron_init_mod.parse_args() + + assert result_args is backend_args + assert result_init is backend_args + finally: + megatron_args_mod.parse_args = original_parse_args_args + megatron_init_mod.parse_args = original_parse_args_init diff --git a/tests/unit_tests/core/backend/test_backend_adapter.py b/tests/unit_tests/core/backend/test_backend_adapter.py index ccc4d693c..1a4a18c9c 100644 --- a/tests/unit_tests/core/backend/test_backend_adapter.py +++ b/tests/unit_tests/core/backend/test_backend_adapter.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -70,7 +70,7 @@ def convert_config(self, params: Any) -> Any: model_name=params.get("model") if isinstance(params, dict) else getattr(params, "model", None), ) - def load_trainer_class(self, stage: str = "pretrain"): + def load_trainer_class(self, stage: str = "pretrain", trainer_class=None): self.load_trainer_calls += 1 return DummyTrainer @@ -99,21 +99,6 @@ def primus_config(): return SimpleNamespace(exp_name="unit-test-exp") -def test_create_trainer_orchestrates_flow(monkeypatch, primus_config, module_config): - adapter = DummyBackendAdapter(framework="megatron", version="1.2.3") - - # Abstract methods were called exactly once - adapter.prepare_backend(module_config) - adapter.convert_config(module_config.params) - adapter.load_trainer_class(stage="pretrain") - adapter.detect_backend_version() - - assert adapter.prepare_calls == [module_config] - assert adapter.convert_calls == [module_config.params] - assert adapter.load_trainer_calls == 1 - assert adapter.detect_version_calls == 1 - - def test_adapter_setup_backend_path_with_explicit_path(tmp_path, monkeypatch): adapter = DummyBackendAdapter(framework="test_backend") backend_dir = tmp_path / "explicit_backend" diff --git a/tests/unit_tests/core/backend/test_backend_registry.py b/tests/unit_tests/core/backend/test_backend_registry.py index 42ea02a15..7b08d553f 100644 --- a/tests/unit_tests/core/backend/test_backend_registry.py +++ b/tests/unit_tests/core/backend/test_backend_registry.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -15,11 +15,6 @@ import primus.core.backend.backend_registry as registry_module from primus.core.backend.backend_adapter import BackendAdapter -_SUPPORTS_TRAINER_CLASS_REGISTRY = all( - hasattr(registry_module.BackendRegistry, attr) - for attr in ("_trainer_classes", "register_trainer_class", "get_trainer_class", "has_trainer_class") -) - class MockAdapter(BackendAdapter): """Mock adapter for testing.""" @@ -45,12 +40,6 @@ def setup_method(self): # Save original state self._original_adapters = registry_module.BackendRegistry._adapters.copy() registry_module.BackendRegistry._adapters.clear() - if hasattr(registry_module.BackendRegistry, "_trainer_classes"): - self._original_trainer_classes = registry_module.BackendRegistry._trainer_classes.copy() - registry_module.BackendRegistry._trainer_classes.clear() - else: - self._original_trainer_classes = None - # Silence logging dependencies (logger may not be initialized in tests) self._orig_log_rank_0 = registry_module.log_rank_0 registry_module.log_rank_0 = lambda *args, **kwargs: None @@ -58,11 +47,6 @@ def setup_method(self): def teardown_method(self): """Restore registry after each test.""" registry_module.BackendRegistry._adapters = self._original_adapters - if ( - hasattr(registry_module.BackendRegistry, "_trainer_classes") - and self._original_trainer_classes is not None - ): - registry_module.BackendRegistry._trainer_classes = self._original_trainer_classes registry_module.log_rank_0 = self._orig_log_rank_0 def test_get_adapter_not_found_helpful_error(self): @@ -76,13 +60,6 @@ def test_get_adapter_not_found_helpful_error(self): with pytest.raises(ModuleNotFoundError, match="No module named 'primus.backends.non_existent'"): registry_module.BackendRegistry.get_adapter("non_existent") - def test_get_adapter_empty_registry_error(self): - """Test error message when no backends are registered.""" - # get_adapter first calls _load_backend, which fails with ModuleNotFoundError - # when the backend module doesn't exist. - with pytest.raises(ModuleNotFoundError, match="No module named 'primus.backends.any_backend'"): - registry_module.BackendRegistry.get_adapter("any_backend") - def test_get_adapter_creation_failure(self): """Test error handling when adapter creation fails.""" @@ -120,11 +97,6 @@ def setup_method(self): # Reset adapter registry state self._original_adapters = registry_module.BackendRegistry._adapters.copy() registry_module.BackendRegistry._adapters.clear() - if hasattr(registry_module.BackendRegistry, "_trainer_classes"): - self._original_trainer_classes = registry_module.BackendRegistry._trainer_classes.copy() - registry_module.BackendRegistry._trainer_classes.clear() - else: - self._original_trainer_classes = None # Ensure backend module can be re-imported so that lazy loading # re-runs registration even if other tests imported it earlier. @@ -139,11 +111,6 @@ def setup_method(self): def teardown_method(self): """Restore registry after each test.""" registry_module.BackendRegistry._adapters = self._original_adapters - if ( - hasattr(registry_module.BackendRegistry, "_trainer_classes") - and self._original_trainer_classes is not None - ): - registry_module.BackendRegistry._trainer_classes = self._original_trainer_classes registry_module.log_rank_0 = self._orig_log_rank_0 # Restore original backend module to avoid impacting other tests @@ -157,11 +124,6 @@ def test_try_load_backend_non_existent(self): with pytest.raises(ImportError): registry_module.BackendRegistry._try_load_backend("definitely_not_a_backend") - def test_try_load_backend_returns_bool(self): - """Test that _try_load_backend returns boolean.""" - with pytest.raises(ImportError): - registry_module.BackendRegistry._try_load_backend("non_existent") - def test_get_adapter_with_lazy_loading(self): """Test that get_adapter triggers lazy loading.""" # Don't pre-register, let it lazy load @@ -189,25 +151,6 @@ def test_list_available_backends(self): assert "backend2" in available assert len(available) == 2 - def test_list_available_backends_empty(self): - """Test listing when no backends registered.""" - available = registry_module.BackendRegistry.list_available_backends() - assert available == [] - - -class TestBackendRegistryPathNames: - """Deprecated: path-name mapping removed; backend path resolution is owned by adapters.""" - - def test_path_name_mapping_removed(self): - pytest.skip("BackendRegistry path-name mapping removed; use adapter.third_party_dir_name.") - - -class TestBackendRegistrySetupPath: - """Deprecated: setup_backend_path moved to BackendAdapter.setup_backend_path().""" - - def test_setup_backend_path_removed(self): - pytest.skip("BackendRegistry.setup_backend_path removed; use adapter.setup_backend_path().") - class TestBackendRegistryGetAdapterIntegration: """Test get_adapter with automatic path setup.""" @@ -262,95 +205,6 @@ def test_get_adapter_path_not_found_error(self): registry_module.BackendRegistry.get_adapter("test_backend", backend_path="/non/existent/path") -class TestBackendRegistryHasAdapter: - """Test has_adapter functionality.""" - - def setup_method(self): - """Clear registry before each test.""" - self._original_adapters = registry_module.BackendRegistry._adapters.copy() - registry_module.BackendRegistry._adapters.clear() - - def teardown_method(self): - """Restore registry after each test.""" - registry_module.BackendRegistry._adapters = self._original_adapters - - def test_has_adapter_true(self): - """Test has_adapter returns True for registered adapter.""" - registry_module.BackendRegistry.register_adapter("test_backend", MockAdapter) - - assert registry_module.BackendRegistry.has_adapter("test_backend") is True - - def test_has_adapter_false(self): - """Test has_adapter returns False for non-registered adapter.""" - assert registry_module.BackendRegistry.has_adapter("non_existent") is False - - -@pytest.mark.skipif( - not _SUPPORTS_TRAINER_CLASS_REGISTRY, - reason="Trainer class registry is not available on BackendRegistry in this version.", -) -class TestBackendRegistryTrainerClasses: - """Test trainer class registration and retrieval.""" - - def setup_method(self): - """Clear trainer classes before each test.""" - if not hasattr(registry_module.BackendRegistry, "_trainer_classes"): - pytest.skip("BackendRegistry has no _trainer_classes in this version.") - self._original_trainer_classes = registry_module.BackendRegistry._trainer_classes.copy() - registry_module.BackendRegistry._trainer_classes.clear() - - def teardown_method(self): - """Restore trainer classes after each test.""" - registry_module.BackendRegistry._trainer_classes = self._original_trainer_classes - - def test_register_and_get_trainer_class(self): - """Test registering and retrieving trainer classes.""" - - class DummyTrainer: - pass - - registry_module.BackendRegistry.register_trainer_class(DummyTrainer, "test_backend") - - trainer_cls = registry_module.BackendRegistry.get_trainer_class("test_backend") - assert trainer_cls is DummyTrainer - - def test_get_trainer_class_not_found(self): - """Test error when trainer class not registered.""" - with pytest.raises(ValueError) as exc_info: - registry_module.BackendRegistry.get_trainer_class("non_existent_backend") - - assert "No trainer class registered for backend 'non_existent_backend'" in str(exc_info.value) - - def test_has_trainer_class(self): - """Test has_trainer_class reflects registration state.""" - - class DummyTrainer: - pass - - assert registry_module.BackendRegistry.has_trainer_class("test_backend") is False - registry_module.BackendRegistry.register_trainer_class(DummyTrainer, "test_backend") - assert registry_module.BackendRegistry.has_trainer_class("test_backend") is True - - def test_register_and_get_trainer_class_with_stage(self): - """Test trainer registration with explicit stage.""" - - class DummyTrainer: - pass - - # Register with explicit stage "sft" - registry_module.BackendRegistry.register_trainer_class(DummyTrainer, "test_backend", stage="sft") - - # Get with explicit stage "sft" should work - trainer_cls = registry_module.BackendRegistry.get_trainer_class("test_backend", stage="sft") - assert trainer_cls is DummyTrainer - - # has_trainer_class with explicit stage should work - assert registry_module.BackendRegistry.has_trainer_class("test_backend", stage="sft") is True - - # Default stage "pretrain" should not find it - assert registry_module.BackendRegistry.has_trainer_class("test_backend") is False - - class TestBackendRegistrySetupHooks: """Test setup hook registration and execution.""" @@ -404,5 +258,67 @@ def failing_hook(): assert "Error in setup hook" in captured.out +class TestBackendRegistryTrainerClass: + """Test the trainer-class registry API. + + ``register_trainer_class`` / ``get_trainer_class`` / ``has_trainer_class`` + back the stage-based fallback in ``MegatronAdapter.load_trainer_class`` (and + the other backend adapters), so they must keep (backend, stage) keying and + raise on lookups for unregistered combinations. + """ + + def setup_method(self): + self._original_trainer_classes = registry_module.BackendRegistry._trainer_classes.copy() + registry_module.BackendRegistry._trainer_classes.clear() + + def teardown_method(self): + registry_module.BackendRegistry._trainer_classes = self._original_trainer_classes + + def test_register_and_get_trainer_class(self): + class DummyTrainer: + pass + + registry_module.BackendRegistry.register_trainer_class(DummyTrainer, "megatron", stage="pretrain") + + assert registry_module.BackendRegistry.get_trainer_class("megatron", stage="pretrain") is DummyTrainer + assert registry_module.BackendRegistry.has_trainer_class("megatron", stage="pretrain") is True + + def test_register_defaults_to_pretrain_stage(self): + class DummyTrainer: + pass + + # Default stage is "pretrain" for both register and lookup. + registry_module.BackendRegistry.register_trainer_class(DummyTrainer, "megatron") + assert registry_module.BackendRegistry.get_trainer_class("megatron") is DummyTrainer + + def test_stage_is_part_of_the_key(self): + class PretrainTrainer: + pass + + registry_module.BackendRegistry.register_trainer_class(PretrainTrainer, "megatron", stage="pretrain") + + # A different stage for the same backend must NOT resolve to the + # pretrain class; it is a distinct registry key. + assert registry_module.BackendRegistry.has_trainer_class("megatron", stage="sft") is False + with pytest.raises(ValueError): + registry_module.BackendRegistry.get_trainer_class("megatron", stage="sft") + + def test_get_unregistered_trainer_class_raises(self): + assert registry_module.BackendRegistry.has_trainer_class("nonexistent") is False + with pytest.raises(ValueError, match="No trainer class registered"): + registry_module.BackendRegistry.get_trainer_class("nonexistent") + + def test_register_overwrites_same_key(self): + class TrainerA: + pass + + class TrainerB: + pass + + registry_module.BackendRegistry.register_trainer_class(TrainerA, "megatron", stage="pretrain") + registry_module.BackendRegistry.register_trainer_class(TrainerB, "megatron", stage="pretrain") + assert registry_module.BackendRegistry.get_trainer_class("megatron", stage="pretrain") is TrainerB + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/core/config/test_primus_config.py b/tests/unit_tests/core/config/test_primus_config.py new file mode 100644 index 000000000..b9e606359 --- /dev/null +++ b/tests/unit_tests/core/config/test_primus_config.py @@ -0,0 +1,61 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +from pathlib import Path + +from primus.core.config.primus_config import get_module_config, load_primus_config +from tests.utils import PrimusUT + +# Top-level attributes that `_normalize_module_for_runtime` keeps in place; every +# other public attribute is moved into `params` on the normalized copy. +_RESERVED_KEYS = {"name", "framework", "config", "model", "params", "trainer_class"} + +_EXAMPLE_CONFIG = "examples/megatron/exp_pretrain.yaml" + + +class TestLoadPrimusConfig(PrimusUT): + def test_exposes_legacy_primus_config(self): + """load_primus_config attaches the underlying legacy PrimusConfig as + `_legacy` so the core runtime can reuse it without re-parsing.""" + cfg = load_primus_config(Path(_EXAMPLE_CONFIG), None) + + self.assertTrue(hasattr(cfg, "_legacy")) + legacy = cfg._legacy + # The legacy object must still provide the PrimusConfig interface that + # BaseModule relies on. + self.assertTrue(callable(getattr(legacy, "get_module_config", None))) + + def test_legacy_config_is_pristine_after_normalization(self): + """Regression guard for the Option A deepcopy nuance: normalization must + operate on deep copies, leaving `_legacy`'s module configs un-mutated. + + With the previous in-place normalization, `delattr` stripped non-reserved + training params off the legacy module config (they were moved into + `params`). That made the exposed legacy config unusable for BaseModule. + Here we assert the legacy module configs still carry their original + top-level training params. + """ + cfg = load_primus_config(Path(_EXAMPLE_CONFIG), None) + legacy = cfg._legacy + + module_keys = list(getattr(legacy, "module_keys", [])) + self.assertTrue(module_keys, "expected at least one module in example config") + + for name in module_keys: + legacy_mod = legacy.get_module_config(name) + top_level = {k for k in vars(legacy_mod) if not k.startswith("_")} + # The legacy module must retain non-reserved top-level params; if the + # normalization had mutated it in place, these would be gone. + self.assertTrue( + top_level - _RESERVED_KEYS, + f"legacy module '{name}' was mutated/stripped by normalization", + ) + + # The normalized SimpleNamespace copy must still expose those same + # params under `.params` for the new runtime. + normalized = get_module_config(cfg, name) + self.assertIsNotNone(normalized) + self.assertTrue(hasattr(normalized, "params")) diff --git a/tests/unit_tests/core/runtime/test_train_runtime.py b/tests/unit_tests/core/runtime/test_train_runtime.py index 51fd5b996..abf257078 100644 --- a/tests/unit_tests/core/runtime/test_train_runtime.py +++ b/tests/unit_tests/core/runtime/test_train_runtime.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -126,7 +126,7 @@ def test_run_trainer_lifecycle_calls_trainer_methods_in_order(self): # Use a dummy trainer that records the call order. class DummyTrainer: - def __init__(self, backend_args=None): + def __init__(self, backend_args=None, **kwargs): self.calls = [] self.backend_args = backend_args @@ -174,7 +174,7 @@ def _fake_run_patches(**kwargs): with patch("primus.core.runtime.train_runtime.run_patches", side_effect=_fake_run_patches): # Dummy trainer class DummyTrainer: - def __init__(self, backend_args=None): + def __init__(self, backend_args=None, **kwargs): self.backend_args = backend_args def setup(self): @@ -204,5 +204,349 @@ def cleanup(self, on_error: bool = False): assert phases == ["build_args", "setup", "before_train", "after_train"] +class TestPrimusRuntimeTrainerClassSelection(PrimusUT): + """Test trainer class selection from config.""" + + def _build_args(self, config: str = "examples/megatron/exp_pretrain.yaml"): + """Build args for testing.""" + import argparse + + return argparse.Namespace(config=config, data_path="./data", backend_path=None) + + def test_initialize_trainer_extracts_trainer_class_from_module_config_top_level(self): + """Test that trainer_class is extracted from module_config (top-level).""" + args = self._build_args() + runtime = PrimusRuntime(args=args) + + # Create a mock adapter that records calls + mock_adapter = Mock() + mock_adapter.prepare_backend = Mock() + mock_adapter.convert_config = Mock(return_value=SimpleNamespace()) + + # Mock trainer class + mock_trainer_class = Mock() + mock_trainer_class.__name__ = "FluxPretrainTrainer" + mock_trainer_instance = Mock() + mock_trainer_class.return_value = mock_trainer_instance + mock_adapter.load_trainer_class = Mock(return_value=mock_trainer_class) + + # Setup context with trainer_class in module_config (top-level) + runtime.ctx = SimpleNamespace( + config_path=Path("dummy.yaml"), + data_path=Path("./data"), + module_name="pre_trainer", + primus_config=SimpleNamespace(), + rank=0, + world_size=1, + master_addr="localhost", + master_port="12345", + module_config=SimpleNamespace( + framework="megatron", + trainer_class="FluxPretrainTrainer", # Top-level attribute + params=SimpleNamespace(stage="pretrain"), + ), + framework="megatron", + adapter=mock_adapter, + ) + + # Mock patches and logging + with patch("primus.core.runtime.train_runtime.log_dict_aligned"), patch( + "primus.core.runtime.train_runtime.log_rank_0" + ) as mock_log, patch("primus.core.runtime.train_runtime.merge_namespace"), patch.object( + runtime, "_run_phase_patches" + ): + + runtime._initialize_trainer() + + # Verify adapter.load_trainer_class was called with trainer_class + mock_adapter.load_trainer_class.assert_called_once_with( + stage="pretrain", trainer_class="FluxPretrainTrainer" + ) + + # Verify logging indicates trainer_class usage + log_calls = [str(call) for call in mock_log.call_args_list] + assert any("Using trainer_class: FluxPretrainTrainer" in str(call) for call in log_calls) + + def test_initialize_trainer_extracts_trainer_class_from_params_when_not_top_level(self): + """Test that trainer_class is extracted from params when not in top-level.""" + args = self._build_args() + runtime = PrimusRuntime(args=args) + + mock_adapter = Mock() + mock_adapter.prepare_backend = Mock() + # Critically, backend_args does NOT carry trainer_class. This isolates the + # params-extraction branch: the only way trainer_class can be resolved is + # from module_config.params (the post-merge backend_args fallback cannot + # mask a regression in that branch). + backend_args = SimpleNamespace(stage="pretrain") + mock_adapter.convert_config = Mock(return_value=backend_args) + mock_trainer_class = Mock() + mock_trainer_class.__name__ = "FluxPretrainTrainer" + mock_trainer_class.return_value = Mock() + mock_adapter.load_trainer_class = Mock(return_value=mock_trainer_class) + + # trainer_class in params, NOT top-level (no top-level attribute) + runtime.ctx = SimpleNamespace( + config_path=Path("dummy.yaml"), + data_path=Path("./data"), + module_name="pre_trainer", + primus_config=SimpleNamespace(), + rank=0, + world_size=1, + master_addr="localhost", + master_port="12345", + module_config=SimpleNamespace( + framework="megatron", + # No trainer_class attribute here + params=SimpleNamespace( + stage="pretrain", + trainer_class="FluxPretrainTrainer", # In params + ), + ), + framework="megatron", + adapter=mock_adapter, + ) + + with patch("primus.core.runtime.train_runtime.log_dict_aligned"), patch( + "primus.core.runtime.train_runtime.log_rank_0" + ), patch("primus.core.runtime.train_runtime.merge_namespace"), patch.object( + runtime, "_run_phase_patches" + ): + runtime._initialize_trainer() + + # trainer_class must have been resolved from module_config.params and forwarded. + mock_adapter.load_trainer_class.assert_called_once_with( + stage="pretrain", trainer_class="FluxPretrainTrainer" + ) + + def test_initialize_trainer_top_level_takes_precedence_over_params(self): + """Test that top-level trainer_class takes precedence over params.trainer_class.""" + args = self._build_args() + runtime = PrimusRuntime(args=args) + + mock_adapter = Mock() + mock_adapter.prepare_backend = Mock() + mock_adapter.convert_config = Mock(return_value=SimpleNamespace()) + mock_trainer_class = Mock() + mock_trainer_class.__name__ = "TopLevelTrainer" + mock_trainer_class.return_value = Mock() + mock_adapter.load_trainer_class = Mock(return_value=mock_trainer_class) + + # Both top-level and params have trainer_class (top-level should win due to if/elif) + runtime.ctx = SimpleNamespace( + config_path=Path("dummy.yaml"), + data_path=Path("./data"), + module_name="pre_trainer", + primus_config=SimpleNamespace(), + rank=0, + world_size=1, + master_addr="localhost", + master_port="12345", + module_config=SimpleNamespace( + framework="megatron", + trainer_class="TopLevelTrainer", # Top-level (checked first) + params=SimpleNamespace( + stage="pretrain", + trainer_class="ParamsTrainer", # In params (should be ignored) + ), + ), + framework="megatron", + adapter=mock_adapter, + ) + + with patch("primus.core.runtime.train_runtime.log_dict_aligned"), patch( + "primus.core.runtime.train_runtime.log_rank_0" + ), patch("primus.core.runtime.train_runtime.merge_namespace"), patch.object( + runtime, "_run_phase_patches" + ): + + runtime._initialize_trainer() + + # Verify top-level trainer_class was used (not params) + mock_adapter.load_trainer_class.assert_called_once_with( + stage="pretrain", trainer_class="TopLevelTrainer" # Top-level, not ParamsTrainer + ) + + def test_initialize_trainer_falls_back_to_stage_when_no_trainer_class(self): + """Test fallback to stage-based selection when trainer_class not specified.""" + args = self._build_args() + runtime = PrimusRuntime(args=args) + + mock_adapter = Mock() + mock_adapter.prepare_backend = Mock() + mock_adapter.convert_config = Mock(return_value=SimpleNamespace()) + mock_trainer_class = Mock() + mock_trainer_class.__name__ = "MockTrainer" + mock_trainer_class.return_value = Mock() + mock_adapter.load_trainer_class = Mock(return_value=mock_trainer_class) + + # No trainer_class specified anywhere + runtime.ctx = SimpleNamespace( + config_path=Path("dummy.yaml"), + data_path=Path("./data"), + module_name="pre_trainer", + primus_config=SimpleNamespace(), + rank=0, + world_size=1, + master_addr="localhost", + master_port="12345", + module_config=SimpleNamespace( + framework="megatron", + params=SimpleNamespace(stage="pretrain"), + ), + framework="megatron", + adapter=mock_adapter, + ) + + with patch("primus.core.runtime.train_runtime.log_dict_aligned"), patch( + "primus.core.runtime.train_runtime.log_rank_0" + ) as mock_log, patch("primus.core.runtime.train_runtime.merge_namespace"), patch.object( + runtime, "_run_phase_patches" + ): + + runtime._initialize_trainer() + + # Verify fallback to stage-based selection. When trainer_class is falsy + # the kwarg is omitted entirely (see train_runtime.py:390-392), so the + # adapter is called with stage only. + mock_adapter.load_trainer_class.assert_called_once_with(stage="pretrain") + + # Verify logging indicates fallback + log_calls = [str(call) for call in mock_log.call_args_list] + assert any("trainer_class not found" in str(call) for call in log_calls) + + def test_initialize_trainer_handles_empty_string_trainer_class(self): + """Test that empty string trainer_class falls back to stage.""" + args = self._build_args() + runtime = PrimusRuntime(args=args) + + mock_adapter = Mock() + mock_adapter.prepare_backend = Mock() + mock_adapter.convert_config = Mock(return_value=SimpleNamespace()) + mock_trainer_class = Mock() + mock_trainer_class.__name__ = "MockTrainer" + mock_trainer_class.return_value = Mock() + mock_adapter.load_trainer_class = Mock(return_value=mock_trainer_class) + + # Empty string trainer_class (falsy, should fall back) + runtime.ctx = SimpleNamespace( + config_path=Path("dummy.yaml"), + data_path=Path("./data"), + module_name="pre_trainer", + primus_config=SimpleNamespace(), + rank=0, + world_size=1, + master_addr="localhost", + master_port="12345", + module_config=SimpleNamespace( + framework="megatron", + trainer_class="", # Empty string (falsy) + params=SimpleNamespace(stage="pretrain"), + ), + framework="megatron", + adapter=mock_adapter, + ) + + with patch("primus.core.runtime.train_runtime.log_dict_aligned"), patch( + "primus.core.runtime.train_runtime.log_rank_0" + ), patch("primus.core.runtime.train_runtime.merge_namespace"), patch.object( + runtime, "_run_phase_patches" + ): + + runtime._initialize_trainer() + + # Should fall back to stage (empty string is falsy): the trainer_class + # kwarg is omitted entirely (see train_runtime.py:390-392). + mock_adapter.load_trainer_class.assert_called_once_with(stage="pretrain") + + +class TestPrimusRuntimeLifecycle(PrimusUT): + """Tests for PrimusRuntime trainer lifecycle execution.""" + + def _build_args(self, config: str = "examples/megatron/exp_pretrain.yaml") -> argparse.Namespace: + return argparse.Namespace(config=config, data_path="./data", backend_path=None) + + def test_run_trainer_lifecycle_applies_patches_before_train(self): + """Test that patches are applied in before_train phase before train() is called.""" + args = self._build_args() + runtime = PrimusRuntime(args=args) + + train_called = [] + patch_applied = [] + + class MockTrainer: + def setup(self): + pass + + def init(self): + pass + + def train(self): + train_called.append(1) + # Verify patch was applied before train + assert len(patch_applied) > 0 + + def cleanup(self, on_error=False): + pass + + mock_trainer = MockTrainer() + runtime.ctx = SimpleNamespace( + trainer=mock_trainer, backend_args=SimpleNamespace(), runtime_state=None + ) + + def mock_run_phase_patches(phase, backend_args=None, runtime_state=None): + if phase == "before_train": + patch_applied.append(1) + + runtime._run_phase_patches = mock_run_phase_patches + + with patch("primus.core.runtime.train_runtime.log_rank_0"): + runtime._run_trainer_lifecycle() + + # Verify patch was applied before train + assert len(patch_applied) == 1 + assert len(train_called) == 1 + + def test_run_trainer_lifecycle_passes_backend_args_to_patches(self): + """Test that backend_args are passed to patch phases.""" + args = self._build_args() + runtime = PrimusRuntime(args=args) + + backend_args_received = [] + + class MockTrainer: + def setup(self): + pass + + def init(self): + pass + + def train(self): + pass + + def cleanup(self, on_error=False): + pass + + mock_trainer = MockTrainer() + test_backend_args = SimpleNamespace(test_param="value") + runtime.ctx = SimpleNamespace( + trainer=mock_trainer, + backend_args=test_backend_args, + runtime_state=None, + ) + + def mock_run_phase_patches(phase, backend_args=None, runtime_state=None): + backend_args_received.append(backend_args) + + runtime._run_phase_patches = mock_run_phase_patches + + with patch("primus.core.runtime.train_runtime.log_rank_0"): + runtime._run_trainer_lifecycle() + + # Verify backend_args were passed to all patch phases + assert len(backend_args_received) == 3 + assert all(args is test_backend_args for args in backend_args_received) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/core/trainer/test_base_trainer.py b/tests/unit_tests/core/trainer/test_base_trainer.py index 6e3b5c4be..9595b71c1 100644 --- a/tests/unit_tests/core/trainer/test_base_trainer.py +++ b/tests/unit_tests/core/trainer/test_base_trainer.py @@ -1,95 +1,117 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### -"""Unit tests for BaseTrainer.""" - -from types import SimpleNamespace - -from primus.core.trainer.base_trainer import BaseTrainer +""" +Unit tests for BaseTrainer. +Tests initialization, MRO handling, distributed environment setup, and +backend_args handling for the universal base trainer class. +""" -class DummyTrainer(BaseTrainer): - """Minimal concrete implementation of BaseTrainer for testing.""" - - def __init__(self, backend_args=None): - super().__init__(backend_args=backend_args) - self.train_calls: int = 0 - self.setup_calls: int = 0 - self.init_calls: int = 0 - - def init(self, *args, **kwargs): - """No-op init for testing.""" - self.init_calls += 1 - return None +from types import SimpleNamespace - def setup(self, *args, **kwargs): - """No-op setup for testing.""" - self.setup_calls += 1 - return None +import pytest - def train(self): - self.train_calls += 1 +from primus.core.trainer.base_trainer import BaseTrainer class TestBaseTrainer: - """Verify BaseTrainer interface and lifecycle.""" + """Tests for BaseTrainer initialization and behavior.""" + + def test_init_sets_distributed_environment(self, monkeypatch: pytest.MonkeyPatch): + """Test that distributed environment attributes are set correctly.""" + mock_dist_env = { + "rank": 2, + "world_size": 8, + "local_rank": 1, + "master_addr": "192.168.1.1", + "master_port": 54321, + } + + monkeypatch.setattr( + "primus.core.trainer.base_trainer.get_torchrun_env", + lambda: mock_dist_env, + ) + monkeypatch.setattr("builtins.print", lambda *args, **kwargs: None) + monkeypatch.setattr("sys.stderr.write", lambda *args, **kwargs: None) + + # Create a concrete subclass for testing + class ConcreteBaseTrainer(BaseTrainer): + def setup(self): + pass + + def init(self): + pass + + def train(self): + pass + + trainer = ConcreteBaseTrainer(backend_args=SimpleNamespace()) + + assert trainer.rank == 2 + assert trainer.world_size == 8 + assert trainer.local_rank == 1 + assert trainer.master_addr == "192.168.1.1" + assert trainer.master_port == 54321 + + def test_init_mro_with_base_module(self, monkeypatch: pytest.MonkeyPatch): + """Test MRO handling when BaseModule IS in inheritance chain (legacy pattern).""" + from primus.modules.base_module import BaseModule + + # Create a class that inherits from both BaseTrainer and BaseModule + class LegacyTrainer(BaseTrainer, BaseModule): + def setup(self): + pass + + def init(self): + pass + + def train(self): + pass + + def run(self): + pass # BaseModule requires run() method + + mock_dist_env = { + "rank": 0, + "world_size": 1, + "local_rank": 0, + "master_addr": "localhost", + "master_port": 12345, + } + + monkeypatch.setattr( + "primus.core.trainer.base_trainer.get_torchrun_env", + lambda: mock_dist_env, + ) + monkeypatch.setattr("builtins.print", lambda *args, **kwargs: None) + monkeypatch.setattr("sys.stderr.write", lambda *args, **kwargs: None) + + # Mock BaseModule.__init__ to track if it's called with kwargs and provide required args + base_module_init_called = [] + + def tracked_init( + self, module_name="test", primus_config=None, module_rank=0, module_world_size=1, **kwargs + ): + base_module_init_called.append(("kwargs", kwargs)) + # Don't call original_init, just track the call + + monkeypatch.setattr(BaseModule, "__init__", tracked_init) - def test_trainer_stores_backend_args(self): - """Backend args should be stored on the trainer instance.""" - backend_args = SimpleNamespace(lr=1e-4, batch_size=32) - trainer = DummyTrainer(backend_args=backend_args) - - assert trainer.backend_args is backend_args - assert trainer.backend_args.lr == 1e-4 - assert trainer.backend_args.batch_size == 32 - - def test_trainer_allows_none_backend_args(self): - """Trainer should accept None as backend_args.""" - trainer = DummyTrainer(backend_args=None) - assert trainer.backend_args is None - - def test_trainer_train_method_executes(self): - """The train() method should execute correctly.""" - backend_args = SimpleNamespace(lr=1e-4) - trainer = DummyTrainer(backend_args=backend_args) - - trainer.train() - - assert trainer.train_calls == 1 - - def test_trainer_lifecycle_methods_exist(self): - """Trainer should have setup, init, train, cleanup methods.""" backend_args = SimpleNamespace() - trainer = DummyTrainer(backend_args=backend_args) - - # All lifecycle methods should be callable - trainer.setup() - trainer.init() - trainer.train() - trainer.cleanup() - - assert trainer.setup_calls == 1 - assert trainer.init_calls == 1 - assert trainer.train_calls == 1 - - def test_cleanup_accepts_on_error_flag(self): - """cleanup() should accept on_error parameter.""" - trainer = DummyTrainer(backend_args=None) - - # Should not raise - trainer.cleanup(on_error=False) - trainer.cleanup(on_error=True) - - def test_trainer_has_distributed_env_attributes(self): - """Trainer should expose distributed environment attributes.""" - trainer = DummyTrainer(backend_args=None) - - # These attributes should exist (values depend on env) - assert hasattr(trainer, "rank") - assert hasattr(trainer, "world_size") - assert hasattr(trainer, "local_rank") - assert hasattr(trainer, "master_addr") - assert hasattr(trainer, "master_port") + trainer = LegacyTrainer( + backend_args=backend_args, + some_kwarg="value", + module_name="test", + primus_config=SimpleNamespace(), + module_rank=0, + module_world_size=1, + ) + + # Verify BaseModule.__init__ was called with kwargs + assert len(base_module_init_called) > 0 + assert "some_kwarg" in base_module_init_called[0][1] + assert trainer.backend_args is backend_args diff --git a/tests/utils.py b/tests/utils.py index 9255b75a6..83490432b 100755 --- a/tests/utils.py +++ b/tests/utils.py @@ -10,17 +10,58 @@ import sys import time import unittest -import warnings from typing import Optional from primus.core.utils import logger TRAINING_COMPLETED_MARKER = "Training completed." -# run_patcher's failure line ("[Patch] \u2717 Patch '' failed..."); distinct -# from a patch's own graceful "[SKIP]". run_patches doesn't fail training on it -# (stop_on_error=False), so we surface it ourselves. -PATCH_FAILURE_MARKER = "\u2717 Patch '" + +def skip_if_no_cuda(reason: str = "requires GPU (primus_turbo initializes CUDA at import)") -> None: + """Skip the calling test module at collection time when CUDA is unavailable. + + Several Flux/diffusion test modules import ``primus_turbo`` (directly or + transitively), which initializes CUDA at import and raises on CPU-only + hosts. Call this at module scope *before* those imports so collection + succeeds without a GPU. It is a no-op when CUDA is present, so GPU CI still + runs every test. + """ + import pytest + import torch + + if not torch.cuda.is_available(): + pytest.skip(reason, allow_module_level=True) + + +def install_aiter_deepbind_hook() -> None: + """Install Primus' production RTLD_DEEPBIND import hook for aiter's mha kernels. + + This invokes the exact mechanism the ``megatron.turbo.aiter_deepbind`` + before_train patch uses in real training: it wraps ``importlib.import_module`` + so aiter's pinned mha extensions bind their own ``aiter::mha_bwd`` instead of + transformer_engine's stale vendored ``libmha`` (ROCm/aiter#1332). On + gfx942/gfx950 that stale-symbol interposition otherwise makes Turbo's hd128 + backward launch with an invalid grid config and crash the process. + + Unit tests call ``flash_attn_func`` directly and never run the before_train + phase, so without this hook they hit the same crash the production patch + prevents. Call this at diffusion ``conftest`` import time so the hook is in + place before any test first imports the aiter mha modules (which happens + lazily on the first attention op). No-op when CUDA is unavailable or the hook + cannot be installed (e.g. CPU-only host). + """ + try: + import torch + + if not torch.cuda.is_available(): + return + from primus.backends.megatron.patches.turbo.aiter_deepbind_patches import ( + _install_deepbind_import_hook, + ) + except Exception: + return + + _install_deepbind_import_hook() class PrimusUT(unittest.TestCase): @@ -51,25 +92,6 @@ def tearDown(self): pass -def _warn_on_patch_failures(tag: str, stdout_output: str) -> None: - """Warn (don't fail) when patches failed to apply during training. - - A failed patch doesn't crash training and a 3-step smoke run can't prove it - was harmless, so we emit a GitHub Actions warning to keep it visible without - blocking CI. No-op for backends that don't emit the marker. - """ - if PATCH_FAILURE_MARKER not in stdout_output: - return - failed = [ln.strip() for ln in stdout_output.splitlines() if PATCH_FAILURE_MARKER in ln] - msg = ( - f"[{tag}] {len(failed)} patch(es) failed to apply during training " - f"(training still completed): " + " | ".join(failed[:10]) - ) - # GitHub Actions annotation: shows on the run/PR without failing the job. - print(f"::warning title=Primus patch failed to apply::{msg}") - warnings.warn(msg) - - def run_training_script( tag: str, cmd: list[str], @@ -129,8 +151,6 @@ def run_training_script( f"Log file: {train_log_path}" ) - _warn_on_patch_failures(tag, stdout_output) - return stdout_output, "" except subprocess.CalledProcessError as e: From 6e8914ff7859fb408918a312aeda57d1d9bd8722 Mon Sep 17 00:00:00 2001 From: luiza-amd Date: Tue, 7 Jul 2026 11:27:12 +0300 Subject: [PATCH 006/127] merge to main of docs(flux): diffusion training documentation (#858) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **Important:** The content of this PR was approved and merged https://github.com/AMD-AGI/Primus/pull/823, but not to main (auto-target to main was not triggered and the lesson was taken into account for further PRs). The purpose of this PR is to finalize the merge to main and not to introduce any new changes. Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Content-independent — can review/merge in any order. Cut from `feat/flux/ci-env` (so its head carries the draft guard) and targets that branch; it auto-retargets to `main` when the CI-pins PR merges. ## What this changes The diffusion documentation set: the `docs/backends/megatron/diffusion/*` pages, the CLI / top-level / examples READMEs, the example `run_pretrain.sh`, and the requirement-check runner hook. Content reflects the current curated layout (rewritten fp8/mxfp4/structure/data docs). ## Dependencies No content parents (docs-only); cut from `feat/flux/ci-env` only to carry the draft guard, with no functional dependency on the turbo bump. ## Test plan Lint/pre-commit plus a link / relocated-path check; no runtime tests. ## Files 17 (diffusion docs, READMEs, example launcher, requirement-check hook). Co-authored-by: Flux Split Trial --- docs/README.md | 2 + docs/backends/megatron/diffusion/README.md | 371 ++++++ docs/backends/megatron/diffusion/STRUCTURE.md | 254 ++++ .../megatron/diffusion/adding_new_models.md | 711 +++++++++++ .../megatron/diffusion/api_reference.md | 1059 +++++++++++++++++ .../diffusion/architecture_overview.md | 497 ++++++++ .../megatron/diffusion/data_preprocessing.md | 629 ++++++++++ .../megatron/diffusion/energon_integration.md | 514 ++++++++ .../megatron/diffusion/flux_architecture.md | 926 ++++++++++++++ .../megatron/diffusion/fp8_training.md | 514 ++++++++ .../megatron/diffusion/mxfp4_training.md | 210 ++++ docs/cli/README.md | 16 + examples/README.md | 7 + examples/megatron/diffusion/README.md | 317 +++++ examples/run_pretrain.sh | 12 +- .../hooks/05_check_primus_requirements.sh | 68 ++ .../hooks/train/pretrain/megatron/prepare.py | 46 +- 17 files changed, 6145 insertions(+), 8 deletions(-) create mode 100644 docs/backends/megatron/diffusion/README.md create mode 100644 docs/backends/megatron/diffusion/STRUCTURE.md create mode 100644 docs/backends/megatron/diffusion/adding_new_models.md create mode 100644 docs/backends/megatron/diffusion/api_reference.md create mode 100644 docs/backends/megatron/diffusion/architecture_overview.md create mode 100644 docs/backends/megatron/diffusion/data_preprocessing.md create mode 100644 docs/backends/megatron/diffusion/energon_integration.md create mode 100644 docs/backends/megatron/diffusion/flux_architecture.md create mode 100644 docs/backends/megatron/diffusion/fp8_training.md create mode 100644 docs/backends/megatron/diffusion/mxfp4_training.md create mode 100644 examples/megatron/diffusion/README.md create mode 100755 runner/helpers/hooks/05_check_primus_requirements.sh diff --git a/docs/README.md b/docs/README.md index 48c6c271e..eb68c4fa5 100644 --- a/docs/README.md +++ b/docs/README.md @@ -36,6 +36,8 @@ In-depth technical documentation: - **[Backend Extension Guide](./backends/extending-backends.md)** - How to add a new backend using the current adapter/trainer architecture - **[Megatron Model Extension Guide](./backends/adding-megatron-models.md)** - How to add a new Megatron model config - **[TorchTitan Model Extension Guide](./backends/adding-torchtitan-models.md)** - How to add a new TorchTitan model config +- **[Flux Diffusion Models](./backends/megatron/diffusion/README.md)** - Flux diffusion model architecture, training, and API reference +- **[FP8 Training Guide](./backends/megatron/diffusion/fp8_training.md)** - FP8 precision training on AMD MI300X/MI355X: configuration, benchmarks, and tuning ### 💡 Help & Support diff --git a/docs/backends/megatron/diffusion/README.md b/docs/backends/megatron/diffusion/README.md new file mode 100644 index 000000000..606a25b3a --- /dev/null +++ b/docs/backends/megatron/diffusion/README.md @@ -0,0 +1,371 @@ +# Diffusion Models in Primus - Developer & Architecture Guide + +**Purpose:** Developer-focused documentation for understanding Primus diffusion architecture, design decisions, and implementation details. + +**For training/usage instructions, see:** [examples/megatron/diffusion/README.md](../../../../examples/megatron/diffusion/README.md) + +**For test documentation, see:** [tests/unit_tests/backends/megatron/diffusion/](../../../../tests/unit_tests/backends/megatron/diffusion/) + +--- + +## Architecture Philosophy + +Primus diffusion models are built as **Megatron-Core native implementations**, designed for: +- Production-scale distributed training +- Seamless integration with Megatron parallelism strategies (TP, PP, DP, EP) +- Advanced checkpoint management with heterogeneous layers +- Clean separation of concerns (no framework dependencies like PyTorch Lightning) + +### Key Design Decisions + +**1. Megatron-Core Integration** +- Models in `core/models/diffusion/` follow Megatron-Core patterns +- Extends `TransformerConfig` for configurations (inherits all Megatron features) +- Uses `TransformerBlock` with heterogeneous layer support +- Compatible with Megatron's distributed checkpointing + +**2. Unified TransformerBlock Architecture** +- Unlike HuggingFace's ModuleLists, uses Megatron's unified TransformerBlock +- Simplifies checkpoint management +- More efficient gradient synchronization +- Note: pipeline parallelism is not supported for diffusion models (`pipeline_model_parallel_size` must be 1) + +**3. No Framework Dependencies** +- Direct PyTorch implementation (no PyTorch Lightning) +- Uses Megatron's distributed primitives directly +- Simpler debugging and profiling +- Better control over distributed training + +**4. Extensibility First** +- Base classes designed for multiple diffusion models (Flux, DiT, MovieGen) +- Clear shared vs model-specific separation +- Hierarchical encoder registry for easy extension + +--- + +## Supported Models + +### Flux ✅ Production Ready +Flow-based diffusion model with MMDiT (Multimodal Diffusion Transformer) architecture. + +- **Architecture**: Dual-stream with joint and single transformer blocks +- **Sizes**: 535M (testing) and 12B (production) +- **Reference**: [Black Forest Labs FLUX.1](https://huggingface.co/black-forest-labs/FLUX.1-dev) +- **Status**: Fully implemented and tested (390 tests) + +### Future Models ⏳ Planned +- **DiT**: Diffusion Transformer for image generation +- **MovieGen**: Video diffusion models +- **Custom Models**: Extensible framework for new architectures + +--- + +## Project Structure + +``` +primus/backends/megatron/ +├── core/models/ +│ ├── common/diffusion_module/ # DiffusionModule (base class with sharded state dict) +│ │ └── diffusion_module.py +│ └── diffusion/ # Model implementations (Megatron-Core style) +│ ├── common/ # Shared components (MMDiT layers, attention) +│ │ ├── config.py # BaseDiffusionConfig (extends TransformerConfig) +│ │ └── layers.py # Shared layers (if any) +│ └── flux/ # Flux-specific code +│ ├── config.py # FluxConfig with factory methods (535M, 12B) +│ ├── model.py # Flux model (extends DiffusionModule) +│ └── layer_spec.py # Flux layer specifications +│ +├── training/diffusion/ # Training utilities +│ ├── noise_utils.py # Noise application (flow matching, DDPM) +│ ├── loss_computation.py # Loss functions (flow matching, epsilon, v-prediction) +│ ├── timestep_sampling.py # Timestep sampling strategies +│ └── schedulers/ +│ ├── base.py # BaseScheduler +│ └── flow_matching.py # FlowMatchEulerDiscreteScheduler +│ +└── data/ + ├── energon/ # Shared Energon infrastructure + └── diffusion/ # Diffusion-specific data + ├── encoders/ # Hierarchical encoder registry + │ ├── image/vae/ # VAE variants (SD VAE, custom VAEs) + │ ├── text/t5/ # T5 variants (XXL, etc.) + │ └── text/clip/ # CLIP variants (L, H, etc.) + ├── preprocessing/ # Data preprocessing utilities + │ ├── download.py # Reusable download utils (retry, MD5, manifests) + │ ├── finalize.py # Energon dataset finalization + │ ├── validate.py # Dataset structure validation + │ └── pipelines/ # Dataset preparation pipelines + │ ├── base.py # DatasetPipeline abstract base class + │ ├── raw.py # Raw image pipeline + │ ├── encoded.py # Pre-encoded pipeline + │ └── ingest.py # StreamingIngestPipeline (MLPerf Arrow->WDS) + └── task_encoders/ # Energon TaskEncoders for diffusion + +primus/configs/models/megatron/diffusion/ +├── flux_535m.yaml # Flux 535M config +├── flux_12b.yaml # Flux 12B config +└── encoders.yaml # Encoder configuration + +tests/unit_tests/backends/megatron/diffusion/ # Comprehensive test suite (390 tests) +├── models/ # Model-level tests +├── layers/ # Layer-level tests +├── unit/ # Unit tests for utilities +├── distributed/ # Distributed training tests +├── functional/ # End-to-end functional tests +└── checkpointing/ # Checkpoint tests + +docs/backends/megatron/diffusion/ # This directory +├── README.md # This file (developer guide) +├── architecture_overview.md # Detailed architecture +├── data_preprocessing.md # Data pipeline guide (includes Flux-specific section) +├── energon_integration.md # Energon patterns +├── flux_architecture.md # Flux deep dive +├── fp8_training.md # FP8 training guide (benchmarks, tuning, troubleshooting) +├── api_reference.md # API documentation +├── adding_new_models.md # Extension guide +└── STRUCTURE.md # Directory tree and organization +``` + +--- + +## Key Technical Features + +### 1. DiffusionModule Base Class + +All diffusion models inherit from `DiffusionModule`, which provides: +- Megatron-Core integration (process groups, parallelism) +- Sharded state dict support for distributed checkpointing +- Gradient checkpointing +- Mixed precision support +- Device placement utilities + +**Location:** `primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py` + +### 2. BaseDiffusionConfig + +Configuration class extending `TransformerConfig`: +- Inherits all Megatron-Core configuration (TP, PP, sequence_parallel, etc.) +- Adds diffusion-specific parameters (channels, patch_size, etc.) +- Factory methods for common presets + +**Location:** `primus/backends/megatron/core/models/diffusion/common/config.py` + +### 3. Hierarchical Encoder Registry + +Organized by modality → type → variant: +``` +encoders/ +├── image/vae/ +│ ├── sd_vae.py # Standard SD VAE +│ └── (future: custom VAEs) +├── text/t5/ +│ ├── t5_xxl.py # T5-XXL encoder +│ └── (future: T5 variants) +└── text/clip/ + ├── clip_l.py # CLIP-L encoder + └── (future: CLIP-H, etc.) +``` + +Benefits: +- Easy to add new encoder variants (5+ planned per modality) +- Config-driven selection via `encoders.yaml` +- Lazy loading (encoders loaded only when needed) +- Shared base classes for common functionality + +### 4. Training Utilities Structure + +**Noise Application** (`noise_utils.py`): +- `apply_flow_matching_noise()`: For flow matching models (Flux) +- `apply_ddpm_noise()`: For DDPM-based models +- Support for different noise schedules + +**Loss Computation** (`loss_computation.py`): +- `compute_flow_matching_loss()`: For flow matching +- `compute_epsilon_loss()`: For epsilon prediction (DDPM) +- `compute_v_prediction_loss()`: For v-prediction +- Unified interface for different loss types + +**Timestep Sampling** (`timestep_sampling.py`): +- `LogitNormalSampler`: Logit-normal distribution +- `UniformSampler`: Uniform distribution +- `ModeSampler`: Mode-focused sampling +- Base class for custom samplers + +### 5. Shared Energon Infrastructure + +Located in `data/energon/` for reusability across models: +- Shared data loading utilities +- Common preprocessing functions +- WebDataset integration +- Model-specific TaskEncoders in `data/diffusion/task_encoders/` + +### 6. Precalculated Data Support + +**Performance**: 5-10x faster training than on-the-fly encoding + +**Supported encodings**: +- `preencoded` -- Primus-encoded PyTorch `.pth` format (VAE latents + text embeddings) +- `preencoded_numpy` -- MLPerf NumPy uint16 format (bfloat16 tensors as `.bytes` entries) + +**Workflow**: +1. Precompute VAE latents and text embeddings offline +2. Store in WebDataset/Energon format +3. Load directly during training (no encoder overhead) + +**Benefits**: +- Faster training iteration +- Consistent encoder versions across runs +- Lower GPU memory (no encoders loaded during training) +- Better reproducibility + +### 7. MLPerf Streaming Ingest Pipeline + +**Location:** `data/diffusion/preprocessing/pipelines/ingest.py` + +The `StreamingIngestPipeline` downloads Apache Arrow IPC files from MLCommons R2 storage and converts them directly into Energon WebDataset tar shards in a single streaming pass. This avoids storing the full ~6 TB raw Arrow dataset on disk. + +**Architecture**: Producer-consumer with concurrent download and sequential conversion: +- **Producer thread**: Acquires a semaphore permit, submits downloads to a `ThreadPoolExecutor`, passes completed futures to a drain thread +- **Drain thread**: Processes futures in submission order and feeds the prefetch queue +- **Consumer (main thread)**: Converts Arrow data to tar shards, deletes temporary files, releases semaphore permits + +**Key properties**: +- Bounded disk usage: `threading.Semaphore(prefetch_depth)` limits Arrow files on disk +- Deterministic shard ordering preserved via in-order future draining +- Retry with exponential backoff for HTTP 429/503 and MD5 mismatches (`download.py`) +- Skip-and-log: individual failures are recorded in `failed_files.json` +- Resume: re-running skips shards that already exist on disk + +**Related modules**: +- `download.py`: `download_with_backoff()`, `fetch_manifest()`, `parse_md5_manifest()` +- `pipelines/base.py`: `DatasetPipeline` ABC (shared by `raw.py`, `encoded.py`, `ingest.py`) +- `finalize.py`: Energon dataset finalization (`.nv-meta/dataset.yaml` + `energon prepare`) +- `validate.py`: Post-finalization structural validation + +--- + +## Implementation Status + +### Core Infrastructure ✅ +- ✅ Directory structure with 25+ directories +- ✅ Base classes (DiffusionModule, BaseDiffusionConfig, BaseScheduler) +- ✅ DiffusionModule with Megatron-Core integration +- ✅ FluxConfig with factory methods (flux_535m, flux_12b) +- ✅ FlowMatchEulerDiscreteScheduler +- ✅ Configuration system (YAML files) +- ✅ Testing framework (390 tests) +- ✅ Comprehensive documentation + +### Flux Model Implementation ✅ +- ✅ Flux model architecture (dual-stream MMDiT) +- ✅ MMDiT layers and attention (joint + single blocks) +- ✅ Embeddings (3D RoPE, timestep, vector) +- ✅ Hierarchical encoder registry +- ✅ Data pipeline and TaskEncoders +- ✅ Training utilities (noise, loss, sampling) +- ✅ Checkpoint conversion (HF <-> Megatron) + +--- + +## Documentation Map + +### Core Guides + +📖 **[Architecture Overview](architecture_overview.md)** +High-level design, directory structure, and architectural decisions. + +📖 **[Directory Structure](STRUCTURE.md)** +Complete directory tree and file organization. + +📖 **[Data Preprocessing Guide](data_preprocessing.md)** +How to prepare datasets, precalculate latents, and use Energon. + +📖 **[Energon Integration](energon_integration.md)** +Megatron-Energon patterns and TaskEncoder implementation. + +📖 **[Adding New Models](adding_new_models.md)** +Step-by-step guide for implementing new diffusion models. + +### Advanced Documentation + +📖 **[Flux Architecture Deep Dive](flux_architecture.md)** +Mathematical formulation, detailed component descriptions, and performance optimizations. + +📖 **[API Reference](api_reference.md)** +Complete API documentation with function signatures and usage examples. + +📖 **[FP8 Training Guide](fp8_training.md)** +FP8 precision training on AMD MI300X: configuration, benchmarks, tuning recipes, and troubleshooting. + +### Related Documentation + +📖 **[Training Guide](../../../../examples/megatron/diffusion/README.md)** +User-facing guide for training Flux models (quick start, configurations, troubleshooting). + +📖 **[Test Directory](../../../../tests/unit_tests/backends/megatron/diffusion/)** +Test suite for diffusion models. + +--- + +## Testing Architecture + +**Test Organization** (following Megatron-LM patterns): +- One comprehensive file per model (`test_flux_model.py`) +- Unit tests for utilities (`unit/test_utils.py`, etc.) +- Distributed tests in separate directory (`distributed/`) +- Functional tests for workflows (`functional/`) + +**Test Status**: ✅ 390 tests passing + +See [tests/unit_tests/backends/megatron/diffusion/](../../../../tests/unit_tests/backends/megatron/diffusion/) for details. + +--- + +## Hardware Requirements + +### Flux 535M (Testing) +- **Training**: 1x MI300X 192GB (compatible with H100/A100) +- **Inference**: 1x MI300X 192GB +- **Batch Size**: 1-8 per GPU + +### Flux 12B (Production) +- **Training**: 8x MI300X 192GB (recommended) or 4x MI300X 192GB with TP=2 +- **Inference**: 1x MI300X 192GB +- **Batch Size**: 1-2 per GPU for training, 1-4 for inference + +--- + +## Contributing + +See the main guide: [Adding New Models](adding_new_models.md) + +**To contribute:** +1. Follow the established directory structure +2. Extend base classes (DiffusionModule, BaseDiffusionConfig) +3. Add comprehensive tests in `tests/unit_tests/backends/megatron/diffusion/` +4. Update documentation (architecture guide + API reference) +5. Submit PR with clear description + +--- + +## License + +- **Primus Code**: AMD Copyright 2025, Apache License 2.0 +- **Flux Encoders**: + - FLUX.1 [dev]: Non-commercial license + - FLUX.1 [schnell]: Apache 2.0 (commercial use allowed) + - Individual components (T5, CLIP, VAE): Check respective licenses + +--- + +## Resources + +- **Megatron-Core**: [nvidia/Megatron-LM](https://github.com/NVIDIA/Megatron-LM) - Core framework +- **Flux Model**: [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) +- **Flow Matching**: Rectified flow and flow matching papers +- **NeMo**: [nvidia/NeMo](https://github.com/NVIDIA/NeMo) - Alternative diffusion implementation + +--- + +**Last Updated**: January 2026 diff --git a/docs/backends/megatron/diffusion/STRUCTURE.md b/docs/backends/megatron/diffusion/STRUCTURE.md new file mode 100644 index 000000000..382ab7eb4 --- /dev/null +++ b/docs/backends/megatron/diffusion/STRUCTURE.md @@ -0,0 +1,254 @@ +# Flux Diffusion Infrastructure - Directory Structure + +**Created**: December 5, 2025 +**Status**: ✓ Implementation Complete + +## Overview + +This document describes the directory structure created for Flux diffusion model support in Primus, following Megatron-Core conventions with production-ready enhancements. + +--- + +## Directory Tree + +``` +Primus/ +├── primus/backends/megatron/ +│ ├── core/models/ +│ │ ├── common/diffusion_module/ # DiffusionModule base class +│ │ │ └── diffusion_module.py +│ │ └── diffusion/ # Diffusion models (Megatron-Core convention) +│ │ ├── common/ # Shared components (MMDiT layers, attention) +│ │ │ ├── __init__.py +│ │ │ ├── config.py # ✓ BaseDiffusionConfig +│ │ │ ├── embeddings.py # ✓ TimeStepEmbedder, MLPEmbedder +│ │ │ └── normalization.py # ✓ AdaLN, AdaLNContinuous, RMSNorm +│ │ ├── flux/ # Flux-specific components +│ │ │ ├── __init__.py +│ │ │ ├── config.py # ✓ FluxConfig (with factory methods) +│ │ │ ├── model.py # ✓ Flux model +│ │ │ ├── layers.py # ✓ EmbedND, embedders +│ │ │ ├── layer_spec.py # ✓ get_flux_layer_spec, get_flux_*_spec_for_backend, MMDiTLayer +│ │ │ ├── attention.py # ✓ JointSelfAttention, FluxSingleAttention +│ │ │ ├── utils.py # ✓ generate_image_position_ids +│ │ │ ├── checkpoint_utils.py # ✓ Checkpoint utilities +│ │ │ └── checkpoint_converter.py # ✓ HF <-> Megatron conversion +│ │ └── __init__.py +│ │ +│ ├── training/diffusion/ # Training utilities +│ │ ├── schedulers/ +│ │ │ ├── __init__.py +│ │ │ ├── base.py # ✓ BaseScheduler +│ │ │ └── flow_matching.py # ✓ FlowMatchEulerDiscreteScheduler +│ │ ├── noise_utils.py # ✓ apply_flow_matching_noise, apply_ddpm_noise +│ │ ├── loss_computation.py # ✓ compute_flow_matching_loss, etc. +│ │ ├── timestep_sampling.py # ✓ LogitNormalSampler, UniformSampler +│ │ └── __init__.py +│ │ +│ └── data/ +│ ├── energon/ # Shared Energon infrastructure +│ │ └── __init__.py # ✓ Energon wrappers +│ │ +│ └── diffusion/ # Diffusion-specific data +│ ├── encoders/ # Hierarchical encoder registry +│ │ ├── image/ +│ │ │ ├── vae/ # VAE variants +│ │ │ │ └── __init__.py # ✓ AutoencoderKL, VQVAE, etc. +│ │ │ └── __init__.py +│ │ ├── text/ +│ │ │ ├── t5/ # T5 variants +│ │ │ │ └── __init__.py # ✓ T5-XXL, T5-Large, etc. +│ │ │ ├── clip/ # CLIP variants +│ │ │ │ └── __init__.py # ✓ CLIP-L, CLIP-H, etc. +│ │ │ └── __init__.py +│ │ └── __init__.py # ✓ EncoderRegistry +│ │ +│ ├── preprocessing/ +│ │ ├── image/ +│ │ │ └── __init__.py # ✓ Resizing, augmentation +│ │ └── __init__.py +│ │ +│ ├── task_encoders/ # Energon TaskEncoders +│ │ ├── __init__.py +│ │ └── image.py # ✓ EncodedDiffusionTaskEncoder, RawDiffusionTaskEncoder +│ │ +│ └── __init__.py +│ +├── primus/modules/trainer/megatron/ +│ └── diffusion/ # Diffusion trainer +│ └── __init__.py # ✓ DiffusionTrainer +│ +├── primus/configs/models/megatron/ +│ └── diffusion/ # YAML configs +│ ├── __init__.py +│ ├── flux_535m.yaml # ✓ Flux 535M config +│ ├── flux_12b.yaml # ✓ Flux 12B config +│ └── encoders.yaml # ✓ Encoder configs +│ +├── examples/megatron/ +│ ├── diffusion/ +│ │ └── README.md # ✓ Training guide (consolidated) +│ ├── configs/MI300X/diffusion/ # MI300X training configs +│ │ ├── flux_535m_pretrain.yaml +│ │ ├── flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml +│ │ ├── flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml +│ │ └── ... +│ ├── configs/MI355X/diffusion/ # MI355X training configs (mirrors MI300X + MXFP4/MLPerf) +│ │ ├── flux_12b_ddp_energon_schnell_resample_*.yaml +│ │ ├── flux_12b_fsdp2_energon_schnell_resample_*.yaml +│ │ └── ... +│ └── prepare.py +│ +├── examples/run_pretrain.sh # Main training script +│ +├── tests/ +│ ├── unit_tests/backends/megatron/diffusion/ # Unit test suite +│ │ ├── test_flux_model.py +│ │ ├── test_flux_config.py +│ │ ├── test_flux_layers.py +│ │ ├── test_flux_embeddings.py +│ │ ├── test_flux_normalization.py +│ │ ├── test_flux_utils.py +│ │ ├── test_flux_checkpoint_converter.py +│ │ ├── test_flux_checkpoint_utils.py +│ │ ├── test_flux_layer_spec_backend_selection.py +│ │ ├── test_flux_compile_checkpoint_keys.py +│ │ ├── training/ +│ │ ├── data/ +│ │ └── distributed/ +│ └── integration_tests/backends/megatron/diffusion/ +│ ├── data/ +│ └── distributed/ +│ +└── docs/backends/megatron/ + └── diffusion/ # Documentation + ├── README.md # ✓ Overview + ├── STRUCTURE.md # ✓ This file + ├── architecture_overview.md # ✓ Design details + ├── data_preprocessing.md # ✓ Data guide (includes Flux-specific section) + ├── energon_integration.md # ✓ Energon patterns + ├── flux_architecture.md # ✓ Flux deep dive + ├── fp8_training.md # ✓ FP8 training guide + ├── api_reference.md # ✓ API documentation + └── adding_new_models.md # ✓ Extension guide +``` + +--- + +## Completed Components + +### ✓ Base Classes + +1. **DiffusionModule** (`core/models/common/diffusion_module/diffusion_module.py`) + - Base class for all diffusion models (extends MegatronModule) + - Provides Megatron-Core integration + - Required methods: `forward()` + - Loss computation: Use standalone functions from `loss_computation.py` + - Utility methods: `get_num_params()`, `set_requires_grad()` + +2. **BaseDiffusionConfig** (`common/config.py`) + - Extends `megatron.core.transformer.transformer_config.TransformerConfig` + - Common parameters: `in_channels`, `out_channels`, `patch_size` + - Validation method for configuration integrity + +3. **FluxConfig** (`flux/config.py`) + - Flux-specific configuration + - Parameters: `num_joint_layers`, `num_single_layers`, `context_dim`, `vec_in_dim` + - Factory methods: `flux_535m()`, `flux_12b()` + - 3D RoPE configuration: `axes_dim`, `theta` + +3. **BaseScheduler** (`schedulers/base.py`) + - Abstract base for diffusion schedulers + - Required: `add_noise()`, `get_velocity_target()`, `sample_timesteps()` + - Optional: `scale_model_input()`, `get_snr()`, `get_alpha()`, `get_sigma()` + +4. **FlowMatchEulerDiscreteScheduler** (`schedulers/flow_matching.py`) + - Concrete implementation for Flux + - Linear interpolation: `x_t = (1-t)*noise + t*data` + - Velocity target: `v = data - noise` + +### ✓ Directory Structure + +- **25 `__init__.py` files** with comprehensive docstrings +- **Multiple implementation files** (models, configs, schedulers, data pipeline) +- **Complete test suite** with fixtures and helpers + +--- + +## Architectural Decisions + +### 1. Models under `core/models/` +- Follows Megatron-Core convention (`megatron/core/models/gpt/`, etc.) +- Easier upstream tracking when Megatron-Core adds diffusion support + +### 2. Shared Components in `common/` +- Standard approach stores shared code in model-specific directories +- Primus: `common/` for MMDiT layers, attention, shared utilities +- Flux-specific: Only `EmbedND` and Flux model class + +### 3. Hierarchical Encoder Structure +- `encoders/image/vae/`, `encoders/text/t5/`, `encoders/text/clip/` +- Registry pattern for config-driven selection +- Easy to add new encoder variants (5+ planned per modality) + +### 4. Shared Energon Infrastructure +- `data/energon/` for cross-model utilities (VLM, diffusion, future) +- `data/diffusion/task_encoders/` for diffusion-specific TaskEncoders +- Traditional approach nests Energon under model-specific directories + +### 5. Mock Data in Tests +- `tests/fixtures/diffusion/` (not production code) +- Traditional approach mixes test utilities with production code +- Follows pytest best practices + +### 6. No PyTorch Lightning +- Pure Megatron patterns (no PTL DataModules) +- Better integration with Megatron training loop + +--- + +## Import Examples + +```python +# Base classes +from primus.backends.megatron.core.models.diffusion.common import ( + BaseDiffusionConfig, +) + +# Flux configuration +from primus.backends.megatron.core.models.diffusion.flux import FluxConfig + +# Create configs +config_535m = FluxConfig.flux_535m() +config_12b = FluxConfig.flux_12b() + +# Schedulers +from primus.backends.megatron.training.diffusion.schedulers import ( + BaseScheduler, + FlowMatchEulerDiscreteScheduler, +) + +# Create scheduler +scheduler = FlowMatchEulerDiscreteScheduler() +timesteps = scheduler.sample_timesteps(batch_size=8, device='cuda') +``` + +--- + +## Validation Status + +✓ All Python files syntactically correct +✓ No linter errors detected +✓ All imports properly structured +✓ Comprehensive docstrings +✓ Copyright headers applied (AMD 2025, Apache 2.0) + +--- + +## Files Summary + +All infrastructure files, model implementations, data pipeline components, tests, and documentation are complete and ready for production use. + +--- + +**End of Structure Document** diff --git a/docs/backends/megatron/diffusion/adding_new_models.md b/docs/backends/megatron/diffusion/adding_new_models.md new file mode 100644 index 000000000..a29ed8fd0 --- /dev/null +++ b/docs/backends/megatron/diffusion/adding_new_models.md @@ -0,0 +1,711 @@ +# Adding New Diffusion Models + +This guide explains how to add new diffusion models to Primus, following the established patterns and architecture. + +--- + +## Table of Contents + +1. [Overview](#overview) +2. [Prerequisites](#prerequisites) +3. [Step-by-Step Guide](#step-by-step-guide) +4. [Example: Adding DiT](#example-adding-dit) +5. [Testing Your Model](#testing-your-model) +6. [Best Practices](#best-practices) + +--- + +## Overview + +Adding a new diffusion model involves: +1. Creating model configuration +2. Implementing model class +3. Adding necessary layers +4. Creating data pipeline components +5. Writing tests +6. Updating documentation + +**Time Estimate**: 5-10 days depending on model complexity + +--- + +## Prerequisites + +Before adding a new model, ensure you have: +- ✅ Understanding of the model architecture (paper, reference implementation) +- ✅ Access to pretrained weights (if applicable) +- ✅ Sample dataset for testing +- ✅ Familiarity with Primus diffusion architecture +- ✅ Development environment setup + +**Required Reading**: +- [Architecture Overview](architecture_overview.md) +- [Data Preprocessing Guide](data_preprocessing.md) +- [Energon Integration](energon_integration.md) + +--- + +## Step-by-Step Guide + +### Step 1: Create Model Directory + +Create a directory for your model under `core/models/diffusion/`: + +```bash +mkdir -p primus/backends/megatron/core/models/diffusion/dit +cd primus/backends/megatron/core/models/diffusion/dit +``` + +Create files: +```bash +touch __init__.py +touch config.py +touch model.py +touch layers.py # If model-specific layers needed +``` + +### Step 2: Implement Configuration + +**File**: `config.py` + +```python +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Configuration for DiT (Diffusion Transformer) model.""" + +from dataclasses import dataclass +from typing import Optional +from ..common.config import BaseDiffusionConfig + + +@dataclass +class DiTConfig(BaseDiffusionConfig): + """ + DiT-specific configuration. + + DiT uses a standard transformer architecture for diffusion. + """ + + # Model identification + model_type: str = "dit" + + # Architecture: Number of layers + num_layers: int = 28 # DiT-XL/2 default + + # Architecture: Dimensions + hidden_size: int = 1152 + num_attention_heads: int = 16 + + # Input dimensions + in_channels: int = 4 # VAE latent channels (standard SD VAE) + + # Context dimensions + context_dim: int = 768 # CLIP text embedding dimension + + # Patchification + patch_size: int = 2 # DiT uses 2x2 patches + + # Class conditioning (for conditional generation) + num_classes: int = 1000 # ImageNet classes + class_dropout_prob: float = 0.1 + + # Adaptive LayerNorm (DiT-specific) + use_adaptive_layernorm: bool = True + + def validate(self): + """Validate DiT-specific configuration.""" + # Call parent validation + super().validate() + + # DiT-specific validations + if self.num_layers <= 0: + raise ValueError(f"num_layers must be positive, got {self.num_layers}") + + if self.patch_size <= 0: + raise ValueError(f"patch_size must be positive, got {self.patch_size}") + + if self.num_classes < 0: + raise ValueError(f"num_classes must be non-negative, got {self.num_classes}") + + @classmethod + def dit_xl_2(cls, **kwargs): + """ + Create configuration for DiT-XL/2. + + Args: + **kwargs: Override default parameters + + Returns: + DiTConfig instance + """ + defaults = { + 'num_layers': 28, + 'hidden_size': 1152, + 'num_attention_heads': 16, + 'patch_size': 2, + } + defaults.update(kwargs) + return cls(**defaults) + + @classmethod + def dit_l_2(cls, **kwargs): + """Create configuration for DiT-L/2.""" + defaults = { + 'num_layers': 24, + 'hidden_size': 1024, + 'num_attention_heads': 16, + 'patch_size': 2, + } + defaults.update(kwargs) + return cls(**defaults) +``` + +### Step 3: Implement Model Class + +**File**: `model.py` + +```python +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""DiT model implementation.""" + +import torch +import torch.nn as nn +from primus.backends.megatron.core.models.common.diffusion_module.diffusion_module import DiffusionModule +from megatron.core.process_groups_config import ProcessGroupCollection +from ..common.layers import MMDiTLayer # Reuse shared components if applicable +from .config import DiTConfig + + +class DiT(DiffusionModule): + """ + DiT (Diffusion Transformer) model. + + Reference: "Scalable Diffusion Models with Transformers" (Peebles & Xie, 2023) + + Note: Inherits from DiffusionModule for Megatron-Core integration + (process groups, distributed checkpointing, attention backend config) + """ + + def __init__( + self, + config: DiTConfig, + encoder_configs: Optional[Dict[str, Any]] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + ): + """ + Initialize DiT model. + + Args: + config: DiT configuration + encoder_configs: Optional encoder configurations (VAE, T5, CLIP) + pg_collection: Process group collection for distributed training + """ + super().__init__(config, pg_collection=pg_collection, encoder_configs=encoder_configs) + + self.config = config + + # Input projection (patchify) + self.input_proj = nn.Conv2d( + config.in_channels, + config.hidden_size, + kernel_size=config.patch_size, + stride=config.patch_size, + ) + + # Positional embedding + self.pos_embed = nn.Parameter( + torch.zeros(1, (config.seq_length // config.patch_size) ** 2, config.hidden_size) + ) + + # Timestep embedding + self.time_embed = nn.Sequential( + nn.Linear(config.hidden_size, config.hidden_size * 4), + nn.SiLU(), + nn.Linear(config.hidden_size * 4, config.hidden_size), + ) + + # Class embedding (for conditional generation) + if config.num_classes > 0: + self.class_embed = nn.Embedding(config.num_classes, config.hidden_size) + + # Transformer blocks + self.blocks = nn.ModuleList([ + DiTBlock(config) for _ in range(config.num_layers) + ]) + + # Output projection + self.output_proj = nn.Sequential( + nn.LayerNorm(config.hidden_size), + nn.Linear(config.hidden_size, config.patch_size ** 2 * config.out_channels), + ) + + # Initialize weights + self._init_weights() + + def forward(self, x, timesteps, context=None, class_labels=None, **kwargs): + """ + Forward pass through DiT. + + Args: + x: Noisy latents [B, C, H, W] + timesteps: Diffusion timesteps [B] + context: Text conditioning [B, S, D] (optional) + class_labels: Class labels [B] (optional) + + Returns: + Model prediction [B, C, H, W] + """ + B, C, H, W = x.shape + + # Patchify input + x = self.input_proj(x) # [B, hidden_size, H/p, W/p] + x = x.flatten(2).transpose(1, 2) # [B, N, hidden_size] + + # Add positional embedding + x = x + self.pos_embed + + # Embed timesteps + t_emb = self.time_embed(self._timestep_embedding(timesteps)) # [B, hidden_size] + + # Embed class labels (if provided) + if class_labels is not None and self.config.num_classes > 0: + c_emb = self.class_embed(class_labels) # [B, hidden_size] + # Combine with timestep embedding + cond = t_emb + c_emb + else: + cond = t_emb + + # Transformer blocks + for block in self.blocks: + x = block(x, cond, context) + + # Output projection + x = self.output_proj(x) # [B, N, p^2 * C] + + # Unpatchify + x = self._unpatchify(x, H, W) # [B, C, H, W] + + return x + + target: Ground truth target [B, C, H, W] + + Returns: + Loss scalar + """ + # Simple MSE loss (can be extended) + loss = nn.functional.mse_loss(model_output, target) + return loss + + def _timestep_embedding(self, timesteps): + """Create sinusoidal timestep embeddings.""" + # Standard sinusoidal embedding + half_dim = self.config.hidden_size // 2 + emb = torch.exp( + -torch.arange(half_dim, device=timesteps.device) * + (torch.log(torch.tensor(10000.0)) / (half_dim - 1)) + ) + emb = timesteps[:, None] * emb[None, :] + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + return emb + + def _unpatchify(self, x, H, W): + """Convert patched tensor back to image.""" + p = self.config.patch_size + h = H // p + w = W // p + x = x.reshape(x.shape[0], h, w, p, p, self.config.out_channels) + x = x.permute(0, 5, 1, 3, 2, 4).contiguous() + x = x.reshape(x.shape[0], self.config.out_channels, H, W) + return x + + def _init_weights(self): + """Initialize model weights.""" + # Standard initialization (customize as needed) + pass + + +class DiTBlock(nn.Module): + """DiT transformer block with adaptive LayerNorm.""" + + def __init__(self, config): + super().__init__() + # Implementation details... + pass + + def forward(self, x, cond, context=None): + # Block forward pass + pass +``` + +### Step 4: Add Model-Specific Layers (If Needed) + +If your model has unique layers not shared with other models, add them to `layers.py`. + +### Step 5: Export Model + +**File**: `__init__.py` + +```python +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""DiT model implementation.""" + +from .config import DiTConfig +from .model import DiT + +__all__ = [ + 'DiTConfig', + 'DiT', +] +``` + +### Step 6: Create Configuration Files + +**File**: `primus/configs/models/megatron/diffusion/dit_xl_2.yaml` + +```yaml +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# DiT-XL/2 Configuration + +model_type: dit + +# Architecture: Layers +num_layers: 28 + +# Architecture: Dimensions +hidden_size: 1152 +num_attention_heads: 16 + +# Input/Output Channels +in_channels: 4 # VAE latent channels +out_channels: 4 + +# Patchification +patch_size: 2 + +# Class Conditioning +num_classes: 1000 # ImageNet classes +class_dropout_prob: 0.1 + +# Adaptive LayerNorm +use_adaptive_layernorm: true + +# Precision Settings +bf16: true +fp16: false + +# Training +seq_length: 4096 +micro_batch_size: 8 +global_batch_size: 256 +learning_rate: 1.0e-4 +``` + +### Step 7: Write Tests + +**File**: `tests/unit_tests/backends/megatron/diffusion/test_dit_model.py` + +```python +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Unit tests for DiT model.""" + +import pytest +import torch +from primus.backends.megatron.core.models.diffusion.dit import DiT, DiTConfig + + +class TestDiTConfig: + """Tests for DiT configuration.""" + + def test_dit_xl_2_factory(self): + """Test DiT-XL/2 configuration factory method.""" + config = DiTConfig.dit_xl_2() + + assert config.model_type == "dit" + assert config.num_layers == 28 + assert config.hidden_size == 1152 + assert config.num_attention_heads == 16 + assert config.patch_size == 2 + + def test_dit_config_validation(self): + """Test configuration validation.""" + config = DiTConfig.dit_xl_2() + config.validate() # Should not raise + + # Invalid configuration + with pytest.raises(ValueError): + config = DiTConfig(num_layers=-1) + config.validate() + + +class TestDiTModel: + """Tests for DiT model.""" + + def test_dit_initialization(self): + """Test DiT model initialization.""" + config = DiTConfig.dit_xl_2() + model = DiT(config) + + assert model is not None + assert isinstance(model, DiffusionModule) + + def test_dit_forward_shapes(self): + """Test forward pass produces correct output shapes.""" + config = DiTConfig.dit_xl_2() + model = DiT(config) + + batch_size = 2 + latents = torch.randn(batch_size, 4, 32, 32) + timesteps = torch.rand(batch_size) + class_labels = torch.randint(0, 1000, (batch_size,)) + + output = model(latents, timesteps, class_labels=class_labels) + + assert output.shape == latents.shape + +### Loss Computation + +Use standalone loss functions from `loss_computation.py` instead of implementing loss as a method: + +```python +from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss + +# In your forward_step_func +target = noise - clean_latents +loss = compute_flow_matching_loss(prediction, clean_latents, noise) +``` + +Models do NOT implement loss as a method. Loss computation is: +- Separate from model architecture +- Reusable across models +- Testable independently + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) +``` + +### Step 8: Update Documentation + +1. Add model to `README.md` supported models list +2. Update `architecture_overview.md` with model-specific details +3. Create model-specific training guide (e.g., `dit_training.md`) + +### Step 9: Add Example Scripts + +**File**: Use `examples/run_pretrain.sh` with appropriate config + +```python +#!/usr/bin/env python3 +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Training script for DiT model.""" + +import argparse +import yaml + +from primus.backends.megatron.core.models.diffusion.dit import DiT, DiTConfig +from primus.backends.megatron.training.diffusion.schedulers import DDPMScheduler +from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper +# ... other imports + + +def main(): + # Use examples/run_pretrain.sh with config from examples/megatron/configs/MI300X/diffusion/ + # MegatronDataloaderWrapper wraps an existing iterable (from dataset provider): + # dataloader = MegatronDataloaderWrapper(energon_loader_or_pytorch_loader) + # ... + + +if __name__ == "__main__": + main() +``` + +--- + +## Example: Adding DiT + +See the complete example in the step-by-step guide above. + +**Key Files Created**: +1. `core/models/diffusion/dit/config.py` - DiTConfig +2. `core/models/diffusion/dit/model.py` - DiT model +3. `configs/models/megatron/diffusion/dit_xl_2.yaml` - Config file +4. `tests/unit_tests/backends/megatron/diffusion/test_dit_model.py` - Tests +5. `examples/run_pretrain.sh` - Use with config from `examples/megatron/configs/MI300X/diffusion/` + +--- + +## Testing Your Model + +### Unit Tests + +Run tests to verify implementation: + +```bash +# Run all DiT tests +pytest tests/unit_tests/backends/megatron/diffusion/test_dit_model.py -v + +# Run specific test +pytest tests/unit_tests/backends/megatron/diffusion/test_dit_model.py::TestDiTModel::test_dit_forward_shapes -v +``` + +### Integration Tests + +Test with actual data: + +```bash +# Small dataset test +./examples/run_pretrain.sh --config examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml +``` + +### Validation + +Compare with reference implementation: +1. Load reference weights +2. Run same inputs through both models +3. Compare outputs (should match within tolerance) + +--- + +## Best Practices + +### 1. Code Organization +- ✅ Separate configuration from model code +- ✅ Reuse shared components from `common/` +- ✅ Keep model-specific code minimal +- ✅ Follow existing naming conventions + +### 2. Configuration +- ✅ Extend `BaseDiffusionConfig` +- ✅ Add factory methods for common sizes +- ✅ Implement validation +- ✅ Document all parameters + +### 3. Model Implementation +- ✅ Extend `DiffusionModule` from `primus.backends.megatron.core.models.common.diffusion_module.diffusion_module` +- ✅ Implement required method: `forward()` +- ✅ Use standalone loss functions from `loss_computation.py` +- ✅ Add comprehensive docstrings +- ✅ Use type hints +- ✅ Include `pg_collection` parameter in `__init__` for distributed training + +### 4. Testing +- ✅ Test configuration validation +- ✅ Test model initialization +- ✅ Test forward pass shapes +- ✅ Test loss computation +- ✅ Test with mock data first + +### 5. Documentation +- ✅ Update README with new model +- ✅ Document architecture specifics +- ✅ Provide usage examples +- ✅ Reference original paper + +### 6. Performance +- ✅ Profile memory usage +- ✅ Optimize critical paths +- ✅ Support mixed precision +- ✅ Enable gradient checkpointing + +--- + +## Common Pitfalls + +### 1. Import Errors +❌ **Wrong**: Absolute imports +```python +from primus.backends.megatron.core.models.common.diffusion_module.diffusion_module import DiffusionModule +``` + +✅ **Right**: Relative imports +```python +from ...common.diffusion_module.diffusion_module import DiffusionModule +``` + +### 2. Configuration Validation +❌ **Wrong**: No validation +```python +class DiTConfig(BaseDiffusionConfig): + pass # No validation +``` + +✅ **Right**: Validate parameters +```python +def validate(self): + super().validate() + if self.num_layers <= 0: + raise ValueError(f"num_layers must be positive") +``` + +### 3. Shape Mismatches +❌ **Wrong**: Assuming fixed shapes +```python +def forward(self, x): + # Assumes x is always [B, 4, 32, 32] + pass +``` + +✅ **Right**: Handle variable shapes +```python +def forward(self, x): + B, C, H, W = x.shape + # Handle any valid shape + pass +``` + +### 4. Testing +❌ **Wrong**: No tests +```python +# Just implement and hope it works +``` + +✅ **Right**: Comprehensive tests +```python +def test_forward_shapes(self): + # Test various input shapes + pass +``` + +--- + +## Checklist + +Before submitting your new model: + +- [ ] Configuration class implemented and validated +- [ ] Model class extends `DiffusionModule` +- [ ] Required method implemented: `forward()` +- [ ] Loss computation uses standalone functions from `loss_computation.py` +- [ ] Process group support added (`pg_collection` parameter) +- [ ] YAML configuration files created +- [ ] Unit tests written and passing +- [ ] Integration tests run successfully +- [ ] Documentation updated +- [ ] Example training script provided +- [ ] Code follows Primus style guide +- [ ] No linter errors +- [ ] PR description includes architecture details + +--- + +## Getting Help + +If you encounter issues: +1. Review existing models (Flux) for patterns +2. Check documentation in `docs/backends/megatron/diffusion/` +3. Run tests in debug mode: `pytest --pdb` +4. Consult architecture overview for design principles + +--- + +**Last Updated**: December 2025 diff --git a/docs/backends/megatron/diffusion/api_reference.md b/docs/backends/megatron/diffusion/api_reference.md new file mode 100644 index 000000000..f2f84651a --- /dev/null +++ b/docs/backends/megatron/diffusion/api_reference.md @@ -0,0 +1,1059 @@ +# Flux Model API Reference + +## Overview + +This document provides comprehensive API reference for the Flux diffusion model implementation in Primus. Flux is a flow-based diffusion model that uses MMDiT (Multimodal Diffusion Transformer) architecture for high-quality text-to-image generation. + +--- + +## Base Classes + +### DiffusionModule + +**Location**: `primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py` + +Base class for all diffusion models, providing Megatron-Core integration. + +```python +from primus.backends.megatron.core.models.common.diffusion_module import DiffusionModule +``` + +**Key Features**: +- Process group management (TP, PP, CP, DP) +- Attention backend configuration +- Distributed checkpointing support +- Common loss computation utilities + +**Inherited Methods** (available to all diffusion models): +- `get_num_params()` - Count trainable/total parameters +- `set_requires_grad()` - Freeze/unfreeze model +- `compute_diffusion_loss()` - Common loss helper (MSE, MAE, Huber) +- `sharded_state_dict()` - Distributed checkpointing + +> Note: in-model encoder loading (`load_encoders()`, `get_encoder_by_type()`, +> `get_encoders_by_type()`) is not implemented and raises `NotImplementedError`. +> Encode VAE/T5/CLIP inputs via the offline diffusion preprocessing pipeline +> (`primus.backends.megatron.data.diffusion.preprocessing`) instead. + +--- + +## Model Architecture + +### Flux Class + +**Location**: `primus/backends/megatron/core/models/diffusion/flux/model.py` + +```python +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig + +# Create model +config = FluxConfig.flux_535m() +model = Flux(config) +``` + +#### Architecture Diagram + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ Flux Model │ +├─────────────────────────────────────────────────────────────────┤ +│ │ +│ Input Processing: │ +│ ┌────────────┐ ┌────────────┐ │ +│ │ Image │──┐ ┌──│ Text │ │ +│ │ Latents │ │ │ │ Embeddings │ │ +│ │ [B,64,H,W] │ │ │ │ [B,S,4096] │ │ +│ └────────────┘ │ │ └────────────┘ │ +│ ▼ ▼ │ +│ ┌──────────────┐ │ +│ │ Linear Embed │ │ +│ └──────┬───────┘ │ +│ │ [B,seq,3072] │ +│ ▼ │ +│ ┌──────────────┐ │ +│ │ 3D RoPE │ │ +│ │ Position Emb │ │ +│ └──────┬───────┘ │ +│ │ │ +│ Conditioning: │ │ +│ ┌────────────┐ │ │ +│ │ Timestep │──┼──┐ │ +│ │ Embedding │ │ │ │ +│ └────────────┘ │ │ │ +│ ┌────────────┐ │ │ │ +│ │ CLIP │──┼──┤ │ +│ │ Pooled │ │ │ vec_emb [B,3072] │ +│ └────────────┘ │ │ │ +│ ┌────────────┐ │ │ │ +│ │ Guidance │──┼──┘ │ +│ │ (optional) │ │ │ +│ └────────────┘ │ │ +│ │ │ +│ Double Blocks │ │ +│ (Joint): │ │ +│ ┌────────────────▼──────────┐ │ +│ │ MMDiTLayer x N_joint │ N_joint = 1 (535M), 19 (12B) │ +│ │ ┌──────────────────────┐ │ │ +│ │ │ Joint Self-Attention │ │ (image + text together) │ +│ │ └──────────────────────┘ │ │ +│ │ ┌──────────────────────┐ │ │ +│ │ │ Image MLP │ │ │ +│ │ └──────────────────────┘ │ │ +│ │ ┌──────────────────────┐ │ │ +│ │ │ Text MLP │ │ │ +│ │ └──────────────────────┘ │ │ +│ └────────────┬──────────────┘ │ +│ │ │ +│ Single Blocks│ │ +│ (Combined): │ │ +│ ┌────────────▼──────────────┐ │ +│ │ FluxSingleTransformer │ N_single = 1 (535M), 38 (12B) │ +│ │ Block x N_single │ │ +│ │ ┌──────────────────────┐ │ │ +│ │ │ Self-Attention │ │ (image + text concatenated) │ +│ │ └──────────────────────┘ │ │ +│ │ ┌──────────────────────┐ │ │ +│ │ │ MLP │ │ │ +│ │ └──────────────────────┘ │ │ +│ └────────────┬──────────────┘ │ +│ │ (extract image tokens) │ +│ ▼ │ +│ Output Processing: │ +│ ┌────────────────────────┐ │ +│ │ AdaLNContinuous │ │ +│ │ (timestep conditioned) │ │ +│ └────────────┬───────────┘ │ +│ ▼ │ +│ ┌────────────────────────┐ │ +│ │ Linear Projection │ │ +│ └────────────┬───────────┘ │ +│ ▼ │ +│ [B,64,H,W] │ +│ Predicted Velocity │ +└─────────────────────────────────────────────────────────────────┘ +``` + +#### Parameter Counts + +| Variant | Joint Layers | Single Layers | Total Parameters | Use Case | +|---------|-------------|---------------|------------------|----------| +| Flux 535M | 1 | 1 | ~535 million | Testing, debugging, prototyping | +| Flux 12B | 19 | 38 | ~12 billion | Production training | + +--- + +## Configuration + +### FluxConfig + +**Location**: `primus/backends/megatron/core/models/diffusion/flux/config.py` + +Complete configuration class for Flux models, inheriting from `BaseDiffusionConfig`. + +#### Key Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `num_joint_layers` | int | 19 | Number of joint (MMDiT) transformer layers | +| `num_single_layers` | int | 38 | Number of single transformer layers | +| `hidden_size` | int | 3072 | Hidden dimension size | +| `num_attention_heads` | int | 24 | Number of attention heads | +| `in_channels` | int | 64 | Input channels (VAE latent dimension) | +| `context_dim` | int | 4096 | Text context dimension (T5-XXL) | +| `vec_in_dim` | int | 768 | Vector input dimension (CLIP pooled) | +| `model_channels` | int | 256 | Channels for timestep embedding | +| `guidance_embed` | bool | False | Enable guidance embedding for CFG | +| `guidance_scale` | float | 3.5 | Guidance scale for classifier-free guidance | +| `theta` | int | 10000 | Base for RoPE frequency computation | +| `axes_dim` | tuple | (16, 56, 56) | Dimensions for 3D RoPE axes | +| `patch_size` | int | 1 | Patch size for image tokens | +| `add_qkv_bias` | bool | True | Add bias to QKV projections | +| `rotary_interleaved` | bool | True | Interleave RoPE dimensions | +| `layernorm_epsilon` | float | 1e-6 | Epsilon for layer normalization | +| `hidden_dropout` | float | 0.0 | Hidden layer dropout rate | +| `attention_dropout` | float | 0.0 | Attention dropout rate | + +#### Configuration Examples + +**Flux 535M (Testing)**: +```python +config = FluxConfig.flux_535m() +# Equivalent to: +config = FluxConfig( + num_joint_layers=1, + num_single_layers=1, + hidden_size=3072, + num_attention_heads=24, +) +``` + +**Flux 12B (Production)**: +```python +config = FluxConfig.flux_12b() +# Equivalent to: +config = FluxConfig( + num_joint_layers=19, + num_single_layers=38, + hidden_size=3072, + num_attention_heads=24, +) +``` + +**Custom Configuration**: +```python +config = FluxConfig( + num_joint_layers=4, + num_single_layers=8, + hidden_size=2048, + num_attention_heads=16, + guidance_embed=True, + patch_size=2, +) +``` + +--- + +## Components API + +### Embeddings + +#### TimeStepEmbedder + +**Location**: `primus/backends/megatron/core/models/diffusion/common/embeddings.py` + +Converts scalar timesteps to high-dimensional embeddings using sinusoidal encoding. + +```python +from primus.backends.megatron.core.models.diffusion.common.embeddings import TimeStepEmbedder + +embedder = TimeStepEmbedder(embedding_dim=256, hidden_dim=3072) +timesteps = torch.tensor([0, 100, 500, 999]) # [B] +t_emb = embedder(timesteps) # [B, 3072] +``` + +**Input**: Timesteps [B] in range [0, 1000] +**Output**: Embeddings [B, hidden_dim] + +#### MLPEmbedder + +Embeds vector conditioning (e.g., CLIP pooled embeddings) via 2-layer MLP. + +```python +from primus.backends.megatron.core.models.diffusion.common.embeddings import MLPEmbedder + +embedder = MLPEmbedder(in_dim=768, hidden_dim=3072) +clip_pooled = torch.randn(4, 768) # [B, 768] +embedded = embedder(clip_pooled) # [B, 3072] +``` + +**Input**: Vectors [B, in_dim] +**Output**: Embeddings [B, hidden_dim] + +--- + +### Normalization + +#### AdaLN (Adaptive Layer Normalization) + +**Location**: `primus/backends/megatron/core/models/diffusion/common/normalization.py` + +Applies layer normalization conditioned on timestep embeddings. + +```python +from megatron.core.transformer.transformer_config import TransformerConfig +from primus.backends.megatron.core.models.diffusion.common.normalization import AdaLN + +config = TransformerConfig(hidden_size=3072) +adaln = AdaLN(config, n_adaln_chunks=6) + +timestep_emb = torch.randn(4, 3072) +shift, scale, gate, shift_mlp, scale_mlp, gate_mlp = adaln(timestep_emb) +# Each output: [B, 3072] +``` + +**Methods**: +- `forward(timestep_emb)`: Generate modulation parameters +- `modulate(x, shift, scale)`: Apply adaptive modulation +- `scale_add(residual, x, gate)`: Gated residual addition + +#### AdaLNContinuous + +Continuous variant of AdaLN for Flux output normalization. + +```python +from primus.backends.megatron.core.models.diffusion.common.normalization import AdaLNContinuous + +adaln = AdaLNContinuous(config, conditioning_embedding_dim=3072) +x = torch.randn(4, 256, 3072) # [B, seq, hidden] +cond = torch.randn(4, 3072) # [B, cond_dim] +x_norm = adaln(x, cond) # [B, seq, hidden] +``` + +#### RMSNorm + +Root Mean Square Layer Normalization (simpler, faster than LayerNorm). + +```python +from primus.backends.megatron.core.models.diffusion.common.normalization import RMSNorm + +norm = RMSNorm(hidden_size=3072) +x = torch.randn(4, 256, 3072) +x_norm = norm(x) +``` + +--- + +### Position Embeddings + +#### EmbedND (3D RoPE) + +**Location**: `primus/backends/megatron/core/models/diffusion/flux/layers.py` + +Multi-dimensional Rotary Position Embedding for image patches. + +```python +from primus.backends.megatron.core.models.diffusion.flux.layers import ( + EmbedND, +) +from primus.backends.megatron.core.models.diffusion.flux.utils import ( + generate_image_position_ids, +) + +# Initialize +embed_nd = EmbedND(dim=3072, theta=10000, axes_dim=[16, 56, 56]) + +# Generate position IDs for 56x56 image patches +batch_size = 2 +height, width = 112, 112 # Unpacked dimensions (56*2, 56*2) +img_ids = generate_image_position_ids(batch_size, height, width) +# img_ids: [B, H*W/4, 3] where dimension 0 is always 0 + +# Get RoPE frequencies +rope_freqs = embed_nd(img_ids) # [3, B, H*W, 3072] +``` + +**Axes**: +- Axis 0: Channel groups (16 for 64 channels) +- Axis 1: Height positions (56 for 1024px image) +- Axis 2: Width positions (56 for 1024px image) + +--- + +### Attention Mechanisms + +#### JointSelfAttention + +**Location**: `primus/backends/megatron/core/models/diffusion/flux/attention.py` + +Joint attention over image and text tokens (MMDiT architecture). + +```python +from primus.backends.megatron.core.models.diffusion.flux.attention import ( + JointSelfAttention, + JointSelfAttentionSubmodules, +) + +submodules = JointSelfAttentionSubmodules(...) +joint_attn = JointSelfAttention(config, submodules, layer_number=0) + +# Forward +img_tokens = torch.randn(3136, 2, 3072) # [seq_img, B, hidden] +txt_tokens = torch.randn(512, 2, 3072) # [seq_txt, B, hidden] +img_out, txt_out = joint_attn( + img_tokens, + attention_mask=None, + additional_hidden_states=txt_tokens, +) +``` + +**Input**: +- `hidden_states`: Image tokens [seq_img, B, hidden] +- `additional_hidden_states`: Text tokens [seq_txt, B, hidden] + +**Output**: Tuple of (img_output, txt_output) + +#### FluxSingleAttention + +Single-stream self-attention for image tokens only. + +```python +from primus.backends.megatron.core.models.diffusion.flux.attention import FluxSingleAttention + +single_attn = FluxSingleAttention(config, submodules, layer_number=0) + +img_tokens = torch.randn(3136, 2, 3072) +output = single_attn(img_tokens, attention_mask=None) +``` + +--- + +### Layer Specifications + +#### MMDiTLayer + +**Location**: `primus/backends/megatron/core/models/diffusion/flux/layer_spec.py` + +Joint image-text transformer block. + +```python +from primus.backends.megatron.core.models.diffusion.flux.layer_spec import ( + MMDiTLayer, + get_flux_double_transformer_spec_for_backend, +) + +# Using factory function (recommended) +spec = get_flux_double_transformer_spec_for_backend(backend) +mmdit_layer = MMDiTLayer( + config=config, + submodules=spec.submodules, + layer_number=0, +) + +# Forward +img_tokens = torch.randn(3136, 2, 3072) +txt_tokens = torch.randn(512, 2, 3072) +emb = torch.randn(2, 3072) +img_out, txt_out = mmdit_layer(img_tokens, txt_tokens, emb=emb) +``` + +#### FluxSingleTransformerBlock + +Image-only transformer block. + +```python +from primus.backends.megatron.core.models.diffusion.flux.layer_spec import ( + FluxSingleTransformerBlock, + get_flux_single_transformer_spec_for_backend, +) + +spec = get_flux_single_transformer_spec_for_backend(backend) +single_block = FluxSingleTransformerBlock( + config=config, + submodules=spec.submodules, + layer_number=0, +) + +# Forward +img_tokens = torch.randn(3136, 2, 3072) +emb = torch.randn(2, 3072) +output, _ = single_block(img_tokens, emb=emb) +``` + +--- + +## Training Utilities + +### Noise Application + +**Location**: `primus/backends/megatron/training/diffusion/noise_utils.py` + +Pure functions for applying noise according to different diffusion forward processes. + +#### apply_flow_matching_noise() + +```python +from primus.backends.megatron.training.diffusion.noise_utils import apply_flow_matching_noise + +clean_latents = torch.randn(2, 16, 64, 64) +noise = torch.randn(2, 16, 64, 64) +sigma = torch.tensor([0.3, 0.7]).reshape(2, 1, 1, 1) + +noisy = apply_flow_matching_noise(clean_latents, noise, sigma) +# Formula: noisy = (1 - sigma) * clean + sigma * noise +``` + +**Parameters**: +- `clean_latents` (Tensor): Clean latents [any shape] +- `noise` (Tensor): Sampled noise [same shape as clean_latents] +- `sigma` (Tensor): Noise schedule values [broadcast compatible], range [0, 1] + +**Returns**: Noisy latents [same shape as clean_latents] + +**Reference**: [Flow Matching for Generative Modeling](https://arxiv.org/abs/2210.02747) + +#### apply_ddpm_noise() + +```python +from primus.backends.megatron.training.diffusion.noise_utils import apply_ddpm_noise + +clean = torch.randn(2, 3, 256, 256) +noise = torch.randn(2, 3, 256, 256) +alpha_bar = torch.tensor([0.9, 0.5]).reshape(2, 1, 1, 1) + +noisy = apply_ddpm_noise(clean, noise, alpha_bar) +# Formula: noisy = sqrt(alpha_bar) * clean + sqrt(1 - alpha_bar) * noise +``` + +**Parameters**: +- `clean_latents` (Tensor): Clean latents +- `noise` (Tensor): Sampled noise (same shape) +- `alpha_bar` (Tensor): Cumulative product of alphas, range (0, 1] + +**Returns**: Noisy latents + +**Reference**: [Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239) + +--- + +### Loss Computation + +**Location**: `primus/backends/megatron/training/diffusion/loss_computation.py` + +Reusable loss computation logic for different diffusion training objectives. + +#### compute_flow_matching_loss() + +```python +from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss + +prediction = torch.randn(2, 16, 64, 64) # Model output +clean = torch.randn(2, 16, 64, 64) +noise = torch.randn(2, 16, 64, 64) + +loss = compute_flow_matching_loss(prediction, clean, noise) +# Formula: target = noise - clean +# loss = MSE(prediction, target) +``` + +**Parameters**: +- `prediction` (Tensor): Model output (predicted velocity) [any shape] +- `clean_latents` (Tensor): Original clean latents [same shape] +- `noise` (Tensor): Sampled noise [same shape] + +**Returns**: Scalar loss value (mean squared error) + +**Used by**: Flux, SD3, video models + +#### compute_epsilon_loss() + +```python +from primus.backends.megatron.training.diffusion.loss_computation import compute_epsilon_loss + +prediction = torch.randn(2, 3, 256, 256) +noise = torch.randn(2, 3, 256, 256) + +loss = compute_epsilon_loss(prediction, noise) +# Formula: loss = MSE(prediction, noise) +``` + +**Parameters**: +- `prediction` (Tensor): Model output (predicted noise) +- `noise` (Tensor): Sampled noise (ground truth) + +**Returns**: Scalar loss value + +**Used by**: DDPM and older models + +#### compute_v_prediction_loss() + +```python +from primus.backends.megatron.training.diffusion.loss_computation import compute_v_prediction_loss + +prediction = torch.randn(2, 16, 64, 64) +clean = torch.randn(2, 16, 64, 64) +noise = torch.randn(2, 16, 64, 64) +sigma = torch.tensor([0.3, 0.7]).reshape(2, 1, 1, 1) + +loss = compute_v_prediction_loss(prediction, clean, noise, sigma) +# Formula: v = sigma * noise - (1 - sigma) * clean +# loss = MSE(prediction, v) +``` + +**Parameters**: +- `prediction` (Tensor): Model output +- `clean_latents` (Tensor): Clean latents +- `noise` (Tensor): Sampled noise +- `sigma` (Tensor): Noise schedule values [broadcast compatible] + +**Returns**: Scalar loss value + +**Reference**: [Progressive Distillation for Fast Sampling](https://arxiv.org/abs/2202.00512) + +--- + +### Timestep Sampling + +**Location**: `primus/backends/megatron/training/diffusion/timestep_sampling.py` + +Sampling strategies for training timesteps (hyperparameter optimization, separate from inference). + +#### LogitNormalSampler + +```python +from primus.backends.megatron.training.diffusion.timestep_sampling import LogitNormalSampler + +sampler = LogitNormalSampler(mean=0.0, std=1.0) +timesteps, sigmas = sampler.sample( + batch_size=32, + device='cuda', + scheduler=flow_scheduler +) +``` + +**Description**: Logit-normal distribution sampling, emphasizes boundary timesteps (t≈0 and t≈1000). + +**Parameters**: +- `mean` (float): Mean of normal distribution (default: 0.0) +- `std` (float): Standard deviation (default: 1.0) + +**Returns**: Tuple of (timesteps [B], sigmas [B]) + +**Reference**: [Stable Diffusion 3](https://arxiv.org/abs/2403.03206v1), Section 3.1 + +**Used by**: Flux, SD3 + +#### UniformSampler + +```python +from primus.backends.megatron.training.diffusion.timestep_sampling import UniformSampler + +sampler = UniformSampler() +timesteps, sigmas = sampler.sample(batch_size=32, device='cuda', scheduler=flow_scheduler) +``` + +**Description**: Uniform timestep sampling (baseline approach for comparison). + +**Returns**: Tuple of (timesteps [B], sigmas [B]) + +#### ModeSampler + +```python +from primus.backends.megatron.training.diffusion.timestep_sampling import ModeSampler + +sampler = ModeSampler(mode_scale=1.29) +timesteps, sigmas = sampler.sample(batch_size=32, device='cuda', scheduler=flow_scheduler) +``` + +**Description**: Mode-based sampling from SD3 paper (alternative to logit-normal). + +**Parameters**: +- `mode_scale` (float): Scaling factor (default: 1.29 from SD3 paper) + +**Returns**: Tuple of (timesteps [B], sigmas [B]) + +#### create_timestep_sampler() + +```python +from primus.backends.megatron.training.diffusion.timestep_sampling import create_timestep_sampler + +# Factory function for easy experimentation +sampler = create_timestep_sampler("logit_normal", mean=0.0, std=1.0) +sampler = create_timestep_sampler("uniform") +sampler = create_timestep_sampler("mode", mode_scale=1.5) +``` + +**Parameters**: +- `strategy` (str): Sampling strategy ("logit_normal", "uniform", "mode") +- `**kwargs`: Additional arguments for the sampler + +**Returns**: TimestepSampler instance + +--- + +### Training Workflow Example + +Complete training step using the utilities: + +```python +import torch +from torch.optim import AdamW +from primus.backends.megatron.core.models.diffusion.flux import Flux, FluxConfig +from primus.backends.megatron.training.diffusion.noise_utils import apply_flow_matching_noise +from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss +from primus.backends.megatron.training.diffusion.timestep_sampling import LogitNormalSampler +from primus.backends.megatron.training.diffusion.schedulers.flow_match_euler import ( + FlowMatchEulerDiscreteScheduler +) + +# Setup +config = FluxConfig.flux_535m() +model = Flux(config).cuda() +optimizer = AdamW(model.parameters(), lr=1e-4) + +# Initialize scheduler and timestep sampler +scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000) +timestep_sampler = LogitNormalSampler(mean=0.0, std=1.0) + +# Training loop +for batch in dataloader: + clean_latents = batch['latents'].cuda() # [B, 16, 64, 64] + txt_embeddings = batch['text'].cuda() # [B, 512, 4096] + clip_pooled = batch['clip'].cuda() # [B, 768] + + batch_size = clean_latents.shape[0] + + # 1. Sample timesteps + timesteps, sigmas = timestep_sampler.sample( + batch_size=batch_size, + device='cuda', + scheduler=scheduler + ) + + # 2. Sample noise + noise = torch.randn_like(clean_latents) + + # 3. Apply noise + sigma_reshaped = sigmas.view(-1, 1, 1, 1) + noisy_latents = apply_flow_matching_noise(clean_latents, noise, sigma_reshaped) + + # 4. Prepare position IDs + img_ids = generate_image_position_ids(batch_size, 128, 128).cuda() + txt_ids = torch.zeros(batch_size, 512, 3).cuda() + + # 5. Forward pass + predicted_velocity = model( + img=noisy_latents, + txt=txt_embeddings, + y=clip_pooled, + timesteps=sigmas, + img_ids=img_ids, + txt_ids=txt_ids, + ) + + # 6. Compute loss + loss = compute_flow_matching_loss(predicted_velocity, clean_latents, noise) + + # 7. Backward and optimize + optimizer.zero_grad() + loss.backward() + optimizer.step() + + print(f"Step {step}, Loss: {loss.item():.4f}") +``` + +--- + +## Usage Examples + +### Basic Inference + +```python +import torch +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.utils import generate_image_position_ids + +# 1. Setup model +config = FluxConfig.flux_535m() +model = Flux(config) +model.eval() + +# 2. Prepare inputs +batch_size = 1 +img_latents = torch.randn(batch_size, 64, 128, 128) # VAE latents +txt_embeddings = torch.randn(batch_size, 512, 4096) # T5-XXL embeddings +clip_pooled = torch.randn(batch_size, 768) # CLIP-L pooled +timesteps = torch.tensor([0.5]) # Diffusion timestep [0, 1] + +# 3. Generate position IDs +img_ids = generate_image_position_ids(batch_size, 256, 256) +txt_ids = torch.zeros(batch_size, 512, 3) + +# 4. Forward pass +with torch.no_grad(): + predicted_velocity = model( + img=img_latents, + txt=txt_embeddings, + y=clip_pooled, + timesteps=timesteps, + img_ids=img_ids, + txt_ids=txt_ids, + ) + +print(f"Output shape: {predicted_velocity.shape}") # [1, 64, 128, 128] +``` + +### Training Step + +```python +import torch +from torch.optim import AdamW + +# Setup +config = FluxConfig.flux_535m() +model = Flux(config) +model.train() + +optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.01) + +# Prepare batch +batch_size = 4 +img = torch.randn(batch_size, 64, 128, 128) +txt = torch.randn(batch_size, 512, 4096) +y = torch.randn(batch_size, 768) +timesteps = torch.rand(batch_size) + +img_ids = generate_image_position_ids(batch_size, 256, 256) +txt_ids = torch.zeros(batch_size, 512, 3) + +# Add noise for flow matching +original = img.clone() +noise = torch.randn_like(img) +t = timesteps.view(-1, 1, 1, 1) +noisy_img = original + noise * t + +# Velocity target for flow matching +velocity_target = noise - original + +# Training step +optimizer.zero_grad() + +# Forward pass +output = model(noisy_img, txt, y, timesteps, img_ids, txt_ids) + +# Loss computation (using standalone function) +from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss +target = noise - clean_latents +loss = compute_flow_matching_loss(output, clean_latents, noise) + +# Backward pass +loss.backward() + +# Optimizer step +optimizer.step() + +print(f"Loss: {loss.item():.4f}") +``` + +### With Guidance (Classifier-Free Guidance) + +```python +# Enable guidance in config +config = FluxConfig.flux_535m(guidance_embed=True) +model = Flux(config) +model.eval() + +# Prepare inputs with guidance +guidance_scale = torch.tensor([3.5]) # Typical guidance scale + +with torch.no_grad(): + output = model( + img=img_latents, + txt=txt_embeddings, + y=clip_pooled, + timesteps=timesteps, + img_ids=img_ids, + txt_ids=txt_ids, + guidance=guidance_scale, # Add guidance + ) +``` + +### Different Resolutions + +```python +# Flux can handle different resolutions +resolutions = [ + (64, 64), # 512x512 pixels + (128, 128), # 1024x1024 pixels + (192, 192), # 1536x1536 pixels +] + +for height, width in resolutions: + img = torch.randn(1, 64, height, width) + img_ids = generate_image_position_ids(1, height, width) + + with torch.no_grad(): + output = model(img, txt, y, timesteps, img_ids, txt_ids) + + assert output.shape == img.shape +``` + +--- + +## Methods Reference + +### Flux.forward() + +```python +def forward( + self, + img: Tensor, # [B, C, H, W] Image latents from VAE + txt: Tensor, # [B, S_txt, D_txt] T5-XXL embeddings + y: Tensor, # [B, D_pool] CLIP pooled embeddings + timesteps: Tensor, # [B] Timesteps in [0, 1] + img_ids: Tensor, # [B, H*W, 3] Image position IDs + txt_ids: Tensor, # [B, S_txt, 3] Text position IDs + guidance: Optional[Tensor] = None, # [B] Guidance scale + controlnet_double_block_samples: Optional[Tensor] = None, + controlnet_single_block_samples: Optional[Tensor] = None, +) -> Tensor: # Returns: [B, C, H, W] Predicted velocity +``` + +## Loss Computation + +Loss is computed using standalone functions from `loss_computation.py`: + +### compute_flow_matching_loss() + +```python +from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss + +def compute_flow_matching_loss( + prediction: Tensor, # [any shape] Model prediction + clean_latents: Tensor, # [same shape] Original clean latents + noise: Tensor, # [same shape] Sampled noise +) -> Tensor: # Returns: Scalar loss +``` + +### Flux.get_num_params() + +```python +def get_num_params( + self, + trainable_only: bool = True, +) -> int: # Returns: Number of parameters +``` + +--- + +## Testing + +### Running Tests + +```bash +# Run all Flux tests +pytest tests/unit_tests/backends/megatron/diffusion/ -v + +# Run specific test files +pytest tests/unit_tests/backends/megatron/diffusion/test_flux_embeddings.py -v +pytest tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py -v +pytest tests/unit_tests/backends/megatron/diffusion/test_flux_config.py -v +pytest tests/unit_tests/backends/megatron/diffusion/test_flux_layers.py -v +pytest tests/unit_tests/backends/megatron/diffusion/test_flux_model.py -v +pytest tests/unit_tests/backends/megatron/diffusion/test_flux_layer_spec_backend_selection.py -v + +# Run with coverage +pytest tests/unit_tests/backends/megatron/diffusion/ --cov=primus.backends.megatron.core.models.diffusion --cov-report=html +``` + +### Test Coverage + +- **Component tests**: 80+ tests covering all components +- **Model tests**: 17+ tests for full model +- **Integration tests**: 11+ tests for complete workflows + +--- + +## Performance Considerations + +### Memory Usage + +| Configuration | Model Weights | Training (bf16) | Training (fp32) | +|--------------|---------------|-----------------|-----------------| +| Flux 535M | ~2 GB | ~6-8 GB | ~10-12 GB | +| Flux 12B | ~24 GB | ~60-80 GB | ~100-120 GB | + +### Throughput (Estimated) + +On MI300X (192GB): +- **Flux 535M**: ~5-10 samples/sec (depends on resolution) +- **Flux 12B**: ~0.5-1 samples/sec (requires multi-GPU) + +### Optimization Tips + +1. **Use mixed precision**: `torch.autocast(device_type='cuda', dtype=torch.bfloat16)` +2. **Enable CUDA graphs**: Set `enable_cuda_graph=True` in config +3. **Use Transformer Engine**: Automatically used with factory functions +4. **Gradient checkpointing**: Can be enabled for memory savings + +--- + +## Common Issues + +### Issue: Import Errors + +```python +# Old import style +from some_other_library import Flux + +# Correct Primus import +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +``` + +### Issue: Position ID Shape Mismatch + +```python +# Position IDs must be [B, seq, 3] for 3D RoPE +img_ids = generate_image_position_ids(batch_size, height, width) +# Not: img_ids = torch.randn(batch_size, height * width, 2) # Wrong! +``` + +### Issue: Timestep Range + +```python +# Flux expects timesteps in [0, 1] +timesteps = torch.rand(batch_size) # Correct: [0, 1] +# Not: timesteps = torch.randint(0, 1000, (batch_size,)) # Wrong range! +``` + +--- + +## API Compatibility + +### Primus Architecture Features + +| Aspect | Primus Implementation | +|--------|----------------------| +| Import path | `primus.backends.megatron.core.models.diffusion.flux` | +| Base class | `DiffusionModule` (extends MegatronModule) | +| Config parent | `BaseDiffusionConfig` (extends TransformerConfig) | +| Layer organization | Unified `TransformerBlock` with heterogeneous specs | +| Checkpoint format | `transformer.layers.{0-56}` unified namespace | +| Process groups | Via `pg_collection` parameter | + +### Key Design Choices + +**TransformerBlock Architecture**: +- Primus uses Megatron-Core's `TransformerBlock` with heterogeneous layer specifications +- Unified checkpoint format for simpler distributed training +- Note: pipeline parallelism is not supported for diffusion models (`pipeline_model_parallel_size` must be 1) + +**Example Usage**: + +```python +# Primus native approach +from primus.backends.megatron.core.models.diffusion.flux import Flux, FluxConfig + +# Primus +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +config = FluxConfig() +model = Flux(config) +``` + +For more advanced examples, see `examples/run_pretrain.sh`. + +--- + +## References + +### Papers +- **Flux**: "Flux: A Scalable Diffusion Model for High-Resolution Image Synthesis" +- **MMDiT**: "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis" +- **Flow Matching**: "Flow Matching for Generative Modeling" +- **RoPE**: "RoFormer: Enhanced Transformer with Rotary Position Embedding" +- **DiT**: "Scalable Diffusion Models with Transformers" (Peebles & Xie, 2023) + +### Source Code +- **Primus Implementation**: `primus/backends/megatron/core/models/diffusion/flux/` +- **Megatron-Core**: `megatron/core/transformer/` +- **Official Flux**: Black Forest Labs (HuggingFace) + +--- + +## Version Information + +- **Primus Version**: Current +- **Megatron-Core Version**: Latest +- **Transformer Engine**: Optional (recommended for performance) + +--- + +## Support + +For issues or questions: +1. Check this API reference +2. Read architecture guide: `docs/backends/megatron/diffusion/flux_architecture.md` +3. See examples in docstrings +4. Check test files for usage patterns diff --git a/docs/backends/megatron/diffusion/architecture_overview.md b/docs/backends/megatron/diffusion/architecture_overview.md new file mode 100644 index 000000000..c58ffe47b --- /dev/null +++ b/docs/backends/megatron/diffusion/architecture_overview.md @@ -0,0 +1,497 @@ +# Architecture Overview + +This document provides a detailed overview of the diffusion model architecture in Primus, including design decisions, directory structure, and implementation patterns. + +--- + +## Table of Contents + +1. [Design Philosophy](#design-philosophy) +2. [Directory Structure](#directory-structure) +3. [Architectural Decisions](#architectural-decisions) +4. [Component Hierarchy](#component-hierarchy) +5. [Data Flow](#data-flow) +6. [Comparison with Alternative Implementations](#comparison-with-alternative-implementations) + +--- + +## Design Philosophy + +### Core Principles + +1. **Megatron-Core Native**: Built on Megatron-Core patterns and conventions +2. **Extensibility**: Easy to add new models (DiT, MovieGen, custom) +3. **Clarity**: Clear separation between shared and model-specific code +4. **Reusability**: Shared components across multiple models +5. **Performance**: Support for precalculated data and multi-GPU training + +### Architectural Advantages + +| Aspect | Primus Design Choice | +|--------|---------------------| +| Model Location | `core/models/diffusion/` (Megatron-Core convention) | +| Layer Organization | Unified `TransformerBlock` with heterogeneous specs | +| Shared Code | Dedicated `common/` directory | +| Encoders | Hierarchical `encoders/{type}/{variant}/` | +| Energon | Shared `data/energon/` for all models | +| Mock Data | Separated in `tests/fixtures/` | +| Framework | Pure Megatron (no PyTorch Lightning dependency) | + +--- + +## Directory Structure + +### High-Level Organization + +``` +primus/backends/megatron/ +├── core/models/diffusion/ # Core model implementations +├── training/diffusion/ # Training utilities (schedulers, etc.) +└── data/ + ├── energon/ # Shared Energon infrastructure + └── diffusion/ # Diffusion-specific data +``` + +### Detailed Breakdown + +#### 1. Core Models (`core/models/diffusion/`) + +Following Megatron-Core convention (`megatron/core/models/gpt/`, `megatron/core/models/multimodal/`): + +``` +core/models/diffusion/ +├── common/ # Shared across all diffusion models +│ ├── __init__.py +│ ├── config.py # BaseDiffusionConfig +│ ├── attention.py # JointSelfAttention, FluxSingleAttention +│ └── layers.py # MMDiTLayer, FluxSingleTransformerBlock +│ +└── flux/ # Flux-specific components + ├── __init__.py + ├── config.py # FluxConfig + ├── model.py # Flux (main model class) + └── layers.py # EmbedND (3D RoPE), embedders +``` + +**Rationale**: +- `common/`: Shared MMDiT patterns (used by both DiT and Flux) +- `flux/`: Only truly Flux-specific code (3D RoPE, Flux model class) +- Clear distinction enables easy DiT implementation (reuse `common/`) + +#### 2. Training (`training/diffusion/`) + +``` +training/diffusion/ +├── __init__.py +└── schedulers/ + ├── __init__.py + ├── base.py # BaseScheduler + ├── flow_matching.py # FlowMatchEulerDiscreteScheduler + ├── ddpm.py # [Future] DDPMScheduler + ├── edm.py # [Future] EDMScheduler + └── euler.py # [Future] EulerDiscreteScheduler +``` + +**Rationale**: +- Schedulers define noise schedules and training targets +- Separate from model code for clarity +- Easy to add new schedulers (DDPM, EDM, etc.) + +#### 3. Data Pipeline (`data/`) + +``` +data/ +├── energon/ # Shared Energon utilities +│ └── __init__.py +├── dataloader.py # MegatronDataloaderWrapper (wraps any iterable) +│ +└── diffusion/ + ├── __init__.py + ├── encoders/ # Hierarchical encoder registry + │ ├── __init__.py + │ ├── registry.py # EncoderRegistry + │ ├── image/ + │ │ └── vae/ + │ │ ├── __init__.py + │ │ ├── autoencoder_kl.py # AutoencoderKL + │ │ └── vqvae.py # VQVAE + │ └── text/ + │ ├── t5/ + │ │ ├── __init__.py + │ │ ├── t5_xxl.py # T5-XXL + │ │ └── t5_large.py # T5-Large + │ └── clip/ + │ ├── __init__.py + │ ├── clip_l.py # CLIP-L + │ └── clip_h.py # CLIP-H + │ + ├── preprocessing/ + │ ├── __init__.py + │ └── image/ + │ ├── __init__.py + │ ├── transforms.py # Resizing, normalization + │ └── augmentation.py # Data augmentation + │ + └── task_encoders/ + ├── __init__.py + └── image.py # EncodedDiffusionTaskEncoder +``` + +**Rationale**: +- **Energon shared**: VLM and other models can reuse Energon utilities +- **Hierarchical encoders**: Organized by modality and variant type +- **Registry pattern**: Config-driven encoder selection +- **Preprocessing separated**: Clear pipeline stages + +--- + +## Architectural Decisions + +### Decision 1: Models Under `core/models/` + +**Choice**: `primus/backends/megatron/core/models/diffusion/` +**Not**: `primus/backends/megatron/models/diffusion/` + +**Reasoning**: +- Aligns with Megatron-Core structure +- Easier upstream tracking +- Clear that these are Megatron-Core compatible +- Consistent with existing Primus structure (`core/models/gpt/`) + +### Decision 2: Separate `common/` Directory + +**Choice**: Shared components in `common/` +**Not**: Everything in `flux/` or flat structure + +**Reasoning**: +- Standard approach puts shared code in model-specific directories +- Primus: `common/` makes it explicit what's shared +- Future DiT implementation trivial (reuse from `common/`) +- Clear contract: if in `common/`, must work for all models + +**Shared Components**: +- `JointSelfAttention`: Used by DiT and Flux joint layers +- `FluxSingleAttention`: Used by DiT and Flux single layers +- `MMDiTLayer`: Joint (multimodal) transformer block +- `FluxSingleTransformerBlock`: Single-modality transformer block + +**Flux-Only Components**: +- `EmbedND`: 3D RoPE position embedding (Flux-specific) +- `Flux` model class + +### Decision 3: Hierarchical Encoder Structure + +**Choice**: `encoders/image/vae/`, `encoders/text/t5/`, `encoders/text/clip/` +**Not**: Flat `encoders/conditioner.py` + +**Reasoning**: +- Traditional approach uses flat structure with all encoders in one file +- User requirement: support 5+ variants per modality +- Registry pattern enables config-driven selection +- Easy to add new encoders without modifying existing files + +**Registry Pattern**: +```python +from primus.backends.megatron.data.diffusion.encoders import get_encoder + +# Config-driven selection +vae = get_encoder('autoencoder_kl', config=vae_config) +t5 = get_encoder('t5_xxl', config=t5_config) +clip = get_encoder('clip_l', config=clip_config) +``` + +### Decision 4: Shared Energon Infrastructure + +**Choice**: `data/energon/` for shared utilities +**Not**: Nested under `data/diffusion/` + +**Reasoning**: +- Energon is general-purpose (VLM, diffusion, future models) +- Traditional approach nests under model-specific directories +- Megatron-LM has Energon at example level (not in core) +- Primus approach: shared infra + model-specific TaskEncoders + +**Pattern**: +```python +# Shared: data/dataloader.py +# MegatronDataloaderWrapper wraps any iterable (Energon loader, PyTorch DataLoader, etc.) +wrapper = MegatronDataloaderWrapper(dataloader) + +# Model-specific: data/diffusion/task_encoders/image.py +class EncodedDiffusionTaskEncoder: + """Diffusion-specific encoding logic""" + pass +``` + +### Decision 5: Model Provider at Adapter Level + +**Choice**: `diffusion_model_provider.py` at `primus/backends/megatron/` +**Not**: Under `core/` + +**Reasoning**: +- Follows existing pattern (`primus/backends/megatron/model_provider.py`) +- Model providers are adapter/wrapper functions +- Sit above core models to add Primus-specific functionality +- Example: Wrap model with logit softcapping, custom loss, etc. + +### Decision 6: No PyTorch Lightning + +**Choice**: Pure Megatron patterns +**Not**: PyTorch Lightning DataModules + +**Reasoning**: +- Primus doesn't use PyTorch Lightning +- Framework-specific implementations reduce flexibility +- Better integration with Megatron training loop +- Follows Megatron-LM's `MegatronDataloaderWrapper` wrapper pattern + +--- + +## Component Hierarchy + +### 1. Model Hierarchy + +``` +nn.Module (PyTorch) +└── MegatronModule + └── DiffusionModule (abstract) + └── Flux (concrete) + ├── Joint layers: MMDiTLayer × num_joint_layers + │ └── JointSelfAttention (shared) + ├── Single layers: FluxSingleTransformerBlock × num_single_layers + │ └── FluxSingleAttention (shared) + └── Embeddings: + ├── TimeStepEmbedder (shared) + ├── MLPEmbedder (shared) + └── EmbedND (Flux-specific, 3D RoPE) +``` + +### 2. Configuration Hierarchy + +``` +TransformerConfig (Megatron-Core) +└── BaseDiffusionConfig + └── FluxConfig + ├── flux_535m() factory + └── flux_12b() factory +``` + +### 3. Scheduler Hierarchy + +``` +BaseScheduler (abstract) +├── FlowMatchEulerDiscreteScheduler (Flux) +├── DDPMScheduler (future) +├── EDMScheduler (future) +└── EulerDiscreteScheduler (future) +``` + +### 4. Encoder Hierarchy + +``` +BaseEncoder (abstract) +├── ImageEncoder +│ └── VAE +│ ├── AutoencoderKL (Flux) +│ └── VQVAE (future) +└── TextEncoder + ├── T5 + │ ├── T5-XXL (Flux) + │ └── T5-Large (future) + └── CLIP + ├── CLIP-L (Flux) + └── CLIP-H (future) +``` + +--- + +## Data Flow + +### Training Pipeline + +``` +Raw Data (images + captions) + ↓ +[Optional] Precalculation + ├─ VAE → latents [B, 64, H/8, W/8] + ├─ T5-XXL → embeddings [B, 512, 4096] + └─ CLIP-L → pooled [B, 768] + ↓ +WebDataset/Energon Format (.tar files) + ↓ +MegatronDataloaderWrapper (from data/dataloader.py) + ↓ +EncodedDiffusionTaskEncoder (from data/diffusion/task_encoders/) + ├─ Load precalculated data, OR + └─ Encode on-the-fly (slower) + ↓ +Training Batch: + ├─ latents: [B, 64, H, W] + ├─ t5_embeddings: [B, S, 4096] + └─ clip_pooled: [B, 768] + ↓ +FlowMatchEulerDiscreteScheduler + ├─ Sample timesteps: t ~ U(0, 1) + ├─ Sample noise: ε ~ N(0, I) + ├─ Add noise: x_t = (1-t)*ε + t*x_0 + └─ Compute target: v = x_0 - ε + ↓ +Flux Model Forward Pass + ├─ Embed timesteps + ├─ Embed pooled text (CLIP) + ├─ Joint layers (process latents + T5 embeddings) + ├─ Single layers (process latents only) + └─ Output: v_pred [B, 64, H, W] + ↓ +Loss Computation: MSE(v_pred, v_target) + ↓ +Backward Pass & Optimizer Step +``` + +### Inference Pipeline + +``` +Text Prompt + ↓ +Text Encoders + ├─ T5-XXL → embeddings [1, S, 4096] + └─ CLIP-L → pooled [1, 768] + ↓ +Initialize Noise: x_0 ~ N(0, I) + ↓ +Sampling Loop (t = 1.0 → 0.0) + ├─ Model forward: v_t = Flux(x_t, t, embeddings) + ├─ Update: x_{t-dt} = x_t + v_t * dt + └─ Repeat until t = 0 + ↓ +Latents: x_0 [1, 64, H, W] + ↓ +VAE Decoder + ↓ +Generated Image [1, 3, H*8, W*8] +``` + +--- + +## Comparison with Alternative Implementations + +### Primus Architectural Advantages + +#### 1. TransformerBlock Architecture (Primus Innovation) +- **Primus**: Unified `TransformerBlock` with heterogeneous layer specs +- **Others**: Separate `nn.ModuleList` containers for double/single blocks +- **Benefit**: Better PP slicing, unified checkpointing, future-proof + +#### 2. Megatron-Core Native +- **Primus**: Pure Megatron-Core, no framework dependencies +- **Others**: Often integrated with PyTorch Lightning or other frameworks +- **Benefit**: Tighter integration, simpler training loops + +#### 3. Checkpoint Format +- **Primus**: Unified `transformer.layers.{0-56}` structure +- **Others**: Separate `double_blocks.{i}` and `single_blocks.{j}` +- **Benefit**: Simpler distributed checkpointing + +#### 4. Encoder Architecture +- **Primus**: Registry-based, hierarchical organization +- **Others**: Direct imports from monolithic files +- **Benefit**: Easy extensibility for new encoder variants + +### File Organization Comparison + +| Component | Standard Location | Primus Location | Improvement | +|-----------|------------------|-----------------|-------------| +| Model | `models/diffusion/flux/` | `core/models/diffusion/flux/` | Megatron-Core convention | +| Layers | Mixed locations | `common/` for shared, `flux/` for specific | Clear boundaries | +| Encoders | Single file | `data/diffusion/encoders/{type}/{variant}/` | Hierarchical, extensible | +| Tests | Mixed with code | `tests/unit_tests/backends/megatron/diffusion/` | Proper separation | + +--- + +## Implementation Status + +### Core Infrastructure ✅ +- ✅ Directory structure +- ✅ Base classes (DiffusionModule, BaseDiffusionConfig, BaseScheduler) +- ✅ DiffusionModule with Megatron-Core integration +- ✅ FluxConfig with factory methods +- ✅ FlowMatchEulerDiscreteScheduler implementation +- ✅ Configuration files (YAML) +- ✅ Testing framework (290+ tests) +- ✅ Documentation structure + +### Flux Model Implementation ✅ +- ✅ Flux model architecture +- ✅ MMDiT layers and attention +- ✅ Embeddings (RoPE, timestep, vector) +- ✅ Encoder registry and loaders +- ✅ TaskEncoder for Energon +- ✅ Data pipeline + +--- + +## Extension Points + +### Adding a New Model (e.g., DiT) + +1. **Create model directory**: `core/models/diffusion/dit/` +2. **Add config**: Extend `BaseDiffusionConfig` +3. **Implement model**: Extend `DiffusionModule` (which extends MegatronModule) + - Inherit process group management + - Get distributed checkpointing support + - Access attention backend configuration +4. **Reuse shared components**: Import from `common/` +5. **Add tests**: `tests/unit_tests/backends/megatron/diffusion/test_dit_model.py` +6. **Update configs**: Add `dit_config.yaml` + +### Adding a New Encoder Variant + +1. **Create encoder file**: e.g., `data/diffusion/encoders/text/t5/t5_large.py` +2. **Implement encoder class**: Extend `BaseEncoder` +3. **Register**: Add to `ENCODER_REGISTRY` +4. **Add config**: Update `encoders.yaml` +5. **Add tests**: Test in `tests/unit_tests/backends/megatron/diffusion/data/encoders/` + +### Adding a New Scheduler + +1. **Create scheduler file**: `training/diffusion/schedulers/ddpm.py` +2. **Implement**: Extend `BaseScheduler` +3. **Export**: Add to `__init__.py` +4. **Add tests**: `tests/unit_tests/backends/megatron/diffusion/training/test_scheduler.py` +5. **Document**: Update this file and README + +--- + +## Performance Considerations + +### Memory Optimization +- **Precalculated data**: 5-10x faster, lower memory +- **Frozen encoders**: Only train diffusion model +- **Gradient checkpointing**: Trade compute for memory +- **Mixed precision**: bf16 on MI300X (compatible with H100/A100) + +### Multi-GPU Scaling +- **Tensor Parallelism**: Split model across GPUs +- **Pipeline Parallelism**: Split layers across GPUs +- **Data Parallelism**: Replicate model, split data +- **Sequence Parallelism**: For very long sequences + +### Best Practices +1. Use precalculated mode for training +2. Freeze encoders (standard practice) +3. Use bf16 on modern hardware +4. Start with TP=1, PP=1, scale as needed +5. Profile before optimizing + +--- + +## References + +- **Megatron-Core**: [nvidia/Megatron-LM](https://github.com/NVIDIA/Megatron-LM) transformer patterns +- **Flux**: [black-forest-labs/FLUX.1](https://huggingface.co/black-forest-labs/FLUX.1-dev) +- **Flow Matching**: Rectified flow and flow matching papers +- **NeMo**: [nvidia/NeMo](https://github.com/NVIDIA/NeMo) - Alternative diffusion implementation + +--- + +**Last Updated**: December 2025 diff --git a/docs/backends/megatron/diffusion/data_preprocessing.md b/docs/backends/megatron/diffusion/data_preprocessing.md new file mode 100644 index 000000000..b5a592729 --- /dev/null +++ b/docs/backends/megatron/diffusion/data_preprocessing.md @@ -0,0 +1,629 @@ +# Data Preprocessing Guide + +This guide explains how to prepare datasets for Flux and other diffusion models in Primus, including pre-encoding of VAE latents and text embeddings into Energon WebDataset format. + +--- + +## Table of Contents + +1. [Overview](#overview) +2. [Quick Start](#quick-start) +3. [Two Pipelines](#two-pipelines) +4. [Running Preprocessing](#running-preprocessing) +5. [Configuration](#configuration) +6. [Authentication](#authentication) +7. [Finalization](#finalization) +8. [Output Format](#output-format) +9. [Validation](#validation) +10. [Troubleshooting](#troubleshooting) + +--- + +## Overview + +Diffusion models require three types of encodings: +1. **VAE latents**: Images encoded to latent space +2. **Text embeddings**: Captions encoded with T5-XXL (sequence) +3. **Pooled embeddings**: Captions encoded with CLIP-L (pooled) + +### Why Pre-encode? + +**Benefits**: +- 5-10x faster training (no online encoding) +- Lower GPU memory usage (encoders not loaded during training) +- Deterministic inputs (same preprocessing for all runs) +- Eliminates encoder differences as a variable + +**When to Use**: +- Training (highly recommended) +- Fine-tuning on fixed datasets +- Benchmarking and reproducibility + +**When NOT to Use**: +- Interactive data augmentation needed +- Dataset too large to store pre-encoded +- Rapid prototyping with changing data + +--- + +## Quick Start + +Preprocess the Pokemon dataset with a single command: + +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml \ + --hf-token-file /path/to/.hf_token +``` + +This will: +1. Download the `diffusers/pokemon-gpt4-captions` dataset from HuggingFace +2. Encode all images with VAE and captions with T5/CLIP +3. Write Energon WebDataset tar shards to `/workspace/Primus/data/quickstart_pokemon` +4. Automatically finalize the dataset (create `dataset.yaml`, run `energon prepare`, validate) + +The default encoder model (`black-forest-labs/FLUX.1-dev`) is gated and requires a HuggingFace token. Get one at https://huggingface.co/settings/tokens and save it to a file. + +--- + +## Two Pipelines + +Primus provides two preprocessing pipelines via the `primus data` CLI: + +### `diffusion-encoded` (Recommended for Training) + +Pre-encodes images with VAE and text with T5/CLIP. Produces larger datasets but enables faster training since encoders are not needed at training time. + +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token +``` + +**Output**: WebDataset shards containing `latents.pth`, `prompt_embeds.pth`, `pooled_prompt_embeds.pth`, `caption.txt` + +**Use when**: Training on production datasets, maximum training speed is needed, storage is available. + +### `diffusion-raw` + +Stores raw images and captions without encoding. Smaller datasets but encoding happens on-the-fly during training (requires GPU + encoders loaded in memory). + +```bash +primus-cli direct -- data diffusion-raw \ + --source-type huggingface \ + --hf-dataset diffusers/pokemon-gpt4-captions \ + --output-dir /workspace/Primus/data/raw_pokemon \ + --hf-token-file /path/to/.hf_token +``` + +**Output**: WebDataset shards containing `jpg` (or `png`/`webp`) and `txt` files. + +**Use when**: Storage is limited, experimenting with different encoders, rapid prototyping. + +**Note**: `diffusion-raw` does not support `--config` files. All options must be passed as CLI arguments. + +--- + +## Running Preprocessing + +### Using a Config File (Recommended) + +The `--config` flag is supported by `diffusion-encoded` only. The simplest approach uses a YAML config file: + +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token +``` + +Available example configs in `primus/configs/data/megatron/diffusion/preprocessing/`: + +| Config | Source | Description | +|--------|--------|-------------| +| `quickstart_pokemon.yaml` | HuggingFace | Minimal config, 256px, fast | +| `example_huggingface.yaml` | HuggingFace | Full example with all options | +| `example_directory.yaml` | Local directory | Images + captions from disk | +| `example_webdataset.yaml` | WebDataset | Existing tar archives | +| `example_base.yaml` | N/A | Comprehensive reference with all fields | +| `text_to_image_2m_10k.yaml` | HuggingFace | 10K subset of text-to-image-2M (1024px) | + +### Using CLI Arguments Directly + +All config values can be provided as CLI arguments: + +```bash +primus-cli direct -- data diffusion-encoded \ + --source-type huggingface \ + --hf-dataset diffusers/pokemon-gpt4-captions \ + --output-dir /workspace/Primus/data/encoded_pokemon \ + --model-path black-forest-labs/FLUX.1-dev \ + --batch-size 8 \ + --precision bf16 \ + --hf-token-file /path/to/.hf_token +``` + +### CLI Overrides Config Values + +When using both `--config` and CLI arguments, CLI arguments take priority: + +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token \ + --output-dir /my/custom/path \ + --batch-size 16 \ + --max-samples 1000 +``` + +Priority order (highest to lowest): +1. Explicitly provided CLI arguments +2. YAML config file values +3. CLI default values + +### Multi-GPU Processing + +Use `--nproc-per-node` for data-parallel preprocessing across multiple GPUs: + +```bash +primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token +``` + +Each GPU processes a subset of the data. Shards are named to avoid conflicts across ranks. + +--- + +## Configuration + +### YAML Config Structure + +Preprocessing configs have four sections: + +```yaml +source: + type: huggingface # huggingface | directory | webdataset + hf_dataset: diffusers/pokemon-gpt4-captions + hf_split: train + +output: + output_dir: /workspace/Primus/data/encoded_pokemon + shard_size: 1000 # samples per tar shard + max_samples: null # null = process all + compress: false + +model: + model_path: black-forest-labs/FLUX.1-dev # HF repo or local path + precision: bf16 # bf16 | fp16 | fp32 + batch_size: 8 + # Optional per-encoder overrides: + vae_path: null + t5_path: null + clip_path: null + +image: + image_size: 1024 + center_crop: false +``` + +### Source Types + +**HuggingFace** (`type: huggingface`): +```yaml +source: + type: huggingface + hf_dataset: diffusers/pokemon-gpt4-captions + hf_split: train + hf_data_files: null # optional: specific files within dataset +``` + +**Local Directory** (`type: directory`): +```yaml +source: + type: directory + input_dir: /data/my_images +``` + +Expected directory structure for `directory` source: +``` +my_images/ +├── images/ +│ ├── 00001.jpg +│ ├── 00002.png +│ └── ... +└── captions/ + ├── 00001.txt + ├── 00002.txt + └── ... +``` + +**WebDataset** (`type: webdataset`): +```yaml +source: + type: webdataset + input_path: /data/existing_shards/*.tar +``` + +### Model Configuration + +The `model_path` defaults to `black-forest-labs/FLUX.1-dev`, which downloads VAE, T5-XXL, and CLIP-L encoders from HuggingFace. This model is gated and requires authentication (see [Authentication](#authentication)). + +Individual encoder paths can be overridden: + +```yaml +model: + model_path: black-forest-labs/FLUX.1-dev + vae_path: /local/models/vae # use local VAE instead + t5_path: null # falls back to model_path + clip_path: null # falls back to model_path +``` + +--- + +## Authentication + +The default encoder model (`FLUX.1-dev`) is gated on HuggingFace and requires authentication. Primus supports three authentication methods, checked in priority order: + +### 1. Token File (Recommended) + +```bash +primus-cli direct -- data diffusion-encoded \ + --config your_config.yaml \ + --hf-token-file /path/to/.hf_token +``` + +The token file must have secure permissions (600 or 400). Create it with: + +```bash +echo "hf_your_token_here" > /path/to/.hf_token +chmod 600 /path/to/.hf_token +``` + +### 2. Environment Variable + +```bash +export HF_TOKEN=hf_your_token_here +primus-cli direct -- data diffusion-encoded --config your_config.yaml +``` + +### 3. HuggingFace CLI Login + +```bash +huggingface-cli login +primus-cli direct -- data diffusion-encoded --config your_config.yaml +``` + +If authentication fails, Primus provides a clear error message indicating which encoder failed and how to fix it. + +--- + +## Finalization + +Finalization is **automatic by default**. After preprocessing completes, Primus automatically: + +1. **Creates `.nv-meta/dataset.yaml`** with `CrudeWebdataset` sample type and encoding subflavor +2. **Runs `energon prepare`** to index the tar shards and create split assignments +3. **Validates the dataset** using Primus's custom validation (metadata checks, sample count verification, energon API spot-check) + +### Skipping Finalization + +To skip automatic finalization (e.g., for manual post-processing): + +```bash +primus-cli direct -- data diffusion-encoded \ + --config your_config.yaml \ + --hf-token-file /path/to/.hf_token \ + --no-finalize +``` + +### Custom Train/Val/Test Splits + +By default, 100% of data goes to the training split. To create validation and test splits: + +```bash +primus-cli direct -- data diffusion-encoded \ + --config your_config.yaml \ + --hf-token-file /path/to/.hf_token \ + --train-split 0.8 +``` + +This creates an 80% train / 10% val / 10% test split. + +--- + +## Output Format + +### Directory Structure + +After preprocessing and finalization, the output directory contains: + +``` +encoded_pokemon/ +├── 000000.tar # WebDataset shard +├── 000001.tar +├── 000000.tar.idx # Energon index files +├── 000001.tar.idx +└── .nv-meta/ # Energon metadata + ├── dataset.yaml # Dataset type configuration + ├── split.yaml # Train/val/test split assignments + └── .info.json # Shard counts and sample counts +``` + +### Pre-encoded Shard Contents (`diffusion-encoded`) + +Each tar shard contains samples with these keys: + +``` +000000.tar: +├── 0000000000.latents.pth # VAE latents tensor +├── 0000000000.prompt_embeds.pth # T5-XXL embeddings tensor +├── 0000000000.pooled_prompt_embeds.pth # CLIP-L pooled embeddings tensor +├── 0000000000.caption.txt # Original caption text +├── 0000000001.latents.pth +├── 0000000001.prompt_embeds.pth +├── ... +``` + +Tensor shapes: +- `latents.pth`: `[64, H/8, W/8]` (e.g., `[64, 128, 128]` for 1024x1024 images) +- `prompt_embeds.pth`: `[seq_len, 4096]` (T5-XXL hidden dim) +- `pooled_prompt_embeds.pth`: `[768]` (CLIP-L pooled dim) + +### Raw Shard Contents (`diffusion-raw`) + +``` +000000.tar: +├── 0000000000.jpg # Preprocessed image +├── 0000000000.txt # Caption text +├── 0000000001.jpg +├── 0000000001.txt +├── ... +``` + +### Dataset YAML + +The auto-generated `.nv-meta/dataset.yaml` uses `CrudeWebdataset` format: + +```yaml +__module__: megatron.energon +__class__: CrudeWebdataset +subflavors: + encoding: preencoded # or 'raw' for diffusion-raw +``` + +--- + +## Validation + +### Automatic Validation + +Validation runs automatically as part of finalization. It performs four checks: + +1. **Metadata check**: Verifies `.nv-meta/` files exist and are valid (`.info.json`, `split.yaml`, `dataset.yaml`) +2. **Sample count check**: Spot-checks that tar shard entry counts match `.info.json` +3. **Sample load check**: Loads one sample through Energon's Python API (same code path as training) +4. **Summary report**: Prints dataset statistics (encoding, total samples, splits, data shapes, size) + +### Standalone Validation + +To validate a dataset independently: + +```bash +python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset + +# For raw datasets: +python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset --encoding raw +``` + +### Programmatic Validation + +```python +from primus.backends.megatron.data.diffusion.preprocessing.validate import validate_energon_dataset + +ok = validate_energon_dataset('/path/to/dataset', encoding='preencoded') +``` + +### Known Energon CLI Limitations + +The standard Energon CLI tools (`energon info`, `energon preview`, `energon lint`) do **not** work correctly with `CrudeWebdataset` format. Primus uses custom validation instead: + +- `energon info` raises `KeyError: 'sample_type'` +- `energon preview` raises `TypeError` (expects dataclass, `CrudeSample` is a dict) +- `energon lint` raises `AssertionError` (expects registered cookers) + +Use Primus's built-in validation or the standalone script above. + +--- + +## Troubleshooting + +### HuggingFace Authentication Failure + +**Symptoms**: Error mentioning "token", "gated", "401", or "403" when downloading encoders. + +**Solutions**: +1. Provide a token: `--hf-token-file /path/to/.hf_token` +2. Set environment variable: `export HF_TOKEN=hf_xxx` +3. Run `huggingface-cli login` +4. Accept the model's license on https://huggingface.co/black-forest-labs/FLUX.1-dev + +### Out of Memory During Preprocessing + +**Symptoms**: CUDA out of memory error during encoding. + +**Solutions**: +1. Reduce batch size: `--batch-size 4` or `--batch-size 1` +2. Use smaller image size: `--image-size 512` +3. Use fp16 precision: `--precision fp16` +4. Use multi-GPU to distribute work: `--nproc-per-node=8` + +### Missing Dependencies + +**Symptoms**: `ModuleNotFoundError` for `webdataset`, `megatron-energon`, `tqdm`, etc. + +**Solution**: +```bash +pip install -r requirements.txt +``` + +Or set `PRIMUS_AUTO_INSTALL=1` in the container to auto-install missing packages. + +### Finalization Fails + +**Symptoms**: Error during `energon prepare` or validation after preprocessing. + +**Solutions**: +1. Check that tar shards exist in the output directory +2. Ensure `megatron-energon` is installed (`pip install megatron-energon`) +3. Re-run with `--no-finalize`, then manually inspect the output before finalizing + +### Slow Preprocessing + +**Symptoms**: Low throughput (< 10 samples/sec). + +**Solutions**: +1. Increase batch size (limited by VRAM): `--batch-size 16` +2. Use multiple GPUs: `--nproc-per-node=8` +3. Use bf16 precision (faster on supported hardware): `--precision bf16` +4. For raw pipeline, reduce image quality: `--image-quality 85` + +--- + +## Best Practices + +1. **Always pre-encode for production training**: 5-10x speedup is worth the storage +2. **Test on small dataset first**: Use `--max-samples 100` to verify the pipeline works +3. **Use bf16 precision**: Good balance of speed, storage, and quality +4. **Start with quickstart_pokemon.yaml**: Verify your setup before processing large datasets +5. **Keep raw data**: Pre-encoding is a one-way transformation +6. **Version your configs**: Track which preprocessing config produced each dataset + +--- + +## Flux-Specific Data Preparation (torchrun) + +This section covers preparing datasets for Flux training using `torchrun` directly +inside a Docker container, as an alternative to the `primus-cli` workflow above. + +### Prerequisites + +Start the container and optionally install requirements: + +```bash +bash tools/docker/start_container.sh +docker exec dev_primus bash -c 'pip install -r /workspace/Primus/requirements.txt' +``` + +This mounts the repository to `/workspace/Primus` inside the container. +Override the image with `DOCKER_IMAGE`: + +```bash +DOCKER_IMAGE=docker.io/rocm/primus:v26.1 bash tools/docker/start_container.sh +``` + +### Input Directory Structure + +When using `--source-type directory`, organize your data as follows: + +``` +dataset/ +├── images/ +│ ├── 0000000.png +│ ├── 0000001.png +│ └── ... +└── captions/ + ├── 0000000.txt + ├── 0000001.txt + └── ... +``` + +The file stem links each image to its caption (e.g., `images/0000000.png` pairs +with `captions/0000000.txt`). Images can be `.jpg`, `.jpeg`, `.png`, or `.webp`. +Captions must be UTF-8 `.txt` files. Samples without a matching caption are skipped. + +Other supported source types: +- **huggingface** -- Load directly from HuggingFace Hub. Requires `--hf-dataset`. +- **webdataset** -- Read from existing WebDataset tar archives. Requires `--input-path`. + +### Image Sizing + +**Fixed size (default):** Every image is resized to `--image-size` pixels square +(default 1024). Use `--center-crop` to control center-cropping. + +**Variable size (`--variable-size`):** Preserves aspect ratio by scaling the +longest side to `--max-size` (default 1024) and rounding dimensions to multiples +of 16. Only use this when all images share the same dimensions -- mixed tensor +sizes cause load imbalance across GPUs. + +### Key Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `--source-type` | — | `directory`, `huggingface`, or `webdataset` | +| `--input-dir` | — | Dataset root (for `directory` source) | +| `--image-size` | 1024 | Square target size | +| `--variable-size` | off | Preserve aspect ratio | +| `--max-size` | 1024 | Maximum dimension in variable-size mode | +| `--model-path` | `black-forest-labs/FLUX.1-dev` | Base model for encoder weights | +| `--t5-max-length` | 512 | T5 token limit (use 256 for FLUX.1-schnell) | +| `--batch-size` | 8 | Encoding batch size per GPU | +| `--output-dir` | — | Destination for encoded WebDataset shards | +| `--shard-size` | 1000 | Samples per tar shard | + +### Examples + +**From host (via `docker exec`):** + +```bash +docker exec dev_primus bash -c '\ + export HF_HOME=/workspace/Primus/checkpoints/flux; \ +PYTHONPATH=/workspace/Primus:/workspace/Primus/third_party/Megatron-LM:$PYTHONPATH \ + torchrun --nproc_per_node=8 /workspace/Primus/primus/cli/main.py \ + data diffusion-encoded \ + --source-type directory \ + --input-dir /workspace/Primus/data/dataset \ + --output-dir /workspace/Primus/data/dataset_encoded_256 \ + --image-size 256 \ + --t5-max-length 256' +``` + +**Inside the container:** + +```bash +export HF_HOME=/workspace/Primus/checkpoints/flux +PYTHONPATH=/workspace/Primus:/workspace/Primus/third_party/Megatron-LM:$PYTHONPATH \ + torchrun --nproc_per_node=8 /workspace/Primus/primus/cli/main.py \ + data diffusion-encoded \ + --source-type directory \ + --input-dir /workspace/Primus/data/dataset \ + --output-dir /workspace/Primus/data/dataset_encoded_256 \ + --image-size 256 \ + --t5-max-length 256 +``` + +Both commands encode a local directory dataset at 256px on 8 GPUs, then create an Energon dataset by default. + +### Limitations + +**Uniform output size only.** When using `--variable-size`, images may produce +tensors of different shapes. The current data pipeline does not support mixed +tensor sizes because samples are distributed evenly across GPUs and unequal +shapes cause load imbalance and eventual timeout. Use fixed `--image-size` if +your images have various sizes. + +--- + +## Next Steps + +After preprocessing: +1. **Training**: Use the preprocessed dataset path in your training config +2. **Validation**: The dataset is ready for training immediately after finalization + +See: +- [Energon Integration](energon_integration.md) for TaskEncoder and dataloader details +- [Config Directory Guide](../../../../primus/configs/data/megatron/diffusion/README.md) for config file reference +- Example configs in `primus/configs/data/megatron/diffusion/preprocessing/` + +--- + +**Last Updated**: June 2026 diff --git a/docs/backends/megatron/diffusion/energon_integration.md b/docs/backends/megatron/diffusion/energon_integration.md new file mode 100644 index 000000000..032e7222f --- /dev/null +++ b/docs/backends/megatron/diffusion/energon_integration.md @@ -0,0 +1,514 @@ +# Energon Integration Guide + +This guide explains how Megatron-Energon is integrated with Primus diffusion models, including the Cooker/TaskEncoder pattern, dataloader configuration, and dataset format. + +--- + +## Table of Contents + +1. [Overview](#overview) +2. [Energon Architecture](#energon-architecture) +3. [Dataset Format](#dataset-format) +4. [TaskEncoder and Cooker Pattern](#taskencoder-and-cooker-pattern) +5. [Dataloader Setup](#dataloader-setup) +6. [Dataset Configuration](#dataset-configuration) +7. [Implementation Examples](#implementation-examples) +8. [Best Practices](#best-practices) + +--- + +## Overview + +### What is Megatron-Energon? + +Megatron-Energon is NVIDIA's data loading framework for large-scale multimodal training. It provides: +- **WebDataset integration**: Efficient streaming from .tar archives +- **Task encoding**: Flexible data transformation pipeline via Cookers +- **Multi-worker support**: Parallel data loading +- **Deterministic iteration**: Reproducible training +- **Checkpoint resumption**: Resume from any step + +### Why Energon for Diffusion? + +- **Proven at scale**: Used by NVIDIA for LLM and multimodal training +- **Flexible**: Supports various data formats and transformations +- **Efficient**: Optimized for multi-GPU training +- **Compatible**: Works with Megatron parallelism (TP, PP, DP) + +### Primus Integration Strategy + +``` +Shared Infrastructure (data/) + ├─ MegatronDataloaderWrapper + ├─ Dataset configuration parsers + └─ Common utilities + ↓ +Model-Specific TaskEncoders (data/diffusion/task_encoders/) + ├─ EncodedDiffusionTaskEncoder (preencoded data) + ├─ RawDiffusionTaskEncoder (raw images + text) + └─ Custom task encoders +``` + +**Key Principle**: Share infrastructure, separate domain logic. + +--- + +## Energon Architecture + +### Component Stack + +``` +Training Loop + ↓ +MegatronDataloaderWrapper (Primus wrapper, cyclic iteration) + ↓ +Megatron-Energon Core + ├─ WebDataset Reader (reads .tar shards) + ├─ Cooker (transforms raw dict → typed Sample) + ├─ TaskEncoder.batch() (stacks samples → batch) + └─ Worker Pool + ↓ +.tar Shards (CrudeWebdataset format) +``` + +### Data Flow + +``` +1. Load from .tar shard + → raw_dict = {'__key__': '0000000000', + 'latents.pth': bytes, + 'prompt_embeds.pth': bytes, + 'pooled_prompt_embeds.pth': bytes, + 'caption.txt': bytes} + +2. Cooker function (e.g., cook_preencoded_diffusion) + → Deserializes bytes to tensors + → Returns typed DiffusionSample dataclass + → DiffusionSample(latents=Tensor[64,128,128], + prompt_embeds=Tensor[512,4096], + pooled_prompt_embeds=Tensor[768], + caption="a photo of...") + +3. TaskEncoder.batch() + → Stacks list of DiffusionSample into batch dict + → batch = {'latents': [B, 64, H, W], + 'prompt_embeds': [B, seq_len, 4096], + 'pooled_prompt_embeds': [B, 768]} + +4. Return to training loop via MegatronDataloaderWrapper + → Forward pass, loss, backprop +``` + +--- + +## Dataset Format + +### CrudeWebdataset + +Primus uses Energon's `CrudeWebdataset` format for preprocessed datasets. This is the simplest Energon format -- it stores raw key-value pairs in tar shards without requiring a strict schema. + +The dataset type is configured in `.nv-meta/dataset.yaml`: + +```yaml +__module__: megatron.energon +__class__: CrudeWebdataset +subflavors: + encoding: preencoded # or 'raw' +``` + +The `subflavors.encoding` field tells the TaskEncoder which Cooker to use: +- `preencoded`: Uses `cook_preencoded_diffusion` -- loads pre-encoded tensors +- `raw`: Uses `cook_raw_images` -- loads raw images and text + +### Pre-encoded Shard Contents + +Each tar shard contains samples with these keys: + +| Key | Type | Shape | Description | +|-----|------|-------|-------------| +| `latents.pth` | Tensor | `[64, H/8, W/8]` | VAE-encoded image latents | +| `prompt_embeds.pth` | Tensor | `[seq_len, 4096]` | T5-XXL text embeddings | +| `pooled_prompt_embeds.pth` | Tensor | `[768]` | CLIP-L pooled embeddings | +| `caption.txt` | str | N/A | Original caption text | + +### Raw Shard Contents + +| Key | Type | Description | +|-----|------|-------------| +| `jpg` / `png` / `webp` | bytes | Preprocessed image | +| `txt` | str | Caption text | + +### Known Energon CLI Limitations + +The standard Energon CLI tools have issues with `CrudeWebdataset`: +- `energon info` raises `KeyError: 'sample_type'` +- `energon preview` raises `TypeError` (expects dataclass, `CrudeSample` is a dict) +- `energon lint` raises `AssertionError` (expects registered cookers) + +Primus includes custom validation (`primus.backends.megatron.data.diffusion.preprocessing.validate`) as a replacement. See the [Data Preprocessing Guide](data_preprocessing.md#validation) for details. + +--- + +## TaskEncoder and Cooker Pattern + +Primus uses Energon's Cooker pattern rather than the older `encode_sample()` approach. Cookers are `@stateless` functions that transform raw sample dicts into typed dataclass instances, dispatched based on `subflavors`. + +### DiffusionSample Dataclass + +Location: `primus/backends/megatron/data/diffusion/task_encoders/image.py` + +```python +# (imports like torch omitted for brevity) +from dataclasses import dataclass +from megatron.energon import Sample + +@dataclass +class DiffusionSample(Sample): + """ + Diffusion training sample with framework-standard field names. + + Inherits from megatron.energon.Sample to ensure __key__, __restore_key__, + and __subflavors__ are properly tracked for deterministic training resumption. + """ + latents: torch.Tensor # [C, H, W] VAE latents + prompt_embeds: torch.Tensor # [seq_len, hidden_dim] T5 embeddings + pooled_prompt_embeds: torch.Tensor # [hidden_dim] CLIP pooled + caption: str = "" +``` + +### Cooker Functions + +Cookers are `@stateless` functions registered with a `Cooker` wrapper that specifies which `subflavors` they handle: + +```python +# (imports like torch, io omitted for brevity) +from megatron.energon import Cooker, basic_sample_keys, stateless + +@stateless +def cook_preencoded_diffusion(sample: dict) -> DiffusionSample: + """Load precalculated VAE latents and text embeddings from disk.""" + + def load_tensor(data): + if isinstance(data, bytes): + return torch.load(io.BytesIO(data), map_location='cpu') + return data + + latents = load_tensor(sample.get('latents.pth')) + prompt_embeds = load_tensor(sample.get('prompt_embeds.pth')) + pooled_prompt_embeds = load_tensor(sample.get('pooled_prompt_embeds.pth')) + + caption = sample.get('caption.txt', b'') + if isinstance(caption, bytes): + caption = caption.decode('utf-8') + + return DiffusionSample( + **basic_sample_keys(sample), + latents=latents, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + caption=caption, + ) + + +@stateless +def cook_raw_images(sample: dict) -> Dict[str, Any]: + """Load raw images and text -- NO encoding, just data loading.""" + return { + **basic_sample_keys(sample), + 'images': sample.get('images'), + 'txt': sample.get('txt'), + } +``` + +The cooker code above is simplified for clarity. See `primus/backends/megatron/data/diffusion/task_encoders/image.py` for the full implementation, which includes additional input type handling and validation. + +### EncodedDiffusionTaskEncoder + +The TaskEncoder registers Cookers and provides the `batch()` method: + +```python +from megatron.energon import DefaultTaskEncoder, SampleDecoder, Cooker, WorkerConfig + +class EncodedDiffusionTaskEncoder(DefaultTaskEncoder[DiffusionSample, DiffusionSample, dict, dict]): + """TaskEncoder for PRE-ENCODED diffusion data.""" + + decoder = SampleDecoder(image_decode="pil") + + cookers = [ + Cooker(cook_preencoded_diffusion, has_subflavors={"encoding": "preencoded"}), + ] + + def __init__(self, worker_config: Optional[WorkerConfig] = None): + super().__init__() + self.worker_config = worker_config + + def batch(self, samples: List[DiffusionSample]) -> Dict[str, torch.Tensor]: + return { + 'latents': torch.stack([s.latents for s in samples]), + 'prompt_embeds': torch.stack([s.prompt_embeds for s in samples]), + 'pooled_prompt_embeds': torch.stack([s.pooled_prompt_embeds for s in samples]), + } +``` + +### RawDiffusionTaskEncoder + +For raw (on-the-fly encoding) datasets: + +```python +class RawDiffusionTaskEncoder(DefaultTaskEncoder): + """TaskEncoder for RAW diffusion data (images and text).""" + + decoder = SampleDecoder(image_decode="pil") + + cookers = [ + Cooker(cook_raw_images, has_subflavors={"encoding": "raw"}), + ] + + def __init__(self, worker_config: Optional[WorkerConfig] = None): + super().__init__() + self.worker_config = worker_config + + def batch(self, samples: List[Dict]) -> Dict[str, Any]: + return { + 'images': [s['images'] for s in samples], + 'txt': [s['txt'] for s in samples], + } +``` + +### How Cooker Dispatch Works + +The Cooker framework matches samples to cooker functions based on `subflavors`: + +1. `dataset.yaml` specifies `subflavors: { encoding: preencoded }` +2. Energon reads a sample from the tar shard +3. The `has_subflavors` on each `Cooker` is checked against the sample's subflavors +4. The matching cooker function is called to transform the raw dict into a typed sample +5. `TaskEncoder.batch()` stacks multiple samples into a training batch + +This decouples data format (what's in the tar) from data loading logic (how to interpret it). + +--- + +## Dataloader Setup + +### MegatronDataloaderWrapper + +Location: `primus/backends/megatron/data/dataloader.py` + +The `MegatronDataloaderWrapper` is a generic wrapper that makes any iterable compatible with Megatron's training loop. It provides: +- Cyclic iteration (never raises StopIteration) +- Optional checkpoint support via duck typing (`save_state_rank()` / `restore_state_rank()`) +- Works with PyTorch DataLoader, Energon loaders, and synthetic data + +```python +from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper + +wrapper = MegatronDataloaderWrapper(pytorch_or_energon_loader) +``` + +Note: Originally named `EnergonDataloader`, renamed to `MegatronDataloaderWrapper` to reflect its generic nature (it has no Energon dependencies). The old name is available as a deprecated alias. + +### Usage Example + +```python +from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper + +# Create Energon loader (with TaskEncoder configured) +energon_loader = get_loader(...) +dataloader = MegatronDataloaderWrapper(energon_loader) + +# Use in training loop +for batch in dataloader: + latents = batch['latents'] # [B, 64, H, W] + prompt_embeds = batch['prompt_embeds'] # [B, seq_len, 4096] + pooled_prompt_embeds = batch['pooled_prompt_embeds'] # [B, 768] + + output = model(latents, timesteps, prompt_embeds, pooled_prompt_embeds) + loss = criterion(output, target) + loss.backward() +``` + +--- + +## Dataset Configuration + +### Per-Dataset dataset.yaml + +Each dataset directory has a `.nv-meta/dataset.yaml` that specifies the Energon dataset type: + +```yaml +__module__: megatron.energon +__class__: CrudeWebdataset +subflavors: + encoding: preencoded +``` + +This file is auto-generated by Primus during finalization. See [Data Preprocessing Guide](data_preprocessing.md#finalization). + +### Metadataset (Multi-Dataset Mixing) + +For combining multiple datasets with different weights: + +```yaml +__module__: megatron.energon +__class__: Metadataset + +splits: + train: + datasets: + - weight: 0.7 + path: /data/laion_precalculated/ + - weight: 0.3 + path: /data/coco_precalculated/ +``` + +Energon samples proportionally to weights (70% from LAION, 30% from COCO). + +--- + +## Implementation Examples + +### Example 1: Pre-encoded Training + +```python +from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper +from primus.backends.megatron.data.diffusion.task_encoders import EncodedDiffusionTaskEncoder + +# TaskEncoder is configured by the dataset provider +# The cooker automatically handles subflavors dispatch +energon_loader = get_loader(...) +dataloader = MegatronDataloaderWrapper(energon_loader) + +for batch in dataloader: + output = model( + batch['latents'], + batch['prompt_embeds'], + batch['pooled_prompt_embeds'], + ) +``` + +### Example 2: Raw Data Training (On-the-Fly Encoding) + +```python +from primus.backends.megatron.data.diffusion.task_encoders import RawDiffusionTaskEncoder + +# Raw TaskEncoder passes through images and text without encoding +# Encoding happens in the model's forward_step +energon_loader = get_loader(...) +dataloader = MegatronDataloaderWrapper(energon_loader) + +for batch in dataloader: + images = batch['images'] # List of PIL Images + captions = batch['txt'] # List of caption strings + # Model handles VAE/T5/CLIP encoding in forward_step +``` + +### Example 3: Multi-GPU Training + +```python +import torch.distributed as dist + +dist.init_process_group(backend='nccl') + +energon_loader = get_loader(...) +dataloader = MegatronDataloaderWrapper(energon_loader) + +for batch in dataloader: + output = model(batch['latents'], ...) +``` + +--- + +## Best Practices + +### 1. Pre-encoded Mode +- **Always use for production training**: 5-10x faster +- **Validate first**: Test with small dataset using `quickstart_pokemon.yaml` +- **Version datasets**: Track which preprocessing config produced each dataset + +### 2. Cooker Design +- **Use `@stateless`**: Cookers must be stateless and side-effect-free +- **Use `basic_sample_keys()`**: Always forward `__key__`, `__restore_key__`, `__subflavors__` +- **Keep it simple**: Cookers should only deserialize and restructure data, not transform it +- **Use subflavors for dispatch**: Let the framework choose the right cooker + +### 3. TaskEncoder Design +- **Single responsibility**: One task encoder per data format family +- **Minimal `batch()`**: Only stack tensors, avoid computation in the batch method +- **Separate concerns**: Data loading in Cooker, encoding in model forward_step + +### 4. Dataloader Configuration +- **Num workers**: Match CPU cores (typically 4-8) +- **Batch size**: Max out GPU memory +- **Shuffle**: Always true for training +- **Drop last**: True to avoid irregular batches + +### 5. Debugging +- **Small dataset**: Test with 100 samples first (`--max-samples 100`) +- **Single worker**: Set `num_workers=1` for debugging +- **Validate**: Run `python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset` + +--- + +## Customization Patterns + +### Custom Cooker Function + +Add a new cooker for a different data format: + +```python +from megatron.energon import Cooker, basic_sample_keys, stateless + +@stateless +def cook_my_custom_format(sample: dict) -> DiffusionSample: + """Custom cooker for a different data layout.""" + latents = torch.load(io.BytesIO(sample['vae_output.pth']), map_location='cpu') + prompt = torch.load(io.BytesIO(sample['text_embed.pth']), map_location='cpu') + pooled = torch.load(io.BytesIO(sample['clip_embed.pth']), map_location='cpu') + + return DiffusionSample( + **basic_sample_keys(sample), + latents=latents, + prompt_embeds=prompt, + pooled_prompt_embeds=pooled, + ) + +class CustomDiffusionTaskEncoder(DefaultTaskEncoder[DiffusionSample, DiffusionSample, dict, dict]): + decoder = SampleDecoder(image_decode="pil") + cookers = [ + Cooker(cook_my_custom_format, has_subflavors={"encoding": "custom_v2"}), + ] + + def batch(self, samples): + return { + 'latents': torch.stack([s.latents for s in samples]), + 'prompt_embeds': torch.stack([s.prompt_embeds for s in samples]), + 'pooled_prompt_embeds': torch.stack([s.pooled_prompt_embeds for s in samples]), + } +``` + +### Multiple Cookers in One TaskEncoder + +A single TaskEncoder can register multiple cookers for different subflavors: + +```python +class MultiFormatTaskEncoder(DefaultTaskEncoder[DiffusionSample, DiffusionSample, dict, dict]): + cookers = [ + Cooker(cook_preencoded_diffusion, has_subflavors={"encoding": "preencoded"}), + Cooker(cook_my_custom_format, has_subflavors={"encoding": "custom_v2"}), + ] +``` + +--- + +## References + +- **Megatron-Energon**: [nvidia/Megatron-LM](https://github.com/NVIDIA/Megatron-LM) multimodal examples +- **WebDataset**: [webdataset/webdataset](https://github.com/webdataset/webdataset) +- **NeMo Implementation**: `nemo/collections/diffusion/data/` +- **Primus TaskEncoders**: `primus/backends/megatron/data/diffusion/task_encoders/image.py` +- **Primus Dataloader**: `primus/backends/megatron/data/dataloader.py` + +--- + +**Last Updated**: March 2026 diff --git a/docs/backends/megatron/diffusion/flux_architecture.md b/docs/backends/megatron/diffusion/flux_architecture.md new file mode 100644 index 000000000..02189af37 --- /dev/null +++ b/docs/backends/megatron/diffusion/flux_architecture.md @@ -0,0 +1,926 @@ +# Flux Architecture Deep Dive + +## Table of Contents + +1. [Overview](#overview) +2. [Architecture Principles](#architecture-principles) +3. [Model Components](#model-components) +4. [Data Flow](#data-flow) +5. [Mathematical Formulation](#mathematical-formulation) +6. [Implementation Details](#implementation-details) +7. [Megatron-Core Integration](#megatron-core-integration) +8. [Performance Optimizations](#performance-optimizations) + +--- + +## Overview + +Flux is a **flow-based diffusion model** for high-quality text-to-image generation. It uses an innovative **MMDiT (Multimodal Diffusion Transformer)** architecture that jointly processes image and text tokens through shared transformer blocks. + +### Key Innovations + +1. **Flow Matching**: Uses rectified flow instead of traditional diffusion +2. **MMDiT Architecture**: Joint image-text attention in early layers +3. **3D RoPE**: Multi-dimensional rotary position embeddings for spatial awareness +4. **Two-Stage Processing**: Joint layers followed by image-only layers + +### Model Variants + +| Variant | Joint Layers | Single Layers | Parameters | Use Case | +|---------|--------------|---------------|------------|----------| +| Flux 535M | 1 | 1 | ~535M | Development, testing | +| Flux 12B | 19 | 38 | ~12B | Production deployment | + +--- + +## Architecture Principles + +### 1. Flow Matching Framework + +Unlike traditional diffusion (which adds Gaussian noise), Flux uses **rectified flow**: + +``` +Forward process: z_t = (1-t) * z_0 + t * z_1 +where: + - z_0 = original image (latent) + - z_1 = random noise + - t ∈ [0, 1] is the flow timestep +``` + +The model predicts the **velocity field** v_θ: + +``` +v_θ(z_t, t, c) ≈ z_1 - z_0 +``` + +**Advantages**: +- Straight-line interpolation paths (more efficient than diffusion curves) +- Faster sampling (fewer steps needed) +- Better training stability + +### 2. MMDiT (Multimodal Diffusion Transformer) + +Traditional DiT processes image tokens independently. MMDiT jointly processes image and text: + +``` +┌─────────────┐ ┌─────────────┐ +│ Image Tokens│ │ Text Tokens │ +└──────┬──────┘ └──────┬──────┘ + │ │ + └───────┬───────────┘ + │ + ┌──────▼──────┐ + │ Joint Attn │ ← Cross-attend image ↔ text + └──────┬──────┘ + │ + ┌──────▼──────┐ + │ Split │ + └──┬───────┬──┘ + │ │ + ┌───────▼──┐ ┌─▼────────┐ + │ Image MLP│ │ Text MLP │ ← Separate processing + └──────────┘ └──────────┘ +``` + +**Benefits**: +- Better text-image alignment +- Richer cross-modal interactions +- Improved compositional understanding + +### 3. Two-Stage Processing + +Flux uses a unique two-stage architecture: + +**Stage 1: Joint Processing (Double Blocks)** +- Both image and text tokens +- Multi-modal attention +- Rich semantic understanding + +**Stage 2: Image Refinement (Single Blocks)** +- Image tokens only (text concatenated but not separated) +- Focus on spatial coherence +- Fine-grained detail generation + +--- + +## Model Components + +### 1. Input Embeddings + +#### Image Path + +``` +Image (RGB) → VAE Encoder → Latents [B, 64, H/8, W/8] + ↓ + Patchify + Linear + ↓ + Image Tokens [B, H*W, 3072] +``` + +**VAE**: AutoencoderKL (8x downsampling) +- Input: 1024×1024 RGB image +- Output: 64×128×128 latent + +**Linear Projection**: Maps 64 channels → 3072 hidden dim + +#### Text Path + +``` +Caption → T5-XXL Encoder → Embeddings [B, S, 4096] + ↓ + Linear Projection + ↓ + Text Tokens [B, S, 3072] + +Caption → CLIP-L Encoder → Pooled [B, 768] + ↓ + MLP Embedder + ↓ + Vector Embedding [B, 3072] +``` + +**T5-XXL**: Context-rich text embeddings (max 512 tokens) +**CLIP-L**: Global style/semantic vector (768 dim) + +#### Conditioning + +``` +Timestep t ∈ [0, 1] → Sinusoidal Encoding → [B, 256] + ↓ + MLP + ↓ + Timestep Embedding [B, 3072] + +(Optional) Guidance scale g → Linear → [B, 3072] +``` + +**Combined Conditioning Vector**: +``` +vec = timestep_emb + clip_pooled_emb + [guidance_emb] +``` + +### 2. Position Embeddings (3D RoPE) + +Flux uses **3D Rotary Position Embeddings** for spatial awareness: + +**Axes**: +1. **Axis 0**: Channel groups (16 groups for 64 channels) +2. **Axis 1**: Height positions (e.g., 128 for 1024px) +3. **Axis 2**: Width positions (e.g., 128 for 1024px) + +**Implementation**: +```python +# Generate position IDs for each axis +pos_ids = [ + (h * W + w) // (16 * patch_size), # Axis 0: channel group + h, # Axis 1: height + w, # Axis 2: width +] + +# Compute frequencies for each axis +theta_i = theta ^ (2i / dim_axis) +freqs = pos_id / theta_i + +# Apply RoPE rotation +cos_freq = cos(freqs) +sin_freq = sin(freqs) +``` + +**Advantages**: +- Encodes spatial structure (height × width) +- Encodes channel relationships +- Works for any resolution (generalization) + +### 3. MMDiT Layer (Double Block) + +Each MMDiT layer performs: + +``` +Input: img [B, H*W, D], txt [B, S, D], vec [B, D] + +1. Pre-normalization (AdaLN with timestep conditioning) + img_norm = AdaLN(img, vec) + txt_norm = AdaLN(txt, vec) + +2. Joint Self-Attention + img_qkv = Linear(img_norm) # [B, H*W, 3*D] + txt_qkv = Linear(txt_norm) # [B, S, 3*D] + + # Concatenate for joint attention + joint_qkv = concat([img_qkv, txt_qkv], dim=1) # [B, H*W+S, 3*D] + + # Apply attention + joint_out = Attention(joint_qkv) # [B, H*W+S, D] + + # Split back + img_attn, txt_attn = split(joint_out, [H*W, S]) + +3. Gated Residual Addition + img = img + gate_img * img_attn + txt = txt + gate_txt * txt_attn + +4. Feed-Forward Networks (separate for img and txt) + img_mlp = AdaLN(img, vec) → Linear → GELU → Linear + txt_mlp = AdaLN(txt, vec) → Linear → GELU → Linear + + img = img + gate_img_mlp * img_mlp + txt = txt + gate_txt_mlp * txt_mlp + +Output: img [B, H*W, D], txt [B, S, D] +``` + +**Key Features**: +- **Shared attention space**: Image and text attend to each other +- **Adaptive gating**: Timestep-conditioned residual connections +- **Separate MLPs**: Modality-specific processing + +### 4. Flux Single Block + +After joint processing, image tokens go through single blocks: + +``` +Input: img [B, H*W, D], txt [B, S, D], vec [B, D] + +1. Concatenate (but don't split later) + combined = concat([img, txt], dim=1) # [B, H*W+S, D] + +2. Pre-normalization (AdaLN) + combined_norm = AdaLN(combined, vec) + +3. Self-Attention + qkv = Linear(combined_norm) # [B, H*W+S, 3*D] + attn_out = Attention(qkv) # [B, H*W+S, D] + +4. Gated Residual + combined = combined + gate * attn_out + +5. Feed-Forward + mlp_norm = AdaLN(combined, vec) + mlp_out = Linear → GELU → Linear + combined = combined + gate_mlp * mlp_out + +6. Extract image tokens + img = combined[:, :H*W, :] # Only use image part + +Output: img [B, H*W, D] (text tokens discarded) +``` + +**Rationale**: +- Text still influences attention (as keys/values) +- Output focuses on image generation +- More efficient than full joint processing + +### 5. Output Processing + +``` +Image Tokens [B, H*W, D] + ↓ + AdaLNContinuous(vec) ← Final timestep conditioning + ↓ + Linear Projection + ↓ + [B, H*W, 64] + ↓ + Reshape + ↓ + [B, 64, H, W] ← Predicted velocity field +``` + +--- + +## Data Flow + +### Complete Forward Pass + +``` + Input + │ + ┌───────────────────┼───────────────────┐ + │ │ │ + Image Text Timestep + [B,64,H,W] [B,S,4096] [B] + │ │ │ + ▼ ▼ ▼ + img_linear txt_linear time_emb + │ │ │ + ├─────── + ─────────┴─────── vec ──────┤ + │ (conditioning) │ + ▼ │ + img [B,H*W,3072] │ + txt [B,S,3072] │ + │ │ + │ ┌──────────┐ │ + └──────────────► MMDiT ├────◄───────┤ + ┌──────────────◄ Layer 1 ├────────────┘ + │ └──────────┘ + │ ┌──────────┐ + └──────────────► MMDiT ├────◄───────┐ + ┌──────────────◄ Layer 2 ├────────────┤ + │ └──────────┘ │ + ... vec + │ ┌──────────┐ │ + └──────────────► MMDiT ├────◄───────┘ + ┌──────────────◄ Layer N ├────────────┐ + │ └──────────┘ │ + │ │ + img [B,H*W,3072] │ + txt [B,S,3072] │ + │ │ + │ ┌──────────┐ │ + └──────────────► Single ├────◄───────┤ + ┌──────────────◄ Block 1 ├────────────┘ + │ └──────────┘ + │ ┌──────────┐ + └──────────────► Single ├────◄───────┐ + ┌──────────────◄ Block 2 ├────────────┤ + │ └──────────┘ │ + ... vec + │ ┌──────────┐ │ + └──────────────► Single ├────◄───────┘ + ┌──────────────◄ Block M ├────────────┐ + │ └──────────┘ │ + │ │ + img [B,H*W,3072] │ + │ │ + ▼ │ + AdaLNContinuous ◄───────────────────────────┘ + │ + ▼ + Linear Projection + │ + ▼ + Reshape to [B,64,H,W] + │ + ▼ + Predicted Velocity +``` + +### Training Data Flow + +``` +Original Image + │ + ▼ + VAE Encode → z_0 [B,64,H,W] + │ + ├─────────────┐ + │ │ + ▼ ▼ + z_0 Sample Noise → z_1 + │ │ + │ Sample t ~ Uniform(0,1) + │ │ + └──────┬──────┘ + │ + z_t = (1-t)*z_0 + t*z_1 ← Noisy latent + │ + ▼ + Flux Model(z_t, text, t) + │ + ▼ + v_pred [B,64,H,W] ← Predicted velocity + │ + ▼ + Loss = MSE(v_pred, z_1 - z_0) ← Flow matching loss +``` + +--- + +## Mathematical Formulation + +### Flow Matching Objective + +**Forward Process**: +``` +z_t = (1 - t) * z_0 + t * z_1, where t ~ Uniform(0, 1) +``` + +**Velocity Target**: +``` +v* = dz_t/dt = z_1 - z_0 +``` + +**Training Loss**: +``` +L = E_{z_0, z_1, t, c} [ ||v_θ(z_t, t, c) - (z_1 - z_0)||² ] + +where: + z_0 = VAE(image) # Original latent + z_1 ~ N(0, I) # Random noise + c = (text_emb, clip_pooled) # Conditioning + v_θ = Flux model # Predicted velocity +``` + +### Sampling (Inference) + +**Euler Integration** (first-order ODE solver): +``` +z_0 = z_1 # Start from noise +for t in [1.0, 0.9, ..., 0.1, 0.0]: + v_t = Flux(z_t, t, c) + z_{t-Δt} = z_t - Δt * v_t +``` + +**Higher-Order Solvers** (optional): +- Heun's method (2nd order) +- DPM-Solver (adaptive) + +### Classifier-Free Guidance + +During inference, use guidance scale `w`: + +``` +v_guided = v_uncond + w * (v_cond - v_uncond) + +where: + v_cond = Flux(z_t, t, c_text) # With text + v_uncond = Flux(z_t, t, c_empty) # Without text (null prompt) + w = guidance scale (typically 3-5) +``` + +**Implementation**: Use guidance embedding in config: +```python +config = FluxConfig(guidance_embed=True, guidance_scale=3.5) +``` + +### 3D RoPE Mathematics + +For position `(h, w)` in image: + +**Position IDs**: +``` +pid = [floor((h*W + w) / 16), h, w] # [channel_group, height, width] +``` + +**Frequencies**: +``` +θ_i = θ_base ^ (2i / d_axis), for i = 0, ..., d_axis/2 + +freq_{axis,i} = pid[axis] / θ_i +``` + +**RoPE Rotation**: +``` +q_rot = [q[:d/2] * cos(freq) - q[d/2:] * sin(freq), + q[:d/2] * sin(freq) + q[d/2:] * cos(freq)] + +k_rot = [k[:d/2] * cos(freq) - k[d/2:] * sin(freq), + k[:d/2] * sin(freq) + k[d/2:] * cos(freq)] +``` + +--- + +## Implementation Details + +### Memory Layout + +Megatron-Core uses **sequence-first format**: `[seq, batch, hidden]` + +**Conversions**: +```python +# User format: [B, C, H, W] +img_latents = torch.randn(B, 64, H, W) + +# Reshape to [B, H*W, 64] +img_seq = rearrange(img_latents, 'b c h w -> b (h w) c') + +# Project to hidden_size +img_tokens = linear(img_seq) # [B, H*W, 3072] + +# Convert to Megatron format: [seq, batch, hidden] +img_megatron = rearrange(img_tokens, 'b s d -> s b d') +``` + +### Adaptive Layer Normalization + +**Standard AdaLN**: +```python +class AdaLN: + def forward(self, timestep_emb): + modulation = MLP(SiLU(timestep_emb)) + shift, scale, gate = split(modulation, 3) + return shift, scale, gate + + @staticmethod + def modulate(x, shift, scale): + return LayerNorm(x) * (1 + scale) + shift +``` + +**Usage in Layer**: +```python +shift, scale, gate = adaln(timestep_emb) +x_norm = AdaLN.modulate(x, shift, scale) +x_attn = attention(x_norm) +x = x + gate * x_attn # Gated residual +``` + +### Attention Implementation + +**Using Megatron SelfAttention**: +```python +from megatron.core.transformer.attention import SelfAttention + +attn = SelfAttention( + config=config, + submodules=submodules, + layer_number=layer_idx, + attn_mask_type=AttnMaskType.no_mask, # Flux uses no mask +) + +# Megatron expects [seq, batch, hidden] +output = attn(hidden_states) +``` + +**Joint Attention Trick**: +```python +# Concatenate image and text +combined = torch.cat([img, txt], dim=0) # [seq_img+seq_txt, batch, hidden] + +# Single attention call processes both +joint_output = attention(combined) + +# Split back +img_out = joint_output[:seq_img] +txt_out = joint_output[seq_img:] +``` + +--- + +## Megatron-Core Integration + +### TransformerConfig + +Flux uses Megatron's `TransformerConfig`: + +```python +from megatron.core.transformer.transformer_config import TransformerConfig + +config = TransformerConfig( + num_layers=19 + 38, # joint + single + hidden_size=3072, + num_attention_heads=24, + ffn_hidden_size=3072 * 4, # Standard 4x expansion + layernorm_epsilon=1e-6, + hidden_dropout=0.0, + attention_dropout=0.0, + add_qkv_bias=True, + # ... other Megatron params +) +``` + +### Layer Specs + +Flux provides factory functions for layer specs: + +```python +from primus.backends.megatron.core.models.diffusion.flux.layer_spec import ( + get_flux_double_transformer_spec_for_backend, + get_flux_single_transformer_spec_for_backend, + get_flux_layer_spec, +) + +# For MMDiT layers (pass backend from config) +double_spec = get_flux_double_transformer_spec_for_backend(backend) + +# For single blocks +single_spec = get_flux_single_transformer_spec_for_backend(backend) + +# Or use high-level API for full TransformerBlock +layer_specs = get_flux_layer_spec(config, backend=backend) +``` + +### Distributed Training Support + +Flux inherits Megatron's parallelism: + +**Tensor Parallelism** (TP): +```python +config = TransformerConfig( + tensor_model_parallel_size=8, # 8-way TP + sequence_parallel=True, # Sequence parallelism +) +``` + +**Pipeline Parallelism** (PP): not supported for diffusion models. The forward +path runs embeddings/output head on every rank and does not relay activations +between stages, so `pipeline_model_parallel_size` must be 1 (PP > 1 is rejected +at config construction). + +**Data Parallelism**: Handled automatically by trainer + +--- + +## Performance Optimizations + +### 1. Transformer Engine + +Flux uses NVIDIA Transformer Engine for FP8 training: + +```python +from megatron.core.extensions.transformer_engine import ( + TELayerNormColumnParallelLinear, + TERowParallelLinear, +) + +# Automatically used in layer specs +linear_qkv = TELayerNormColumnParallelLinear(...) +linear_proj = TERowParallelLinear(...) +``` + +**Benefits**: +- FP8 matmul (faster, less memory) +- Fused operations (LayerNorm + Linear) +- Automatic scaling for numerical stability + +### 2. Flash Attention + +Enabled via Megatron: + +```python +config = TransformerConfig( + attention_type='flash_attention', # Use Flash Attention 2 +) +``` + +**Speedup**: 2-3x faster attention, 4x less memory + +### 3. Gradient Checkpointing + +For large models: + +```python +config = TransformerConfig( + recompute_granularity='selective', # Checkpoint expensive ops + recompute_method='uniform', + recompute_num_layers=19, # Checkpoint all joint layers +) +``` + +**Memory Savings**: ~40% reduction, ~20% slower + +### 4. Fused Operations + +```python +config = TransformerConfig( + bias_activation_fusion=True, # Fuse bias + activation + masked_softmax_fusion=True, # Fuse mask + softmax + gradient_accumulation_fusion=True, # Fuse grad accumulation +) +``` + +### 5. Mixed Precision + +```python +from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss + +# Using PyTorch autocast +with torch.autocast(device_type='cuda', dtype=torch.bfloat16): + output = model(img, txt, y, timesteps, img_ids, txt_ids) + target = noise - clean_latents + loss = compute_flow_matching_loss(output, clean_latents, noise) + +# Backward pass (handled automatically) +scaler = torch.cuda.amp.GradScaler() +scaler.scale(loss).backward() +scaler.step(optimizer) +scaler.update() +``` + +--- + +## Comparison with Other Models + +### Flux vs DiT + +| Aspect | DiT | Flux | +|--------|-----|------| +| Architecture | Image-only transformer | MMDiT (image + text joint) | +| Conditioning | AdaLN (injected) | Joint attention | +| Position Encoding | Learned 2D | 3D RoPE | +| Diffusion Type | DDPM | Flow matching | + +### Flux vs Stable Diffusion + +| Aspect | Stable Diffusion (UNet) | Flux (Transformer) | +|--------|-------------------------|---------------------| +| Backbone | UNet with ResNet blocks | Full transformer | +| Text Integration | Cross-attention | Joint self-attention | +| Scalability | Limited (U-Net bottleneck) | Excellent (transformer scaling) | +| Efficiency | Fast (fewer params) | Slower but higher quality | + +--- + +## Design Decisions + +### Why Two-Stage (Joint + Single)? + +**Joint Blocks**: +- Deep semantic understanding +- Text-image alignment +- Compositional reasoning + +**Single Blocks**: +- Spatial refinement +- Detail generation +- Efficient (no text processing) + +**Alternative**: All joint blocks → slower, marginal quality gain + +### Why Flow Matching? + +**Advantages over DDPM**: +1. **Simpler training**: Straight-line interpolation (no schedule design) +2. **Faster sampling**: Fewer steps (10-20 vs 50-100) +3. **Better mode coverage**: Straighter paths → less error accumulation + +**Math**: Rectified flow is ODE-based (vs SDE for DDPM) + +### Why 3D RoPE? + +**Advantages**: +1. **Spatial awareness**: Encodes (height, width) structure +2. **Resolution flexibility**: Works for any image size +3. **Channel grouping**: Models relationships between channels + +**Alternative**: Learned absolute embeddings → less flexible + +--- + +## Primus Implementation Highlights + +### TransformerBlock Architecture + +Primus's key architectural enhancement is the use of Megatron-Core's **TransformerBlock with heterogeneous layer specifications**: + +```mermaid +graph TD + subgraph traditional[Traditional Approach] + A[Model] --> B[double_blocks: ModuleList] + A --> C[single_blocks: ModuleList] + B --> D[Manual iteration] + C --> D + D --> E[Manual PP splitting] + end + + subgraph primus[Primus Approach] + F[Model] --> G[transformer: TransformerBlock] + G --> H[layer_specs: heterogeneous] + H --> J[Unified checkpoint format] + end +``` + +**Benefits**: +1. **Unified Checkpointing**: Single `transformer.layers.{0-56}` namespace +2. **Future-Proof**: Native support for new Megatron-Core features +3. **Cleaner Code**: No manual iteration over separate block lists + +### Layer Specification Pattern + +```python +# Primus approach +layer_specs = get_flux_layer_spec(config, backend=backend) +# Or manually: +# layer_specs = [ +# *[get_flux_double_transformer_spec_for_backend(backend) for _ in range(19)], +# *[get_flux_single_transformer_spec_for_backend(backend) for _ in range(38)], +# ] + +transformer = TransformerBlock( + config=config, + spec=TransformerBlockSubmodules(layer_specs=layer_specs) +) + +# Automatic PP slicing handled by TransformerBlock +# No manual offset calculation needed +``` + +### Checkpoint Format Comparison + +**Traditional Format**: +``` +double_blocks.0.attn.qkv.weight +double_blocks.18.mlp.fc2.bias +single_blocks.0.attn.qkv.weight +single_blocks.37.mlp.fc2.bias +``` + +**Primus Format** (TransformerBlock): +``` +transformer.layers.0.self_attention.linear_qkv.weight # Joint layer 0 +transformer.layers.18.mlp.linear_fc2.bias # Joint layer 18 +transformer.layers.19.self_attention.linear_qkv.weight # Single layer 0 +transformer.layers.56.mlp.linear_fc2.bias # Single layer 37 +``` + +Benefits: Simpler distributed checkpointing, easier layer inspection, consistent with Megatron GPT models. + +--- + +## Future Enhancements + +### Planned Features + +1. **ControlNet Support**: + - Spatial conditioning (pose, depth, edges) + - Residual connections from control encoder + +2. **Multi-Resolution Training**: + - Dynamic image sizes during training + - Aspect ratio bucketing + +3. **Efficient Sampling**: + - DPM-Solver integration + - Distillation for 1-step generation + +4. **LoRA Fine-Tuning**: + - Low-rank adaptation for custom styles + - Efficient personalization + +### Research Directions + +- **Sparse Attention**: Reduce quadratic complexity for high-res +- **Mixture of Experts**: Conditional computation for efficiency +- **3D Extension**: Video generation with Flux architecture + +--- + +## References + +### Papers + +1. **Flow Matching**: Lipman et al., "Flow Matching for Generative Modeling", 2022 +2. **Rectified Flow**: Liu et al., "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow", 2022 +3. **DiT**: Peebles & Xie, "Scalable Diffusion Models with Transformers", 2023 +4. **RoPE**: Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding", 2021 +5. **Transformer Engine**: NVIDIA, "Transformer Engine: Accelerating Transformer Training", 2022 + +### Code References + +- **Primus Flux**: `primus/backends/megatron/core/models/diffusion/flux/` +- **Megatron-Core**: `megatron/core/transformer/` +- **Transformer Engine**: `transformer_engine/pytorch/` +- **Official Flux**: Black Forest Labs (HuggingFace) + +--- + +## Appendix: Hyperparameters + +### Flux 535M Training + +```yaml +model: + num_joint_layers: 1 + num_single_layers: 1 + hidden_size: 3072 + num_attention_heads: 24 + +training: + batch_size: 256 # global + learning_rate: 1e-4 + weight_decay: 0.01 + lr_schedule: cosine + warmup_steps: 10000 + total_steps: 500000 + +optimization: + optimizer: AdamW + beta1: 0.9 + beta2: 0.999 + epsilon: 1e-8 + gradient_clip: 1.0 +``` + +### Flux 12B Training + +```yaml +model: + num_joint_layers: 19 + num_single_layers: 38 + hidden_size: 3072 + num_attention_heads: 24 + +training: + batch_size: 2048 # global, multi-node + learning_rate: 1e-4 + weight_decay: 0.01 + lr_schedule: cosine + warmup_steps: 10000 + total_steps: 1000000 + +optimization: + optimizer: AdamW + beta1: 0.9 + beta2: 0.95 + epsilon: 1e-8 + gradient_clip: 1.0 + +parallelism: + tensor_parallel: 8 + pipeline_parallel: 4 + data_parallel: 8 + sequence_parallel: True +``` + +--- + +*For API usage, see [api_reference.md](api_reference.md).* diff --git a/docs/backends/megatron/diffusion/fp8_training.md b/docs/backends/megatron/diffusion/fp8_training.md new file mode 100644 index 000000000..8f846dfc5 --- /dev/null +++ b/docs/backends/megatron/diffusion/fp8_training.md @@ -0,0 +1,514 @@ +# FP8 Training for Flux Models + +Complete guide for training Flux diffusion models with FP8 (8-bit floating point) precision on AMD MI300X GPUs using Transformer Engine's delayed scaling recipe. + +## Overview + +FP8 training provides significant memory and speed improvements while maintaining numerical stability through delayed scaling: + +- **~2x memory reduction** (activations and weights) +- **1.5-2x training speedup** on AMD MI300X +- **Maintains numerical stability** via delayed scaling +- **Enables larger batch sizes** or higher resolutions + +## Table of Contents + +- [Prerequisites](#prerequisites) +- [Quick Start](#quick-start) +- [Configuration](#configuration) +- [Performance Benchmarks](#performance-benchmarks) +- [Troubleshooting](#troubleshooting) +- [Best Practices](#best-practices) +- [AMD MI300X Specific](#amd-mi300x-specific) + +--- + +## Prerequisites + +### Hardware Requirements + +- **AMD MI300X GPUs** with ROCm 6.0+ support +- **Minimum GPUs:** + - Flux 535M: 1x MI300X (testing) + - Flux 12B: 2x MI300X with TP=2 (can train with FP8) + +### Software Requirements + +1. **ROCm 6.0+** with FP8 tensor core support +2. **Transformer Engine 2.1.0+** with ROCm backend +3. **PyTorch** with ROCm support +4. **Megatron-LM** (included in Primus) + +### Verification + +Verify your environment has: +- Transformer Engine 2.1.0+ with ROCm backend +- FP8 support (run `python3 -c "import transformer_engine.pytorch as te; print(te.fp8.is_fp8_available())"`) + +--- + +## Quick Start + +### Test FP8 with Flux 535M (Recommended First Step) + +```bash +# 1. Prepare test dataset (or use existing) +# See primus/configs/data/megatron/diffusion/README.md + +# 2. Train Flux 535M with FP8 +EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml \ +GPUS_PER_NODE=1 \ +bash examples/run_pretrain.sh +``` + +### Production Training with Flux 12B + +```bash +# After validating with 535M, scale to 12B (TransformerEngine FP8) +EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml \ +GPUS_PER_NODE=8 \ +NNODES=4 \ +bash examples/run_slurm_pretrain.sh + +# Or local-spec FP8 (no TransformerEngine dependency) +EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml \ +GPUS_PER_NODE=8 \ +NNODES=4 \ +bash examples/run_slurm_pretrain.sh +``` + +--- + +## Configuration + +### FP8 Model Configuration + +FP8 is configured at the model level. Two pre-configured files are available: + +- `primus/configs/models/megatron/diffusion/flux_535m_fp8.yaml` +- `primus/configs/models/megatron/diffusion/flux_12b_fp8.yaml` + +**Key FP8 Parameters:** + +```yaml +# Enable FP8 +fp8: "e4m3" # E4M3 format (recommended) + # Alternative: "hybrid" (E4M3 activations + E5M2 gradients) + +# FP8 Recipe +fp8_recipe: "delayed" # Delayed scaling (most stable) + # Alternatives: "tensorwise", "blockwise", "mxfp8" + +# Scaling Configuration +fp8_margin: 0 # Margin for scaling factor (0 = no margin) +fp8_amax_history_len: 1024 # History window for delayed scaling + # Larger = more stable, smaller = adapts faster +fp8_amax_compute_algo: "most_recent" # or "max" + +# Gradient Precision +fp8_wgrad: true # Enable FP8 for weight gradients (recommended) + +# Attention Precision +fp8_dot_product_attention: false # Keep attention in higher precision +fp8_multi_head_attention: false # Keep MHA in higher precision +``` + +### Training Configuration Adjustments + +**Batch Sizes with FP8:** + +```yaml +# Flux 12B - can increase batch size with FP8 memory savings +micro_batch_size: 2 # vs 1 for BF16 +global_batch_size: 256 # same as BF16 + +# Flux 535M - can increase significantly +micro_batch_size: 4 # vs 2 for BF16 +global_batch_size: 32 +``` + +**Optimizer Settings (Same as BF16):** + +```yaml +optimizer: adamw +lr: 1.0e-4 +min_lr: 1.0e-5 +weight_decay: 0.01 +clip_grad: 1.0 # Gradient clipping still important! +``` + +**Parallelism with FP8:** + +```yaml +# Flux 12B - can potentially reduce TP with FP8 +tensor_model_parallel_size: 2 # or reduce to 1 with FP8 +pipeline_model_parallel_size: 1 +context_parallel_size: 1 +``` + +--- + +## Performance Benchmarks + +### Memory Usage + +| Model | Precision | Memory/GPU | Batch Size | Notes | +|-------|-----------|------------|------------|-------| +| Flux 535M | BF16 | ~7-10GB | 2 | Baseline | +| Flux 535M | FP8 | ~3-5GB | 4 | ~50% reduction | +| Flux 12B | BF16 | ~40-50GB | 1 | TP=2 required | +| Flux 12B | FP8 | ~20-25GB | 2 | TP=2, ~50% reduction | + +### Training Speed + +| Model | Precision | Steps/sec | Speedup | Hardware | +|-------|-----------|-----------|---------|----------| +| Flux 535M | BF16 | ~20-30 | 1.0x | 1x MI300X | +| Flux 535M | FP8 | ~30-50 | 1.5-2x | 1x MI300X | +| Flux 12B | BF16 | ~0.5-1.0 | 1.0x | 32x MI300X | +| Flux 12B | FP8 | ~0.8-1.5 | 1.5-2x | 32x MI300X | + +### Expected Results + +- **Memory:** ~50% reduction vs BF16 +- **Speed:** 1.5-2x faster training +- **Quality:** Loss curves within 5% of BF16 +- **Convergence:** Similar or faster than BF16 + +--- + +## Troubleshooting + +### NaN or Inf in Losses + +**Problem:** Training becomes unstable with NaN/Inf values + +**Solutions:** + +1. **Increase scaling history:** + ```yaml + fp8_amax_history_len: 2048 # or 4096 + ``` + +2. **Disable FP8 for weight gradients:** + ```yaml + fp8_wgrad: false + ``` + +3. **Use more conservative scaling:** + ```yaml + fp8_amax_compute_algo: "max" # instead of "most_recent" + ``` + +4. **Add scaling margin:** + ```yaml + fp8_margin: 1 # or 2 + ``` + +5. **Keep first/last layers in BF16:** + ```yaml + first_last_layers_bf16: true + num_layers_at_start_in_bf16: 2 + num_layers_at_end_in_bf16: 2 + ``` + +### Out of Memory Even with FP8 + +**Problem:** Still hitting OOM errors with FP8 enabled + +**Solutions:** + +1. **Reduce micro batch size:** + ```yaml + micro_batch_size: 1 # back to minimum + ``` + +2. **Enable gradient checkpointing:** + ```yaml + recompute_granularity: "selective" # or "full" + recompute_method: "block" + ``` + +3. **Increase tensor parallelism:** + ```yaml + tensor_model_parallel_size: 4 # distribute more + ``` + +4. **Reduce sequence length:** + ```yaml + seq_length: 2048 # if applicable + ``` + +### FP8 Not Available + +**Problem:** Setup script shows "FP8 not available" + +**Checks:** + +1. **Verify GPU model:** + ```bash + rocm-smi --showproductname + # Should show MI300X + ``` + +2. **Check ROCm version:** + ```bash + rocm-smi --showversion + # Should be 6.0+ + ``` + +3. **Verify Transformer Engine:** + ```bash + python3 -c "import transformer_engine; print(transformer_engine.__version__)" + # Should be 2.1.0+ + ``` + +4. **Test FP8 directly:** + ```python + import transformer_engine.pytorch as te + print(te.fp8.is_fp8_available()) # Should be True + ``` + +### Slower Than Expected + +**Problem:** FP8 training is not faster than BF16 + +**Checks:** + +1. **Verify FP8 is actually enabled:** + - Check logs for FP8 context messages + - Run with `NCCL_DEBUG=INFO` to see precision info + +2. **Check batch size:** + - Ensure you increased micro_batch_size with FP8 + - Small batches may not show speedup + +3. **Verify tensor cores:** + - FP8 requires tensor core support + - Check ROCm driver configuration + +4. **Profile training:** + ```yaml + log_timers_to_tensorboard: true + ``` + - Compare FP8 vs BF16 step times + +--- + +## Best Practices + +### Recommended Workflow + +1. **Start with 535M:** + - Validate FP8 works correctly + - Test for 100-1000 steps + - Verify no NaN/Inf + +2. **Validate on small 12B run:** + - Train for 1000-5000 steps + - Compare loss with BF16 baseline + - Check memory and speed improvements + +3. **Production training:** + - Monitor closely for first 10K steps + - Watch for numerical issues + - Compare checkpoints with BF16 + +### Training Configuration + +**Conservative (stable):** +```yaml +fp8_recipe: "delayed" +fp8_amax_history_len: 2048 +fp8_amax_compute_algo: "max" +fp8_wgrad: false +``` + +**Balanced (recommended):** +```yaml +fp8_recipe: "delayed" +fp8_amax_history_len: 1024 +fp8_amax_compute_algo: "most_recent" +fp8_wgrad: true +``` + +**Aggressive (maximum performance):** +```yaml +fp8_recipe: "tensorwise" # Requires TE 2.2.0+ +fp8_amax_history_len: 512 +fp8_amax_compute_algo: "most_recent" +fp8_wgrad: true +``` + +### Monitoring + +**Key metrics to watch:** + +1. **Loss curves:** + - Should be smooth (no spikes) + - Should decrease normally + - Compare with BF16 baseline + +2. **Gradient norms:** + - Should be stable + - No sudden jumps to infinity + +3. **Memory usage:** + - Should be ~50% of BF16 + - Check with `rocm-smi` + +4. **Training speed:** + - Should be 1.5-2x faster + - Measure steps/second + +### Checkpointing + +- **FP8 checkpoints are compatible with BF16** +- Can switch between FP8/BF16 training +- Optimizer state includes FP8 scaling factors +- Checkpoints are same size as BF16 + +--- + +## Autotune (Local Spec FP8) + +> **Scope:** This section covers the **local-spec** FP8 path (`PrimusTurboFloat8LocalSpecProvider`, no TransformerEngine), e.g. `flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml`. The TransformerEngine prerequisites and checks elsewhere in this guide (`te.fp8.is_fp8_available()`, "FP8 Not Available") do **not** apply here -- this path quantizes via Primus Turbo directly. + +### Enable autotune (and do not pin the FP8 GEMM backend) + +The local-spec FP8 kernels benefit from the Primus-Turbo autotuner, which picks the best backend per GEMM shape. Enable it with `PRIMUS_TURBO_AUTO_TUNE=1`. + +`PRIMUS_TURBO_AUTO_TUNE=1` is necessary but **not sufficient**: an explicit `PRIMUS_TURBO_GEMM_BACKEND` short-circuits autotune (the FP8 kernel dispatcher returns the user-specified backend before the autotune step), so it must be unset (or scoped so it does not cover FP8) for autotune to engage. + +**Note:** some base images bake `PRIMUS_TURBO_GEMM_BACKEND` as an *empty string* rather than leaving it unset. An empty value is not treated as "unset" and can raise `KeyError ''` on the first FP8 GEMM. If you hit this, `unset PRIMUS_TURBO_GEMM_BACKEND` before launching. + +```bash +unset PRIMUS_TURBO_GEMM_BACKEND # or scope it so it does not cover FP8 +export PRIMUS_TURBO_AUTO_TUNE=1 + +EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml \ + bash examples/run_pretrain.sh +``` + +### Contrast with MXFP4 + +For MXFP4/FP4 + AITER with a tuned CSV, do the **opposite**: leave `PRIMUS_TURBO_AUTO_TUNE` unset, because autotune disables the AITER preshuffle fast path. See the [MXFP4 Training Guide](mxfp4_training.md) ("Preshuffle fast path"). Do not copy the MXFP4 env recipe for FP8. + +--- + +## AMD MI300X Specific + +### Environment Variables + +```bash +# Optional: set for better performance +export HSA_FORCE_FINE_GRAIN_PCIE=1 # Better PCIe performance +export NCCL_DEBUG=INFO # For debugging +export HSA_ENABLE_SDMA=0 # Disable SDMA for stability +``` + +### ROCm Optimization + +1. **HipBLASLt tuning:** + ```bash + # Generate optimal GEMM kernels for your hardware + # See ROCm documentation for details + ``` + +2. **NCCL configuration:** + ```bash + export NCCL_IB_DISABLE=0 # Enable InfiniBand if available + export NCCL_NET_GDR_LEVEL=3 # GPU Direct RDMA + ``` + +3. **Memory management:** + ```bash + export HSA_OVERRIDE_GFX_VERSION=9.4.2 # For MI300X + ``` + +### Known Issues + +1. **Transformer Engine ROCm support:** + - Verify TE version supports ROCm FP8 + - Some recipes may require specific TE versions + +2. **Numerical stability:** + - MI300X may require longer history (fp8_amax_history_len) + - Start conservative and tune + +3. **Multi-node training:** + - Ensure RCCL/NCCL properly configured + - Test single-node first + +--- + +## Testing + +### Integration Test + +```bash +# Quick 100-step validation run +EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml \ +GPUS_PER_NODE=1 \ +bash examples/run_pretrain.sh +``` + +### Convergence Test + +1. Train both BF16 and FP8 for 5000 steps +2. Compare loss curves (should be within 5%) +3. Generate images from checkpoints +4. Compare quality visually + +--- + +## Numerical verification status + +To set expectations for external users, the precision/convergence claims in this +codebase fall into two tiers: + +- **Backed by in-repo tests** (CI-runnable on supported hardware): structural and + convention checks — attention TE-vs-local-spec equivalence, RNG/seed + determinism, chimera init, VAE resample reproducibility, fused delayed-scale + update, and MLPerf warmup FP8 state. +- **Asserted, not yet backed by an in-repo test:** end-to-end *tensor parity* + against HuggingFace/Diffusers FLUX from a real checkpoint, and the exact MLPerf + v5.1 eval sample-count / validation-timestep semantics. These are validated by + internal reference runs but no committed test reproduces them. + +Tracked follow-ups (file as public-repo issues): + +1. A real-checkpoint forward-parity test (535M minimum) comparing Primus Flux + against HF/Diffusers within a documented tolerance. +2. A robustness test for the MLPerf validation-timestep fallback path. + +Treat any "bit-exact / matches NeMo / matches MLPerf / within X%" statement in +source comments as *asserted, unverified* until the parity test above lands. + +## References + +- [Transformer Engine Documentation](https://docs.nvidia.com/deeplearning/transformer-engine/) +- [AMD ROCm Documentation](https://rocm.docs.amd.com/) +- [Megatron-LM FP8 Guide](https://github.com/NVIDIA/Megatron-LM/blob/main/docs/llm/fp8.md) +- [FP8 Formats Explained (E4M3 vs E5M2)](https://arxiv.org/abs/2209.05433) + +--- + +## Support + +For issues or questions: + +1. Check this guide first +2. Verify Transformer Engine and FP8 support (see Prerequisites) +3. Review logs for error messages +4. File issue with: + - Hardware specs (GPU model, ROCm version) + - Software versions (TE, PyTorch, Megatron-LM) + - Config files used + - Error logs + +--- + +**Happy FP8 Training! 🚀** + +*Last updated: 2026-01-10* diff --git a/docs/backends/megatron/diffusion/mxfp4_training.md b/docs/backends/megatron/diffusion/mxfp4_training.md new file mode 100644 index 000000000..b0dfbf86d --- /dev/null +++ b/docs/backends/megatron/diffusion/mxfp4_training.md @@ -0,0 +1,210 @@ +# MXFP4 Training for Flux Models + +Guide for training Flux diffusion models in **MXFP4** (E2M1 mantissa + E8M0 block-of-32 scales) on AMD MI355X GPUs using Primus's local-spec MXFP4 implementation backed by Primus-Turbo and AITER. + +## Overview + +MXFP4 stores activations and weights in 4-bit microscale floating-point with one E8M0 exponent shared per block of 32 elements. The Primus integration: + +- Uses a **local spec** (`PrimusTurboMXFP4LocalSpecProvider`) with **no Transformer Engine dependency** — MXFP4 linear layers are self-contained autograd `Function`s that call Primus-Turbo's `gemm_fp4_impl` directly, so the path is `torch.compile`-friendly with minimal graph breaks. +- Keeps **attention, optimizer state / main params, and inter-rank communication in BF16**. Only the MMA inputs of the column- and row-parallel linears are quantized. +- Supports two backward modes via `mxfp4_backward_precision`: pure **MXFP4** (default) or **FP8** hybrid (E5M2 backward with tensorwise scaling on HipBLASLt). +- Dispatches the FP4 GEMM through Primus-Turbo's pluggable backend layer, which can route to either AITER (recommended for MI355X) or HipBLASLt. + +## Table of Contents + +- [Prerequisites](#prerequisites) +- [Quick Start](#quick-start) +- [Configuration](#configuration) +- [Primus-Turbo Backend Selection](#primus-turbo-backend-selection) +- [Tuned GEMMs](#tuned-gemms) +- [Troubleshooting](#troubleshooting) +- [Verification Status](#verification-status) + +--- + +## Prerequisites + +### Hardware + +- **AMD Instinct MI355X** (gfx950) with FP4 tensor-core support. The MXFP4 linear-layer modules assert `check_mxfp4_support()` at construction and will refuse to initialize on unsupported devices ([`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`](../../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). +- Single node (the local-spec layers require `tensor_model_parallel_size: 1`). + +### Software + +- ROCm-compatible install of `aiter` (provides `aiter.gemm_a4w4` and the tuned-config loader in `aiter/jit/core.py`). +- Primus-Turbo with FP4 backend registered (`primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl`). +- `enable_primus_turbo: true` and `use_turbo_attention: true` in the training config. + +--- + +## Quick Start + +The verified MXFP4 config is `examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml`. Launch with the AITER backend and the pre-tuned GEMM CSV: + +```bash +# Path to a checkout of the `tuned_gemm_configs` directory. +# Set TUNED_GEMM_DIR to wherever you have the tuned configs available. +export TUNED_GEMM_DIR=${TUNED_GEMM_DIR:-/path/to/tuned_gemm_configs} + +export EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml +export PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER +export AITER_CONFIG_GEMM_A4W4=$TUNED_GEMM_DIR/mi355x/flux_12b.csv +export AITER_LOG_TUNED_CONFIG=1 # recommended: confirms each shape hits the CSV + +bash examples/run_pretrain.sh +``` + +The pre-tuned CSV is distributed via an internal tuned-config source (`tuned_gemm_configs/mi355x/flux_12b.csv`). If you do not have access, omit `AITER_CONFIG_GEMM_A4W4` and AITER will fall back to its bundled `a4w4_blockscale_tuned_gemm.csv` (slower for Flux 12B shapes). + +--- + +## Configuration + +The relevant overrides in [`examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml`](../../../../examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml): + +```yaml +# MXFP4 precision +fp4: "mxfp4" +fp4_recipe: "mxfp4" # default is "nvfp4" in trainer_base.yaml; must override +mxfp4_backward_precision: "mxfp4" # "mxfp4" (pure) or "fp8" (hybrid) + +# Local spec + Primus-Turbo +transformer_impl: "local" +enable_primus_turbo: true +use_turbo_attention: true + +# Required by the MXFP4 linear-layer modules +tensor_model_parallel_size: 1 +gradient_accumulation_fusion: false +# sequence_parallel must remain false +``` + +### Knob semantics + +| Knob | Values | Notes | +|------|--------|-------| +| `fp4` | `"mxfp4"` | Top-level switch to enable FP4. | +| `fp4_recipe` | `"mxfp4"` for this guide | Default in [`primus/configs/modules/megatron/trainer_base.yaml`](../../../../primus/configs/modules/megatron/trainer_base.yaml) is `nvfp4`; the MXFP4 config overrides it. | +| `mxfp4_backward_precision` | `"mxfp4"` or `"fp8"` | Exhaustive set (checked by branch in [`primus_turbo_mxfp4_local.py`](../../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). `"fp8"` uses E5M2 with tensorwise HipBLASLt for backward. | +| `mxfp4_gradient_stochastic_rounding` | `true` / `false` | Optional. Enables SR on FP4 gradient quantization. | + +--- + +## Primus-Turbo Backend Selection + +The FP4 GEMM call is routed by `GEMMFP4KernelDispatcher` in `Primus-Turbo/primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py`. Backends are selected with the precision-scoped env var `PRIMUS_TURBO_GEMM_BACKEND` (declared in `Primus-Turbo/primus_turbo/common/constants.py`): + +```bash +# Single backend for every precision: +export PRIMUS_TURBO_GEMM_BACKEND=AITER + +# Precision-scoped (recommended): route FP4 GEMMs to AITER, leave others to defaults: +export PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER + +# Per-precision routing: +export PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER,FP8:HIPBLASLT +``` + +The dispatcher (`GlobalBackendManager` / `AutoKernelDispatcher.dispatch` in `Primus-Turbo/primus_turbo/pytorch/core/backend.py`) resolves the backend in this order: **explicit env > code-set > auto-tune > registered default > fallback**. + +### Preshuffle fast path + +When **all** of the following are true, MXFP4 GEMMs take the preshuffled fast path with no per-call shuffle overhead: + +- `PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER` (or `AITER`) is set. +- `PRIMUS_TURBO_AUTO_TUNE` is unset or `0`. + +The `_enable_preshuffle()` helper in `primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py` returns `True` under these conditions (FP4 backend pinned to AITER and auto-tune off), and the call becomes `aiter.gemm_a4w4(..., bpreshuffle=True)`. This helper reproduces the upstream `enable_preshuffle()` that Primus-Turbo removed in PR #383 ("refactor preshuffle ..."), which moved per-call preshuffle control onto `Float4QuantConfig.use_preshuffle`; Primus keeps the runtime probe locally because `MXFP4LinearFunction` passes a plain `bool` into its custom ops. + +> **Do not combine `PRIMUS_TURBO_AUTO_TUNE=1` with a tuned CSV.** Auto-tune disables the preshuffle fast path, so each call pays the shuffle cost while AITER still picks the same kernel internally. For production runs, leave `PRIMUS_TURBO_AUTO_TUNE` unset. + +--- + +## Tuned GEMMs + +AITER reads its tuned-GEMM CSV from the `AITER_CONFIG_GEMM_A4W4` env var (handled in `aiter/jit/core.py`; the default is the bundled `aiter/configs/a4w4_blockscale_tuned_gemm.csv`). Each row maps `(cu_num, M, N, K)` to a profiled kernel and split-K factor. + +For Flux 12B on MI355X, the pre-tuned CSV is provided by an internal tuned-config source (`tuned_gemm_configs/mi355x/flux_12b.csv`). See that directory's `README.md` for the tuning runbook, CSV schema, ASM-vs-CK kernel distinction, and re-tuning triggers. + +### Verifying the CSV is being used + +Set `AITER_LOG_TUNED_CONFIG=1`. AITER will log one line per **hit**: + +``` +shape is M:16384, N:9216, K:3072, found padded_M: 16384, N:9216, K:3072 is tuned on cu_num = 256 in /path/to/tuned_gemm_configs/mi355x/flux_12b.csv, kernel name is _ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E, splitK is 0! +``` + +**Miss** lines are printed unconditionally (no env var required) and look like: + +``` +shape is M:..., N:..., K:..., not found tuned config in /path/to/flux_12b.csv, will use default config! +``` + +Any miss line means the CSV needs re-tuning for that shape — follow the runbook in `tuned_gemm_configs/README.md`. + +### First-run JIT compile + +The two `a4w4_blockscale_*_intrawave_v3` CK kernels are JIT-compiled on first use (~2-5 min). Subsequent runs reuse the cached `.so` files. The ASM `f4gemm_bf16_per1x32Fp4_BpreShuffle_*` kernels are pre-compiled blobs shipped with AITER and incur no JIT cost. + +--- + +## Troubleshooting + +### `not found tuned config in {file}, will use default config!` + +The (M, N, K) shape is missing from your CSV. AITER will fall back to its compiled-in default, which is typically slow. Re-tune for that shape per the internal `tuned_gemm_configs/README.md` runbook (capture the shape via the same log line, append to the untuned CSV, re-run the AITER tuner, commit the new CSV). + +### Slow first iteration (~minutes), normal afterwards + +Expected — the first call to a CK-based `a4w4_blockscale_*` kernel triggers JIT compilation. Cached `.so` files are reused on subsequent starts. + +### `User specified backend AITER cannot handle the given inputs` + +Raised by `AutoKernelDispatcher.dispatch` when `GEMMFP4AITERBackend.can_handle` rejects the input. Common causes: + +- `M` not a multiple of 16, or `N` not a multiple of 16 (constants `AITER_FP4GEMM_M_MULTIPLE` / `AITER_FP4GEMM_N_MULTIPLE` in `gemm_fp4_impl.py`). +- Unsupported dtype combination (only `(float4_e2m1fn_x2, float4_e2m1fn_x2, fp16/bf16)` is supported). +- Non-NT layout (`trans_a=False, trans_b=True, trans_c=False`). + +Workaround: switch to `PRIMUS_TURBO_GEMM_BACKEND=FP4:HIPBLASLT` for unsupported shapes, or pad/reshape inputs. + +### `MXFP4ColumnParallelLinear requires tensor_model_parallel_size=1` + +The MXFP4 linear-layer modules assert on `tensor_model_parallel_size == 1`, `gradient_accumulation_fusion == False`, and `sequence_parallel == False` ([`primus_turbo_mxfp4_local.py`](../../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). Adjust the config accordingly. + +### NaN losses + +Switch to the hybrid backward mode, which keeps the FP4 forward but does the gradient GEMM in FP8 (E5M2 tensorwise on HipBLASLt): + +```yaml +mxfp4_backward_precision: "fp8" +``` + +If NaNs persist, also try `mxfp4_gradient_stochastic_rounding: true`. + +--- + +## Verification Status + +The public config has been smoke-tested end-to-end: 1000 iters on 8x MI355X (single node, micro-batch 64 / global 512, sequence length 512) completes in ~16-20 minutes with `PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER` and the tuned CSV. No errors across ranks; `pretrain() completed successfully`. + +Formal A/B benchmarks vs BF16 and FP8 (delayed and tensorwise) are pending and will be published once a representative suite is run; do not rely on the wall-clock numbers above as performance characterizations. + +--- + +## Source Code Pointers + +- MXFP4 spec provider: [`primus/backends/megatron/core/extensions/primus_turbo_local_spec.py`](../../../../primus/backends/megatron/core/extensions/primus_turbo_local_spec.py) (`PrimusTurboMXFP4LocalSpecProvider`). +- MXFP4 linear-layer autograd / fwd-bwd: [`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`](../../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py). +- Config schema defaults: [`primus/configs/modules/megatron/trainer_base.yaml`](../../../../primus/configs/modules/megatron/trainer_base.yaml). +- Dataclass field `mxfp4_backward_precision`: [`primus/backends/megatron/core/models/diffusion/common/config.py`](../../../../primus/backends/megatron/core/models/diffusion/common/config.py). +- FP4 backend selection (Primus-Turbo): `primus_turbo/common/constants.py`, `primus_turbo/pytorch/core/backend.py`, `primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py`. +- AITER tuned-config loader: `aiter/jit/core.py` (`AITER_CONFIG_GEMM_A4W4`). +- AITER A4W4 dispatch + hit/miss logging: `aiter/ops/gemm_op_a4w4.py`. + +## Related Documentation + +- [FP8 Training Guide](fp8_training.md) — companion guide for FP8. +- [Diffusion Architecture / Developer Guide](README.md). +- [Diffusion Examples README](../../../../examples/megatron/diffusion/README.md). diff --git a/docs/cli/README.md b/docs/cli/README.md index bdb5e618e..bd5061462 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -43,6 +43,22 @@ If you're running from the Primus repo root (after `git clone ... && cd Primus`) primus-cli direct -- benchmark gemm -M 4096 -N 4096 -K 4096 ``` +### Data Preparation Commands + +```bash +# Prepare a raw WebDataset (smaller, on-the-fly encoding during training) +primus-cli direct -- data diffusion-raw \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml + +# Prepare a pre-encoded dataset from HuggingFace +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml + +# Ingest MLPerf Flux1 pre-encoded data (streaming download + conversion) +primus-cli direct -- data diffusion-ingest \ + --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml +``` + ## 🎯 Three Execution Modes | Mode | Use Case | Command Example | diff --git a/examples/README.md b/examples/README.md index ce0e36460..7067527c8 100644 --- a/examples/README.md +++ b/examples/README.md @@ -208,6 +208,13 @@ The following models are supported out of the box via provided configuration fil | Mixtral-8x7B-v0.1 | [mistralai/Mixtral-8x7B-v0.1 ](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1) | [mixtral_8x7B_v0.1-BF16-pretrain.yaml](https://github.com/AMD-AGI/Primus/blob/main/examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-BF16-pretrain.yaml) | | | Mixtral-8x22B-v0.1 | [mistralai/Mixtral-8x22B-v0.1 ](https://huggingface.co/mistralai/Mixtral-8x22B-v0.1) | [mixtral_8x22B_v0.1-BF16-pretrain.yaml](https://github.com/AMD-AGI/Primus/blob/main/examples/megatron/configs/MI300X/mixtral_8x22B_v0.1-BF16-pretrain.yaml) | | +### Diffusion Models + +- **Flux** - Flow-based diffusion model for text-to-image generation + - Training guide: [examples/megatron/diffusion/README.md](megatron/diffusion/README.md) (Flux 535M and 12B) + - Architecture & developer docs: [docs/backends/megatron/diffusion/README.md](../docs/backends/megatron/diffusion/README.md) + - FP8 training: [docs/backends/megatron/diffusion/fp8_training.md](../docs/backends/megatron/diffusion/fp8_training.md) + --- ### 🏃‍♂️ How to Run a Supported Model diff --git a/examples/megatron/diffusion/README.md b/examples/megatron/diffusion/README.md new file mode 100644 index 000000000..d03f7babd --- /dev/null +++ b/examples/megatron/diffusion/README.md @@ -0,0 +1,317 @@ +# Flux Diffusion Model Training Examples + +Training examples for Flux diffusion models with Primus-Megatron on AMD GPUs. + +## Related Documentation + +- **Architecture & Developer Guide:** [docs/backends/megatron/diffusion/README.md](../../../docs/backends/megatron/diffusion/README.md) +- **API Reference:** [docs/backends/megatron/diffusion/api_reference.md](../../../docs/backends/megatron/diffusion/api_reference.md) +- **FP8 Training Guide:** [docs/backends/megatron/diffusion/fp8_training.md](../../../docs/backends/megatron/diffusion/fp8_training.md) +- **MXFP4 Training Guide:** [docs/backends/megatron/diffusion/mxfp4_training.md](../../../docs/backends/megatron/diffusion/mxfp4_training.md) +- **Dataset Preparation:** [primus/configs/data/megatron/diffusion/README.md](../../../primus/configs/data/megatron/diffusion/README.md) +- **Tests:** [tests/unit_tests/backends/megatron/diffusion/](../../../tests/unit_tests/backends/megatron/diffusion/) + +--- + +## Quick Start + +### Prerequisites + +- AMD Instinct GPU(s) (MI300X, MI325X, MI355X) +- Docker or Podman with ROCm support +- Primus Docker image: `docker.io/rocm/primus:v26.1` +- Prepared dataset (see [Dataset Preparation](../../../primus/configs/data/megatron/diffusion/README.md)) + +### 5-Minute Test Run + +1. **Prepare a small test dataset:** + +```bash +mkdir -p /tmp/flux_test_data/raw && cd /tmp/flux_test_data/raw + +for i in {000..099}; do + convert -size 512x512 xc:blue sample_${i}.jpg + echo "A blue square" > sample_${i}.txt +done + +tar -cf train-000000.tar sample_*.jpg sample_*.txt + +cat > dataset.yaml << 'EOF' +__module__: megatron.energon +__class__: CrudeWebdataset +subflavors: + encoding: raw +EOF + +energon prepare . --num-workers 4 +``` + +2. **Launch training:** + +```bash +EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \ +DATA_PATH=/tmp/flux_test_data \ +GPUS_PER_NODE=1 \ +bash examples/run_pretrain.sh +``` + +--- + +## Model Variants + +| Feature | Flux 535M | Flux 12B | +|---------|-----------|----------| +| Parameters | 535M | 12B | +| Joint Layers | 1 | 19 | +| Single Layers | 1 | 38 | +| Min GPUs | 1 | 8 (FSDP2 / DDP) | +| Recommended GPUs | 1-8 | 8-64 | +| Sharding | DP | FSDP2 (ZeRO-2/3) or DDP + distributed optimizer | +| Best For | Testing | Production | + +--- + +## Available Configurations + +The same configs are provided for both MI300X (`examples/megatron/configs/MI300X/diffusion/`) +and MI355X (`examples/megatron/configs/MI355X/diffusion/`). The only differences +are hardware-tuned batch sizes (MI300X has 192GB HBM3, MI355X has 256GB), so MI300X +uses smaller default micro/global batch sizes on the 12B DDP configs. + +### Shared (MI300X and MI355X) + +| Config | Description | +|--------|-------------| +| `flux_535m_pretrain.yaml` | Flux 535M, BF16, single/multi-GPU | +| `flux_535m_pretrain_fp8.yaml` | Flux 535M with FP8 precision | +| `flux_535m_with_guidance_embed.yaml` | Flux 535M with guidance embedding | +| `flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml` | Flux 12B, FSDP2, BF16, local spec | +| `flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml` | Flux 12B, FSDP2, FP8, local spec | +| `flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml` | Flux 12B, DDP, FP8, local spec (delayed scaling) | +| `flux_12b_ddp_energon_schnell_resample_te_spec.yaml` | Flux 12B, DDP, BF16, TransformerEngine spec | +| `flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml` | Flux 12B, DDP, FP8, TransformerEngine spec | + +### MI355X only + +| Config | Description | +|--------|-------------| +| `flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml` | Flux 12B, DDP, MXFP4, local spec | +| `flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml` | MLPerf benchmark reproduction (local spec FP8) | +| `flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml` | MLPerf benchmark reproduction (TE spec FP8) | + +> The `*_mlperf.yaml` configs reproduce the MLPerf Training Flux.1 benchmark +> (MLPerf logging + convergence target). Use the non-MLPerf configs above for +> general training. + +--- + +## Training Modes + +### Single-Node Training + +```bash +# Flux 535M (1-8 GPUs) +EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \ +GPUS_PER_NODE=8 \ +bash examples/run_pretrain.sh + +# Flux 12B (FSDP2, BF16, 8 GPUs) +EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml \ +GPUS_PER_NODE=8 \ +bash examples/run_pretrain.sh +``` + +### Multi-Node Training (SLURM) + +```bash +export DOCKER_IMAGE="docker.io/rocm/primus:v26.1" +export NNODES=8 +export GPUS_PER_NODE=8 + +EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml \ +bash examples/run_slurm_pretrain.sh +``` + +### MLPerf Benchmark Reproduction (MI355X) + +The `*_mlperf.yaml` configs reproduce the MLPerf Training Flux.1 benchmark and are +intended for benchmark reproduction rather than general training. + +```bash +# Step 1: Ingest MLPerf data +primus-cli direct -- data diffusion-ingest \ + --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml + +# Step 2: Train (configs already set vae_latent_mode: resample, vae_scale: 0.3611, vae_shift: 0.1159) +EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml \ +GPUS_PER_NODE=8 \ +bash examples/run_pretrain.sh +``` + +--- + +## FP8 Training + +FP8 provides ~2x memory reduction and 1.5-2x training speedup on AMD MI300X/MI355X GPUs. + +```bash +# Quick validation with 535M +EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml \ +GPUS_PER_NODE=1 \ +bash examples/run_pretrain.sh + +# Production with 12B (TransformerEngine FP8) +EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml \ +GPUS_PER_NODE=8 NNODES=4 \ +bash examples/run_slurm_pretrain.sh + +# Production with 12B (local-spec FP8, no TransformerEngine dependency) +EXP=examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml \ +GPUS_PER_NODE=8 NNODES=4 \ +bash examples/run_slurm_pretrain.sh +``` + +| Model | Precision | Memory/GPU | Batch Size | Speed | +|-------|-----------|------------|------------|-------| +| Flux 535M | BF16 | ~7-10GB | 2 | 1.0x | +| Flux 535M | FP8 | ~3-5GB | 4 | 1.5-2x | +| Flux 12B | BF16 | ~40-50GB | 1 | 1.0x | +| Flux 12B | FP8 | ~20-25GB | 2 | 1.5-2x | + +For configuration details, tuning recipes, benchmarks, and troubleshooting, see the [FP8 Training Guide](../../../docs/backends/megatron/diffusion/fp8_training.md). + +--- + +## MXFP4 Training + +MXFP4 (E2M1 + E8M0 block-of-32 scales) Flux 12B training on MI355X is supported via the local-spec provider (`PrimusTurboMXFP4LocalSpecProvider`, no TransformerEngine dependency). Forward and weight GEMMs run in FP4 through Primus-Turbo + AITER; attention, the optimizer state, and inter-rank communication stay in BF16. + +```bash +# Path to a checkout of the `tuned_gemm_configs` directory. +# Set TUNED_GEMM_DIR to wherever you have the tuned configs available. +export TUNED_GEMM_DIR=${TUNED_GEMM_DIR:-/path/to/tuned_gemm_configs} + +EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml \ +PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER \ +AITER_CONFIG_GEMM_A4W4=$TUNED_GEMM_DIR/mi355x/flux_12b.csv \ +AITER_LOG_TUNED_CONFIG=1 \ +bash examples/run_pretrain.sh +``` + +For configuration knobs, backend-selector semantics, tuned-GEMM verification, and troubleshooting, see the [MXFP4 Training Guide](../../../docs/backends/megatron/diffusion/mxfp4_training.md). + +--- + +## Converting HuggingFace Checkpoints + +Convert pre-trained HuggingFace Flux checkpoints to Primus/Megatron-Core format: + +```bash +python tools/checkpoint_conversion/convert_flux_hf_to_primus.py \ + --input black-forest-labs/FLUX.1-dev \ + --output checkpoints/primus_flux_12b.safetensors \ + --variant flux_12b +``` + +Supported variants: `flux_535m`, `flux_12b`, `custom` (with `--num-joint-layers` / `--num-single-layers`). + +For gated models (FLUX.1-dev), set `export HF_TOKEN="your_token"` or run `huggingface-cli login`. + +Primus also auto-detects tokens from `.hf_token` at the project root or `~/.cache/huggingface/token`. + +--- + +## Configuration Reference + +### Key Training Parameters + +```yaml +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_535m.yaml + trainer_class: FluxPretrainTrainer + + overrides: + train_iters: 100000 + micro_batch_size: 2 + global_batch_size: 16 + lr: 1.0e-4 + min_lr: 1.0e-5 + weight_decay: 0.01 + clip_grad: 1.0 + lr_decay_style: cosine + lr_warmup_iters: 1000 +``` + +### Parallelism + +The Flux 12B configs scale with FSDP2 (ZeRO-2/3) or DDP + distributed optimizer +rather than tensor/pipeline parallelism. + +| Setting | Flux 535M | Flux 12B | +|---------|-----------|----------| +| `tensor_model_parallel_size` | 1 | 1 | +| `pipeline_model_parallel_size` | 1 | 1 | +| `context_parallel_size` | 1 | 1 | +| Sharding | DP | FSDP2 (`use_torch_fsdp2: true`) or DDP (`use_distributed_optimizer: true`) | + +### Memory Optimization + +```yaml +modules: + pre_trainer: + overrides: + recompute_granularity: selective # or 'full' + recompute_method: block + sequence_parallel: false +``` + +--- + +## Troubleshooting + +### Dataset Not Found + +Verify `data_path` in your config, ensure `energon prepare` was run, and check that `dataset.yaml` exists. + +### Out of Memory (OOM) + +Reduce `micro_batch_size`, increase `tensor_model_parallel_size`, enable `recompute_granularity: full`, or switch to pre-encoded data mode. + +### NaN Loss + +Reduce learning rate to `1.0e-5`, ensure `clip_grad: 1.0`, increase `lr_warmup_iters`, check dataset for corruption. + +### NaN Loss with MLPerf Data + +Ensure config includes the required normalization constants: + +```yaml +vae_latent_mode: resample +vae_scale: 0.3611 +vae_shift: 0.1159 +``` + +### Encoder Download Fails + +Set `export HF_TOKEN=your_token` or use a local model path via `encoder_model_path` in config overrides. + +### Slow Training + +Use pre-encoded data (2-3x faster), increase `num_workers`, enable `use_flash_attn: true`. + +--- + +## Source Code Pointers + +- **Model Architecture:** `primus/backends/megatron/core/models/diffusion/flux/` +- **Trainer:** `primus/modules/trainer/megatron/diffusion/flux_pretrain_trainer.py` +- **Data Pipeline:** `primus/backends/megatron/data/diffusion/` +- **Model Configs:** `primus/configs/models/megatron/diffusion/` + +## Getting Help + +- [GitHub Issues](https://github.com/AMD-AGI/Primus/issues) +- [GitHub Discussions](https://github.com/AMD-AGI/Primus/discussions) diff --git a/examples/run_pretrain.sh b/examples/run_pretrain.sh index d93d251cb..03af31528 100755 --- a/examples/run_pretrain.sh +++ b/examples/run_pretrain.sh @@ -1,6 +1,6 @@ #!/bin/bash ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -89,6 +89,7 @@ PRIMUS_PATH=$(realpath "$(dirname "$0")/..") export DATA_PATH=${DATA_PATH:-"${PRIMUS_PATH}/data"} export HF_HOME=${HF_HOME:-"${DATA_PATH}/huggingface"} +# shellcheck source=/dev/null source "${PRIMUS_PATH}/runner/helpers/envs/path_utils.sh" LOG_INFO_RANK0 "Pip installing required packages ..." @@ -117,7 +118,14 @@ if [ ! -f "${EXP}" ]; then exit 1 fi -TRAIN_LOG=${TRAIN_LOG:-"output/log_mp_pretrain_$(basename "$EXP" .yaml).txt"} +if [ -z "${TRAIN_LOG:-}" ]; then + RUN_FOLDER=$(python3 -c " +import yaml, sys +d = yaml.safe_load(open(sys.argv[1])) +print(f\"{d.get('workspace','./output')}/{d.get('work_group','default')}/{d.get('user_name','unknown')}/{d.get('exp_name','experiment')}\") +" "$EXP" 2>/dev/null || echo "output") + TRAIN_LOG="${RUN_FOLDER}/train.log" +fi LOG_INFO_RANK0 "==========Training info==========" LOG_INFO_RANK0 "EXP: $EXP" diff --git a/runner/helpers/hooks/05_check_primus_requirements.sh b/runner/helpers/hooks/05_check_primus_requirements.sh new file mode 100755 index 000000000..f1e8be670 --- /dev/null +++ b/runner/helpers/hooks/05_check_primus_requirements.sh @@ -0,0 +1,68 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# Check that Primus requirements.txt packages are installed. +# Warns on missing packages by default. Set PRIMUS_STRICT_REQUIREMENTS=1 to +# make it a fatal error. +############################################################################### + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PRIMUS_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +REQ_FILE="${PRIMUS_ROOT}/requirements.txt" + +if [[ ! -f "$REQ_FILE" ]]; then + exit 0 +fi + +MISSING=$(python3 - "$REQ_FILE" << 'PYEOF' +import importlib.metadata, re, sys, pathlib + +req_file = pathlib.Path(sys.argv[1]) +missing = [] + +for line in req_file.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or line.startswith("-"): + continue + pkg = re.split(r"[><=!;\[\s]", line)[0].strip() + if not pkg: + continue + normalized = re.sub(r"[-_.]+", "-", pkg).lower() + try: + importlib.metadata.distribution(normalized) + except importlib.metadata.PackageNotFoundError: + missing.append(pkg) + +for m in missing: + print(m) +PYEOF +) + +if [[ -n "$MISSING" ]]; then + echo "" + echo "[WARN] Primus requirements not satisfied. Missing packages:" + while IFS= read -r pkg; do + echo " - $pkg" + done <<< "$MISSING" + echo "" + echo " To install: pip install -r ${REQ_FILE}" + echo "" + + if [[ "${PRIMUS_STRICT_REQUIREMENTS:-0}" == "1" ]]; then + echo "[ERROR] PRIMUS_STRICT_REQUIREMENTS=1 — aborting. Install missing packages first." + exit 1 + fi + + if [[ "${PRIMUS_AUTO_INSTALL:-0}" == "1" ]]; then + echo "[INFO] PRIMUS_AUTO_INSTALL=1 — installing missing packages..." + if pip install -r "${REQ_FILE}" --quiet; then + echo "[INFO] Successfully installed missing packages." + else + echo "[ERROR] Failed to install packages. Please install manually." + exit 1 + fi + fi +fi diff --git a/runner/helpers/hooks/train/pretrain/megatron/prepare.py b/runner/helpers/hooks/train/pretrain/megatron/prepare.py index 78873703a..694ad024e 100644 --- a/runner/helpers/hooks/train/pretrain/megatron/prepare.py +++ b/runner/helpers/hooks/train/pretrain/megatron/prepare.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -103,11 +103,41 @@ def prepare_dataset( def prepare_dataset_if_needed(primus_config: PrimusConfig, data_path: Path, env=None): pre_trainer_cfg = primus_config.get_module_config("pre_trainer") + + # Skip dataset preparation if train_data_path is explicitly set if pre_trainer_cfg.train_data_path is not None: return - tokenizer_type = pre_trainer_cfg.tokenizer_type - tokenizer_model = pre_trainer_cfg.tokenizer_model + # Check if this is a diffusion model (uses Energon datasets, not tokenized datasets) + model_type = getattr(pre_trainer_cfg, "model_type", None) + trainer_class = getattr(pre_trainer_cfg, "trainer_class", None) + data_path_config = getattr(pre_trainer_cfg, "data_path", None) + + # Determine if this is a diffusion model + is_diffusion = False + if model_type == "diffusion_model": + is_diffusion = True + elif trainer_class and ("Flux" in str(trainer_class) or "Diffusion" in str(trainer_class)): + is_diffusion = True + elif hasattr(model_type, "name") and "DIFFUSION" in model_type.name: + is_diffusion = True + + # For diffusion models with data_path set, skip tokenization (they use Energon) + if is_diffusion and data_path_config: + log_info("=" * 80) + log_info("Diffusion model detected with data_path configured.") + log_info("Skipping tokenization (diffusion models use Energon datasets).") + log_info(f"Data will be loaded from: {data_path_config}") + log_info("=" * 80) + return + + # For language models, proceed with tokenization + tokenizer_type = getattr(pre_trainer_cfg, "tokenizer_type", None) + if not tokenizer_type: + log_info("No tokenizer_type found, skipping dataset preparation.") + return + + tokenizer_model = getattr(pre_trainer_cfg, "tokenizer_model", None) default_tokenized_path = Path(data_path) / f"bookcorpus/{tokenizer_type}/bookcorpus_text_sentence" tokenized_data_path = Path(os.environ.get("TOKENIZED_DATA_PATH", str(default_tokenized_path))) @@ -120,6 +150,12 @@ def prepare_dataset_if_needed(primus_config: PrimusConfig, data_path: Path, env= if not hf_token: log_error_and_exit("Environment variable HF_TOKEN must be set.") + if not tokenizer_model: + log_error_and_exit( + "tokenizer_model not found in configuration. " + "This is required for language model tokenization." + ) + log_info(f"TOKENIZED_DATA_PATH is {tokenized_data_path}") prepare_dataset( @@ -179,9 +215,7 @@ def build_megatron_helper(megatron_path: Path): dataset_cpp_dir = megatron_path / "megatron/core/datasets" log_info(f"Building Megatron dataset helper in {dataset_cpp_dir}") - # `-s` silences make's "Nothing to be done"/recipe echo on no-op rebuilds; - # real compiler errors still surface and are handled below. - ret = subprocess.run(["make", "-s"], cwd=dataset_cpp_dir) + ret = subprocess.run(["make"], cwd=dataset_cpp_dir) if ret.returncode != 0: log_error_and_exit("Building Megatron C++ helper failed.") From bbbcfa8994a02747db8d03f03e81ca66c764a5c1 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Wed, 8 Jul 2026 03:01:03 +0300 Subject: [PATCH 007/127] feat(flux): FSDP2 fp32/bf16 optimizers + fp8 all-gather (#808) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/core` — review after it. The diff here is only this layer. ## What this changes The FSDP2 optimization layer used by Flux training: fp32 and bf16-master-weight optimizer variants, incremental grad-norm, the FSDP2 fp8 all-gather path, and the related torch-FSDP2 / fp8-cache / optimizer-registration patches. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/core`. It is the parent of the turbo layer, whose float8 extension lazily imports this layer's fp8 all-gather. ## Test plan `pytest tests/unit_tests/optimizer tests/unit_tests/backends/megatron/diffusion/distributed`. Validated locally on an AMD GPU container: 87 passed. ## Files 14 (FSDP2 optimizers, fp8 all-gather, optimizer/FSDP2 patches + tests). Co-authored-by: Flux Split Trial --- .../core/distributed/fsdp2_fp8_all_gather.py | 654 +++++++++++++++++ .../torch_fully_sharded_data_parallel.py | 419 ++++++++++- .../fsdp2_bf16_master_weight_optimizer.py | 686 ++++++++++++++++++ .../core/optimizer/fsdp2_fp32_optimizer.py | 368 ++++++++++ .../core/optimizer/incremental_grad_norm.py | 90 +++ .../patches/fsdp2_fp8_cache_patches.py | 98 +++ .../megatron/patches/optimizer_patches.py | 251 +++++++ .../megatron/patches/torch_fsdp2_patches.py | 61 +- .../diffusion/distributed/__init__.py | 8 + .../distributed/test_fsdp2_fp8_all_gather.py | 528 ++++++++++++++ .../test_fsdp2_transformer_impl.py | 109 +++ ...test_fsdp2_bf16_master_weight_optimizer.py | 351 +++++++++ .../optimizer/test_fsdp2_fp32_optimizer.py | 311 ++++++++ .../patches/test_fsdp2_fp32_patches.py | 488 +++++++++++++ 14 files changed, 4383 insertions(+), 39 deletions(-) create mode 100644 primus/backends/megatron/core/distributed/fsdp2_fp8_all_gather.py create mode 100644 primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py create mode 100644 primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py create mode 100644 primus/backends/megatron/core/optimizer/incremental_grad_norm.py create mode 100644 primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py create mode 100644 primus/backends/megatron/patches/optimizer_patches.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/distributed/__init__.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_fp8_all_gather.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_transformer_impl.py create mode 100644 tests/unit_tests/backends/megatron/optimizer/test_fsdp2_bf16_master_weight_optimizer.py create mode 100644 tests/unit_tests/backends/megatron/optimizer/test_fsdp2_fp32_optimizer.py create mode 100644 tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py diff --git a/primus/backends/megatron/core/distributed/fsdp2_fp8_all_gather.py b/primus/backends/megatron/core/distributed/fsdp2_fp8_all_gather.py new file mode 100644 index 000000000..9ef79d67a --- /dev/null +++ b/primus/backends/megatron/core/distributed/fsdp2_fp8_all_gather.py @@ -0,0 +1,654 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +FP8 all-gather tensor subclass for FSDP2. + +Wraps BF16 or FP32 parameters so that FSDP2 communicates FP8 data +(1 byte/element) instead of BF16/FP32, reducing all-gather volume. + +Approach A: Keep Weight in FP8 After All-Gather. +Uses a precomputed global scale (all-reduced amax across all shards) so that +the unsharded FP8 weight has a single consistent scale, compatible with +tensorwise GEMM on HIPBLASLT. After all-gather, the unsharded weight stays +in FP8 via FP8UnshardedWeightTensor, saving ~50% memory vs BF16 baseline. + +Reference: torchao WeightWithDynamicFloat8CastTensor +(ao/torchao/float8/fsdp_utils.py) +""" + +import math +from typing import Any, Optional, Tuple + +import torch +import torch.nn as nn +import torch.utils._pytree as pytree +import triton +import triton.language as tl +from primus_turbo.pytorch.core.low_precision import ( + Float8QuantConfig, + Format, + ScalingGranularity, + float8_e4m3, +) +from torch.distributed._tensor import DTensor +from torch.distributed._tensor.placement_types import Partial, Replicate +from torch.library import triton_op, wrap_triton + +_ops_to_preserve_subclass = { + torch.ops.aten.empty_like.default, + torch.ops.aten.new_zeros.default, + torch.ops.aten.slice.Tensor, + torch.ops.aten.copy_.default, + torch.ops.aten.view.default, + torch.ops.aten.as_strided.default, + torch.ops.aten._to_copy.default, + torch.ops.aten._pin_memory.default, + torch.ops.aten.split.Tensor, + torch.ops.aten.clone.default, +} + + +def _get_fp8_dtype(fmt: Format) -> torch.dtype: + if fmt == Format.E4M3: + return float8_e4m3 + elif fmt == Format.HYBRID: + return float8_e4m3 + else: + raise ValueError(f"Unsupported FP8 format for all-gather: {fmt}") + + +@triton.jit +def _quantize_fp8_prescaled_kernel( + input_ptr, + output_ptr, + scale_ptr, + n_elements, + FP8_MAX: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(input_ptr + offsets, mask=mask).to(tl.float32) + scale = tl.load(scale_ptr) + v = x * scale + v = tl.clamp(v, min=-FP8_MAX, max=FP8_MAX) + tl.store(output_ptr + offsets, v.to(output_ptr.dtype.element_ty), mask=mask) + + +@triton_op("primus::quantize_fp8_prescaled", mutates_args=()) +def quantize_fp8_prescaled( + x: torch.Tensor, + fp8_dtype: torch.dtype, + scale: torch.Tensor, + scale_inv: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + fp8_max = torch.finfo(fp8_dtype).max + output = torch.empty_like(x, dtype=fp8_dtype) + n_elements = x.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + wrap_triton(_quantize_fp8_prescaled_kernel)[grid]( + x, + output, + scale, + n_elements, + fp8_max, + BLOCK_SIZE=1024, + ) + return output, scale_inv.clone() + + +@quantize_fp8_prescaled.register_fake +def _quantize_fp8_prescaled_fake(x, fp8_dtype, scale, scale_inv): + return torch.empty_like(x, dtype=fp8_dtype), scale_inv.clone() + + +# --------------------------------------------------------------------------- +# Stochastic rounding variant: adds uniform [-0.5, 0.5) noise before +# truncation to make quantization error unbiased in expectation. +# --------------------------------------------------------------------------- + + +@triton.jit +def _quantize_fp8_prescaled_stochastic_kernel( + input_ptr, + output_ptr, + scale_ptr, + seed, + n_elements, + FP8_MAX: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(input_ptr + offsets, mask=mask).to(tl.float32) + scale = tl.load(scale_ptr) + v = x * scale + v = tl.clamp(v, min=-FP8_MAX, max=FP8_MAX) + noise = tl.rand(seed, offsets) - 0.5 + v_sr = v + noise + tl.store(output_ptr + offsets, v_sr.to(output_ptr.dtype.element_ty), mask=mask) + + +@triton_op("primus::quantize_fp8_prescaled_stochastic", mutates_args=()) +def quantize_fp8_prescaled_stochastic( + x: torch.Tensor, + fp8_dtype: torch.dtype, + scale: torch.Tensor, + scale_inv: torch.Tensor, + seed: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + fp8_max = torch.finfo(fp8_dtype).max + output = torch.empty_like(x, dtype=fp8_dtype) + n_elements = x.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + wrap_triton(_quantize_fp8_prescaled_stochastic_kernel)[grid]( + x, + output, + scale, + seed, + n_elements, + fp8_max, + BLOCK_SIZE=1024, + ) + return output, scale_inv.clone() + + +@quantize_fp8_prescaled_stochastic.register_fake +def _quantize_fp8_prescaled_stochastic_fake(x, fp8_dtype, scale, scale_inv, seed): + return torch.empty_like(x, dtype=fp8_dtype), scale_inv.clone() + + +@torch.compile(mode="max-autotune-no-cudagraphs") +def _foreach_fp8_quantize( + inner_tensors: list[torch.Tensor], + fp8_outputs: list[torch.Tensor], + scales: torch.Tensor, + fp8_max: float, +): + """Batch-quantize all FP8 parameters in a single fused operation. + + Uses _foreach_* ops fused by torch.compile to replace 456 individual + kernel launches with ~1-2 fused kernels. Float32 intermediate ensures + bitwise equivalence with the per-tensor quantize_fp8_prescaled path. + """ + scales_list = list(scales.unbind()) + f32 = [t.float() for t in inner_tensors] + scaled = torch._foreach_mul(f32, scales_list) + clamped = torch._foreach_clamp_min(torch._foreach_clamp_max(scaled, fp8_max), -fp8_max) + torch._foreach_copy_(fp8_outputs, clamped) + + +@torch.compile(mode="max-autotune-no-cudagraphs") +def _foreach_fp8_quantize_stochastic( + inner_tensors: list[torch.Tensor], + fp8_outputs: list[torch.Tensor], + scales: torch.Tensor, + fp8_max: float, +): + """Batch-quantize with stochastic rounding for unbiased FP8 quantization. + + Same as _foreach_fp8_quantize but adds uniform [-0.5, 0.5) noise before + the truncating cast to FP8, making quantization error zero in expectation. + """ + scales_list = list(scales.unbind()) + f32 = [t.float() for t in inner_tensors] + scaled = torch._foreach_mul(f32, scales_list) + noise = [torch.rand_like(s) - 0.5 for s in scaled] + noisy = torch._foreach_add(scaled, noise) + clamped = torch._foreach_clamp_min(torch._foreach_clamp_max(noisy, fp8_max), -fp8_max) + torch._foreach_copy_(fp8_outputs, clamped) + + +class WeightWithFP8AllGatherTensor(torch.Tensor): + """Tensor subclass that quantizes to FP8 before FSDP2 all-gather. + + Wraps a BF16 or FP32 parameter tensor. FSDP2 calls fsdp_pre_all_gather + to get FP8-quantized data for the collective, then fsdp_post_all_gather + to return an FP8UnshardedWeightTensor (keeping weight in FP8). + + Uses a precomputed global scale (set by precompute_fp8_scales_for_fsdp) + so all shards quantize with the same scale, producing a single consistent + scale for the unsharded weight. + """ + + @staticmethod + def __new__(cls, tensor: torch.Tensor, fp8_config: Float8QuantConfig): + if fp8_config.granularity != ScalingGranularity.TENSORWISE: + raise ValueError( + f"FP8 all-gather only supports TENSORWISE granularity, " f"got {fp8_config.granularity}" + ) + return torch.Tensor._make_wrapper_subclass( + cls, + tensor.size(), + strides=tensor.stride(), + storage_offset=tensor.storage_offset(), + dtype=tensor.dtype, + layout=tensor.layout, + device=tensor.device, + requires_grad=tensor.requires_grad, + ) + + def __init__(self, tensor: torch.Tensor, fp8_config: Float8QuantConfig): + self._fp8_config = fp8_config + self._tensor = tensor + # Transient: set by precompute_fp8_scales_for_fsdp(), NOT in + # __tensor_flatten__ (must be hashable/static for torch.compile guards). + self._precomputed_scale = None + self._precomputed_scale_inv = None + self._cached_fp8_data = None + self._use_cpp_quantize = False + self._stochastic_rounding = False + self._sr_counter = 0 + self._deq_after_ag = False + + def inner_data(self) -> torch.Tensor: + """Return the underlying data tensor (bypassing subclass dispatch).""" + return self._tensor + + def fsdp_pre_all_gather(self, mesh): + if self._cached_fp8_data is not None: + return (self._cached_fp8_data,), (self._precomputed_scale_inv, self._tensor.numel()) + fp8_dtype = _get_fp8_dtype(self._fp8_config.format) + if self._precomputed_scale is None: + raise RuntimeError( + "precompute_fp8_scales_for_fsdp() must be called before the first forward pass" + ) + with torch.no_grad(): + if self._use_cpp_quantize: + from primus_turbo.pytorch.ops.quantization import quantize_fp8 + + fp8_data, scale_inv = quantize_fp8( + self._tensor, + fp8_dtype, + self._fp8_config.granularity, + scale=self._precomputed_scale, + ) + elif self._stochastic_rounding: + if self._precomputed_scale_inv is None: + raise RuntimeError( + "precompute_fp8_scales_for_fsdp() must set scale_inv before stochastic-rounding quantize" + ) + if not self._tensor.is_contiguous(): + raise RuntimeError("FP8 stochastic-rounding quantize requires contiguous input") + self._sr_counter += 1 + fp8_data, scale_inv = quantize_fp8_prescaled_stochastic( + self._tensor, + fp8_dtype, + self._precomputed_scale, + self._precomputed_scale_inv, + seed=self._sr_counter, + ) + else: + if self._precomputed_scale_inv is None: + raise RuntimeError("precompute_fp8_scales_for_fsdp() must set scale_inv for Triton path") + if not self._tensor.is_contiguous(): + raise RuntimeError("FP8 prescaled quantize requires contiguous input") + fp8_data, scale_inv = quantize_fp8_prescaled( + self._tensor, + fp8_dtype, + self._precomputed_scale, + self._precomputed_scale_inv, + ) + return (fp8_data,), (scale_inv, self._tensor.numel()) + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional[torch.Tensor] = None, + ): + scale_inv, shard_numel = metadata + (fp8_gathered,) = all_gather_outputs + + if self._deq_after_ag: + bf16_weight = fp8_gathered.to(torch.bfloat16) * scale_inv + if out is not None: + out.data.copy_(bf16_weight) + return + return bf16_weight, (fp8_gathered,) + + if out is not None: + # Reshard path: FSDP already filled the FP8 buffer via + # all-gather into tracked storage. Just update the scale. + target = out.data if isinstance(out, nn.Parameter) else out + if isinstance(target, FP8UnshardedWeightTensor): + target._scale_inv = scale_inv + return + return FP8UnshardedWeightTensor(fp8_gathered, scale_inv, torch.bfloat16, self._fp8_config), ( + fp8_gathered, + ) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + if func == torch.ops.aten.detach.default: + return WeightWithFP8AllGatherTensor(args[0]._tensor.detach(), args[0]._fp8_config) + + if func == torch.ops.aten.copy_.default: + src = args[1] + if isinstance(src, WeightWithFP8AllGatherTensor): + src = src._tensor + args[0]._tensor.copy_(src) + return args[0] + + fp8_config = None + inner_dtype = None + + def unwrap(t): + nonlocal fp8_config, inner_dtype + if fp8_config is None: + fp8_config = t._fp8_config + inner_dtype = t._tensor.dtype + return t._tensor + + args, kwargs = pytree.tree_map_only(WeightWithFP8AllGatherTensor, unwrap, (args, kwargs or {})) + out = func(*args, **kwargs) + if func not in _ops_to_preserve_subclass: + return out + + if func == torch.ops.aten._to_copy.default: + target_dtype = (kwargs or {}).get("dtype", None) + if target_dtype is not None and target_dtype != inner_dtype: + return out + + return pytree.tree_map_only( + torch.Tensor, + lambda x: WeightWithFP8AllGatherTensor(x, fp8_config), + out, + ) + + def __tensor_flatten__(self): + # Guard for torch.compile tracing: dynamo may inspect this tensor + # via __tensor_flatten__ during __init__ before attributes are set. + config = getattr(self, "_fp8_config", None) + if config is None or not hasattr(self, "_tensor"): + return [], {} + # Float8QuantConfig is a mutable dataclass (not hashable), + # so store fields individually as hashable metadata. + return ["_tensor"], { + "format": config.format, + "granularity": config.granularity, + "strategy": config.strategy, + "scale_dtype": config.scale_dtype, + "block_size": config.block_size, + } + + @staticmethod + def __tensor_unflatten__(inner_tensors, metadata, outer_size, outer_stride): + config = Float8QuantConfig( + format=metadata["format"], + granularity=metadata["granularity"], + strategy=metadata["strategy"], + scale_dtype=metadata["scale_dtype"], + block_size=metadata["block_size"], + ) + return WeightWithFP8AllGatherTensor(inner_tensors["_tensor"], config) + + def __repr__(self): + return ( + f"WeightWithFP8AllGatherTensor(" + f"shape={list(self._tensor.shape)}, " + f"dtype={self._tensor.dtype}, " + f"device={self._tensor.device}, " + f"granularity={self._fp8_config.granularity})" + ) + + +_unsharded_ops_to_preserve = { + torch.ops.aten.as_strided.default, + torch.ops.aten.view.default, + torch.ops.aten.slice.Tensor, + torch.ops.aten.clone.default, +} + + +class FP8UnshardedWeightTensor(torch.Tensor): + """Lightweight wrapper for the unsharded FP8 weight after all-gather. + + Holds raw FP8 data + scalar inverse scale. Declares dtype=orig_dtype + (bfloat16) so PyTorch shape/dtype inference sees BF16, but stores 1 + byte/element. FP8 linear layers detect this subclass and extract the + pre-quantized data directly, skipping redundant quantization. + """ + + @staticmethod + def __new__( + cls, + fp8_data: torch.Tensor, + scale_inv: torch.Tensor, + orig_dtype: torch.dtype, + fp8_config: Float8QuantConfig, + ): + return torch.Tensor._make_wrapper_subclass( + cls, + fp8_data.size(), + strides=fp8_data.stride(), + storage_offset=fp8_data.storage_offset(), + dtype=orig_dtype, + layout=fp8_data.layout, + device=fp8_data.device, + requires_grad=False, + ) + + def __init__( + self, + fp8_data: torch.Tensor, + scale_inv: torch.Tensor, + orig_dtype: torch.dtype, + fp8_config: Float8QuantConfig, + ): + self._fp8_data = fp8_data + self._scale_inv = scale_inv + self._orig_dtype = orig_dtype + self._fp8_config = fp8_config + + def get_fp8_data_and_scale_inv(self): + return self._fp8_data, self._scale_inv + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + if func == torch.ops.aten.detach.default: + # Must preserve subclass: nn.Parameter(data) calls detach() + # and asserts type(detach_result) == type(data). + self = args[0] + return FP8UnshardedWeightTensor( + self._fp8_data.detach(), + self._scale_inv.detach(), + self._orig_dtype, + self._fp8_config, + ) + + if func in _unsharded_ops_to_preserve: + self = args[0] + new_data = func(self._fp8_data, *args[1:], **(kwargs or {})) + return FP8UnshardedWeightTensor(new_data, self._scale_inv, self._orig_dtype, self._fp8_config) + + raise NotImplementedError( + f"FP8UnshardedWeightTensor does not support {func}. " + f"Only FP8-aware module weights should be wrapped with FP8 all-gather." + ) + + def __tensor_flatten__(self): + config = getattr(self, "_fp8_config", None) + if config is None or not hasattr(self, "_fp8_data"): + return [], {} + return ["_fp8_data", "_scale_inv"], { + "orig_dtype": self._orig_dtype, + "format": config.format, + "granularity": config.granularity, + "strategy": config.strategy, + "scale_dtype": config.scale_dtype, + "block_size": config.block_size, + } + + @staticmethod + def __tensor_unflatten__(inner_tensors, metadata, outer_size, outer_stride): + config = Float8QuantConfig( + format=metadata["format"], + granularity=metadata["granularity"], + strategy=metadata["strategy"], + scale_dtype=metadata["scale_dtype"], + block_size=metadata["block_size"], + ) + return FP8UnshardedWeightTensor( + inner_tensors["_fp8_data"], + inner_tensors["_scale_inv"], + metadata["orig_dtype"], + config, + ) + + def __repr__(self): + return ( + f"FP8UnshardedWeightTensor(" + f"shape={list(self._fp8_data.shape)}, " + f"fp8_dtype={self._fp8_data.dtype}, " + f"orig_dtype={self._orig_dtype}, " + f"device={self._fp8_data.device})" + ) + + +def _wrap_fp8_weights_for_all_gather( + module: nn.Module, + fp8_config: Float8QuantConfig, + stochastic_rounding: bool = False, + deq_after_ag: bool = False, +) -> int: + """Wrap FP8-eligible weight parameters with WeightWithFP8AllGatherTensor. + + Must be called BEFORE fully_shard(). FSDP2 natively handles tensor + subclasses through fsdp_pre_all_gather/fsdp_post_all_gather. + + Args: + stochastic_rounding: If True, use stochastic rounding during FP8 + quantization to make quantization noise unbiased. + deq_after_ag: If True, dequantize FP8 back to BF16 after all-gather + and return a plain tensor. Downstream dynamic quantization will + produce a fresh scale from the actual assembled weight, avoiding + the numerical issues of keeping weight in FP8 with a global scale. + + Returns the number of wrapped parameters. + """ + wrapped_count = 0 + for child in module.modules(): + if not (hasattr(child, "_fp8_config") and hasattr(child, "weight")): + continue + w = child.weight + if w is not None and w.dtype in (torch.bfloat16, torch.float32) and w.requires_grad: + wrapper = WeightWithFP8AllGatherTensor(w.data, fp8_config) + wrapper._stochastic_rounding = stochastic_rounding + wrapper._deq_after_ag = deq_after_ag + child.weight = nn.Parameter(wrapper, requires_grad=True) + wrapped_count += 1 + return wrapped_count + + +@torch.no_grad() +def precompute_fp8_scales_for_fsdp( + module: nn.Module, + cache_data: bool = True, + use_cpp_quantize: bool = False, + stochastic_rounding: bool = False, +): + """Precompute global FP8 scales for all FP8-all-gather parameters. + + Call after optimizer.step(), before next forward. Uses _foreach_norm + for fused amax computation and DTensor Partial("max") -> Replicate() + redistribution for a single batched all-reduce across all parameters. + + Args: + module: The model module containing FP8-wrapped parameters. + cache_data: If True, batch-quantize all weights and cache the FP8 data + so fsdp_pre_all_gather skips per-layer quantization. If False, only + precompute scales; quantization happens on-demand per layer. + use_cpp_quantize: If True, use C++ quantize_fp8 from primus_turbo for + on-demand quantization (matching run_35). Skips fp8_outputs buffer + allocation and scale_inv computation when cache_data is False. + stochastic_rounding: If True, use stochastic rounding during batch + quantization to make quantization noise unbiased. + """ + need_buffers = cache_data or not use_cpp_quantize + need_scale_inv = cache_data or not use_cpp_quantize + + cache = getattr(module, "_fp8_scale_precompute_cache", None) + if cache is None: + fp8_dtensor_params = [] + fp8_plain_params = [] + inner_tensors = [] + for param in module.parameters(): + if isinstance(param, DTensor) and isinstance(param._local_tensor, WeightWithFP8AllGatherTensor): + fp8_dtensor_params.append(param) + inner_tensors.append(param._local_tensor._tensor) + elif isinstance(param, WeightWithFP8AllGatherTensor): + fp8_plain_params.append(param) + inner_tensors.append(param._tensor) + + fp8_config = ( + fp8_dtensor_params[0]._local_tensor._fp8_config + if fp8_dtensor_params + else fp8_plain_params[0]._fp8_config if fp8_plain_params else None + ) + fp8_dtype = _get_fp8_dtype(fp8_config.format) if fp8_config else None + + if need_buffers: + fp8_outputs = [torch.empty_like(t, dtype=fp8_dtype) for t in inner_tensors] + else: + fp8_outputs = [] + all_local = [ + p._local_tensor if isinstance(p, DTensor) else p for p in fp8_dtensor_params + fp8_plain_params + ] + cache = ( + fp8_dtensor_params, + fp8_plain_params, + inner_tensors, + fp8_outputs, + all_local, + fp8_config, + ) + module._fp8_scale_precompute_cache = cache + + ( + fp8_dtensor_params, + fp8_plain_params, + inner_tensors, + fp8_outputs, + all_local, + fp8_config, + ) = cache + if not inner_tensors: + return + + local_amaxes = torch.stack(torch._foreach_norm(inner_tensors, ord=math.inf)) + + if fp8_dtensor_params: + mesh = fp8_dtensor_params[0].device_mesh + partial_amaxes = DTensor.from_local(local_amaxes, device_mesh=mesh, placements=[Partial("max")]) + global_amaxes = partial_amaxes.redistribute(device_mesh=mesh, placements=[Replicate()]).to_local() + else: + global_amaxes = local_amaxes + + fp8_max = torch.finfo(_get_fp8_dtype(fp8_config.format)).max + global_amaxes = global_amaxes.to(torch.float64).clamp(min=1e-12) + scales = (fp8_max / global_amaxes).to(torch.float32) + scale_invs = (1.0 / scales) if need_scale_inv else None + + if cache_data: + if stochastic_rounding: + _foreach_fp8_quantize_stochastic(inner_tensors, fp8_outputs, scales, fp8_max) + else: + _foreach_fp8_quantize(inner_tensors, fp8_outputs, scales, fp8_max) + + for i, local_t in enumerate(all_local): + local_t._precomputed_scale = scales[i] + local_t._precomputed_scale_inv = scale_invs[i] if scale_invs is not None else None + local_t._cached_fp8_data = fp8_outputs[i] if cache_data else None + local_t._use_cpp_quantize = use_cpp_quantize + + +torch.serialization.add_safe_globals([WeightWithFP8AllGatherTensor]) diff --git a/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py b/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py index 40e6dafe5..8e7e75118 100644 --- a/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py +++ b/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py @@ -1,32 +1,68 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### -from typing import List, Optional +from typing import Optional, Set import torch -from megatron.core import tensor_parallel +import torch.distributed as dist +from megatron.core import parallel_state, tensor_parallel +from megatron.core.distributed.data_parallel_base import _BaseDataParallel from megatron.core.distributed.distributed_data_parallel_config import ( DistributedDataParallelConfig, ) -from megatron.core.distributed.torch_fully_sharded_data_parallel import ( - TorchFullyShardedDataParallel, -) +from megatron.core.fp8_utils import is_float8tensor from megatron.core.models.common.embeddings.language_model_embedding import ( LanguageModelEmbedding, ) from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import TransformerLayer +from torch.distributed import ProcessGroup from primus.modules.module_utils import warning_rank_0 +try: + from torch.distributed import DeviceMesh + from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard + + HAVE_FSDP = True +except ImportError: + HAVE_FSDP = False + + +def _validate_mesh_ranks(rank_tensor, shard_group, replicate_group): + """Validate that mesh tensor slices match the actual process group ranks.""" + current_rank = dist.get_rank() -class PrimusTorchFullyShardedDataParallel(TorchFullyShardedDataParallel): + shard_group_ranks = dist.get_process_group_ranks(shard_group) + for row in rank_tensor.tolist(): + if current_rank in row: + if sorted(row) != sorted(shard_group_ranks): + raise RuntimeError(f"Mesh shard ranks {row} != group ranks {shard_group_ranks}") + + replicate_group_ranks = dist.get_process_group_ranks(replicate_group) + for col_idx in range(rank_tensor.shape[1]): + col = rank_tensor[:, col_idx].tolist() + if current_rank in col: + if sorted(col) != sorted(replicate_group_ranks): + raise RuntimeError(f"Mesh replicate ranks {col} != group ranks {replicate_group_ranks}") + + +class PrimusTorchFullyShardedDataParallel(_BaseDataParallel): """ - Customized FSDP implementation for Primus framework, with pre-defined submodules to wrap. + Customized FSDP implementation for Primus framework with support for TransformerBlock. + + For models using TransformerBlock (e.g., Flux with heterogeneous layers), FSDP wraps + the TransformerBlock itself rather than individual TransformerLayer subclasses to avoid + duplicate mesh_dim_names errors. + + Key difference from base Megatron class: + - Prevents wrapping of modules that are descendants of already-wrapped modules + - This avoids the "Invalid mesh_dim_names ('dp', 'dp')" error when using TransformerBlock + with heterogeneous TransformerLayer subclasses. """ def __init__( @@ -34,20 +70,373 @@ def __init__( config: TransformerConfig, ddp_config: DistributedDataParallelConfig, module: torch.nn.Module, - sub_modules_to_wrap: Optional[List[torch.nn.Module]] = None, + sub_modules_to_wrap: Optional[Set[torch.nn.Module]] = None, + disable_bucketing: bool = False, + process_group: Optional[ProcessGroup] = None, **kwargs, ): + if not HAVE_FSDP: + raise RuntimeError("TorchFullyShardedDataParallel requires PyTorch >= 2.4.0 with FSDP 2 support.") + + super().__init__(config=config, module=module) + + # Store ddp_config for later access + self.ddp_config = ddp_config + + if process_group is None: + self.process_group = parallel_state.get_data_parallel_group(with_context_parallel=True) + else: + self.process_group = process_group + if sub_modules_to_wrap is None: - sub_modules_to_wrap = [ + sub_modules_to_wrap = { TransformerLayer, LanguageModelEmbedding, RotaryEmbedding, tensor_parallel.ColumnParallelLinear, - ] + } if kwargs: - warning_rank_0(f"PrimusTorchFullyShardedDataParallel: not use args: {kwargs}") + warning_rank_0(f"PrimusTorchFullyShardedDataParallel: unused args: {kwargs}") + + # Build DeviceMesh from Megatron's process groups + from megatron.training import get_args + + from primus.modules.module_utils import log_rank_0 + + args = get_args() + replicate_degree = getattr(args, "data_parallel_replicate_degree", 1) + dp_size = dist.get_world_size(self.process_group) + + if replicate_degree > 1: + shard_group = parallel_state.get_data_parallel_group( + with_context_parallel=True, partial_data_parallel=True + ) + replicate_group = parallel_state.get_inter_distributed_optimizer_instance_group() + + shard_size = dist.get_world_size(shard_group) + + full_dp_ranks = dist.get_process_group_ranks(self.process_group) + rank_tensor = torch.tensor(full_dp_ranks).reshape(replicate_degree, shard_size) + + _validate_mesh_ranks(rank_tensor, shard_group, replicate_group) + + mesh = DeviceMesh.from_group( + [replicate_group, shard_group], + device_type="cuda", + mesh=rank_tensor.tolist(), + mesh_dim_names=("dp_replicate", "dp_shard"), + ) + else: + mesh = DeviceMesh.from_group( + self.process_group, + device_type="cuda", + mesh_dim_names=("dp",), + ) + + reshard_after_forward = getattr(self.ddp_config, "reshard_after_forward", True) + + kwargs = { + "mesh": mesh, + "reshard_after_forward": reshard_after_forward, + } + + # When params_dtype is FP32 but training is BF16 (FSDP2 FP32 param optimizer), + # add MixedPrecisionPolicy so FSDP2 casts to BF16 for forward/backward. + # param_dtype=bf16 is always required: it ensures forward inputs + # (activations) get cast to BF16. For params with FP8 AG extensions, + # FSDP2 uses fsdp_post_all_gather (unaffected by param_dtype). + # This matches TorchTitan where param_dtype is always set from config. + use_fp8_all_gather = getattr(args, "use_fsdp2_fp8_all_gather", False) + if config.params_dtype == torch.float32 and config.bf16: + mp_policy = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + ) + kwargs["mp_policy"] = mp_policy + log_rank_0( + "FSDP2: MixedPrecisionPolicy(param_dtype=bf16, reduce_dtype=bf16) " + "[FP32 param optimizer: FSDP casts activations to BF16, BF16 reduce]" + ) + elif ( + config.params_dtype == torch.bfloat16 + and config.bf16 + and getattr(args, "use_fsdp2_bf16_master_weight_optimizer", False) + ): + # BF16 master weight optimizer: params are already BF16, so no + # param_dtype casting is needed. Setting param_dtype=None skips + # cast_forward_inputs entirely (avoids FP32→BF16 activation cast + # at every FSDP module boundary). reduce_dtype=bf16 halves + # ReduceScatter bandwidth; FP32 master weights in the optimizer + # still ensure full-precision parameter updates. + mp_policy = MixedPrecisionPolicy( + param_dtype=None, + reduce_dtype=torch.bfloat16, + ) + kwargs["mp_policy"] = mp_policy + log_rank_0( + "FSDP2: MixedPrecisionPolicy(param_dtype=None, reduce_dtype=bf16) " + "[BF16 master weight optimizer: no cast_forward_inputs, BF16 reduce]" + ) + + if replicate_degree > 1: + log_rank_0(f"FSDP2 Configuration:") + log_rank_0(f" mode: HSDP (replicate={replicate_degree}, shard={dp_size // replicate_degree})") + log_rank_0( + f" reshard_after_forward: {reshard_after_forward} " + f"(ZeRO-{'3' if reshard_after_forward else '2'})" + ) + log_rank_0( + f" Data parallel size: {dp_size} " + f"(replicate={replicate_degree} x shard={dp_size // replicate_degree})" + ) + else: + log_rank_0(f"FSDP2 Configuration:") + log_rank_0(f" mode: FSDP") + log_rank_0( + f" reshard_after_forward: {reshard_after_forward} " + f"(ZeRO-{'3' if reshard_after_forward else '2'})" + ) + log_rank_0(f" Data parallel size: {dp_size}") + + # Helper functions to save/restore custom parameter attributes + def save_custom_attrs(module): + custom_attrs = {} + for name, param in module.named_parameters(): + attrs = vars(param) + if is_float8tensor(param): + # disable fp8 transpose cache and perform transposing fp8 weights + # at each micro-batch because torch-FSDP doesn't recognize the + # micro-batch id, thus removing unnecessary memory stores + attrs["_fp8_attrs"]["transpose_invalid"] = False + del attrs["_fp8_attrs"]["transpose"] + custom_attrs[name] = {k: v for k, v in attrs.items()} + return custom_attrs + + def restore_custom_attrs(module, custom_attrs): + for name, param in module.named_parameters(): + if name in custom_attrs: + for attr_name, attr_value in custom_attrs[name].items(): + setattr(param, attr_name, attr_value) + + # Save custom attributes that might be removed by FSDP + attrs = save_custom_attrs(self.module) + + # FP8 all-gather validation + if ( + use_fp8_all_gather + and isinstance(reshard_after_forward, int) + and not isinstance(reshard_after_forward, bool) + ): + raise ValueError( + "FP8 all-gather is incompatible with reshard_after_forward=int (partial reshard)" + ) + + # FSDP2 FP8 all-gather is incompatible with delayed scaling: the delayed + # forward re-quantizes the weight and does not handle the all-gathered + # FP8UnshardedWeightTensor subclass produced by the all-gather path, so + # the combination would miscompute. Reject it explicitly. + if use_fp8_all_gather: + from megatron.core.enums import Fp8Recipe + + uses_delayed = ( + getattr(config, "fp8_scaling_strategy", "dynamic") == "delayed" + or getattr(config, "fp8_recipe", None) == Fp8Recipe.delayed + ) + if uses_delayed: + raise ValueError( + "use_fsdp2_fp8_all_gather is incompatible with delayed FP8 scaling " + "(fp8_recipe='delayed' / fp8_scaling_strategy='delayed'). Use a " + "non-delayed recipe (e.g. tensorwise) with FSDP2 FP8 all-gather, or " + "disable use_fsdp2_fp8_all_gather." + ) + + # Local transformer implementation does not support ColumnParallelLinear. + if config.transformer_impl == "local": + sub_modules_to_wrap = { + sub_module + for sub_module in sub_modules_to_wrap + if sub_module != tensor_parallel.ColumnParallelLinear + } + sub_modules_to_wrap = set(sub_modules_to_wrap) + + # Process _fsdp_modules attribute + for sub_module in self.module.modules(): + fsdp_modules = getattr(sub_module, "_fsdp_modules", []) + for f in fsdp_modules: + sub_modules_to_wrap.add(f) + + fp8_ag_stochastic_rounding = getattr(args, "fp8_all_gather_stochastic_rounding", False) + fp8_ag_deq_after_ag = getattr(args, "fp8_all_gather_deq_requant", False) + if use_fp8_all_gather: + from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + _wrap_fp8_weights_for_all_gather, + ) + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _build_fp8_config, + ) + + fp8_ag_config = _build_fp8_config(config) + wrapped_count = _wrap_fp8_weights_for_all_gather( + self.module, + fp8_ag_config, + stochastic_rounding=fp8_ag_stochastic_rounding, + deq_after_ag=fp8_ag_deq_after_ag, + ) + sr_tag = " [stochastic rounding]" if fp8_ag_stochastic_rounding else "" + deq_tag = " [deq+requant]" if fp8_ag_deq_after_ag else "" + log_rank_0( + f"FSDP2: Wrapped {wrapped_count} params with FP8 all-gather " + f"(granularity={fp8_ag_config.granularity}){sr_tag}{deq_tag}" + ) + + # ============================================================================ + # CUSTOM WRAPPING LOGIC: Skip descendants of already-wrapped modules + # ============================================================================ + wrapped_modules_set = set() + wrapped_list = [] + + for sub_module in self.module.modules(): + should_wrap = any( + isinstance(sub_module, sub_module_to_wrap) for sub_module_to_wrap in sub_modules_to_wrap + ) + if not should_wrap: + continue + + is_descendant = False + for wrapped_module in wrapped_modules_set: + if sub_module is not wrapped_module: + for child in wrapped_module.modules(): + if child is sub_module: + is_descendant = True + break + if is_descendant: + break + + if is_descendant: + continue + + fully_shard(sub_module, **kwargs) + wrapped_modules_set.add(sub_module) + wrapped_list.append(sub_module) + + log_rank_0(f"FSDP2: wrapped {len(wrapped_list)} inner modules + root") + + prefetch_depth = getattr(self.config, "fsdp_prefetch_depth", 1) + for i, mod in enumerate(wrapped_list): + fwd_targets = wrapped_list[i + 1 : i + 1 + prefetch_depth] + if fwd_targets: + mod.set_modules_to_forward_prefetch(fwd_targets) + + bwd_start = max(0, i - prefetch_depth) + bwd_targets = list(reversed(wrapped_list[bwd_start:i])) + if bwd_targets: + mod.set_modules_to_backward_prefetch(bwd_targets) + + # Wrap the root module as required by the FSDP API + fully_shard(self.module, **kwargs) + + if use_fp8_all_gather: + from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + precompute_fp8_scales_for_fsdp, + ) + + cache_data = getattr(self.config, "fp8_precompute_data_cache", True) + use_cpp = getattr(self.config, "use_cpp_fp8_quantize", False) + precompute_fp8_scales_for_fsdp( + self.module, + cache_data=cache_data, + use_cpp_quantize=use_cpp, + stochastic_rounding=fp8_ag_stochastic_rounding, + ) + + restore_custom_attrs(self.module, attrs) + + if getattr(args, "overlap_grad_norm", False): + from primus.backends.megatron.core.optimizer.incremental_grad_norm import ( + IncrementalGradNormAccumulator, + ) + + # Use the shard group for the norm all-reduce (correct for both + # FSDP and HSDP). For HSDP, replicate-group ranks hold identical + # gradients after AR, so the full DP group would over-count. + if replicate_degree > 1: + norm_reduce_group = shard_group + else: + norm_reduce_group = self.process_group + + accumulator = IncrementalGradNormAccumulator( + shard_process_group=norm_reduce_group, + device=torch.device("cuda"), + ) + n_hooked = 0 + for param in self.module.parameters(): + if param.requires_grad: + param.register_post_accumulate_grad_hook(accumulator.make_hook()) + n_hooked += 1 + args._grad_norm_accumulator = accumulator + log_rank_0(f"FSDP2: Registered incremental grad norm hooks on {n_hooked} params") + + def compile_model(self): + """ + Delegate torch.compile to the underlying model's compile_model() method. + + Traverses the FSDP2 module hierarchy (up to 3 levels deep) to find + the actual model (e.g., Flux) and calls its compile_model() if present. + Skipped if enable_torch_compile is False or no compile_model method is found. + """ + from primus.modules.module_utils import log_rank_0 + + try: + from megatron.training import get_args + + args = get_args() + + # Check if compilation is enabled + if not getattr(args, "enable_torch_compile", False): + return + + except Exception: + # If args not available, skip compilation + log_rank_0(" ℹ Cannot access args for torch.compile settings, skipping") + return + + # FSDP2 wraps models in FSDPFloat16Module, which itself wraps the actual model + # We need to traverse the hierarchy to find the actual model (e.g., Flux) + current_module = self.module + module_path = [] + + # Traverse up to 3 levels deep to find a module with compile_model + for depth in range(3): + module_type = type(current_module).__name__ + module_path.append(module_type) + + if hasattr(current_module, "compile_model") and callable( + getattr(current_module, "compile_model") + ): + log_rank_0(f" Found compile_model at depth {depth}: {' -> '.join(module_path)}") + log_rank_0(f" Calling compile_model on {module_type}...") + current_module.compile_model() + log_rank_0(f" ✓ Model compilation complete") + return + + # Try to go deeper if there's a 'module' attribute + if hasattr(current_module, "module"): + current_module = current_module.module + else: + break + + # No compile_model found - log that we can't compile directly + # Note: We can't reassign self.module after FSDP wrapping, so if the underlying + # model doesn't have compile_model, we skip compilation. The model should implement + # compile_model if torch.compile is desired. + log_rank_0(f" ℹ No compile_model method found in module hierarchy: {' -> '.join(module_path)}") + log_rank_0(f" (Model should implement compile_model() if torch.compile is desired)") + + def finish_grad_sync(self, *args, **kwargs): + """No-op for FSDP2: gradient sync is handled by the FSDP runtime.""" - super().__init__( - config=config, ddp_config=ddp_config, module=module, sub_modules_to_wrap=sub_modules_to_wrap - ) + def load_state_dict(self, state_dict, strict=True): + """ + No-op because tensors are already loaded in-place by + `_load_base_checkpoint` with FSDP2.""" diff --git a/primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py b/primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py new file mode 100644 index 000000000..6b8058ce6 --- /dev/null +++ b/primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py @@ -0,0 +1,686 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +BF16 Master Weight optimizer for FSDP2 mixed precision training. + +Model parameters are stored natively in BF16 (no FP32-to-BF16 cast during +all-gather), while FP32 master weight copies are maintained internally for +optimizer precision. This follows Megatron's Float16OptimizerWithFloat16Params +pattern adapted for FSDP2 DTensor parameters. + +Compared to the FP32 optimizer (fsdp2_fp32_optimizer.py), this eliminates +the per-layer CopyFunctor kernel in every forward all-gather, +at the cost of +2 bytes/param for the FP32 master copy. + +Memory per parameter per GPU (sharded): + FP32 optimizer: 4 (FP32 param) + 4 (exp_avg) + 4 (exp_avg_sq) = 12 bytes + This optimizer: 2 (BF16 param) + 4 (FP32 master) + 4 + 4 = 14 bytes +""" + +from itertools import chain +from typing import TYPE_CHECKING, Callable, List, Optional + +import torch +from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.dist_checkpointing.optimizer import ( + get_param_id_to_sharded_param_map, + make_sharded_optimizer_tensor, + optim_state_to_sharding_state, +) +from megatron.core.optimizer.optimizer import MegatronOptimizer, _zero_grad_group_helper +from megatron.core.optimizer.optimizer_config import OptimizerConfig + +from primus.modules.module_utils import log_rank_0 + +if TYPE_CHECKING: + from megatron.core.process_groups_config import ProcessGroupCollection + + +def _safe_log_rank_0(msg: str): + try: + log_rank_0(msg) + except (AttributeError, TypeError): + import torch.distributed as dist + + if not dist.is_initialized() or dist.get_rank() == 0: + print(msg) + + +class FSDP2BF16MasterWeightOptimizer(MegatronOptimizer): + """BF16 optimizer with FP32 master weights for FSDP2. + + Model parameters live in BF16 (eliminating the FP32->BF16 cast in every + forward all-gather). FP32 master copies are maintained for the optimizer + step, matching Megatron's Float16OptimizerWithFloat16Params pattern. + + Three parameter groups are tracked: + bf16_groups: original BF16 model parameters (what FSDP2 shards/all-gathers) + fp32_from_bf16_groups: FP32 master copies (what the optimizer steps on) + fp32_from_fp32_groups: natively FP32 params (e.g. LayerNorm), no master copy + + Args: + optimizer: Base PyTorch optimizer (param_groups already point to FP32 masters). + config: OptimizerConfig from Megatron. + init_state_fn: Function to initialize optimizer state tensors. + bf16_groups: List of lists of original BF16 model parameters. + fp32_from_bf16_groups: List of lists of FP32 master weight copies. + fp32_from_fp32_groups: List of lists of natively FP32 parameters. + """ + + def __init__( + self, + optimizer: torch.optim.Optimizer, + config: OptimizerConfig, + init_state_fn: Callable, + bf16_groups: List[List[torch.nn.Parameter]], + fp32_from_bf16_groups: List[List[torch.Tensor]], + fp32_from_fp32_groups: List[List[torch.nn.Parameter]], + use_foreach: bool = True, + ): + super().__init__(optimizer, config, init_state_fn) + self._scale = torch.tensor([1.0], dtype=torch.float, device="cuda") + self.is_stub_optimizer = optimizer is None + + self.bf16_groups = bf16_groups + self.fp32_from_bf16_groups = fp32_from_bf16_groups + self.fp32_from_fp32_groups = fp32_from_fp32_groups + + self._use_foreach = use_foreach + self._foreach_ready = False + + from megatron.training import get_args + + args = get_args() + self._grad_norm_accumulator = getattr(args, "_grad_norm_accumulator", None) + self._validate_count = 0 + + @torch.compiler.disable + @torch.no_grad() + def warmup_foreach_cache(self): + """Eagerly initialize the foreach cache before torch.compile wraps step(). + + Creates temporary zero gradients to satisfy _init_foreach_cache's + requirement that grads exist (including DTensor placements for the + assertion at line ~164), then clears them. + Must be called after FSDP wrapping but before torch.compile wraps + optimizer.step. + """ + if self._foreach_ready or not self._use_foreach: + return + + from torch.distributed.tensor import DTensor + + for group in self.bf16_groups: + for p in group: + if p.grad is not None: + continue + if isinstance(p.data, DTensor): + local = torch.zeros( + p.data.to_local().shape, + dtype=p.dtype, + device=p.device, + ) + p.grad = DTensor.from_local( + local, + device_mesh=p.data.device_mesh, + placements=p.data.placements, + ) + else: + p.grad = torch.zeros(p.shape, dtype=p.dtype, device=p.device) + self._init_foreach_cache() + for group in self.bf16_groups: + for p in group: + p.grad = None + + @torch.compiler.disable + @torch.no_grad() + def _init_foreach_cache(self): + """Build cached flat tensor lists for _foreach_copy_ batched operations. + + Called lazily on first prepare_grads() invocation (grads must exist for + DTensor spec extraction). Uses public DTensor APIs only. + """ + from torch.distributed.tensor import DTensor + + try: + from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + WeightWithFP8AllGatherTensor, + ) + + self._has_fp8_subclass = True + except ImportError: + self._has_fp8_subclass = False + + bf16_flat = [p for group in self.bf16_groups for p in group] + fp32_flat = [p for group in self.fp32_from_bf16_groups for p in group] + self._bf16_flat = bf16_flat + self._fp32_flat = fp32_flat + + bf16_inners = [] + for p in bf16_flat: + t = p.data + if isinstance(t, DTensor): + t = t.to_local() + if self._has_fp8_subclass and isinstance(t, WeightWithFP8AllGatherTensor): + t = t.inner_data() + bf16_inners.append(t) + self._bf16_inners = bf16_inners + + fp32_locals = [] + for m in fp32_flat: + t = m + if isinstance(t, DTensor): + t = t.to_local() + fp32_locals.append(t) + self._fp32_locals = fp32_locals + + fp32_grad_locals = [torch.empty_like(fl) for fl in fp32_locals] + self._fp32_grad_locals = fp32_grad_locals + + fp32_grad_dtensors = [] + for m, gl, bp in zip(fp32_flat, fp32_grad_locals, bf16_flat): + if isinstance(m, DTensor): + if m.placements != bp.grad.placements: + raise RuntimeError( + f"FP32 master placements {m.placements} != " + f"BF16 grad placements {bp.grad.placements}" + ) + dg = DTensor.from_local(gl, device_mesh=m.device_mesh, placements=m.placements) + else: + dg = gl + fp32_grad_dtensors.append(dg) + m.grad = dg + self._fp32_grad_dtensors = fp32_grad_dtensors + + self._foreach_ready = True + n_params = len(bf16_flat) + _safe_log_rank_0( + f"[FSDP2BF16MasterWeightOptimizer] foreach cache initialized: " + f"{n_params} params, fp8_subclass={self._has_fp8_subclass}" + ) + + def zero_grad(self, set_to_none=True): + if self.is_stub_optimizer: + return + if self._grad_norm_accumulator is not None: + self._grad_norm_accumulator.reset() + for group in self.bf16_groups: + _zero_grad_group_helper(group, set_to_none) + if not self._foreach_ready: + for group in self.fp32_from_bf16_groups: + _zero_grad_group_helper(group, set_to_none) + for group in self.fp32_from_fp32_groups: + _zero_grad_group_helper(group, set_to_none) + + def get_loss_scale(self): + return self._scale + + @torch.compiler.disable + @torch.no_grad() + def _prepare_grads_fallback(self) -> bool: + """Original per-parameter gradient copy (fallback path).""" + for bf16_group, fp32_group in zip(self.bf16_groups, self.fp32_from_bf16_groups): + for bf16_param, fp32_master in zip(bf16_group, fp32_group): + if bf16_param.grad is not None: + fp32_master.grad = bf16_param.grad.float() + bf16_param.grad = None + return False + + @torch.compiler.disable + def _validate_foreach_cache(self): + """Debug assertion: verify cached tensor refs haven't gone stale.""" + from torch.distributed.tensor import DTensor + + if self._has_fp8_subclass: + from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + WeightWithFP8AllGatherTensor, + ) + for i, p in enumerate(self._bf16_flat): + live = p.data + if isinstance(live, DTensor): + live = live.to_local() + if self._has_fp8_subclass and isinstance(live, WeightWithFP8AllGatherTensor): + live = live.inner_data() + if live.data_ptr() != self._bf16_inners[i].data_ptr(): + raise RuntimeError(f"Cached tensor ref stale for param {i}") + self._validate_count += 1 + + @torch.no_grad() + def prepare_grads(self) -> bool: + """Copy BF16 gradients to FP32 master grads. + + FSDP2 writes gradients to the BF16 model parameter's .grad attribute + (after FP32 reduce-scatter followed by cast to orig_dtype). We cast + them to FP32 for the optimizer step on master weights. + + Falls back to per-parameter loop if foreach cache is not ready or + any grad is missing. + """ + if not self._foreach_ready: + any_grad = any(p.grad is not None for group in self.bf16_groups for p in group) + if any_grad and self._use_foreach: + self._init_foreach_cache() + if not self._foreach_ready: + self._prepare_grads_fallback() + for fp32_group in self.fp32_from_fp32_groups: + for fp32_param in fp32_group: + if hasattr(fp32_param, "main_grad"): + fp32_param.grad = fp32_param.main_grad + return False + + from torch.distributed.tensor import DTensor + + bf16_grad_locals = [] + for p in self._bf16_flat: + if p.grad is None: + self._prepare_grads_fallback() + for fp32_group in self.fp32_from_fp32_groups: + for fp32_param in fp32_group: + if hasattr(fp32_param, "main_grad"): + fp32_param.grad = fp32_param.main_grad + return False + g = p.grad + if isinstance(g, DTensor): + g = g.to_local() + bf16_grad_locals.append(g) + + if self._validate_count < 3: + self._validate_foreach_cache() + + torch._foreach_copy_(self._fp32_grad_locals, bf16_grad_locals) + + for p in self._bf16_flat: + p.grad = None + + for fp32_group in self.fp32_from_fp32_groups: + for fp32_param in fp32_group: + if hasattr(fp32_param, "main_grad"): + fp32_param.grad = fp32_param.main_grad + + return False + + @torch.no_grad() + def clip_grad_norm(self, clip_grad: float) -> float | torch.Tensor: + """DTensor-native gradient clipping matching TorchTitan. + + Operates on FP32 master grads (from prepare_grads) plus any natively + FP32 parameter grads. + + When overlap_grad_norm is enabled, the squared norms have already been + accumulated in the RS stream via post_accumulate_grad_hooks on the + BF16 model params. The pre-computed norm is correct because + prepare_grads copies BF16->FP32 without scaling. The clipping is + applied to the FP32 master params. + """ + all_params = list(chain.from_iterable(self.fp32_from_bf16_groups)) + all_params.extend(chain.from_iterable(self.fp32_from_fp32_groups)) + + if self._grad_norm_accumulator is not None: + return self._grad_norm_accumulator.finalize(clip_grad, all_params) + + from torch.distributed.tensor import DTensor + + all_grads = [] + for fp32_group in self.fp32_from_bf16_groups: + for p in fp32_group: + if p.grad is not None: + all_grads.append(p.grad) + for fp32_group in self.fp32_from_fp32_groups: + for p in fp32_group: + if p.grad is not None: + all_grads.append(p.grad) + + if not all_grads: + return 0.0 + + total_norm = torch.nn.utils.get_total_norm(all_grads, norm_type=2.0, foreach=True) + if isinstance(total_norm, DTensor): + total_norm = total_norm.full_tensor() + + torch.nn.utils.clip_grads_with_norm_(all_params, clip_grad, total_norm, foreach=True) + return total_norm + + @torch.no_grad() + def step_with_ready_grads(self) -> bool: + """Step the optimizer on FP32 masters, then copy back to BF16 params.""" + if self.is_stub_optimizer: + return True + timers = self.config.timers + + if timers is not None: + timers("optimizer-inner-step", log_level=1).start(barrier=self.config.barrier_with_L1_time) + self.optimizer.step() + if timers is not None: + timers("optimizer-inner-step").stop() + + if timers is not None: + timers("optimizer-copy-main-to-model", log_level=1).start( + barrier=self.config.barrier_with_L1_time + ) + self._copy_main_params_to_model_params() + if timers is not None: + timers("optimizer-copy-main-to-model").stop() + + return True + + def _copy_main_params_to_model_params(self): + """Copy FP32 master weights back to BF16 model parameters.""" + if self._foreach_ready: + self._copy_main_to_model_foreach() + else: + self._copy_main_to_model_fallback() + + def _copy_main_to_model_foreach(self): + torch._foreach_copy_(self._bf16_inners, self._fp32_locals) + + @torch.compiler.disable + def _copy_main_to_model_fallback(self): + for bf16_group, fp32_group in zip(self.bf16_groups, self.fp32_from_bf16_groups): + for bf16_param, fp32_master in zip(bf16_group, fp32_group): + bf16_param.data.copy_(fp32_master.data) + + @torch.compiler.disable + def _copy_model_params_to_main_params(self): + """Copy BF16 model parameters to FP32 masters (for checkpoint reload).""" + if self._foreach_ready: + torch._foreach_copy_(self._fp32_locals, self._bf16_inners) + else: + self._copy_model_to_main_fallback() + + @torch.compiler.disable + def _copy_model_to_main_fallback(self): + for bf16_group, fp32_group in zip(self.bf16_groups, self.fp32_from_bf16_groups): + for bf16_param, fp32_master in zip(bf16_group, fp32_group): + fp32_master.data.copy_(bf16_param.data) + + @torch.no_grad() + def step(self): + """Clip gradients and step. Always succeeds (no overflow for BF16).""" + timers = self.config.timers + + found_inf_flag = self.prepare_grads() + if found_inf_flag: + return False, None, None + + if timers is not None: + timers("optimizer-clip-main-grad", log_level=1).start(barrier=self.config.barrier_with_L1_time) + grad_norm = None + if self.config.clip_grad > 0.0: + grad_norm = self.clip_grad_norm(self.config.clip_grad) + if timers is not None: + timers("optimizer-clip-main-grad").stop() + + if timers is not None: + timers("optimizer-count-zeros", log_level=1).start(barrier=self.config.barrier_with_L1_time) + num_zeros_in_grad = self.count_zeros() if self.config.log_num_zeros_in_grad else None + if timers is not None: + timers("optimizer-count-zeros").stop() + + success = self.step_with_ready_grads() + + return success, grad_norm, num_zeros_in_grad + + def reload_model_params(self, state_dict=None): + """After loading a checkpoint, copy BF16 model params to FP32 masters.""" + self._copy_model_params_to_main_params() + + def state_dict(self, is_loading: bool = False): + if is_loading: + self.init_state_fn(self.optimizer, self.config) + + state_dict = {} + state_dict["optimizer"] = self.optimizer.state_dict() + state_dict["fp32_from_fp16_params"] = self.fp32_from_bf16_groups + return state_dict + + def load_state_dict(self, state_dict): + optimizer_key = "optimizer" + if optimizer_key not in state_dict: + optimizer_key = "optimizer_state_dict" + + if "common_step" in state_dict[optimizer_key].get("state", {}): + common_step = state_dict[optimizer_key]["state"].pop("common_step") + self._restore_common_per_param_step(state_dict[optimizer_key], common_step) + + state_dict[optimizer_key]["param_groups"] = self._filter_and_reorder_param_groups( + self.optimizer.param_groups, state_dict[optimizer_key]["param_groups"] + ) + self.optimizer.load_state_dict(state_dict[optimizer_key]) + + # Restore FP32 master weights + if "fp32_from_fp16_params" in state_dict: + for current_group, saved_group in zip( + self.fp32_from_bf16_groups, state_dict["fp32_from_fp16_params"] + ): + for current_param, saved_param in zip(current_group, saved_group): + current_param.data.copy_(saved_param.data) + + def sharded_state_dict( + self, + model_sharded_state_dict: ShardedStateDict, + is_loading: bool = False, + metadata: Optional[dict] = None, + ): + if is_loading: + self.init_state_fn(self.optimizer, self.config) + + state_dict = self.state_dict() + + # Key on BF16 model params (not FP32 masters) for checkpoint sharding + # alignment with the model's sharded state dict + id_to_sharded_param_map = get_param_id_to_sharded_param_map( + model_sharded_state_dict, + chain.from_iterable(g for g in self.bf16_groups), + ) + + # Convert fp32_from_fp16_params to sharded tensors + if len(state_dict["fp32_from_fp16_params"]) != len(state_dict["optimizer"]["param_groups"]): + raise ValueError( + "state_dict fp32_from_fp16_params length does not match optimizer param_groups length" + ) + state_dict["fp32_from_fp16_params"] = [ + [ + make_sharded_optimizer_tensor( + id_to_sharded_param_map[param_id], + fp32_param, + prefix="optimizer.state.fp32_param", + ) + for param_id, fp32_param in zip(state_group["params"], fp32_group) + ] + for fp32_group, state_group in zip( + state_dict["fp32_from_fp16_params"], + state_dict["optimizer"]["param_groups"], + ) + ] + + step = self._extract_common_per_param_step(state_dict["optimizer"]) + + optim_state_to_sharding_state(state_dict["optimizer"], id_to_sharded_param_map, exclude_keys="step") + if step: + state_dict["optimizer"]["state"]["common_step"] = step + return state_dict + + def finalize_dist_ckpt_load(self, iteration): + """Restore optimizer step counter and sync FP32 masters to BF16 params. + + After dist_checkpointing in-place load (skip_load_to_model_and_opt=True), + load_state_dict is not called, so we manually set steps and copy masters. + """ + step_val = float(iteration) + for fp32_group in self.fp32_from_bf16_groups: + for p in fp32_group: + if p in self.optimizer.state and "step" in self.optimizer.state[p]: + self.optimizer.state[p]["step"].fill_(step_val) + for fp32_group in self.fp32_from_fp32_groups: + for p in fp32_group: + if p in self.optimizer.state and "step" in self.optimizer.state[p]: + self.optimizer.state[p]["step"].fill_(step_val) + + self._copy_main_params_to_model_params() + + +def get_fsdp2_bf16_master_weight_optimizer( + config: OptimizerConfig, + model_chunks: List[torch.nn.Module], + no_weight_decay_cond: Optional[Callable] = None, + scale_lr_cond: Optional[Callable] = None, + lr_mult: float = 1.0, + use_gloo_process_groups: bool = True, + default_skip_embedding_weight_decay: bool = False, + pg_collection: Optional["ProcessGroupCollection"] = None, + base_optimizer_cls=torch.optim.AdamW, + use_foreach: bool = True, + **optimizer_kwargs, +) -> FSDP2BF16MasterWeightOptimizer: + """Factory function to create FSDP2 BF16 master weight optimizer. + + Collects trainable parameters, creates FP32 master copies for BF16 params, + replaces param_group entries with FP32 masters, builds a fused AdamW, and + wraps in FSDP2BF16MasterWeightOptimizer. + """ + all_params = [] + for model_chunk in model_chunks: + for param in model_chunk.parameters(): + if param.requires_grad: + all_params.append(param) + + if not all_params: + raise ValueError("No trainable parameters found in model chunks!") + + weight_decay = config.weight_decay + lr = config.lr + param_groups = [] + + if default_skip_embedding_weight_decay and no_weight_decay_cond is None: + embedding_params = [] + non_embedding_params = [] + for param in all_params: + is_embedding = False + for model_chunk in model_chunks: + for name, p in model_chunk.named_parameters(): + if p is param and "embed" in name.lower(): + is_embedding = True + break + if is_embedding: + break + if is_embedding: + embedding_params.append(param) + else: + non_embedding_params.append(param) + + if embedding_params: + param_groups.append({"params": embedding_params, "weight_decay": 0.0, "lr": lr}) + if non_embedding_params: + param_groups.append({"params": non_embedding_params, "weight_decay": weight_decay, "lr": lr}) + elif no_weight_decay_cond is not None: + no_wd_params = [] + wd_params = [] + for param in all_params: + if no_weight_decay_cond(param): + no_wd_params.append(param) + else: + wd_params.append(param) + if no_wd_params: + param_groups.append({"params": no_wd_params, "weight_decay": 0.0, "lr": lr}) + if wd_params: + param_groups.append({"params": wd_params, "weight_decay": weight_decay, "lr": lr}) + else: + param_groups.append({"params": all_params, "weight_decay": weight_decay, "lr": lr}) + + if scale_lr_cond is not None and lr_mult != 1.0: + new_groups = [] + for param_group in param_groups: + scaled = [p for p in param_group["params"] if scale_lr_cond(p)] + normal = [p for p in param_group["params"] if not scale_lr_cond(p)] + if normal: + new_groups.append({**param_group, "params": normal}) + if scaled: + new_groups.append({**param_group, "params": scaled, "lr": param_group["lr"] * lr_mult}) + param_groups = new_groups + + # Create FP32 master copies and replace in param_groups + bf16_groups = [] + fp32_from_bf16_groups = [] + fp32_from_fp32_groups = [] + + for param_group in param_groups: + bf16_params_this_group = [] + fp32_from_bf16_this_group = [] + fp32_from_fp32_this_group = [] + + for i, param in enumerate(param_group["params"]): + if param.dtype in (torch.float16, torch.bfloat16): + bf16_params_this_group.append(param) + main_param = param.detach().clone().float() + param.main_param = main_param + param_group["params"][i] = main_param + fp32_from_bf16_this_group.append(main_param) + elif param.dtype == torch.float32: + fp32_from_fp32_this_group.append(param) + else: + _safe_log_rank_0(f"WARNING: unexpected param dtype {param.dtype}, treating as FP32") + fp32_from_fp32_this_group.append(param) + + bf16_groups.append(bf16_params_this_group) + fp32_from_bf16_groups.append(fp32_from_bf16_this_group) + fp32_from_fp32_groups.append(fp32_from_fp32_this_group) + + bf16_count = sum(len(g) for g in bf16_groups) + fp32_count = sum(len(g) for g in fp32_from_fp32_groups) + master_count = sum(len(g) for g in fp32_from_bf16_groups) + + _safe_log_rank_0( + f"Creating FSDP2 BF16 master weight optimizer with " + f"{bf16_count + fp32_count:,} parameters in {len(param_groups)} param groups" + ) + + base_optimizer = base_optimizer_cls( + param_groups, + betas=(config.adam_beta1, config.adam_beta2), + eps=config.adam_eps, + fused=True, + **optimizer_kwargs, + ) + + for param_group in base_optimizer.param_groups: + param_group.setdefault("wd_mult", 1.0) + param_group.setdefault("lr_mult", 1.0) + param_group.setdefault("is_expert_parallel", False) + param_group.setdefault("is_decoupled_lr", False) + param_group.setdefault("default_config", True) + + def init_state_fn(opt, config=None): + for group in opt.param_groups: + for p in group["params"]: + if len(opt.state[p]) == 0: + opt.state[p]["step"] = torch.zeros((), dtype=torch.float32, device=p.device) + opt.state[p]["exp_avg"] = torch.zeros_like(p.data) + opt.state[p]["exp_avg_sq"] = torch.zeros_like(p.data) + + optimizer = FSDP2BF16MasterWeightOptimizer( + optimizer=base_optimizer, + config=config, + init_state_fn=init_state_fn, + bf16_groups=bf16_groups, + fp32_from_bf16_groups=fp32_from_bf16_groups, + fp32_from_fp32_groups=fp32_from_fp32_groups, + use_foreach=use_foreach, + ) + + _safe_log_rank_0("=" * 80) + _safe_log_rank_0("[FSDP2BF16MasterWeightOptimizer Initialized]") + _safe_log_rank_0(" BF16 model params + FP32 master weights") + _safe_log_rank_0(" No FP32->BF16 cast in forward all-gather") + _safe_log_rank_0(" FP32 optimizer step + copy-back to BF16 per iteration") + _safe_log_rank_0(f" BF16 parameters: {bf16_count:,} ({master_count:,} FP32 master copies)") + _safe_log_rank_0(f" FP32 parameters: {fp32_count:,} (no master copy needed)") + _safe_log_rank_0(f" Total trainable parameters: {bf16_count + fp32_count:,}") + _safe_log_rank_0("=" * 80) + + optimizer.warmup_foreach_cache() + + return optimizer diff --git a/primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py b/primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py new file mode 100644 index 000000000..6c00caa9c --- /dev/null +++ b/primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py @@ -0,0 +1,368 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +FP32 optimizer for FSDP2 mixed precision training. + +Uses the TorchTitan approach: model parameters are initialized in FP32, +FSDP2's MixedPrecisionPolicy casts to BF16 for forward/backward, and +the optimizer operates on FP32 parameters with FP32 states. + +This eliminates the "stale weights" problem of BF16 optimizer states +while avoiding the master-copy duplication of Float16OptimizerWithFloat16Params. + +Memory impact vs BF16 optimizer (Flux 12B, 8 GPUs): + +3 GB/GPU for FP32 parameters (vs BF16) + +6 GB/GPU for FP32 optimizer states (vs BF16) + = +9 GB/GPU total + +Gradient clipping uses PyTorch-native DTensor-aware APIs matching TorchTitan. +""" + +from typing import TYPE_CHECKING, Callable, List, Optional + +import torch +from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.dist_checkpointing.optimizer import ( + get_param_id_to_sharded_param_map, + optim_state_to_sharding_state, +) +from megatron.core.optimizer.optimizer import MegatronOptimizer +from megatron.core.optimizer.optimizer_config import OptimizerConfig + +from primus.modules.module_utils import log_rank_0 + +if TYPE_CHECKING: + from megatron.core.process_groups_config import ProcessGroupCollection + + +def _safe_log_rank_0(msg: str): + try: + log_rank_0(msg) + except (AttributeError, TypeError): + import torch.distributed as dist + + if not dist.is_initialized() or dist.get_rank() == 0: + print(msg) + + +class FSDP2FP32Optimizer(MegatronOptimizer): + """FP32 optimizer for FSDP2 mixed precision training. + + Extends MegatronOptimizer directly (not MixedPrecisionOptimizer). + Modeled on Megatron's own FP32Optimizer with two key differences: + + 1. prepare_grads is a no-op: FSDP2 writes gradients directly to param.grad + (no main_grad -> grad copy needed). + 2. clip_grad_norm uses TorchTitan-style DTensor-native APIs: + torch.nn.utils.get_total_norm + clip_grads_with_norm_ for correct + norm computation across FSDP2's sharded DTensor parameters. + + Args: + optimizer: Base PyTorch optimizer (e.g., AdamW with fused=True). + config: OptimizerConfig from Megatron. + init_state_fn: Function to initialize optimizer state tensors. + """ + + def __init__( + self, + optimizer: torch.optim.Optimizer, + config: OptimizerConfig, + init_state_fn: Callable, + ): + super().__init__(optimizer, config, init_state_fn) + self._scale = torch.tensor([1.0], dtype=torch.float, device="cuda") + self.is_stub_optimizer = optimizer is None + + from megatron.training import get_args + + args = get_args() + self._grad_norm_accumulator = getattr(args, "_grad_norm_accumulator", None) + + def zero_grad(self, set_to_none=True): + if self.is_stub_optimizer: + return + if self._grad_norm_accumulator is not None: + self._grad_norm_accumulator.reset() + self.optimizer.zero_grad(set_to_none=set_to_none) + + def get_loss_scale(self): + return self._scale + + @torch.no_grad() + def prepare_grads(self) -> bool: + """No-op: FSDP2 writes gradients directly to param.grad.""" + return False + + @torch.no_grad() + def clip_grad_norm(self, clip_grad: float) -> float | torch.Tensor: + """DTensor-native gradient clipping matching TorchTitan. + + Uses torch.nn.utils.get_total_norm which natively handles DTensor + gradients (returns a DTensor with _NormPartial placement that is + reduced via full_tensor()), then clips with foreach-optimized + clip_grads_with_norm_. + + When overlap_grad_norm is enabled, the squared norms have already been + accumulated in the RS stream via post_accumulate_grad_hooks. Only a + single all-reduce + sqrt + clip is needed. + """ + params = self.get_parameters() + + if self._grad_norm_accumulator is not None: + return self._grad_norm_accumulator.finalize(clip_grad, params) + + from torch.distributed.tensor import DTensor + + grads = [p.grad for p in params if p.grad is not None] + + if not grads: + return 0.0 + + total_norm = torch.nn.utils.get_total_norm(grads, norm_type=2.0, foreach=True) + if isinstance(total_norm, DTensor): + total_norm = total_norm.full_tensor() + torch.nn.utils.clip_grads_with_norm_(params, clip_grad, total_norm, foreach=True) + return total_norm + + @torch.no_grad() + def step_with_ready_grads(self) -> bool: + if self.is_stub_optimizer: + return True + timers = self.config.timers + + if timers is not None: + timers("optimizer-inner-step", log_level=1).start(barrier=self.config.barrier_with_L1_time) + self.optimizer.step() + if timers is not None: + timers("optimizer-inner-step").stop() + + return True + + @torch.no_grad() + def step(self): + """Clip gradients and step. Always succeeds (no overflow for FP32).""" + timers = self.config.timers + + found_inf_flag = self.prepare_grads() + if found_inf_flag: + return False, None, None + + if timers is not None: + timers("optimizer-clip-main-grad", log_level=1).start(barrier=self.config.barrier_with_L1_time) + grad_norm = None + if self.config.clip_grad > 0.0: + grad_norm = self.clip_grad_norm(self.config.clip_grad) + if timers is not None: + timers("optimizer-clip-main-grad").stop() + + if timers is not None: + timers("optimizer-count-zeros", log_level=1).start(barrier=self.config.barrier_with_L1_time) + num_zeros_in_grad = self.count_zeros() if self.config.log_num_zeros_in_grad else None + if timers is not None: + timers("optimizer-count-zeros").stop() + + success = self.step_with_ready_grads() + + return success, grad_norm, num_zeros_in_grad + + def reload_model_params(self, state_dict=None): + pass + + def state_dict(self): + return self.optimizer.state_dict() + + def load_state_dict(self, state_dict): + if "common_step" in state_dict.get("state", {}): + common_step = state_dict["state"].pop("common_step") + self._restore_common_per_param_step(state_dict, common_step) + + state_dict["param_groups"] = self._filter_and_reorder_param_groups( + self.optimizer.param_groups, state_dict["param_groups"] + ) + self.optimizer.load_state_dict(state_dict) + + def sharded_state_dict( + self, + model_sharded_state_dict: ShardedStateDict, + is_loading: bool = False, + metadata: Optional[dict] = None, + ): + if is_loading: + self.init_state_fn(self.optimizer, self.config) + + state_dict = self.state_dict() + id_to_sharded_param_map = get_param_id_to_sharded_param_map( + model_sharded_state_dict, self.get_parameters() + ) + step = self._extract_common_per_param_step(state_dict) + + optim_state_to_sharding_state(state_dict, id_to_sharded_param_map, exclude_keys="step") + if step: + state_dict["state"]["common_step"] = step + return state_dict + + def finalize_dist_ckpt_load(self, iteration): + """Restore optimizer step counter after dist_checkpointing in-place load. + + When skip_load_to_model_and_opt=True (FSDP2), load_state_dict is not + called, so common_step is never fanned out to per-parameter step + entries. This fills step from the training iteration. + """ + step_val = float(iteration) + for p in self.get_parameters(): + if p in self.optimizer.state and "step" in self.optimizer.state[p]: + self.optimizer.state[p]["step"].fill_(step_val) + + +def get_fsdp2_fp32_optimizer( + config: OptimizerConfig, + model_chunks: List[torch.nn.Module], + no_weight_decay_cond: Optional[Callable] = None, + scale_lr_cond: Optional[Callable] = None, + lr_mult: float = 1.0, + use_gloo_process_groups: bool = True, + default_skip_embedding_weight_decay: bool = False, + pg_collection: Optional["ProcessGroupCollection"] = None, + base_optimizer_cls=torch.optim.AdamW, + use_foreach: bool = False, + **optimizer_kwargs, +) -> FSDP2FP32Optimizer: + """Factory function to create FSDP2 FP32 param optimizer from model chunks. + + Collects trainable FP32 parameters, builds param groups with weight decay + and LR scaling, creates AdamW (fused or foreach), and wraps in + FSDP2FP32Optimizer. + + Args: + config: OptimizerConfig from Megatron. + model_chunks: List of model modules (FSDP2-wrapped). + no_weight_decay_cond: Optional predicate for zero weight decay. + scale_lr_cond: Optional predicate for scaled learning rate. + lr_mult: Learning rate multiplier for scaled params. + use_gloo_process_groups: Unused (kept for API compatibility). + default_skip_embedding_weight_decay: Skip weight decay for embeddings + if no_weight_decay_cond not provided. + pg_collection: Unused (kept for API compatibility). + base_optimizer_cls: PyTorch optimizer class (default: AdamW). + use_foreach: If True, use foreach mode; if False (default), use fused. + **optimizer_kwargs: Additional kwargs for base optimizer. + + Returns: + FSDP2FP32Optimizer instance. + """ + all_params = [] + for model_chunk in model_chunks: + for param in model_chunk.parameters(): + if param.requires_grad: + all_params.append(param) + + if not all_params: + raise ValueError("No trainable parameters found in model chunks!") + + weight_decay = config.weight_decay + lr = config.lr + param_groups = [] + + if default_skip_embedding_weight_decay and no_weight_decay_cond is None: + embedding_params = [] + non_embedding_params = [] + for param in all_params: + is_embedding = False + for model_chunk in model_chunks: + for name, p in model_chunk.named_parameters(): + if p is param and "embed" in name.lower(): + is_embedding = True + break + if is_embedding: + break + if is_embedding: + embedding_params.append(param) + else: + non_embedding_params.append(param) + + if embedding_params: + param_groups.append({"params": embedding_params, "weight_decay": 0.0, "lr": lr}) + if non_embedding_params: + param_groups.append({"params": non_embedding_params, "weight_decay": weight_decay, "lr": lr}) + elif no_weight_decay_cond is not None: + no_wd_params = [] + wd_params = [] + for param in all_params: + if no_weight_decay_cond(param): + no_wd_params.append(param) + else: + wd_params.append(param) + if no_wd_params: + param_groups.append({"params": no_wd_params, "weight_decay": 0.0, "lr": lr}) + if wd_params: + param_groups.append({"params": wd_params, "weight_decay": weight_decay, "lr": lr}) + else: + param_groups.append({"params": all_params, "weight_decay": weight_decay, "lr": lr}) + + if scale_lr_cond is not None and lr_mult != 1.0: + new_groups = [] + for param_group in param_groups: + scaled = [p for p in param_group["params"] if scale_lr_cond(p)] + normal = [p for p in param_group["params"] if not scale_lr_cond(p)] + if normal: + new_groups.append({**param_group, "params": normal}) + if scaled: + new_groups.append({**param_group, "params": scaled, "lr": param_group["lr"] * lr_mult}) + param_groups = new_groups + + _safe_log_rank_0( + f"Creating FSDP2 FP32 optimizer with {len(all_params):,} parameters " + f"in {len(param_groups)} param groups" + ) + + base_optimizer = base_optimizer_cls( + param_groups, + betas=(config.adam_beta1, config.adam_beta2), + eps=config.adam_eps, + fused=not use_foreach, + foreach=use_foreach, + **optimizer_kwargs, + ) + + for param_group in base_optimizer.param_groups: + param_group.setdefault("wd_mult", 1.0) + param_group.setdefault("lr_mult", 1.0) + param_group.setdefault("is_expert_parallel", False) + param_group.setdefault("is_decoupled_lr", False) + param_group.setdefault("default_config", True) + + def init_state_fn(opt, config=None): + for group in opt.param_groups: + for p in group["params"]: + if len(opt.state[p]) == 0: + opt.state[p]["step"] = torch.zeros((), dtype=torch.float32, device=p.device) + opt.state[p]["exp_avg"] = torch.zeros_like(p.data) + opt.state[p]["exp_avg_sq"] = torch.zeros_like(p.data) + + optimizer = FSDP2FP32Optimizer( + optimizer=base_optimizer, + config=config, + init_state_fn=init_state_fn, + ) + + fp32_count = sum(1 for p in all_params if p.dtype == torch.float32) + bf16_count = sum(1 for p in all_params if p.dtype == torch.bfloat16) + other_count = len(all_params) - fp32_count - bf16_count + + _safe_log_rank_0("=" * 80) + _safe_log_rank_0("[FSDP2FP32ParamOptimizer Initialized]") + _safe_log_rank_0(" TorchTitan-style: FP32 params + FSDP2 MixedPrecisionPolicy") + _safe_log_rank_0(" DTensor-native gradient clipping") + _safe_log_rank_0(f" AdamW mode: {'foreach' if use_foreach else 'fused'}") + _safe_log_rank_0(f" FP32 parameters: {fp32_count:,}") + _safe_log_rank_0(f" BF16 parameters: {bf16_count:,}") + if other_count > 0: + _safe_log_rank_0(f" Other dtype parameters: {other_count:,}") + _safe_log_rank_0(f" Total trainable parameters: {len(all_params):,}") + _safe_log_rank_0("=" * 80) + + return optimizer diff --git a/primus/backends/megatron/core/optimizer/incremental_grad_norm.py b/primus/backends/megatron/core/optimizer/incremental_grad_norm.py new file mode 100644 index 000000000..9320c3244 --- /dev/null +++ b/primus/backends/megatron/core/optimizer/incremental_grad_norm.py @@ -0,0 +1,90 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Incremental gradient norm accumulator for overlapping grad norm with FSDP2 +reduce-scatter. + +Registers a `register_post_accumulate_grad_hook` on each sharded parameter. +FSDP2 fires these hooks inside the post-reduce stream (RS or AR stream) +after the sharded gradient is assigned. Each hook accumulates ||grad||_2^2 +into a shared tensor. By the time the default stream synchronises +(`_wait_for_post_backward`), the total squared norm is already computed. +The optimizer's `clip_grad_norm` then only needs a single all-reduce + sqrt ++ clip — eliminating the per-parameter norm recomputation that would +otherwise serialise after the last reduce-scatter. +""" + +from typing import Callable + +import torch +import torch.distributed as dist + + +class IncrementalGradNormAccumulator: + """Accumulates squared L2 norms of sharded gradients in the RS stream. + + Args: + shard_process_group: The FSDP shard process group used for the final + all-reduce of the accumulated squared norm. For HSDP this must be + the shard-only group (not the full DP group) because replicate-group + ranks hold identical gradients after the HSDP all-reduce. + device: CUDA device for the accumulator tensor. + """ + + def __init__( + self, + shard_process_group: dist.ProcessGroup, + device: torch.device, + ) -> None: + self._shard_pg = shard_process_group + self._total_norm_sq = torch.zeros((), dtype=torch.float32, device=device) + + def reset(self) -> None: + """Zero the accumulator. Called from optimizer.zero_grad().""" + self._total_norm_sq.zero_() + + def make_hook(self) -> Callable[[torch.Tensor], None]: + """Return a closure suitable for register_post_accumulate_grad_hook. + + The hook extracts the local shard data (bypassing DTensor overhead) + and accumulates its squared L2 norm into ``_total_norm_sq``. + """ + acc = self + + def hook(param: torch.Tensor) -> None: + grad = param.grad + if grad is None: + return + local_grad = grad._local_tensor if hasattr(grad, "_local_tensor") else grad + acc._total_norm_sq.add_(local_grad.float().norm(2).pow_(2)) + + return hook + + @torch.no_grad() + def finalize( + self, + clip_grad: float, + params, + ) -> torch.Tensor: + """All-reduce the accumulated norm, sqrt, clip, and return total_norm. + + No additional stream synchronisation is needed here: by the time the + optimizer calls this method the default stream has already waited on + ``post_reduce_event`` via ``_wait_for_post_backward``, so the + accumulated value is visible. + + Args: + clip_grad: Maximum gradient norm for clipping. + params: Iterable of parameters whose ``.grad`` will be clipped. + + Returns: + The global L2 gradient norm (scalar tensor). + """ + dist.all_reduce(self._total_norm_sq, op=dist.ReduceOp.SUM, group=self._shard_pg) + total_norm = self._total_norm_sq.sqrt_() + torch.nn.utils.clip_grads_with_norm_(params, clip_grad, total_norm, foreach=True) + return total_norm diff --git a/primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py b/primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py new file mode 100644 index 000000000..dbbbcc056 --- /dev/null +++ b/primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py @@ -0,0 +1,98 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Megatron train_step patch for FP8 all-gather cache refresh. + +The FP8 all-gather cache (populated by precompute_fp8_scales_for_fsdp) must be +refreshed after every optimizer.step() so that the cached FP8 data reflects the +updated weights. Without this patch the cache is only populated once during +FSDP setup and goes stale after the first iteration. + +This patch wraps Megatron's standalone train_step() to call +precompute_fp8_scales_for_fsdp(model[0]) after the original train_step returns. +""" + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.modules.module_utils import log_rank_0 + + +def _needs_fp8_cache_update(ctx: PatchContext) -> bool: + args = get_args(ctx) + return args is not None and getattr(args, "use_fsdp2_fp8_all_gather", False) + + +@register_patch( + "megatron.training.train_step_fp8_cache_update", + backend="megatron", + phase="before_train", + description="Patch train_step to refresh FP8 all-gather cache after optimizer.step()", + condition=_needs_fp8_cache_update, + priority=45, +) +def patch_train_step_fp8_cache(ctx: PatchContext): + import megatron.training.training as megatron_training + + from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + precompute_fp8_scales_for_fsdp, + ) + from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched + + _PATCH_KEY = "megatron.training.train_step_fp8_cache_update" + if is_patched(megatron_training, _PATCH_KEY): + log_rank_0("[Patch:train_step_fp8_cache_update] Already applied; skipping re-wrap.") + return + + args = get_args(ctx) + cache_data = getattr(args, "fp8_precompute_data_cache", True) + use_cpp = getattr(args, "use_cpp_fp8_quantize", False) + stochastic_rounding = getattr(args, "fp8_all_gather_stochastic_rounding", False) + + _original_train_step = megatron_training.train_step + + def _patched_train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=None, + ): + result = _original_train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, + ) + if _refresh_scales: + precompute_fp8_scales_for_fsdp( + model[0], + cache_data=cache_data, + use_cpp_quantize=use_cpp, + stochastic_rounding=stochastic_rounding, + ) + return result + + # When using the C++ quantize kernel without data caching, scales were + # already set once during FSDP setup and the C++ kernel uses them + # directly -- no per-step refresh needed (matches run_35 behavior). + _refresh_scales = cache_data or not use_cpp + + megatron_training.train_step = _patched_train_step + mark_patched(megatron_training, _PATCH_KEY) + log_rank_0( + "[Patch:train_step_fp8_cache_update] " + f"Patched train_step to refresh FP8 all-gather cache after optimizer.step() " + f"(cache_data={cache_data}, use_cpp_quantize={use_cpp}, " + f"stochastic_rounding={stochastic_rounding}, " + f"refresh_scales={_refresh_scales})" + ) diff --git a/primus/backends/megatron/patches/optimizer_patches.py b/primus/backends/megatron/patches/optimizer_patches.py new file mode 100644 index 000000000..5667c7cd9 --- /dev/null +++ b/primus/backends/megatron/patches/optimizer_patches.py @@ -0,0 +1,251 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Megatron Optimizer Patches + +This module contains patches that modify Megatron's optimizer creation to use +Primus-specific implementations when requested. +""" + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.modules.module_utils import log_rank_0 + + +@register_patch( + "megatron.optimizer.fsdp2_fp32_param", + backend="megatron", + phase="before_train", + description="Patch get_megatron_optimizer for FSDP2 FP32 param optimizer (TorchTitan-style)", + priority=50, + condition=lambda ctx: ( + getattr(get_args(ctx), "use_fsdp2_fp32_param_optimizer", False) + and getattr(get_args(ctx), "use_torch_fsdp2", False) + and getattr(get_args(ctx), "bf16", False) + ), +) +def patch_fsdp2_fp32_optimizer(ctx: PatchContext): + """Patch Megatron to use FSDP2 FP32 optimizer. + + This optimizer uses FP32 parameters + FP32 optimizer states with FSDP2's + MixedPrecisionPolicy for BF16 forward/backward. Eliminates the stale + weights problem of BF16 optimizer states. + """ + try: + import megatron.core.optimizer as optimizer_module + from megatron.training import training + + from primus.backends.megatron.core.optimizer.fsdp2_fp32_optimizer import ( + get_fsdp2_fp32_optimizer, + ) + + args = get_args(ctx) + use_foreach = getattr(args, "optimizer_foreach", False) + + _MEGATRON_ONLY_KWARGS = { + "config_overrides", + "use_gloo_process_groups", + "dump_param_to_param_group_map", + } + + def patched_get_megatron_optimizer(config, model_chunks, *args, **kwargs): + log_rank_0("=" * 80) + log_rank_0("[Using FSDP2 FP32 Param Optimizer (TorchTitan-style)]") + log_rank_0(" FP32 params + FP32 optimizer states") + log_rank_0(" FSDP2 MixedPrecisionPolicy handles BF16 compute") + log_rank_0(" DTensor-native gradient clipping") + log_rank_0(f" AdamW mode: {'foreach' if use_foreach else 'fused'}") + log_rank_0("=" * 80) + + filtered_kwargs = {k: v for k, v in kwargs.items() if k not in _MEGATRON_ONLY_KWARGS} + + return get_fsdp2_fp32_optimizer( + config=config, + model_chunks=model_chunks, + use_foreach=use_foreach, + **filtered_kwargs, + ) + + patched_count = 0 + if hasattr(training, "get_megatron_optimizer"): + training.get_megatron_optimizer = patched_get_megatron_optimizer + log_rank_0( + "[Patch:megatron.optimizer.fsdp2_fp32_param] " "Patched training.get_megatron_optimizer" + ) + patched_count += 1 + + if hasattr(optimizer_module, "get_megatron_optimizer"): + optimizer_module.get_megatron_optimizer = patched_get_megatron_optimizer + log_rank_0( + "[Patch:megatron.optimizer.fsdp2_fp32_param] " + "Patched optimizer_module.get_megatron_optimizer" + ) + patched_count += 1 + + if patched_count == 0: + log_rank_0( + "[Patch:megatron.optimizer.fsdp2_fp32_param] " + "WARNING: get_megatron_optimizer not found in either location!" + ) + else: + log_rank_0( + "[Patch:megatron.optimizer.fsdp2_fp32_param] " + f"Patched get_megatron_optimizer in {patched_count} location(s) " + "to use FSDP2FP32Optimizer" + ) + + except Exception as e: + log_rank_0( + f"[Patch:megatron.optimizer.fsdp2_fp32_param] " + f"WARNING: Failed to patch get_megatron_optimizer: {type(e).__name__}: {e}" + ) + import traceback + + log_rank_0(f"Traceback: {traceback.format_exc()}") + + +@register_patch( + "megatron.optimizer.fsdp2_bf16_master_weight", + backend="megatron", + phase="before_train", + description="Patch get_megatron_optimizer for FSDP2 BF16 master weight optimizer", + priority=50, + condition=lambda ctx: ( + getattr(get_args(ctx), "use_fsdp2_bf16_master_weight_optimizer", False) + and getattr(get_args(ctx), "use_torch_fsdp2", False) + and getattr(get_args(ctx), "bf16", False) + ), +) +def patch_fsdp2_bf16_master_weight_optimizer(ctx: PatchContext): + """Patch Megatron to use FSDP2 BF16 master weight optimizer. + + This optimizer keeps model parameters in BF16 (no FP32->BF16 cast in + forward all-gather) and maintains FP32 master copies for optimizer + precision, matching Megatron's Float16OptimizerWithFloat16Params pattern. + """ + try: + import megatron.core.optimizer as optimizer_module + from megatron.training import training + + from primus.backends.megatron.core.optimizer.fsdp2_bf16_master_weight_optimizer import ( + get_fsdp2_bf16_master_weight_optimizer, + ) + + args = get_args(ctx) + use_foreach = getattr(args, "optimizer_foreach", True) + + _MEGATRON_ONLY_KWARGS = { + "config_overrides", + "use_gloo_process_groups", + "dump_param_to_param_group_map", + } + + def patched_get_megatron_optimizer(config, model_chunks, *args, **kwargs): + log_rank_0("=" * 80) + log_rank_0("[Using FSDP2 BF16 Master Weight Optimizer]") + log_rank_0(" BF16 model params + FP32 master weights") + log_rank_0(" No FP32->BF16 cast in forward all-gather") + log_rank_0(" DTensor-native gradient clipping") + log_rank_0(f" Foreach batching: {use_foreach}") + log_rank_0("=" * 80) + + filtered_kwargs = {k: v for k, v in kwargs.items() if k not in _MEGATRON_ONLY_KWARGS} + + return get_fsdp2_bf16_master_weight_optimizer( + config=config, + model_chunks=model_chunks, + use_foreach=use_foreach, + **filtered_kwargs, + ) + + patched_count = 0 + if hasattr(training, "get_megatron_optimizer"): + training.get_megatron_optimizer = patched_get_megatron_optimizer + log_rank_0( + "[Patch:megatron.optimizer.fsdp2_bf16_master_weight] " + "Patched training.get_megatron_optimizer" + ) + patched_count += 1 + + if hasattr(optimizer_module, "get_megatron_optimizer"): + optimizer_module.get_megatron_optimizer = patched_get_megatron_optimizer + log_rank_0( + "[Patch:megatron.optimizer.fsdp2_bf16_master_weight] " + "Patched optimizer_module.get_megatron_optimizer" + ) + patched_count += 1 + + if patched_count == 0: + log_rank_0( + "[Patch:megatron.optimizer.fsdp2_bf16_master_weight] " + "WARNING: get_megatron_optimizer not found in either location!" + ) + else: + log_rank_0( + "[Patch:megatron.optimizer.fsdp2_bf16_master_weight] " + f"Patched get_megatron_optimizer in {patched_count} location(s) " + "to use FSDP2BF16MasterWeightOptimizer" + ) + + except Exception as e: + log_rank_0( + f"[Patch:megatron.optimizer.fsdp2_bf16_master_weight] " + f"WARNING: Failed to patch get_megatron_optimizer: {type(e).__name__}: {e}" + ) + import traceback + + log_rank_0(f"Traceback: {traceback.format_exc()}") + + +@register_patch( + "megatron.optimizer.precision_aware_fp8_tensorwise", + backend="megatron", + phase="before_train", + description=( + "Override precision-aware optimizer flag for local spec + tensorwise FP8, " + "enabling BF16 decoupled_grad path (no BF16-to-FP32 gradient cast)." + ), + priority=35, + condition=lambda ctx: ( + getattr(get_args(ctx), "use_precision_aware_optimizer", False) + and getattr(get_args(ctx), "fp8_recipe", None) == "tensorwise" + and getattr(get_args(ctx), "transformer_impl", None) == "local" + ), +) +def patch_precision_aware_for_tensorwise(ctx: PatchContext): + """Enable precision-aware optimizer's decoupled_grad path for local spec + tensorwise FP8. + + With local spec, parameters are BF16 (not Float8Tensor) -- FP8 quantization + happens inside per-module autograd Functions, transparent to the optimizer. + The upstream condition incorrectly excludes tensorwise FP8 because it was + designed for TE spec where Float8Tensor storage requires special handling. + """ + try: + from megatron.core.optimizer.optimizer_config import OptimizerConfig + + original_post_init = OptimizerConfig.__post_init__ + + def patched_post_init(self): + original_post_init(self) + if self.use_precision_aware_optimizer and self.fp8_recipe == "tensorwise": + self.use_precision_aware_optimizer_no_fp8_or_ds_fp8 = True + + OptimizerConfig.__post_init__ = patched_post_init + + log_rank_0( + "[Patch:megatron.optimizer.precision_aware_fp8_tensorwise] " + "Patched OptimizerConfig.__post_init__ to enable " + "decoupled_grad path for local spec + tensorwise FP8" + ) + + except Exception as e: + log_rank_0( + f"[Patch:megatron.optimizer.precision_aware_fp8_tensorwise] " + f"WARNING: Failed to patch OptimizerConfig: {type(e).__name__}: {e}" + ) + import traceback + + log_rank_0(f"Traceback: {traceback.format_exc()}") diff --git a/primus/backends/megatron/patches/torch_fsdp2_patches.py b/primus/backends/megatron/patches/torch_fsdp2_patches.py index 1ac3c56cf..7eaf39e98 100644 --- a/primus/backends/megatron/patches/torch_fsdp2_patches.py +++ b/primus/backends/megatron/patches/torch_fsdp2_patches.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -59,27 +59,40 @@ def patch_torch_fsdp(ctx: PatchContext): f"[Patch:megatron.fsdp.torch_fsdp2] Patched megatron.training.training.torch_FSDP " f"-> {PrimusTorchFullyShardedDataParallel.__name__}" ) - # Megatron Core 0.16 may pass new kwargs (e.g., force_all_reduce) into - # model_chunk.finish_grad_sync(). Keep FSDP2 path forward-compatible by - from megatron.core.distributed import data_parallel_base - original_start_grad_sync = data_parallel_base._BaseDataParallel.start_grad_sync - original_finish_grad_sync = data_parallel_base._BaseDataParallel.finish_grad_sync - - if not getattr(original_start_grad_sync, "_primus_grad_sync_compat", False): - - def _patched_start_grad_sync(self, *unused, **unused_kwargs): - return original_start_grad_sync(self, *unused) - - def _patched_finish_grad_sync(self, *unused, **unused_kwargs): - return original_finish_grad_sync(self) - - setattr(_patched_start_grad_sync, "_primus_grad_sync_compat", True) - setattr(_patched_finish_grad_sync, "_primus_grad_sync_compat", True) - data_parallel_base._BaseDataParallel.start_grad_sync = _patched_start_grad_sync - data_parallel_base._BaseDataParallel.finish_grad_sync = _patched_finish_grad_sync - - log_rank_0( - "[Patch:megatron.fsdp.torch_fsdp2] Patched _BaseDataParallel " - "start_grad_sync/finish_grad_sync to accept extra kwargs " - ) + # Patch get_data_parallel_group_if_dtensor to handle 2D HSDP meshes. + # The upstream implementation calls tensor.device_mesh.get_group() without mesh_dim, + # which fails for 2D meshes. For HSDP (dp_replicate, dp_shard), we return the + # shard group (innermost dim) since that's where parameters are actually sharded; + # replicas have identical gradients so we must not all-reduce across them. + import megatron.core.optimizer.clip_grads as clip_grads_module + import megatron.core.utils as mcore_utils + from megatron.training import utils as training_utils + + try: + from torch.distributed._tensor import DTensor + + HAVE_DTENSOR = True + except ImportError: + HAVE_DTENSOR = False + + def _get_data_parallel_group_if_dtensor(tensor, data_parallel_group=None): + if HAVE_DTENSOR and isinstance(tensor, DTensor): + mesh = tensor.device_mesh + if mesh.ndim > 1: + current_group = mesh.get_group(mesh_dim=-1) + else: + current_group = mesh.get_group() + if data_parallel_group is not None and current_group != data_parallel_group: + raise RuntimeError("DTensor mesh group does not match the expected data_parallel_group") + return current_group + return None + + mcore_utils.get_data_parallel_group_if_dtensor = _get_data_parallel_group_if_dtensor + clip_grads_module.get_data_parallel_group_if_dtensor = _get_data_parallel_group_if_dtensor + if hasattr(training_utils, "get_data_parallel_group_if_dtensor"): + training_utils.get_data_parallel_group_if_dtensor = _get_data_parallel_group_if_dtensor + log_rank_0( + "[Patch:megatron.fsdp.torch_fsdp2] Patched get_data_parallel_group_if_dtensor " + "for 2D HSDP mesh support" + ) diff --git a/tests/unit_tests/backends/megatron/diffusion/distributed/__init__.py b/tests/unit_tests/backends/megatron/diffusion/distributed/__init__.py new file mode 100644 index 000000000..b22d06ccf --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/distributed/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Distributed tests for Flux diffusion model. + +Tests for FSDP2, pipeline parallelism, and other distributed training features. +""" diff --git a/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_fp8_all_gather.py b/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_fp8_all_gather.py new file mode 100644 index 000000000..c313644ea --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_fp8_all_gather.py @@ -0,0 +1,528 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for FP8 all-gather tensor subclasses. + +Tests the FP8 all-gather subclasses in isolation (single GPU, no +multi-GPU / NCCL required). Verifies: +- WeightWithFP8AllGatherTensor: creation, pre/post all-gather with + precomputed scale, dispatch, flatten/unflatten +- FP8UnshardedWeightTensor: creation, dispatch propagation, detach + dequantizes to BF16, unsupported ops raise NotImplementedError, + flatten/unflatten roundtrip, get_fp8_data_and_scale_inv() +""" + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus_turbo.pytorch.core.low_precision import ( + Float8QuantConfig, + Format, + ScalingGranularity, +) + +from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + FP8UnshardedWeightTensor, + WeightWithFP8AllGatherTensor, + _foreach_fp8_quantize, + _get_fp8_dtype, + quantize_fp8_prescaled, +) + +_has_cuda = torch.cuda.is_available() +requires_cuda = pytest.mark.skipif(not _has_cuda, reason="CUDA required") + + +def _make_config(): + return Float8QuantConfig(format=Format.E4M3, granularity=ScalingGranularity.TENSORWISE) + + +def _make_wrapped(shape=(256, 512), device="cpu"): + tensor = torch.randn(shape, dtype=torch.bfloat16, device=device) + wrapped = WeightWithFP8AllGatherTensor(tensor, _make_config()) + return wrapped, tensor + + +def _set_precomputed_scale(wrapped): + """Compute and set precomputed scale, scale_inv, and cached FP8 data for a single-rank scenario.""" + fp8_dtype = _get_fp8_dtype(wrapped._fp8_config.format) + fp8_max = torch.finfo(fp8_dtype).max + amax = wrapped._tensor.abs().amax().float() + scale = fp8_max / amax.clamp(min=1e-12) + wrapped._precomputed_scale = scale + wrapped._precomputed_scale_inv = 1.0 / scale + fp8_data, _ = quantize_fp8_prescaled( + wrapped._tensor, + fp8_dtype, + scale, + wrapped._precomputed_scale_inv, + ) + wrapped._cached_fp8_data = fp8_data + + +def _make_fp8_unsharded(shape=(256, 512), device="cpu"): + """Create an FP8UnshardedWeightTensor from a BF16 source tensor.""" + config = _make_config() + bf16_tensor = torch.randn(shape, dtype=torch.bfloat16, device=device) + fp8_dtype = _get_fp8_dtype(config.format) + fp8_max = torch.finfo(fp8_dtype).max + amax = bf16_tensor.abs().amax().float() + scale = fp8_max / amax.clamp(min=1e-12) + scale_inv = 1.0 / scale + fp8_data = (bf16_tensor.float() * scale).clamp(-fp8_max, fp8_max).to(fp8_dtype) + return ( + FP8UnshardedWeightTensor(fp8_data, scale_inv, torch.bfloat16, config), + bf16_tensor, + fp8_data, + scale_inv, + ) + + +class TestSubclassCreation: + def test_shape_dtype_device_preserved(self): + wrapped, orig = _make_wrapped() + assert wrapped.shape == orig.shape + assert wrapped.dtype == torch.bfloat16 + assert wrapped.device == orig.device + + def test_isinstance_tensor(self): + wrapped, _ = _make_wrapped() + assert isinstance(wrapped, torch.Tensor) + assert isinstance(wrapped, WeightWithFP8AllGatherTensor) + + def test_inner_tensor_accessible(self): + wrapped, orig = _make_wrapped() + assert wrapped._tensor is orig + + def test_tensorwise_only_validation(self): + config = Float8QuantConfig(format=Format.E4M3, granularity=ScalingGranularity.ROWWISE) + with pytest.raises(ValueError, match="TENSORWISE"): + WeightWithFP8AllGatherTensor(torch.randn(4, 4, dtype=torch.bfloat16), config) + + def test_precomputed_scale_initially_none(self): + wrapped, _ = _make_wrapped() + assert wrapped._precomputed_scale is None + + +@requires_cuda +class TestPreAllGather: + def test_returns_one_input(self): + wrapped, _ = _make_wrapped(device="cuda") + _set_precomputed_scale(wrapped) + all_gather_inputs, metadata = wrapped.fsdp_pre_all_gather(None) + assert len(all_gather_inputs) == 1 + + def test_fp8_data_dtype(self): + wrapped, _ = _make_wrapped(device="cuda") + _set_precomputed_scale(wrapped) + (fp8_data,), _ = wrapped.fsdp_pre_all_gather(None) + assert fp8_data.dtype in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ) + + def test_metadata_contains_scale_inv_and_numel(self): + wrapped, orig = _make_wrapped((128, 64), device="cuda") + _set_precomputed_scale(wrapped) + _, metadata = wrapped.fsdp_pre_all_gather(None) + scale_inv, shard_numel = metadata + assert scale_inv.dtype == torch.float32 + assert shard_numel == 128 * 64 + + def test_fp8_data_shape_matches_input(self): + wrapped, orig = _make_wrapped((32, 64), device="cuda") + _set_precomputed_scale(wrapped) + (fp8_data,), _ = wrapped.fsdp_pre_all_gather(None) + assert fp8_data.numel() == orig.numel() + + def test_asserts_without_precomputed_scale(self): + wrapped, _ = _make_wrapped(device="cuda") + with pytest.raises(RuntimeError, match="precompute_fp8_scales_for_fsdp"): + wrapped.fsdp_pre_all_gather(None) + + +@requires_cuda +class TestPostAllGather: + def _simulate_single_rank(self, wrapped): + """Simulate all-gather with a single rank (data passes through).""" + _set_precomputed_scale(wrapped) + (fp8_data,), metadata = wrapped.fsdp_pre_all_gather(None) + return (fp8_data,), metadata + + def test_first_call_returns_fp8_unsharded_tensor(self): + wrapped, _ = _make_wrapped(device="cuda") + outputs, metadata = self._simulate_single_rank(wrapped) + result = wrapped.fsdp_post_all_gather(outputs, metadata, torch.bfloat16) + assert result is not None + tensor, inner_tensors = result + assert isinstance(tensor, FP8UnshardedWeightTensor) + assert tensor.dtype == torch.bfloat16 + assert isinstance(inner_tensors, tuple) + assert len(inner_tensors) == 1 + + def test_steady_state_updates_scale(self): + wrapped, _ = _make_wrapped(device="cuda") + outputs, metadata = self._simulate_single_rank(wrapped) + # First call to get an FP8UnshardedWeightTensor + result = wrapped.fsdp_post_all_gather(outputs, metadata, torch.bfloat16) + fp8_unsharded, _ = result + # Simulate FSDP's reshard path: pass the FP8UnshardedWeightTensor + # directly as `out` (FSDP uses the unsharded_param, not nn.Parameter) + new_scale_inv = metadata[0] * 2.0 + new_metadata = (new_scale_inv, metadata[1]) + result = wrapped.fsdp_post_all_gather(outputs, new_metadata, torch.bfloat16, out=fp8_unsharded) + assert result is None + assert torch.equal(fp8_unsharded._scale_inv, new_scale_inv) + + +@requires_cuda +class TestRoundtripAccuracy: + def test_quantize_keeps_values(self): + torch.manual_seed(42) + wrapped, orig = _make_wrapped((128, 256), device="cuda") + _set_precomputed_scale(wrapped) + outputs, metadata = wrapped.fsdp_pre_all_gather(None) + result = wrapped.fsdp_post_all_gather(outputs, metadata, torch.bfloat16) + fp8_unsharded, _ = result + fp8_data, scale_inv = fp8_unsharded.get_fp8_data_and_scale_inv() + dequantized = fp8_data.to(torch.bfloat16) * scale_inv + dequantized = dequantized.view(orig.shape) + torch.testing.assert_close(dequantized, orig, rtol=0.05, atol=0.1) + + +class TestTorchDispatch: + def test_preserves_subclass_for_standard_ops(self): + wrapped, _ = _make_wrapped((16, 16)) + for op in [ + lambda t: torch.empty_like(t), + lambda t: t.view(-1), + lambda t: t.clone(), + lambda t: t[0:8], + ]: + result = op(wrapped) + assert isinstance( + result, WeightWithFP8AllGatherTensor + ), f"Op should preserve subclass but got {type(result)}" + + def test_unwraps_for_compute_ops(self): + wrapped, _ = _make_wrapped((16, 16)) + plain = torch.randn(16, 16, dtype=torch.bfloat16) + result = wrapped + plain + assert type(result) is torch.Tensor + assert not isinstance(result, WeightWithFP8AllGatherTensor) + + def test_copy_inplace_modifies_inner(self): + """copy_ modifies inner tensor in-place; target retains subclass type.""" + wrapped, _ = _make_wrapped((8, 8)) + target = torch.empty_like(wrapped) + target.copy_(wrapped) + assert isinstance(target, WeightWithFP8AllGatherTensor) + torch.testing.assert_close(target._tensor, wrapped._tensor) + + def test_to_copy_dtype_change_unwraps(self): + """`.float()` should return a plain tensor (for FP32 master copies).""" + wrapped, _ = _make_wrapped((8, 8)) + fp32 = wrapped.float() + assert fp32.dtype == torch.float32 + assert not isinstance(fp32, WeightWithFP8AllGatherTensor) + + def test_to_copy_same_dtype_preserves(self): + wrapped, _ = _make_wrapped((8, 8)) + result = wrapped.to(torch.bfloat16) + assert isinstance(result, WeightWithFP8AllGatherTensor) + + def test_detach_preserves_subclass(self): + """detach preserves subclass so nn.Parameter() and FSDP2 work correctly.""" + wrapped, orig = _make_wrapped((8, 8)) + detached = wrapped.detach() + assert isinstance(detached, WeightWithFP8AllGatherTensor) + assert detached._tensor.data_ptr() == wrapped._tensor.data_ptr() + assert detached._fp8_config.format == wrapped._fp8_config.format + + +class TestTensorFlattenUnflatten: + def test_roundtrip(self): + wrapped, orig = _make_wrapped((32, 64)) + inner_tensors_names, metadata = wrapped.__tensor_flatten__() + assert inner_tensors_names == ["_tensor"] + assert "format" in metadata + assert "granularity" in metadata + + inner_tensors = {"_tensor": wrapped._tensor} + restored = WeightWithFP8AllGatherTensor.__tensor_unflatten__( + inner_tensors, metadata, wrapped.shape, wrapped.stride() + ) + assert isinstance(restored, WeightWithFP8AllGatherTensor) + assert restored._tensor is orig + assert restored._fp8_config.format == Format.E4M3 + assert restored._fp8_config.granularity == ScalingGranularity.TENSORWISE + + def test_metadata_is_hashable(self): + wrapped, _ = _make_wrapped() + _, metadata = wrapped.__tensor_flatten__() + for key, value in metadata.items(): + hash(value) + + def test_precomputed_scale_not_in_metadata(self): + """_precomputed_scale is transient and must not appear in flatten metadata.""" + wrapped, _ = _make_wrapped() + wrapped._precomputed_scale = torch.tensor(1.0) + _, metadata = wrapped.__tensor_flatten__() + assert "precomputed_scale" not in metadata + assert "_precomputed_scale" not in metadata + + +# ============================================================================ +# FP8UnshardedWeightTensor tests +# ============================================================================ + + +class TestFP8UnshardedWeightTensorCreation: + def test_shape_and_declared_dtype(self): + fp8_unsharded, _, fp8_data, _ = _make_fp8_unsharded((64, 32)) + assert fp8_unsharded.shape == (64, 32) + assert fp8_unsharded.dtype == torch.bfloat16 + + def test_inner_fp8_data_accessible(self): + fp8_unsharded, _, fp8_data, scale_inv = _make_fp8_unsharded() + assert fp8_unsharded._fp8_data is fp8_data + assert torch.equal(fp8_unsharded._scale_inv, scale_inv) + + def test_get_fp8_data_and_scale_inv(self): + fp8_unsharded, _, fp8_data, scale_inv = _make_fp8_unsharded() + data, sinv = fp8_unsharded.get_fp8_data_and_scale_inv() + assert data is fp8_data + assert torch.equal(sinv, scale_inv) + + +class TestFP8UnshardedDispatch: + def test_as_strided_preserves_subclass(self): + fp8_unsharded, _, _, _ = _make_fp8_unsharded((8, 8)) + result = torch.as_strided(fp8_unsharded, (64,), (1,), 0) + assert isinstance(result, FP8UnshardedWeightTensor) + assert result.shape == (64,) + + def test_view_preserves_subclass(self): + fp8_unsharded, _, _, _ = _make_fp8_unsharded((8, 8)) + result = fp8_unsharded.view(-1) + assert isinstance(result, FP8UnshardedWeightTensor) + assert result.shape == (64,) + + def test_slice_preserves_subclass(self): + fp8_unsharded, _, _, _ = _make_fp8_unsharded((16, 8)) + result = fp8_unsharded[0:8] + assert isinstance(result, FP8UnshardedWeightTensor) + assert result.shape == (8, 8) + + def test_clone_preserves_subclass(self): + fp8_unsharded, _, _, _ = _make_fp8_unsharded((8, 8)) + result = fp8_unsharded.clone() + assert isinstance(result, FP8UnshardedWeightTensor) + + def test_detach_preserves_subclass(self): + """detach must preserve subclass (required by nn.Parameter wrapping).""" + fp8_unsharded, _, _, _ = _make_fp8_unsharded((8, 8)) + detached = fp8_unsharded.detach() + assert isinstance(detached, FP8UnshardedWeightTensor) + assert detached.dtype == torch.bfloat16 + + def test_unsupported_op_raises(self): + """Unsupported ops raise NotImplementedError (safety guard).""" + fp8_unsharded, _, _, _ = _make_fp8_unsharded((8, 8)) + with pytest.raises(NotImplementedError, match="does not support"): + fp8_unsharded + torch.zeros(8, 8, dtype=torch.bfloat16) + + +class TestFP8UnshardedFlattenUnflatten: + def test_roundtrip(self): + fp8_unsharded, _, _, _ = _make_fp8_unsharded((32, 64)) + inner_names, metadata = fp8_unsharded.__tensor_flatten__() + assert set(inner_names) == {"_fp8_data", "_scale_inv"} + assert "orig_dtype" in metadata + assert "format" in metadata + + inner_tensors = { + "_fp8_data": fp8_unsharded._fp8_data, + "_scale_inv": fp8_unsharded._scale_inv, + } + restored = FP8UnshardedWeightTensor.__tensor_unflatten__( + inner_tensors, metadata, fp8_unsharded.shape, fp8_unsharded.stride() + ) + assert isinstance(restored, FP8UnshardedWeightTensor) + assert restored._orig_dtype == torch.bfloat16 + # FP8 dtypes don't support torch.equal on CPU; compare via uint8 view + assert torch.equal( + restored._fp8_data.view(torch.uint8), + fp8_unsharded._fp8_data.view(torch.uint8), + ) + assert torch.equal(restored._scale_inv, fp8_unsharded._scale_inv) + + def test_metadata_is_hashable(self): + fp8_unsharded, _, _, _ = _make_fp8_unsharded() + _, metadata = fp8_unsharded.__tensor_flatten__() + for key, value in metadata.items(): + hash(value) + + +class TestFP8DataCache: + def test_cache_initially_none(self): + wrapped, _ = _make_wrapped() + assert wrapped._cached_fp8_data is None + + @requires_cuda + def test_pre_all_gather_returns_cached_data(self): + wrapped, _ = _make_wrapped(device="cuda") + _set_precomputed_scale(wrapped) + assert wrapped._cached_fp8_data is not None + (fp8_data,), _ = wrapped.fsdp_pre_all_gather(mesh=None) + assert fp8_data.data_ptr() == wrapped._cached_fp8_data.data_ptr() + + @requires_cuda + def test_pre_all_gather_fallback_without_cache(self): + wrapped, _ = _make_wrapped(device="cuda") + _set_precomputed_scale(wrapped) + wrapped._cached_fp8_data = None + (fp8_data,), (scale_inv, _) = wrapped.fsdp_pre_all_gather(mesh=None) + fp8_dtype = _get_fp8_dtype(wrapped._fp8_config.format) + assert fp8_data.dtype == fp8_dtype + assert fp8_data.shape == wrapped._tensor.shape + + @requires_cuda + def test_cache_matches_direct_quantize(self): + wrapped, _ = _make_wrapped(device="cuda") + _set_precomputed_scale(wrapped) + cached = wrapped._cached_fp8_data + fp8_dtype = _get_fp8_dtype(wrapped._fp8_config.format) + fresh, _ = quantize_fp8_prescaled( + wrapped._tensor, + fp8_dtype, + wrapped._precomputed_scale, + wrapped._precomputed_scale_inv, + ) + assert torch.equal(cached.view(torch.uint8), fresh.view(torch.uint8)) + + @requires_cuda + def test_cache_not_in_flatten_metadata(self): + # _set_precomputed_scale runs the Triton FP8 quantize kernel which + # requires a CUDA-resident tensor (no CPU dispatch). Match the other + # cache tests in this class which already use device="cuda". + wrapped, _ = _make_wrapped(device="cuda") + _set_precomputed_scale(wrapped) + inner_names, _ = wrapped.__tensor_flatten__() + assert "_cached_fp8_data" not in inner_names + + +class TestForeachOptimizerOps: + """Tests for _foreach_copy_ batched operations used by the optimizer.""" + + @requires_cuda + def test_foreach_copy_bf16_to_fp32_matches_float(self): + """_foreach_copy_ BF16->FP32 matches per-tensor .float() (prepare_grads).""" + sizes = [(128, 256), (64,), (512, 512), (32, 16)] + bf16_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in sizes] + ref_fp32 = [t.float() for t in bf16_tensors] + + fp32_targets = [torch.empty(s, dtype=torch.float32, device="cuda") for s in sizes] + torch._foreach_copy_(fp32_targets, bf16_tensors) + + for i in range(len(sizes)): + torch.testing.assert_close( + fp32_targets[i], ref_fp32[i], rtol=0, atol=0, msg=f"Mismatch at tensor {i} shape {sizes[i]}" + ) + + @requires_cuda + def test_foreach_copy_fp32_to_bf16_matches_copy(self): + """_foreach_copy_ FP32->BF16 matches per-tensor .copy_() (copy-back).""" + sizes = [(256, 128), (64,), (1024,)] + fp32_tensors = [torch.randn(s, dtype=torch.float32, device="cuda") for s in sizes] + ref_bf16 = [torch.empty(s, dtype=torch.bfloat16, device="cuda") for s in sizes] + for r, f in zip(ref_bf16, fp32_tensors): + r.copy_(f) + + bf16_targets = [torch.empty(s, dtype=torch.bfloat16, device="cuda") for s in sizes] + torch._foreach_copy_(bf16_targets, fp32_tensors) + + for i in range(len(sizes)): + torch.testing.assert_close( + bf16_targets[i], ref_bf16[i], rtol=0, atol=0, msg=f"Mismatch at tensor {i} shape {sizes[i]}" + ) + + @requires_cuda + def test_cache_extraction_fp8_wrapped(self): + """to_local() + inner_data() produces the correct inner tensor for FP8-wrapped params.""" + wrapped, orig = _make_wrapped((128, 64), device="cuda") + inner = wrapped.inner_data() + assert inner is orig + assert inner.dtype == torch.bfloat16 + assert inner.data_ptr() == orig.data_ptr() + + @requires_cuda + def test_mixed_sizes_foreach_copy(self): + """_foreach_copy_ works with mixed tensor sizes in the same list.""" + sizes = [(1,), (7,), (1024, 1024), (3, 5, 7), (131072,)] + bf16_list = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in sizes] + fp32_list = [torch.empty(s, dtype=torch.float32, device="cuda") for s in sizes] + torch._foreach_copy_(fp32_list, bf16_list) + for i, (b, f) in enumerate(zip(bf16_list, fp32_list)): + torch.testing.assert_close( + f, b.float(), rtol=0, atol=0, msg=f"Mixed-size mismatch at tensor {i} shape {sizes[i]}" + ) + + +class TestForeachFP8Quantize: + """Tests for the compiled _foreach_fp8_quantize batch quantization helper.""" + + @requires_cuda + def test_matches_per_tensor_quantize(self): + """Verify _foreach_fp8_quantize is bitwise identical to quantize_fp8_prescaled.""" + config = _make_config() + fp8_dtype = _get_fp8_dtype(config.format) + fp8_max = torch.finfo(fp8_dtype).max + + sizes = [1024, 4096, 131072, 524288] + tensors = [torch.randn(s, device="cuda", dtype=torch.bfloat16) for s in sizes] + amaxes = torch.stack([t.abs().amax().float() for t in tensors]) + scales = (fp8_max / amaxes.clamp(min=1e-12)).to(torch.float32) + scale_invs = 1.0 / scales + + ref_results = [] + for i, t in enumerate(tensors): + fp8_data, _ = quantize_fp8_prescaled( + t, + fp8_dtype, + scales[i], + scale_invs[i], + ) + ref_results.append(fp8_data) + + fp8_outputs = [torch.empty_like(t, dtype=fp8_dtype) for t in tensors] + _foreach_fp8_quantize.__wrapped__(tensors, fp8_outputs, scales, fp8_max) + + for i in range(len(tensors)): + assert torch.equal( + fp8_outputs[i].view(torch.uint8), + ref_results[i].view(torch.uint8), + ), f"Mismatch at tensor {i} (size {sizes[i]})" + + @requires_cuda + def test_clamps_correctly(self): + """Verify out-of-range values are clamped to fp8_max, not NaN.""" + config = _make_config() + fp8_dtype = _get_fp8_dtype(config.format) + fp8_max = torch.finfo(fp8_dtype).max + + tensor = torch.tensor([1.0, 1000.0, -1000.0, 0.001], device="cuda", dtype=torch.bfloat16) + scale = torch.tensor(fp8_max, device="cuda", dtype=torch.float32) + scales = scale.unsqueeze(0) + + fp8_out = [torch.empty_like(tensor, dtype=fp8_dtype)] + _foreach_fp8_quantize.__wrapped__([tensor], fp8_out, scales, fp8_max) + + result = fp8_out[0] + assert not result.float().isnan().any(), "NaN found in FP8 output" + assert result.float().abs().max() <= fp8_max diff --git a/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_transformer_impl.py b/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_transformer_impl.py new file mode 100644 index 000000000..64849c8f0 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/distributed/test_fsdp2_transformer_impl.py @@ -0,0 +1,109 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for FSDP2 transformer_impl wrapping logic. + +Tests that PrimusTorchFullyShardedDataParallel correctly excludes/includes +ColumnParallelLinear based on transformer_impl setting, aligning with Megatron's pattern. +""" + +from unittest.mock import patch + +import pytest +import torch +from megatron.core import tensor_parallel +from megatron.core.distributed.distributed_data_parallel_config import ( + DistributedDataParallelConfig, +) +from megatron.core.transformer.transformer_config import TransformerConfig + +from primus.backends.megatron.core.distributed import ( + torch_fully_sharded_data_parallel as fsdp_mod, +) +from primus.backends.megatron.core.distributed.torch_fully_sharded_data_parallel import ( + PrimusTorchFullyShardedDataParallel, +) +from tests.utils import PrimusUT + + +class TestFSDP2TransformerImpl(PrimusUT): + """Tests for FSDP2 wrapping logic with transformer_impl.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state for FSDP2 tests.""" + + def _build_and_collect_wrapped(self, transformer_impl): + """Construct the FSDP2 wrapper with fully_shard mocked at the call site. + + Returns the list of module types that fully_shard was invoked on. The + mock must patch the name bound INSIDE the production module (it does + ``from torch.distributed.fsdp import fully_shard`` at import time), not + ``torch.distributed.fsdp.fully_shard`` -- otherwise the real kernel runs + and the wrapping decision is never exercised. + """ + + class MockModel(torch.nn.Module): + def __init__(self): + super().__init__() + from megatron.core.tensor_parallel.layers import ColumnParallelLinear + + inner = TransformerConfig( + hidden_size=128, + num_attention_heads=2, + num_layers=1, + transformer_impl=transformer_impl, + ) + self.linear = ColumnParallelLinear( + inner.hidden_size, + inner.hidden_size, + config=inner, + init_method=lambda x: None, + ) + + model = MockModel().cuda() + config = TransformerConfig( + hidden_size=128, + num_attention_heads=2, + num_layers=1, + transformer_impl=transformer_impl, + ) + ddp_config = DistributedDataParallelConfig() + + wrapped_modules = [] + + def mock_fully_shard(module, **kwargs): + wrapped_modules.append(type(module)) + + with patch.object(fsdp_mod, "fully_shard", side_effect=mock_fully_shard): + PrimusTorchFullyShardedDataParallel( + config=config, + ddp_config=ddp_config, + module=model, + ) + return wrapped_modules + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_fsdp2_excludes_column_parallel_linear_when_local(self): + """ColumnParallelLinear must be excluded when transformer_impl == 'local'.""" + if not fsdp_mod.HAVE_FSDP: + pytest.skip("torch.distributed.fsdp (FSDP2) is unavailable") + + wrapped_modules = self._build_and_collect_wrapped("local") + + assert ( + tensor_parallel.ColumnParallelLinear not in wrapped_modules + ), "ColumnParallelLinear should be excluded when transformer_impl == 'local'" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_fsdp2_includes_column_parallel_linear_when_not_local(self): + """Positive control: ColumnParallelLinear IS wrapped for non-local impl.""" + if not fsdp_mod.HAVE_FSDP: + pytest.skip("torch.distributed.fsdp (FSDP2) is unavailable") + + wrapped_modules = self._build_and_collect_wrapped("transformer_engine") + + assert ( + tensor_parallel.ColumnParallelLinear in wrapped_modules + ), "ColumnParallelLinear should be wrapped when transformer_impl != 'local'" diff --git a/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_bf16_master_weight_optimizer.py b/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_bf16_master_weight_optimizer.py new file mode 100644 index 000000000..35f7ac02c --- /dev/null +++ b/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_bf16_master_weight_optimizer.py @@ -0,0 +1,351 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for FSDP2BF16MasterWeightOptimizer. + +Tests focus on the core design properties: +- FP32 master weights are created from BF16 model parameters +- prepare_grads copies BF16 grads to FP32 master grads +- Optimizer step updates FP32 masters and copies back to BF16 params +- state_dict includes fp32_from_fp16_params key +- Checkpoint save/load round-trip preserves values +- sharded_state_dict integrates with Megatron's torch_dist checkpointing +- finalize_dist_ckpt_load fills step counter and syncs params +- Factory param group splitting (weight decay, LR scaling) works +- Multi-step training actually reduces loss +""" + +import pytest +import torch +from megatron.core.dist_checkpointing.mapping import ShardedTensor +from megatron.core.optimizer.optimizer_config import OptimizerConfig +from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint + +from primus.backends.megatron.core.optimizer.fsdp2_bf16_master_weight_optimizer import ( + get_fsdp2_bf16_master_weight_optimizer, +) +from tests.utils import PrimusUT + + +def _make_optimizer_config(**overrides): + defaults = dict( + bf16=True, + fp16=False, + optimizer="adam", + lr=1e-4, + min_lr=1e-5, + weight_decay=0.01, + adam_beta1=0.9, + adam_beta2=0.999, + adam_eps=1e-8, + clip_grad=1.0, + use_precision_aware_optimizer=False, + ) + defaults.update(overrides) + return OptimizerConfig(**defaults) + + +def _make_model_and_optimizer(device="cpu"): + model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.bfloat16, device=device) + config = _make_optimizer_config() + optimizer = get_fsdp2_bf16_master_weight_optimizer( + config=config, + model_chunks=[model], + ) + return model, optimizer, config + + +def _build_model_sharded_state_dict(model, prefix=""): + model_sd = model.state_dict(prefix=prefix, keep_vars=True) + return make_sharded_tensors_for_checkpoint(model_sd, prefix, {}, ()) + + +class TestFSDP2BF16MasterWeightOptimizer(PrimusUT): + """Tests for FSDP2BF16MasterWeightOptimizer core functionality.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_master_weights_are_fp32_clones_of_bf16(self): + """FP32 master weights must match BF16 params in value and shape.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + for bf16_group, fp32_group in zip(optimizer.bf16_groups, optimizer.fp32_from_bf16_groups): + for bf16_param, fp32_master in zip(bf16_group, fp32_group): + assert bf16_param.dtype == torch.bfloat16 + assert fp32_master.dtype == torch.float32 + assert fp32_master.shape == bf16_param.shape + torch.testing.assert_close(fp32_master, bf16_param.float(), atol=0, rtol=0) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_main_param_attribute_set(self): + """BF16 params must have .main_param pointing to FP32 master.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + for bf16_group, fp32_group in zip(optimizer.bf16_groups, optimizer.fp32_from_bf16_groups): + for bf16_param, fp32_master in zip(bf16_group, fp32_group): + assert hasattr(bf16_param, "main_param") + assert bf16_param.main_param is fp32_master + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_optimizer_operates_on_fp32_masters(self): + """Base optimizer's param_groups must contain FP32 masters, not BF16 params.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + for group in optimizer.optimizer.param_groups: + for p in group["params"]: + assert p.dtype == torch.float32, f"Optimizer param dtype is {p.dtype}, expected float32" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_prepare_grads_copies_to_fp32(self): + """prepare_grads must cast BF16 grads to FP32 on masters and clear BF16 grads.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + x = torch.randn(2, 8, dtype=torch.bfloat16, device="cuda") + loss = model(x).sum() + loss.backward() + + for bf16_group in optimizer.bf16_groups: + for p in bf16_group: + assert p.grad is not None, "BF16 param should have grad before prepare_grads" + + optimizer.prepare_grads() + + for bf16_group in optimizer.bf16_groups: + for p in bf16_group: + assert p.grad is None, "BF16 grad should be cleared after prepare_grads" + + for fp32_group in optimizer.fp32_from_bf16_groups: + for p in fp32_group: + assert p.grad is not None, "FP32 master should have grad after prepare_grads" + assert p.grad.dtype == torch.float32 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_step_updates_both_master_and_model(self): + """step() must update FP32 masters and copy back to BF16 params. + + Uses high LR to ensure BF16 copy-back produces visible changes + (small updates can round to zero in BF16). + """ + model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.bfloat16, device="cuda") + config = _make_optimizer_config(lr=1e-1, clip_grad=1.0) + optimizer = get_fsdp2_bf16_master_weight_optimizer(config=config, model_chunks=[model]) + + bf16_before = {} + for name, p in model.named_parameters(): + bf16_before[name] = p.data.clone() + + x = torch.randn(2, 8, dtype=torch.bfloat16, device="cuda") + loss = (model(x) * 10.0).sum() + loss.backward() + success, grad_norm, num_zeros = optimizer.step() + + assert success is True + + for name, p in model.named_parameters(): + assert not torch.equal( + p.data, bf16_before[name] + ), f"BF16 param '{name}' was not updated after step" + + for bf16_group, fp32_group in zip(optimizer.bf16_groups, optimizer.fp32_from_bf16_groups): + for bf16_param, fp32_master in zip(bf16_group, fp32_group): + expected = fp32_master.data.to(torch.bfloat16) + torch.testing.assert_close(bf16_param.data, expected, atol=0, rtol=0) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_optimizer_states_are_fp32(self): + """All optimizer states (exp_avg, exp_avg_sq) must be FP32.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + x = torch.randn(2, 8, dtype=torch.bfloat16, device="cuda") + loss = model(x).sum() + loss.backward() + optimizer.step() + + for fp32_group in optimizer.fp32_from_bf16_groups: + for p in fp32_group: + state = optimizer.optimizer.state[p] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_state_dict_has_fp32_from_fp16_params_key(self): + """state_dict must include 'fp32_from_fp16_params' with FP32 master weights.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + x = torch.randn(2, 8, dtype=torch.bfloat16, device="cuda") + loss = model(x).sum() + loss.backward() + optimizer.step() + + saved = optimizer.state_dict() + assert "optimizer" in saved + assert "fp32_from_fp16_params" in saved + + for group in saved["fp32_from_fp16_params"]: + for param in group: + assert param.dtype == torch.float32 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_multi_step_loss_decreases(self): + """Run 10 steps and verify loss decreases.""" + torch.manual_seed(42) + model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.bfloat16, device="cuda") + config = _make_optimizer_config(lr=1e-2, clip_grad=0.0) + optimizer = get_fsdp2_bf16_master_weight_optimizer(config=config, model_chunks=[model]) + + x = torch.randn(4, 8, dtype=torch.bfloat16, device="cuda") + target = torch.randn(4, 8, dtype=torch.bfloat16, device="cuda") + + loss_initial = None + loss_final = None + for _ in range(10): + optimizer.zero_grad() + output = model(x) + loss = ((output - target).float() ** 2).mean() + loss.backward() + optimizer.step() + + if loss_initial is None: + loss_initial = loss.item() + loss_final = loss.item() + + assert ( + loss_final < loss_initial + ), f"Loss did not decrease: initial={loss_initial:.6f}, final={loss_final:.6f}" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_state_dict_round_trip(self): + """Verify state_dict save/load preserves FP32 master weights and optimizer states.""" + model, optimizer, config = _make_model_and_optimizer(device="cuda") + + x = torch.randn(2, 8, dtype=torch.bfloat16, device="cuda") + loss = model(x).sum() + loss.backward() + optimizer.step() + + saved_state = optimizer.state_dict() + + original_masters = [] + for group in saved_state["fp32_from_fp16_params"]: + original_masters.append([p.clone() for p in group]) + + original_opt_states = {} + for pid, pstate in saved_state["optimizer"]["state"].items(): + original_opt_states[pid] = { + "exp_avg": pstate["exp_avg"].clone(), + "exp_avg_sq": pstate["exp_avg_sq"].clone(), + } + + fresh_optimizer = get_fsdp2_bf16_master_weight_optimizer(config=config, model_chunks=[model]) + fresh_optimizer.load_state_dict(saved_state) + + loaded_state = fresh_optimizer.state_dict() + for saved_group, loaded_group in zip(original_masters, loaded_state["fp32_from_fp16_params"]): + for saved_p, loaded_p in zip(saved_group, loaded_group): + torch.testing.assert_close(loaded_p, saved_p) + assert loaded_p.dtype == torch.float32 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_sharded_state_dict_creates_sharded_tensors(self): + """Verify sharded_state_dict produces ShardedTensor entries.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + model_ssd = _build_model_sharded_state_dict(model) + opt_ssd = optimizer.sharded_state_dict(model_ssd, is_loading=True) + + assert "optimizer" in opt_ssd + assert "fp32_from_fp16_params" in opt_ssd + + # Check optimizer states are ShardedTensors + state_entries = opt_ssd["optimizer"]["state"] + assert len(state_entries) > 0 + for pid, pstate in state_entries.items(): + if isinstance(pstate, dict): + for key in ("exp_avg", "exp_avg_sq"): + assert key in pstate + assert isinstance(pstate[key], ShardedTensor) + + # Check fp32_from_fp16_params are ShardedTensors + for group in opt_ssd["fp32_from_fp16_params"]: + for entry in group: + assert isinstance(entry, ShardedTensor) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_finalize_dist_ckpt_load(self): + """Verify finalize_dist_ckpt_load sets step counter and syncs params.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + optimizer.init_state_fn(optimizer.optimizer) + + # Modify FP32 masters to differ from BF16 params + for fp32_group in optimizer.fp32_from_bf16_groups: + for p in fp32_group: + p.data.add_(1.0) + + iteration = 42 + optimizer.finalize_dist_ckpt_load(iteration) + + for fp32_group in optimizer.fp32_from_bf16_groups: + for p in fp32_group: + if p in optimizer.optimizer.state: + step = optimizer.optimizer.state[p].get("step") + if step is not None: + assert step.item() == float(iteration) + + # BF16 params should now match FP32 masters (copy-back happened) + for bf16_group, fp32_group in zip(optimizer.bf16_groups, optimizer.fp32_from_bf16_groups): + for bf16_param, fp32_master in zip(bf16_group, fp32_group): + expected = fp32_master.data.to(torch.bfloat16) + torch.testing.assert_close(bf16_param.data, expected, atol=0, rtol=0) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_reload_model_params(self): + """Verify reload_model_params copies BF16 params to FP32 masters.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + # Manually change BF16 params (simulating checkpoint load) + for p in model.parameters(): + p.data.fill_(0.5) + + optimizer.reload_model_params() + + for bf16_group, fp32_group in zip(optimizer.bf16_groups, optimizer.fp32_from_bf16_groups): + for bf16_param, fp32_master in zip(bf16_group, fp32_group): + torch.testing.assert_close(fp32_master, bf16_param.float(), atol=0, rtol=0) + + def test_factory_applies_weight_decay_condition(self): + """Verify no_weight_decay_cond splits parameters into correct groups.""" + model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.bfloat16) + config = _make_optimizer_config(weight_decay=0.1) + + def no_wd_for_bias(param): + return param.ndim == 1 + + optimizer = get_fsdp2_bf16_master_weight_optimizer( + config=config, + model_chunks=[model], + no_weight_decay_cond=no_wd_for_bias, + ) + + assert len(optimizer.optimizer.param_groups) == 2 + wds = {pg["weight_decay"] for pg in optimizer.optimizer.param_groups} + assert 0.0 in wds + assert 0.1 in wds + + def test_factory_handles_fp32_params(self): + """Verify FP32 params (e.g. LayerNorm) go to fp32_from_fp32_groups without master copy.""" + model = torch.nn.Sequential( + torch.nn.Linear(8, 8, bias=False).to(dtype=torch.bfloat16), + torch.nn.LayerNorm(8).to(dtype=torch.float32), + ) + config = _make_optimizer_config() + optimizer = get_fsdp2_bf16_master_weight_optimizer(config=config, model_chunks=[model]) + + total_bf16 = sum(len(g) for g in optimizer.bf16_groups) + total_fp32 = sum(len(g) for g in optimizer.fp32_from_fp32_groups) + + assert total_bf16 == 1, f"Expected 1 BF16 param, got {total_bf16}" + assert total_fp32 == 2, f"Expected 2 FP32 params (LN weight+bias), got {total_fp32}" diff --git a/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_fp32_optimizer.py b/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_fp32_optimizer.py new file mode 100644 index 000000000..9bda55efa --- /dev/null +++ b/tests/unit_tests/backends/megatron/optimizer/test_fsdp2_fp32_optimizer.py @@ -0,0 +1,311 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for FSDP2FP32Optimizer. + +Tests focus on the core design properties of this optimizer: +- Optimizer states are FP32 (the whole point) +- state_dict format is raw PyTorch (not Megatron-wrapped) +- Gradient clipping actually clips gradients +- Optimizer step actually updates parameters and reduces loss +- Checkpoint save/load round-trip preserves state values and dtypes +- sharded_state_dict integrates with Megatron's torch_dist checkpointing +- finalize_dist_ckpt_load fills step counter (FSDP2-specific) +- Factory param group splitting (weight decay, LR scaling) works +""" + +import pytest +import torch +from megatron.core.dist_checkpointing.mapping import ShardedTensor +from megatron.core.optimizer.optimizer_config import OptimizerConfig +from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint + +from primus.backends.megatron.core.optimizer.fsdp2_fp32_optimizer import ( + get_fsdp2_fp32_optimizer, +) +from tests.utils import PrimusUT + + +def _make_optimizer_config(**overrides): + defaults = dict( + bf16=True, + fp16=False, + optimizer="adam", + lr=1e-4, + min_lr=1e-5, + weight_decay=0.01, + adam_beta1=0.9, + adam_beta2=0.999, + adam_eps=1e-8, + clip_grad=1.0, + use_precision_aware_optimizer=False, + ) + defaults.update(overrides) + return OptimizerConfig(**defaults) + + +def _make_model_and_optimizer(device="cpu"): + model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.float32, device=device) + config = _make_optimizer_config() + optimizer = get_fsdp2_fp32_optimizer( + config=config, + model_chunks=[model], + ) + return model, optimizer, config + + +def _build_model_sharded_state_dict(model, prefix=""): + model_sd = model.state_dict(prefix=prefix, keep_vars=True) + return make_sharded_tensors_for_checkpoint(model_sd, prefix, {}, ()) + + +def _compute_grad_norm(model): + """Compute L2 gradient norm across all model parameters.""" + total = 0.0 + for p in model.parameters(): + if p.grad is not None: + total += p.grad.float().norm().item() ** 2 + return total**0.5 + + +class TestFSDP2FP32Optimizer(PrimusUT): + """Tests for FSDP2FP32Optimizer core functionality.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_optimizer_states_are_fp32(self): + """The core invariant: all optimizer states must be FP32.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + x = torch.randn(2, 8, dtype=torch.float32, device="cuda") + loss = model(x).sum() + loss.backward() + optimizer.step() + + for p in optimizer.get_parameters(): + state = optimizer.optimizer.state[p] + assert ( + state["exp_avg"].dtype == torch.float32 + ), f"exp_avg dtype is {state['exp_avg'].dtype}, expected float32" + assert ( + state["exp_avg_sq"].dtype == torch.float32 + ), f"exp_avg_sq dtype is {state['exp_avg_sq'].dtype}, expected float32" + assert state["step"].dtype == torch.float32 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_state_dict_format_is_raw(self): + """state_dict must return raw PyTorch format, NOT a Megatron-wrapped + format ({"optimizer": {...}}). + + This distinction is critical for checkpoint compatibility. + """ + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + x = torch.randn(2, 8, dtype=torch.float32, device="cuda") + loss = model(x).sum() + loss.backward() + optimizer.step() + + saved = optimizer.state_dict() + assert "state" in saved, "state_dict missing 'state' key" + assert "param_groups" in saved, "state_dict missing 'param_groups' key" + assert "optimizer" not in saved, ( + "state_dict has 'optimizer' wrapper key -- this is the BFloat16Optimizer " + "format, not the raw format expected by FSDP2FP32Optimizer" + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_clip_grad_norm_clips_gradients(self): + """Verify clipping reduces grad norm to at most clip_value.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + x = torch.randn(2, 8, dtype=torch.float32, device="cuda") + loss = (model(x) * 100.0).sum() + loss.backward() + + grad_norm_before = _compute_grad_norm(model) + assert grad_norm_before > 1.0, "Gradients too small to test clipping" + + clip_value = 0.01 + returned_norm = optimizer.clip_grad_norm(clip_value) + assert isinstance( + returned_norm, torch.Tensor + ), "clip_grad_norm should return a GPU tensor (deferred .item())" + returned_norm_f = returned_norm.item() + + grad_norm_after = _compute_grad_norm(model) + + assert ( + abs(returned_norm_f - grad_norm_before) < 1e-3 + ), f"Returned norm {returned_norm_f} should approximate pre-clip norm {grad_norm_before}" + assert ( + grad_norm_after <= clip_value + 1e-6 + ), f"Post-clip norm {grad_norm_after} exceeds clip_value {clip_value}" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_step_actually_updates_parameters(self): + """Verify optimizer.step() changes parameter values and returns + the (success, grad_norm, num_zeros) 3-tuple that Megatron's + training loop unpacks.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + params_before = {name: p.clone() for name, p in model.named_parameters()} + + x = torch.randn(2, 8, dtype=torch.float32, device="cuda") + loss = model(x).sum() + loss.backward() + success, grad_norm, num_zeros = optimizer.step() + + assert success is True + assert isinstance( + grad_norm, (float, torch.Tensor) + ), f"grad_norm should be float or Tensor, got {type(grad_norm)}" + grad_norm_f = grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm + assert grad_norm_f > 0.0 + + for name, p in model.named_parameters(): + assert not torch.equal( + p, params_before[name] + ), f"Parameter '{name}' was not updated by optimizer.step()" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_multi_step_loss_decreases(self): + """Run 10 steps and verify loss decreases, catching silent update failures.""" + torch.manual_seed(42) + model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.float32, device="cuda") + config = _make_optimizer_config(lr=1e-2, clip_grad=0.0) + optimizer = get_fsdp2_fp32_optimizer(config=config, model_chunks=[model]) + + x = torch.randn(4, 8, dtype=torch.float32, device="cuda") + target = torch.randn(4, 8, dtype=torch.float32, device="cuda") + + loss_initial = None + loss_final = None + for _ in range(10): + optimizer.zero_grad() + output = model(x) + loss = ((output - target) ** 2).mean() + loss.backward() + optimizer.step() + + if loss_initial is None: + loss_initial = loss.item() + loss_final = loss.item() + + assert ( + loss_final < loss_initial + ), f"Loss did not decrease: initial={loss_initial:.6f}, final={loss_final:.6f}" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_state_dict_round_trip(self): + """Verify state_dict save/load preserves values and FP32 dtypes.""" + model, optimizer, config = _make_model_and_optimizer(device="cuda") + + x = torch.randn(2, 8, dtype=torch.float32, device="cuda") + loss = model(x).sum() + loss.backward() + optimizer.step() + + saved_state = optimizer.state_dict() + + original_states = {} + for pid, pstate in saved_state["state"].items(): + original_states[pid] = { + "exp_avg": pstate["exp_avg"].clone(), + "exp_avg_sq": pstate["exp_avg_sq"].clone(), + } + + fresh_optimizer = get_fsdp2_fp32_optimizer( + config=config, + model_chunks=[model], + ) + fresh_optimizer.load_state_dict(saved_state) + + loaded_state = fresh_optimizer.state_dict() + for pid in original_states: + assert pid in loaded_state["state"] + torch.testing.assert_close( + loaded_state["state"][pid]["exp_avg"], + original_states[pid]["exp_avg"], + ) + torch.testing.assert_close( + loaded_state["state"][pid]["exp_avg_sq"], + original_states[pid]["exp_avg_sq"], + ) + assert loaded_state["state"][pid]["exp_avg"].dtype == torch.float32 + assert loaded_state["state"][pid]["exp_avg_sq"].dtype == torch.float32 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_sharded_state_dict_creates_sharded_tensors(self): + """Verify sharded_state_dict produces ShardedTensor entries for torch_dist checkpointing.""" + model, optimizer, _ = _make_model_and_optimizer(device="cuda") + + model_ssd = _build_model_sharded_state_dict(model) + opt_ssd = optimizer.sharded_state_dict(model_ssd, is_loading=True) + + assert "state" in opt_ssd + state_entries = opt_ssd["state"] + assert len(state_entries) > 0 + + for pid, pstate in state_entries.items(): + if isinstance(pstate, dict): + for key in ("exp_avg", "exp_avg_sq"): + assert key in pstate, f"Key '{key}' missing from state[{pid}]" + assert isinstance(pstate[key], ShardedTensor) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_finalize_dist_ckpt_load(self): + """Verify finalize_dist_ckpt_load sets step counter from iteration.""" + model, optimizer, config = _make_model_and_optimizer(device="cuda") + + optimizer.init_state_fn(optimizer.optimizer) + + iteration = 42 + optimizer.finalize_dist_ckpt_load(iteration) + + for p in optimizer.get_parameters(): + if p in optimizer.optimizer.state: + step = optimizer.optimizer.state[p].get("step") + if step is not None: + assert step.item() == float(iteration) + + def test_init_state_fn_creates_fp32_states(self): + """Verify init_state_fn creates FP32 placeholders for checkpoint loading.""" + model, optimizer, _ = _make_model_and_optimizer() + + base_opt = optimizer.optimizer + for group in base_opt.param_groups: + for p in group["params"]: + assert len(base_opt.state[p]) == 0 + + optimizer.init_state_fn(base_opt) + + for group in base_opt.param_groups: + for p in group["params"]: + state = base_opt.state[p] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + assert state["exp_avg"].shape == p.data.shape + + def test_factory_applies_weight_decay_condition(self): + """Verify no_weight_decay_cond splits parameters into correct groups.""" + model = torch.nn.Linear(8, 8, bias=True).to(dtype=torch.float32) + config = _make_optimizer_config(weight_decay=0.1) + + def no_wd_for_bias(param): + return param.ndim == 1 + + optimizer = get_fsdp2_fp32_optimizer( + config=config, + model_chunks=[model], + no_weight_decay_cond=no_wd_for_bias, + ) + + assert len(optimizer.optimizer.param_groups) == 2 + wds = {pg["weight_decay"] for pg in optimizer.optimizer.param_groups} + assert 0.0 in wds + assert 0.1 in wds diff --git a/tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py b/tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py new file mode 100644 index 000000000..e4951cfe5 --- /dev/null +++ b/tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py @@ -0,0 +1,488 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for FSDP2 FP32 optimizer patches. + +Tests: +- FSDP2 FP32 optimizer patch condition evaluation and application +- Float16Module skip patch condition evaluation and application +- _FSDP2PassthroughModule preserves FP32 weights and passes through forward +- Mutual exclusivity of bf16 and fsdp2_fp32 optimizer patches +""" + +import sys +import types +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from primus.core.patches.context import PatchContext +from primus.core.patches.patch_runner import run_patches + + +def _build_patch_context(backend_args=None): + if backend_args is None: + backend_args = SimpleNamespace() + module_config = SimpleNamespace(params=backend_args) + return PatchContext( + backend="megatron", + phase="before_train", + extra={"backend_args": backend_args, "module_config": module_config}, + ) + + +def _install_fake_megatron_modules(monkeypatch): + training_mod = types.ModuleType("megatron.training.training") + training_pkg = types.ModuleType("megatron.training") + training_pkg.training = training_mod + + core_pkg = types.ModuleType("megatron.core") + optimizer_mod = types.ModuleType("megatron.core.optimizer") + core_pkg.optimizer = optimizer_mod + + megatron_pkg = types.ModuleType("megatron") + megatron_pkg.training = training_pkg + megatron_pkg.core = core_pkg + + monkeypatch.setitem(sys.modules, "megatron", megatron_pkg) + monkeypatch.setitem(sys.modules, "megatron.training", training_pkg) + monkeypatch.setitem(sys.modules, "megatron.training.training", training_mod) + monkeypatch.setitem(sys.modules, "megatron.core", core_pkg) + monkeypatch.setitem(sys.modules, "megatron.core.optimizer", optimizer_mod) + + return training_mod, optimizer_mod + + +def _install_fsdp2_fp32_stub(monkeypatch, optimizer_fn): + stub = types.ModuleType("primus.backends.megatron.core.optimizer.fsdp2_fp32_optimizer") + stub.get_fsdp2_fp32_optimizer = optimizer_fn + monkeypatch.setitem( + sys.modules, + "primus.backends.megatron.core.optimizer.fsdp2_fp32_optimizer", + stub, + ) + + +class TestFSDP2FP32OptimizerPatch: + """Tests for the FSDP2 FP32 optimizer patch.""" + + def test_patch_applies_when_all_conditions_met(self, monkeypatch): + from primus.backends.megatron.patches.optimizer_patches import ( + patch_fsdp2_fp32_optimizer, + ) + + backend_args = SimpleNamespace( + use_fsdp2_fp32_param_optimizer=True, + use_torch_fsdp2=True, + bf16=True, + ) + ctx = _build_patch_context(backend_args) + + training_mod, optimizer_mod = _install_fake_megatron_modules(monkeypatch) + original = lambda *args, **kwargs: Mock() + training_mod.get_megatron_optimizer = original + optimizer_mod.get_megatron_optimizer = original + + mock_optimizer = Mock() + _install_fsdp2_fp32_stub(monkeypatch, lambda *args, **kwargs: mock_optimizer) + + monkeypatch.setattr( + "primus.backends.megatron.patches.optimizer_patches.log_rank_0", + lambda *args, **kwargs: None, + ) + + patch_fsdp2_fp32_optimizer(ctx) + + assert training_mod.get_megatron_optimizer is not original + assert optimizer_mod.get_megatron_optimizer is not original + + result = training_mod.get_megatron_optimizer(Mock(), [Mock()]) + assert result is mock_optimizer + + @pytest.mark.parametrize( + "backend_args", + [ + SimpleNamespace(use_fsdp2_fp32_param_optimizer=False, use_torch_fsdp2=True, bf16=True), + SimpleNamespace(use_fsdp2_fp32_param_optimizer=True, use_torch_fsdp2=False, bf16=True), + SimpleNamespace(use_fsdp2_fp32_param_optimizer=True, use_torch_fsdp2=True, bf16=False), + ], + ) + def test_patch_skipped_when_any_condition_false(self, monkeypatch, backend_args): + import primus.backends.megatron.patches.optimizer_patches # noqa: F401 + + monkeypatch.setattr( + "primus.core.patches.patch_runner.log_rank_0", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "primus.core.patches.patch_runner.error_rank_0", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "primus.core.patches.patch.log_rank_0", + lambda *args, **kwargs: None, + ) + + training_mod, optimizer_mod = _install_fake_megatron_modules(monkeypatch) + original = lambda *args, **kwargs: Mock() + training_mod.get_megatron_optimizer = original + optimizer_mod.get_megatron_optimizer = original + + count = run_patches( + backend="megatron", + phase="before_train", + extra={"module_config": SimpleNamespace(params=backend_args)}, + enabled_ids=["megatron.optimizer.fsdp2_fp32_param"], + ) + + assert count == 0 + assert training_mod.get_megatron_optimizer is original + + def test_patched_optimizer_filters_megatron_kwargs(self, monkeypatch): + from primus.backends.megatron.patches.optimizer_patches import ( + patch_fsdp2_fp32_optimizer, + ) + + backend_args = SimpleNamespace( + use_fsdp2_fp32_param_optimizer=True, + use_torch_fsdp2=True, + bf16=True, + ) + ctx = _build_patch_context(backend_args) + + training_mod, optimizer_mod = _install_fake_megatron_modules(monkeypatch) + training_mod.get_megatron_optimizer = lambda *a, **kw: Mock() + optimizer_mod.get_megatron_optimizer = lambda *a, **kw: Mock() + + received_kwargs = {} + + def tracking_factory(**kwargs): + received_kwargs.update(kwargs) + return Mock() + + _install_fsdp2_fp32_stub(monkeypatch, lambda *args, **kwargs: tracking_factory(**kwargs)) + + monkeypatch.setattr( + "primus.backends.megatron.patches.optimizer_patches.log_rank_0", + lambda *args, **kwargs: None, + ) + + patch_fsdp2_fp32_optimizer(ctx) + + training_mod.get_megatron_optimizer( + Mock(), + [Mock()], + config_overrides={"x": 1}, + use_gloo_process_groups=True, + no_weight_decay_cond=Mock(), + ) + + assert "config_overrides" not in received_kwargs + assert "use_gloo_process_groups" not in received_kwargs + assert "no_weight_decay_cond" in received_kwargs + + +class TestSkipFloat16ModulePatch: + """Tests for the Float16Module skip patch.""" + + def test_patch_replaces_float16_module(self, monkeypatch): + from primus.backends.megatron.patches.build_model_patches import ( + patch_skip_float16_module, + ) + + backend_args = SimpleNamespace( + use_fsdp2_fp32_param_optimizer=True, + use_torch_fsdp2=True, + bf16=True, + ) + ctx = _build_patch_context(backend_args) + + monkeypatch.setattr( + "primus.backends.megatron.patches.build_model_patches.log_rank_0", + lambda *args, **kwargs: None, + ) + + _install_fake_megatron_modules(monkeypatch) + + transformer_pkg = types.ModuleType("megatron.core.transformer") + transformer_module_mod = types.ModuleType("megatron.core.transformer.module") + + class FakeFloat16Module: + pass + + class FakeMegatronModule(torch.nn.Module): + def __init__(self, config=None): + super().__init__() + + transformer_module_mod.Float16Module = FakeFloat16Module + transformer_module_mod.MegatronModule = FakeMegatronModule + transformer_pkg.module = transformer_module_mod + monkeypatch.setitem( + sys.modules, + "megatron.core.transformer", + transformer_pkg, + ) + monkeypatch.setitem( + sys.modules, + "megatron.core.transformer.module", + transformer_module_mod, + ) + sys.modules["megatron.core"].transformer = transformer_pkg + + patch_skip_float16_module(ctx) + + import megatron.training.training as training_mod + + assert training_mod.Float16Module is not FakeFloat16Module + assert training_mod.Float16Module.__name__ == "_FSDP2PassthroughModule" + + def test_passthrough_module_preserves_fp32_weights(self, monkeypatch): + """Verify _FSDP2PassthroughModule does NOT convert parameters to BF16 + and that forward() passes through correctly. + + This is the core contract: Float16Module converts to BF16, but the + passthrough replacement must keep FP32 so FSDP2 MixedPrecisionPolicy + can handle the casting. + """ + from primus.backends.megatron.patches.build_model_patches import ( + patch_skip_float16_module, + ) + + backend_args = SimpleNamespace( + use_fsdp2_fp32_param_optimizer=True, + use_torch_fsdp2=True, + bf16=True, + ) + ctx = _build_patch_context(backend_args) + + monkeypatch.setattr( + "primus.backends.megatron.patches.build_model_patches.log_rank_0", + lambda *args, **kwargs: None, + ) + + _install_fake_megatron_modules(monkeypatch) + + transformer_pkg = types.ModuleType("megatron.core.transformer") + transformer_module_mod = types.ModuleType("megatron.core.transformer.module") + + class FakeMegatronModule(torch.nn.Module): + def __init__(self, config=None): + super().__init__() + + class FakeFloat16Module: + pass + + transformer_module_mod.Float16Module = FakeFloat16Module + transformer_module_mod.MegatronModule = FakeMegatronModule + transformer_pkg.module = transformer_module_mod + monkeypatch.setitem( + sys.modules, + "megatron.core.transformer", + transformer_pkg, + ) + monkeypatch.setitem( + sys.modules, + "megatron.core.transformer.module", + transformer_module_mod, + ) + sys.modules["megatron.core"].transformer = transformer_pkg + + patch_skip_float16_module(ctx) + + import megatron.training.training as training_mod + + PassthroughCls = training_mod.Float16Module + + inner_model = torch.nn.Linear(4, 4, bias=True).to(dtype=torch.float32) + config = SimpleNamespace(virtual_pipeline_model_parallel_size=None) + wrapper = PassthroughCls(config, inner_model) + + # Parameters must remain FP32 + for name, p in wrapper.named_parameters(): + assert ( + p.dtype == torch.float32 + ), f"Parameter '{name}' was converted to {p.dtype}, expected float32" + + # wrapper.module must be the original model + assert wrapper.module is inner_model + + # Forward must pass through correctly + x = torch.randn(2, 4, dtype=torch.float32) + expected = inner_model(x) + actual = wrapper(x) + torch.testing.assert_close(actual, expected) + + @pytest.mark.parametrize( + "backend_args", + [ + SimpleNamespace(use_fsdp2_fp32_param_optimizer=False, use_torch_fsdp2=True, bf16=True), + SimpleNamespace(use_fsdp2_fp32_param_optimizer=True, use_torch_fsdp2=False, bf16=True), + SimpleNamespace(use_fsdp2_fp32_param_optimizer=True, use_torch_fsdp2=True, bf16=False), + ], + ) + def test_patch_skipped_when_conditions_not_met(self, monkeypatch, backend_args): + import primus.backends.megatron.patches.build_model_patches # noqa: F401 + + monkeypatch.setattr( + "primus.core.patches.patch_runner.log_rank_0", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "primus.core.patches.patch_runner.error_rank_0", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "primus.core.patches.patch.log_rank_0", + lambda *args, **kwargs: None, + ) + + count = run_patches( + backend="megatron", + phase="before_train", + extra={"module_config": SimpleNamespace(params=backend_args)}, + enabled_ids=["megatron.training.training.skip_float16_module"], + ) + + assert count == 0 + + +class TestFSDP2BF16MasterWeightOptimizerPatch: + """Tests for the FSDP2 BF16 master weight optimizer patch.""" + + def test_patch_applies_when_all_conditions_met(self, monkeypatch): + from primus.backends.megatron.patches.optimizer_patches import ( + patch_fsdp2_bf16_master_weight_optimizer, + ) + + backend_args = SimpleNamespace( + use_fsdp2_bf16_master_weight_optimizer=True, + use_torch_fsdp2=True, + bf16=True, + ) + ctx = _build_patch_context(backend_args) + + training_mod, optimizer_mod = _install_fake_megatron_modules(monkeypatch) + original = lambda *args, **kwargs: Mock() + training_mod.get_megatron_optimizer = original + optimizer_mod.get_megatron_optimizer = original + + mock_optimizer = Mock() + bf16_mw_stub = types.ModuleType( + "primus.backends.megatron.core.optimizer.fsdp2_bf16_master_weight_optimizer" + ) + bf16_mw_stub.get_fsdp2_bf16_master_weight_optimizer = lambda *args, **kwargs: mock_optimizer + monkeypatch.setitem( + sys.modules, + "primus.backends.megatron.core.optimizer.fsdp2_bf16_master_weight_optimizer", + bf16_mw_stub, + ) + + monkeypatch.setattr( + "primus.backends.megatron.patches.optimizer_patches.log_rank_0", + lambda *args, **kwargs: None, + ) + + patch_fsdp2_bf16_master_weight_optimizer(ctx) + + assert training_mod.get_megatron_optimizer is not original + assert optimizer_mod.get_megatron_optimizer is not original + + result = training_mod.get_megatron_optimizer(Mock(), [Mock()]) + assert result is mock_optimizer + + @pytest.mark.parametrize( + "backend_args", + [ + SimpleNamespace( + use_fsdp2_bf16_master_weight_optimizer=False, + use_torch_fsdp2=True, + bf16=True, + ), + SimpleNamespace( + use_fsdp2_bf16_master_weight_optimizer=True, + use_torch_fsdp2=False, + bf16=True, + ), + SimpleNamespace( + use_fsdp2_bf16_master_weight_optimizer=True, + use_torch_fsdp2=True, + bf16=False, + ), + ], + ) + def test_patch_skipped_when_any_condition_false(self, monkeypatch, backend_args): + import primus.backends.megatron.patches.optimizer_patches # noqa: F401 + + monkeypatch.setattr( + "primus.core.patches.patch_runner.log_rank_0", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "primus.core.patches.patch_runner.error_rank_0", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "primus.core.patches.patch.log_rank_0", + lambda *args, **kwargs: None, + ) + + training_mod, optimizer_mod = _install_fake_megatron_modules(monkeypatch) + original = lambda *args, **kwargs: Mock() + training_mod.get_megatron_optimizer = original + optimizer_mod.get_megatron_optimizer = original + + count = run_patches( + backend="megatron", + phase="before_train", + extra={"module_config": SimpleNamespace(params=backend_args)}, + enabled_ids=["megatron.optimizer.fsdp2_bf16_master_weight"], + ) + + assert count == 0 + assert training_mod.get_megatron_optimizer is original + + +class TestOptimizerPatchMutualExclusivity: + """Tests for mutual exclusivity of the FSDP2 custom optimizer patches.""" + + def test_two_conflicting_optimizer_flags_raise(self): + """Two FSDP2 custom optimizer flags set must fail loudly: the flags all + patch get_megatron_optimizer at the same priority, so enabling more than + one is ambiguous. validate_fsdp2_optimizer_exclusivity enforces this at + arg-validation time.""" + from primus.modules.trainer.megatron.utils import ( + validate_fsdp2_optimizer_exclusivity, + ) + + args = SimpleNamespace( + use_fsdp2_fp32_param_optimizer=True, + use_fsdp2_bf16_master_weight_optimizer=True, + ) + + with pytest.raises(ValueError, match="Conflicting FSDP2 optimizer selection"): + validate_fsdp2_optimizer_exclusivity(args) + + def test_single_optimizer_flag_is_allowed(self): + """Exactly one FSDP2 optimizer flag is the supported case (no raise).""" + from primus.modules.trainer.megatron.utils import ( + validate_fsdp2_optimizer_exclusivity, + ) + + for flag in ( + "use_fsdp2_fp32_param_optimizer", + "use_fsdp2_bf16_master_weight_optimizer", + ): + args = SimpleNamespace(**{flag: True}) + # Should not raise. + validate_fsdp2_optimizer_exclusivity(args) + + def test_no_optimizer_flag_is_allowed(self): + """No FSDP2 optimizer flag set (default Megatron optimizer) must not raise.""" + from primus.modules.trainer.megatron.utils import ( + validate_fsdp2_optimizer_exclusivity, + ) + + validate_fsdp2_optimizer_exclusivity(SimpleNamespace()) From c22ac17bc3f8ceac8619958a75d02a645ca62ca6 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Wed, 8 Jul 2026 03:01:40 +0300 Subject: [PATCH 008/127] feat(flux): curated diffusion example/model/data configs (#821) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/core` — review after it. Can open early; its example runs go green once the runtime PRs land. ## What this changes The curated diffusion config set — the MI300X/MI355X example configs plus the `primus/configs/{data,models,modules}/megatron/diffusion*` definitions and the torch_compile / trainer_base module configs. ## Why it's stacked here It adds files under `primus/configs/data/**`, which public `main`'s broad `data` `.gitignore` rule would otherwise ignore — so it needs the core PR's `.gitignore` fix. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/core`. ## Test plan `yaml.safe_load` parse-check per file (no pytest); optionally one `run_pretrain.sh --dry-run` per GPU family. ## Files ~43 (example + model/data/module diffusion configs). Co-authored-by: Flux Split Trial --- ...ergon_schnell_resample_local_spec_fp8.yaml | 209 ++++++ ..._ddp_energon_schnell_resample_te_spec.yaml | 187 +++++ ..._energon_schnell_resample_te_spec_fp8.yaml | 210 ++++++ ...2_energon_schnell_resample_local_spec.yaml | 177 +++++ ...ergon_schnell_resample_local_spec_fp8.yaml | 190 +++++ .../MI300X/diffusion/flux_535m_pretrain.yaml | 170 +++++ .../diffusion/flux_535m_pretrain_fp8.yaml | 199 ++++++ .../flux_535m_with_guidance_embed.yaml | 58 ++ ...ergon_schnell_resample_local_spec_fp8.yaml | 209 ++++++ ...chnell_resample_local_spec_fp8_mlperf.yaml | 202 ++++++ ...gon_schnell_resample_local_spec_mxfp4.yaml | 198 ++++++ ..._ddp_energon_schnell_resample_te_spec.yaml | 184 +++++ ..._energon_schnell_resample_te_spec_fp8.yaml | 209 ++++++ ...n_schnell_resample_te_spec_fp8_mlperf.yaml | 202 ++++++ ...2_energon_schnell_resample_local_spec.yaml | 175 +++++ ...ergon_schnell_resample_local_spec_fp8.yaml | 198 ++++++ .../MI355X/diffusion/flux_535m_pretrain.yaml | 170 +++++ .../diffusion/flux_535m_pretrain_fp8.yaml | 199 ++++++ .../flux_535m_with_guidance_embed.yaml | 58 ++ .../configs/data/megatron/diffusion/README.md | 659 ++++++++++++++++++ .../coco2014_train_schnell_256.yaml | 31 + .../preprocessing/coco_schnell_256.yaml | 36 + .../diffusion/preprocessing/example_base.yaml | 196 ++++++ .../preprocessing/example_directory.yaml | 130 ++++ .../preprocessing/example_huggingface.yaml | 124 ++++ .../preprocessing/example_webdataset.yaml | 129 ++++ .../diffusion/preprocessing/mlperf_flux1.yaml | 66 ++ .../preprocessing/quickstart_pokemon.yaml | 34 + .../preprocessing/text_to_image_2m_10k.yaml | 93 +++ .../templates/dataset_preencoded.yaml | 79 +++ .../templates/dataset_preencoded_numpy.yaml | 16 + .../diffusion/templates/dataset_raw.yaml | 112 +++ .../diffusion/templates/metadataset.yaml | 81 +++ .../models/megatron/diffusion/encoders.yaml | 126 ++++ .../models/megatron/diffusion/flux_12b.yaml | 141 ++++ .../megatron/diffusion/flux_12b_fp8.yaml | 122 ++++ .../diffusion/flux_12b_rope_fusion.yaml | 50 ++ .../models/megatron/diffusion/flux_535m.yaml | 85 +++ .../megatron/diffusion/flux_535m_fp8.yaml | 107 +++ .../models/megatron/diffusion/flux_base.yaml | 110 +++ .../models/megatron/diffusion_model.yaml | 135 ++++ .../modules/megatron/torch_compile.yaml | 13 + .../modules/megatron/trainer_base.yaml | 17 + 43 files changed, 6096 insertions(+) create mode 100644 examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml create mode 100644 examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml create mode 100644 examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml create mode 100644 examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml create mode 100644 examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml create mode 100644 examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml create mode 100644 examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml create mode 100644 examples/megatron/configs/MI300X/diffusion/flux_535m_with_guidance_embed.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml create mode 100644 examples/megatron/configs/MI355X/diffusion/flux_535m_with_guidance_embed.yaml create mode 100644 primus/configs/data/megatron/diffusion/README.md create mode 100644 primus/configs/data/megatron/diffusion/preprocessing/coco2014_train_schnell_256.yaml create mode 100644 primus/configs/data/megatron/diffusion/preprocessing/coco_schnell_256.yaml create mode 100644 primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml create mode 100644 primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml create mode 100644 primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml create mode 100644 primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml create mode 100644 primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml create mode 100644 primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml create mode 100644 primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml create mode 100644 primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml create mode 100644 primus/configs/data/megatron/diffusion/templates/dataset_preencoded_numpy.yaml create mode 100644 primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml create mode 100644 primus/configs/data/megatron/diffusion/templates/metadataset.yaml create mode 100644 primus/configs/models/megatron/diffusion/encoders.yaml create mode 100644 primus/configs/models/megatron/diffusion/flux_12b.yaml create mode 100644 primus/configs/models/megatron/diffusion/flux_12b_fp8.yaml create mode 100644 primus/configs/models/megatron/diffusion/flux_12b_rope_fusion.yaml create mode 100644 primus/configs/models/megatron/diffusion/flux_535m.yaml create mode 100644 primus/configs/models/megatron/diffusion/flux_535m_fp8.yaml create mode 100644 primus/configs/models/megatron/diffusion/flux_base.yaml create mode 100644 primus/configs/models/megatron/diffusion_model.yaml create mode 100644 primus/configs/modules/megatron/torch_compile.yaml diff --git a/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml new file mode 100644 index 000000000..dbd76aa26 --- /dev/null +++ b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml @@ -0,0 +1,209 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — DDP + Local Spec + FP8 Tensorwise (MI300X) +# +# Combines Megatron DDP + distributed optimizer with the local spec provider +# (PrimusTurboFloat8LocalSpecProvider) for FP8 tensorwise training. +# +# Key configuration: +# - Megatron DDP with overlap_grad_reduce + overlap_param_gather +# - PrimusTurboFloat8LocalSpecProvider (NO TransformerEngine dependency) +# - FP8 hybrid tensorwise (per-module FP8 via Primus Turbo) +# - Primus Turbo attention +# - torch.compile enabled (per_block strategy, compatible with local spec + overlap) +# - Energon pre-encoded dataset with stored VAE mean/logvar (resample mode) + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_fp8} + +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + # Short LR warmup ramp over the first optimizer steps for FP8 stability. + nemo_aligned_lr_warmup: true + warmup_train_steps: 2 + + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + + # Resample mode: re-draw latents from mean+logvar each step + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboFloat8LocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # DATASET CONFIGURATION + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # Batch configuration (scaled down for MI300X 192GB; tune to your hardware) + micro_batch_size: 32 + global_batch_size: 256 + seq_length: 512 # 256 img tokens + 256 text tokens (schnell) + + # ========================================== + # BF16 + FP8 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # Optimizer settings + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader settings + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # Megatron DDP + Distributed Optimizer + # ========================================== + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: false + use_megatron_fsdp: false + + use_distributed_optimizer: true + overlap_grad_reduce: true + overlap_param_gather: true + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: true + main_params_dtype: fp32 + main_grads_dtype: bf16 + exp_avg_dtype: bf16 + exp_avg_sq_dtype: bf16 + + ddp_bucket_size: 256000000 + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + gradient_accumulation_fusion: false + + check_for_nan_in_loss_and_grad: true + + # ========================================== + # FP8 Configuration — Tensorwise + Delayed FP8 via Primus Turbo + # ========================================== + + use_flash_attn: true + + fp8: "hybrid" + # TE-compatible unified form: fp8: hybrid + fp8_recipe: delayed selects + # tensorwise FP8 with delayed scaling. Resolution path: + # primus/backends/megatron/core/extensions/primus_turbo_float8_local.py + # :: Float8{Column,Row}ParallelLinear._use_delayed_scaling. + fp8_recipe: "delayed" + fp8_margin: 0 + fp8_amax_history_len: 1024 + fp8_amax_compute_algo: "max" + fp8_wgrad: true + fp8_dot_product_attention: false + fp8_multi_head_attention: false + fp8_reduce_amax: true + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 180 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_ddp_local_fp8 + wandb_project: flux_12b_ddp_local_fp8 + log_throughput: true + wall_clock_step_timer: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # PyTorch Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo Configuration + enable_primus_turbo: true + use_turbo_attention: true + use_dual_fp8_output_projection: false + + seed: 2025 + per_step_rng_reseed: false + nemo_chimera_init: false + + # Torch Compile — compatible with Float8 local spec (per-module FP8) + torch_compile: + enable: true + strategy: "per_block" + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false + emulate_precision_casts: false + fused_ln_modulate: true diff --git a/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml new file mode 100644 index 000000000..12b11d34a --- /dev/null +++ b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml @@ -0,0 +1,187 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — DDP + TransformerEngine Spec + BF16 (MI300X) +# +# BF16 training using Megatron DDP + distributed optimizer with +# TransformerEngine modules (TEColumnParallelLinear / TERowParallelLinear / +# TEDotProductAttention / TENorm). This is the BF16 baseline on the +# TransformerEngine path; see the *_te_spec_fp8 variant for FP8. +# +# Key settings: +# bf16: true +# params_dtype: bfloat16 +# micro_batch_size: 32 / global_batch_size: 256 +# +# Batch sizes are scaled down from the MI355X recipe (MBS=64/GBS=512) for +# MI300X (192GB) headroom; tune to your hardware. + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_bf16} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b_rope_fusion.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + + # Resample mode: re-draw latents from mean+logvar each step + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # TransformerEngine Spec (default) + # ========================================== + transformer_impl: "transformer_engine" + + # ========================================== + # ENERGON DATASET CONFIGURATION + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # Batch configuration (scaled down for MI300X; tune to your hardware) + micro_batch_size: 32 + global_batch_size: 256 + seq_length: 512 + + # ========================================== + # BF16 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # Optimizer settings (identical to BF16 baseline) + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader settings + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # Megatron DDP + Distributed Optimizer + # ========================================== + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: false + use_megatron_fsdp: false + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + use_distributed_optimizer: true + + overlap_grad_reduce: true + overlap_param_gather: true + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: false + + ddp_bucket_size: 256000000 + + gradient_accumulation_fusion: false + + # ========================================== + # Memory Optimizations + # ========================================== + + use_flash_attn: true + + fp8: null + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_ddp_te_bf16 + wandb_project: flux_12b_ddp_te_bf16 + log_throughput: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # PyTorch Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo Configuration + enable_primus_turbo: false + use_turbo_attention: false + + seed: 2025 + te_rng_tracker: true + + # Torch Compile — stack strategy (compiles the double/single DiT block + # stacks as inductor regions). Matches the MI355X te_spec recipe and the + # MLPerf NeMo reference (COMPILE_DIT strategy=stack); reduces activation + # memory and improves throughput on the eager BF16 TE path. + torch_compile: + enable: true + strategy: "stack" + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false diff --git a/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml new file mode 100644 index 000000000..2df61605f --- /dev/null +++ b/examples/megatron/configs/MI300X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml @@ -0,0 +1,210 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — DDP + TE Spec + FP8 Delayed Scaling (MI300X) +# +# FP8 training using Megatron DDP + distributed optimizer with TransformerEngine +# modules and FP8 hybrid delayed scaling: +# - Megatron DDP with overlap_grad_reduce + overlap_param_gather +# - TEColumnParallelLinear / TERowParallelLinear / TEDotProductAttention / TENorm +# - FP8 hybrid (E4M3 fwd, E5M2 bwd) with delayed scaling (amax history 1024) +# - RoPE fusion via apply_rope_fusion: true +# - Energon pre-encoded dataset +# +# Required environment variables for MI300X (set before launch): +# export NVTE_FUSED_ATTN=1 +# export NVTE_FUSED_ATTN_CK=1 +# export NVTE_FP8_DPA_BWD=1 +# export NVTE_USE_HIPBLASLT=1 +# export USE_HIPBLASLT=1 +# export TORCH_BLAS_PREFER_HIPBLASLT=1 +# export NVTE_USE_CAST_TRANSPOSE_TRITON=1 + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_fp8} + +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + vae_scale: 0.3611 + vae_shift: 0.1159 + vae_latent_mode: resample + + # ========================================== + # RoPE Fusion + # ========================================== + rotary_interleaved: true + apply_rope_fusion: true + position_embedding_type: rope + + # ========================================== + # TransformerEngine Spec + # ========================================== + transformer_impl: "transformer_engine" + + # ========================================== + # FP8 — hybrid delayed scaling + # ========================================== + fp8: "hybrid" + fp8_recipe: "delayed" + fp8_margin: 0 + fp8_amax_history_len: 1024 + fp8_amax_compute_algo: "max" + fp8_wgrad: true + fp8_dot_product_attention: false + fp8_multi_head_attention: false + + # ========================================== + # Energon Dataset + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # ========================================== + # Batch Configuration + # ========================================== + # MI300X has 192GB HBM3 (vs 256GB on MI355X), so MBS reduced + # from 64 to 32 and GBS from 512 to 256 to avoid OOM with FP8. + micro_batch_size: 32 + global_batch_size: 256 + seq_length: 512 + + # ========================================== + # BF16 + FP8 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # ========================================== + # Optimizer + # ========================================== + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # Megatron DDP + Distributed Optimizer + # ========================================== + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: false + use_megatron_fsdp: false + + use_distributed_optimizer: true + overlap_grad_reduce: true + overlap_param_gather: true + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: false + + ddp_bucket_size: 256000000 + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + gradient_accumulation_fusion: false + + check_for_nan_in_loss_and_grad: true + + # ========================================== + # Memory / Misc + # ========================================== + use_flash_attn: true + empty_unused_memory_level: 0 + + # Manual GC — align GC timing across ranks to avoid stragglers + manual_gc: true + manual_gc_interval: 1000 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_ddp_te_fp8 + wandb_project: flux_12b_ddp_te_fp8 + log_throughput: true + wall_clock_step_timer: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo — disabled for pure TE path + enable_primus_turbo: false + use_turbo_attention: false + + seed: 2025 + te_rng_tracker: true + + # torch.compile — selective stack compilation for TE spec + torch_compile: + enable: true + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false + strategy: "stack" + replace_qk_rmsnorm: true + disable_inductor_cudagraphs: false diff --git a/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml b/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml new file mode 100644 index 000000000..a4375f9a6 --- /dev/null +++ b/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml @@ -0,0 +1,177 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — FSDP2 + Local Spec + BF16, VAE Resample Mode (MI300X) +# +# FSDP2 (ZeRO-2) BF16 training with the local spec provider and +# vae_latent_mode: resample. +# +# In resample mode, latents are re-drawn from stored mean+logvar via +# reparameterization (mean + exp(0.5*logvar) * randn) at every training step. +# This introduces per-step stochasticity in the VAE latents. +# +# Dataset must be an Energon pre-encoded dataset containing mean.pth and +# logvar.pth per sample. + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_fsdp2_local_bf16} + +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + + # Resample mode: re-draw latents from mean+logvar each step + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboLocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # ENERGON DATASET CONFIGURATION + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # Batch configuration + micro_batch_size: 32 + global_batch_size: 256 + seq_length: 512 # 256 img tokens + 256 text tokens (schnell) + + # Mixed precision + bf16: true + fp16: false + + # Optimizer settings + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader settings + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # PyTorch FSDP2 Configuration (ZeRO-2) + # ========================================== + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: true + use_megatron_fsdp: false + + torch_fsdp2_reshard_after_forward: false # ZeRO-2 + + use_fsdp2_fp32_param_optimizer: true + + ckpt_format: torch_dist + + use_distributed_optimizer: false + + overlap_grad_reduce: false + overlap_param_gather: false + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: false + + gradient_accumulation_fusion: false + + # ========================================== + # Memory Optimizations + # ========================================== + + use_flash_attn: true + + fp8: null + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_fsdp2_local_bf16 + wandb_project: flux_12b_fsdp2_local_bf16 + log_throughput: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # PyTorch Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo Configuration + enable_primus_turbo: true + use_turbo_attention: true + + seed: 2025 + + # Torch Compile + torch_compile: + enable: true + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: true diff --git a/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml b/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml new file mode 100644 index 000000000..64c2a6d1d --- /dev/null +++ b/examples/megatron/configs/MI300X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml @@ -0,0 +1,190 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — FSDP2 + Local Spec + FP8 + FP8 All-Gather (MI300X) +# +# FP8 training on the FSDP2 (ZeRO-3) path with: +# - FP8 training (tensorwise, local spec, dual FP8 output projection) +# - FP8 all-gather (keep weight in FP8 after FSDP2 all-gather) +# - BF16 master weight optimizer (BF16 params, FP32 optimizer states) +# - torch.compile enabled + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_fsdp2_local_fp8} + +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + + # Resample mode: re-draw latents from mean+logvar each step + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboLocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # ENERGON DATASET CONFIGURATION + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # Batch configuration + micro_batch_size: 52 + global_batch_size: 416 + seq_length: 512 # 256 img tokens + 256 text tokens (schnell) + + # Mixed precision + bf16: true + fp16: false + + # Optimizer settings + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader settings + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # PyTorch FSDP2 Configuration (ZeRO-2) + # ========================================== + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: true + use_megatron_fsdp: false + + torch_fsdp2_reshard_after_forward: true # ZeRO-3 + use_fsdp2_fp8_all_gather: true + # fp8_all_gather_deq_requant: true # Dequant FP8->BF16 after AG, fresh dynamic requant downstream + # use_triton_ops: true # Use Triton @triton_op modulate/LN+modulate (compile-transparent) + fsdp_prefetch_depth: 1 + fp8_precompute_data_cache: false + optimizer_foreach: false + use_cpp_fp8_quantize: true + + # Optimizer mode: BF16 params + FP32 master weights + use_fsdp2_fp32_param_optimizer: false + use_fsdp2_bf16_master_weight_optimizer: true + + ckpt_format: torch_dist + + use_distributed_optimizer: false + + overlap_grad_reduce: false + overlap_param_gather: false + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: false + + gradient_accumulation_fusion: false + + # ========================================== + # FP8 Configuration — Tensorwise via Primus Turbo + # ========================================== + + use_flash_attn: true + + fp8: "e4m3" + # Dynamic (tensorwise) scaling: the FSDP2 path does not yet exercise the + # delayed amax allreduce, so use tensorwise here. Switch to + # `fp8_recipe: "delayed"` once FSDP2 + delayed is wired up. + fp8_recipe: "tensorwise" + fp8_wgrad: true + fp8_dot_product_attention: false + fp8_multi_head_attention: false + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_fsdp2_local_fp8 + wandb_project: flux_12b_fsdp2_local_fp8 + log_throughput: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # PyTorch Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo Configuration + enable_primus_turbo: true + use_turbo_attention: true + use_dual_fp8_output_projection: true + + seed: 2025 + + # Torch Compile — compatible with Float8 local spec (per-module FP8) + torch_compile: + enable: true + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: true diff --git a/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml b/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml new file mode 100644 index 000000000..865cfe832 --- /dev/null +++ b/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml @@ -0,0 +1,170 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 535M Pretraining Configuration (Pre-encoded Data Mode) +# +# This config demonstrates Flux 535M training with pre-encoded features. +# Pre-encoded mode is faster and recommended for production training. +# +# Usage: +# EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \ +# bash examples/run_pretrain.sh + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_535m_pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_535m.yaml + + # Trainer class for diffusion models + trainer_class: FluxPretrainTrainer + + overrides: + # ============================================================================ + # Model Configuration + # ============================================================================ + + model_type: diffusion_model + + # ============================================================================ + # Dataset Configuration — Synthetic / Mock Data (default) + # ============================================================================ + # + # This 535M config defaults to in-memory synthetic data so it runs + # standalone for sanity checks, CI, and training-pipeline validation + # without a prepared Energon dataset. Samples are random tensors with + # Flux-correct shapes (latent_channels=16, T5 seq=512, CLIP pooled=768). + # + # To train on real data instead: set `mock_data: false` and point + # `data_path` at a prepared Energon dataset (see notes at the bottom). + mock_data: true + mock_dataset: + class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset" + params: + num_samples: 256 # In-memory synthetic samples (iterated cyclically) + image_size: 256 # 256 -> 32x32 latents (light/fast for testing) + # data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Batch configuration + micro_batch_size: 2 # Per-GPU batch size (adjust based on VRAM) + global_batch_size: 16 # Total batch size across all GPUs + seq_length: 4096 # Sequence length for latent features + + # DataLoader settings (used only with real Energon data, i.e. mock_data: false) + num_workers: 4 # Number of data loading workers per GPU + dataloader_type: external # Required for Energon dataloaders + + # ============================================================================ + # Training Parameters + # ============================================================================ + + # Total training steps + train_iters: 100000 # Total training iterations + eval_interval: 1000 # Evaluate every N steps + eval_iters: 50 # Number of evaluation iterations + + # Logging + log_interval: 10 # Log every N steps + tensorboard_dir: output/tensorboard/flux_535m_pretrain + log_throughput: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + + # Checkpointing (disabled) + save_interval: 5000 # Save checkpoint every N steps + save: null + load: null # Path to checkpoint for resuming (optional) + finetune: false + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Precision + bf16: true # Use bfloat16 (recommended for AMD MI300X) + fp16: false + + # ============================================================================ + # Optimizer Configuration + # ============================================================================ + + # Optimizer + optimizer: adam + lr: 1.0e-4 # Peak learning rate + min_lr: 1.0e-5 # Minimum learning rate + weight_decay: 0.01 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-8 + + # Gradient clipping + clip_grad: 1.0 + + # ============================================================================ + # Learning Rate Scheduler + # ============================================================================ + + lr_decay_style: cosine + lr_warmup_iters: 1000 # Warmup iterations + lr_decay_iters: 100000 # Total decay steps (typically = train_iters) + + # ============================================================================ + # Distributed Training Configuration + # ============================================================================ + + # Parallelism settings (Flux 535M fits on 1 GPU, but can use DP for speed) + tensor_model_parallel_size: 1 # Tensor parallelism (no need for 535M) + pipeline_model_parallel_size: 1 # Pipeline parallelism + + # Advanced settings + overlap_grad_reduce: true # Overlap gradient communication + use_flash_attn: true # Use Flash Attention 2 + distributed_timeout_minutes: 60 + + # ============================================================================ + # Seed and Reproducibility + # ============================================================================ + + seed: 42 + + # ============================================================================ + # Monitoring and Logging + # ============================================================================ + + wandb_project: flux_535m_pretrain + wandb_exp_name: flux_535m_preencoded + +# ============================================================================ +# Notes +# ============================================================================ +# +# Dataset Preparation: +# 1. Prepare pre-encoded dataset: +# tools/docker/primus data diffusion-encoded \ +# --source-type directory --input-dir /data/raw \ +# --output-dir /data/encoded --model-path black-forest-labs/FLUX.1-dev +# 2. Copy dataset template to output directory: +# cp primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml \ +# /data/encoded/dataset.yaml +# 3. Run Energon indexing: +# energon prepare /data/encoded --num-workers 8 +# 4. Update data_path above to: /data/encoded/dataset.yaml +# +# For more information, see: +# - primus/configs/data/megatron/diffusion/README.md +# - examples/megatron/diffusion/README.md +# +# Single-node training (8 GPUs): +# EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \ +# GPUS_PER_NODE=8 bash examples/run_pretrain.sh +# +# Multi-node training (4 nodes, 8 GPUs each): +# EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \ +# NNODES=4 bash examples/run_slurm_pretrain.sh +# diff --git a/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml b/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml new file mode 100644 index 000000000..83eda2b58 --- /dev/null +++ b/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml @@ -0,0 +1,199 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 535M FP8 Pretraining Configuration (Testing/Development) +# +# This config is for testing FP8 functionality with the minimal Flux 535M model +# before scaling to the full 12B model. Use this to validate: +# - FP8 setup and configuration +# - Numerical stability +# - Memory and speed improvements +# - Transformer Engine compatibility +# +# Target Hardware: Single AMD MI300X GPU with ROCm 6.0+ +# Requires: Transformer Engine 2.1.0+ with ROCm backend +# +# Usage: +# EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml \ +# GPUS_PER_NODE=1 bash examples/run_pretrain.sh + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_535m_pretrain_fp8} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_535m_fp8.yaml # Use FP8-enabled 535M config + + # Trainer class for diffusion models + trainer_class: FluxPretrainTrainer + + overrides: + # ============================================================================ + # Model Configuration + # ============================================================================ + + model_type: diffusion_model + + # ============================================================================ + # Dataset Configuration — Synthetic / Mock Data (default) + # ============================================================================ + # + # This 535M config defaults to in-memory synthetic data so it runs + # standalone for sanity checks, CI, and training-pipeline validation + # without a prepared Energon dataset. Samples are random tensors with + # Flux-correct shapes (latent_channels=16, T5 seq=512, CLIP pooled=768). + # + # To train on real data instead: set `mock_data: false` and point + # `data_path` at a prepared Energon dataset (see notes at the bottom). + mock_data: true + mock_dataset: + class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset" + params: + num_samples: 256 # In-memory synthetic samples (iterated cyclically) + image_size: 256 # 256 -> 32x32 latents (light/fast for testing) + # data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Batch configuration (can use larger batch with FP8) + micro_batch_size: 4 # Can increase from 2 to 4 with FP8 on 535M + global_batch_size: 32 # Small batch for quick testing + seq_length: 4096 # Sequence length for latent features + + # DataLoader settings (used only with real Energon data, i.e. mock_data: false) + num_workers: 4 # Number of data loading workers per GPU + dataloader_type: external # Required for Energon dataloaders + + # ============================================================================ + # Training Parameters (Quick Testing) + # ============================================================================ + + # Short training for validation + train_iters: 1000 # Just 1K steps for FP8 validation + eval_interval: 100 # Evaluate every 100 steps + eval_iters: 10 # Quick evaluation + + # Logging + log_interval: 10 # Log every N steps + tensorboard_dir: output/tensorboard/flux_535m_pretrain_fp8 + log_throughput: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + + # Checkpointing (disabled) + save_interval: 500 # Save more frequently for testing + save: null + load: null # Path to checkpoint for resuming (optional) + finetune: false + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Precision + bf16: true # Use bfloat16 for non-FP8 ops + fp16: false + + # ============================================================================ + # Optimizer Configuration + # ============================================================================ + + # Optimizer + optimizer: adam + lr: 1.0e-4 # Peak learning rate + min_lr: 1.0e-5 # Minimum learning rate + weight_decay: 0.01 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-8 + + # Gradient clipping (important for FP8 stability) + clip_grad: 1.0 + + # ============================================================================ + # Learning Rate Scheduler + # ============================================================================ + + lr_decay_style: cosine + lr_warmup_iters: 100 # Short warmup for testing + lr_decay_iters: 1000 # Match train_iters + + # ============================================================================ + # Distributed Training Configuration + # ============================================================================ + + # Single GPU configuration + tensor_model_parallel_size: 1 # No TP needed for 535M + pipeline_model_parallel_size: 1 # No PP needed + context_parallel_size: 1 # No CP needed + + # Distributed settings + use_distributed_optimizer: false # Not needed for single GPU + overlap_grad_reduce: false # Not applicable for single GPU + use_flash_attn: true # Use Flash Attention 2 + distributed_timeout_minutes: 60 + + # ============================================================================ + # Memory Optimization + # ============================================================================ + + # Activation checkpointing (not needed for 535M with FP8) + recompute_granularity: null # No recompute needed + recompute_method: null + recompute_num_layers: null + + # Sequence parallelism + sequence_parallel: false # Not needed for single GPU + + # ============================================================================ + # Seed and Reproducibility + # ============================================================================ + + seed: 42 + + # ============================================================================ + # Monitoring and Logging + # ============================================================================ + + wandb_project: flux_535m_pretrain_fp8 + wandb_exp_name: flux_535m_fp8_test + +# ============================================================================ +# Notes - FP8 Validation with 535M +# ============================================================================ +# +# Hardware Requirements (with FP8): +# - Minimum: 1× MI300X 192GB +# - Memory per GPU: ~3-5GB (vs ~7-10GB BF16) +# - Training time: Minutes +# +# Validation Checklist: +# [ ] Setup FP8 environment (see docs/backends/megatron/diffusion/fp8_training.md) +# [ ] Verify TE FP8 support is available +# [ ] Run this config to validate FP8 training +# [ ] Check logs for NaN/Inf (should be none) +# [ ] Verify memory usage is ~50% of BF16 +# [ ] Verify training speed is 1.5-2x faster than BF16 +# [ ] Check loss decreases normally +# +# Expected Results: +# - Training completes 1000 steps in 5-15 minutes +# - No NaN/Inf in losses +# - Memory usage: ~3-5GB +# - Speed: ~10-50 steps/sec (depending on hardware) +# - Loss should decrease normally +# +# If validation passes, proceed to one of the 12B FP8 configs, e.g.: +# - flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml (TransformerEngine FP8) +# - flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml (local-spec FP8) +# +# Troubleshooting: +# - If NaN/Inf: Check Transformer Engine FP8 support +# - If OOM: Reduce micro_batch_size +# - If slow: Verify ROCm FP8 tensor cores are being used +# - If unstable: Try fp8_wgrad: false in model config +# +# For more information: See docs/backends/megatron/diffusion/fp8_training.md diff --git a/examples/megatron/configs/MI300X/diffusion/flux_535m_with_guidance_embed.yaml b/examples/megatron/configs/MI300X/diffusion/flux_535m_with_guidance_embed.yaml new file mode 100644 index 000000000..38f10f88d --- /dev/null +++ b/examples/megatron/configs/MI300X/diffusion/flux_535m_with_guidance_embed.yaml @@ -0,0 +1,58 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 535M with Guidance Embedding (Advanced Configuration) +# +# This config demonstrates Flux training with guidance embedding enabled. +# This is an OPTIONAL advanced feature that allows for faster single-pass CFG +# during inference, but requires training with guidance embedding enabled. +# +# IMPORTANT: Most users should use the standard flux_535m_pretrain.yaml config. +# Only use this if you specifically need guidance embedding support. +# +# Usage: +# EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_with_guidance_embed.yaml \ +# bash examples/run_pretrain.sh + +# Extend standard 535M config +extends: + - flux_535m_pretrain.yaml + +modules: + pre_trainer: + overrides: + # ============================================================================ + # Guidance Embedding Configuration (ADVANCED) + # ============================================================================ + + # Enable guidance embedding for single-pass CFG + # This adds a learned MLPEmbedder layer that conditions on guidance scale + guidance_embed: true + + # Guidance scale used during training + # Model learns to adapt its predictions based on this scale + guidance_scale: 3.5 + + # ============================================================================ + # Notes + # ============================================================================ + # + # Training with guidance embedding: + # - Adds ~1-2% more parameters (guidance MLPEmbedder) + # - Allows single-pass CFG during inference (faster) + # - Requires more training data/iterations to converge + # - Model learns guidance as a conditioning signal + # + # Inference with guidance embedding: + # - Pipeline automatically detects guidance_embed layer + # - Uses single forward pass instead of batch doubling + # - ~2x faster CFG compared to explicit CFG + # - Guidance scale can be varied at inference time + # + # Standard approach (guidance_embed: false): + # - Default for most Primus training + # - Uses explicit CFG (batch doubling) at inference + # - More compatible with existing checkpoints + # - Slightly slower but more flexible + # + # See examples/megatron/diffusion/README.md for more details. diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml new file mode 100644 index 000000000..8ba4c13fb --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml @@ -0,0 +1,209 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — DDP + Local Spec + FP8 Tensorwise (MI355X) +# +# Combines Megatron DDP + distributed optimizer with the local spec provider +# (PrimusTurboFloat8LocalSpecProvider) for FP8 tensorwise training. +# +# Key configuration: +# - Megatron DDP with overlap_grad_reduce + overlap_param_gather +# - PrimusTurboFloat8LocalSpecProvider (NO TransformerEngine dependency) +# - FP8 hybrid tensorwise (per-module FP8 via Primus Turbo) +# - Primus Turbo attention +# - torch.compile enabled (per_block strategy, compatible with local spec + overlap) +# - Energon pre-encoded dataset with stored VAE mean/logvar (resample mode) + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_fp8} + +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + # Short LR warmup ramp over the first optimizer steps for FP8 stability. + nemo_aligned_lr_warmup: true + warmup_train_steps: 2 + + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + + # Resample mode: re-draw latents from mean+logvar each step + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboFloat8LocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # DATASET CONFIGURATION + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # Batch configuration + micro_batch_size: 64 + global_batch_size: 512 + seq_length: 512 # 256 img tokens + 256 text tokens (schnell) + + # ========================================== + # BF16 + FP8 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # Optimizer settings + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader settings + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # Megatron DDP + Distributed Optimizer + # ========================================== + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: false + use_megatron_fsdp: false + + use_distributed_optimizer: true + overlap_grad_reduce: true + overlap_param_gather: true + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: true + main_params_dtype: fp32 + main_grads_dtype: bf16 + exp_avg_dtype: bf16 + exp_avg_sq_dtype: bf16 + + ddp_bucket_size: 256000000 + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + gradient_accumulation_fusion: false + + check_for_nan_in_loss_and_grad: true + + # ========================================== + # FP8 Configuration — Tensorwise + Delayed FP8 via Primus Turbo + # ========================================== + + use_flash_attn: true + + fp8: "hybrid" + # TE-compatible unified form: fp8: hybrid + fp8_recipe: delayed selects + # tensorwise FP8 with delayed scaling. Resolution path: + # primus/backends/megatron/core/extensions/primus_turbo_float8_local.py + # :: Float8{Column,Row}ParallelLinear._use_delayed_scaling. + fp8_recipe: "delayed" + fp8_margin: 0 + fp8_amax_history_len: 1024 + fp8_amax_compute_algo: "max" + fp8_wgrad: true + fp8_dot_product_attention: false + fp8_multi_head_attention: false + fp8_reduce_amax: true + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 180 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_ddp_local_fp8 + wandb_project: flux_12b_ddp_local_fp8 + log_throughput: true + wall_clock_step_timer: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # PyTorch Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo Configuration + enable_primus_turbo: true + use_turbo_attention: true + use_dual_fp8_output_projection: false + + seed: 2025 + per_step_rng_reseed: false + nemo_chimera_init: false + + # Torch Compile — compatible with Float8 local spec (per-module FP8) + torch_compile: + enable: true + strategy: "per_block" + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false + emulate_precision_casts: false + fused_ln_modulate: true diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml new file mode 100644 index 000000000..d67662f76 --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml @@ -0,0 +1,202 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — DDP + Local Spec + FP8 Tensorwise (MLPerf Mode) +# +# NOTE: This is a benchmark-reproduction config for the MLPerf Training Flux.1 +# benchmark. It mirrors MLPerf logging/convergence conventions and is intended +# for reproducing benchmark results, not as a general-purpose training starting +# point. For everyday FP8 training use +# flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml. +# +# MLPerf-compliant variant of the DDP + local spec FP8 config: +# - mlperf_mode: true — enables MLPerf logging via mlperf_logging.mllog +# - warmup_train_steps: 2 — synthetic data warmup for torch.compile + FP8 +# - target_val_loss: 0.586 — convergence target for early stopping +# - DDP + distributed optimizer without precision-aware optimizer +# - PrimusTurboFloat8LocalSpecProvider with FP8 tensorwise +# - Suppresses TensorBoard/WandB/print_rank_last during training +# - Emits structured MLPerf events (INIT_START, RUN_START, EVAL, etc.) + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_fp8_mlperf} + +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # MLPerf Mode + # ========================================== + mlperf_mode: true + warmup_train_steps: 2 + target_val_loss: 0.586 + + # ========================================== + # MLPerf Training v5.1 Alignment + # ========================================== + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboFloat8LocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # DATA — Real Energon data (not mock) + # ========================================== + mock_data: false + dataloader_type: external + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + num_workers: 16 + prefetch_factor: 4 + max_samples_per_sequence: null + + # Training iterations + train_iters: 5000 + eval_interval: 512 + eval_iters: 10 + log_interval: 10 + save_interval: 10000 + + # Batch configuration + micro_batch_size: 64 + global_batch_size: 512 + seq_length: 512 + + # ========================================== + # BF16 + FP8 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # Optimizer settings (MLPerf v5.1) + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # ========================================== + # Megatron DDP + Distributed Optimizer + # ========================================== + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: false + use_megatron_fsdp: false + + use_distributed_optimizer: true + overlap_grad_reduce: true + overlap_param_gather: true + overlap_param_gather_with_optimizer_step: false + # Use FP32 optimizer states (m,v): strictly better time-to-train with + # no throughput trade-off vs BF16 on this FP8 + emulate_precision_casts recipe. + use_precision_aware_optimizer: false + main_params_dtype: fp32 + main_grads_dtype: fp32 + exp_avg_dtype: fp32 + exp_avg_sq_dtype: fp32 + + ddp_bucket_size: 256000000 + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + gradient_accumulation_fusion: false + check_for_nan_in_loss_and_grad: false + + # ========================================== + # FP8 — Tensorwise via Primus Turbo + # ========================================== + use_flash_attn: true + + fp8: "hybrid" + # TE-compatible unified form: fp8: hybrid + fp8_recipe: delayed selects + # tensorwise FP8 with delayed scaling. + fp8_recipe: "delayed" + fp8_wgrad: true + fp8_dot_product_attention: false + fp8_multi_head_attention: false + fp8_amax_history_len: 1 + fp8_amax_compute_algo: "most_recent" + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled for MLPerf) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging — suppressed by mlperf_mode + tensorboard_dir: null + wandb_project: null + log_throughput: false + wall_clock_step_timer: true + log_timers_to_tensorboard: false + log_batch_size_to_tensorboard: false + log_learning_rate_to_tensorboard: false + log_memory_to_tensorboard: false + + # Profiler — disabled for MLPerf + profile: false + + # Primus Turbo + enable_primus_turbo: true + use_turbo_attention: true + use_dual_fp8_output_projection: false + + seed: 42 + # MLPerf-aligned per-step CUDA RNG reseed (defaults off elsewhere; MLPerf + # reproduction must opt in for run-to-run determinism). + per_step_rng_reseed: true + + # Torch Compile — per_block: compile each transformer layer individually + torch_compile: + enable: true + strategy: "per_block" + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false + emulate_precision_casts: true + fused_ln_modulate: true diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml new file mode 100644 index 000000000..10283f997 --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml @@ -0,0 +1,198 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — DDP + Local Spec + MXFP4 (MI355X) +# +# Combines Megatron DDP + distributed optimizer with the local spec provider +# (PrimusTurboMXFP4LocalSpecProvider) for MXFP4 block-scaled training. +# +# Key configuration: +# - Megatron DDP with overlap_grad_reduce + overlap_param_gather +# - PrimusTurboMXFP4LocalSpecProvider (NO TransformerEngine dependency) +# - MXFP4 (E2M1 + E8M0 block-of-32 scales) via Primus Turbo + AITER +# - Primus Turbo attention +# - torch.compile enabled (per_block strategy, compatible with local spec) +# - Energon pre-encoded dataset with stored VAE mean/logvar (resample mode) + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_local_mxfp4} + +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + # Short LR warmup ramp over the first optimizer steps for stability. + nemo_aligned_lr_warmup: true + warmup_train_steps: 2 + + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + + # Resample mode: re-draw latents from mean+logvar each step + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboMXFP4LocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # DATASET CONFIGURATION + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # Batch configuration + micro_batch_size: 64 + global_batch_size: 512 + seq_length: 512 # 256 img tokens + 256 text tokens (schnell) + + # ========================================== + # BF16 + MXFP4 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # Optimizer settings + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader settings + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # Megatron DDP + Distributed Optimizer + # ========================================== + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: false + use_megatron_fsdp: false + + use_distributed_optimizer: true + overlap_grad_reduce: true + overlap_param_gather: true + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: true + main_params_dtype: fp32 + main_grads_dtype: bf16 + exp_avg_dtype: bf16 + exp_avg_sq_dtype: bf16 + + ddp_bucket_size: 256000000 + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + gradient_accumulation_fusion: false + + check_for_nan_in_loss_and_grad: true + + # ========================================== + # MXFP4 Configuration — Block-scaled via Primus Turbo + AITER + # ========================================== + + use_flash_attn: true + + fp4: "mxfp4" + fp4_recipe: "mxfp4" + mxfp4_backward_precision: "mxfp4" # "mxfp4" (pure) or "fp8" (hybrid) + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 180 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_ddp_local_mxfp4 + wandb_project: flux_12b_ddp_local_mxfp4 + log_throughput: true + wall_clock_step_timer: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # PyTorch Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo Configuration + enable_primus_turbo: true + use_turbo_attention: true + + seed: 2025 + per_step_rng_reseed: false + nemo_chimera_init: false + + # Torch Compile — compatible with MXFP4 local spec (per-module FP4) + torch_compile: + enable: true + strategy: "per_block" + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false + emulate_precision_casts: false + fused_ln_modulate: true diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml new file mode 100644 index 000000000..d4bddef34 --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec.yaml @@ -0,0 +1,184 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — DDP + TransformerEngine Spec + BF16 (MI355X) +# +# BF16 training using Megatron DDP + distributed optimizer with +# TransformerEngine modules (TEColumnParallelLinear / TERowParallelLinear / +# TEDotProductAttention / TENorm). This is the BF16 baseline on the +# TransformerEngine path; see the *_te_spec_fp8 variant for FP8. +# +# Key settings: +# bf16: true +# params_dtype: bfloat16 +# micro_batch_size: 64 / global_batch_size: 512 + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_bf16} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b_rope_fusion.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + + # Resample mode: re-draw latents from mean+logvar each step + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # TransformerEngine Spec (default) + # ========================================== + transformer_impl: "transformer_engine" + + # ========================================== + # ENERGON DATASET CONFIGURATION + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # Batch configuration + micro_batch_size: 64 + global_batch_size: 512 + seq_length: 512 + + # ========================================== + # BF16 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # Optimizer settings (identical to BF16 baseline) + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader settings + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # Megatron DDP + Distributed Optimizer + # ========================================== + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: false + use_megatron_fsdp: false + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + use_distributed_optimizer: true + + overlap_grad_reduce: true + overlap_param_gather: true + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: false + + ddp_bucket_size: 256000000 + + gradient_accumulation_fusion: false + + # ========================================== + # Memory Optimizations + # ========================================== + + use_flash_attn: true + + fp8: null + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_ddp_te_bf16 + wandb_project: flux_12b_ddp_te_bf16 + log_throughput: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # PyTorch Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo Configuration + enable_primus_turbo: false + use_turbo_attention: false + + seed: 2025 + te_rng_tracker: true + + # Torch Compile — stack strategy (compiles the double/single DiT block + # stacks as inductor regions). Required to fit micro_batch_size 64 in + # BF16 on the TE path; matches the MLPerf NeMo reference (COMPILE_DIT + # strategy=stack). + torch_compile: + enable: true + strategy: "stack" + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml new file mode 100644 index 000000000..4cb6ad043 --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml @@ -0,0 +1,209 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — DDP + TE Spec + FP8 Delayed Scaling (MI355X) +# +# FP8 training using Megatron DDP + distributed optimizer with TransformerEngine +# modules and FP8 hybrid delayed scaling: +# - Megatron DDP with overlap_grad_reduce + overlap_param_gather +# - TEColumnParallelLinear / TERowParallelLinear / TEDotProductAttention / TENorm +# - FP8 hybrid (E4M3 fwd, E5M2 bwd) with delayed scaling (amax history 1024) +# - RoPE fusion via apply_rope_fusion: true +# - Energon pre-encoded dataset +# +# Required environment variables for the TransformerEngine path (set before launch): +# export NVTE_FUSED_ATTN=1 +# export NVTE_FUSED_ATTN_CK=1 +# export NVTE_FP8_DPA_BWD=1 +# export NVTE_USE_HIPBLASLT=1 +# export USE_HIPBLASLT=1 +# export TORCH_BLAS_PREFER_HIPBLASLT=1 +# export NVTE_USE_CAST_TRANSPOSE_TRITON=1 + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_fp8} + +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + vae_scale: 0.3611 + vae_shift: 0.1159 + vae_latent_mode: resample + + # ========================================== + # RoPE Fusion + # ========================================== + rotary_interleaved: true + apply_rope_fusion: true + position_embedding_type: rope + + # ========================================== + # TransformerEngine Spec + # ========================================== + transformer_impl: "transformer_engine" + + # ========================================== + # FP8 — hybrid delayed scaling + # ========================================== + fp8: "hybrid" + fp8_recipe: "delayed" + fp8_margin: 0 + fp8_amax_history_len: 1024 + fp8_amax_compute_algo: "max" + fp8_wgrad: true + fp8_dot_product_attention: false + fp8_multi_head_attention: false + + # ========================================== + # Energon Dataset + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # ========================================== + # Batch Configuration + # ========================================== + # MI355X (256GB HBM3) fits the full MBS=64/GBS=512; tune to your hardware. + micro_batch_size: 64 + global_batch_size: 512 + seq_length: 512 + + # ========================================== + # BF16 + FP8 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # ========================================== + # Optimizer + # ========================================== + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # Megatron DDP + Distributed Optimizer + # ========================================== + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: false + use_megatron_fsdp: false + + use_distributed_optimizer: true + overlap_grad_reduce: true + overlap_param_gather: true + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: false + + ddp_bucket_size: 256000000 + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + gradient_accumulation_fusion: false + + check_for_nan_in_loss_and_grad: true + + # ========================================== + # Memory / Misc + # ========================================== + use_flash_attn: true + empty_unused_memory_level: 0 + + # Manual GC — align GC timing across ranks to avoid stragglers + manual_gc: true + manual_gc_interval: 1000 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_ddp_te_fp8 + wandb_project: flux_12b_ddp_te_fp8 + log_throughput: true + wall_clock_step_timer: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo — disabled for pure TE path + enable_primus_turbo: false + use_turbo_attention: false + + seed: 2025 + te_rng_tracker: true + + # torch.compile — selective stack compilation for TE spec + torch_compile: + enable: true + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false + strategy: "stack" + replace_qk_rmsnorm: true + disable_inductor_cudagraphs: false diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml new file mode 100644 index 000000000..cd01468c8 --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_te_spec_fp8_mlperf.yaml @@ -0,0 +1,202 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — DDP + TE Spec + FP8 Delayed Scaling (MLPerf Mode) +# +# NOTE: This is a benchmark-reproduction config for the MLPerf Training Flux.1 +# benchmark. It mirrors MLPerf logging/convergence conventions and is intended +# for reproducing benchmark results, not as a general-purpose training starting +# point. For everyday FP8 training on the TransformerEngine path use +# flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml. +# +# MLPerf-compliant variant of the TE spec FP8 config: +# - mlperf_mode: true — enables MLPerf logging via mlperf_logging.mllog +# - warmup_train_steps: 2 — synthetic data warmup for torch.compile + FP8 +# - target_val_loss: 0.586 — convergence target for early stopping +# - DDP + distributed optimizer +# - TransformerEngine modules with FP8 hybrid delayed scaling + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_ddp_te_fp8_mlperf} + +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # MLPerf Mode + # ========================================== + mlperf_mode: true + warmup_train_steps: 2 + target_val_loss: 0.586 + + # ========================================== + # MLPerf Training v5.1 Alignment + # ========================================== + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + vae_scale: 0.3611 + vae_shift: 0.1159 + vae_latent_mode: resample + + # ========================================== + # RoPE Fusion + # ========================================== + rotary_interleaved: true + apply_rope_fusion: true + position_embedding_type: rope + + # ========================================== + # TransformerEngine Spec + # ========================================== + transformer_impl: "transformer_engine" + adaln_plain_ops: true + adaln_always_jit_fuser: true + + # ========================================== + # FP8 — hybrid delayed scaling + # ========================================== + fp8: "hybrid" + fp8_recipe: "delayed" + fp8_margin: 0 + fp8_amax_history_len: 1024 + fp8_amax_compute_algo: "max" + fp8_wgrad: true + fp8_dot_product_attention: false + fp8_multi_head_attention: false + + # ========================================== + # DATA — Real Energon data (not mock) + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + mock_data: false + dataloader_type: external + num_workers: 8 + max_samples_per_sequence: null + + train_iters: 5000 + eval_interval: 512 + eval_iters: 10 + log_interval: 10 + save_interval: 10000 + + # ========================================== + # Batch Configuration + # ========================================== + micro_batch_size: 64 + global_batch_size: 512 + seq_length: 512 + + # ========================================== + # BF16 + FP8 Precision + # ========================================== + bf16: true + fp16: false + params_dtype: bfloat16 + grad_reduce_in_bf16: true + + # ========================================== + # Optimizer (MLPerf v5.1) + # ========================================== + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # ========================================== + # Megatron DDP + Distributed Optimizer + # ========================================== + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: false + use_megatron_fsdp: false + + use_distributed_optimizer: true + overlap_grad_reduce: true + overlap_param_gather: true + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: false + + ddp_bucket_size: 256000000 + use_fsdp2_fp32_param_optimizer: false + + ckpt_format: torch_dist + + gradient_accumulation_fusion: false + check_for_nan_in_loss_and_grad: false + + # ========================================== + # Memory / Misc + # ========================================== + use_flash_attn: true + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled for MLPerf) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging — suppressed by mlperf_mode + tensorboard_dir: null + wandb_project: null + log_throughput: false + wall_clock_step_timer: true + log_timers_to_tensorboard: false + log_batch_size_to_tensorboard: false + log_learning_rate_to_tensorboard: false + log_memory_to_tensorboard: false + + # Profiler — disabled + profile: false + + # Primus Turbo — disabled for pure TE path + enable_primus_turbo: false + use_turbo_attention: false + + seed: 2025 + te_rng_tracker: true + # MLPerf-aligned per-step CUDA RNG reseed (defaults off elsewhere; MLPerf + # reproduction must opt in for run-to-run determinism). + per_step_rng_reseed: true + + # torch.compile — selective stack compilation for TE spec + torch_compile: + enable: true + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: false + strategy: "stack" + replace_qk_rmsnorm: true + disable_inductor_cudagraphs: true + emulate_precision_casts: false diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml new file mode 100644 index 000000000..567b58aeb --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec.yaml @@ -0,0 +1,175 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — FSDP2 + Local Spec + BF16, VAE Resample Mode (MI355X) +# +# FSDP2 (ZeRO-2) BF16 training with the local spec provider and +# vae_latent_mode: resample. +# +# In resample mode, latents are re-drawn from stored mean+logvar via +# reparameterization (mean + exp(0.5*logvar) * randn) at every training step. +# This introduces per-step stochasticity in the VAE latents. +# +# Dataset must be an Energon pre-encoded dataset containing mean.pth and +# logvar.pth per sample. + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_fsdp2_local_bf16} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + + # Resample mode: re-draw latents from mean+logvar each step + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboLocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # ENERGON DATASET CONFIGURATION + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 2000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # Batch configuration + micro_batch_size: 64 + global_batch_size: 512 # 64 * 8 GPUs = 512 + seq_length: 512 # 256 img tokens + 256 text tokens (schnell) + + # Mixed precision + bf16: true + fp16: false + + # Optimizer settings + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader settings + num_workers: 4 + dataloader_type: external + + # ========================================== + # PyTorch FSDP2 Configuration (ZeRO-2) + # ========================================== + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: true + use_megatron_fsdp: false + + torch_fsdp2_reshard_after_forward: false # ZeRO-2 + + use_fsdp2_fp32_param_optimizer: true + + ckpt_format: torch_dist + + use_distributed_optimizer: false + + overlap_grad_reduce: false + overlap_param_gather: false + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: false + + gradient_accumulation_fusion: false + + # ========================================== + # Memory Optimizations + # ========================================== + + use_flash_attn: true + + fp8: null + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_fsdp2_local_bf16 + wandb_project: flux_12b_fsdp2_local_bf16 + log_throughput: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # PyTorch Profiler (disabled) + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo Configuration + enable_primus_turbo: true + use_turbo_attention: true + + seed: 42 + + # Torch Compile + torch_compile: + enable: true + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: true diff --git a/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml b/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml new file mode 100644 index 000000000..1ee15a5fb --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml @@ -0,0 +1,198 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Schnell Training — FSDP2 + Local Spec + FP8 + FP32 Param Optimizer (MI355X) +# +# FP8 training on the FSDP2 path with: +# - FP8 training (tensorwise, local spec) +# - BF16 all-gather (use_fsdp2_fp8_all_gather: false) +# - FP32 param optimizer (FP32 params + FP32 optimizer states) +# - No overlap grad norm (overlap_grad_norm: false) +# - torch.compile enabled +# - ZeRO-2 sharding (reshard_after_forward: false) +# +# Uses an Energon pre-encoded dataset with VAE resample mode. + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_12b_fsdp2_local_fp8} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_12b.yaml + + trainer_class: FluxPretrainTrainer + + overrides: + model_type: flux_schnell + + # ========================================== + # Flux Schnell training recipe + # ========================================== + timestep_sampling_strategy: "direct_uniform" + activation_func: "openai_gelu" + cfg_dropout_prob: 0.1 + + # VAE latent normalization (required for resample mode) + vae_scale: 0.3611 + vae_shift: 0.1159 + + # Resample mode: re-draw latents from mean+logvar each step + vae_latent_mode: resample + + # RoPE: must be interleaved to match Flux's EmbedND doubled-frequency layout + rotary_interleaved: true + + # ========================================== + # PrimusTurboLocalSpecProvider + # ========================================== + transformer_impl: "local" + + # ========================================== + # DATASET CONFIGURATION + # ========================================== + data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Training iterations (example value; increase for a full training run) + train_iters: 1000 + eval_interval: 10000 + eval_iters: 0 + log_interval: 1 + save_interval: 10000 + + # Batch configuration + micro_batch_size: 64 + global_batch_size: 512 + seq_length: 512 # 256 img tokens + 256 text tokens (schnell) + + # Mixed precision + bf16: true + fp16: false + + # Optimizer settings + optimizer: adam + lr: 2.0e-4 + min_lr: 2.0e-4 + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-8 + clip_grad: 1.0 + + # Learning rate scheduler (warmup-hold, no decay) + lr_warmup_iters: 1600 + lr_decay_iters: 4000 + lr_decay_style: constant + + # DataLoader settings + num_workers: 8 + dataloader_type: external + max_samples_per_sequence: null + + # ========================================== + # PyTorch FSDP2 Configuration (ZeRO-2) + # ========================================== + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + context_parallel_size: 1 + + use_torch_fsdp2: true + use_megatron_fsdp: false + + torch_fsdp2_reshard_after_forward: false # ZeRO-2 + use_fsdp2_fp8_all_gather: false + fp8_all_gather_stochastic_rounding: true + # fp8_all_gather_deq_requant: true # Dequant FP8->BF16 after AG, fresh dynamic requant downstream + use_triton_ops: true # Triton @triton_op modulate/LN+modulate — eliminates ~510us dispatch overhead per graph + fsdp_prefetch_depth: 5 + fp8_precompute_data_cache: false + optimizer_foreach: false + use_cpp_fp8_quantize: true + overlap_grad_norm: false + + # Optimizer mode: FP32 params + FP32 optimizer states + use_fsdp2_fp32_param_optimizer: true + use_fsdp2_bf16_master_weight_optimizer: false + + ckpt_format: torch_dist + + use_distributed_optimizer: false + + overlap_grad_reduce: false + overlap_param_gather: false + overlap_param_gather_with_optimizer_step: false + use_precision_aware_optimizer: false + + gradient_accumulation_fusion: false + + # ========================================== + # FP8 Configuration — Tensorwise via Primus Turbo + # ========================================== + + use_flash_attn: true + + fp8: "hybrid" + # Dynamic (tensorwise) scaling: the FSDP2 path does not yet exercise the + # delayed amax allreduce, so use tensorwise here. Switch to + # `fp8_recipe: "delayed"` once FSDP2 + delayed is wired up. + fp8_recipe: "tensorwise" + fp8_wgrad: true + fp8_dot_product_attention: false + fp8_multi_head_attention: false + + empty_unused_memory_level: 0 + + distributed_timeout_minutes: 60 + distributed_backend: nccl + + # Checkpointing (disabled) + finetune: false + save: null + load: null + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Logging + tensorboard_dir: output/tensorboard/flux_12b_fsdp2_local_fp8 + wandb_project: flux_12b_fsdp2_local_fp8 + log_throughput: true + wall_clock_step_timer: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_memory_to_tensorboard: true + + # PyTorch Profiler + profile: false + use_pytorch_profiler: true + profile_ranks: [0] + profile_step_start: 9 + profile_step_end: 11 + torch_profiler_record_shapes: true + torch_profiler_with_stack: true + torch_profiler_use_gzip: false + disable_profiler_activity_cpu: false + + # Primus Turbo Configuration + enable_primus_turbo: true + use_turbo_attention: true + use_dual_fp8_output_projection: false + + seed: 42 + + # Torch Compile — compatible with Float8 local spec (per-module FP8) + torch_compile: + enable: true + backend: "inductor" + mode: "default" + fullgraph: false + compile_optimizer: true + emulate_precision_casts: false + fused_ln_modulate: true diff --git a/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml b/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml new file mode 100644 index 000000000..0ec08c870 --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml @@ -0,0 +1,170 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 535M Pretraining Configuration (Pre-encoded Data Mode) +# +# This config demonstrates Flux 535M training with pre-encoded features. +# Pre-encoded mode is faster and recommended for production training. +# +# Usage: +# EXP=examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml \ +# bash examples/run_pretrain.sh + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_535m_pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_535m.yaml + + # Trainer class for diffusion models + trainer_class: FluxPretrainTrainer + + overrides: + # ============================================================================ + # Model Configuration + # ============================================================================ + + model_type: diffusion_model + + # ============================================================================ + # Dataset Configuration — Synthetic / Mock Data (default) + # ============================================================================ + # + # This 535M config defaults to in-memory synthetic data so it runs + # standalone for sanity checks, CI, and training-pipeline validation + # without a prepared Energon dataset. Samples are random tensors with + # Flux-correct shapes (latent_channels=16, T5 seq=512, CLIP pooled=768). + # + # To train on real data instead: set `mock_data: false` and point + # `data_path` at a prepared Energon dataset (see notes at the bottom). + mock_data: true + mock_dataset: + class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset" + params: + num_samples: 256 # In-memory synthetic samples (iterated cyclically) + image_size: 256 # 256 -> 32x32 latents (light/fast for testing) + # data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Batch configuration + micro_batch_size: 2 # Per-GPU batch size (adjust based on VRAM) + global_batch_size: 16 # Total batch size across all GPUs + seq_length: 4096 # Sequence length for latent features + + # DataLoader settings (used only with real Energon data, i.e. mock_data: false) + num_workers: 4 # Number of data loading workers per GPU + dataloader_type: external # Required for Energon dataloaders + + # ============================================================================ + # Training Parameters + # ============================================================================ + + # Total training steps + train_iters: 100000 # Total training iterations + eval_interval: 1000 # Evaluate every N steps + eval_iters: 50 # Number of evaluation iterations + + # Logging + log_interval: 10 # Log every N steps + tensorboard_dir: output/tensorboard/flux_535m_pretrain + log_throughput: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + + # Checkpointing (disabled) + save_interval: 5000 # Save checkpoint every N steps + save: null + load: null # Path to checkpoint for resuming (optional) + finetune: false + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Precision + bf16: true # Use bfloat16 (recommended for AMD MI355X) + fp16: false + + # ============================================================================ + # Optimizer Configuration + # ============================================================================ + + # Optimizer + optimizer: adam + lr: 1.0e-4 # Peak learning rate + min_lr: 1.0e-5 # Minimum learning rate + weight_decay: 0.01 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-8 + + # Gradient clipping + clip_grad: 1.0 + + # ============================================================================ + # Learning Rate Scheduler + # ============================================================================ + + lr_decay_style: cosine + lr_warmup_iters: 1000 # Warmup iterations + lr_decay_iters: 100000 # Total decay steps (typically = train_iters) + + # ============================================================================ + # Distributed Training Configuration + # ============================================================================ + + # Parallelism settings (Flux 535M fits on 1 GPU, but can use DP for speed) + tensor_model_parallel_size: 1 # Tensor parallelism (no need for 535M) + pipeline_model_parallel_size: 1 # Pipeline parallelism + + # Advanced settings + overlap_grad_reduce: true # Overlap gradient communication + use_flash_attn: true # Use Flash Attention 2 + distributed_timeout_minutes: 60 + + # ============================================================================ + # Seed and Reproducibility + # ============================================================================ + + seed: 42 + + # ============================================================================ + # Monitoring and Logging + # ============================================================================ + + wandb_project: flux_535m_pretrain + wandb_exp_name: flux_535m_preencoded + +# ============================================================================ +# Notes +# ============================================================================ +# +# Dataset Preparation: +# 1. Prepare pre-encoded dataset: +# tools/docker/primus data diffusion-encoded \ +# --source-type directory --input-dir /data/raw \ +# --output-dir /data/encoded --model-path black-forest-labs/FLUX.1-dev +# 2. Copy dataset template to output directory: +# cp primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml \ +# /data/encoded/dataset.yaml +# 3. Run Energon indexing: +# energon prepare /data/encoded --num-workers 8 +# 4. Update data_path above to: /data/encoded/dataset.yaml +# +# For more information, see: +# - primus/configs/data/megatron/diffusion/README.md +# - examples/megatron/diffusion/README.md +# +# Single-node training (8 GPUs): +# EXP=examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml \ +# GPUS_PER_NODE=8 bash examples/run_pretrain.sh +# +# Multi-node training (4 nodes, 8 GPUs each): +# EXP=examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain.yaml \ +# NNODES=4 bash examples/run_slurm_pretrain.sh +# diff --git a/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml b/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml new file mode 100644 index 000000000..5db547f21 --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml @@ -0,0 +1,199 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 535M FP8 Pretraining Configuration (Testing/Development) +# +# This config is for testing FP8 functionality with the minimal Flux 535M model +# before scaling to the full 12B model. Use this to validate: +# - FP8 setup and configuration +# - Numerical stability +# - Memory and speed improvements +# - Transformer Engine compatibility +# +# Target Hardware: Single AMD MI355X GPU with ROCm 6.0+ +# Requires: Transformer Engine 2.1.0+ with ROCm backend +# +# Usage: +# EXP=examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml \ +# GPUS_PER_NODE=1 bash examples/run_pretrain.sh + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:flux_535m_pretrain_fp8} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + model: diffusion/flux_535m_fp8.yaml # Use FP8-enabled 535M config + + # Trainer class for diffusion models + trainer_class: FluxPretrainTrainer + + overrides: + # ============================================================================ + # Model Configuration + # ============================================================================ + + model_type: diffusion_model + + # ============================================================================ + # Dataset Configuration — Synthetic / Mock Data (default) + # ============================================================================ + # + # This 535M config defaults to in-memory synthetic data so it runs + # standalone for sanity checks, CI, and training-pipeline validation + # without a prepared Energon dataset. Samples are random tensors with + # Flux-correct shapes (latent_channels=16, T5 seq=512, CLIP pooled=768). + # + # To train on real data instead: set `mock_data: false` and point + # `data_path` at a prepared Energon dataset (see notes at the bottom). + mock_data: true + mock_dataset: + class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset" + params: + num_samples: 256 # In-memory synthetic samples (iterated cyclically) + image_size: 256 # 256 -> 32x32 latents (light/fast for testing) + # data_path: ${PRIMUS_DIFFUSION_DATA_PATH:/path/to/energon/dataset} + + # Batch configuration (can use larger batch with FP8) + micro_batch_size: 4 # Can increase from 2 to 4 with FP8 on 535M + global_batch_size: 32 # Small batch for quick testing + seq_length: 4096 # Sequence length for latent features + + # DataLoader settings (used only with real Energon data, i.e. mock_data: false) + num_workers: 4 # Number of data loading workers per GPU + dataloader_type: external # Required for Energon dataloaders + + # ============================================================================ + # Training Parameters (Quick Testing) + # ============================================================================ + + # Short training for validation + train_iters: 1000 # Just 1K steps for FP8 validation + eval_interval: 100 # Evaluate every 100 steps + eval_iters: 10 # Quick evaluation + + # Logging + log_interval: 10 # Log every N steps + tensorboard_dir: output/tensorboard/flux_535m_pretrain_fp8 + log_throughput: true + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + + # Checkpointing (disabled) + save_interval: 500 # Save more frequently for testing + save: null + load: null # Path to checkpoint for resuming (optional) + finetune: false + no_save_optim: true + no_save_rng: true + auto_continue_train: false + disable_last_saving: true + + # Precision + bf16: true # Use bfloat16 for non-FP8 ops + fp16: false + + # ============================================================================ + # Optimizer Configuration + # ============================================================================ + + # Optimizer + optimizer: adam + lr: 1.0e-4 # Peak learning rate + min_lr: 1.0e-5 # Minimum learning rate + weight_decay: 0.01 + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-8 + + # Gradient clipping (important for FP8 stability) + clip_grad: 1.0 + + # ============================================================================ + # Learning Rate Scheduler + # ============================================================================ + + lr_decay_style: cosine + lr_warmup_iters: 100 # Short warmup for testing + lr_decay_iters: 1000 # Match train_iters + + # ============================================================================ + # Distributed Training Configuration + # ============================================================================ + + # Single GPU configuration + tensor_model_parallel_size: 1 # No TP needed for 535M + pipeline_model_parallel_size: 1 # No PP needed + context_parallel_size: 1 # No CP needed + + # Distributed settings + use_distributed_optimizer: false # Not needed for single GPU + overlap_grad_reduce: false # Not applicable for single GPU + use_flash_attn: true # Use Flash Attention 2 + distributed_timeout_minutes: 60 + + # ============================================================================ + # Memory Optimization + # ============================================================================ + + # Activation checkpointing (not needed for 535M with FP8) + recompute_granularity: null # No recompute needed + recompute_method: null + recompute_num_layers: null + + # Sequence parallelism + sequence_parallel: false # Not needed for single GPU + + # ============================================================================ + # Seed and Reproducibility + # ============================================================================ + + seed: 42 + + # ============================================================================ + # Monitoring and Logging + # ============================================================================ + + wandb_project: flux_535m_pretrain_fp8 + wandb_exp_name: flux_535m_fp8_test + +# ============================================================================ +# Notes - FP8 Validation with 535M +# ============================================================================ +# +# Hardware Requirements (with FP8): +# - Minimum: 1× MI355X 256GB +# - Memory per GPU: ~3-5GB (vs ~7-10GB BF16) +# - Training time: Minutes +# +# Validation Checklist: +# [ ] Setup FP8 environment (see docs/backends/megatron/diffusion/fp8_training.md) +# [ ] Verify TE FP8 support is available +# [ ] Run this config to validate FP8 training +# [ ] Check logs for NaN/Inf (should be none) +# [ ] Verify memory usage is ~50% of BF16 +# [ ] Verify training speed is 1.5-2x faster than BF16 +# [ ] Check loss decreases normally +# +# Expected Results: +# - Training completes 1000 steps in 5-15 minutes +# - No NaN/Inf in losses +# - Memory usage: ~3-5GB +# - Speed: ~10-50 steps/sec (depending on hardware) +# - Loss should decrease normally +# +# If validation passes, proceed to one of the 12B FP8 configs, e.g.: +# - flux_12b_ddp_energon_schnell_resample_te_spec_fp8.yaml (TransformerEngine FP8) +# - flux_12b_fsdp2_energon_schnell_resample_local_spec_fp8.yaml (local-spec FP8) +# +# Troubleshooting: +# - If NaN/Inf: Check Transformer Engine FP8 support +# - If OOM: Reduce micro_batch_size +# - If slow: Verify ROCm FP8 tensor cores are being used +# - If unstable: Try fp8_wgrad: false in model config +# +# For more information: See docs/backends/megatron/diffusion/fp8_training.md diff --git a/examples/megatron/configs/MI355X/diffusion/flux_535m_with_guidance_embed.yaml b/examples/megatron/configs/MI355X/diffusion/flux_535m_with_guidance_embed.yaml new file mode 100644 index 000000000..49cb41194 --- /dev/null +++ b/examples/megatron/configs/MI355X/diffusion/flux_535m_with_guidance_embed.yaml @@ -0,0 +1,58 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 535M with Guidance Embedding (Advanced Configuration) +# +# This config demonstrates Flux training with guidance embedding enabled. +# This is an OPTIONAL advanced feature that allows for faster single-pass CFG +# during inference, but requires training with guidance embedding enabled. +# +# IMPORTANT: Most users should use the standard flux_535m_pretrain.yaml config. +# Only use this if you specifically need guidance embedding support. +# +# Usage: +# EXP=examples/megatron/configs/MI355X/diffusion/flux_535m_with_guidance_embed.yaml \ +# bash examples/run_pretrain.sh + +# Extend standard 535M config +extends: + - flux_535m_pretrain.yaml + +modules: + pre_trainer: + overrides: + # ============================================================================ + # Guidance Embedding Configuration (ADVANCED) + # ============================================================================ + + # Enable guidance embedding for single-pass CFG + # This adds a learned MLPEmbedder layer that conditions on guidance scale + guidance_embed: true + + # Guidance scale used during training + # Model learns to adapt its predictions based on this scale + guidance_scale: 3.5 + + # ============================================================================ + # Notes + # ============================================================================ + # + # Training with guidance embedding: + # - Adds ~1-2% more parameters (guidance MLPEmbedder) + # - Allows single-pass CFG during inference (faster) + # - Requires more training data/iterations to converge + # - Model learns guidance as a conditioning signal + # + # Inference with guidance embedding: + # - Pipeline automatically detects guidance_embed layer + # - Uses single forward pass instead of batch doubling + # - ~2x faster CFG compared to explicit CFG + # - Guidance scale can be varied at inference time + # + # Standard approach (guidance_embed: false): + # - Default for most Primus training + # - Uses explicit CFG (batch doubling) at inference + # - More compatible with existing checkpoints + # - Slightly slower but more flexible + # + # See examples/megatron/diffusion/README.md for more details. diff --git a/primus/configs/data/megatron/diffusion/README.md b/primus/configs/data/megatron/diffusion/README.md new file mode 100644 index 000000000..a674ba69c --- /dev/null +++ b/primus/configs/data/megatron/diffusion/README.md @@ -0,0 +1,659 @@ +# Diffusion Data Configuration Guide + +This directory contains configuration files for preparing and organizing datasets for Megatron-based diffusion model training in Primus. + +## Directory Structure + +``` +primus/configs/data/megatron/diffusion/ +├── README.md # This file +├── templates/ # Dataset YAML templates +│ ├── dataset_preencoded.yaml # Pre-encoded dataset configuration +│ ├── dataset_preencoded_numpy.yaml # Pre-encoded NumPy/MLPerf dataset configuration +│ ├── dataset_raw.yaml # Raw dataset configuration +│ └── metadataset.yaml # Multi-dataset mixing configuration +└── preprocessing/ # CLI preprocessing configuration examples + ├── quickstart_pokemon.yaml # Minimal quickstart (small HuggingFace dataset) + ├── example_base.yaml # Comprehensive example with all options + ├── example_huggingface.yaml # HuggingFace dataset example + ├── example_directory.yaml # Local directory example + ├── example_webdataset.yaml # WebDataset input example + ├── text_to_image_2m_10k.yaml # 10K subset of text-to-image-2M (1024px) + └── mlperf_flux1.yaml # MLPerf Flux1 streaming ingest configuration +``` + +## Quick Reference + +### What Goes Where? + +| File Type | Location | Purpose | When to Use | +|-----------|----------|---------|-------------| +| **Dataset Templates** | `templates/` | Define dataset structure | Auto-applied during finalization (or manually copy if using `--no-finalize`) | +| **Preprocessing Configs** | `preprocessing/` | Configure data preparation | Pass to `primus data` CLI commands | +| **Model Configs** | `primus/configs/models/` | Define model architecture | Reference in training configs | +| **Training Configs** | `examples/megatron/configs/MI300X/diffusion/` | Configure training runs | Main config for training scripts | +| **MLPerf Ingest Configs** | `preprocessing/` | Configure MLPerf streaming ingest | Pass to `primus data diffusion-ingest` | +| **Path Override** | Command line | Customize output location | Add `--output-dir /your/path` to any command | +| **Skip Finalize** | Command line | Skip automatic dataset setup | Add `--no-finalize` flag to skip finalization | + +**Tip**: Finalization (creating `dataset.yaml`, running `energon prepare`, validation) runs automatically by default. Use `--no-finalize` to skip it. + +--- + +## Docker Volume Mounting and Data Paths + +When using `primus-cli direct --` commands inside containers, understanding Docker volume mounts is important for data persistence. + +### Default Docker Setup + +From [`tools/docker/start_container.sh`](../../../../tools/docker/start_container.sh): +```bash +DATA_PATH=${DATA_PATH:-"${PRIMUS_PATH}/data"} # Default: ./data relative to repo +# Mounted as: -v "${DATA_PATH}:${DATA_PATH}" +``` + +**What this means:** +- **Default mount**: `/data` on host → `/data` (or `/workspace/Primus/data`) in container +- **Config file paths**: Use `/workspace/Primus/data/...` by default +- **Result**: Data automatically persists to your host machine's `data/` directory + +### Using Default Paths (Recommended) + +Config files already use the correct paths. Finalization runs automatically: + +```bash +# Single command - data preparation + finalization + validation! +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token + +# Data is saved to: /workspace/Primus/data/encoded_pokemon/ (persisted to host) +# Dataset is automatically indexed, validated, and ready for training! +``` + +Automatic finalization: +- Creates `.nv-meta/dataset.yaml` with `CrudeWebdataset` configuration +- Runs `energon prepare` to index the dataset +- Validates with Primus's custom validation (metadata, sample counts, API spot-check) +- No manual post-processing needed + +### Easy Path Override + +**Override any path using `--output-dir` flag:** + +```bash +# Customize output location +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token \ + --output-dir /workspace/Primus/data/my_custom_output + +# Or use a different mounted volume +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token \ + --output-dir /mnt/shared_storage/datasets/encoded +``` + +### Advanced: Using /data Paths + +If you prefer to use `/data/...` paths (as shown in some examples): + +```bash +# Set DATA_PATH before starting container +export DATA_PATH=/data + +# Then start/restart container with new mount +# Now all /data/... paths will work as written +``` + +### Quick Comparison + +| Approach | Command | Persistence | Notes | +|----------|---------|-------------|-------| +| **Default** | Use config as-is | ✅ Yes | Easiest, works immediately | +| **Override** | Add `--output-dir /path` | ✅ Yes (if path is mounted) | Flexible, per-command | +| **DATA_PATH** | `export DATA_PATH=/data` | ✅ Yes | Cleaner paths, requires setup | + +--- + +## Dataset Templates (`templates/`) + +These files define how Energon loads and processes your prepared datasets. With automatic finalization (the default), these templates are applied automatically -- you only need them for manual workflows or customization. + +### [`dataset_preencoded.yaml`](templates/dataset_preencoded.yaml) + +**Purpose**: Configure a pre-encoded dataset (contains VAE latents and text embeddings). + +**When to use**: Production training (2-3x faster than raw mode). + +**Usage** (automatic finalization handles steps 2-3): +```bash +# Single command -- finalization is automatic: +primus-cli direct -- data diffusion-encoded \ + --source-type directory \ + --input-dir /data/raw_images \ + --output-dir /workspace/Primus/data/encoded_dataset \ + --hf-token-file /path/to/.hf_token + +# Use in training config: +# data: +# dataset_path: /workspace/Primus/data/encoded_dataset + +# Manual workflow (if using --no-finalize): +# 1. Copy template: cp templates/dataset_preencoded.yaml /.nv-meta/dataset.yaml +# 2. Run indexing: energon prepare --num-workers 8 +``` + +### [`dataset_raw.yaml`](templates/dataset_raw.yaml) + +**Purpose**: Configure a raw dataset (contains original images and captions). + +**When to use**: Experimentation, rapid prototyping, when disk space is limited. + +**Usage** (automatic finalization handles steps 2-3): +```bash +# Single command -- finalization is automatic: +primus-cli direct -- data diffusion-raw \ + --source-type directory \ + --input-dir /data/raw_images \ + --output-dir /workspace/Primus/data/raw_dataset + +# Use in training config with encoder configs: +# data: +# dataset_path: /workspace/Primus/data/raw_dataset +# encoder_configs: +# vae: {...} +# t5: {...} +# clip: {...} +``` + +### [`dataset_preencoded_numpy.yaml`](templates/dataset_preencoded_numpy.yaml) + +**Purpose**: Configure an MLPerf pre-encoded numpy dataset (bfloat16 tensors stored as NumPy uint16 bytes). + +**When to use**: Training with MLPerf Flux1 pre-encoded data ingested via `diffusion-ingest`. + +**Sample keys**: `t5.bytes`, `clip.bytes`, `mean.bytes`, `logvar.bytes`, `.json` + +**Usage**: Applied automatically by `diffusion-ingest` finalization. For manual use: +```bash +cp templates/dataset_preencoded_numpy.yaml /.nv-meta/dataset.yaml +energon prepare --num-workers 8 +``` + +### [`metadataset.yaml`](templates/metadataset.yaml) + +**Purpose**: Combine multiple datasets with different weights. + +**When to use**: +- Mixing pre-encoded and raw datasets +- Combining multiple data sources +- Creating weighted train/val splits + +**Usage**: +```bash +# 1. Prepare individual datasets first (pre-encoded and/or raw) + +# 2. Copy and customize metadataset template +cp primus/configs/data/megatron/diffusion/templates/metadataset.yaml \ + /workspace/Primus/data/combined_dataset.yaml + +# 3. Edit paths in combined_dataset.yaml to point to your datasets + +# 4. Use in training config +# data: +# dataset_path: /workspace/Primus/data/combined_dataset.yaml +``` + +--- + +## Preprocessing Configs (`preprocessing/`) + +These files are passed to the `primus data` CLI commands to configure dataset preparation. They are **NOT copied to dataset directories**. + +### [`quickstart_pokemon.yaml`](preprocessing/quickstart_pokemon.yaml) + +**Purpose**: Minimal quickstart config for the Pokemon dataset (256px). + +**Usage**: +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml \ + --hf-token-file /path/to/.hf_token +``` + +### [`example_base.yaml`](preprocessing/example_base.yaml) + +**Purpose**: Comprehensive example showing all available preprocessing options. + +**Contents**: All configuration sections with inline documentation. + +**Usage**: +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \ + --hf-token-file /path/to/.hf_token +``` + +### [`example_huggingface.yaml`](preprocessing/example_huggingface.yaml) + +**Purpose**: Prepare a dataset from HuggingFace Hub. + +**Usage**: +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token +``` + +### [`example_directory.yaml`](preprocessing/example_directory.yaml) + +**Purpose**: Prepare a dataset from a local directory of images and captions. + +**Usage**: +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \ + --hf-token-file /path/to/.hf_token +``` + +### [`example_webdataset.yaml`](preprocessing/example_webdataset.yaml) + +**Purpose**: Prepare a dataset from existing WebDataset tar files. + +**Usage**: +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \ + --hf-token-file /path/to/.hf_token +``` + +### [`text_to_image_2m_10k.yaml`](preprocessing/text_to_image_2m_10k.yaml) + +**Purpose**: Prepare the 10K high-resolution subset (1024x1024) from the text-to-image-2M dataset. + +**Usage**: +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml \ + --hf-token-file /path/to/.hf_token +``` + +### [`mlperf_flux1.yaml`](preprocessing/mlperf_flux1.yaml) + +**Purpose**: MLPerf Flux1 streaming ingest -- downloads pre-encoded Arrow files from MLCommons R2 and converts to Energon WebDataset tar shards. + +**Datasets**: CC12M (train) and COCO (val), plus empty encodings for CFG dropout. + +**Usage**: +```bash +primus-cli direct -- data diffusion-ingest \ + --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml + +# Limit files for testing: +primus-cli direct -- data diffusion-ingest \ + --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml \ + --max-files 5 +``` + +Note: `--hf-token-file` is only required when downloading gated models from HuggingFace (e.g., the default FLUX.1-dev). If using local encoder paths via `--model-path`, it can be omitted. + +--- + +## Path Configuration Patterns + +### Pattern 1: Default (Recommended) + +Use config file paths as-is -- they work out of the box: + +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token + +# Output automatically goes to: /workspace/Primus/data/encoded_pokemon +# This persists to: /data/encoded_pokemon on host +``` + +### Pattern 2: Custom Path Override (Flexible) + +Override output location for any command: + +```bash +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token \ + --output-dir /workspace/Primus/data/custom_output + +# Or use a different mounted volume +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token \ + --output-dir /mnt/shared_storage/datasets +``` + +### Pattern 3: DATA_PATH Environment Variable (Advanced) + +Set `DATA_PATH` for cleaner paths: + +```bash +# Before starting container +export DATA_PATH=/data + +# Start/restart container (mounts /data) +# Now configs using /data/... will work as written +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token +``` + +### Verification + +Check where your data was written: + +```bash +# For default paths +ls -lh /workspace/Primus/data/encoded_pokemon/ + +# On host machine +ls -lh /data/encoded_pokemon/ + +# For custom paths +ls -lh /your/custom/path/ +``` + +--- + +## Complete Workflow + +### Pre-encoded Mode (Recommended for Production) + +```bash +# Step 1: Prepare, finalize, and validate dataset (all automatic) +primus-cli direct -- data diffusion-encoded \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ + --hf-token-file /path/to/.hf_token + +# Output: /workspace/Primus/data/encoded_pokemon (persisted to host) +# Finalization (dataset.yaml + energon prepare + validation) runs automatically. + +# Step 2: Configure training +# Edit examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml: +# data: +# dataset_path: /workspace/Primus/data/encoded_pokemon + +# Step 3: Train +EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \ + bash examples/run_pretrain.sh +``` + +### Raw Mode (On-the-fly Encoding) + +```bash +# Step 1: Prepare, finalize, and validate dataset (all automatic) +primus-cli direct -- data diffusion-raw \ + --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml + +# Output: /workspace/Primus/data/raw_dataset (persisted to host) + +# Step 2: Configure training with encoders +# Edit examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml: +# data: +# dataset_path: /workspace/Primus/data/raw_dataset +# encoder_configs: +# vae: { model_path: black-forest-labs/FLUX.1-dev, ... } +# t5: { model_path: black-forest-labs/FLUX.1-dev, ... } +# clip: { model_path: black-forest-labs/FLUX.1-dev, ... } + +# Step 3: Train +EXP=examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain.yaml \ + bash examples/run_pretrain.sh +``` + +### MLPerf Ingest Mode (Streaming Download + Conversion) + +```bash +# Step 1: Download and convert MLPerf Arrow data to Energon WebDataset +# Downloads ~1.2 TB of pre-encoded data with minimal temporary disk usage. +primus-cli direct -- data diffusion-ingest \ + --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml + +# Output: /workspace/Primus/data/mlperf_flux1/ (persisted to host) +# Finalization (dataset.yaml + energon prepare + validation) runs automatically. +# Pipeline supports resume -- re-run the same command to continue after interruption. + +# Step 2: Point the training config at the prepared data. +# The MLPerf reproduction configs already set the required VAE normalization +# (vae_latent_mode: resample, vae_scale: 0.3611, vae_shift: 0.1159); just set +# data_path (or the PRIMUS_DIFFUSION_DATA_PATH env var) to /workspace/Primus/data/mlperf_flux1. + +# Step 3: Train (MLPerf benchmark reproduction) +EXP=examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_fp8_mlperf.yaml \ + bash examples/run_pretrain.sh +``` + +**Output directory structure:** +``` +mlperf_flux1/ +├── train/ # CC12M shards (shard_000000.tar, ...) +├── val/ # COCO shards +├── empty_encodings/ # empty_t5_encodings.npy, empty_clip_encodings.npy +└── .nv-meta/ # Energon index + dataset.yaml (auto-generated) +``` + +**Pipeline features:** +- Parallel downloads with configurable worker threads (`max_workers`) +- Bounded disk usage via semaphore (`prefetch_depth` Arrow files on disk at once) +- Automatic retry with exponential backoff for HTTP 429/503 and MD5 mismatches +- Resume support -- re-run to skip already-completed shards +- Skip-and-log for individual file failures (`failed_files.json`) + +--- + +## Configuration Hierarchy + +Understanding how different configs relate to each other: + +```mermaid +graph TB + subgraph preprocessing [Preprocessing Phase] + prepConfig[preprocessing/example_*.yaml] + cliCommand[primus data CLI] + prepConfig --> cliCommand + end + + subgraph datasetDir [Dataset Directory] + tarFiles[WebDataset *.tar files] + datasetYaml["dataset.yaml (auto-generated)"] + energonIndex[.nv-meta/ Energon index] + + tarFiles --> energonIndex + datasetYaml --> energonIndex + end + + subgraph training [Training Phase] + trainConfig[Training Config examples/*.yaml] + modelConfig[Model Config primus/configs/models/] + encoderConfig[Encoder Config encoders.yaml] + + trainConfig --> modelConfig + trainConfig -.raw mode only.-> encoderConfig + end + + cliCommand --> tarFiles + energonIndex --> trainConfig +``` + +### Key Points: + +1. **Preprocessing configs** -- Used once during data preparation +2. **Dataset templates** -- Auto-applied during finalization (or manually copied if using `--no-finalize`) +3. **Training configs** -- Main config file that references dataset path +4. **Model configs** -- Define architecture, referenced by training configs +5. **Encoder configs** -- Only needed for raw (on-the-fly) mode + +--- + +## Common Patterns + +### Pattern 1: Single Pre-encoded Dataset + +```yaml +# Training config +data: + dataset_path: /workspace/Primus/data/my_dataset/dataset.yaml # Points to dataset_preencoded.yaml copy + micro_batch_size: 2 + global_batch_size: 16 +``` + +### Pattern 2: Mixed Pre-encoded + Raw + +```yaml +# Copy and customize templates/metadataset.yaml +data: + dataset_path: /workspace/Primus/data/mixed_dataset.yaml # Points to customized metadataset.yaml + encoder_configs: # Needed for raw datasets + vae: {...} + t5: {...} + clip: {...} +``` + +### Pattern 3: Multiple Pre-encoded Datasets + +```yaml +# Use metadataset.yaml to combine +data: + dataset_path: /workspace/Primus/data/combined.yaml # Metadataset config + # No encoder_configs needed if all datasets are pre-encoded +``` + +--- + +## File Formats + +### Dataset YAML (in dataset directories) + +```yaml +__module__: megatron.energon +__class__: CrudeWebdataset +subflavors: + encoding: preencoded # or 'raw' or 'preencoded_numpy' +# Optional filters and settings +``` + +**Encoding types:** +- `preencoded` -- Primus-encoded PyTorch `.pth` format (VAE latents + text embeddings) +- `preencoded_numpy` -- MLPerf NumPy uint16 format (bfloat16 as `.bytes` entries) +- `raw` -- Original images and captions (encoded on-the-fly during training) + +### Metadataset YAML + +```yaml +__module__: megatron.energon +__class__: Metadataset +splits: + train: + datasets: + - weight: 0.8 + path: /path/to/dataset1/ + - weight: 0.2 + path: /path/to/dataset2/ +``` + +### Preprocessing Config YAML + +```yaml +source: + type: huggingface # or directory, webdataset + # source-specific options +output: + output_dir: /data/output + shard_size: 1000 +model: + model_path: black-forest-labs/FLUX.1-dev + batch_size: 8 +``` + +--- + +## Additional Resources + +### Documentation +- **Training Guide**: [`examples/megatron/diffusion/README.md`](../../../../examples/megatron/diffusion/README.md) +- **Energon Integration**: [`docs/backends/megatron/diffusion/energon_integration.md`](../../../../docs/backends/megatron/diffusion/energon_integration.md) +- **FP8 Training Guide**: [`docs/backends/megatron/diffusion/fp8_training.md`](../../../../docs/backends/megatron/diffusion/fp8_training.md) + +### Related Configs +- **Encoder Configs**: [`primus/configs/models/megatron/diffusion/encoders.yaml`](../../models/megatron/diffusion/encoders.yaml) +- **Model Architecture**: [`primus/configs/models/megatron/diffusion/`](../../models/megatron/diffusion/) +- **Training Configs (MI300X)**: [`examples/megatron/configs/MI300X/diffusion/`](../../../../examples/megatron/configs/MI300X/diffusion/) +- **Training Configs (MI355X)**: [`examples/megatron/configs/MI355X/diffusion/`](../../../../examples/megatron/configs/MI355X/diffusion/) + +--- + +## Troubleshooting + +### "No such file or directory: dataset.yaml" + +**Problem**: Training config references a dataset path that doesn't exist. + +**Solution**: If you used `--no-finalize`, you need to manually copy the appropriate template from `templates/` to your dataset's `.nv-meta/` directory and run `energon prepare`. Otherwise, re-run preprocessing without `--no-finalize` (the default). + +### "Unknown subflavor: encoding" + +**Problem**: Old dataset.yaml format or missing subflavors field. + +**Solution**: Use the updated templates from this directory (dataset_preencoded.yaml or dataset_raw.yaml). + +### "Encoder not found" (raw mode) + +**Problem**: Training with raw dataset but encoder_configs not specified. + +**Solution**: Add encoder_configs to your training config or switch to pre-encoded mode. + +### NaN loss with MLPerf data + +**Problem**: Training on MLPerf-ingested data produces NaN or diverging loss. + +**Solution**: Ensure your training config includes the required normalization constants: +```yaml +vae_latent_mode: resample +vae_scale: 0.3611 +vae_shift: 0.1159 +``` +These are mandatory because `forward_step.py` unconditionally applies `vae_scale * (latents - vae_shift)` after resampling from mean/logvar. + +### Missing mean/logvar keys in MLPerf data + +**Problem**: Training errors about missing `mean` or `logvar` fields. + +**Solution**: Verify your dataset's `.nv-meta/dataset.yaml` uses `encoding: preencoded_numpy` (not `preencoded`). The `preencoded_numpy` cooker expects `.bytes` keys while `preencoded` expects `.pth` keys. + +### Preprocessing config not found + +**Problem**: Trying to use old config paths that have been moved. + +**Solution**: Use new paths: +- Old: `examples/megatron/diffusion/configs/data_preprocessing.yaml` +- New: `primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml` + +--- + +## Migration from Old Structure + +If you have configs using old paths, update them: + +| Old Path | New Path | +|----------|----------| +| `examples/megatron/diffusion/configs/preencoded_dataset.yaml` | `primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml` | +| `examples/megatron/diffusion/configs/flux_dataset_preencoded.yaml` | `primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml` | +| `examples/megatron/diffusion/configs/raw_dataset.yaml` | `primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml` | +| `examples/megatron/diffusion/configs/flux_dataset_raw.yaml` | `primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml` | +| `examples/megatron/diffusion/configs/metadataset.yaml` | `primus/configs/data/megatron/diffusion/templates/metadataset.yaml` | +| `examples/megatron/diffusion/data/test_pokemon_config.yaml` | `primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml` | +| `examples/megatron/diffusion/configs/data_preprocessing.yaml` | `primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml` | + +--- + +**Questions or issues?** See the main documentation or open an issue on GitHub. diff --git a/primus/configs/data/megatron/diffusion/preprocessing/coco2014_train_schnell_256.yaml b/primus/configs/data/megatron/diffusion/preprocessing/coco2014_train_schnell_256.yaml new file mode 100644 index 000000000..4031e2783 --- /dev/null +++ b/primus/configs/data/megatron/diffusion/preprocessing/coco2014_train_schnell_256.yaml @@ -0,0 +1,31 @@ +# Pre-encode COCO 2014 train (82,783 images) for Flux-Schnell training +# 256px images, 256-token T5 sequences (pass --t5-max-length 256 on CLI) +# +# Dataset: AbdoTW/COCO_2014 train split (full COCO 2014 train2014) +# Fields: image (PIL), caption (list[str] -> first caption taken automatically) +# +# Usage (4 GPU on devices 4-7): +# HIP_VISIBLE_DEVICES=4,5,6,7 torchrun --nproc_per_node=4 \ +# /workspace/Primus/primus/cli/main.py data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/coco2014_train_schnell_256.yaml \ +# --t5-max-length 256 --hf-token-file /workspace/Primus/.hf_token + +source: + type: huggingface + hf_dataset: AbdoTW/COCO_2014 + hf_split: train + +output: + output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/coco2014_train_schnell_256_encoded + shard_size: 1000 + compress: false + +model: + model_path: black-forest-labs/FLUX.1-schnell + precision: bf16 + device: cuda + batch_size: 8 + +image: + image_size: 256 + center_crop: false diff --git a/primus/configs/data/megatron/diffusion/preprocessing/coco_schnell_256.yaml b/primus/configs/data/megatron/diffusion/preprocessing/coco_schnell_256.yaml new file mode 100644 index 000000000..f033192be --- /dev/null +++ b/primus/configs/data/megatron/diffusion/preprocessing/coco_schnell_256.yaml @@ -0,0 +1,36 @@ +# Pre-encode COCO for Flux-Schnell training (VAE resample mode) +# 256px images, 256-token T5 sequences (pass --t5-max-length 256 on CLI) +# vae_latent_mode: resample — stores mean+logvar so latents are re-drawn +# via reparameterization (mean + exp(0.5*logvar) * randn) at every training step. +# +# Dataset: UCSC-VLAA/Recap-COCO-30K (30K COCO images with captions) +# Fields: image (PIL), caption (str) +# +# Usage (8 GPU): +# torchrun --nproc_per_node=8 /workspace/Primus/primus/cli/main.py \ +# data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/coco_schnell_256.yaml \ +# --t5-max-length 256 \ +# --hf-token-file /workspace/Primus/.hf_token + +source: + type: huggingface + hf_dataset: UCSC-VLAA/Recap-COCO-30K + hf_split: train + +output: + output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/coco_schnell_256_resample_encoded + shard_size: 1000 + max_samples: null + compress: false + +model: + model_path: black-forest-labs/FLUX.1-schnell + precision: bf16 + device: cuda + batch_size: 8 + vae_latent_mode: resample + +image: + image_size: 256 + center_crop: false diff --git a/primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml b/primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml new file mode 100644 index 000000000..ee6cf4332 --- /dev/null +++ b/primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml @@ -0,0 +1,196 @@ +# Data Preprocessing Configuration for Diffusion Models +# +# This comprehensive config file shows all available preprocessing parameters. +# +# ============================================================================ +# Quick Start (Recommended - Single Command) +# ============================================================================ +# Prepare and finalize dataset in one command: +# +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \ +# --hf-token-file /path/to/.hf_token +# +# The default encoder model (FLUX.1-dev) is gated and requires a +# HuggingFace token. Get one at https://huggingface.co/settings/tokens +# +# Finalization (dataset.yaml + energon prepare) runs automatically. +# To skip it, pass --no-finalize. +# +# ============================================================================ +# Advanced Usage +# ============================================================================ +# Override specific parameters: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --batch-size 16 \ +# --output-dir /your/custom/path +# +# Priority: CLI arguments > YAML config > defaults +# ============================================================================ + +# ============================================================================ +# Path Configuration +# ============================================================================ +# Default paths use Docker mounted paths that persist to host: +# /workspace/Primus/data/... → /data/... on host +# +# EASY OVERRIDE: Customize output with --output-dir flag: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --output-dir /your/custom/path +# +# For /data/... paths, set: export DATA_PATH=/data (before starting container) +# ============================================================================ + +# ============================================================================ +# Source Configuration +# ============================================================================ +source: + # Type of data source (required) + type: huggingface # Options: huggingface, directory, webdataset + + # HuggingFace source options (when type: huggingface) + hf_dataset: diffusers/pokemon-gpt4-captions # HuggingFace dataset name + hf_split: train # Dataset split to use + + # Directory source options (when type: directory) + # input_dir: /path/to/images # Directory containing images and captions + + # WebDataset source options (when type: webdataset) + # input_path: '/path/to/dataset/*.tar' # Glob pattern for WebDataset tars + +# ============================================================================ +# Data Format Configuration (Optional) +# ============================================================================ +# Specify how to extract image and caption from dataset samples. +# If not specified, uses automatic detection with common field names: +# - Images: 'image', 'jpg', 'jpeg', 'png', 'img', 'photo' +# - Captions: 'caption', 'text', 'txt', 'description', 'prompt' +# +# Use this section for datasets with non-standard field names or JSON-encoded data. +# data_format: +# # Single field names (simplest approach) +# image_key: jpg # e.g., images are in 'jpg' field +# caption_key: json.caption # e.g., captions in JSON at 'caption' key +# +# # Alternative: specify multiple fallback keys (tried in order) +# # image_keys: ['image', 'jpg', 'png'] +# # caption_keys: ['caption', 'json.caption', 'json.text'] +# +# Examples: +# # For text-to-image-2M dataset: +# data_format: +# image_key: jpg +# caption_key: json.caption +# +# # For datasets with nested JSON metadata: +# data_format: +# image_key: image +# caption_key: json.metadata.description + +# ============================================================================ +# Output Configuration +# ============================================================================ +output: + output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/encoded_pokemon # Output directory (persisted to host) + shard_size: 1000 # Number of samples per output shard + max_samples: null # Maximum samples to process (null = all) + compress: false # Whether to compress output with gzip + +# ============================================================================ +# Model Configuration +# ============================================================================ +model: + # Base model path (HuggingFace repo or local path) + model_path: black-forest-labs/FLUX.1-dev + + # Override individual encoder paths (null means use model_path) + vae_path: null # Custom VAE path + t5_path: null # Custom T5-XXL path + clip_path: null # Custom CLIP-L path + + # Model settings + precision: bf16 # Options: bf16, fp16, fp32 + device: cuda # Device for encoding + batch_size: 8 # Encoding batch size + +# ============================================================================ +# Image Preprocessing +# ============================================================================ +image: + image_size: 1024 # Target image size (height and width) + center_crop: false # Whether to center crop before resize + +# ============================================================================ +# Examples +# ============================================================================ +# +# Example 1: From HuggingFace with custom batch size +# --------------------------------------------------- +# source: +# type: huggingface +# hf_dataset: diffusers/pokemon-gpt4-captions +# output: +# output_dir: /workspace/Primus/data/pokemon +# model: +# batch_size: 16 +# precision: bf16 +# +# Example 2: From local directory +# -------------------------------- +# source: +# type: directory +# input_dir: /data/raw_images +# output: +# output_dir: /workspace/Primus/data/encoded_images +# model: +# model_path: black-forest-labs/FLUX.1-dev +# +# Example 3: From WebDataset with custom encoders +# ------------------------------------------------ +# source: +# type: webdataset +# input_path: '/data/dataset/*.tar' +# output: +# output_dir: /workspace/Primus/data/encoded +# model: +# vae_path: madebyollin/sdxl-vae-fp16-fix +# t5_path: google/t5-v1_1-xxl +# clip_path: openai/clip-vit-large-patch14 +# +# ============================================================================ +# Usage +# ============================================================================ +# +# Single GPU (uses config paths): +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \ +# --hf-token-file /path/to/.hf_token +# +# Multi-GPU (8 GPUs): +# primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \ +# --hf-token-file /path/to/.hf_token +# +# Override output directory: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --output-dir /workspace/Primus/data/my_output +# +# Override multiple values: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_base.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --output-dir /your/custom/path \ +# --batch-size 16 \ +# --max-samples 1000 +# +# For more examples, see: +# - example_huggingface.yaml (HuggingFace dataset) +# - example_directory.yaml (Local directory) +# - example_webdataset.yaml (WebDataset input) +# diff --git a/primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml b/primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml new file mode 100644 index 000000000..cb80838be --- /dev/null +++ b/primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml @@ -0,0 +1,130 @@ +# Preprocessing Configuration Example: Local Directory +# +# This example shows how to prepare a dataset from a local directory. +# +# ============================================================================ +# Quick Start (Recommended - Single Command) +# ============================================================================ +# Prepare and finalize dataset in one command: +# +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \ +# --hf-token-file /path/to/.hf_token +# +# The default encoder model (FLUX.1-dev) is gated and requires a +# HuggingFace token. Get one at https://huggingface.co/settings/tokens +# +# Finalization (dataset.yaml + energon prepare) runs automatically. +# To skip it, pass --no-finalize. +# +# ============================================================================ +# Advanced Usage +# ============================================================================ +# Custom paths: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --input-dir /your/input/path \ +# --output-dir /your/output/path +# ============================================================================ + +# ============================================================================ +# Path Configuration +# ============================================================================ +# Default paths use Docker mounted paths that persist to host: +# /workspace/Primus/data/... → /data/... on host +# +# EASY OVERRIDE: Customize paths with command-line flags: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --input-dir /your/input/path \ +# --output-dir /your/output/path +# +# For /data/... paths, set: export DATA_PATH=/data (before starting container) +# ============================================================================ + +# ============================================================================ +# Source Configuration (Directory) +# ============================================================================ +source: + type: directory + input_dir: /data/raw_images # Directory containing images and captions + +# ============================================================================ +# Output Configuration +# ============================================================================ +output: + output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/encoded_images # Output directory (persisted to host) + shard_size: 1000 # Number of samples per output shard + max_samples: null # Maximum samples to process (null = all) + compress: false # Whether to compress output with gzip + +# ============================================================================ +# Model Configuration +# ============================================================================ +model: + # Base model path (HuggingFace repo or local path) + model_path: black-forest-labs/FLUX.1-dev + + # Override individual encoder paths (null means use model_path) + vae_path: null # Custom VAE path + t5_path: null # Custom T5-XXL path + clip_path: null # Custom CLIP-L path + + # Model settings + precision: bf16 # Options: bf16, fp16, fp32 + device: cuda # Device for encoding + batch_size: 8 # Encoding batch size + +# ============================================================================ +# Image Preprocessing +# ============================================================================ +image: + image_size: 1024 # Target image size (height and width) + center_crop: false # Whether to center crop before resize + +# ============================================================================ +# Input Directory Structure +# ============================================================================ +# +# Your input_dir should contain image-caption pairs: +# +# /data/raw_images/ +# ├── image_001.jpg +# ├── image_001.txt +# ├── image_002.jpg +# ├── image_002.txt +# └── ... +# +# Each .jpg file should have a corresponding .txt file with the same name. +# Supported image formats: .jpg, .jpeg, .png, .webp +# +# ============================================================================ +# Usage +# ============================================================================ +# +# Single GPU (uses config paths): +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \ +# --hf-token-file /path/to/.hf_token +# +# Multi-GPU (8 GPUs): +# primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \ +# --hf-token-file /path/to/.hf_token +# +# Override paths: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --input-dir /your/input/path \ +# --output-dir /your/output/path +# +# Override specific values: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_directory.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --output-dir /workspace/Primus/data/my_output \ +# --batch-size 16 +# diff --git a/primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml b/primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml new file mode 100644 index 000000000..eed538b31 --- /dev/null +++ b/primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml @@ -0,0 +1,124 @@ +# Preprocessing Configuration Example: HuggingFace Dataset +# +# This example shows how to prepare a dataset from HuggingFace Hub. +# +# ============================================================================ +# Quick Start (Recommended - Single Command) +# ============================================================================ +# Prepare and finalize dataset in one command (no manual steps!): +# +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ +# --hf-token-file /path/to/.hf_token +# +# The default encoder model (FLUX.1-dev) is gated and requires a +# HuggingFace token. Get one at https://huggingface.co/settings/tokens +# +# Finalization (dataset.yaml + energon prepare) runs automatically. +# To skip it, pass --no-finalize. +# +# ============================================================================ +# Advanced Usage +# ============================================================================ +# Skip auto-finalization (manual post-processing): +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ +# --hf-token-file /path/to/.hf_token --no-finalize +# +# Custom output path: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --output-dir /your/custom/path +# +# Custom train/val/test split (default: 100% train): +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --train-split 0.8 # 80% train, 10% val, 10% test +# ============================================================================ + +# ============================================================================ +# Path Configuration +# ============================================================================ +# Default output_dir uses Docker mounted path that persists to host: +# /workspace/Primus/data/... → /data/... on host +# +# EASY OVERRIDE: Customize output location with --output-dir flag: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --output-dir /your/custom/path +# +# For /data/... paths, set: export DATA_PATH=/data (before starting container) +# ============================================================================ + +# ============================================================================ +# Source Configuration (HuggingFace) +# ============================================================================ +source: + type: huggingface + hf_dataset: diffusers/pokemon-gpt4-captions # HuggingFace dataset name + hf_split: train # Dataset split to use + +# ============================================================================ +# Output Configuration +# ============================================================================ +output: + output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/encoded_pokemon # Output directory (persisted to host) + shard_size: 1000 # Number of samples per output shard + max_samples: null # Maximum samples to process (null = all) + compress: false # Whether to compress output with gzip + +# ============================================================================ +# Model Configuration +# ============================================================================ +model: + # Base model path (HuggingFace repo or local path) + model_path: black-forest-labs/FLUX.1-dev + + # Override individual encoder paths (null means use model_path) + vae_path: null # Custom VAE path + t5_path: null # Custom T5-XXL path + clip_path: null # Custom CLIP-L path + + # Model settings + precision: bf16 # Options: bf16, fp16, fp32 + device: cuda # Device for encoding + batch_size: 8 # Encoding batch size + +# ============================================================================ +# Image Preprocessing +# ============================================================================ +image: + image_size: 1024 # Target image size (height and width) + center_crop: false # Whether to center crop before resize + +# ============================================================================ +# Usage +# ============================================================================ +# +# Single GPU (uses default output_dir from config): +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ +# --hf-token-file /path/to/.hf_token +# +# Multi-GPU (8 GPUs): +# primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ +# --hf-token-file /path/to/.hf_token +# +# Override output directory: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --output-dir /workspace/Primus/data/custom_output +# +# Override multiple values: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --output-dir /my/custom/path \ +# --batch-size 16 \ +# --max-samples 1000 +# diff --git a/primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml b/primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml new file mode 100644 index 000000000..ab5a1e296 --- /dev/null +++ b/primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml @@ -0,0 +1,129 @@ +# Preprocessing Configuration Example: WebDataset Input +# +# This example shows how to prepare a dataset from existing WebDataset tar files. +# +# ============================================================================ +# Quick Start (Recommended - Single Command) +# ============================================================================ +# Prepare and finalize dataset in one command: +# +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \ +# --hf-token-file /path/to/.hf_token +# +# The default encoder model (FLUX.1-dev) is gated and requires a +# HuggingFace token. Get one at https://huggingface.co/settings/tokens +# +# Finalization (dataset.yaml + energon prepare) runs automatically. +# To skip it, pass --no-finalize. +# +# ============================================================================ +# Advanced Usage +# ============================================================================ +# Custom paths: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --input-path '/your/input/*.tar' \ +# --output-dir /your/output/path +# ============================================================================ + +# ============================================================================ +# Path Configuration +# ============================================================================ +# Default paths use Docker mounted paths that persist to host: +# /workspace/Primus/data/... → /data/... on host +# +# EASY OVERRIDE: Customize paths with command-line flags: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --input-path '/your/input/*.tar' \ +# --output-dir /your/output/path +# +# For /data/... paths, set: export DATA_PATH=/data (before starting container) +# ============================================================================ + +# ============================================================================ +# Source Configuration (WebDataset) +# ============================================================================ +source: + type: webdataset + input_path: '/data/input_dataset/*.tar' # Glob pattern for WebDataset tars + +# ============================================================================ +# Output Configuration +# ============================================================================ +output: + output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/encoded_dataset # Output directory (persisted to host) + shard_size: 1000 # Number of samples per output shard + max_samples: null # Maximum samples to process (null = all) + compress: false # Whether to compress output with gzip + +# ============================================================================ +# Model Configuration +# ============================================================================ +model: + # Base model path (HuggingFace repo or local path) + model_path: black-forest-labs/FLUX.1-dev + + # Override individual encoder paths (null means use model_path) + vae_path: null # Custom VAE path + t5_path: null # Custom T5-XXL path + clip_path: null # Custom CLIP-L path + + # Model settings + precision: bf16 # Options: bf16, fp16, fp32 + device: cuda # Device for encoding + batch_size: 8 # Encoding batch size + +# ============================================================================ +# Image Preprocessing +# ============================================================================ +image: + image_size: 1024 # Target image size (height and width) + center_crop: false # Whether to center crop before resize + +# ============================================================================ +# Input WebDataset Format +# ============================================================================ +# +# Your input tar files should contain image-caption pairs: +# +# sample_000001.jpg +# sample_000001.txt +# sample_000002.jpg +# sample_000002.txt +# ... +# +# Supported image formats: .jpg, .jpeg, .png, .webp +# Caption files should be plain text UTF-8 (.txt) +# +# ============================================================================ +# Usage +# ============================================================================ +# +# Single GPU (uses config paths): +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \ +# --hf-token-file /path/to/.hf_token +# +# Multi-GPU (8 GPUs): +# primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \ +# --hf-token-file /path/to/.hf_token +# +# Override paths: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --input-path '/workspace/Primus/data/input/*.tar' \ +# --output-dir /workspace/Primus/data/output +# +# Override specific values: +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/example_webdataset.yaml \ +# --hf-token-file /path/to/.hf_token \ +# --output-dir /your/custom/path \ +# --batch-size 16 +# diff --git a/primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml b/primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml new file mode 100644 index 000000000..7d6b5e33e --- /dev/null +++ b/primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml @@ -0,0 +1,66 @@ +# MLPerf Flux1 streaming ingest configuration +# +# Downloads pre-encoded Arrow files from MLCommons R2 and converts them +# into Energon WebDataset tar shards in a single streaming pass. +# +# ============================================================================ +# Usage +# ============================================================================ +# +# primus data diffusion-ingest \ +# --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml +# +# Override output: +# primus data diffusion-ingest \ +# --config primus/configs/data/megatron/diffusion/preprocessing/mlperf_flux1.yaml \ +# --output-dir /custom/path --max-files 5 +# +# ============================================================================ +# Datasets +# ============================================================================ +# Each entry runs a StreamingIngestPipeline: download Arrow files in parallel, +# convert to WebDataset tar shards, delete Arrow files after conversion. + +datasets: + - name: cc12m + manifest_url: https://training.mlcommons-storage.org/metadata/flux-1-cc12m-preprocessed.uri + split_name: train + + - name: coco + manifest_url: https://training.mlcommons-storage.org/metadata/flux-1-coco-preprocessed.uri + split_name: val + +# ============================================================================ +# Empty Encodings (for CFG dropout during training) +# ============================================================================ +# Small .npy files downloaded separately (not Arrow). + +empty_encodings: + manifest_url: https://training.mlcommons-storage.org/metadata/flux-1-empty-encodings.uri + output_subdir: empty_encodings + +# ============================================================================ +# Output Configuration +# ============================================================================ + +output: + output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/mlperf_flux1 + +# ============================================================================ +# Pipeline Configuration +# ============================================================================ + +pipeline: + max_workers: 4 # concurrent download threads + prefetch_depth: 6 # max Arrow files buffered on disk (~1.2 GB) + max_files: null # null = all files (set to small number for testing) + +# ============================================================================ +# Required Training-Time Settings (not consumed by ingest, for reference) +# ============================================================================ +# When training with this dataset, configure: +# vae_latent_mode: resample +# vae_scale: 0.3611 +# vae_shift: 0.1159 +# These are required because forward_step.py unconditionally applies +# vae_scale * (latents - vae_shift) after resampling from mean/logvar. diff --git a/primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml b/primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml new file mode 100644 index 000000000..d4ca2c275 --- /dev/null +++ b/primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml @@ -0,0 +1,34 @@ +# Quickstart: Pokemon Dataset (256px, works out of the box) +# +# Minimal config for getting started quickly with a small public dataset. +# Uses 256px images for fast processing, even on limited hardware. +# +# ============================================================================ +# Usage +# ============================================================================ +# +# primus-cli direct -- data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml \ +# --hf-token-file /path/to/.hf_token +# +# The default encoder model (FLUX.1-dev) is gated and requires a +# HuggingFace token. Get one at https://huggingface.co/settings/tokens +# ============================================================================ + +source: + type: huggingface + hf_dataset: diffusers/pokemon-gpt4-captions + hf_split: train + +output: + output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/quickstart_pokemon + shard_size: 1000 + +model: + model_path: black-forest-labs/FLUX.1-dev + precision: bf16 + batch_size: 4 + +image: + image_size: 256 + center_crop: false diff --git a/primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml b/primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml new file mode 100644 index 000000000..7fd77aabb --- /dev/null +++ b/primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml @@ -0,0 +1,93 @@ +# Text-to-Image-2M 10K Subset Configuration +# +# This config uses the 10K high-resolution subset (1024x1024) from the +# text-to-image-2M dataset, perfect for fine-tuning Flux models. +# +# Dataset: https://huggingface.co/datasets/jackyhate/text-to-image-2M +# Subset: data_1024_10K (exactly 10,000 samples at 1024x1024 resolution) +# +# ============================================================================ +# Quick Start (Recommended - Single Command) +# ============================================================================ +# Prepare and finalize dataset in one command: +# +# primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml +# +# Finalization (dataset.yaml + energon prepare) runs automatically. +# To skip it, pass --no-finalize. +# +# Processing time: ~30-60 minutes on 8 GPUs +# Output size: ~50-80 GB (pre-encoded tensors) +# ============================================================================ + +# ============================================================================ +# Source Configuration (HuggingFace with data_files) +# ============================================================================ +source: + type: huggingface + hf_dataset: jackyhate/text-to-image-2M + hf_split: train + hf_data_files: data_1024_10K/*.tar # Load only the 10K subset + +# ============================================================================ +# Data Format Configuration +# ============================================================================ +# This dataset stores images as 'jpg' files and captions in JSON format. +# Specify how to extract image and caption from dataset samples. +data_format: + image_key: jpg # Images are stored in 'jpg' field + caption_key: json.prompt # Captions are in JSON at 'prompt' key (NOT 'caption') + + # Alternative: specify multiple fallback keys (tried in order) + # image_keys: ['image', 'jpg'] + # caption_keys: ['prompt', 'json.prompt', 'caption', 'json.caption'] + +# ============================================================================ +# Output Configuration +# ============================================================================ +output: + output_dir: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/text2img_10k_encoded + shard_size: 1000 + max_samples: null + compress: false + +# ============================================================================ +# Model Configuration +# ============================================================================ +model: + model_path: black-forest-labs/FLUX.1-dev + vae_path: null + t5_path: null + clip_path: null + precision: bf16 + device: cuda + batch_size: 8 + +# ============================================================================ +# Image Preprocessing +# ============================================================================ +image: + image_size: 1024 + center_crop: false + +# ============================================================================ +# Authentication +# ============================================================================ +auth: + hf_token_file: "" # Public dataset, no token needed + +# ============================================================================ +# Dataset Information +# ============================================================================ +# Name: Text-to-Image-2M (10K subset) +# Size: 10,000 high-resolution images +# Resolution: 1024x1024 +# Content: Complex scenes with rich, detailed annotations +# License: MIT +# Use case: Fine-tuning Flux diffusion models +# +# Full command example: +# primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \ +# --config primus/configs/data/megatron/diffusion/preprocessing/text_to_image_2m_10k.yaml +# ============================================================================ diff --git a/primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml b/primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml new file mode 100644 index 000000000..890f6fa1c --- /dev/null +++ b/primus/configs/data/megatron/diffusion/templates/dataset_preencoded.yaml @@ -0,0 +1,79 @@ +# Local Dataset Configuration for Pre-encoded Diffusion Data +# +# This file should be placed in your pre-encoded dataset directory +# as dataset.yaml (alongside the WebDataset tar files). +# +# Directory structure: +# /path/to/dataset_preencoded/ +# ├── dataset.yaml <- This file +# ├── train-000000.tar +# ├── train-000001.tar +# ├── ... +# └── .nv-meta/ <- Created by `energon prepare` +# +# After creating this file, run: +# energon prepare /path/to/dataset_preencoded --num-workers 8 + +# Use CrudeWebdataset to enable cooker-based processing +__module__: megatron.energon +__class__: CrudeWebdataset + +# Subflavor indicates this dataset contains pre-encoded features +# This routes samples to the preencoded cooker in the task encoder +subflavors: + encoding: preencoded + +# Optional: Filter which tar files to load +# Uncomment and adjust to use specific shards +# part_filter: +# - train-*.tar # Load all training shards +# - val-*.tar # Or load validation shards + +# Optional: Sample filtering +# max_samples: 100000 # Limit total samples (useful for testing) + +# ============================================================================ +# Expected WebDataset Format +# ============================================================================ +# +# Each sample in the tar files should contain: +# - __key__: unique identifier (e.g., "sample_000001") +# - latent.npy: VAE-encoded image latents [C, H, W] float32 +# - t5_embed.npy: T5-XXL text embeddings [seq_len, 4096] float32 +# - clip_embed.npy: CLIP-L pooled embeddings [768] float32 +# - metadata.json: sample metadata (optional) +# +# Example tar content: +# sample_000001.latent.npy +# sample_000001.t5_embed.npy +# sample_000001.clip_embed.npy +# sample_000001.metadata.json +# +# ============================================================================ +# Preparation Steps +# ============================================================================ +# +# 1. Prepare pre-encoded features using the Primus CLI: +# tools/docker/primus data diffusion-encoded \ +# --source-type directory \ +# --input-dir /path/to/raw_images \ +# --output-dir /path/to/dataset_preencoded \ +# --model-path black-forest-labs/FLUX.1-dev \ +# --batch-size 8 +# +# 2. Copy this template to the output directory as dataset.yaml +# +# 3. Run Energon indexing: +# energon prepare /path/to/dataset_preencoded --num-workers 8 +# +# 4. Verify the dataset: +# energon info /path/to/dataset_preencoded +# +# 5. Use in training config: +# data: +# dataset_path: /path/to/dataset_preencoded/dataset.yaml +# +# For more information, see: +# - primus/configs/data/megatron/diffusion/README.md +# - examples/megatron/diffusion/README.md +# diff --git a/primus/configs/data/megatron/diffusion/templates/dataset_preencoded_numpy.yaml b/primus/configs/data/megatron/diffusion/templates/dataset_preencoded_numpy.yaml new file mode 100644 index 000000000..30cd15051 --- /dev/null +++ b/primus/configs/data/megatron/diffusion/templates/dataset_preencoded_numpy.yaml @@ -0,0 +1,16 @@ +# Dataset Configuration for Pre-encoded NumPy Diffusion Data (MLPerf format) +# +# This file is auto-generated by finalize_energon_dataset(encoding="preencoded_numpy") +# and placed in /.nv-meta/dataset.yaml. +# +# Data format: numpy-serialized bfloat16 tensors stored as .bytes entries +# in WebDataset tar shards. Each sample contains: +# - t5.bytes: T5 text embeddings (numpy uint16, viewed as bfloat16) +# - clip.bytes: CLIP pooled embeddings (numpy uint16, viewed as bfloat16) +# - mean.bytes: VAE posterior mean (numpy uint16, viewed as bfloat16) +# - logvar.bytes: VAE posterior log-variance (numpy uint16, viewed as bfloat16) + +__module__: megatron.energon +__class__: CrudeWebdataset +subflavors: + encoding: preencoded_numpy diff --git a/primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml b/primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml new file mode 100644 index 000000000..d24eef25f --- /dev/null +++ b/primus/configs/data/megatron/diffusion/templates/dataset_raw.yaml @@ -0,0 +1,112 @@ +# Local Dataset Configuration for Raw Diffusion Data (On-the-fly Encoding) +# +# This file should be placed in your raw dataset directory +# as dataset.yaml (alongside the WebDataset tar files). +# +# Directory structure: +# /path/to/dataset_raw/ +# ├── dataset.yaml <- This file +# ├── train-000000.tar +# ├── train-000001.tar +# ├── ... +# └── .nv-meta/ <- Created by `energon prepare` +# +# After creating this file, run: +# energon prepare /path/to/dataset_raw --num-workers 8 + +# Use CrudeWebdataset to enable cooker-based processing +__module__: megatron.energon +__class__: CrudeWebdataset + +# Subflavor indicates this dataset contains raw images/text +# This routes samples to the on-the-fly encoding cooker in the task encoder +subflavors: + encoding: raw + +# Optional: Filter which tar files to load +# Uncomment and adjust to use specific shards +# part_filter: +# - train-*.tar # Load all training shards + +# Optional: Sample filtering +# max_samples: 10000 # Limit total samples (useful for testing) + +# ============================================================================ +# Expected WebDataset Format +# ============================================================================ +# +# Each sample in the tar files should contain: +# - __key__: unique identifier (e.g., "sample_000001") +# - image.jpg: raw image file (JPEG/PNG/WebP) +# - caption.txt: text caption describing the image +# - metadata.json: sample metadata (optional) +# +# Example tar content: +# sample_000001.jpg +# sample_000001.txt +# sample_000001.metadata.json +# +# Image requirements: +# - Format: JPEG, PNG, or WebP +# - Resolution: Any (will be resized to target size during training) +# - Color: RGB (grayscale will be converted) +# +# Caption requirements: +# - Format: Plain text UTF-8 +# - Length: Any (will be truncated to model max length) +# - Content: Descriptive text for the image +# +# ============================================================================ +# Preparation Steps +# ============================================================================ +# +# 1. Prepare raw dataset using the Primus CLI: +# tools/docker/primus data diffusion-raw \ +# --source-type directory \ +# --input-dir /path/to/raw_images \ +# --output-dir /path/to/dataset_raw +# +# 2. Copy this template to the output directory as dataset.yaml +# +# 3. Run Energon indexing: +# energon prepare /path/to/dataset_raw --num-workers 8 +# +# 4. Verify the dataset: +# energon info /path/to/dataset_raw +# +# 5. Configure encoders in your training config: +# data: +# dataset_path: /path/to/dataset_raw/dataset.yaml +# encoder_configs: +# vae: +# model_path: black-forest-labs/FLUX.1-dev +# device: cuda +# precision: bf16 +# t5: +# model_path: black-forest-labs/FLUX.1-dev +# device: cuda +# precision: bf16 +# clip: +# model_path: black-forest-labs/FLUX.1-dev +# device: cuda +# precision: bf16 +# +# ============================================================================ +# Performance Notes +# ============================================================================ +# +# On-the-fly encoding is slower than pre-encoded mode: +# - Encoding adds ~100-200ms per image on MI300X +# - Requires additional GPU memory for encoders (~10-15GB) +# - Recommended for experimentation, not production +# +# For production training, use pre-encoded mode: +# 1. Pre-encode your dataset once using diffusion-encoded +# 2. Use dataset_preencoded.yaml configuration +# 3. Enjoy 2-3x faster training +# +# For more information, see: +# - primus/configs/data/megatron/diffusion/README.md +# - examples/megatron/diffusion/README.md +# - primus/configs/models/megatron/diffusion/encoders.yaml +# diff --git a/primus/configs/data/megatron/diffusion/templates/metadataset.yaml b/primus/configs/data/megatron/diffusion/templates/metadataset.yaml new file mode 100644 index 000000000..b6563b368 --- /dev/null +++ b/primus/configs/data/megatron/diffusion/templates/metadataset.yaml @@ -0,0 +1,81 @@ +# Metadataset Configuration for Mixed Training +# +# Metadataset combines multiple datasets with different weights. +# This is useful for: +# - Mixing pre-encoded and raw datasets +# - Combining multiple data sources +# - Creating train/val splits from different datasets +# +# This file can be used directly or referenced in training configs: +# data: +# dataset_path: /path/to/metadataset.yaml + +# ============================================================================ +# Path Configuration +# ============================================================================ +# Dataset paths should point to directories on mounted volumes for persistence. +# Default Docker mount: /workspace/Primus/data/ (persists to host) +# +# Adjust paths below to match your dataset locations. +# All paths should be absolute and point to prepared datasets. +# ============================================================================ + +# Use Energon's Metadataset class +__module__: megatron.energon +__class__: Metadataset + +splits: + # Training split with mixed data sources + train: + datasets: + # 80% pre-encoded data (faster training) + - weight: 0.8 + path: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/diffusion_preencoded/ + # This dataset should have: subflavors: {encoding: preencoded} + + # 20% raw data (flexibility for new data) + - weight: 0.2 + path: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/diffusion_raw/ + # This dataset should have: subflavors: {encoding: raw} + + # Validation split (typically all pre-encoded for speed) + val: + datasets: + - weight: 1.0 + path: ${PRIMUS_DATA_ROOT:/workspace/Primus/data}/diffusion_preencoded_val/ + # This dataset should have: subflavors: {encoding: preencoded} + +# Global settings applied to all datasets +shuffle_buffer_size: 1000 +max_samples_per_sequence: 100 + +# ============================================================================ +# Usage Notes +# ============================================================================ +# +# Each dataset path should point to a directory containing: +# 1. WebDataset tar files (*.tar) +# 2. Energon index (.nv-meta/ directory) +# 3. Local dataset.yaml with correct subflavors +# +# Example structure (with default Docker mount): +# /workspace/Primus/data/diffusion_preencoded/ +# ├── dataset.yaml (with subflavors: {encoding: preencoded}) +# ├── train-*.tar +# └── .nv-meta/ +# +# /workspace/Primus/data/diffusion_raw/ +# ├── dataset.yaml (with subflavors: {encoding: raw}) +# ├── train-*.tar +# └── .nv-meta/ +# +# Dataset weights: +# - Weights are relative, not absolute (0.8 + 0.2 = 1.0) +# - Higher weight = more samples from that dataset +# - Weights can sum to any positive number +# +# For more information, see: +# - primus/configs/data/megatron/diffusion/README.md +# - examples/megatron/diffusion/README.md +# - docs/backends/megatron/diffusion/energon_integration.md +# diff --git a/primus/configs/models/megatron/diffusion/encoders.yaml b/primus/configs/models/megatron/diffusion/encoders.yaml new file mode 100644 index 000000000..43d771fdb --- /dev/null +++ b/primus/configs/models/megatron/diffusion/encoders.yaml @@ -0,0 +1,126 @@ +# Encoder Configuration for Flux Diffusion Models +# This file defines configurations for VAE, T5-XXL, and CLIP-L encoders + +# ============================================================================ +# Subfolder Configuration Guide +# ============================================================================ +# Different model repositories organize files differently: +# +# 1. FLUX.1-dev (complex structure): +# - VAE: model in vae/ +# - T5: model in text_encoder_2/, tokenizer in tokenizer_2/ +# - CLIP: model in text_encoder/, tokenizer in tokenizer/ +# +# 2. Standalone models (simple structure): +# - google/t5-v1_1-xxl: all files in root (subfolder: null) +# - openai/clip-vit-large-patch14: all files in root (subfolder: null) +# +# Configuration fields: +# - subfolder: Where model weights are located +# - tokenizer_subfolder: Where tokenizer files are located (text encoders only) +# If tokenizer_subfolder is null, uses subfolder value +# ============================================================================ + +# VAE Encoder (AutoencoderKL) +vae: + type: autoencoder_kl + model_path: black-forest-labs/FLUX.1-dev # HuggingFace repo or local path + subfolder: vae # Subfolder containing VAE weights + precision: bf16 # bf16, fp16, or fp32 + device: cuda + use_cached: true # Use pre-encoded latents if available + freeze_weights: true # Freeze VAE weights during training + cache_dir: null # Optional: custom cache directory + + # Flux-specific VAE parameters + scale_factor: 0.3611 # Latent scaling factor + shift_factor: 0.1159 # Latent shift factor + in_channels: 3 # RGB input + out_channels: 16 # Latent channels + latent_downsample_factor: 8 # H/8, W/8 + +# T5-XXL Text Encoder +t5: + type: t5_xxl + model_path: black-forest-labs/FLUX.1-dev + subfolder: text_encoder_2 # Subfolder containing T5 model weights + tokenizer_path: null # Uses model_path if null + tokenizer_subfolder: tokenizer_2 # Subfolder containing T5 tokenizer files + precision: bf16 + device: cuda + use_cached: true + freeze_weights: true + cache_dir: null # Optional: custom cache directory + + # T5-XXL parameters + max_length: 512 # Maximum sequence length + embedding_dim: 4096 # T5-XXL hidden size + return_pooled: false # T5 doesn't use pooled embeddings + +# CLIP-L Text Encoder +clip: + type: clip_l + model_path: black-forest-labs/FLUX.1-dev + subfolder: text_encoder # Subfolder containing CLIP model weights + tokenizer_path: null # Uses model_path if null + tokenizer_subfolder: tokenizer # Subfolder containing CLIP tokenizer files + precision: bf16 + device: cuda + use_cached: true + freeze_weights: true + cache_dir: null # Optional: custom cache directory + + # CLIP-L parameters + max_length: 77 # CLIP max sequence length + embedding_dim: 768 # CLIP-L hidden size + pooled_dim: 768 # CLIP-L pooled embedding size + return_pooled: true # CLIP returns both sequence and pooled + +# ============================================================================ +# Data Mode Configuration +# ============================================================================ +# Data mode (preencoded vs on-the-fly) is now determined by dataset subflavors +# in dataset.yaml, not by a config parameter. +# +# For preencoded datasets (RECOMMENDED): +# - Create dataset.yaml with: subflavors: {encoding: preencoded} +# - Encoders above are NOT loaded during training +# - Dataset serves precalculated latents and embeddings +# +# For raw/on-the-fly datasets: +# - Create dataset.yaml with: subflavors: {encoding: raw} +# - Encoders above ARE loaded during training +# - Images/text are encoded on-the-fly (slower but saves disk space) +# +# See: docs/backends/megatron/diffusion/data_preprocessing.md +# ============================================================================ + +# ============================================================================ +# Example: Standalone Models (simpler configuration) +# ============================================================================ +# For standalone models, subfolder and tokenizer_subfolder should be null: +# +# t5_standalone: +# type: t5_xxl +# model_path: google/t5-v1_1-xxl +# subfolder: null # All files in root +# tokenizer_subfolder: null # Not needed +# precision: bf16 +# device: cuda +# +# clip_standalone: +# type: clip_l +# model_path: openai/clip-vit-large-patch14 +# subfolder: null # All files in root +# tokenizer_subfolder: null # Not needed +# precision: bf16 +# device: cuda +# ============================================================================ + +# Notes: +# - Model paths can be: +# 1. HuggingFace repo: "black-forest-labs/FLUX.1-dev" +# 2. Local path: "/path/to/models/flux-dev" +# +# - License note: FLUX.1-dev is non-commercial only +# Make sure to comply with license restrictions diff --git a/primus/configs/models/megatron/diffusion/flux_12b.yaml b/primus/configs/models/megatron/diffusion/flux_12b.yaml new file mode 100644 index 000000000..0a3dcf556 --- /dev/null +++ b/primus/configs/models/megatron/diffusion/flux_12b.yaml @@ -0,0 +1,141 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B Configuration (Production/Standard Model) +# +# This is the full Flux configuration matching FLUX.1 [dev] and FLUX.1 [schnell] +# from Black Forest Labs (black-forest-labs/FLUX.1-dev on Hugging Face). +# +# Use cases: +# - Production image generation +# - Large-scale training +# - Fine-tuning from pretrained weights +# - Research and benchmarking +# - High-quality image synthesis + +extends: + - flux_base.yaml + +# ============================================================================== +# Architecture: Layer Counts (ONLY difference from base) +# ============================================================================== + +num_joint_layers: 19 # Multimodal (joint) transformer blocks +num_single_layers: 38 # Image-only transformer blocks +num_layers: 57 # Total layers for Megatron compatibility (joint + single) + +# Total layers: 19 + 38 = 57 transformer blocks +# Total parameters: ~12 billion +# Model matches: FLUX.1 [dev] and FLUX.1 [schnell] architectures + +# ============================================================================== +# Guidance Configuration (Optional Override) +# ============================================================================== + +# guidance_embed: false # Default from flux_base.yaml + +# To train FLUX.1 [dev] style (with CFG support): +# guidance_embed: true +# +# To train FLUX.1 [schnell] style (distilled, no CFG): +# guidance_embed: false (default) + +# ============================================================================== +# All Other Parameters Inherited from flux_base.yaml +# ============================================================================== + +# The following are inherited and should NOT be redefined here: +# - hidden_size: 3072 +# - num_attention_heads: 24 +# - context_dim: 4096 (T5-XXL) +# - vec_in_dim: 768 (CLIP-L) +# - in_channels: 64 (VAE latents) +# - activation_func: openai_gelu (from FluxConfig default) +# - add_qkv_bias: true (from FluxConfig default) +# - All other architectural parameters + +# ============================================================================== +# Hardware Requirements +# ============================================================================== + +# Training: +# - Minimum: 8x H100 80GB or 8x MI300X 192GB +# - Recommended: 8x H100 80GB with NVLink +# - Batch size: Start with global_batch_size=256, micro_batch_size=1 +# - Memory: ~70GB per GPU with micro_batch=1, TP=1 + +# Inference: +# - Minimum: 1x H100 80GB or 1x MI300X 192GB +# - Typical: 1x A100 80GB works for batch_size=1 +# - Memory: ~24GB for model weights + activations + +# ============================================================================== +# Training Recommendations +# ============================================================================== + +# Parallelism Strategy: +# - Data Parallel (DP): 8 (default, simplest) +# - Tensor Parallel (TP): 1-2 (use if memory constrained) +# - Pipeline Parallel (PP): 1 (Flux doesn't benefit from PP) +# - Sequence Parallel (SP): false (not needed for typical resolutions) + +# Training Duration: +# - From scratch: 100K-1M steps (weeks on 8x H100) +# - Fine-tuning: 10K-50K steps (hours to days) +# - LoRA adaptation: 1K-10K steps (hours) + +# Suggested Training Config: +# micro_batch_size: 1 +# global_batch_size: 256 +# learning_rate: 1.0e-4 +# weight_decay: 0.01 +# warmup_steps: 1000 +# total_steps: 100000 + +# ============================================================================== +# Comparison with 535M Variant +# ============================================================================== + +# Flux 535M (testing): +# - Layers: 1 joint + 1 single = 2 total +# - Parameters: ~535M +# - Training: Single GPU, minutes to hours +# - Use: Testing, debugging, quick iteration + +# Flux 12B (this config): +# - Layers: 19 joint + 38 single = 57 total +# - Parameters: ~12B +# - Training: Multi-GPU, days to weeks +# - Use: Production, research, deployment + +# ============================================================================== +# NeMo Compatibility +# ============================================================================== + +# This config matches NeMo's FluxConfig defaults: +# nemo/collections/diffusion/models/flux/model.py:74-114 +# +# @dataclass +# class FluxConfig(TransformerConfig, io.IOMixin): +# num_joint_layers: int = 19 +# num_single_layers: int = 38 +# hidden_size: int = 3072 +# num_attention_heads: int = 24 +# activation_func: Callable = openai_gelu +# add_qkv_bias: bool = True + +# Checkpoints are fully interoperable between Primus and NeMo + +# ============================================================================== +# Pretrained Weights +# ============================================================================== + +# Compatible with weights from: +# - black-forest-labs/FLUX.1-dev (HuggingFace) +# - black-forest-labs/FLUX.1-schnell (HuggingFace) +# +# Note: FLUX.1-dev is non-commercial license +# FLUX.1-schnell is Apache 2.0 license +# +# To load pretrained weights, use checkpoint conversion utilities in: +# primus/backends/megatron/core/models/diffusion/flux/checkpoint_converter.py diff --git a/primus/configs/models/megatron/diffusion/flux_12b_fp8.yaml b/primus/configs/models/megatron/diffusion/flux_12b_fp8.yaml new file mode 100644 index 000000000..3fd6c19d3 --- /dev/null +++ b/primus/configs/models/megatron/diffusion/flux_12b_fp8.yaml @@ -0,0 +1,122 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B with FP8 Configuration for AMD MI300X +# +# This configuration enables FP8 (8-bit floating point) training for Flux 12B, +# providing ~2x memory reduction and 1.5-2x training speedup while maintaining +# numerical stability through delayed scaling. +# +# Target Hardware: +# - AMD MI300X GPUs with ROCm 6.0+ +# - Requires Transformer Engine 2.1.0+ with ROCm backend +# +# Benefits: +# - ~2x memory reduction (activations and weights) +# - 1.5-2x training speedup on MI300X +# - Maintains numerical stability via delayed scaling +# - Allows larger batch sizes or higher resolutions +# +# Use cases: +# - Memory-constrained training scenarios +# - Faster iteration during experimentation +# - Training with limited GPU resources +# - Production training with improved efficiency + +extends: + - flux_12b.yaml + +# ============================================================================== +# FP8 Configuration +# ============================================================================== + +# Enable FP8 precision for training +fp8: "e4m3" # Use E4M3 format for all FP8 tensors + # Alternative: "hybrid" (E4M3 for activations/weights, E5M2 for grad outputs) + +# FP8 Recipe: Controls scaling strategy +fp8_recipe: "delayed" # Delayed scaling for numerical stability (recommended) + # Alternatives: "tensorwise", "blockwise", "mxfp8" + # Note: delayed scaling is most stable for training + +# Scaling Factor Configuration +fp8_margin: 0 # Margin for scaling factor computation (0 = no margin) +fp8_amax_history_len: 1024 # History window length for delayed scaling + # Larger = more stable, smaller = adapts faster +fp8_amax_compute_algo: "most_recent" # Algorithm for amax computation + # Options: "most_recent", "max" + +# FP8 Gradient Configuration +fp8_wgrad: true # Enable FP8 for weight gradients (recommended for memory) + # Set to false if numerical issues occur + +# Attention Precision (keep in higher precision for stability) +fp8_dot_product_attention: false # Keep dot product attention in BF16/FP32 +fp8_multi_head_attention: false # Keep multi-head attention in BF16/FP32 + +# ============================================================================== +# Optimization Adjustments for FP8 +# ============================================================================== + +# Gradient clipping is still important with FP8 +# May need adjustment if training becomes unstable +clip_grad: 1.0 + +# ============================================================================== +# Hardware Requirements (FP8) +# ============================================================================== + +# Training with FP8: +# - Minimum: 4x MI300X 192GB (vs 8x without FP8) +# - Recommended: 8x MI300X 192GB with RDMA +# - Batch size: Can increase by ~2x vs BF16 +# - Memory: ~35-40GB per GPU (vs ~70GB BF16) + +# Inference with FP8: +# - Minimum: 1x MI300X 192GB +# - Memory: ~12-15GB for model weights + activations (vs ~24GB BF16) + +# ============================================================================== +# Training Recommendations for FP8 +# ============================================================================== + +# Suggested Training Config (adjusted for FP8): +# micro_batch_size: 2 (vs 1 for BF16) +# global_batch_size: 256 (same as BF16) +# learning_rate: 1.0e-4 (same as BF16) +# weight_decay: 0.01 (same as BF16) +# warmup_steps: 1000 (same as BF16) +# total_steps: 100000 (same as BF16) + +# Parallelism Strategy (FP8): +# - Data Parallel (DP): 8 +# - Tensor Parallel (TP): 1 (less memory pressure with FP8) +# - Pipeline Parallel (PP): 1 +# - Sequence Parallel (SP): false + +# ============================================================================== +# Troubleshooting +# ============================================================================== + +# If training becomes unstable (NaN/Inf losses): +# 1. Increase fp8_amax_history_len to 2048 or 4096 +# 2. Set fp8_wgrad: false (use higher precision for weight gradients) +# 3. Increase fp8_margin to 1 or 2 +# 4. Try fp8_amax_compute_algo: "max" instead of "most_recent" +# 5. Consider keeping first/last layers in BF16 (set first_last_layers_bf16: true) + +# If insufficient memory even with FP8: +# - Enable gradient checkpointing (recompute_granularity: "selective") +# - Increase tensor parallelism (tensor_model_parallel_size: 2) +# - Reduce batch size further + +# ============================================================================== +# Verification +# ============================================================================== + +# Before production use, verify: +# - Transformer Engine FP8 support: Check with setup_fp8_rocm.sh +# - Forward pass stability: Run test_flux_precision.py::test_flux_fp8_forward +# - Training convergence: Compare loss curves with BF16 baseline +# - Memory usage: Should be ~50% of BF16 training +# - Training speed: Should be 1.5-2x faster than BF16 diff --git a/primus/configs/models/megatron/diffusion/flux_12b_rope_fusion.yaml b/primus/configs/models/megatron/diffusion/flux_12b_rope_fusion.yaml new file mode 100644 index 000000000..e25d902b2 --- /dev/null +++ b/primus/configs/models/megatron/diffusion/flux_12b_rope_fusion.yaml @@ -0,0 +1,50 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 12B with RoPE Fusion Optimization +# +# This configuration enables Transformer Engine's fused RoPE kernels for improved performance. +# +# ============================================================================== +# CRITICAL REQUIREMENT: Same-Resolution Batches +# ============================================================================== +# +# RoPE fusion requires ALL images in each training batch to have the SAME +# resolution (height and width). This constraint exists because the fused kernel +# expects position IDs with shape [1, H*W/4, 3] that broadcast across the batch. +# +# Variable-resolution batches will produce INCORRECT positional encodings and +# degrade model quality. You will receive a warning at model initialization. +# +# Recommended data pipeline configurations: +# 1. Fixed resolution: All images resized to same size (e.g., 512x512, 1024x1024) +# 2. Resolution bucketing: Group images by resolution in separate batches +# 3. Aspect ratio bucketing: Use bucketing with consistent dimensions per batch +# +# Technical details: +# - Follows NVIDIA's MLPerf Flux implementation strategy (hsg_ngpu1152 config) +# - Uses batch_size=1 for position ID generation (broadcasting handles actual batch) +# - Compatible with both AMD ROCm and NVIDIA CUDA Transformer Engine forks +# - Fused kernel requires freqs tensor shape: [S, 1, 1, D] +# +# ============================================================================== + +extends: + - flux_12b.yaml + +# ============================================================================== +# RoPE Fusion Optimization - ENABLED +# ============================================================================== + +apply_rope_fusion: true +rotary_interleaved: true +position_embedding_type: rope # Prevents Megatron validate_args from overriding apply_rope_fusion + +# ============================================================================== +# Transformer Implementation +# ============================================================================== +# Options: +# - transformer_impl: "transformer_engine" (default) - Use TransformerEngine +# - transformer_impl: "local" - Use native Megatron + Primus Turbo (NO TransformerEngine) +# ============================================================================== +# Override in training config if needed (default is "transformer_engine") diff --git a/primus/configs/models/megatron/diffusion/flux_535m.yaml b/primus/configs/models/megatron/diffusion/flux_535m.yaml new file mode 100644 index 000000000..28293cb65 --- /dev/null +++ b/primus/configs/models/megatron/diffusion/flux_535m.yaml @@ -0,0 +1,85 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 535M Configuration (Minimal Testing Variant) +# +# This is a scaled-down Flux configuration with only 1 joint + 1 single layer +# for rapid prototyping, testing, and validation. +# +# Use cases: +# - Unit testing and CI/CD +# - Quick iteration during development +# - Validating training pipeline +# - Comparing against NeMo baseline +# - Single GPU training/testing + +extends: + - flux_base.yaml + +# ============================================================================== +# Architecture: Layer Counts (ONLY difference from base) +# ============================================================================== + +num_joint_layers: 1 # Multimodal (joint) transformer blocks +num_single_layers: 1 # Image-only transformer blocks +num_layers: 2 # Total layers for Megatron compatibility (joint + single) + +# Total layers: 1 + 1 = 2 (vs. 57 for Flux 12B) +# Total parameters: ~535 million (vs. ~12 billion for Flux 12B) + +# ============================================================================== +# All Other Parameters Inherited from flux_base.yaml +# ============================================================================== + +# The following are inherited and should NOT be redefined here: +# - hidden_size: 3072 +# - num_attention_heads: 24 +# - context_dim: 4096 +# - vec_in_dim: 768 +# - in_channels: 64 +# - activation_func: openai_gelu (from FluxConfig default) +# - add_qkv_bias: true (from FluxConfig default) +# - All other architectural parameters + +# ============================================================================== +# Training Recommendations +# ============================================================================== + +# Suggested training config for 535M: +# - micro_batch_size: 1-2 +# - global_batch_size: 8-16 +# - GPUs: 1-2 (fits on single H100 80GB) +# - Training time: ~10-100x faster than 12B +# - Use for: sanity checks, debugging, quick experiments + +# ============================================================================== +# Comparison with Full Flux +# ============================================================================== + +# Flux 535M (this config): +# - Layers: 1 joint + 1 single = 2 total +# - Parameters: ~535M +# - Training: Single GPU capable +# - Purpose: Testing/development + +# Flux 12B (production): +# - Layers: 19 joint + 38 single = 57 total +# - Parameters: ~12B +# - Training: Multi-GPU required +# - Purpose: Production deployment + +# ============================================================================== +# NeMo Compatibility +# ============================================================================== + +# This config matches NeMo's unit_test() recipe: +# nemo/scripts/flux/flux_training.py:252-278 +# +# recipe.model.flux_params.flux_config = run.Config( +# FluxConfig, +# num_joint_layers=1, +# num_single_layers=1, +# ... +# ) + +# Checkpoints are interoperable between Primus and NeMo diff --git a/primus/configs/models/megatron/diffusion/flux_535m_fp8.yaml b/primus/configs/models/megatron/diffusion/flux_535m_fp8.yaml new file mode 100644 index 000000000..ff31f5150 --- /dev/null +++ b/primus/configs/models/megatron/diffusion/flux_535m_fp8.yaml @@ -0,0 +1,107 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux 535M with FP8 Configuration (Testing/Development) +# +# This is a minimal Flux variant with FP8 enabled, designed for: +# - Quick validation of FP8 functionality +# - Testing FP8 integration before scaling to 12B +# - Debugging FP8-related issues +# - Development and experimentation +# +# Target Hardware: +# - Single AMD MI300X GPU with ROCm 6.0+ +# - Requires Transformer Engine 2.1.0+ with ROCm backend +# +# Use cases: +# - Validating FP8 setup and configuration +# - Quick iteration during FP8 feature development +# - Testing FP8 numerical stability +# - Verifying Transformer Engine compatibility + +extends: + - flux_535m.yaml + +# ============================================================================== +# FP8 Configuration (Same as flux_12b_fp8.yaml) +# ============================================================================== + +# Enable FP8 precision for training +fp8: "e4m3" # Use E4M3 format for all FP8 tensors + +# FP8 Recipe: Controls scaling strategy +fp8_recipe: "delayed" # Delayed scaling for numerical stability (recommended) + +# Scaling Factor Configuration +fp8_margin: 0 # Margin for scaling factor computation +fp8_amax_history_len: 1024 # History window length for delayed scaling +fp8_amax_compute_algo: "most_recent" # Algorithm for amax computation + +# FP8 Gradient Configuration +fp8_wgrad: true # Enable FP8 for weight gradients + +# Attention Precision (keep in higher precision for stability) +fp8_dot_product_attention: false # Keep dot product attention in BF16/FP32 +fp8_multi_head_attention: false # Keep multi-head attention in BF16/FP32 + +# ============================================================================== +# Optimization Adjustments for FP8 +# ============================================================================== + +# Gradient clipping +clip_grad: 1.0 + +# ============================================================================== +# Hardware Requirements (FP8) +# ============================================================================== + +# Training with FP8: +# - Minimum: 1x MI300X 192GB +# - Memory: ~5-7GB per GPU (vs ~10-15GB BF16) +# - Training time: Minutes (vs hours for 12B) + +# Inference with FP8: +# - Minimum: 1x MI300X 192GB +# - Memory: ~2-3GB for model weights + activations + +# ============================================================================== +# Training Recommendations for FP8 +# ============================================================================== + +# Suggested Training Config (testing): +# micro_batch_size: 4 (vs 2 for BF16) +# global_batch_size: 32 +# learning_rate: 1.0e-4 +# weight_decay: 0.01 +# warmup_steps: 100 +# total_steps: 1000 (for quick validation) + +# Parallelism Strategy: +# - Data Parallel (DP): 1 (single GPU testing) +# - Tensor Parallel (TP): 1 (not needed for 535M) +# - Pipeline Parallel (PP): 1 +# - Sequence Parallel (SP): false + +# ============================================================================== +# Testing Strategy +# ============================================================================== + +# Use this config to verify: +# 1. FP8 forward pass works without NaN/Inf +# 2. FP8 backward pass computes gradients correctly +# 3. Training loop runs stable for multiple iterations +# 4. Memory usage is reduced as expected +# 5. Training speed improvement is measurable + +# Run tests with: +# pytest tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py -v + +# ============================================================================== +# Troubleshooting +# ============================================================================== + +# Same troubleshooting steps as flux_12b_fp8.yaml: +# - Increase fp8_amax_history_len if unstable +# - Set fp8_wgrad: false if gradient issues +# - Adjust fp8_margin if scaling problems +# - Use fp8_amax_compute_algo: "max" for more conservative scaling diff --git a/primus/configs/models/megatron/diffusion/flux_base.yaml b/primus/configs/models/megatron/diffusion/flux_base.yaml new file mode 100644 index 000000000..5f76aae7a --- /dev/null +++ b/primus/configs/models/megatron/diffusion/flux_base.yaml @@ -0,0 +1,110 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Flux Base Configuration +# Common parameters for all Flux model variants (535M, 12B, custom) + +extends: + - ../diffusion_model.yaml + +# ============================================================================== +# Core Architecture (Flux Standard) +# ============================================================================== + +# Layer configuration (Flux uses joint + single layers, not standard num_layers) +# num_layers will be computed in code as num_joint_layers + num_single_layers +num_joint_layers: null # To be set by specific variant (535M, 12B, etc.) +num_single_layers: null # To be set by specific variant (535M, 12B, etc.) + +# Dimensions (same across all Flux variants) +hidden_size: 3072 # Transformer hidden dimension +num_attention_heads: 24 # Number of attention heads +ffn_hidden_size: null # Auto-computed as hidden_size * 4 = 12288 + +# Head dimension (computed): hidden_size / num_attention_heads = 128 + +# ============================================================================== +# Context Dimensions (Text Encoder Integration) +# ============================================================================== + +context_dim: 4096 # T5-XXL hidden dimension (text embeddings) +vec_in_dim: 768 # CLIP-L pooled dimension (pooled text) +model_channels: 256 # Channels for timestep embedding MLPs + +# ============================================================================== +# Input/Output Channels +# ============================================================================== + +in_channels: 64 # VAE latent channels (16 * 4 from 2x2 packing) +out_channels: 64 # Output channels (same as input for latent space) + +# ============================================================================== +# Patchification +# ============================================================================== + +patch_size: 1 # Flux uses 1x1 patches (no spatial patchification) + +# ============================================================================== +# Position Embeddings (3D RoPE) +# ============================================================================== + +theta: 10000 # Base frequency for RoPE +axes_dim: [16, 56, 56] # 3D RoPE dimensions: [channels, height, width] + # Sum = 128 = head_dim + +# ============================================================================== +# Transformer Configuration +# ============================================================================== + +# LayerNorm +layernorm_epsilon: 1.0e-6 # Epsilon for numerical stability + +# Residual connections +apply_residual_connection_post_layernorm: false # Flux uses pre-norm + +# Dropout (Flux typically uses none) +hidden_dropout: 0.0 +attention_dropout: 0.0 + +# ============================================================================== +# Guidance (Classifier-Free Guidance) +# ============================================================================== + +guidance_embed: false # Set true to enable CFG embedding + # FLUX.1 [dev]: true (supports CFG) + # FLUX.1 [schnell]: false (distilled, no CFG) + +# ============================================================================== +# CRITICAL: Activation Function & QKV Configuration +# ============================================================================== + +# Activation function: defaults to openai_gelu_no_jit (non-JIT tanh GELU, ROCm-safe; +# see FluxConfig.activation_func). Override via activation_func in your config: +# activation_func: openai_gelu # fused F.gelu(approximate="tanh") +# activation_func: erf_gelu # erf-based GELU +# Changing the activation changes numerics: keep it consistent with the activation +# used to train any checkpoint you load. + +# FluxConfig defaults that we intentionally inherit: +# activation_func: openai_gelu_no_jit (FluxConfig default, ROCm-safe) +# add_qkv_bias: true +# single_block_bias: true (new parameter) +# rotary_interleaved: true +# apply_rope_fusion: false + +# New parameter: single_block_bias +# Controls bias in single block linear layers independently from joint blocks +# Default: true (bias enabled in single blocks) +# Set to false to disable bias in single blocks only + +# These defaults are compatible with Primus training conventions + +# ============================================================================== +# Notes +# ============================================================================== + +# - This config defines Flux architecture shared across all variants +# - Specific variants (535M, 12B) only override num_joint_layers, num_single_layers +# - Activation function (openai_gelu) comes from FluxConfig class default +# - All Flux variants share same hidden_size (3072) and num_heads (24) +# - For custom variants, extend this file and override layer counts diff --git a/primus/configs/models/megatron/diffusion_model.yaml b/primus/configs/models/megatron/diffusion_model.yaml new file mode 100644 index 000000000..6af2c1583 --- /dev/null +++ b/primus/configs/models/megatron/diffusion_model.yaml @@ -0,0 +1,135 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +# Base configuration for diffusion models +# This file defines ARCHITECTURE ONLY - training params go in trainer config + +extends: + - primus_megatron_model.yaml + +# ============================================================================== +# Model Parallelism +# ============================================================================== + +# Parallelism configuration +model_parallel_size: null +tensor_model_parallel_size: 1 +encoder_tensor_model_parallel_size: 0 +context_parallel_size: 1 +cp_comm_type: p2p # p2p, a2a, allgather or a2a+p2p +hierarchical_context_parallel_sizes: null + +pipeline_model_parallel_size: 1 +pipeline_model_parallel_layout: null +pipeline_model_parallel_comm_backend: null # str: nccl, ucc +encoder_pipeline_model_parallel_size: 0 +pipeline_model_parallel_split_rank: null +decoder_first_pipeline_num_layers: null # int +decoder_last_pipeline_num_layers: null # int +num_layers_per_virtual_pipeline_stage: null # int +num_virtual_stages_per_pipeline_rank: null # int +microbatch_group_size_per_vp_stage: null # int +sequence_parallel: true +expert_model_parallel_size: 1 +expert_tensor_parallel_size: null # int + +# ============================================================================== +# Architecture Parameters +# ============================================================================== + +# Model tying (required for FSDP2) +untie_embeddings_and_output_weights: true + +# Layer configuration (will be overridden by specific model configs) +num_layers: null # Must be set by specific diffusion model configs +encoder_num_layers: null +decoder_num_layers: null + +# Position embeddings (diffusion models may use different positional encoding) +max_position_embeddings: 8192 # Maximum sequence length + +# Tokenizer (diffusion models don't use tokenizers - text is pre-encoded) +tokenizer_type: NullTokenizer # Placeholder tokenizer for diffusion models +vocab_size: 1 # Placeholder vocab size + +# Initialization (architecture-related) +init_method_std: 0.02 # Standard deviation for weight initialization + +# MoE settings (diffusion models typically don't use MoE, but fields needed for compatibility) +num_experts: null +moe_layer_freq: 1 # int +moe_ffn_hidden_size: null # int +moe_shared_expert_overlap: false +moe_shared_expert_intermediate_size: null # int +moe_grouped_gemm: false +moe_router_load_balancing_type: "aux_loss" +moe_router_dtype: null +moe_router_score_function: softmax +moe_router_topk: 2 +moe_router_pre_softmax: false +moe_router_num_groups: null +moe_router_group_topk: null +moe_router_topk_scaling_factor: null +moe_router_enable_expert_bias: false +moe_router_bias_update_rate: 1.0e-03 +moe_use_legacy_grouped_gemm: false +moe_aux_loss_coeff: 0.0 +moe_z_loss_coeff: null +moe_input_jitter_eps: null +moe_token_dispatcher_type: allgather +moe_enable_deepep: false +moe_per_layer_logging: false +moe_expert_capacity_factor: null +moe_pad_expert_input_to_capacity: false +moe_token_drop_policy: probs +moe_layer_recompute: false +moe_extended_tp: false +moe_use_upcycling: false +moe_permute_fusion: false +disable_primus_topk_router: false +moe_router_force_load_balancing: false +use_deprecated_20241209_moe_layer: false +delay_wgrad_compute: false + +# Tensor parallelism communication overlap settings +tp_comm_overlap: false +tp_comm_overlap_cfg: null +tp_comm_overlap_ag: true +tp_comm_overlap_rs: true +tp_comm_overlap_rs_dgrad: false +tp_comm_bulk_wgrad: true +tp_comm_bulk_dgrad: true + +# Optimization flags +gradient_accumulation_fusion: true # Should be disabled for FSDP2 + +# Other compatibility fields +fused_padded_mla_attention: false + +# ============================================================================== +# Important: Tokenizer Not Applicable +# ============================================================================== + +# Diffusion models process image latents, not text tokens +# DO NOT define: tokenizer_type, tokenizer_model, vocab_size, etc. +# These parameters are only for language models + +# ============================================================================== +# Training Parameters Removed +# ============================================================================== + +# The following were REMOVED and belong in trainer config or overrides: +# - seq_length, micro_batch_size, global_batch_size (trainer config) +# - learning_rate, min_lr, weight_decay, clip_grad (trainer config) +# - scheduler_type (trainer config) +# - bf16, fp16 (trainer config) +# +# Override these in your training config's overrides section as needed. + +# ============================================================================== +# Notes +# ============================================================================== + +# - This config provides architecture defaults for all diffusion models +# - Specific architectures (Flux, DiT) extend this and add their parameters +# - Training parameters should be set in pre_trainer.yaml or overrides diff --git a/primus/configs/modules/megatron/torch_compile.yaml b/primus/configs/modules/megatron/torch_compile.yaml new file mode 100644 index 000000000..8b5550feb --- /dev/null +++ b/primus/configs/modules/megatron/torch_compile.yaml @@ -0,0 +1,13 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +# Torch Compile Configuration +# Default values for torch.compile settings (Primus-specific) +torch_compile: + enable: false + backend: "inductor" + mode: "default" + fullgraph: false diff --git a/primus/configs/modules/megatron/trainer_base.yaml b/primus/configs/modules/megatron/trainer_base.yaml index ba89bd7d1..003e54eec 100755 --- a/primus/configs/modules/megatron/trainer_base.yaml +++ b/primus/configs/modules/megatron/trainer_base.yaml @@ -2,6 +2,7 @@ extends: - ../module_base.yaml - primus_megatron_module.yaml - primus_turbo.yaml + - torch_compile.yaml - zero_bubble.yaml - primus_pipeline.yaml @@ -190,6 +191,7 @@ gradient_reduce_div_fusion: true suggested_communication_unit_size: 400000000 # int keep_fp8_transpose_cache_when_using_custom_fsdp: false num_distributed_optimizer_instances: 1 # int +data_parallel_replicate_degree: 1 use_torch_fsdp2: false nccl_communicator_config_path: null use_tp_pp_dp_mapping: false @@ -235,6 +237,20 @@ data_args_path: null # str per_split_data_args_path: null # str data_cache_path: null mock_data: false +# mock_dataset: Configuration for synthetic dataset (used when mock_data=true) +# class: Fully qualified class name (optional, defaults based on model_type) +# params: Dict of parameters passed to dataset constructor +# Example: +# mock_dataset: +# class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset" +# params: +# num_samples: 1000 +# image_size: 512 +mock_dataset: + class: null + params: + num_samples: 1000 + image_size: 512 merge_file: null seq_length: 4096 encoder_seq_length: null @@ -269,6 +285,7 @@ log_avg_reset_interval: 10 log_params_norm: false log_num_zeros_in_grad: false log_throughput: false +wall_clock_step_timer: false log_progress: false timing_log_level: 0 timing_log_option: minmax From f8bee51085ab3883b7e49d6f6d4803a3eb2eb7c6 Mon Sep 17 00:00:00 2001 From: HuangWei-95 Date: Wed, 8 Jul 2026 09:31:34 +0800 Subject: [PATCH 009/127] feat(megatron): migrate MLPerf GPT-OSS-20B pretrain trainer & patches into Primus (#847) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Migrate the MLPerf GPT-OSS-20B pretraining flow and its optimizations from the standalone mlperf source tree into Primus, so it runs through the native `primus-cli ... train pretrain` path (`stage: mlperf_pretrain`) instead of a separate entrypoint/wheel. - **MLPerf trainer & logging** integrated into the BaseTrainer architecture (`primus/backends/megatron/mlperf/`: `mlperf_pretrain_trainer.py`, `mlperf_logger.py`, `warmup.py`), registered as the `mlperf_pretrain` stage. - **Source patches migrated to `register_patch`** (`primus/backends/megatron/patches/`): MoE skip-identity-sort, SDMA param all-gather, TE BSHD-layout, turbo fused-residual-norm. - Honor `MLLOG_TRAIN_LOSS_LOG_FREQ`; add MLPerf log suppression (`mlperf_log_suppression.py`); fix a tensor-keyed `WeakKeyDictionary` in the MoE skip-identity-sort patch. ## Changes - 13 files, +2926 (additive). New `primus/backends/megatron/mlperf/` and `primus/backends/megatron/patches/{moe,parallelism,te,turbo}_patches` modules; `sdma_param_gather.py`, `fused_residual_rmsnorm.py`; `cli/main.py` wiring. ## Test plan - [x] End-to-end on MI355X (1 node × 8 GPUs), image `tasimage/primus:pr-830`, config `gpt_oss_20B-pretrain-fp8.yaml`, EP=1, fp8(e4m3, tensorwise), `use_turbo_grouped_gemm=false` (TE grouped GEMM). - [x] Trains cleanly, ~580 TFLOP/s/GPU, no NaN; **train loss 11.85 → 3.34**, **eval loss 4.57 → 3.35** (approaching the MLPerf target 3.34); eval + `:::MLLOG` events emitted correctly. ## Notes / known limitation (not in this PR) - With `use_turbo_grouped_gemm=true` on gfx950 (MI350/MI355), the Primus-Turbo fp8 tensorwise grouped-GEMM backward hits `K mismatch (5760 vs 2880)` on the non-square expert fc1, because the gfx950 NT-layout backward consumes the extension's pre-quantized `b_t` (col-wise, non-transposed) directly. The raw op is fine standalone; this is a framework/turbo interop issue tracked separately. TE grouped GEMM is the working fp8 path on MI355 for now. --------- Co-authored-by: HuangWei-95 Co-authored-by: Cursor Co-authored-by: Wei Huang Co-authored-by: Wei Huang --- .../gpt_oss_20B-FP8-mlperf-pretrain.yaml | 222 +++++ primus/backends/megatron/__init__.py | 2 + .../core/distributed/sdma_param_gather.py | 248 ++++++ .../core/extensions/fused_residual_rmsnorm.py | 453 ++++++++++ primus/backends/megatron/mlperf/__init__.py | 26 + .../backends/megatron/mlperf/mlperf_logger.py | 159 ++++ .../mlperf/mlperf_pretrain_trainer.py | 417 ++++++++++ primus/backends/megatron/mlperf/warmup.py | 776 ++++++++++++++++++ .../moe_patches/skip_identity_sort_patches.py | 121 +++ .../sdma_param_all_gather_patches.py | 200 +++++ .../patches/te_patches/bshd_layout_patches.py | 145 ++++ .../turbo/fused_residual_norm_patches.py | 72 ++ primus/cli/main.py | 6 + primus/mlperf_log_suppression.py | 273 ++++++ requirements.txt | 1 + 15 files changed, 3121 insertions(+) create mode 100644 examples/megatron/configs/MI355X/gpt_oss_20B-FP8-mlperf-pretrain.yaml create mode 100644 primus/backends/megatron/core/distributed/sdma_param_gather.py create mode 100644 primus/backends/megatron/core/extensions/fused_residual_rmsnorm.py create mode 100644 primus/backends/megatron/mlperf/__init__.py create mode 100644 primus/backends/megatron/mlperf/mlperf_logger.py create mode 100644 primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py create mode 100644 primus/backends/megatron/mlperf/warmup.py create mode 100644 primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py create mode 100644 primus/backends/megatron/patches/parallelism/sdma_param_all_gather_patches.py create mode 100644 primus/backends/megatron/patches/te_patches/bshd_layout_patches.py create mode 100644 primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py create mode 100644 primus/mlperf_log_suppression.py diff --git a/examples/megatron/configs/MI355X/gpt_oss_20B-FP8-mlperf-pretrain.yaml b/examples/megatron/configs/MI355X/gpt_oss_20B-FP8-mlperf-pretrain.yaml new file mode 100644 index 000000000..5b62b5265 --- /dev/null +++ b/examples/megatron/configs/MI355X/gpt_oss_20B-FP8-mlperf-pretrain.yaml @@ -0,0 +1,222 @@ +work_group: ${TEAM:amd} +user_name: ${USER:root} +exp_name: ${EXP_NAME:gpt_oss_20b} +workspace: ./output + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: ${PRIMUS_MODEL:gpt_oss_20B}.yaml + overrides: + + # Activate the migrated MLPerf pretrain trainer (mllog + MLPerf hooks). + stage: mlperf_pretrain + + # tokenizer + tokenizer_type: Llama3Tokenizer + tokenizer_model: ${MODEL:meta-llama/Llama-3.1-8B} + + # model + num_layers: 24 + hidden_size: 2880 + ffn_hidden_size: 2880 + num_attention_heads: 64 + num_query_groups: 8 # Group Query Attention (GQA) - matches HF num_key_value_heads + num_experts: 32 + activation_func: swiglu # SiLU activation (matches HF hidden_act: "silu") + + # rotary + position_embedding_type: rope + rotary_base: 150000 + + # mixed-precision + attention_softmax_in_fp32: false + grad_reduce_in_bf16: ${PRIMUS_GRAD_REDUCE_IN_BF16:true} + + # log + wandb_project: "Primus_GPT_OSS_20B" + stderr_sink_level: DEBUG + log_interval: ${LOG_INTERVAL:10} + + # debug + # moe_router_force_load_balancing: true + # log_avg_skip_iterations: 2 + # log_avg_reset_interval: 50 + + # profile + profile: ${PRIMUS_PROFILE:false} + 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] + + # enable fp8 training + fp8: e4m3 + fp8_recipe: tensorwise + clip_grad: 1.0 # Gradient clipping (already default, but explicit) + check_for_nan_in_loss_and_grad: false + + # hyper parameters + train_iters: ${PRIMUS_TRAIN_ITERS:1200000} + micro_batch_size: ${PRIMUS_MICRO_BATCH_SIZE:2} + global_batch_size: ${PRIMUS_GLOBAL_BATCH_SIZE:16} + seq_length: ${PRIMUS_SEQ_LENGTH:8192} + max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:131072} + seed: ${SEED:1234} # Random seed for reproducibility + lr: ${PRIMUS_LR:8.0e-4} # Reduced from 8e-4 for FP8 stability + min_lr: ${PRIMUS_MIN_LR:8.0e-5} # Set to 10% of max LR + lr_warmup_iters: ${PRIMUS_LR_WARMUP_ITERS:128} + lr_decay_iters: ${PRIMUS_LR_DECAY_ITERS:1199872} + lr_decay_style: cosine + weight_decay: 0.1 + optimizer: adam + use_distributed_optimizer: true # use distributed optimizer + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-5 + eod_mask_loss: true + init_method_std: 0.008 + norm_epsilon: 1.0e-6 + layernorm_epsilon: 1.0e-05 # RMSNorm epsilon (matches HF rms_norm_eps) + + # Dropout (disabled for training) + hidden_dropout: 0.0 + attention_dropout: 0.0 + + # parallel + tensor_model_parallel_size: ${PRIMUS_TP:1} + pipeline_model_parallel_size: ${PRIMUS_PP:1} + expert_model_parallel_size: ${PRIMUS_EP:8} + overlap_grad_reduce: true + overlap_param_gather: true + ddp_num_buckets: 8 + ddp_average_in_collective: true + + # data + mock_data: false + num_workers: ${PRIMUS_NUM_WORKERS:0} + train_data_path: "10 /data/c4-train.en_6_text_document" + valid_data_path: "/data/c4-validation-91205-samples.en_text_document" + test_data_path: "/data/c4-validation-91205-samples.en_text_document" + # Avoid copying a dense (B, 1, S, S) CPU attention mask every step. + # TE receives causal/sliding-window metadata from attn_mask_type + window_size. + # Use numeric 0/1 because Primus env expansion only type-casts numbers. + create_attention_mask_in_dataloader: ${PRIMUS_CREATE_ATTENTION_MASK_IN_DATALOADER:0} + + # fusion + moe_permute_fusion: true + gradient_accumulation_fusion: true + moe_use_legacy_grouped_gemm: false # Sync-Free MoE stage 2 or 3 require PrimusTurboGroupedMLP, please set `moe_use_legacy_grouped_gemm=True + moe_use_fused_router_with_aux_score: true + multi_latent_attention: false # Flag config.ENABLE_EXPERIMENTAL not enabled + apply_rope_fusion: true + + + # sliding window attention (GPT-OSS-20B model definition; matches HF sliding_window: 128) + # use_turbo_attention is false so non-turbo attention (which supports sliding window) is used. + # Pattern: alternating sliding_attention (1) and full_attention (0) for 24 layers + # window_size must be a tuple (left_window, right_window) for Transformer Engine + # For causal attention: left = past tokens, right = 0 (no future tokens) + # HF sliding_window: 128 means 128 past tokens, so use (128, 0) + window_size: [128, 0] # Left window: 128 past tokens, Right: 0 (causal) + window_attn_skip_freq: [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0] + + # MoE settings + moe_apply_probs_on_input: false + moe_aux_loss_coeff: 0.0 #0.9 + moe_deepep_num_sms: 20 + moe_enable_deepep: false + moe_expert_capacity_factor: null + moe_extended_tp: false + moe_ffn_hidden_size: 2880 + moe_flex_dispatcher_backend: deepep + moe_grouped_gemm: true + moe_hybridep_num_sms: 16 + moe_input_jitter_eps: null + moe_latent_size: null + moe_layer_freq: 1 + moe_layer_recompute: false + moe_pad_expert_input_to_capacity: false + moe_per_layer_logging: false + moe_router_bias_update_rate: 0.001 + moe_router_dtype: fp32 # DeepEP only supports float32 probs + moe_router_enable_expert_bias: false + moe_router_force_load_balancing: false + moe_router_fusion: true + moe_router_group_topk: null + moe_router_load_balancing_type: none + moe_router_num_groups: null + moe_router_padding_for_fp8: false + moe_router_padding_for_quantization: false + moe_router_pre_softmax: false + moe_router_score_function: softmax + moe_router_topk: 4 + moe_router_topk_limited_devices: null + moe_router_topk_scaling_factor: null + moe_shared_expert_gate: false + moe_shared_expert_intermediate_size: null + moe_shared_expert_overlap: false + moe_token_dispatcher_type: alltoall + moe_token_drop_policy: probs + moe_token_dropping: false + moe_z_loss_coeff: null + + # ckpt + finetune: false + auto_continue_train: false + load: null + no_load_optim: null + no_load_rng: null + save: null + save_interval: 100000 + no_save_optim: null + no_save_rng: null + disable_last_saving: true + exit_on_missing_checkpoint: false + ckpt_format: torch + eval_iters: ${EVAL_ITERS:64} # eval_samples = eval_iters * GBS = 1024; set EVAL_ITERS in config shell (1024/GBS). + eval_interval: ${PRIMUS_EVAL_INTERVAL:768} + + # Turbo + enable_primus_turbo: true + use_turbo_attention: false + use_turbo_grouped_gemm: true + use_turbo_rms_norm: ${USE_TURBO_RMS_NORM: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 + + # 64 or 80 for ep8, 32 for ep16-64 is best practice + turbo_deepep_num_cu: 64 + turbo_deepep_use_comm_stream: false + + # sync-free moe support stage 0-3, 0 means not use sync-free moe + # stage 3 is completely no gpu-cpu sync in MoE, but cost more memory + # stage 2 is recommended for better performance + turbo_sync_free_moe_stage: 0 + + # Cross entropy flags + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true + + # tensorboard logging, set 'disable_tensorboard: false' to enable tensorboard logging + disable_tensorboard: true + tensorboard_dir: /workspace/code/tensorboard + tensorboard_log_interval: 1 + tensorboard_queue_size: 1000 + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_validation_ppl_to_tensorboard: true + log_memory_to_tensorboard: true + log_world_size_to_tensorboard: true + log_loss_scale_to_tensorboard: true diff --git a/primus/backends/megatron/__init__.py b/primus/backends/megatron/__init__.py index 939a1c4ba..b1576a233 100644 --- a/primus/backends/megatron/__init__.py +++ b/primus/backends/megatron/__init__.py @@ -7,11 +7,13 @@ from primus.backends.megatron.megatron_adapter import MegatronAdapter from primus.backends.megatron.megatron_pretrain_trainer import MegatronPretrainTrainer from primus.backends.megatron.megatron_sft_trainer import MegatronSFTTrainer +from primus.backends.megatron.mlperf import MLPerfMegatronPretrainTrainer from primus.core.backend.backend_registry import BackendRegistry BackendRegistry.register_adapter("megatron", MegatronAdapter) BackendRegistry.register_trainer_class(MegatronPretrainTrainer, "megatron") BackendRegistry.register_trainer_class(MegatronSFTTrainer, "megatron", "sft") +BackendRegistry.register_trainer_class(MLPerfMegatronPretrainTrainer, "megatron", "mlperf_pretrain") # Export trainers for convenience # Use lazy import for FluxPretrainTrainer to avoid Megatron dependency diff --git a/primus/backends/megatron/core/distributed/sdma_param_gather.py b/primus/backends/megatron/core/distributed/sdma_param_gather.py new file mode 100644 index 000000000..4f2f5ebb8 --- /dev/null +++ b/primus/backends/megatron/core/distributed/sdma_param_gather.py @@ -0,0 +1,248 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""SDMA (copy-engine) param all-gather helpers for the distributed optimizer. + +Migrated from the source patch ``megatron_sdma_allgather.patch`` (mpo branch). + +These helpers implement a ``torch.distributed.all_gather_into_tensor``-compatible +all-gather that routes the gather through Primus-Turbo symmetric memory + HIP +copy-engine (SDMA) memcpys instead of RCCL kernels. The intent is to free up +CU/compute resources by performing the param all-gather purely on the copy +engine, overlapping it with forward compute. + +The implementation is consumed by +:mod:`primus.backends.megatron.patches.parallelism.sdma_param_all_gather_patches`, +which swaps Megatron's ``_ParamAndGradBucketGroup.start_param_sync`` distributed +optimizer path to dispatch per-bucket all-gathers through +:func:`all_gather_into_tensor_sdma`. + +Activation: + ``ENABLE_SDMA_ALLGATHER=1`` (gated by the patch). When the required + Primus-Turbo / ``hip`` primitives are unavailable, every call falls back to + ``torch.distributed.all_gather_into_tensor`` so behaviour is preserved. +""" + +import importlib +import os +import warnings +from typing import Callable, Dict, Optional, Tuple + +import torch + +# Minimum symmetric-memory workspace size (bytes) requested per group. The +# workspace is reused across calls, so it is sized for the largest expected +# bucket. Overridable via env for A/B sizing. +_SDMA_SYMM_MEM_MIN_BYTES = int(os.getenv("MEGATRON_SDMA_SYMM_MEM_MIN_BYTES", str(749887296))) + + +class _WaitableHandle: + """Lightweight waitable handle that mimics ``torch.distributed.Work``.""" + + def __init__(self, wait_fn: Optional[Callable[[], None]] = None, work=None): + self._wait_fn = wait_fn + self._work = work + self._done = False + + def wait(self): + if self._done: + return True + if self._work is not None: + self._work.wait() + if self._wait_fn is not None: + self._wait_fn() + self._done = True + return True + + +def _all_gather_into_tensor_waitable_fallback( + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: Optional[torch.distributed.ProcessGroup] = None, + async_op: bool = False, +): + """Fallback to ``torch.distributed.all_gather_into_tensor`` with a waitable handle.""" + work = torch.distributed.all_gather_into_tensor( + output_tensor, input_tensor, group=group, async_op=async_op + ) + if async_op and work is not None: + return _WaitableHandle(work=work) + return _WaitableHandle() + + +def _get_sdma_peer_copy_stream_count(world_size: int) -> int: + """Resolve the number of peer-copy streams for SDMA param all-gather.""" + max_streams = max(world_size - 1, 1) + default_streams = min(max_streams, 8) + env_value = os.getenv("MEGATRON_SDMA_PEER_COPY_STREAMS") + if env_value is None: + return default_streams + + try: + configured_streams = int(env_value) + except ValueError: + warnings.warn( + f"Invalid MEGATRON_SDMA_PEER_COPY_STREAMS={env_value!r}; using default " f"{default_streams}." + ) + return default_streams + + if configured_streams < 1: + warnings.warn( + f"MEGATRON_SDMA_PEER_COPY_STREAMS must be >= 1, got {configured_streams}; " + f"using default {default_streams}." + ) + return default_streams + + return min(configured_streams, max_streams) + + +class _SDMAGroupRuntime: + """Shared SDMA runtime for a process group (comm + peer-copy streams).""" + + def __init__(self, world_size: int): + # Barrier/publish/local-copy run on one shared communication stream. + self.comm_stream = torch.cuda.Stream() + # Limit peer-copy fan-out by default to reduce SDMA/memory-system + # pressure and improve overlap with forward compute. + peer_copy_stream_count = _get_sdma_peer_copy_stream_count(world_size) + self.peer_copy_streams = [torch.cuda.Stream() for _ in range(peer_copy_stream_count)] + + +_sdma_group_runtime_cache: Dict[Tuple[str, int], _SDMAGroupRuntime] = {} + + +def _get_sdma_group_runtime(group_name: str, world_size: int) -> _SDMAGroupRuntime: + cache_key = (group_name, world_size) + runtime = _sdma_group_runtime_cache.get(cache_key) + if runtime is None: + runtime = _SDMAGroupRuntime(world_size=world_size) + _sdma_group_runtime_cache[cache_key] = runtime + return runtime + + +def all_gather_into_tensor_sdma( + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: Optional[torch.distributed.ProcessGroup] = None, + async_op: bool = False, +): + """SDMA all-gather with ``all_gather_into_tensor``-compatible signature. + + Follows the Primus-Turbo async-tp style (symmetric memory + DMA copies) and + always returns a ``.wait()``-able handle. Falls back to + ``torch.distributed.all_gather_into_tensor`` when SDMA primitives are + unavailable. + """ + if not torch.distributed.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before all_gather_into_tensor_sdma") + + if group is None: + group = torch.distributed.distributed_c10d._get_default_group() + + world_size = torch.distributed.get_world_size(group=group) + rank = torch.distributed.get_rank(group=group) + + if output_tensor.numel() != input_tensor.numel() * world_size: + raise ValueError( + "output_tensor.numel() must equal input_tensor.numel() * world_size " "for all_gather_into_tensor" + ) + if output_tensor.dtype != input_tensor.dtype: + raise ValueError("output_tensor.dtype must match input_tensor.dtype for SDMA all_gather_into_tensor") + + # SDMA path is CUDA-only and expects a contiguous flattened layout. + if ( + not output_tensor.is_cuda + or not input_tensor.is_cuda + or output_tensor.device != input_tensor.device + or output_tensor.device.index != torch.cuda.current_device() + or not output_tensor.is_contiguous() + ): + return _all_gather_into_tensor_waitable_fallback( + output_tensor, input_tensor, group=group, async_op=async_op + ) + 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 + except Exception: + return _all_gather_into_tensor_waitable_fallback( + output_tensor, input_tensor, group=group, async_op=async_op + ) + + group_name = getattr(group, "group_name", None) + if group_name is None: + return _all_gather_into_tensor_waitable_fallback( + output_tensor, input_tensor, group=group, async_op=async_op + ) + + 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)) + gather_buffers = [ + symm_mem.get_buffer(r, input_tensor.shape, input_tensor.dtype) for r in range(world_size) + ] + runtime = _get_sdma_group_runtime(group_name=group_name, world_size=world_size) + + current_stream = torch.cuda.current_stream() + # Per-call barrier event so each handle tracks its own all-gather point. + barrier_event = torch.cuda.Event() + + # Ensure input is ready on the communication stream. + runtime.comm_stream.wait_stream(current_stream) + + with torch.cuda.stream(runtime.comm_stream): + # Barrier (workspace safety), publish local shard, post-publish barrier. + symm_mem.barrier() + gather_buffers[rank].copy_(input_tensor, True) + symm_mem.barrier() + # 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, + ) + ) + + for peer_idx, src_rank in enumerate(r for r in range(world_size) if r != rank): + src_buf = gather_buffers[src_rank] + 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, + ) + ) + + def _wait_impl(): + wait_stream = torch.cuda.current_stream() + wait_stream.wait_stream(runtime.comm_stream) + for copy_stream in runtime.peer_copy_streams: + wait_stream.wait_stream(copy_stream) + + handle = _WaitableHandle(wait_fn=_wait_impl) + if not async_op: + handle.wait() + return handle diff --git a/primus/backends/megatron/core/extensions/fused_residual_rmsnorm.py b/primus/backends/megatron/core/extensions/fused_residual_rmsnorm.py new file mode 100644 index 000000000..d9778d352 --- /dev/null +++ b/primus/backends/megatron/core/extensions/fused_residual_rmsnorm.py @@ -0,0 +1,453 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Runtime monkeypatch — fuse residual+add into RMSNorm at two sites. + +This is the runtime install hook (callable from a trainer entry point such +as ``small_llm_moe_pretraining/primus/src/train.py``). The Triton RMSNorm +kernels live in Primus-Turbo (``primus_turbo.pytorch.ops.normalization``, +which exposes ``rmsnorm`` and the fused ``rmsnorm_residual`` variant). + +V1 — IN-LAYER fuse (``PRIMUS_FUSED_RESIDUAL_NORM=1``) + self_attn_bda(residual + attn_out) → pre_mlp_layernorm + Kills 1 of 2 ``vectorized_elementwise_kernel>`` + launches per layer (the in-layer ADD#1). + +V2 — CROSS-LAYER fuse (``PRIMUS_FUSED_RESIDUAL_NORM_V2=1``, implies V1) + layer N: skip mlp_bda(mlp_out + residual_post_attn), instead stash + ``(mlp_out, residual_post_attn)`` as a "carry" on layer N+1 + layer N+1: input_layernorm(carry) does the fused add+norm + last layer: carry goes to final_layernorm via ``_v2_pending_carry``. + Kills the remaining ADD#2 in the bf16 add tax (24 → 1 launches per step; + the last layer keeps its add when final_layernorm isn't a + ``PrimusTurboRMSNorm``, otherwise that one is fused too). + +Activation +---------- +* V1: ``PRIMUS_FUSED_RESIDUAL_NORM=1`` (default 0). +* V2: ``PRIMUS_FUSED_RESIDUAL_NORM_V2=1`` (default 0). Implies V1. +* Requires ``use_turbo_rms_norm=true`` so the norms are ``PrimusTurboRMSNorm`` + instances we can extend with a ``residual=`` arg. + +Falls back silently to the original ``TransformerLayer.forward`` whenever +any precondition fails (recompute paths, fp32 residual, inference fused TP, +non-zero hidden_dropout, cross-attention layers, non-PrimusTurboRMSNorm +pre-norm). + +Safety +------ +* Both gates default off; A/B is always reproducible. +* On any exception in the fused path the patch reverts permanently and the + rest of the run uses the original Megatron forward. +* V2 keeps the original mlp_bda when the next consumer cannot fuse (e.g. + final_layernorm is not a PrimusTurboRMSNorm, or the next layer doesn't + pass ``_can_fuse``), so partial coverage degrades gracefully. +""" +from __future__ import annotations + +import os +import sys +from typing import Any + + +def _enabled() -> bool: + """V1 gate (or implied by V2).""" + if _v2_enabled(): + return True + v = os.environ.get("PRIMUS_FUSED_RESIDUAL_NORM", "0").strip().lower() + return v in ("1", "true", "yes", "on") + + +def _v2_enabled() -> bool: + v = os.environ.get("PRIMUS_FUSED_RESIDUAL_NORM_V2", "0").strip().lower() + return v in ("1", "true", "yes", "on") + + +def _log(msg: str) -> None: + rank = os.environ.get("RANK", "0") + if rank == "0": + print(f"[fused_residual_rmsnorm] {msg}", file=sys.stderr, flush=True) + + +_INSTALLED = False + + +def install() -> bool: + """Install the monkeypatch. Returns True on success.""" + global _INSTALLED + if _INSTALLED: + return True + if not _enabled(): + _log("disabled (set PRIMUS_FUSED_RESIDUAL_NORM=1 or " "PRIMUS_FUSED_RESIDUAL_NORM_V2=1 to enable)") + return False + + try: + from primus_turbo.pytorch.ops.normalization import ( + rmsnorm_residual as triton_rmsnorm_residual, + ) + except ImportError as exc: + _log(f"could not import primus_turbo rmsnorm_residual: {exc}; abort install") + return False + + try: + from primus.backends.megatron.core.extensions.primus_turbo import ( + PrimusTurboRMSNorm, + ) + except ImportError as exc: + _log(f"PrimusTurboRMSNorm not importable: {exc}; abort install") + return False + + try: + from megatron.core.transformer.transformer_block import TransformerBlock + from megatron.core.transformer.transformer_layer import TransformerLayer + except ImportError as exc: + _log(f"Megatron transformer modules not importable: {exc}; abort install") + return False + + # ----- 1. Extend PrimusTurboRMSNorm to accept residual ------------------ + if not getattr(PrimusTurboRMSNorm, "_fused_residual_patched", False): + _orig_norm_forward = PrimusTurboRMSNorm.forward + + def _new_norm_forward(self, x, residual=None): + # V2 path: an upstream layer stashed a carry for *this* norm to + # consume on its next forward call. The carry takes priority + # over an explicit ``residual=`` arg (which is V1-style and is + # only used when the layer's own _do_fused_forward already + # picked the right pair). + pending = getattr(self, "_v2_pending_carry", None) + if pending is not None: + self._v2_pending_carry = None + x_pending, r_pending = pending + gamma = self.weight + if getattr(self, "zero_centered_gamma", False): + gamma = gamma + 1 + norm_out, _xpr = triton_rmsnorm_residual(x_pending, r_pending, gamma, self.eps) + return norm_out + if residual is None: + return _orig_norm_forward(self, x) + gamma = self.weight + if getattr(self, "zero_centered_gamma", False): + gamma = gamma + 1 + return triton_rmsnorm_residual(x, residual, gamma, self.eps) + + PrimusTurboRMSNorm.forward = _new_norm_forward + PrimusTurboRMSNorm._fused_residual_patched = True + _log("patched PrimusTurboRMSNorm.forward (residual= arg + V2 _pending_carry)") + + # ----- 2. Patch TransformerBlock.__init__ to wire V2 layer links -------- + if _v2_enabled() and not getattr(TransformerBlock, "_fused_residual_v2_init_patched", False): + _orig_block_init = TransformerBlock.__init__ + + def _v2_init(self, *args, **kwargs): + _orig_block_init(self, *args, **kwargs) + try: + _wire_v2_layer_links(self) + except Exception as exc: + _log(f"V2 layer-link wiring raised {type(exc).__name__}: {exc}; " "block will run V1-only") + + TransformerBlock.__init__ = _v2_init + TransformerBlock._fused_residual_v2_init_patched = True + _log("patched TransformerBlock.__init__ (V2 layer-link wiring active)") + + # ----- 3. Replace TransformerLayer.forward with fused variant ----------- + if not getattr(TransformerLayer, "_fused_residual_patched", False): + _orig_layer_forward = TransformerLayer.forward + + def _fused_layer_forward(self, *args, **kwargs): + if not _can_fuse(self): + # Per-layer fallback. If a carry was queued for *this* layer + # but the layer itself can't fuse, drop it back to a real + # add so the layer sees a correct hidden_states. + _drain_carry_into_hidden_states(self, args, kwargs) + return _orig_layer_forward(self, *args, **kwargs) + try: + return _do_fused_forward(self, *args, **kwargs) + except Exception as exc: + _log( + f"fused forward raised {type(exc).__name__}: {exc}; " + f"falling back to original forward for all layers" + ) + TransformerLayer.forward = _orig_layer_forward + return _orig_layer_forward(self, *args, **kwargs) + + TransformerLayer.forward = _fused_layer_forward + TransformerLayer._fused_residual_patched = True + if _v2_enabled(): + _log("patched TransformerLayer.forward (V1 ADD#1 + V2 cross-layer ADD#2 fusion active)") + else: + _log("patched TransformerLayer.forward (V1 in-layer ADD#1+norm fusion active)") + + _INSTALLED = True + return True + + +# --------------------------------------------------------------------------- +# V2 layer-link wiring +# --------------------------------------------------------------------------- +def _wire_v2_layer_links(block: Any) -> None: + """Annotate each layer in ``block`` with V2 navigation pointers. + + Sets per-layer attributes: + * ``_v2_block`` -> back-reference to the owning block + * ``_v2_next_layer`` -> the next TransformerLayer or None + * ``_v2_is_last_layer`` -> True for the final layer in the block + * ``_v2_final_layernorm`` -> reference to block.final_layernorm if any + (only set on the last layer) + """ + layers = getattr(block, "layers", None) + if layers is None: + return + n = len(layers) + if n == 0: + return + final_ln = getattr(block, "final_layernorm", None) + # CRITICAL: nn.Module.__setattr__ auto-registers any nn.Module-valued + # attribute as a child module. Setting layer._v2_next_layer = next_layer + # would create cycles (each layer becomes a child of the previous one) + # and blow up tree-walking ops like .cuda() / .children() with + # RecursionError. Bypass with object.__setattr__ so these are plain + # Python attributes invisible to nn.Module bookkeeping. + for i, layer in enumerate(layers): + nxt = layers[i + 1] if (i + 1) < n else None + is_last = (i + 1) == n + object.__setattr__(layer, "_v2_block", block) + object.__setattr__(layer, "_v2_next_layer", nxt) + object.__setattr__(layer, "_v2_is_last_layer", is_last) + object.__setattr__(layer, "_v2_final_layernorm", final_ln if is_last else None) + object.__setattr__(layer, "_v2_carry", None) + + +def _drain_carry_into_hidden_states(layer: Any, args: tuple, kwargs: dict) -> None: + """If a V2 carry was left for ``layer`` but we are about to take the + *original* (unpatched) forward, materialise the carry as a regular add + into ``hidden_states`` so the original forward still gets the right + activations. + """ + carry = getattr(layer, "_v2_carry", None) + if carry is None: + return + layer._v2_carry = None + mlp_out, residual = carry + new_hs = mlp_out + residual + if args: + # python tuples are immutable; the caller (TransformerBlock loop) + # passes hidden_states via kwargs. We can't mutate the caller's + # args tuple from here, so route via kwargs (the kwargs path is + # always taken for hidden_states in current Megatron). + kwargs["hidden_states"] = new_hs + return + kwargs["hidden_states"] = new_hs + + +# --------------------------------------------------------------------------- +# Per-layer "can we fuse?" guard +# --------------------------------------------------------------------------- +def _can_fuse(layer: Any) -> bool: + """Return True when ``layer`` matches the assumptions of the fused path.""" + cfg = getattr(layer, "config", None) + if cfg is None: + return False + if getattr(cfg, "fp32_residual_connection", False): + return False + if getattr(cfg, "inference_fuse_tp_communication", False): + return False + if getattr(layer, "recompute_input_layernorm", False): + return False + if getattr(layer, "recompute_pre_mlp_layernorm", False): + return False + if getattr(layer, "hidden_dropout", 0.0) != 0.0: + return False + if getattr(layer, "offload_attn_norm", False): + return False + if getattr(layer, "offload_mlp_norm", False): + return False + + pre_mlp = getattr(layer, "pre_mlp_layernorm", None) + if pre_mlp is None: + return False + try: + from primus.backends.megatron.core.extensions.primus_turbo import ( + PrimusTurboRMSNorm, + ) + except ImportError: + return False + if not isinstance(pre_mlp, PrimusTurboRMSNorm): + return False + + from megatron.core.transformer.identity_op import IdentityOp + + cross_attn = getattr(layer, "cross_attention", None) + if cross_attn is not None and not isinstance(cross_attn, IdentityOp): + return False + + return True + + +def _v2_next_can_consume_carry(layer: Any) -> bool: + """Return True when layer N can stash an unfused carry instead of + running mlp_bda. Requires that the next consumer (next layer's + input_layernorm OR this block's final_layernorm) is a + ``PrimusTurboRMSNorm`` so it can absorb the residual at no extra cost. + """ + if not _v2_enabled(): + return False + + try: + from primus.backends.megatron.core.extensions.primus_turbo import ( + PrimusTurboRMSNorm, + ) + except ImportError: + return False + + is_last = getattr(layer, "_v2_is_last_layer", False) + if not is_last: + nxt = getattr(layer, "_v2_next_layer", None) + if nxt is None: + return False + nxt_in_ln = getattr(nxt, "input_layernorm", None) + if not isinstance(nxt_in_ln, PrimusTurboRMSNorm): + return False + # If the next layer can't fuse for its own reasons, we shouldn't + # leave a carry it has to drain — drain it ourselves via the + # _drain_carry_into_hidden_states fallback in the wrapper. + return _can_fuse(nxt) + # Last layer: route through final_layernorm if it's a PrimusTurboRMSNorm. + final_ln = getattr(layer, "_v2_final_layernorm", None) + if final_ln is None: + return False + return isinstance(final_ln, PrimusTurboRMSNorm) + + +# --------------------------------------------------------------------------- +# The fused forward +# --------------------------------------------------------------------------- +def _do_fused_forward(layer: Any, hidden_states=None, *args, **kwargs): + """Mirror of ``TransformerLayer.forward`` with V1 + V2 fusion paths. + + V1: pre_mlp_layernorm receives ``(attn_out, residual)`` and emits both + the normed activation AND ``x_plus_r`` for the next bda. + V2: input_layernorm consumes a ``_v2_carry`` left by the previous + layer, and at exit either stashes a new carry (next layer / final + layernorm) or falls back to the explicit mlp_bda add. + """ + from megatron.core.utils import ( + deprecate_inference_params, + make_viewless_tensor, + nvtx_range_pop, + nvtx_range_push, + ) + + if hidden_states is None: + hidden_states = kwargs.pop("hidden_states") + else: + kwargs.pop("hidden_states", None) + + inference_context = deprecate_inference_params( + kwargs.get("inference_context"), kwargs.get("inference_params") + ) + + v2_active = _v2_enabled() + + # ---- input_layernorm (with optional V2 carry consume) ---------------- + nvtx_range_push(suffix="input_layernorm") + carry = getattr(layer, "_v2_carry", None) if v2_active else None + if carry is not None: + # Fused path: input_layernorm absorbs the previous layer's deferred + # mlp_bda add. The returned x_plus_r becomes our residual base. + layer._v2_carry = None + prev_mlp_out, prev_residual = carry + # We bypass PrimusTurboRMSNorm.forward to access x_plus_r directly, + # which is needed as the residual for ADD#1. + from primus_turbo.pytorch.ops.normalization import ( + rmsnorm_residual as triton_rmsnorm_residual, + ) + + in_ln = layer.input_layernorm + gamma = in_ln.weight + if getattr(in_ln, "zero_centered_gamma", False): + gamma = gamma + 1 + input_layernorm_output, hidden_states = triton_rmsnorm_residual( + prev_mlp_out, prev_residual, gamma, in_ln.eps + ) + else: + input_layernorm_output = layer.input_layernorm(hidden_states) + nvtx_range_pop(suffix="input_layernorm") + + residual = hidden_states # base for ADD#1 (post input_layernorm) + + # ---- self attention -------------------------------------------------- + nvtx_range_push(suffix="self_attention") + attention_output_with_bias = layer.self_attention( + input_layernorm_output, + attention_mask=kwargs.get("attention_mask"), + inference_context=inference_context, + rotary_pos_emb=kwargs.get("rotary_pos_emb"), + rotary_pos_cos=kwargs.get("rotary_pos_cos"), + rotary_pos_sin=kwargs.get("rotary_pos_sin"), + rotary_pos_cos_sin=kwargs.get("rotary_pos_cos_sin"), + attention_bias=kwargs.get("attention_bias"), + packed_seq_params=kwargs.get("packed_seq_params"), + sequence_len_offset=kwargs.get("sequence_len_offset"), + ) + nvtx_range_pop(suffix="self_attention") + + attn_out, attn_bias = attention_output_with_bias + if attn_bias is not None: + attn_out = attn_out + attn_bias.to(attn_out.dtype) + + # ---- V1 fuse: bda(attn_out, residual) + pre_mlp_layernorm ------------ + nvtx_range_push(suffix="fused_residual_pre_mlp_layernorm") + pre_mlp_layernorm_output, hidden_states = layer.pre_mlp_layernorm(attn_out, residual=residual) + nvtx_range_pop(suffix="fused_residual_pre_mlp_layernorm") + # hidden_states now == attn_out + residual (= residual_post_attn). + + # ---- mlp ------------------------------------------------------------- + padding_mask = kwargs.get("padding_mask", None) + try: + mlp_output_with_bias = layer.mlp(pre_mlp_layernorm_output, padding_mask=padding_mask) + except TypeError: + mlp_output_with_bias = layer.mlp(pre_mlp_layernorm_output) + + # ---- mlp_bda OR V2 carry stash --------------------------------------- + if v2_active and _v2_next_can_consume_carry(layer): + # Skip mlp_bda. Stash carry on the next consumer, return + # residual_post_attn as the layer's "output" tensor (used only by + # TransformerBlock for offload commit / make_viewless plumbing; + # the next layer / final_layernorm reads the carry instead). + mlp_out, mlp_bias = mlp_output_with_bias + if mlp_bias is not None: + mlp_out = mlp_out + mlp_bias.to(mlp_out.dtype) + + is_last = getattr(layer, "_v2_is_last_layer", False) + if is_last: + final_ln = layer._v2_final_layernorm + final_ln._v2_pending_carry = (mlp_out, hidden_states) + else: + layer._v2_next_layer._v2_carry = (mlp_out, hidden_states) + + nvtx_range_push(suffix="v2_skip_mlp_bda") + output = make_viewless_tensor( + inp=hidden_states, + requires_grad=hidden_states.requires_grad, + keep_graph=True, + ) + nvtx_range_pop(suffix="v2_skip_mlp_bda") + return output, None + + # V1-only path (or V2 last-layer-without-norm-final): explicit mlp_bda. + nvtx_range_push(suffix="mlp_bda") + with layer.bias_dropout_add_exec_handler(): + hidden_states = layer.mlp_bda(layer.training, layer.config.bias_dropout_fusion)( + mlp_output_with_bias, hidden_states, layer.hidden_dropout + ) + nvtx_range_pop(suffix="mlp_bda") + + output = make_viewless_tensor( + inp=hidden_states, + requires_grad=hidden_states.requires_grad, + keep_graph=True, + ) + return output, None diff --git a/primus/backends/megatron/mlperf/__init__.py b/primus/backends/megatron/mlperf/__init__.py new file mode 100644 index 000000000..1ddb2d92f --- /dev/null +++ b/primus/backends/megatron/mlperf/__init__.py @@ -0,0 +1,26 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Primus MLPerf logging integration for the Megatron pretrain backend. + +This package integrates the (formerly external) ``primus_mllog`` thin wrapper +directly into Primus, mapping it onto the new ``BaseTrainer`` lifecycle. + +``mlperf_logging`` and other MLPerf-only dependencies are imported lazily +inside methods so that importing this package (which happens whenever the +Megatron backend is loaded) never breaks non-MLPerf runs. +""" + +from primus.backends.megatron.mlperf.mlperf_logger import MLPerfLogger, ThroughputTimer +from primus.backends.megatron.mlperf.mlperf_pretrain_trainer import ( + MLPerfMegatronPretrainTrainer, +) + +__all__ = [ + "MLPerfLogger", + "ThroughputTimer", + "MLPerfMegatronPretrainTrainer", +] diff --git a/primus/backends/megatron/mlperf/mlperf_logger.py b/primus/backends/megatron/mlperf/mlperf_logger.py new file mode 100644 index 000000000..ed60e5de8 --- /dev/null +++ b/primus/backends/megatron/mlperf/mlperf_logger.py @@ -0,0 +1,159 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""MLPERF Logging support for Primus Megatron backend.""" + +import os +import time +from typing import Any, Dict, Optional + + +class MLPerfLogger: + """Wrapper around mllog library with rank-aware logging.""" + + def __init__(self): + from mlperf_logging import mllog + + self.mllogger = mllog.get_mllogger() + self._configured = False + self._rank = None + self._save_to_file = os.getenv("MLLOG_SAVE_TO_FILE", "1").lower() not in ("0", "false", "no") + + def configure(self, filepath: str, args) -> Dict[str, Any]: + from mlperf_logging import mllog + + if self._save_to_file: + mllog.config(filename=filepath, default_stack_offset=3) + else: + mllog.config(default_stack_offset=3) + self._configured = True + self._rank = self._get_rank() + return self.extract_mlperf_configs(args) + + def _get_rank(self) -> int: + import torch + + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank() + return int(os.environ.get("RANK", 0)) + + def log_event_all_ranks(self, key: str, value: Any, metadata: Optional[Dict] = None): + self.mllogger.event(key=key, value=value, metadata=metadata or {}) + + def log_event(self, key: str, value: Any, metadata: Optional[Dict] = None, stack_offset: int = 2): + if self._rank == 0: + self.mllogger.event(key=key, value=value, metadata=metadata or {}, stack_offset=stack_offset) + + def log_start(self, key: str, metadata: Optional[Dict] = None, stack_offset: int = 2): + if self._rank == 0: + self.mllogger.start(key=key, metadata=metadata or {}, stack_offset=stack_offset) + + def log_end(self, key: str, metadata: Optional[Dict] = None, stack_offset: int = 2): + if self._rank == 0: + self.mllogger.end(key=key, metadata=metadata or {}, stack_offset=stack_offset) + + def extract_mlperf_configs(self, args) -> Dict[str, Any]: + """Extract MLPERF config parameters from Megatron args.""" + from mlperf_logging.mllog import constants + + data_parallel_size = getattr(args, "data_parallel_size", 1) + if data_parallel_size == 0: + data_parallel_size = 1 + micro_batches = args.global_batch_size // (args.micro_batch_size * data_parallel_size) + + eval_samples = args.eval_iters * args.global_batch_size if args.eval_iters > 0 else 1024 + + train_samples = getattr(args, "train_samples", None) + if not train_samples: + train_samples = args.train_iters * args.global_batch_size + + optimizer_name = getattr(args, "optimizer", "adam") + if optimizer_name.lower() == "adam": + optimizer_name = "adamw" + else: + optimizer_name = optimizer_name.lower() + + configs = { + constants.GLOBAL_BATCH_SIZE: args.global_batch_size, + constants.GRADIENT_ACCUMULATION_STEPS: micro_batches, + "max_sequence_length": args.seq_length, + constants.TRAIN_SAMPLES: train_samples, + constants.EVAL_SAMPLES: eval_samples, + constants.SEED: args.seed, + "init_checkpoint_step": args.iteration, + constants.OPT_NAME: optimizer_name, + constants.OPT_BASE_LR: args.lr, + constants.OPT_ADAMW_BETA_1: getattr(args, "adam_beta1", 0.9), + constants.OPT_ADAMW_BETA_2: getattr(args, "adam_beta2", 0.999), + constants.OPT_ADAMW_EPSILON: getattr(args, "adam_eps", 1e-8), + constants.OPT_ADAMW_WEIGHT_DECAY: args.weight_decay, + "opt_gradient_clip_norm": getattr(args, "clip_grad", 1.0), + "opt_end_learning_rate": args.min_lr, + "opt_learning_rate_warmup_steps": getattr(args, "lr_warmup_iters", 0), + "opt_learning_rate_decay_steps": args.lr_decay_iters if args.lr_decay_iters else args.train_iters, + "opt_learning_rate_decay_schedule": self._get_lr_schedule_name(args.lr_decay_style), + "max_steps": args.train_iters, + constants.SUBMISSION_BENCHMARK: os.getenv("MLLOG_SUBMISSION_BENCHMARK", ""), + constants.SUBMISSION_DIVISION: os.getenv("MLLOG_SUBMISSION_DIVISION", ""), + constants.SUBMISSION_STATUS: os.getenv("MLLOG_SUBMISSION_STATUS", ""), + constants.SUBMISSION_ORG: os.getenv("MLLOG_SUBMISSION_ORG", ""), + constants.SUBMISSION_PLATFORM: os.getenv("MLLOG_SUBMISSION_PLATFORM", ""), + constants.TENSOR_PARALLELISM: int(os.getenv("MLLOG_TENSOR_PARALLELISM", 1)), + constants.PIPELINE_PARALLELISM: int(os.getenv("MLLOG_PIPELINE_PARALLELISM", 1)), + constants.CONTEXT_PARALLELISM: int(os.getenv("MLLOG_CONTEXT_PARALLELISM", 1)), + constants.EXPERT_PARALLELISM: int(os.getenv("MLLOG_EXPERT_PARALLELISM", 1)), + constants.MICRO_BATCH_SIZE: int(os.getenv("MLLOG_MICRO_BATCH_SIZE", 1)), + constants.CONFIG_FILENAME: os.getenv("MLLOG_CONFIG_FILENAME", ""), + "lowest_numerical_precision_linear": os.getenv("MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR", ""), + } + return configs + + def _get_lr_schedule_name(self, style: str) -> str: + mapping = { + "cosine": "cosine with linear warmup", + "linear": "linear", + "constant": "constant", + } + return mapping.get(style, style) + + +class ThroughputTimer: + """Per-step accumulation timer for training throughput (samples/sec). + + Measures only the time inside each training step, excluding eval time + and inter-step overhead (dataloader, callbacks, Python gaps). This + matches the approach used by nemo's Timer class. + """ + + def __init__(self, global_batch_size: int): + self.gbs = global_batch_size + self._step_start_time = None + self._accumulated_time = 0.0 + self._accumulated_samples = 0 + + def step_start(self): + """Mark the beginning of a training step.""" + self._step_start_time = time.time() + + def step_stop(self): + """Mark the end of a training step and accumulate elapsed time.""" + if self._step_start_time is not None: + self._accumulated_time += time.time() - self._step_start_time + self._accumulated_samples += self.gbs + self._step_start_time = None + + def get_throughput(self) -> Optional[float]: + """Return throughput for the accumulated steps and reset counters. + + Returns: + Samples per second, or None if no steps were accumulated. + """ + if self._accumulated_time <= 0: + return None + throughput = self._accumulated_samples / self._accumulated_time + self._accumulated_time = 0.0 + self._accumulated_samples = 0 + return throughput diff --git a/primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py b/primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py new file mode 100644 index 000000000..360013b2b --- /dev/null +++ b/primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py @@ -0,0 +1,417 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""MLPerf-enabled Megatron pretrain trainer for the Primus new architecture. + +Unlike the legacy ``primus_mllog.MLPerfMegatronPretrainTrainer`` (which +subclassed the *old* Primus trainer and overrode Primus-owned ``init`` / +``run`` / ``train`` / ``train_step`` methods), the new architecture delegates +the entire training loop to upstream ``megatron.training.pretrain()`` inside +:meth:`MegatronPretrainTrainer.train`. There are therefore no Primus-owned +loop methods to override. + +This trainer maps the MLPerf logging onto the new ``BaseTrainer`` lifecycle +(``setup -> init -> train -> cleanup``) by **monkey-patching the upstream +Megatron functions** that the loop calls: + + * ``megatron.training.training.train`` (loop entry/exit) + * ``megatron.training.training.train_step`` (per-step hooks) + * ``megatron.training.training.evaluate`` (capture val loss) + * ``megatron.training.training.evaluate_and_print_results`` (eval markers) + +All ``mlperf_logging`` imports are lazy so importing this module never breaks +non-MLPerf runs. +""" + +import os +import time + +from primus.backends.megatron.megatron_pretrain_trainer import MegatronPretrainTrainer +from primus.backends.megatron.mlperf.mlperf_logger import MLPerfLogger, ThroughputTimer + +try: + from rpdTracerControl import rpdTracerControl +except ImportError: + rpdTracerControl = None + + +def _get_arg(args, kwargs, index, name): + """Fetch a positional-or-keyword argument from a wrapped call. + + Upstream Megatron calls ``train`` / ``train_step`` positionally, but we + accept either form so the patch is robust across signature tweaks. + """ + if name in kwargs: + return kwargs[name] + if index < len(args): + return args[index] + return None + + +class MLPerfMegatronPretrainTrainer(MegatronPretrainTrainer): + """MegatronPretrainTrainer with MLPerf (mllog) logging.""" + + def __init__(self, backend_args): + super().__init__(backend_args) + + self.mllogger = MLPerfLogger() + self.throughput_timer = None + self.train_start_time = None + self.train_stop_time = None + self.last_validation_loss = None + self.train_loss_log_freq = int(os.getenv("MLLOG_TRAIN_LOSS_LOG_FREQ", "1")) + self.block_tput_log = os.getenv("MLLOG_BLOCK_TPUT_LOG", "0") == "1" + self.target_eval_loss = float(os.getenv("MLLOG_TARGET_EVAL_LOSS", "0.0")) + self.is_target_reached = False + self._warmup_active = False + + # RPD profiler (optional). Mirrors the legacy wrapper. + profiler = os.getenv("PROFILER", "") + self.profile_segment = os.getenv("PROFILE_SEGMENT", "train") + self.rpd = None + if profiler == "rpd" and rpdTracerControl is not None: + rpdTracerControl.setFilename(name="trace.rpd", append=True) + self.rpd = rpdTracerControl() + enable_python_trace = os.getenv("ENABLE_PYTHON_TRACE", "false").strip().lower() in { + "1", + "true", + "yes", + "on", + } + self.rpd.setPythonTrace(doTrace=enable_python_trace) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def setup(self): + """Megatron setup + MLPerf CACHE_CLEAR / INIT_START. + + ``setup()`` runs before ``init()`` and before ``train()`` (which + calls upstream ``pretrain()`` that performs the heavy init), so this + is the correct place to bracket the start of initialisation. + """ + super().setup() + + # Primus reconfigures stdlib ``logging`` via loguru during setup, + # which resets the per-logger levels installed by the log-suppression + # module at import time. Re-apply the quiet levels (no-op unless + # PRIMUS_LOG_SUPPRESSION=1 and not verbose). + try: + from primus.mlperf_log_suppression import reapply_quiet_logger_levels + + reapply_quiet_logger_levels() + except Exception: + pass + + from mlperf_logging.mllog import constants + + self.mllogger.log_event_all_ranks(key=constants.CACHE_CLEAR, value=True) + self.mllogger.log_event_all_ranks(key=constants.INIT_START, value=None) + + def train(self): + """Install MLPerf hooks on upstream Megatron, then run pretrain().""" + import megatron.training.training as mt + + orig_train = mt.train + orig_train_step = mt.train_step + orig_evaluate = mt.evaluate + orig_eval_and_print = mt.evaluate_and_print_results + + wrapped_train = self._make_wrapped_train(mt, orig_train) + wrapped_train_step = self._make_wrapped_train_step(orig_train_step) + wrapped_evaluate = self._make_wrapped_evaluate(orig_evaluate) + wrapped_eval_and_print = self._make_wrapped_eval_and_print(orig_eval_and_print) + + mt.train = wrapped_train + mt.train_step = wrapped_train_step + mt.evaluate = wrapped_evaluate + mt.evaluate_and_print_results = wrapped_eval_and_print + + try: + return super().train() + finally: + mt.train = orig_train + mt.train_step = orig_train_step + mt.evaluate = orig_evaluate + mt.evaluate_and_print_results = orig_eval_and_print + + # ------------------------------------------------------------------ + # Patch builders + # ------------------------------------------------------------------ + + def _make_wrapped_train(self, mt, orig_train): + def wrapped_train(*args, **kwargs): + from megatron.training import get_args + from mlperf_logging.mllog import constants + + from primus.backends.megatron.mlperf.warmup import run_synthetic_warmup + + megatron_args = get_args() + + # Configure mllog + log hyper-parameters + INIT_STOP. ``get_args`` + # is valid here because upstream pretrain() has finished init by + # the time it calls train(). + output_file = os.getenv("MLLOG_OUTPUT_FILE", "/results/mlperf_output.log") + configs = self.mllogger.configure(output_file, megatron_args) + for key, value in configs.items(): + self.mllogger.log_event(key=key, value=value) + self.mllogger.log_end(key=constants.INIT_STOP) + + self.throughput_timer = ThroughputTimer(megatron_args.global_batch_size) + + forward_step_func = _get_arg(args, kwargs, 0, "forward_step_func") + model = _get_arg(args, kwargs, 1, "model") + optimizer = _get_arg(args, kwargs, 2, "optimizer") + opt_param_scheduler = _get_arg(args, kwargs, 3, "opt_param_scheduler") + config = _get_arg(args, kwargs, 7, "config") + + # Synthetic warmup runs BEFORE RUN_START so its time is excluded + # from the timed run. It calls the (patched) module-level + # train_step; _warmup_active suppresses per-step loss logging. + self._warmup_active = True + try: + run_synthetic_warmup( + mt.train_step, + forward_step_func, + model, + optimizer, + opt_param_scheduler, + config, + megatron_args, + ) + finally: + self._warmup_active = False + # Discard warmup step timings. + self.throughput_timer.get_throughput() + + self.mllogger.log_start(key=constants.RUN_START) + self.train_start_time = time.time() + self.mllogger.log_start(key=constants.EPOCH_START, metadata={constants.SAMPLES_COUNT: 0}) + self.mllogger.log_start(key=constants.BLOCK_START, metadata={constants.SAMPLES_COUNT: 0}) + + if self.rpd and self.profile_segment == "train": + self.rpd.start() + + result = orig_train(*args, **kwargs) + + if self.rpd and self.profile_segment == "train": + self.rpd.stop() + self.rpd = None + + final_args = get_args() + consumed_samples = final_args.consumed_train_samples + self.mllogger.log_end( + key=constants.EPOCH_STOP, metadata={constants.SAMPLES_COUNT: consumed_samples} + ) + + if self.is_target_reached: + self.mllogger.log_end( + key=constants.RUN_STOP, + metadata={ + constants.SAMPLES_COUNT: consumed_samples, + constants.STATUS: constants.SUCCESS, + }, + ) + else: + self.train_stop_time = time.time() + self.mllogger.log_end( + key=constants.RUN_STOP, + metadata={ + constants.SAMPLES_COUNT: consumed_samples, + constants.STATUS: constants.ABORTED, + }, + ) + + # Overall run summary (rank-0 only; informational keys). + duration = (self.train_stop_time or time.time()) - (self.train_start_time or time.time()) + duration_minutes = duration / 60.0 + overall_throughput = consumed_samples / duration if duration > 0 else 0.0 + self.mllogger.log_event( + key="run_duration", + value=f"{round(duration, 2)}s -> {round(duration_minutes, 2)} minutes", + metadata={"samples": consumed_samples}, + ) + self.mllogger.log_event( + key="overall_throughput", + value=round(overall_throughput, 2), + metadata={"samples": consumed_samples}, + ) + + return result + + return wrapped_train + + def _make_wrapped_train_step(self, orig_train_step): + def wrapped_train_step(*args, **kwargs): + if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "train_step": + self.rpd.start() + + if self.throughput_timer is not None: + self.throughput_timer.step_start() + result = orig_train_step(*args, **kwargs) + if self.throughput_timer is not None: + self.throughput_timer.step_stop() + + # Upstream train_step returns an 8-tuple: + # (loss_dict, skipped_iter, should_checkpoint, should_exit, + # exit_code, grad_norm, num_zeros_in_grad, log_max_attention_logit) + loss_dict = result[0] + skipped_iter = result[1] + + if not skipped_iter and loss_dict and not self._warmup_active: + optimizer = _get_arg(args, kwargs, 3, "optimizer") + # Upstream passes the real loop iteration as the 8th positional + # (index 7) / ``iteration`` kwarg of train_step. ``args.iteration`` + # is NOT updated until after train_step returns, so use this. + iteration = _get_arg(args, kwargs, 7, "iteration") + self._on_train_step_end(loss_dict, optimizer, iteration) + + if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "train_step": + self.rpd.stop() + self.rpd = None + + return result + + return wrapped_train_step + + def _make_wrapped_evaluate(self, orig_evaluate): + def wrapped_evaluate(*args, **kwargs): + if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "eval": + self.rpd.start() + + result = orig_evaluate(*args, **kwargs) + + if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "eval": + self.rpd.stop() + self.rpd = None + + total_loss_dict = result[0] + if total_loss_dict and "lm loss" in total_loss_dict: + val_loss = total_loss_dict["lm loss"] + if hasattr(val_loss, "item"): + self.last_validation_loss = val_loss.item() + else: + self.last_validation_loss = float(val_loss) + + return result + + return wrapped_evaluate + + def _make_wrapped_eval_and_print(self, orig_eval_and_print): + def wrapped_eval_and_print(*args, **kwargs): + from megatron.training import get_args + from mlperf_logging.mllog import constants + + eval_args = get_args() + consumed_samples = eval_args.consumed_train_samples + + self._log_tracked_stats(eval_args.iteration, consumed_samples) + + self.mllogger.log_end( + key=constants.BLOCK_STOP, metadata={constants.SAMPLES_COUNT: consumed_samples} + ) + self.mllogger.log_start( + key=constants.EVAL_START, metadata={constants.SAMPLES_COUNT: consumed_samples} + ) + + if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "eval": + self.rpd.start() + + result = orig_eval_and_print(*args, **kwargs) + + if self.rpd and self.mllogger._get_rank() == 0 and self.profile_segment == "eval": + self.rpd.stop() + self.rpd = None + + validation_loss = self._get_validation_loss() + if validation_loss is not None: + self.mllogger.log_event( + key=constants.EVAL_ACCURACY, + value=validation_loss, + metadata={constants.SAMPLES_COUNT: consumed_samples}, + ) + + if self.target_eval_loss > 0.0 and validation_loss <= self.target_eval_loss: + if not self.is_target_reached: + self.is_target_reached = True + self.train_stop_time = time.time() + if self.mllogger._get_rank() == 0: + print( + f"[MLPERF] Target eval loss {self.target_eval_loss} reached " + f"with validation loss {validation_loss}" + ) + eval_args.train_iters = eval_args.iteration + eval_args.do_valid = False + eval_args.do_test = False + + self.mllogger.log_end( + key=constants.EVAL_STOP, metadata={constants.SAMPLES_COUNT: consumed_samples} + ) + + if not self.is_target_reached: + self.mllogger.log_start( + key=constants.BLOCK_START, metadata={constants.SAMPLES_COUNT: consumed_samples} + ) + + return result + + return wrapped_eval_and_print + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _on_train_step_end(self, loss_dict, optimizer, iteration=None): + from megatron.training import get_args + from mlperf_logging.mllog import constants + + args = get_args() + # Prefer the iteration passed into train_step (real loop counter); fall + # back to args.iteration only if it was not provided. args.iteration is + # stale during train_step, which made ``% freq`` true every step. + if iteration is None: + iteration = args.iteration + consumed_samples = args.consumed_train_samples + + if self.train_loss_log_freq <= 0 or iteration % self.train_loss_log_freq != 0: + return + + loss_value = loss_dict.get("lm loss") + if loss_value is None: + return + + if isinstance(loss_value, (tuple, list)): + loss_value = loss_value[0] + if hasattr(loss_value, "item"): + loss_value = loss_value.item() + + learning_rate = None + if optimizer is not None: + for param_group in optimizer.param_groups: + if not param_group.get("is_decoupled_lr", False): + learning_rate = param_group["lr"] + break + + self.mllogger.log_event( + key="train_loss", + value=loss_value, + metadata={constants.SAMPLES_COUNT: consumed_samples, "lr": learning_rate}, + ) + + def _log_tracked_stats(self, iteration, consumed_samples): + if self.throughput_timer is None: + return + throughput = self.throughput_timer.get_throughput() + if throughput is not None and self.block_tput_log: + self.mllogger.log_event( + key="tracked_stats", + value={"throughput": throughput}, + metadata={"step": consumed_samples}, + ) + + def _get_validation_loss(self): + return getattr(self, "last_validation_loss", None) diff --git a/primus/backends/megatron/mlperf/warmup.py b/primus/backends/megatron/mlperf/warmup.py new file mode 100644 index 000000000..f136ceee0 --- /dev/null +++ b/primus/backends/megatron/mlperf/warmup.py @@ -0,0 +1,776 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Synthetic-data warmup for Primus / Megatron training. + +Adds optional FP4 (MXFP4 / NVFP4) warmup support on top of the original FP8 +warmup helpers. The original FP8 / BF16 path is preserved exactly: + + * No autocast wrapping by default — ``train_step`` runs at the model's + native precision, identical to the committed pre-FP4 behaviour. + * After warmup, ``reset_fp8_state`` + ``seed_fp8_amax(1.0)`` always runs + (no-op for BF16 / non-FP8 modules). + * ``ChainedOptimizer`` support (EP < num_gpus) is preserved. + +FP4 mode is **opt-in** via ``WARMUP_RECIPE=fp4_mxfp4`` / ``fp4_nvfp4``. When +set, the warmup loop is wrapped in ``te.fp8_autocast(fp8_recipe=...)`` and +an additional ``reset_fp4_state()`` pass runs after warmup. + +Runs N forward+backward passes with random token sequences before the real +training loop starts. This pre-compiles Triton / CK / hipBLASLt kernels and +amortizes: + + * distributed-optimizer FP32 main-param allocation + * DDP gradient-bucket allocation + * NCCL communicator init + collective autotune + * hipBLASLt heuristic / Triton / CK JIT caches (recipe-dependent) + * PyTorch HIP allocator block layout + +After warmup: + * Model parameters are restored from a pre-warmup snapshot. + * Optimizer state is restored (neutered during warmup so weights never move). + * Adam buffers (exp_avg / exp_avg_sq) are zeroed and grads cleared. + * LR scheduler state is rolled back and re-synced (param_groups['lr']). + * FP8 scaling state (amax_history, scale, scale_inv, fp8_initialized) is + fully reset and seeded with safe defaults. + * If WARMUP_RECIPE selected an FP4 recipe, FP4 state is also reset. + +Controlled by: + + SYNTH_WARMUP_STEPS default 3, 0 disables. + WARMUP_RECIPE default "" (no autocast, original behaviour), + one of: + "" | "bf16" + | "fp8_hybrid" | "fp8_e4m3" + | "fp4_mxfp4" | "fp4_nvfp4" + Legacy alias WARMUP_FP8_RECIPE=hybrid|e4m3 also + accepted (maps to fp8_hybrid / fp8_e4m3). + FP8/BF16 models do NOT need to set this; it + only affects which autocast (if any) wraps the + warmup train_step. + WARMUP_FP8_HISTORY_LEN default 4, only used for fp8_* recipes. + SYNTH_WARMUP_EMPTY_CACHE default 1, call torch.cuda.empty_cache() per step. +""" + +import os +import time +from contextlib import nullcontext + +import torch +import torch.distributed + +# --------------------------------------------------------------------------- +# Logging +# --------------------------------------------------------------------------- + + +def _log(msg): + rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + if rank == 0: + print(f"[SYNTH_WARMUP] {msg}", flush=True) + + +# --------------------------------------------------------------------------- +# Synthetic data +# --------------------------------------------------------------------------- + + +class SyntheticGPTDataIterator: + """Infinite iterator yielding random token batches for Primus/Megatron GPT. + + Primus's ``get_batch_on_this_tp_rank`` expects the iterator to yield a dict + with ``tokens``, ``labels``, ``loss_mask``, and ``position_ids`` tensors, + each of shape ``[mbs, seq_length]``. + """ + + def __init__(self, seq_length, micro_batch_size, vocab_size=32000): + self.seq_length = seq_length + self.micro_batch_size = micro_batch_size + self.vocab_size = vocab_size + + def __iter__(self): + return self + + def __next__(self): + mbs = self.micro_batch_size + sl = self.seq_length + tokens = torch.randint(0, self.vocab_size, (mbs, sl), dtype=torch.int64) + labels = torch.randint(0, self.vocab_size, (mbs, sl), dtype=torch.int64) + loss_mask = torch.ones(mbs, sl, dtype=torch.float32) + position_ids = torch.arange(sl, dtype=torch.int64).unsqueeze(0).expand(mbs, -1) + return { + "tokens": tokens, + "labels": labels, + "loss_mask": loss_mask, + "position_ids": position_ids, + } + + +# --------------------------------------------------------------------------- +# FP8 state management -- ORIGINAL from committed warmup.py. +# +# The only change vs. the committed version is one defensive line at the top +# of ``_is_delayed_scaling_recipe`` that returns False for FP4-flagged +# modules so the FP8 reset path skips them (handled by reset_fp4_state() +# instead). No behavioural change for FP8 / BF16 / current-scaling models. +# --------------------------------------------------------------------------- + + +def _is_delayed_scaling_recipe(module): + """Return True if the module uses delayed scaling (which has scale/amax_history). + + Current scaling (tensorwise) is stateless — no scale or amax_history to + manage — so reset / seed operations must be skipped for it. + FP4 (MXFP4 / NVFP4) modules also have no scale/amax_history; they are + skipped here and handled by reset_fp4_state() if WARMUP_RECIPE is fp4_*. + Returns True when recipe type cannot be determined (safe default). + """ + # Defensive: FP4 modules use a different recipe state class with no + # scale/amax_history. Calling FP8 helpers on them is at best a no-op + # and at worst raises AttributeError. + if bool(getattr(module, "fp4", False)) or hasattr(module, "fp4_initialized"): + return False + + from transformer_engine.common.recipe import DelayedScaling + + fp8_meta = getattr(module, "fp8_meta", None) + if fp8_meta is None: + return True + fwd_state = fp8_meta.get("scaling_fwd", None) + if fwd_state is None: + return True + return isinstance(getattr(fwd_state, "recipe", None), DelayedScaling) + + +def _manual_reset_fp8_meta(module): + """Reset FP8 amax / scale tensors when reset_fp8_meta_tensors() is absent.""" + if not hasattr(module, "fp8_meta"): + return False + if not _is_delayed_scaling_recipe(module): + return False + meta = module.fp8_meta + reset_count = 0 + for key in ("scaling_fwd", "scaling_bwd"): + if key not in meta: + continue + tensor_meta = meta[key] + if hasattr(tensor_meta, "amax_history"): + tensor_meta.amax_history.fill_(0.0) + reset_count += 1 + if hasattr(tensor_meta, "scale"): + tensor_meta.scale.fill_(1.0) + reset_count += 1 + if hasattr(tensor_meta, "scale_inv"): + tensor_meta.scale_inv.fill_(1.0) + reset_count += 1 + return reset_count > 0 + + +def reset_fp8_state(model, reset_meta_tensors=True): + """Clear ``fp8_initialized`` on every TE layer, forcing re-init. + + When *reset_meta_tensors* is True, also zeros out amax_history and + resets scale / scale_inv to 1.0. + + Skips FP4 modules (handled by reset_fp4_state() when WARMUP_RECIPE=fp4_*). + """ + count = 0 + method_count = 0 + manual_count = 0 + + def _reset(m): + nonlocal count, method_count, manual_count + if hasattr(m, "fp8_initialized"): + m.fp8_initialized = False + count += 1 + if reset_meta_tensors: + if hasattr(m, "reset_fp8_meta_tensors"): + if _is_delayed_scaling_recipe(m): + m.reset_fp8_meta_tensors() + method_count += 1 + elif _manual_reset_fp8_meta(m): + manual_count += 1 + + models = model if isinstance(model, (list, tuple)) else [model] + for m in models: + m.apply(_reset) + _log(f"reset_fp8_state: {count} modules, " f"{method_count} via method, {manual_count} via manual reset") + return count + + +def seed_fp8_amax(model, seed_value=1.0): + """Fill amax_history with *seed_value* to prevent scale=inf after reset. + + Delayed scaling computes ``scale = fp8_max / max(amax_history)``. + If amax_history is all-zero after reset, scale becomes inf -> NaN. + Seeding with 1.0 gives scale = fp8_max / 1.0 ~ 448 (E4M3), which is safe. + """ + count = 0 + + def _seed(m): + nonlocal count + if not hasattr(m, "fp8_meta"): + return + if not _is_delayed_scaling_recipe(m): + return + meta = m.fp8_meta + for key in ("scaling_fwd", "scaling_bwd"): + if key not in meta: + continue + tensor_meta = meta[key] + if hasattr(tensor_meta, "amax_history"): + tensor_meta.amax_history.fill_(seed_value) + count += 1 + + models = model if isinstance(model, (list, tuple)) else [model] + for m in models: + m.apply(_seed) + _log(f"seed_fp8_amax: seeded {count} amax_history tensors with {seed_value}") + return count + + +# =========================================================================== +# FP4 state management -- NEW, only invoked when WARMUP_RECIPE=fp4_* +# =========================================================================== +# +# Notes on FP4 vs FP8 in TE: +# * MXFP4BlockScaling and NVFP4BlockScaling derive scales **per tile** from +# the current activation/weight on every forward — there is no +# amax_history and no global scale-inv to bias. So no seeding step. +# * Some TE versions store FP4 quantizer caches under fp8_meta as well +# (under attrs like ``block_scales``, ``mxfp4_quantizer``, etc.). We +# try to clear those defensively; they're allowed to be missing. +# * Different TE versions may track init via ``fp4_initialized``, +# ``fp8_initialized`` (shared with FP8), or both. We touch whichever +# attribute exists. + +# Per-tensor-meta attribute names that may hold FP4-specific cached state. +# Cleared opportunistically if present; missing attrs are silently skipped. +_FP4_META_ATTRS = ( + "block_scales", + "block_scales_inv", + "mxfp4_block_scale", + "fp4_quantizer_state", + "fp4_amax", +) + + +def _is_fp4_module(m): + """True if a TE module is configured for FP4 (MXFP4 / NVFP4). + + FP4 modules in TE share ``fp8_initialized`` and ``fp8_meta`` with FP8 but + populate them with an FP4 recipe state (``MXFP4BlockScalingRecipeState``, + ``NVFP4BlockScalingRecipeState``) that has no ``scale``/``amax_history``. + """ + if bool(getattr(m, "fp4", False)): + return True + if hasattr(m, "fp4_initialized"): + return True + meta = getattr(m, "fp8_meta", None) + if isinstance(meta, dict): + for key in ("scaling_fwd", "scaling_bwd"): + tm = meta.get(key) + if tm is None: + continue + cls = type(tm).__name__ + if "FP4" in cls or "Fp4" in cls: + return True + return False + + +def _manual_reset_fp4_meta(module): + """Reset FP4 per-tile block-scale buffers when reset_fp4_meta_tensors() is absent.""" + if not hasattr(module, "fp8_meta"): + return False + meta = module.fp8_meta + reset_count = 0 + for key in ("scaling_fwd", "scaling_bwd"): + if key not in meta: + continue + tensor_meta = meta[key] + for attr in _FP4_META_ATTRS: + if not hasattr(tensor_meta, attr): + continue + t = getattr(tensor_meta, attr) + if hasattr(t, "zero_"): + try: + t.zero_() + reset_count += 1 + except Exception: + pass + return reset_count > 0 + + +def reset_fp4_state(model, reset_meta_tensors=True): + """Clear ``fp4_initialized`` (or shared ``fp8_initialized``) on FP4 layers. + + Parallel to :func:`reset_fp8_state` but for FP4 modules. Unlike FP8 this + does NOT seed any history — block-scaling recipes have no amax_history, + so TE will recompute per-tile scales from real data on the next forward. + + Returns the number of modules touched. + """ + count = 0 + method_count = 0 + manual_count = 0 + + def _reset(m): + nonlocal count, method_count, manual_count + if not _is_fp4_module(m): + return + + touched = False + if hasattr(m, "fp4_initialized"): + m.fp4_initialized = False + touched = True + if hasattr(m, "fp8_initialized"): + m.fp8_initialized = False + touched = True + if touched: + count += 1 + + if reset_meta_tensors: + if hasattr(m, "reset_fp4_meta_tensors"): + try: + m.reset_fp4_meta_tensors() + method_count += 1 + except Exception: + pass + elif hasattr(m, "reset_fp8_meta_tensors"): + # FP4 modules typically share the meta-reset method with FP8. + try: + m.reset_fp8_meta_tensors() + method_count += 1 + except Exception: + if _manual_reset_fp4_meta(m): + manual_count += 1 + elif _manual_reset_fp4_meta(m): + manual_count += 1 + + models = model if isinstance(model, (list, tuple)) else [model] + for m in models: + m.apply(_reset) + _log(f"reset_fp4_state: {count} modules, " f"{method_count} via method, {manual_count} via manual reset") + return count + + +# =========================================================================== +# Warmup recipe selection -- NEW, opt-in via WARMUP_RECIPE +# =========================================================================== + + +def _resolve_recipe_str(): + """Return the warmup recipe string, honouring legacy WARMUP_FP8_RECIPE. + + Default is "" (empty) → no autocast wrapping, identical to the original + pre-FP4 warmup. FP8/BF16 users do not need to set anything. + """ + val = os.getenv("WARMUP_RECIPE", "").strip().lower() + if val: + return val + legacy = os.getenv("WARMUP_FP8_RECIPE", "").strip().lower() + if legacy in ("hybrid", "e4m3"): + return f"fp8_{legacy}" + return "" + + +def _build_warmup_recipe(): + """Return a TE recipe object if WARMUP_RECIPE selects one, else None. + + Supports FP8 (DelayedScaling) and FP4 (MXFP4BlockScaling, NVFP4BlockScaling). + Both kinds are consumed by ``te.fp8_autocast(fp8_recipe=...)``. + + Returns + ------- + (recipe, kind) where kind is one of ``"fp8"``, ``"fp4"`` or ``None``. + """ + recipe_str = _resolve_recipe_str() + if not recipe_str or recipe_str == "bf16": + if recipe_str == "bf16": + _log("WARMUP_RECIPE=bf16: no autocast wrapping") + return None, None + + # ---- FP8 family ------------------------------------------------------ + if recipe_str.startswith("fp8_"): + try: + from transformer_engine.common.recipe import DelayedScaling, Format + except ImportError: + _log(f"WARMUP_RECIPE={recipe_str!r} but transformer_engine missing; no autocast") + return None, None + fmt = {"fp8_hybrid": Format.HYBRID, "fp8_e4m3": Format.E4M3}.get(recipe_str) + if fmt is None: + _log(f"Unknown FP8 WARMUP_RECIPE={recipe_str!r}; no autocast") + return None, None + history_len = int(os.getenv("WARMUP_FP8_HISTORY_LEN", "4")) + _log( + f"Warmup will wrap train_step in fp8_autocast" + f"(format={recipe_str[4:]}, amax_history_len={history_len}, algo=most_recent)" + ) + return ( + DelayedScaling( + margin=0, + fp8_format=fmt, + amax_history_len=history_len, + amax_compute_algo="most_recent", + ), + "fp8", + ) + + # ---- FP4 family ------------------------------------------------------ + if recipe_str.startswith("fp4_"): + try: + import transformer_engine.common.recipe as te_recipe + except ImportError: + _log(f"WARMUP_RECIPE={recipe_str!r} but transformer_engine missing; no autocast") + return None, None + + # Try to nudge Primus's MXFP4 recipe-state patch into place if the + # model was built without FP4 (so the patch never ran). No-op otherwise. + try: + from primus.backends.megatron.core.fp4_utils import ( + _ensure_mxfp4_recipe_support, + ) + + _ensure_mxfp4_recipe_support() + except Exception: + pass + + cls_name = { + "fp4_mxfp4": "MXFP4BlockScaling", + "fp4_nvfp4": "NVFP4BlockScaling", + }.get(recipe_str) + if cls_name is None: + _log(f"Unknown FP4 WARMUP_RECIPE={recipe_str!r}; no autocast") + return None, None + + recipe_cls = getattr(te_recipe, cls_name, None) + if recipe_cls is None: + _log(f"WARMUP_RECIPE={recipe_str!r} but {cls_name} not in this TE build; no autocast") + return None, None + try: + recipe = recipe_cls() + except Exception as e: + _log(f"Failed to construct {cls_name}() for warmup: {e}; no autocast") + return None, None + _log(f"Warmup will wrap train_step in fp8_autocast(fp8_recipe={cls_name}())") + return recipe, "fp4" + + _log(f"Unknown WARMUP_RECIPE={recipe_str!r}; no autocast") + return None, None + + +def _warmup_autocast(recipe): + """Return a context-manager factory for the warmup loop.""" + if recipe is None: + return nullcontext + import transformer_engine.pytorch as te + + def _ctx(): + return te.fp8_autocast(enabled=True, fp8_recipe=recipe) + + return _ctx + + +# --------------------------------------------------------------------------- +# Optimizer save / restore -- ORIGINAL from committed warmup.py +# (preserves ChainedOptimizer support for EP < num_gpus configs) +# --------------------------------------------------------------------------- + + +def _get_inner_optimizers(optimizer): + """Return a list of leaf-level torch optimizers from any Megatron wrapper. + + When expert_parallel < num_gpus, Megatron creates a ChainedOptimizer with + separate sub-optimizers for dense and expert parameters. The previous code + used ``optimizer.optimizer`` which asserts ``len == 1`` on ChainedOptimizer. + This helper iterates over all chained sub-optimizers when present, and falls + back to the single-optimizer path otherwise. + """ + if hasattr(optimizer, "chained_optimizers"): + inners = [] + for sub_opt in optimizer.chained_optimizers: + inners.append(getattr(sub_opt, "optimizer", sub_opt)) + return inners + return [getattr(optimizer, "optimizer", optimizer)] + + +def _neuter_optimizer(optimizer): + """Set betas=[1,1], weight_decay=0, bias_correction=False. + + With betas=[1,1] the Adam momentum/variance stays at its initial (zero) + value, so the effective weight update is zero. + + Handles ChainedOptimizer by iterating over all sub-optimizers. + """ + all_saved = [] + for inner in _get_inner_optimizers(optimizer): + saved = [] + for group in inner.param_groups: + state = {} + for key in ("betas", "weight_decay", "bias_correction", "pre_mult_wd"): + if key in group: + state[key] = group[key] + saved.append(state) + + if "betas" in group: + group["betas"] = [1.0, 1.0] + if "weight_decay" in group: + group["weight_decay"] = 0.0 + if "bias_correction" in group: + group["bias_correction"] = False + if "pre_mult_wd" in group: + group["pre_mult_wd"] = 0.0 + all_saved.append(saved) + return all_saved + + +def _restore_optimizer(optimizer, all_saved): + """Restore optimizer state, handling ChainedOptimizer.""" + for inner, saved in zip(_get_inner_optimizers(optimizer), all_saved): + for group, state in zip(inner.param_groups, saved): + for key, val in state.items(): + group[key] = val + if "step" in group: + del group["step"] + + +# --------------------------------------------------------------------------- +# Model parameter snapshot -- ORIGINAL from committed warmup.py, unchanged +# --------------------------------------------------------------------------- + + +def _save_model_params(models): + """Snapshot all model parameters to CPU to avoid doubling GPU memory.""" + saved = {} + for m in models: + for name, p in m.named_parameters(): + saved[(id(m), name)] = p.data.to("cpu", copy=True) + return saved + + +def _restore_model_params(models, saved): + restored = 0 + for m in models: + for name, p in m.named_parameters(): + key = (id(m), name) + if key in saved: + p.data.copy_(saved[key].to(p.device)) + restored += 1 + return restored + + +# --------------------------------------------------------------------------- +# LR scheduler save / restore -- ORIGINAL + step(0) re-sync after restore +# --------------------------------------------------------------------------- + +_SCHEDULER_KEYS = ( + "num_steps", + "num_floating_point_operations_so_far", +) + + +def _save_scheduler_state(scheduler): + if scheduler is None: + return None + return {k: getattr(scheduler, k) for k in _SCHEDULER_KEYS if hasattr(scheduler, k)} + + +def _restore_scheduler_state(scheduler, state): + if scheduler is None or state is None: + return + for k, v in state.items(): + setattr(scheduler, k, v) + # Re-sync param_groups['lr'] and ['weight_decay'] to match the restored + # num_steps. Without this, the optimizer still holds the LR from the last + # warmup step even though num_steps has been rewound to its pre-warmup value. + scheduler.step(0) + + +# --------------------------------------------------------------------------- +# Main entry point +# --------------------------------------------------------------------------- + + +def run_synthetic_warmup( + train_step_func, + forward_step_func, + model, + optimizer, + opt_param_scheduler, + config, + megatron_args, +): + """Run *warmup_steps* forward+backward passes with synthetic data. + + Called from ``MLPerfMegatronPretrainTrainer``'s patched ``train()`` + **before** ``RUN_START`` is logged, so warmup time is excluded from the + timed run. All side-effects (model params, optimizer, LR scheduler, + FP8/FP4 state) are reverted after warmup. + + Default behaviour (``WARMUP_RECIPE`` unset): no autocast wrapping, + identical to the pre-FP4 committed version. Set + ``WARMUP_RECIPE=fp4_mxfp4`` or ``fp4_nvfp4`` to opt into FP4 warmup. + + Args: + train_step_func: The (already MLPerf-patched) upstream + ``megatron.training.training.train_step`` function. It is called + with the new-style signature + ``(forward_step_func, data_iterator, model, optimizer, + opt_param_scheduler, config, forward_backward_func, iteration=0)``. + ``forward_backward_func`` is resolved here via + ``get_forward_backward_func()`` (mirrors ``training.train()``). + """ + warmup_steps = int(os.getenv("SYNTH_WARMUP_STEPS", "3")) + if warmup_steps <= 0: + _log(f"Skipped (SYNTH_WARMUP_STEPS={warmup_steps})") + return + + empty_cache_each_step = os.getenv("SYNTH_WARMUP_EMPTY_CACHE", "1") not in ("0", "false", "False") + + t0 = time.time() + _log( + f"Starting {warmup_steps}-step synthetic warmup " + f"(seq_len={megatron_args.seq_length}, " + f"mbs={megatron_args.micro_batch_size}, " + f"empty_cache_each_step={empty_cache_each_step})" + ) + + vocab_size = getattr( + megatron_args, + "padded_vocab_size", + getattr(megatron_args, "vocab_size", 32000), + ) + synth_iter = SyntheticGPTDataIterator( + megatron_args.seq_length, + megatron_args.micro_batch_size, + vocab_size, + ) + + models = model if isinstance(model, (list, tuple)) else [model] + + # Resolve the forward-backward func the same way upstream + # ``megatron.training.training.train()`` does (see training.py). Newer + # Megatron ``train_step`` takes this as an explicit positional argument. + from megatron.core.pipeline_parallel import get_forward_backward_func + + forward_backward_func = get_forward_backward_func() + + # Temporarily set config fields needed by forward_backward_func. + # Megatron's train() normally sets these, but warmup runs before it. + from megatron.core.distributed import finalize_model_grads + from megatron.core.distributed.distributed_data_parallel import ( + DistributedDataParallel as DDP, + ) + + saved_config = {} + for key in ("finalize_model_grads_func", "grad_scale_func", "no_sync_func"): + saved_config[key] = getattr(config, key, None) + if config.finalize_model_grads_func is None: + config.finalize_model_grads_func = finalize_model_grads + if config.grad_scale_func is None: + config.grad_scale_func = optimizer.scale_loss + if megatron_args.overlap_grad_reduce and config.no_sync_func is None: + if isinstance(models[0], DDP): + config.no_sync_func = models[0].no_sync if len(models) == 1 else [m.no_sync for m in models] + + # ---- save state ------------------------------------------------------ + _log(f"Saving {sum(p.numel() for m in models for p in m.parameters())} parameters") + saved_params = _save_model_params(models) + saved_opt = _neuter_optimizer(optimizer) + saved_sched = _save_scheduler_state(opt_param_scheduler) + + # ---- decide warmup precision (default: no autocast = original path) - + recipe, recipe_kind = _build_warmup_recipe() + warmup_ctx_factory = _warmup_autocast(recipe) + + # ---- warmup steps ---------------------------------------------------- + for step in range(1, warmup_steps + 1): + step_t0 = time.time() + with warmup_ctx_factory(): + train_step_func( + forward_step_func, + synth_iter, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=0, + ) + torch.cuda.synchronize() + if empty_cache_each_step: + torch.cuda.empty_cache() + _log(f"Step {step}/{warmup_steps} done in {time.time() - step_t0:.1f}s") + + # ---- restore state --------------------------------------------------- + _log("Restoring optimizer") + _restore_optimizer(optimizer, saved_opt) + + _log("Restoring LR scheduler") + _restore_scheduler_state(opt_param_scheduler, saved_sched) + + _log("Restoring model parameters") + n_restored = _restore_model_params(models, saved_params) + del saved_params + _log(f"Restored {n_restored} parameter tensors") + + if hasattr(optimizer, "reload_model_params"): + optimizer.reload_model_params() + _log("Called optimizer.reload_model_params()") + + # ---- zero Adam state buffers (defensive) ---------------------------- + # With betas=[1,1] during warmup, exp_avg / exp_avg_sq should already be + # zero, but zero them explicitly in case any param group lacked betas. + # Iterate over chained sub-optimizers for ChainedOptimizer support. + zeroed_state_tensors = 0 + for inner_opt in _get_inner_optimizers(optimizer): + for param_states in inner_opt.state.values(): + for k, v in param_states.items(): + if isinstance(v, torch.Tensor) and v.is_floating_point(): + v.zero_() + zeroed_state_tensors += 1 + _log(f"Zeroed {zeroed_state_tensors} optimizer state tensors (exp_avg / exp_avg_sq)") + + for m in models: + m.zero_grad(set_to_none=True) + _log("Zeroed all model gradients") + + # ---- reset FP8 (always; no-op for BF16 / non-FP8 modules) ----------- + _log("Resetting FP8 state") + total_reset = 0 + for m in models: + total_reset += reset_fp8_state(m, reset_meta_tensors=True) + seed_fp8_amax(models, seed_value=1.0) + + # ---- reset FP4 (only when warmup explicitly used an FP4 recipe) ----- + total_fp4_reset = 0 + if recipe_kind == "fp4": + _log("Resetting FP4 state (no seeding)") + for m in models: + total_fp4_reset += reset_fp4_state(m, reset_meta_tensors=True) + + # ---- release cached allocator blocks -------------------------------- + torch.cuda.synchronize() + torch.cuda.empty_cache() + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + # ---- restore Megatron config ---------------------------------------- + for key, val in saved_config.items(): + setattr(config, key, val) + + nan_params = sum( + 1 + for m in models + for _, p in m.named_parameters() + if p.data.is_floating_point() and torch.isnan(p.data).any() + ) + elapsed = time.time() - t0 + _log( + f"Warmup complete in {elapsed:.1f}s " + f"(fp8_reset={total_reset}, fp4_reset={total_fp4_reset}, " + f"recipe={recipe_kind or 'native'}, nan_params={nan_params})" + ) diff --git a/primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py b/primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py new file mode 100644 index 000000000..d072e3d02 --- /dev/null +++ b/primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py @@ -0,0 +1,121 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +MoE identity-sort short-circuit patch. + +Migrated from the source patch ``megatron_moe_skip_identity_sort.patch`` +(mpo branch). + +In the all-to-all MoE token dispatcher, ``sort_chunks_by_idxs`` is called with +``sort_input_by_local_experts`` (dispatch) and ``restore_output_by_local_experts`` +(combine). When the EP/TP topology yields an *identity* permutation +(``[0, 1, ..., N-1]``) -- the common ``EP=1`` / ``TP=1`` case -- the call +degrades into a full-tensor ``split`` + ``cat`` round-trip that produces an +output identical to its input. This patch short-circuits those calls. + +Rather than editing ``third_party/Megatron-LM`` in place, we wrap the +``sort_chunks_by_idxs`` symbol bound inside the ``token_dispatcher`` module. +The identity decision is cached per index tensor (the topology indices are +created once per dispatcher and reused every step), so the ``torch.equal`` +probe -- and the device sync it implies -- happens at most once per tensor. + +Gate: + ``MOE_SKIP_IDENTITY_SORT`` (default ``1`` / enabled; set ``0`` to disable). +""" + +import os +import weakref + +import torch + +from primus.core.patches import PatchContext, register_patch +from primus.modules.module_utils import log_rank_0, warning_rank_0 + +# Cache the identity decision keyed on the index tensor's ``id()``. The topology +# index tensors persist for the dispatcher's lifetime, so caching by identity is +# stable across steps. A ``WeakKeyDictionary`` cannot be used here because weakref +# equality falls back to ``tensor == tensor`` (which yields a tensor, not a bool) +# and raises "Boolean value of Tensor ... is ambiguous". Instead we key on +# ``id(idxs)`` and register a ``weakref.finalize`` callback so the entry is +# evicted when the tensor is garbage-collected (preventing stale hits from +# ``id()`` reuse). +_IDENTITY_CACHE: dict = {} + + +def _skip_enabled(_ctx: PatchContext) -> bool: + return os.environ.get("MOE_SKIP_IDENTITY_SORT", "1") != "0" + + +def _is_identity(idxs) -> bool: + if idxs is None or not torch.is_tensor(idxs) or idxs.dim() != 1 or idxs.numel() == 0: + return False + key = id(idxs) + cached = _IDENTITY_CACHE.get(key) + if cached is not None: + return cached + result = bool(torch.equal(idxs, torch.arange(idxs.numel(), device=idxs.device, dtype=idxs.dtype))) + try: + # Evict on GC so a future tensor reusing this ``id()`` cannot hit a stale + # decision. + weakref.finalize(idxs, _IDENTITY_CACHE.pop, key, None) + _IDENTITY_CACHE[key] = result + except TypeError: + # Some tensors are not weak-referenceable; fall back to no caching. + pass + return result + + +def _make_wrapped_sort(orig_sort): + def sort_chunks_by_idxs(input, split_sizes, sorted_idxs, probs=None, fused=False): + # Identity permutation -> output == input, permuted_probs == probs. + if _is_identity(sorted_idxs): + return input, probs + return orig_sort(input, split_sizes, sorted_idxs, probs=probs, fused=fused) + + return sort_chunks_by_idxs + + +@register_patch( + "megatron.moe.skip_identity_sort", + backend="megatron", + phase="before_train", + description=( + "Short-circuit sort_chunks_by_idxs in the MoE token dispatcher when the " + "local-expert permutation is identity (e.g. EP=1/TP=1); gated by " + "MOE_SKIP_IDENTITY_SORT (default on)." + ), + condition=_skip_enabled, +) +def patch_moe_skip_identity_sort(ctx: PatchContext): + del ctx + + try: + from megatron.core.transformer.moe import token_dispatcher as td_mod + except ImportError as exc: + warning_rank_0( + f"[Patch:megatron.moe.skip_identity_sort] token_dispatcher not " f"importable; skipping: {exc}" + ) + return + + orig_sort = getattr(td_mod, "sort_chunks_by_idxs", None) + if orig_sort is None: + warning_rank_0( + "[Patch:megatron.moe.skip_identity_sort] sort_chunks_by_idxs not bound " + "in token_dispatcher; skipping." + ) + return + + if getattr(orig_sort, "_primus_skip_identity_sort", False): + return + + wrapped = _make_wrapped_sort(orig_sort) + wrapped._primus_skip_identity_sort = True + td_mod.sort_chunks_by_idxs = wrapped + log_rank_0( + "[Patch:megatron.moe.skip_identity_sort] Patched " + "token_dispatcher.sort_chunks_by_idxs to skip identity permutations." + ) 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 new file mode 100644 index 000000000..64f9275f8 --- /dev/null +++ b/primus/backends/megatron/patches/parallelism/sdma_param_all_gather_patches.py @@ -0,0 +1,200 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +SDMA (copy-engine) distributed-optimizer param all-gather patch. + +Migrated from the source patch ``megatron_sdma_allgather.patch`` (mpo branch). + +This replaces the in-place edits the source patch made to +``third_party/Megatron-LM`` with three runtime monkey-patches, all gated by +``ENABLE_SDMA_ALLGATHER=1``: + + 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 + delegated unchanged to the original method. + 2. ``DistributedDataParallel.__init__`` -- after construction, each bucket + group is annotated with ``param_gather_order`` (reverse dispatch order, + mirroring the source patch). + 3. MoE experts ``forward`` -- ``tokens_per_expert`` is moved onto the + dispatched-input device before the experts run (the source patch's + ``moe_layer.py`` one-liner). + +When the SDMA primitives (Primus-Turbo symmetric memory + ``hip``) are +unavailable, :func:`all_gather_into_tensor_sdma` falls back to +``torch.distributed.all_gather_into_tensor``, so enabling the flag is safe even +on images without the kernels. +""" + +import os + +import torch + +from primus.core.patches import PatchContext, register_patch +from primus.modules.module_utils import log_rank_0, warning_rank_0 + + +def _sdma_allgather_enabled(_ctx: PatchContext) -> bool: + return os.environ.get("ENABLE_SDMA_ALLGATHER", "0") == "1" + + +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 + + from primus.backends.megatron.core.distributed.sdma_param_gather import ( + _all_gather_into_tensor_waitable_fallback, + _WaitableHandle, + all_gather_into_tensor_sdma, + ) + + def start_param_sync(self, force_sync: bool = False): + # Layer-wise optimizer path is unchanged; only the distributed-optimizer + # path routes through SDMA. + if not self.ddp_config.use_distributed_optimizer: + return orig_start_param_sync(self, force_sync=force_sync) + + if force_sync: + if self.param_gather_handle is not None: + self.param_gather_handle.wait() + self.param_gather_handle = None + return + else: + assert self.param_gather_handle is None + + 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. + param_gather_order = getattr(self, "param_gather_order", None) + enable_sdma = os.getenv("ENABLE_SDMA_ALLGATHER") == "1" + 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 + ) + + param_gather_handles = [] + for idx, bucket in enumerate(self.buckets): + if self.cached_param_buffer_shard_list[idx] is None: + self.cached_param_buffer_shard_list[idx] = shard_buffer( + bucket.param_data, self.intra_distributed_optimizer_instance_size + ) + local_data_view = self.cached_param_buffer_shard_list[idx][ + self.intra_distributed_optimizer_instance_rank + ] + handle = all_gather_func( + bucket.param_data, + local_data_view, + group=self.intra_distributed_optimizer_instance_group, + async_op=async_op, + ) + if async_op: + param_gather_handles.append(handle) + + if async_op: + + def _wait_all_param_gathers(): + for handle in param_gather_handles: + handle.wait() + + self.param_gather_handle = _WaitableHandle(wait_fn=_wait_all_param_gathers) + else: + self.param_gather_handle = None + self.param_gather_dispatched = True + + return start_param_sync + + +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. + 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)): + bucket_group.param_gather_order = order + + return __init__ + + +def _make_wrapped_experts_forward(orig_forward): + def forward(self, permuted_local_hidden_states, tokens_per_expert, *args, **kwargs): + # Source patch moved tokens_per_expert onto the dispatched-input device + # before the experts run. + if torch.is_tensor(tokens_per_expert): + tokens_per_expert = tokens_per_expert.to(permuted_local_hidden_states.device) + return orig_forward(self, permuted_local_hidden_states, tokens_per_expert, *args, **kwargs) + + return forward + + +@register_patch( + "megatron.distributed.sdma_param_all_gather", + backend="megatron", + phase="before_train", + description=( + "Route the distributed-optimizer param all-gather through Primus-Turbo " + "SDMA (copy-engine) memcpys; gated by ENABLE_SDMA_ALLGATHER=1." + ), + condition=_sdma_allgather_enabled, +) +def patch_sdma_param_all_gather(ctx: PatchContext): + del ctx + + try: + import megatron.core.distributed.param_and_grad_buffer as pgb + from megatron.core.distributed.distributed_data_parallel import ( + DistributedDataParallel, + ) + except ImportError as exc: + warning_rank_0( + f"[Patch:megatron.distributed.sdma_param_all_gather] Megatron distributed " + f"modules not importable; skipping: {exc}" + ) + return + + bucket_group_cls = getattr(pgb, "_ParamAndGradBucketGroup", None) + if bucket_group_cls is None: + warning_rank_0( + "[Patch:megatron.distributed.sdma_param_all_gather] " + "_ParamAndGradBucketGroup not found; skipping." + ) + return + + if not getattr(bucket_group_cls, "_primus_sdma_param_gather_patched", False): + bucket_group_cls.start_param_sync = _make_start_param_sync(bucket_group_cls.start_param_sync) + bucket_group_cls._primus_sdma_param_gather_patched = True + + if not getattr(DistributedDataParallel, "_primus_sdma_param_gather_patched", False): + DistributedDataParallel.__init__ = _make_wrapped_ddp_init(DistributedDataParallel.__init__) + DistributedDataParallel._primus_sdma_param_gather_patched = True + + # MoE experts: move tokens_per_expert onto the dispatched-input device. + try: + from megatron.core.transformer.moe import experts as experts_mod + + for cls_name in ("GroupedMLP", "TEGroupedMLP", "SequentialMLP"): + cls = getattr(experts_mod, cls_name, None) + if cls is None or getattr(cls, "_primus_sdma_tokens_to_device", False): + continue + cls.forward = _make_wrapped_experts_forward(cls.forward) + cls._primus_sdma_tokens_to_device = True + except ImportError as exc: + warning_rank_0( + f"[Patch:megatron.distributed.sdma_param_all_gather] MoE experts module " + f"not importable; skipping tokens_per_expert device fixup: {exc}" + ) + + log_rank_0( + "[Patch:megatron.distributed.sdma_param_all_gather] Installed SDMA param " + "all-gather (start_param_sync + param_gather_order + experts tokens_per_expert)." + ) diff --git a/primus/backends/megatron/patches/te_patches/bshd_layout_patches.py b/primus/backends/megatron/patches/te_patches/bshd_layout_patches.py new file mode 100644 index 000000000..dba74f6e6 --- /dev/null +++ b/primus/backends/megatron/patches/te_patches/bshd_layout_patches.py @@ -0,0 +1,145 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Transformer Engine attention BSHD-layout patch. + +Migrated from the source patch ``megatron_te_bshd_layout.patch`` (mpo branch). + +Some TE FMHA kernels are faster with a ``bshd`` memory layout than with the +default ``sbhd`` layout Megatron feeds them. When ``NVTE_FMHA_USE_BSHD=1`` is +set, this patch transposes the ``sbhd`` ``query``/``key``/``value`` tensors to +``bshd`` before the TE ``DotProductAttention`` call and transposes the result +back to ``sbhd`` afterwards, so the change is transparent to the rest of +Megatron. + +Instead of editing ``third_party/Megatron-LM`` in place, we wrap +``TEDotProductAttention.forward`` and temporarily flip the effective +``qkv_format`` for the duration of the wrapped call so TE interprets the +already-transposed tensors correctly. + +Gate: + ``NVTE_FMHA_USE_BSHD=1`` (off by default; the patch is not installed at + all when unset, so there is zero overhead). +""" + +import functools +import os + +from primus.core.patches import PatchContext, register_patch +from primus.modules.module_utils import log_rank_0, warning_rank_0 + + +def _bshd_enabled(_ctx: PatchContext) -> bool: + return os.environ.get("NVTE_FMHA_USE_BSHD", "0") == "1" + + +def _effective_qkv_format(self, packed_seq_params): + """Mirror TEDotProductAttention.forward's qkv_format resolution.""" + if packed_seq_params is not None: + fmt = getattr(packed_seq_params, "qkv_format", None) + if fmt is not None: + return fmt + return self.qkv_format + + +def _make_wrapped_forward(orig_forward): + @functools.wraps(orig_forward) + def forward( + self, + query, + key, + value, + attention_mask, + attn_mask_type, + attention_bias=None, + packed_seq_params=None, + **kwargs, + ): + # Only convert when the effective layout is sbhd; otherwise behave + # exactly like the original. + if _effective_qkv_format(self, packed_seq_params) != "sbhd": + return orig_forward( + self, + query, + key, + value, + attention_mask, + attn_mask_type, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + **kwargs, + ) + + query = query.transpose(0, 1).contiguous() + key = key.transpose(0, 1).contiguous() + value = value.transpose(0, 1).contiguous() + + # Temporarily flip the qkv_format TE will read so the transposed + # tensors are interpreted as bshd. Restore it afterwards even on error. + if packed_seq_params is not None and hasattr(packed_seq_params, "qkv_format"): + restore_target, restore_value = packed_seq_params, packed_seq_params.qkv_format + packed_seq_params.qkv_format = "bshd" + else: + restore_target, restore_value = self, self.qkv_format + self.qkv_format = "bshd" + + try: + core_attn_out = orig_forward( + self, + query, + key, + value, + attention_mask, + attn_mask_type, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + **kwargs, + ) + finally: + restore_target.qkv_format = restore_value + + return core_attn_out.transpose(0, 1).contiguous() + + return forward + + +@register_patch( + "megatron.te.fmha_bshd_layout", + backend="megatron", + phase="before_train", + description=( + "Run TE DotProductAttention in bshd layout (transpose sbhd<->bshd " + "around the kernel) when NVTE_FMHA_USE_BSHD=1." + ), + condition=_bshd_enabled, +) +def patch_te_fmha_bshd_layout(ctx: PatchContext): + del ctx + + try: + from megatron.core.extensions import transformer_engine as te_ext + except ImportError as exc: + warning_rank_0( + f"[Patch:megatron.te.fmha_bshd_layout] transformer_engine extension " + f"not importable; skipping: {exc}" + ) + return + + cls = getattr(te_ext, "TEDotProductAttention", None) + if cls is None: + warning_rank_0("[Patch:megatron.te.fmha_bshd_layout] TEDotProductAttention not found; skipping.") + return + + if getattr(cls, "_primus_bshd_patched", False): + return + + cls.forward = _make_wrapped_forward(cls.forward) + cls._primus_bshd_patched = True + log_rank_0( + "[Patch:megatron.te.fmha_bshd_layout] Patched TEDotProductAttention.forward " + "to use bshd layout (NVTE_FMHA_USE_BSHD=1)." + ) diff --git a/primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py b/primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py new file mode 100644 index 000000000..86ff432e9 --- /dev/null +++ b/primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py @@ -0,0 +1,72 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Tier 1A — fused (residual + RMSNorm) patches. + +Wraps the runtime install hook +:mod:`primus.backends.megatron.core.extensions.fused_residual_rmsnorm` in a +``@register_patch`` so it runs at the standard ``before_train`` phase via +``run_patches(...)`` instead of being explicitly installed from the trainer +entry point. + +Two gates (mirroring the install hook): + * ``PRIMUS_FUSED_RESIDUAL_NORM=1`` — V1, in-layer ADD#1+norm fusion. + * ``PRIMUS_FUSED_RESIDUAL_NORM_V2=1`` — V2, cross-layer ADD#2 carry + (implies V1). + +Either env var enables the patch. The install hook itself further requires +``PrimusTurboRMSNorm`` (i.e. ``use_turbo_rms_norm=true``) and bails +gracefully otherwise, so it's safe to register unconditionally. +""" + +import os + +from primus.core.patches import PatchContext, register_patch +from primus.modules.module_utils import log_rank_0 + + +def _env_truthy(name: str) -> bool: + v = os.environ.get(name, "0").strip().lower() + return v in ("1", "true", "yes", "on") + + +def _is_fused_residual_norm_enabled(ctx: PatchContext) -> bool: + return _env_truthy("PRIMUS_FUSED_RESIDUAL_NORM") or _env_truthy("PRIMUS_FUSED_RESIDUAL_NORM_V2") + + +@register_patch( + "megatron.turbo.fused_residual_norm", + backend="megatron", + phase="before_train", + description=( + "Fuse residual+add into PrimusTurboRMSNorm via Primus-Turbo " + "rmsnorm_residual; gated by PRIMUS_FUSED_RESIDUAL_NORM(_V2)." + ), + condition=_is_fused_residual_norm_enabled, + # Run after megatron.turbo.rms_norm so PrimusTurboRMSNorm is in place + # before we extend its forward signature. + priority=60, +) +def patch_fused_residual_norm(ctx: PatchContext): + """Install the fused residual+RMSNorm runtime monkeypatch.""" + from primus.backends.megatron.core.extensions import fused_residual_rmsnorm + + log_rank_0( + "[Patch:megatron.turbo.fused_residual_norm] Installing fused " + "residual+RMSNorm (V2={v2}, V1={v1})".format( + v1=_env_truthy("PRIMUS_FUSED_RESIDUAL_NORM"), + v2=_env_truthy("PRIMUS_FUSED_RESIDUAL_NORM_V2"), + ) + ) + ok = fused_residual_rmsnorm.install() + if ok: + log_rank_0("[Patch:megatron.turbo.fused_residual_norm] install() returned True") + else: + log_rank_0( + "[Patch:megatron.turbo.fused_residual_norm] install() returned False " + "(precondition not met; e.g. use_turbo_rms_norm=false)" + ) diff --git a/primus/cli/main.py b/primus/cli/main.py index 600bc7702..26606fa8d 100644 --- a/primus/cli/main.py +++ b/primus/cli/main.py @@ -4,6 +4,12 @@ # See LICENSE for license information. ############################################################################### +# MLPerf log suppression must run before any heavy import (Megatron, TE, +# aiter, ...) that may print/log at import time. Importing this module only +# triggers the light ``primus/__init__.py`` and installs the FD-level filter +# when ``PRIMUS_LOG_SUPPRESSION=1`` is set; it is a complete no-op otherwise. +import primus.mlperf_log_suppression # noqa: F401 # isort: skip + import argparse import importlib import pkgutil diff --git a/primus/mlperf_log_suppression.py b/primus/mlperf_log_suppression.py new file mode 100644 index 000000000..2b11f16e9 --- /dev/null +++ b/primus/mlperf_log_suppression.py @@ -0,0 +1,273 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +# --------------------------------------------------------------------------- +# Non-MLLOG log suppression for MLPerf submission runs. +# +# The MLPerf reference/submission logs (":::MLLOG ...") plus the run banners +# emitted by the launch script are the only required stdout lines for a +# submission run. All other noisy framework output (Primus loguru banners, +# Megatron deprecations, aiter JIT build chatter, TE-RoPE, hipify, torchrun +# warnings, UserWarnings, ...) is suppressed when this module is enabled, to +# keep the run logs clean. +# +# Activation is GATED so that importing this module is a complete no-op for +# normal (non-MLPerf) Primus runs: +# +# PRIMUS_LOG_SUPPRESSION=1 -> enable the suppression (MLPerf launch sets it) +# PRIMUS_LOG_SUPPRESSION=0 -> default: install() is a no-op +# +# When enabled, verbosity is further controlled by: +# +# MLPERF_VERBOSE_LOGS=1 -> restore the full verbose output (old behaviour), +# i.e. enabled-but-not-quiet. +# MLPERF_VERBOSE_LOGS=0 -> quiet mode (default when enabled): only MLLOG + +# training timing/result lines are emitted. +# +# Strategy (strongly preferred over stdout scraping): +# 1. Export env vars that the noisy libraries honour (AITER_LOG_LEVEL, +# PYTHONWARNINGS, TRANSFORMERS_VERBOSITY, HF_HUB_DISABLE_PROGRESS_BARS). +# 2. Raise Python ``logging`` levels for the actual logger names emitting +# noise (``aiter``, ``megatron`` families, ``torch.distributed``, ...). +# Safe to call multiple times via ``reapply_quiet_logger_levels``. +# 3. Silence ``warnings`` category to match ``PYTHONWARNINGS=ignore``. +# +# A handful of sources (unconditional C++ ``std::cout`` in aiter / hipify; +# Primus ``loguru`` sinks that bypass ``logging``; raw ``print()`` in TE +# ``rope.py``; ...) cannot be controlled via the above. For those a narrow +# FD-level line filter is installed on stdout so even native writes get +# matched. Stderr carries no useful signal in quiet mode so FD 2 is +# redirected wholesale to ``/dev/null``. If you need stderr back (e.g. to +# debug a crash) re-run with ``MLPERF_VERBOSE_LOGS=1``. +# +# IMPORTANT: This module must be imported BEFORE any other import that may +# log/print at import time (Megatron, TE, aiter, ...). For the Primus CLI it +# is imported at the very top of ``primus/cli/main.py``. Because it lives +# directly under the ``primus`` package, importing it only triggers the light +# ``primus/__init__.py`` (config/logging utilities) and pulls in none of the +# heavy native libraries that emit the banners we want to hide. +# --------------------------------------------------------------------------- +import logging as _logging +import os as _os + +ENABLED = _os.environ.get("PRIMUS_LOG_SUPPRESSION", "0") == "1" +VERBOSE_LOGS = _os.environ.get("MLPERF_VERBOSE_LOGS", "0") == "1" + +# Logger names that emit the non-MLLOG noise we want to silence. Safe to +# set even when a given logger is not present in the process. +QUIET_LOGGER_NAMES = ( + # aiter: Primus also raises this logger to ERROR inside the trainer, + # but we do it earlier (before the aiter JIT build banner fires). + "aiter", + # Megatron-LM deprecation + rerun-state-machine warnings and the + # "Setting RerunStateMachine mode" warning from rerun_state_machine.py. + "megatron", + "megatron.core", + "megatron.core.utils", + "megatron.core.rerun_state_machine", + "megatron.core.pipeline_parallel", + # Primus / Primus-Turbo stdlib loggers (Primus routes most output + # through loguru instead, which we drop at the FD level below). + "primus", + "primus_turbo", + "primus_mllog", + # TransformerEngine import / RoPE banners that use stdlib logging. + "transformer_engine", + # torch.distributed elastic / launcher chatter. + "torch.distributed", + "torch.distributed.run", + "torch.distributed.elastic", + "torch.distributed.elastic.multiprocessing", + "torch.distributed.launcher.api", + # HuggingFace Hub / transformers progress + warnings. + "transformers", + "huggingface_hub", + # General misc. + "filelock", + "urllib3", +) + + +def reapply_quiet_logger_levels() -> None: + """Raise levels on noisy Python loggers. Safe to call multiple times. + + No-op unless suppression is enabled and quiet mode is active. + """ + if not ENABLED or VERBOSE_LOGS: + return + for _logger_name in QUIET_LOGGER_NAMES: + _logging.getLogger(_logger_name).setLevel(_logging.ERROR) + + +def _configure_non_mllog_logs_quiet() -> None: + """Silence every non-MLLOG log source we can control via env vars or + the standard ``logging`` / ``warnings`` modules. :::MLLOG output is + untouched because ``mlperf_logging`` has its own stdout handler. + """ + _os.environ.setdefault("AITER_LOG_LEVEL", "ERROR") + _os.environ.setdefault("AITER_LOG_MORE", "0") + _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error") + _os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") + _os.environ.setdefault("PYTHONWARNINGS", "ignore") + # Tell Primus-Turbo / GPT-OSS helpers to keep JIT tuning traces off. + _os.environ.setdefault("AITER_LOG_TUNED_CONFIG", "0") + + for _logger_name in QUIET_LOGGER_NAMES: + _logging.getLogger(_logger_name).setLevel(_logging.ERROR) + + # ``warnings.filterwarnings("ignore")`` covers cases where + # ``PYTHONWARNINGS=ignore`` has no effect because Python was started + # before we ran (e.g. re-exec'd by a launcher). + try: + import warnings as _warnings + + _warnings.filterwarnings("ignore") + except Exception: + pass + + +def _install_fd_level_fallback_filter() -> None: + """Last-resort line filter for logs that bypass Python ``logging`` + entirely: unconditional ``std::cout`` writes in aiter / hipify C++, + raw ``print()`` statements in TE / Megatron, and Primus loguru output + that landed on stdout. + + Design: + * FD 2 (stderr) is redirected straight to ``/dev/null``. In quiet + mode the only stderr sources observed in practice are Primus's + ``loguru`` sink, ``warnings.warn`` output, the hipify + ``"Successfully preprocessed all matching files."`` banner, and + torchrun's elastic SIGTERM chatter -- none of which carry a + training or MLLOG signal. + * FD 1 (stdout) is routed through an ``os.pipe`` + reader thread. + The thread runs each line through the suppression regexes and + forwards surviving lines to the saved original stdout FD. + * ``:::MLLOG`` lines always pass through. + """ + import re + import sys + import threading + + ansi_re = re.compile(r"\x1b\[[0-9;]*[ -/]*[@-~]") + + suppress_patterns = tuple( + re.compile(p) + for p in ( + # aiter JIT module-load / build / baton-wait banners. + r"^\[aiter\] ", + # TE-RoPE banner from transformer_engine rope.py (raw print). + r"^\[TE-RoPE\] ", + # "[MLPerf Train] ..." status prints -- all informational, all + # emitted N times in distributed mode. Re-run verbose to see them. + r"^\[MLPerf Train\] ", + # Primus runtime/CLI/patch informational prints that bypass loguru + # (e.g. "[Primus:Runtime] ...", "[Primus:Env] ...", "[Primus] ...", + # "[Primus CLI] ...", "[PrimusPatch] ..."). Covers the new-arch + # runtime/launcher print() calls. Re-run verbose to see them. + r"^\[Primus", + # Gloo C++ peer-connect banner ("[Gloo] Rank N is connected + # to 7 peer ranks. Expected number of connected peer ranks + # is : 7") -- native std::cout from libgloo. + r"^\[Gloo\] Rank \d+ is connected to ", + r"^Expected number of connected peer ranks is\s*:", + # hipify preprocessing banner (mostly fires on stderr, listed + # here as a safety net for any leak to stdout). + r"^Successfully preprocessed all matching files\.", + # Primus loguru lines, in case they end up on stdout for any + # reason (e.g. ``colorize=False`` configurations). The Primus + # format always starts with ``[YYYYMMDD HH:MM:SS]``. + r"^\[\d{8} \d{2}:\d{2}:\d{2}\]\[(?:rank|node)-\d+/", + # "Setting RerunStateMachine mode RerunMode.DISABLED" fires + # from Megatron's rerun_state_machine.py as a root-logger + # warning (the default Python logging handler dumps the message + # unformatted to stderr, but belt-and-braces: cover stdout too). + r"^Setting RerunStateMachine mode ", + ) + ) + + def _should_suppress(line: str) -> bool: + stripped = ansi_re.sub("", line) + if ":::MLLOG" in stripped: + return False + for pat in suppress_patterns: + if pat.search(stripped): + return True + return False + + def _start_reader(read_fd: int, out_fd: int) -> None: + def _run() -> None: + buf = b"" + try: + while True: + chunk = _os.read(read_fd, 4096) + if not chunk: + break + buf += chunk + while b"\n" in buf: + raw, buf = buf.split(b"\n", 1) + line = raw.decode("utf-8", errors="replace") + if not _should_suppress(line): + _os.write(out_fd, raw + b"\n") + except Exception: + # Never let the filter thread bring down the training run. + pass + finally: + if buf: + line = buf.decode("utf-8", errors="replace") + if not _should_suppress(line): + try: + _os.write(out_fd, buf) + except Exception: + pass + + threading.Thread(target=_run, daemon=True).start() + + # Flush any buffered Python stdout/stderr output before we steal the FDs. + sys.stdout.flush() + sys.stderr.flush() + + orig_stdout_fd = _os.dup(1) + + # Drop stderr unconditionally: nothing useful lands on FD 2 in quiet + # mode (Primus loguru, warnings, hipify). Users who need stderr back + # (e.g. crash debug) can re-run with ``MLPERF_VERBOSE_LOGS=1``. + devnull_fd = _os.open(_os.devnull, _os.O_WRONLY) + _os.dup2(devnull_fd, 2) + _os.close(devnull_fd) + + stdout_r, stdout_w = _os.pipe() + _os.dup2(stdout_w, 1) + _os.close(stdout_w) + + _start_reader(stdout_r, orig_stdout_fd) + + sys.stdout = _os.fdopen(1, "w", buffering=1, closefd=False) + sys.stderr = _os.fdopen(2, "w", buffering=1, closefd=False) + + +_INSTALLED = False + + +def install() -> None: + """Install the quiet-mode suppression exactly once per process. + + No-op unless ``PRIMUS_LOG_SUPPRESSION=1``. Calling this a second time is + also a no-op. + """ + global _INSTALLED + if _INSTALLED: + return + _INSTALLED = True + if not ENABLED: + return + if VERBOSE_LOGS: + return + _configure_non_mllog_logs_quiet() + _install_fd_level_fallback_filter() + + +# Auto-install on import (gated by PRIMUS_LOG_SUPPRESSION). +install() diff --git a/requirements.txt b/requirements.txt index 68a16522c..9d3c7ac6e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,3 +15,4 @@ mlflow==3.11.1 pyrsmi plotext hip-python +git+https://github.com/mlcommons/logging.git@6.0.0-rc5 From 7a447fec83a2e04193bc097f480b4655c893f668 Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Wed, 8 Jul 2026 09:43:47 +0800 Subject: [PATCH 010/127] feat: add AITER_LOG_LEVEL to suppress log (#849) # Description This PR adds a new global runner hook that sets `AITER_LOG_LEVEL=ERROR` to suppress the verbose AITER logs during training runs. AITER emits a large amount of log output by default, which clutters the run logs and makes it harder to spot the relevant training information. The new hook raises the AITER log level to `ERROR` so that only errors are surfaced. The hook emits an `env.*` line that is exported by `execute_hooks.sh`, following the existing global-hook convention under `runner/helpers/hooks/`. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [x] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Add `runner/helpers/hooks/03_enable_aiter.sh`, a global hook that exports `AITER_LOG_LEVEL=ERROR` to suppress the verbose AITER logs. # Checklist: - [x] The functionality is complete - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --- runner/helpers/hooks/03_enable_aiter.sh | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 runner/helpers/hooks/03_enable_aiter.sh diff --git a/runner/helpers/hooks/03_enable_aiter.sh b/runner/helpers/hooks/03_enable_aiter.sh new file mode 100644 index 000000000..67a0d0fd8 --- /dev/null +++ b/runner/helpers/hooks/03_enable_aiter.sh @@ -0,0 +1,19 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Global hook: enable AITER. +# +# Trigger: +# export AITER_LOG_LEVEL=ERROR +# +# This hook emits env.* lines which will be exported by execute_hooks.sh. +# + +set -euo pipefail + +# Set AITER log level to ERROR to suppress the verbose logs. +echo "env.AITER_LOG_LEVEL=ERROR" From bd124bf83bf021389cf65f1ba5e8e4b7d9e2045e Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Wed, 8 Jul 2026 09:45:58 +0800 Subject: [PATCH 011/127] [Megatron-LM] fix: duplicated memory footprint when enable turbo grouped gemm (#850) # Description This PR fixes duplicated GPU memory usage in `PrimusTurboGroupedLinear` when turbo grouped GEMM is enabled for MoE expert layers. `PrimusTurboGroupedLinear` consolidates per-expert `weight{i}` parameters into a single `self.weights` tensor for grouped GEMM execution. The previous implementation had two issues that left an extra copy of the consolidated weights resident on GPU: 1. `buffer.clone()` was used when registering `self.weights`, allocating a redundant tensor. 2. Per-expert `weight{i}` views were registered immediately in `__init__`. Those views pinned the pre-DDP-remap storage. After the distributed optimizer remapped `self.weights` into the param buffer, both the old pinned storage and the remapped buffer remained on GPU. This change registers `self.weights` directly from the consolidation buffer and defers `weight{i}` view creation until after DDP param-buffer remapping. Views are created lazily on the first forward pass (via a forward pre-hook) or when `state_dict()` is called, preserving checkpoint and legacy `weight{i}` lookup compatibility without retaining duplicate weight storage. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Remove the unnecessary `buffer.clone()` when registering the consolidated `self.weights` parameter in `PrimusTurboGroupedLinear.__init__`. - Defer per-expert `weight{i}` view registration until after DDP distributed-optimizer param-buffer remapping, avoiding pinned pre-remap storage. - Add `_ensure_weight_views()` with lazy registration triggered by a forward pre-hook and overridden `state_dict()`. - Preserve per-expert weight metadata via `_saved_weight_attrs` so checkpoint and state-dict code paths continue to work. # Checklist: - [x] The functionality is complete - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --- .../megatron/core/extensions/primus_turbo.py | 48 +++++++++++-------- 1 file changed, 28 insertions(+), 20 deletions(-) diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index 8b787044d..b565bf01d 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -1625,9 +1625,7 @@ def __init__( weight = getattr(self, f"weight{i}") buffer[i].copy_(weight) - weights = buffer.clone() - - self.register_parameter("weights", torch.nn.Parameter(weights)) + self.register_parameter("weights", torch.nn.Parameter(buffer)) # Capture the per-expert weights' extra attributes BEFORE deleting them. saved_weight_attrs = [dict(getattr(self, f"weight{i}").__dict__) for i in range(self.num_gemms)] @@ -1643,31 +1641,41 @@ def __init__( name = f"weight{i}" if name in self._parameters: del self._parameters[name] - del buffer gc.collect() torch.cuda.empty_cache() - # Re-expose each expert's slice as a zero-copy weight{i} Parameter view - # of self.weights, so existing code paths and checkpoints that look up - # weight{i} by name keep working without allocating a new buffer. - # ``requires_grad=False`` is required: self.weights is the canonical - # trainable Parameter and these views share its storage, so leaving - # them trainable would make the optimizer (and DDP) update / sync the - # same memory twice. ``.detach()`` strips the view's autograd graph - # so each Parameter ends up as a leaf with ``_base is None`` (which is - # what Megatron's distributed-optimizer param-bucket re-mapping - # expects), while still aliasing the same underlying storage. - # We also restore each weight{i}'s saved extra attributes so checkpoint / - # state-dict code that inspects them keeps seeing the right markers. + # Defer weight{i} view registration until after DDP has remapped + # self.weights into the distributed-optimizer param buffer. Registering + # views here would pin the pre-remap storage and leave a duplicate copy + # of the consolidated weights resident on GPU. + self._saved_weight_attrs = saved_weight_attrs + self._weight_views_registered = False + self.register_forward_pre_hook(self._forward_pre_hook_ensure_weight_views) + + self.register_buffer("quantized_weight_buffer", None, persistent=False) + self.register_buffer("quantized_weight_t_buffer", None, persistent=False) + + def _ensure_weight_views(self) -> None: + """Register per-expert weight{i} views after DDP param-buffer remap.""" + if self._weight_views_registered: + return + for i in range(self.num_gemms): - weight_i = torch.nn.Parameter(self.weights[i].detach(), requires_grad=False) - for attr_name, attr_val in saved_weight_attrs[i].items(): + weight_i = torch.nn.Parameter(self.weights[i], requires_grad=False) + for attr_name, attr_val in self._saved_weight_attrs[i].items(): setattr(weight_i, attr_name, attr_val) self.register_parameter(f"weight{i}", weight_i) - self.register_buffer("quantized_weight_buffer", None, persistent=False) - self.register_buffer("quantized_weight_t_buffer", None, persistent=False) + self._weight_views_registered = True + + @staticmethod + def _forward_pre_hook_ensure_weight_views(module, _inputs): + module._ensure_weight_views() + + def state_dict(self, *args, **kwargs): + self._ensure_weight_views() + return super().state_dict(*args, **kwargs) def forward(self, x: torch.Tensor, m_splits: torch.Tensor): _is_first_microbatch = self.is_first_microbatch From 80e680eed4c3e55708b4db887d97088c741f7a02 Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Wed, 8 Jul 2026 09:50:32 +0800 Subject: [PATCH 012/127] opt: remove grouped mlp d2h sync (#859) # Description This PR removes an unnecessary GPU-to-CPU (d2h) synchronization in `PrimusGroupedMLP.forward()` when applying the final expert output bias. Previously, `tokens_per_expert.tolist()` was called unconditionally at the end of `forward()` before `_apply_bias`, forcing a device sync on every MoE forward pass even when `output_bias` is `None` (the common case when bias is already fused into TE GroupedLinear output). The fix overrides `_apply_bias` to defer the `.tolist()` conversion until bias application is actually needed, and to skip it entirely when `bias_parallel` is `None`. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [x] Code refactoring ## Changes Please list the changes introduced in this PR: - Add a `PrimusGroupedMLP._apply_bias` static override that returns early when `bias_parallel` is `None`, avoiding the d2h sync in the no-bias path. - Move `tokens_per_expert.tolist()` from `forward()` into `_apply_bias`, so the CPU conversion only runs when bias must be applied. - Update the `forward()` call site to pass the GPU `tokens_per_expert` tensor directly to `_apply_bias` instead of a pre-materialized Python list. # Checklist: - [x] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --- .../backends/megatron/core/transformer/experts.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/primus/backends/megatron/core/transformer/experts.py b/primus/backends/megatron/core/transformer/experts.py index eb2d51e05..847014b4d 100644 --- a/primus/backends/megatron/core/transformer/experts.py +++ b/primus/backends/megatron/core/transformer/experts.py @@ -77,6 +77,18 @@ def bias_act_func_with_mask( # use the original bias_act_func from TEGroupedMLP, ignore the tokens_per_experts return self.bias_act_func(intermediate_parallel, bias_parallel, permuted_probs) + @staticmethod + def _apply_bias(intermediate_parallel, bias_parallel, tokens_per_expert, permuted_probs): + if bias_parallel is None: + return intermediate_parallel + + # NOTE: tokens_per_expert is on GPU, so we need to convert it to a list of ints. + tokens_per_expert_cpu = tokens_per_expert.tolist() + + return super()._apply_bias( + intermediate_parallel, bias_parallel, tokens_per_expert_cpu, permuted_probs + ) + def forward( self, permuted_local_hidden_states: torch.Tensor, @@ -151,8 +163,7 @@ def forward( # to make sure the fc1_output is reloaded to GPU before recomputing moe_act. if self.offload_moe_act: output = off_interface.group_commit(output, name="moe_act", forced_released_tensors=[fc1_output]) - # NOTE: tokens_per_expert is on GPU, so we need to convert it to a list of ints. - output = self._apply_bias(output, output_bias, tokens_per_expert.tolist(), permuted_probs) + 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): From 629fb80454590ab26d1d5cdd9a0cde4d3b59f769 Mon Sep 17 00:00:00 2001 From: zirui Date: Wed, 8 Jul 2026 13:27:34 +0800 Subject: [PATCH 013/127] Add diffusion backend & Wan training support (#779) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary This PR integrates diffusion training support into Primus training framework, including WAN2.1 and WAN2.2 models. The goal is to enable stable single-node and multi-node training for WAN diffusion models with minimal framework changes and clean extensibility. --- ## Motivation WAN diffusion models are now a primary workload for video generation training. This integration enables: - Unified training interface for diffusion models - Scalable single-node → multi-node training - Reuse of existing Primus training infra (optimizer, checkpointing, logging) --- ## Scope ### Included - WAN2.1 / WAN2.2 model wrapper integration - Diffusion training loop adaptation - Basic dataset pipeline support - Single-node SFT training validation - Checkpoint save/load compatibility ### Not included (future work) - Inference pipeline - Advanced scheduling strategies - Flash attention / kernel-level optimizations --- ## Current Status ### Completed - [x] WAN model integration (WAN2.1-1.3B / WAN2.2-5B) - [x] Training loop adapted for diffusion objective - [x] Single-node SFT training verified ### In Progress - [ ] Multi-node training validation (2-node cluster test) - [x] Config refactor (clean separation of model/trainer/data) - [x] Dataset pipeline cleanup - [x] Documentation and example scripts --- ## Testing ### Single-node - Wan2.1-1.3B SFT training: ✅ - Wan2.2-5B SFT training: ✅ ### Multi-node - 2-node test: in progress ## Known Issues / Risks - Dataset preprocessing still partially ad-hoc - Config system needs refactor for diffusion-specific parameters - Multi-node stability not fully validated yet ## Next Steps 1. Complete 2-node validation 2. Add training performance benchmarks 3. Refactor config structure for diffusion training 4. Add reproducible example scripts --- ## Notes This is a WIP draft PR. Frequent commits will be pushed as development continues. ## benchmarks ## Primus Wan2.2 TI2V 5B Benchmark - 2026-07-06 10:28:13 UTC - Summary uses median step time after skipping first 5 logged steps. - `batch_mode=local_accum` means effective per-GPU batch size via gradient accumulation with micro batch 1. - Stability/efficiency fixes enabled: `video_backend=decord`, `dataloader_num_workers=0`, `PRIMUS_CACHE_RAW_VIDEO_FRAMES=1`, and `PRIMUS_CACHE_PROCESSED_SAMPLES=1`. | engine | model | data | resolution | frames | batch_size | batch_mode | gpus | steps | gpu_mem_GB | step_time_s | step/s | TPS(samples/s/gpu) | status | | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | | primus-fsdp2-flash_attn_aiter | wan2.2-ti2v-5b | tiny-video-sample | 480P | 121 | 1 | local_accum | 8 | 95 | 27.080 | 1.760000 | 0.568182 | 0.569300 | ok | | primus-fsdp2-flash_attn_aiter | wan2.2-ti2v-5b | tiny-video-sample | 480P | 121 | 8 | local_accum | 8 | 95 | 28.330 | 14.330000 | 0.069784 | 0.558300 | ok | | primus-fsdp2-flash_attn_aiter | wan2.2-ti2v-5b | tiny-video-sample | 480P | 121 | 16 | local_accum | 8 | 95 | 28.330 | 28.430000 | 0.035174 | 0.562800 | ok | --------- Co-authored-by: Cursor Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- examples/diffusion/README.md | 74 ++ .../MI355X/wan2.1_t2v_1.3b-posttrain.yaml | 53 ++ .../MI355X/wan2.1_t2v_1.3b-pretrain.yaml | 53 ++ .../MI355X/wan2.2_ti2v_5b-posttrain.yaml | 53 ++ .../MI355X/wan2.2_ti2v_5b-pretrain.yaml | 53 ++ primus/backends/diffusion/README.md | 237 +++++ primus/backends/diffusion/__init__.py | 12 + primus/backends/diffusion/argument_builder.py | 262 ++++++ .../backends/diffusion/attention/__init__.py | 29 + .../diffusion/attention/_flash_common.py | 130 +++ primus/backends/diffusion/attention/aiter.py | 71 ++ .../backends/diffusion/attention/attention.py | 410 ++++++++ primus/backends/diffusion/attention/flex.py | 280 ++++++ primus/backends/diffusion/data/__init__.py | 7 + .../diffusion/data/registrations/__init__.py | 7 + .../diffusion/data/registrations/wan.py | 176 ++++ .../backends/diffusion/diffusion_adapter.py | 67 ++ .../diffusion/diffusion_pretrain_trainer.py | 131 +++ .../diffusion/distributed/__init__.py | 32 + .../diffusion/distributed/checkpoint.py | 114 +++ primus/backends/diffusion/distributed/mesh.py | 126 +++ .../backends/diffusion/distributed/ulysses.py | 211 +++++ primus/backends/diffusion/models/__init__.py | 13 + primus/backends/diffusion/models/interface.py | 48 + .../models/registrations/__init__.py | 7 + .../diffusion/models/registrations/wan.py | 226 +++++ .../backends/diffusion/models/wan/__init__.py | 20 + .../backends/diffusion/models/wan/adapter.py | 154 +++ .../diffusion/models/wan/attention_backend.py | 228 +++++ .../diffusion/models/wan/components.py | 26 + .../models/wan/configuration_wanvideo.py | 84 ++ primus/backends/diffusion/models/wan/t5.py | 299 ++++++ .../diffusion/models/wan/train_pipeline.py | 314 +++++++ .../backends/diffusion/models/wan/vae2_1.py | 646 +++++++++++++ .../backends/diffusion/models/wan/vae2_2.py | 873 ++++++++++++++++++ .../backends/diffusion/models/wan/wan_dit.py | 580 ++++++++++++ .../diffusion/optim/adamw_fp32_state.py | 111 +++ primus/backends/diffusion/registry.py | 66 ++ .../diffusion/schedulers/flow_match.py | 130 +++ .../backends/diffusion/trainers/__init__.py | 11 + primus/backends/diffusion/trainers/base.py | 551 +++++++++++ primus/backends/diffusion/trainers/fsdp2.py | 395 ++++++++ primus/backends/diffusion/utils/__init__.py | 10 + primus/backends/diffusion/utils/data_utils.py | 56 ++ primus/backends/diffusion/utils/log.py | 42 + .../backends/diffusion/utils/train_utils.py | 201 ++++ .../diffusion/utils/vision_process.py | 583 ++++++++++++ .../models/diffusion/wan2.1_t2v_1.3b.yaml | 13 + .../models/diffusion/wan2.1_t2v_1.3b_sft.yaml | 8 + .../models/diffusion/wan2.2_ti2v_5b.yaml | 13 + .../models/diffusion/wan2.2_ti2v_5b_sft.yaml | 8 + .../modules/diffusion/post_trainer.yaml | 5 + .../modules/diffusion/pre_trainer.yaml | 4 + .../train/posttrain/diffusion/prepare.py | 24 + .../diffusion/00_install_requirements.sh | 40 + .../hooks/train/pretrain/diffusion/prepare.py | 132 +++ .../diffusion/requirements-diffusion.txt | 12 + .../diffusion/test_wan_argument_builder.py | 194 ++++ .../diffusion/test_wan_trainer_optimizer.py | 61 ++ 59 files changed, 8776 insertions(+) create mode 100644 examples/diffusion/README.md create mode 100644 examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml create mode 100644 examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml create mode 100644 examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml create mode 100644 examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml create mode 100644 primus/backends/diffusion/README.md create mode 100644 primus/backends/diffusion/__init__.py create mode 100644 primus/backends/diffusion/argument_builder.py create mode 100644 primus/backends/diffusion/attention/__init__.py create mode 100644 primus/backends/diffusion/attention/_flash_common.py create mode 100644 primus/backends/diffusion/attention/aiter.py create mode 100644 primus/backends/diffusion/attention/attention.py create mode 100644 primus/backends/diffusion/attention/flex.py create mode 100644 primus/backends/diffusion/data/__init__.py create mode 100644 primus/backends/diffusion/data/registrations/__init__.py create mode 100644 primus/backends/diffusion/data/registrations/wan.py create mode 100644 primus/backends/diffusion/diffusion_adapter.py create mode 100644 primus/backends/diffusion/diffusion_pretrain_trainer.py create mode 100644 primus/backends/diffusion/distributed/__init__.py create mode 100644 primus/backends/diffusion/distributed/checkpoint.py create mode 100644 primus/backends/diffusion/distributed/mesh.py create mode 100644 primus/backends/diffusion/distributed/ulysses.py create mode 100644 primus/backends/diffusion/models/__init__.py create mode 100644 primus/backends/diffusion/models/interface.py create mode 100644 primus/backends/diffusion/models/registrations/__init__.py create mode 100644 primus/backends/diffusion/models/registrations/wan.py create mode 100644 primus/backends/diffusion/models/wan/__init__.py create mode 100644 primus/backends/diffusion/models/wan/adapter.py create mode 100644 primus/backends/diffusion/models/wan/attention_backend.py create mode 100644 primus/backends/diffusion/models/wan/components.py create mode 100644 primus/backends/diffusion/models/wan/configuration_wanvideo.py create mode 100644 primus/backends/diffusion/models/wan/t5.py create mode 100644 primus/backends/diffusion/models/wan/train_pipeline.py create mode 100644 primus/backends/diffusion/models/wan/vae2_1.py create mode 100644 primus/backends/diffusion/models/wan/vae2_2.py create mode 100644 primus/backends/diffusion/models/wan/wan_dit.py create mode 100644 primus/backends/diffusion/optim/adamw_fp32_state.py create mode 100644 primus/backends/diffusion/registry.py create mode 100644 primus/backends/diffusion/schedulers/flow_match.py create mode 100644 primus/backends/diffusion/trainers/__init__.py create mode 100644 primus/backends/diffusion/trainers/base.py create mode 100644 primus/backends/diffusion/trainers/fsdp2.py create mode 100644 primus/backends/diffusion/utils/__init__.py create mode 100644 primus/backends/diffusion/utils/data_utils.py create mode 100644 primus/backends/diffusion/utils/log.py create mode 100644 primus/backends/diffusion/utils/train_utils.py create mode 100644 primus/backends/diffusion/utils/vision_process.py create mode 100644 primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml create mode 100644 primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml create mode 100644 primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml create mode 100644 primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml create mode 100644 primus/configs/modules/diffusion/post_trainer.yaml create mode 100644 primus/configs/modules/diffusion/pre_trainer.yaml create mode 100644 runner/helpers/hooks/train/posttrain/diffusion/prepare.py create mode 100755 runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh create mode 100644 runner/helpers/hooks/train/pretrain/diffusion/prepare.py create mode 100644 runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt create mode 100644 tests/unit_tests/backends/diffusion/test_wan_argument_builder.py create mode 100644 tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md new file mode 100644 index 000000000..2df6093f7 --- /dev/null +++ b/examples/diffusion/README.md @@ -0,0 +1,74 @@ +# Wan Examples + +Wan examples exercise the independent PyTorch Diffusion backend under +`primus/backends/diffusion`. For backend details, data/checkpoint layout, and minimal +configs, see `primus/backends/diffusion/README.md`. + +## Data + +The default smoke-test dataset is `zirui3/tiny-video-samples` on Hugging Face: + +```bash +huggingface-cli download zirui3/tiny-video-samples \ + --repo-type dataset \ + --local-dir /data/tiny-video-samples +``` + +Expected layout: + +```text +/data/tiny-video-samples/ + meta.jsonl + data/*.mp4 +``` + +## Run + +Set the shared `torchrun` environment first: + +```bash +export NNODES=${NNODES:-1} +export NODE_RANK=${NODE_RANK:-0} +export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} +export MASTER_PORT=${MASTER_PORT:-29500} +export GPUS_PER_NODE=${GPUS_PER_NODE:-8} +``` + +### Pretrain + +```bash +DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ +DATA_FOLDER=/data/tiny-video-samples/data \ +ATTENTION_BACKEND=flash_attn_aiter \ +SP_SIZE=1 \ +MAX_STEPS=10 \ +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ + -m primus.cli.main train pretrain \ + --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml +``` + +Use `SP_SIZE=4` or `SP_SIZE=8` to enable Ulysses sequence parallelism +when the model head count supports it. + +### Posttrain + +```bash +INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B \ +DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ +DATA_FOLDER=/data/tiny-video-samples/data \ +MAX_STEPS=10 \ +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ + -m primus.cli.main train posttrain \ + --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml +``` + +The MI355X configs use Primus-style override sections such as `training`, +`data`, `parallelism`, `optimizer`, `runtime`, and `metrics`. The diffusion +adapter normalizes those sections into the Wan model/dataset/trainer +arguments at runtime. diff --git a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml new file mode 100644 index 000000000..c7f93c421 --- /dev/null +++ b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml @@ -0,0 +1,53 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:wan2.1_t2v_1.3b-posttrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + post_trainer: + framework: diffusion + config: post_trainer.yaml + + # Model preset to fine-tune from INIT_CHECKPOINT. + model: wan2.1_t2v_1.3b_sft.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: 1 + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/wan2.1_t2v_1.3b-posttrain} + save_steps: 0 + run_name: wan2.1_t2v_1.3b-posttrain + + data: + dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} + data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} + frame_num: 81 + video_backend: imageio + text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.1-T2V-1.3B/google/umt5-xxl} + height: 480 + width: 832 + + parallelism: + sp_size: ${SP_SIZE:1} + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: 5.0e-6 + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} + report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml new file mode 100644 index 000000000..5f17780f1 --- /dev/null +++ b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml @@ -0,0 +1,53 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:wan2.1_t2v_1.3b-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + + # Model preset to run. + model: wan2.1_t2v_1.3b.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: 1 + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/wan2.1_t2v_1.3b-pretrain} + save_steps: 0 + run_name: wan2.1_t2v_1.3b-pretrain + + data: + dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} + data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} + frame_num: 81 + video_backend: imageio + text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.1-T2V-1.3B/google/umt5-xxl} + height: 480 + width: 832 + + parallelism: + sp_size: ${SP_SIZE:1} + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: 1.0e-5 + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} + report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml new file mode 100644 index 000000000..203aa13ff --- /dev/null +++ b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml @@ -0,0 +1,53 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:wan2.2_ti2v_5b-posttrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + post_trainer: + framework: diffusion + config: post_trainer.yaml + + # Model preset to fine-tune from INIT_CHECKPOINT. + model: wan2.2_ti2v_5b_sft.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: 1 + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/wan2.2_ti2v_5b-posttrain} + save_steps: 0 + run_name: wan2.2_ti2v_5b-posttrain + + data: + dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} + data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} + frame_num: 81 + video_backend: imageio + text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.2-TI2V-5B/google/umt5-xxl} + height: 480 + width: 832 + + parallelism: + sp_size: ${SP_SIZE:1} + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: 5.0e-6 + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} + report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml new file mode 100644 index 000000000..7eed61bdd --- /dev/null +++ b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml @@ -0,0 +1,53 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:wan2.2_ti2v_5b-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + + # Model preset to run. + model: wan2.2_ti2v_5b.yaml + overrides: + sink_level: null + file_sink_level: DEBUG + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: 1 + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/wan2.2_ti2v_5b-pretrain} + save_steps: 0 + run_name: wan2.2_ti2v_5b-pretrain + + data: + dataset_path: ${DATASET_PATH:/data/tiny-video-samples/meta.jsonl} + data_folder: ${DATA_FOLDER:/data/tiny-video-samples/data} + frame_num: 81 + video_backend: imageio + text_tokenizer: ${TEXT_TOKENIZER:/models/Wan2.2-TI2V-5B/google/umt5-xxl} + height: 480 + width: 832 + + parallelism: + sp_size: ${SP_SIZE:1} + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: 1.0e-5 + weight_decay: 0.01 + + runtime: + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} + report_to: none diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md new file mode 100644 index 000000000..9b7db9704 --- /dev/null +++ b/primus/backends/diffusion/README.md @@ -0,0 +1,237 @@ +# Primus Diffusion Backend + +`diffusion` integrates PyTorch diffusion-model training as an independent Primus +backend. Primus provides the config/launch entrypoint, while model, dataset, +attention, and FSDP2 training logic are owned by the in-tree Wan implementation +under `primus/backends/diffusion`. + +All runtime code resolves through the Primus namespace, for example +`primus.backends.diffusion.models` and `primus.backends.diffusion.trainers`, +so Wan training is a first-class part of the Primus `diffusion` backend. + +Supported scope: + +- Model implementation: `wan` for Wan2.1 and Wan2.2. +- Trainer: FSDP2 only. +- Sequence parallelism: Ulysses SP via `trainer.args.sp_size`. + +Wan-specific dependencies are kept out of top-level Primus requirements. See +`runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt`. + +## Data + +Wan training reads a jsonl metadata file and a video folder: + +```jsonl +{"prompt": "text prompt", "video": "example.mp4"} +``` + +The default small dataset for smoke tests is the Hugging Face dataset +`zirui3/tiny-video-samples`: + +```bash +huggingface-cli download zirui3/tiny-video-samples \ + --repo-type dataset \ + --local-dir /data/tiny-video-samples +``` + +This produces: + +```text +/data/tiny-video-samples/ + meta.jsonl + data/*.mp4 +``` + +Use these public config fields to point training at another dataset: + +```yaml +data: + dataset_path: /path/to/meta.jsonl + data_folder: /path/to/videos + video_backend: imageio +``` + +## Checkpoints + +Each Wan model is a single Hugging Face repo that already bundles everything Wan +training needs: the DiT weights, the UMT5-XXL text encoder +(`models_t5_umt5-xxl-enc-bf16.pth`), the VAE (`Wan2.1_VAE.pth` / +`Wan2.2_VAE.pth`), and the tokenizer under `google/umt5-xxl`. Download the +model(s) you plan to train: + +| Model | Preset | Hugging Face repo | +| --- | --- | --- | +| Wan2.1-T2V-1.3B | `wan2.1_t2v_1.3b.yaml` | [Wan-AI/Wan2.1-T2V-1.3B](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B) | +| Wan2.2-TI2V-5B | `wan2.2_ti2v_5b.yaml` | [Wan-AI/Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) | +| Wan2.1-T2V-14B | (no shipped preset) | [Wan-AI/Wan2.1-T2V-14B](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B) | + +```bash +# Wan2.1-T2V-1.3B +huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B \ + --local-dir /models/Wan2.1-T2V-1.3B + +# Wan2.2-TI2V-5B +huggingface-cli download Wan-AI/Wan2.2-TI2V-5B \ + --local-dir /models/Wan2.2-TI2V-5B +``` + +A downloaded Wan repo looks like this (the T5 encoder, VAE, and tokenizer are +shipped inside the same repo, so no separate download is required): + +```text +/models/Wan2.1-T2V-1.3B/ + config.json + diffusion_pytorch_model*.safetensors # DiT weights + models_t5_umt5-xxl-enc-bf16.pth # T5 (UMT5-XXL) text encoder + Wan2.1_VAE.pth # VAE (Wan2.2_VAE.pth in the 5B repo) + google/umt5-xxl/ # tokenizer +``` + +The model presets reference these paths by default. Override any asset with +environment variables: + +```bash +export PRETRAINED_PATH=/models/Wan2.1-T2V-1.3B # pretrain DiT init +export INIT_CHECKPOINT=/models/Wan2.1-T2V-1.3B # post-train/SFT DiT init +export TEXT_TOKENIZER=/models/Wan2.1-T2V-1.3B/google/umt5-xxl +export TEXT_ENCODER=/models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth +export VAE_CHECKPOINT=/models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth +``` + +Pretrain presets use `PRETRAINED_PATH`; post-train SFT presets use +`INIT_CHECKPOINT` for DiT initialization. For Wan2.2-5B use the matching +`/models/Wan2.2-TI2V-5B` directory and `Wan2.2_VAE.pth`. + +## Pre-flight validation + +The diffusion prepare hook validates the prepared assets; it does **not** +download them. Download the dataset and checkpoints above first, then run the +hook to confirm the configured dataset, tokenizer, and DiT/T5/VAE paths exist +before launching distributed training: + +```bash +# pretrain config (validates modules.pre_trainer) +python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ + --config examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml + +# post-train config (validates modules.post_trainer) +python3 runner/helpers/hooks/train/posttrain/diffusion/prepare.py \ + --config examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml +``` + +On success it prints `env.PREPARED=1`. Set `SKIP_PREPARE=1` to bypass the check +for debugging. + +## Launch + +Training runs through the Primus CLI (`primus.cli.main train `) +under `torchrun`. Single-node and multi-node share **one** launch command: the +same `torchrun` invocation runs on every node, parameterized by the standard +rendezvous variables. The defaults below give a single-node 8-GPU run. + +```bash +# distributed knobs (defaults = single node, 8 GPUs) +export NNODES=${NNODES:-1} +export NODE_RANK=${NODE_RANK:-0} +export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} +export MASTER_PORT=${MASTER_PORT:-29500} +export GPUS_PER_NODE=${GPUS_PER_NODE:-8} + +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ + -m primus.cli.main train pretrain --config /path/to/wan_config.yaml +``` + +- **Single node**: run as-is (the defaults above). +- **Multi-node**: run the same command on each node with a shared + `MASTER_ADDR`/`MASTER_PORT` (a routable IP of node rank 0) and a distinct + `NODE_RANK` per node. World size is `NNODES * GPUS_PER_NODE`. +- **Post-train**: identical command with `train posttrain` and a posttrain config. + +Useful runtime knobs: + +- `trainer.args.attention_backend`: defaults to `flash_attn_aiter` for Wan training; use `sdpa` as the portable fallback or baseline. +- `trainer.args.sp_size`: Ulysses sequence parallel size. It must divide the + model attention head count; for example Wan2.1-1.3B supports `sp_size=4` but + not `sp_size=8`. +- `trainer.args.dp_replicate`: data parallel replication size. +- `FIXED_TIMESTEP` and `FIXED_SEED`: optional debug variables for reproducible + loss-alignment checks. + +On ROCm clusters, point compiler/cache directories to a large filesystem: + +```bash +export TMPDIR=/path/to/large/tmp +export TRITON_CACHE_DIR=/path/to/large/cache/triton +export TORCHINDUCTOR_CACHE_DIR=/path/to/large/cache/inductor +export AMD_COMGR_CACHE_DIR=/path/to/large/cache/comgr +``` + +## Primus-Style Configs + +New examples should use Primus-style override sections and let the diffusion +adapter normalize them into Wan args. See: + +```text +examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml +examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml +examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml +examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml +``` + +Minimal Wan2.1-1.3B shape: + +```yaml +work_group: local +user_name: local +exp_name: wan2.1_t2v_1.3b-pretrain +workspace: ./output + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + model: wan2.1_t2v_1.3b.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + training: + local_batch_size: 1 + steps: 100 + gradient_accumulation_steps: 1 + output_dir: ./output/wan2.1_t2v_1.3b-pretrain + save_steps: 0 + data: + dataset_path: /data/tiny-video-samples/meta.jsonl + data_folder: /data/tiny-video-samples/data + text_tokenizer: /models/Wan2.1-T2V-1.3B/google/umt5-xxl + height: 480 + width: 832 + parallelism: + sp_size: 1 + dp_replicate: 1 + runtime: + attention_backend: flash_attn_aiter + report_to: none +``` + +Minimal Wan2.2-5B uses the same shape with `model: wan2.2_ti2v_5b.yaml` and +the Wan2.2 tokenizer path: + +```yaml +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + model: wan2.2_ti2v_5b.yaml + overrides: + data: + text_tokenizer: /models/Wan2.2-TI2V-5B/google/umt5-xxl +``` diff --git a/primus/backends/diffusion/__init__.py b/primus/backends/diffusion/__init__.py new file mode 100644 index 000000000..c47954e23 --- /dev/null +++ b/primus/backends/diffusion/__init__.py @@ -0,0 +1,12 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Primus-owned diffusion backend components for t2i/t2v training.""" + +from primus.backends.diffusion.diffusion_adapter import DiffusionAdapter +from primus.core.backend.backend_registry import BackendRegistry + +BackendRegistry.register_adapter("diffusion", DiffusionAdapter) diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py new file mode 100644 index 000000000..b6b444ceb --- /dev/null +++ b/primus/backends/diffusion/argument_builder.py @@ -0,0 +1,262 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import copy +from types import SimpleNamespace +from typing import Any + +from primus.core.utils.yaml_utils import nested_namespace_to_dict + + +class WanArgBuilder: + """Build the compact config object consumed by the Wan trainer.""" + + DEFAULT_DATASET: dict[str, Any] = { + "name": "wan", + "config": { + "dataset_type": "vision", + "dataset_format": "jsonl", + "dataset_path": "/path/to/meta.jsonl", + "data_folder": "/path/to/videos", + "video_sampling_strategy": "frame_num", + "frame_num": 81, + "shuffle": True, + "video_backend": "imageio", + "processor_config": { + "processor_name": "wanvideo", + "processor_type": "wanvideo", + "max_text_length": 512, + "text_tokenizer": "/path/to/umt5-xxl", + "extra_kwargs": { + "do_resize": True, + "size": { + "height": 480, + "width": 832, + }, + "do_normalize": True, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], + }, + }, + }, + } + DEFAULT_TRAINER: dict[str, Any] = { + "name": "fsdp2", + "args": { + "output_dir": "./output/wan", + "per_device_train_batch_size": 1, + "per_device_eval_batch_size": 1, + "gradient_accumulation_steps": 1, + "gradient_checkpointing": True, + "attention_backend": "flash_attn_aiter", + "learning_rate": 1.0e-5, + "lr_scheduler_type": "constant", + "warmup_steps": 0, + "weight_decay": 0.01, + "num_train_epochs": 1, + "max_steps": 100, + "logging_steps": 1, + "save_steps": 0, + "dataloader_num_workers": 4, + "report_to": "none", + "run_name": "wan-fsdp2", + "bf16": True, + "seed": 10007, + "optim": "adamw_torch", + "adam_beta1": 0.9, + "adam_beta2": 0.999, + "adam_epsilon": 1.0e-8, + "max_grad_norm": 1.0, + "fsdp2_wrap_target": "dit", + "fsdp_transformer_layer_cls_to_wrap": "DiTBlock", + "fsdp2_reshard_after_forward": True, + "save_strategy": "dit_only", + "sp_size": 1, + "dp_replicate": 1, + "flow_match_scheduler": { + "shift": 5, + "sigma_min": 0.0, + "extra_one_step": True, + "num_train_timesteps": 1000, + }, + }, + } + + def __init__(self) -> None: + self._params: dict[str, Any] = {} + + def update(self, params: Any) -> None: + if isinstance(params, SimpleNamespace): + self._params = nested_namespace_to_dict(params) + elif isinstance(params, dict): + self._params = copy.deepcopy(params) + else: + raise TypeError(f"WanArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}") + + def finalize(self) -> SimpleNamespace: + params = copy.deepcopy(self._params) + if "model" not in params: + raise ValueError("Wan backend config requires a model preset.") + for legacy_section in ("dataset", "trainer"): + if legacy_section in params: + raise ValueError( + f"Wan backend no longer accepts public `{legacy_section}` overrides. " + "Use Primus-style `data`, `training`, `parallelism`, `optimizer`, " + "`runtime`, and `metrics` sections instead." + ) + + params = self._normalize_primus_style_sections(params) + + return SimpleNamespace( + model=params["model"], + dataset=params["dataset"], + trainer=params["trainer"], + stage=params.get("stage", "pretrain"), + primus=params.get("primus", {}), + ) + + @staticmethod + def _set_nested(target: dict[str, Any], path: tuple[str, ...], value: Any) -> None: + cursor = target + for key in path[:-1]: + next_value = cursor.get(key) + if not isinstance(next_value, dict): + next_value = {} + cursor[key] = next_value + cursor = next_value + cursor[path[-1]] = value + + @staticmethod + def _get_any(source: dict[str, Any], *keys: str) -> Any: + for key in keys: + if key in source: + return source[key] + return None + + def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, Any]: + """Translate concise Primus-style Wan sections into trainer arguments. + + Public configs use high-level sections such as `training`, `data`, + `parallelism`, and `runtime`; internal defaults supply the compact + `dataset` and `trainer` objects consumed by the Wan runtime. + """ + + normalized = { + "model": params["model"], + "dataset": copy.deepcopy(self.DEFAULT_DATASET), + "trainer": copy.deepcopy(self.DEFAULT_TRAINER), + "stage": params.get("stage", "pretrain"), + "primus": params.get("primus", {}), + } + + dataset_cfg = normalized["dataset"]["config"] + trainer_args = normalized["trainer"]["args"] + + training = params.get("training") or {} + data = params.get("data") or {} + parallelism = params.get("parallelism") or {} + optimizer = params.get("optimizer") or {} + runtime = params.get("runtime") or {} + metrics = params.get("metrics") or {} + + training_map = { + ("steps",): ("max_steps",), + ("local_batch_size",): ("per_device_train_batch_size",), + ("global_batch_size",): ("global_batch_size",), + ("gradient_accumulation_steps",): ("gradient_accumulation_steps",), + ("output_dir",): ("output_dir",), + ("save_steps",): ("save_steps",), + ("run_name",): ("run_name",), + ("num_train_epochs",): ("num_train_epochs",), + ("dataloader_num_workers",): ("dataloader_num_workers",), + ("resume_from_checkpoint",): ("resume_from_checkpoint",), + } + for source_path, target_path in training_map.items(): + value = self._get_any(training, *source_path) + if value is not None: + self._set_nested(trainer_args, target_path, value) + + data_map = { + ("dataset_path",): ("dataset_path",), + ("data_folder",): ("data_folder",), + ("frame_num",): ("frame_num",), + ("video_backend",): ("video_backend",), + ("text_tokenizer",): ("processor_config", "text_tokenizer"), + ("processor_name",): ("processor_config", "processor_name"), + ("processor_type",): ("processor_config", "processor_type"), + } + for source_path, target_path in data_map.items(): + value = self._get_any(data, *source_path) + if value is not None: + self._set_nested(dataset_cfg, target_path, value) + + height = data.get("height") + width = data.get("width") + if height is not None: + self._set_nested( + dataset_cfg, + ("processor_config", "extra_kwargs", "size", "height"), + height, + ) + if width is not None: + self._set_nested( + dataset_cfg, + ("processor_config", "extra_kwargs", "size", "width"), + width, + ) + + parallelism_map = { + ("sp_size",): ("sp_size",), + ("dp_replicate",): ("dp_replicate",), + } + for source_path, target_path in parallelism_map.items(): + value = self._get_any(parallelism, *source_path) + if value is not None: + self._set_nested(trainer_args, target_path, value) + + optimizer_map = { + ("lr",): ("learning_rate",), + ("learning_rate",): ("learning_rate",), + ("weight_decay",): ("weight_decay",), + ("adam_beta1",): ("adam_beta1",), + ("adam_beta2",): ("adam_beta2",), + ("adam_epsilon",): ("adam_epsilon",), + ("max_grad_norm",): ("max_grad_norm",), + } + for source_path, target_path in optimizer_map.items(): + value = self._get_any(optimizer, *source_path) + if value is not None: + self._set_nested(trainer_args, target_path, value) + + runtime_map = { + ("attention_backend",): ("attention_backend",), + ("report_to",): ("report_to",), + ("seed",): ("seed",), + ("fsdp2_reshard_after_forward",): ("fsdp2_reshard_after_forward",), + } + for source_path, target_path in runtime_map.items(): + value = self._get_any(runtime, *source_path) + if value is not None: + self._set_nested(trainer_args, target_path, value) + + log_freq = metrics.get("log_freq") + if log_freq is not None: + self._set_nested(trainer_args, ("logging_steps",), log_freq) + + enable_wandb = metrics.get("enable_wandb") + if enable_wandb is False and runtime.get("report_to") is None: + self._set_nested(trainer_args, ("report_to",), "none") + elif enable_wandb is True and runtime.get("report_to") is None: + self._set_nested(trainer_args, ("report_to",), "wandb") + + checkpoint = params.get("checkpoint") or {} + resume_from_checkpoint = checkpoint.get("resume_from_checkpoint") + if resume_from_checkpoint is not None: + self._set_nested(trainer_args, ("resume_from_checkpoint",), resume_from_checkpoint) + + return normalized diff --git a/primus/backends/diffusion/attention/__init__.py b/primus/backends/diffusion/attention/__init__.py new file mode 100644 index 000000000..a1fea3747 --- /dev/null +++ b/primus/backends/diffusion/attention/__init__.py @@ -0,0 +1,29 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from .attention import ( + AITER_FLASH_ATTN_AVAILABLE, + FLASH_ATTN_2_AVAILABLE, + FLASH_ATTN_3_AVAILABLE, + attention, + attention_fused, + flash_attention, + get_attention_backend, + set_attention_backend, +) +from .flex import FLEX_ATTENTION_AVAILABLE + +__all__ = [ + "AITER_FLASH_ATTN_AVAILABLE", + "FLASH_ATTN_2_AVAILABLE", + "FLASH_ATTN_3_AVAILABLE", + "FLEX_ATTENTION_AVAILABLE", + "attention", + "attention_fused", + "flash_attention", + "get_attention_backend", + "set_attention_backend", +] diff --git a/primus/backends/diffusion/attention/_flash_common.py b/primus/backends/diffusion/attention/_flash_common.py new file mode 100644 index 000000000..c1bf0f7a4 --- /dev/null +++ b/primus/backends/diffusion/attention/_flash_common.py @@ -0,0 +1,130 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from collections.abc import Callable + +import torch + + +def _unwrap_output(x: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor: + return x[0] if isinstance(x, tuple) else x + + +def run_flash_attention_backend( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + *, + fixed_attention: Callable[..., torch.Tensor | tuple[torch.Tensor, ...]], + varlen_attention: Callable[..., torch.Tensor | tuple[torch.Tensor, ...]], + window_size_adapter: Callable[[tuple[int, ...]], tuple[int, ...]] | None = None, + fixed_extra_kwargs: dict | None = None, + varlen_extra_kwargs: dict | None = None, +): + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == "cuda" and q.size(-1) <= 256 + + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + call_window_size = window_size_adapter(window_size) if window_size_adapter is not None else window_size + fixed_extra_kwargs = fixed_extra_kwargs or {} + varlen_extra_kwargs = varlen_extra_kwargs or {} + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + if q_lens is None and k_lens is None: + qh = half(q) + kh = half(k) + vh = half(v) + qh = qh.to(vh.dtype) + kh = kh.to(vh.dtype) + if q_scale is not None: + qh = qh * q_scale + + x = fixed_attention( + qh, + kh, + vh, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=call_window_size, + deterministic=deterministic, + **fixed_extra_kwargs, + ) + return _unwrap_output(x).type(out_dtype) + + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens, strict=False)])) + + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens, strict=False)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens, strict=False)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + if q_scale is not None: + q = q * q_scale + + cu_seqlens_q = ( + torch.cat([q_lens.new_zeros([1]), q_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True) + ) + cu_seqlens_k = ( + torch.cat([k_lens.new_zeros([1]), k_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True) + ) + max_sq = int(q_lens.max()) + max_sk = int(k_lens.max()) + + x = varlen_attention( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_sq, + max_seqlen_k=max_sk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=call_window_size, + deterministic=deterministic, + **varlen_extra_kwargs, + ) + # Pad output back to (b, lq) shape to match the input padded layout + out = _unwrap_output(x) + if max_sq == lq: + return out.unflatten(0, (b, lq)).type(out_dtype) + # Variable-length: need to scatter back into padded tensor + result = q.new_zeros(b, lq, *out.shape[1:]) + offset = 0 + for i in range(b): + sl = int(q_lens[i]) + result[i, :sl] = out[offset : offset + sl] + offset += sl + return result.type(out_dtype) diff --git a/primus/backends/diffusion/attention/aiter.py b/primus/backends/diffusion/attention/aiter.py new file mode 100644 index 000000000..dfe3e655a --- /dev/null +++ b/primus/backends/diffusion/attention/aiter.py @@ -0,0 +1,71 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import torch + +from ._flash_common import run_flash_attention_backend + +AITER_FLASH_ATTN_AVAILABLE = False +aiter = None + +try: + import aiter as _aiter # type: ignore + + if hasattr(_aiter, "flash_attn_func") and hasattr(_aiter, "flash_attn_varlen_func"): + aiter = _aiter + AITER_FLASH_ATTN_AVAILABLE = True +except Exception: + AITER_FLASH_ATTN_AVAILABLE = False + + +def _normalize_window_size(window_size: tuple[int, ...]) -> tuple[int, int, int]: + if len(window_size) == 2: + left, right = window_size + return (left, right, 0) + if len(window_size) == 3: + left, right, sink = window_size + return (left, right, sink) + raise ValueError(f"window_size must have 2 or 3 items, got {window_size}") + + +def aiter_flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, +): + assert AITER_FLASH_ATTN_AVAILABLE and aiter is not None + + need_lse = torch.is_grad_enabled() + return run_flash_attention_backend( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + fixed_attention=aiter.flash_attn_func, + varlen_attention=aiter.flash_attn_varlen_func, + window_size_adapter=_normalize_window_size, + fixed_extra_kwargs={"return_lse": need_lse}, + varlen_extra_kwargs={"return_lse": need_lse}, + ) diff --git a/primus/backends/diffusion/attention/attention.py b/primus/backends/diffusion/attention/attention.py new file mode 100644 index 000000000..4e67b59b9 --- /dev/null +++ b/primus/backends/diffusion/attention/attention.py @@ -0,0 +1,410 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Unified attention entry points. + +Public API: +- `attention(q, k, v, ...)` : Wan2.2-style, q/k/v shape [B, L, N, D] +- `attention_fused(q, k, v, ...)` : Legacy Wan-style, q/k/v shape [B, S, N*D] +""" + +from __future__ import annotations + +import warnings + +import torch + +from ._flash_common import run_flash_attention_backend +from .aiter import AITER_FLASH_ATTN_AVAILABLE, aiter_flash_attention + +FLASH_ATTN_3_AVAILABLE = False +FLASH_ATTN_2_AVAILABLE = False +flash_attn_interface = None +flash_attn = None + +try: + import flash_attn_interface as _flash_attn_interface # type: ignore + + flash_attn_interface = _flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except Exception: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn as _flash_attn # type: ignore + + flash_attn = _flash_attn + FLASH_ATTN_2_AVAILABLE = True +except Exception: + FLASH_ATTN_2_AVAILABLE = False + + +__all__ = [ + "AITER_FLASH_ATTN_AVAILABLE", + "FLASH_ATTN_2_AVAILABLE", + "FLASH_ATTN_3_AVAILABLE", + "attention", + "attention_fused", + "flash_attention", + "get_attention_backend", + "set_attention_backend", +] + + +# --------------------------------------------------------------------------- +# Global backend config (set once at startup, e.g. from YAML) +# --------------------------------------------------------------------------- +_ATTENTION_BACKEND: str = "auto" + + +_VALID_BACKENDS = ("auto", "sdpa", "flex_attention", "flash_attn2", "flash_attn3", "flash_attn_aiter") + + +def set_attention_backend(backend: str) -> None: + """ + Set global attention backend. + + Accepted values (case-insensitive): + - "auto" | "sdpa" | "flex_attention" | "flash_attn2" | "flash_attn3" | "flash_attn_aiter" + """ + global _ATTENTION_BACKEND + backend = (backend or "").strip().lower() + if backend not in _VALID_BACKENDS: + hint = "" + if backend.startswith("flash_atten"): + hint = " (did you mean 'flash_attn2' / 'flash_attn3'?)" + raise ValueError(f"attention_backend must be one of {_VALID_BACKENDS}, got '{backend}'{hint}") + _ATTENTION_BACKEND = backend + + +def get_attention_backend() -> str: + return _ATTENTION_BACKEND + + +def _resolve_flash_version(device_type: str) -> int | None: + """ + Returns: + - None: do not use flash attention (use SDPA) + - 2 or 3: use flash attention, prefer that version (3 may fall back to 2 in `flash_attention`) + """ + backend = _ATTENTION_BACKEND + if backend in ("sdpa", "flex_attention", "flash_attn_aiter"): + return None + + if device_type != "cuda": + # auto: CPU should use SDPA + if backend == "auto": + return None + # explicit flash backend: error (prevents silent slow CPU path) + raise RuntimeError(f"attention_backend='{backend}' requires CUDA tensors.") + + # CUDA path + if backend == "auto": + if FLASH_ATTN_3_AVAILABLE: + return 3 + if FLASH_ATTN_2_AVAILABLE: + return 2 + return None + + if backend == "flash_attn2": + if not FLASH_ATTN_2_AVAILABLE: + raise RuntimeError("attention_backend='flash_attn2' but flash-attn 2 is not available.") + return 2 + + # "flash_attn3" + if not FLASH_ATTN_3_AVAILABLE: + raise RuntimeError("attention_backend='flash_attn3' but flash-attn 3 is not available.") + return 3 + + +# --------------------------------------------------------------------------- +# Flash Attention implementation (varlen + fixed-length fast path) +# --------------------------------------------------------------------------- +def _flash_attn_3_fixed(q, k, v, **kwargs): + return flash_attn_interface.flash_attn_func(q, k, v) + + +def _flash_attn_3_varlen( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + causal, + window_size, + deterministic, +): + return flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_q=None, + seqused_k=None, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + ) + + +def _flash_attn_2_fixed(q, k, v, dropout_p, softmax_scale, causal, window_size, deterministic): + try: + return flash_attn.flash_attn_func( + q, + k, + v, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + ) + except TypeError: + return flash_attn.flash_attn_func(q, k, v) + + +def _flash_attn_2_varlen( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + dropout_p, + softmax_scale, + causal, + window_size, + deterministic, +): + return flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + ) + + +def flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + version=None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn( + "Flash attention 3 is not available, use flash attention 2 instead.", + stacklevel=2, + ) + + if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p and window_size are not supported in FA3 now. + return run_flash_attention_backend( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + fixed_attention=_flash_attn_3_fixed, + varlen_attention=_flash_attn_3_varlen, + ) + assert FLASH_ATTN_2_AVAILABLE + return run_flash_attention_backend( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + fixed_attention=_flash_attn_2_fixed, + varlen_attention=_flash_attn_2_varlen, + ) + + +def attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + fa_version=None, +): + """ + Unified attention for Wan2.2-style tensors. + + q: [B, Lq, N, D] + k: [B, Lk, N, D] + v: [B, Lk, N, D] + """ + if _ATTENTION_BACKEND == "flash_attn_aiter": + if q.device.type != "cuda": + raise RuntimeError("attention_backend='flash_attn_aiter' requires CUDA tensors.") + if not AITER_FLASH_ATTN_AVAILABLE: + raise RuntimeError("attention_backend='flash_attn_aiter' but aiter is not available.") + return aiter_flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + ) + + resolved = _resolve_flash_version(q.device.type) + if resolved is not None: + version = resolved + # Only "auto" allows call-site override (useful for debugging). + if _ATTENTION_BACKEND == "auto" and fa_version is not None: + version = fa_version + return flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=version, + ) + + # FlexAttention path (PyTorch, requires torch.compile for perf) + if _ATTENTION_BACKEND == "flex_attention": + from .flex import flex_attention_fn + + return flex_attention_fn( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + ) + + # SDPA fallback (also used for CPU). + # Currently assumes *uniform-length* training sequences, so we do NOT support padding masks on the SDPA path yet + if q_lens is not None or k_lens is not None: + warnings.warn("Padding mask is disabled when using scaled_dot_product_attention.", stacklevel=2) + out_dtype = q.dtype + q_ = q.transpose(1, 2).to(dtype) + k_ = k.transpose(1, 2).to(dtype) + v_ = v.transpose(1, 2).to(dtype) + if q_scale is not None: + q_ = q_ * q_scale + out = torch.nn.functional.scaled_dot_product_attention( + q_, k_, v_, attn_mask=None, is_causal=causal, dropout_p=dropout_p + ) + return out.transpose(1, 2).contiguous().to(out_dtype) + + +def attention_fused( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + num_heads: int, + dropout_p: float = 0.0, + softmax_scale=None, + causal: bool = False, +): + """ + Legacy Wan fused-head attention. + + q/k/v: [B, S, num_heads * head_dim] + returns: [B, S, num_heads * head_dim] + """ + b, s, c = q.shape + if c % num_heads != 0: + raise ValueError(f"hidden dim {c} not divisible by num_heads {num_heads}") + d = c // num_heads + + q4d = q.view(b, s, num_heads, d) + k4d = k.view(b, -1, num_heads, d) + v4d = v.view(b, -1, num_heads, d) + + # For fused calls, q_lens/k_lens are not used (fixed length). + out = attention( + q=q4d, + k=k4d, + v=v4d, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + dtype=q.dtype if q.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16, + ) + return out.flatten(2) diff --git a/primus/backends/diffusion/attention/flex.py b/primus/backends/diffusion/attention/flex.py new file mode 100644 index 000000000..c5f0ad091 --- /dev/null +++ b/primus/backends/diffusion/attention/flex.py @@ -0,0 +1,280 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +FlexAttention backend — pure PyTorch, relies on torch.compile for performance. + +This module provides a drop-in replacement for the SDPA/FlashAttention paths +using ``torch.nn.attention.flex_attention``. FlexAttention compiles a +user-defined ``score_mod`` function into a fused kernel via ``torch.compile``, +giving FlashAttention-level performance without third-party dependencies. + +Key capabilities over SDPA: + - Variable-length padding masks via ``create_block_mask`` + - Sliding-window attention via ``score_mod`` + - Composable attention score modifications +""" + +from __future__ import annotations + +import torch + +# Import flex_attention from PyTorch (available since PyTorch 2.5+). +# We keep the import at module level so that ``FLEX_ATTENTION_AVAILABLE`` +# reflects the actual environment at init time. +FLEX_ATTENTION_AVAILABLE = False +_flex_attention_compiled = None +_create_block_mask = None + +try: + from torch.nn.attention.flex_attention import create_block_mask, flex_attention + + # torch.compile is essential for FlexAttention performance — without it, + # score_mod/block_mask run as eager Python callbacks (separate kernel per op). + # We attempt compile lazily on first call; if Triton/Inductor fails (e.g. + # missing system libs on ROCm), we fall back to eager with a warning. + _flex_attention_eager = flex_attention + _flex_attention_compiled = None # lazy-init on first call + _create_block_mask = create_block_mask + FLEX_ATTENTION_AVAILABLE = True +except Exception: + FLEX_ATTENTION_AVAILABLE = False + + +def _get_flex_attention(): + """ + Return compiled flex_attention, falling back to eager if compile fails. + + torch.compile is lazy: wrapping succeeds immediately, but actual Triton + codegen happens on first invocation and can fail (e.g. missing system libs + on ROCm). We handle this by catching runtime errors on first call and + permanently switching to eager mode. + """ + global _flex_attention_compiled + if _flex_attention_compiled is not None: + return _flex_attention_compiled + + import logging + + logger = logging.getLogger(__name__) + try: + compiled = torch.compile(_flex_attention_eager) + logger.info("flex_attention: torch.compile wrapper created (lazy — actual compile on first call)") + _flex_attention_compiled = compiled + except Exception as e: + logger.warning("flex_attention: torch.compile wrapping failed (%s), using eager mode", e) + _flex_attention_compiled = _flex_attention_eager + return _flex_attention_compiled + + +def _flex_attention_with_fallback(q, k, v, **kwargs): + """Call flex_attention with runtime fallback if compiled version fails.""" + global _flex_attention_compiled + fn = _get_flex_attention() + try: + return fn(q, k, v, **kwargs) + except Exception as e: + if fn is _flex_attention_eager: + raise # already in eager mode, real error + import logging + + logger = logging.getLogger(__name__) + logger.warning( + "flex_attention: compiled call failed (%s: %s), falling back to eager mode", + type(e).__name__, + e, + ) + _flex_attention_compiled = _flex_attention_eager + return _flex_attention_eager(q, k, v, **kwargs) + + +def _build_score_mod( + causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + softmax_scale: float | None = None, +): + """ + Compose a ``score_mod`` function for ``flex_attention``. + + The returned function has signature ``(score, b, h, q_idx, kv_idx) -> score`` + and is compiled by ``torch.compile`` into fused operations. + """ + mods: list = [] + + if softmax_scale is not None: + # flex_attention already applies 1/sqrt(d) by default, so we only + # inject a custom scale if the caller wants a non-default value. + def scale_mod(score, b, h, q_idx, kv_idx): + return score * softmax_scale + + mods.append(scale_mod) + + if causal: + + def causal_mod(score, b, h, q_idx, kv_idx): + return torch.where(q_idx >= kv_idx, score, float("-inf")) + + mods.append(causal_mod) + + if window_size != (-1, -1): + left, right = window_size + + def window_mod(score, b, h, q_idx, kv_idx): + in_window = True + if left >= 0: + in_window = in_window & (q_idx - kv_idx <= left) + if right >= 0: + in_window = in_window & (kv_idx - q_idx <= right) + return torch.where(in_window, score, float("-inf")) + + mods.append(window_mod) + + if not mods: + return None + + # Compose all mods into one function. + def composed_score_mod(score, b, h, q_idx, kv_idx): + for mod in mods: + score = mod(score, b, h, q_idx, kv_idx) + return score + + return composed_score_mod + + +def _build_block_mask( + B: int, + N: int, + Lq: int, + Lk: int, + q_lens: torch.Tensor | None, + k_lens: torch.Tensor | None, + device: torch.device, + causal: bool = False, + window_size: tuple[int, int] = (-1, -1), +): + """ + Build a ``BlockMask`` for ``flex_attention`` from variable-length seqs. + + If both ``q_lens`` and ``k_lens`` are None (uniform-length), returns None + so that ``flex_attention`` uses a full dense mask. + """ + if q_lens is None and k_lens is None and not causal and window_size == (-1, -1): + return None + + # Ensure length tensors are on the same device as the mask indices + # (create_block_mask uses vmap which creates index tensors on `device`) + if q_lens is not None: + q_lens = q_lens.to(device) + if k_lens is not None: + k_lens = k_lens.to(device) + + # Build the mask function. Closures capture the lengths. + def mask_fn(b, h, q_idx, kv_idx): + mask = True + if q_lens is not None: + mask = mask & (q_idx < q_lens[b]) + if k_lens is not None: + mask = mask & (kv_idx < k_lens[b]) + if causal: + mask = mask & (q_idx >= kv_idx) + if window_size != (-1, -1): + left, right = window_size + if left >= 0: + mask = mask & (q_idx - kv_idx <= left) + if right >= 0: + mask = mask & (kv_idx - q_idx <= right) + return mask + + return _create_block_mask(mask_fn, B=B, H=N, Q_LEN=Lq, KV_LEN=Lk, device=device) + + +def flex_attention_fn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: torch.Tensor | None = None, + k_lens: torch.Tensor | None = None, + dropout_p: float = 0.0, + softmax_scale: float | None = None, + q_scale: float | None = None, + causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + deterministic: bool = False, + dtype: torch.dtype = torch.bfloat16, + fa_version=None, +) -> torch.Tensor: + """ + Attention using ``torch.nn.attention.flex_attention``. + + Drop-in replacement for ``attention()`` — same signature, same tensor + layout (q/k/v: ``[B, L, N, D]``). + + Args: + q: Query tensor ``[B, Lq, N, D]`` + k: Key tensor ``[B, Lk, N, D]`` + v: Value tensor ``[B, Lk, N, D]`` + q_lens: Optional valid lengths per batch for queries ``[B]`` + k_lens: Optional valid lengths per batch for keys ``[B]`` + dropout_p: Dropout probability (passed through) + softmax_scale: Custom softmax scale (None = default 1/sqrt(D)) + q_scale: Pre-multiply q by this scalar + causal: Whether to apply causal mask + window_size: (left, right) sliding window; (-1, -1) = disabled + deterministic: Ignored (for API compat with flash_attention) + dtype: Target dtype for computation + fa_version: Ignored (for API compat) + """ + if not FLEX_ATTENTION_AVAILABLE: + raise RuntimeError( + "flex_attention backend requested but torch.nn.attention.flex_attention " + "is not available. Requires PyTorch >= 2.5." + ) + + out_dtype = q.dtype + B, Lq, N, D = q.shape + Lk = k.shape[1] + + # Cast to compute dtype + half_dtypes = (torch.float16, torch.bfloat16) + q_ = q.to(dtype) if q.dtype not in half_dtypes else q + k_ = k.to(dtype) if k.dtype not in half_dtypes else k + v_ = v.to(dtype) if v.dtype not in half_dtypes else v + q_ = q_.to(v_.dtype) + k_ = k_.to(v_.dtype) + + if q_scale is not None: + q_ = q_ * q_scale + + # flex_attention expects [B, N, L, D] (heads-first) + q_ = q_.transpose(1, 2) # [B, N, Lq, D] + k_ = k_.transpose(1, 2) # [B, N, Lk, D] + v_ = v_.transpose(1, 2) # [B, N, Lk, D] + + # Build block mask (handles padding, causal, window in the mask itself) + block_mask = _build_block_mask( + B=B, + N=N, + Lq=Lq, + Lk=Lk, + q_lens=q_lens, + k_lens=k_lens, + device=q.device, + causal=causal, + window_size=window_size, + ) + + # Build score_mod only for softmax_scale override + # (causal and window are handled in block_mask for better sparsity) + score_mod = None + if softmax_scale is not None: + score_mod = _build_score_mod(softmax_scale=softmax_scale) + + # Call flex_attention (with runtime compile-error fallback) + out = _flex_attention_with_fallback(q_, k_, v_, score_mod=score_mod, block_mask=block_mask) + + # Back to [B, L, N, D] + out = out.transpose(1, 2).contiguous() + return out.to(out_dtype) diff --git a/primus/backends/diffusion/data/__init__.py b/primus/backends/diffusion/data/__init__.py new file mode 100644 index 000000000..18ae84901 --- /dev/null +++ b/primus/backends/diffusion/data/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Data builders for the diffusion backend.""" diff --git a/primus/backends/diffusion/data/registrations/__init__.py b/primus/backends/diffusion/data/registrations/__init__.py new file mode 100644 index 000000000..b5a3705a2 --- /dev/null +++ b/primus/backends/diffusion/data/registrations/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Registered diffusion dataset builders.""" diff --git a/primus/backends/diffusion/data/registrations/wan.py b/primus/backends/diffusion/data/registrations/wan.py new file mode 100644 index 000000000..76492cb2b --- /dev/null +++ b/primus/backends/diffusion/data/registrations/wan.py @@ -0,0 +1,176 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image +from torch.utils.data import Dataset +from torchvision.transforms import functional as F +from transformers import AutoTokenizer + +from primus.backends.diffusion.utils.vision_process import fetch_video, strip_file_uri + + +class WanVideoProcessor: + """Tokenize prompts and normalize video tensors for Wan training.""" + + def __init__(self, config: dict[str, Any]): + self.config = config + self.max_text_length = int(config.get("max_text_length", 512)) + tokenizer_path = config.get("text_tokenizer") + if not tokenizer_path: + raise ValueError("Wan dataset processor requires `text_tokenizer`.") + trust_remote_code = bool(config.get("trust_remote_code", False)) + self.tokenizer = AutoTokenizer.from_pretrained( + tokenizer_path, + trust_remote_code=trust_remote_code, + ) + + extra_kwargs = config.get("extra_kwargs", {}) or {} + size = extra_kwargs.get("size", {}) or {} + self.height = int(size.get("height", 480)) + self.width = int(size.get("width", 832)) + self.image_mean = torch.tensor(extra_kwargs.get("image_mean", [0.5, 0.5, 0.5])).view(3, 1, 1, 1) + self.image_std = torch.tensor(extra_kwargs.get("image_std", [0.5, 0.5, 0.5])).view(3, 1, 1, 1) + + def tokenize(self, prompt: str) -> dict[str, torch.Tensor]: + encoded = self.tokenizer( + prompt, + max_length=self.max_text_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ) + return { + "input_ids": encoded["input_ids"].squeeze(0).long(), + "attention_mask": encoded["attention_mask"].squeeze(0).long(), + } + + def normalize_video(self, video_tchw: torch.Tensor) -> torch.Tensor: + if video_tchw.ndim != 4: + raise ValueError(f"Expected video tensor [T,C,H,W], got shape={tuple(video_tchw.shape)}") + video_tchw = video_tchw[:, :3].float() + if video_tchw.max() > 2: + video_tchw = video_tchw / 255.0 + video_tchw = F.resize(video_tchw, [self.height, self.width], antialias=True) + video_cthw = video_tchw.permute(1, 0, 2, 3).contiguous() + return (video_cthw - self.image_mean) / self.image_std + + def prepare_batch( + self, *, batch: dict[str, Any], device: torch.device, dtype: torch.dtype + ) -> dict[str, Any]: + return batch + + +class WanVideoDataset(Dataset): + def __init__(self, config: dict[str, Any], processor: WanVideoProcessor): + self.config = config + self.processor = processor + self.dataset_path = Path(config.get("dataset_path", "")) + self.data_folder = Path(config.get("data_folder", "")) + self.frame_num = int(config.get("frame_num", 81)) + self.video_backend = str(config.get("video_backend", "imageio")).lower() + + if not self.dataset_path.exists(): + raise FileNotFoundError(f"Wan dataset metadata not found: {self.dataset_path}") + with self.dataset_path.open(encoding="utf-8") as f: + self.samples = [json.loads(line) for line in f if line.strip()] + if not self.samples: + raise ValueError(f"Wan dataset metadata is empty: {self.dataset_path}") + + def __len__(self) -> int: + return len(self.samples) + + def _resolve_video_path(self, value: str) -> str: + if value.startswith(("http://", "https://", "file://")): + return value + path = Path(value) + if not path.is_absolute() and self.data_folder: + path = self.data_folder / path + return str(path) + + def _sample_video(self, video_tchw: torch.Tensor) -> torch.Tensor: + total_frames = int(video_tchw.shape[0]) + if total_frames <= 0: + raise ValueError("Video contains no frames.") + idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long() + return video_tchw[idx] + + def _read_video_imageio(self, path: str) -> torch.Tensor: + import imageio.v3 as iio + + path = strip_file_uri(path) + frames = [] + for frame in iio.imiter(path): + image = Image.fromarray(frame).convert("RGB") + frames.append(torch.from_numpy(np.asarray(image).copy())) + if not frames: + raise ValueError(f"Video contains no frames: {path}") + return torch.stack(frames).permute(0, 3, 1, 2) + + def _read_video_decord(self, path: str) -> torch.Tensor: + import decord + + path = strip_file_uri(path) + vr = decord.VideoReader(path) + total_frames = len(vr) + idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long().tolist() + return torch.from_numpy(vr.get_batch(idx).asnumpy()).permute(0, 3, 1, 2) + + def _read_video(self, path: str) -> torch.Tensor: + if self.video_backend == "imageio": + video = self._sample_video(self._read_video_imageio(path)) + elif self.video_backend == "decord": + video = self._read_video_decord(path) + else: + video = fetch_video( + { + "video": path, + "nframes": self.frame_num, + "resized_height": self.processor.height, + "resized_width": self.processor.width, + } + ) + return self.processor.normalize_video(video) + + def __getitem__(self, index: int) -> dict[str, Any]: + sample = self.samples[index] + prompt = sample.get("prompt") or sample.get("text") or sample.get("caption") or "" + video_key = sample.get("video") or sample.get("video_path") + if not video_key: + raise KeyError(f"Wan dataset sample missing `video`: index={index}") + + item = self.processor.tokenize(str(prompt)) + item["video"] = self._read_video(self._resolve_video_path(str(video_key))) + if "seed" in sample: + item["seed"] = int(sample["seed"]) + return item + + @staticmethod + def get_collator(): + def collate(samples: list[dict[str, Any]]) -> dict[str, Any]: + batch = { + "video": torch.stack([sample["video"] for sample in samples]), + "input_ids": torch.stack([sample["input_ids"] for sample in samples]), + "attention_mask": torch.stack([sample["attention_mask"] for sample in samples]), + } + if any("seed" in sample for sample in samples): + batch["seed"] = torch.tensor([sample.get("seed", 0) for sample in samples], dtype=torch.long) + return batch + + return collate + + +def build_wan_dataset(config: dict[str, Any]): + processor = WanVideoProcessor(config.get("processor_config", {}) or {}) + dataset = WanVideoDataset(config, processor) + return dataset, processor diff --git a/primus/backends/diffusion/diffusion_adapter.py b/primus/backends/diffusion/diffusion_adapter.py new file mode 100644 index 000000000..202c3914a --- /dev/null +++ b/primus/backends/diffusion/diffusion_adapter.py @@ -0,0 +1,67 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from primus.backends.diffusion.argument_builder import WanArgBuilder +from primus.core.backend.backend_adapter import BackendAdapter +from primus.modules.module_utils import log_rank_0 + + +class DiffusionAdapter(BackendAdapter): + """Primus adapter for PyTorch diffusion training.""" + + def __init__(self, framework: str = "diffusion"): + super().__init__(framework) + + def setup_backend_path(self, backend_path=None) -> str: + """Validate the Primus-owned in-tree diffusion package.""" + if backend_path: + raise ValueError( + "The diffusion backend is built into Primus and does not support " + f"external backend_path overrides. Got: {backend_path}" + ) + + resolved = Path(__file__).resolve().parent + if not resolved.exists(): + raise FileNotFoundError(f"[Primus:Diffusion] backend package does not exist: {resolved}") + + resolved_str = str(resolved) + try: + log_rank_0(f"[Primus:Diffusion] using in-tree backend package -> {resolved_str}") + except Exception: + # Best-effort startup logging should not block backend setup. + pass + + return resolved_str + + def convert_config(self, params: Any): + builder = WanArgBuilder() + builder.update(params) + wan_args = builder.finalize() + # convert_config is also called by the standalone prepare hook, where the + # Primus logger may not be initialized yet; guard the informational log. + try: + log_rank_0("[Primus:DiffusionAdapter] Converted Primus module params -> Wan args") + except Exception: + # Standalone prepare hooks may run before the Primus logger is bound. + pass + return wan_args + + def load_trainer_class(self, stage: str = "pretrain"): + if stage in ("pretrain", "posttrain", "sft"): + from primus.backends.diffusion.diffusion_pretrain_trainer import ( + DiffusionPretrainTrainer, + ) + + return DiffusionPretrainTrainer + raise ValueError(f"Invalid stage for Diffusion backend: {stage}") + + def detect_backend_version(self) -> str: + return "in-tree" diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py new file mode 100644 index 000000000..ddbb70d49 --- /dev/null +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -0,0 +1,131 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import importlib.util +from typing import Any + +from primus.core.trainer.base_trainer import BaseTrainer +from primus.core.utils.yaml_utils import nested_namespace_to_dict +from primus.modules.module_utils import log_rank_0 + + +class DiffusionPretrainTrainer(BaseTrainer): + """Primus lifecycle wrapper for Wan diffusion training.""" + + def __init__(self, backend_args: Any): + super().__init__(backend_args=backend_args) + self.wan_trainer = None + + @staticmethod + def _as_dict(value: Any) -> dict: + if isinstance(value, dict): + return value + return nested_namespace_to_dict(value) + + def setup(self): + trainer_cfg = self._as_dict(self.backend_args.trainer) + dataset_cfg = self._as_dict(self.backend_args.dataset) + trainer_args = trainer_cfg.get("args", {}) + attention_backend = trainer_args.get("attention_backend") + + missing = [ + package + for package in ( + "torch", + "loguru", + "safetensors", + "transformers", + "PIL", + "torchvision", + "requests", + "packaging", + ) + if importlib.util.find_spec(package) is None + ] + video_backend = (dataset_cfg.get("config", {}) or {}).get("video_backend") + if video_backend == "imageio" and importlib.util.find_spec("imageio") is None: + missing.append("imageio") + if video_backend == "decord" and importlib.util.find_spec("decord") is None: + missing.append("decord") + if missing: + raise RuntimeError( + "Diffusion backend missing required Python packages: " + f"{', '.join(missing)}. Install the Wan diffusion training extras first." + ) + + if attention_backend: + from primus.backends.diffusion.attention import set_attention_backend + + set_attention_backend(attention_backend) + log_rank_0(f"[Primus:Diffusion] attention_backend={attention_backend}") + + if attention_backend == "flash_attn_aiter": + from primus.backends.diffusion.attention.aiter import ( + AITER_FLASH_ATTN_AVAILABLE, + ) + + if not AITER_FLASH_ATTN_AVAILABLE: + raise RuntimeError( + "attention_backend=flash_attn_aiter was requested, but AITER flash attention " + "is unavailable in this environment." + ) + + def init(self): + from primus.backends.diffusion.registry import ( + get_dataset_builder, + get_model_builder, + get_trainer_builder, + ) + + model_cfg = self._as_dict(self.backend_args.model) + dataset_cfg = self._as_dict(self.backend_args.dataset) + trainer_cfg = self._as_dict(self.backend_args.trainer) + + model_name = model_cfg["name"] + dataset_name = dataset_cfg["name"] + trainer_name = trainer_cfg["name"] + + model_config = model_cfg["config"] + dataset_config = dataset_cfg["config"] + trainer_args = trainer_cfg["args"] + + log_rank_0( + f"[Primus:Diffusion] Building model={model_name}, dataset={dataset_name}, trainer={trainer_name}" + ) + model = get_model_builder(model_name)(model_config) + dataset, processor = get_dataset_builder(dataset_name)(dataset_config) + self.wan_trainer = get_trainer_builder(trainer_name)( + model=model, + dataset=dataset, + processor=processor, + trainer_args=trainer_args, + ) + + def train(self): + if self.wan_trainer is None: + raise RuntimeError("DiffusionPretrainTrainer.init() must be called before train().") + + self.wan_trainer.train() + self.wan_trainer.save_model() + + def cleanup(self, on_error: bool = False): + try: + import wandb + + if getattr(wandb, "run", None) is not None: + wandb.finish(exit_code=1 if on_error else 0) + except Exception as exc: + log_rank_0(f"[Primus:Diffusion] wandb cleanup failed: {exc}") + + try: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + except Exception as exc: + log_rank_0(f"[Primus:Diffusion] distributed cleanup failed: {exc}") diff --git a/primus/backends/diffusion/distributed/__init__.py b/primus/backends/diffusion/distributed/__init__.py new file mode 100644 index 000000000..7fea79b62 --- /dev/null +++ b/primus/backends/diffusion/distributed/__init__.py @@ -0,0 +1,32 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Primus diffusion distributed training utilities. + +Modules: + mesh - Device mesh creation, distributed setup + checkpoint - DTCP sharded checkpoint save/load + ulysses - Ulysses Sequence Parallel primitives +""" + +from .checkpoint import load_checkpoint_dtcp, save_checkpoint_dtcp +from .mesh import create_device_mesh, setup_distributed +from .ulysses import distributed_attention, sp_gather, sp_split, sp_unpad + +__all__ = [ + # mesh + "setup_distributed", + "create_device_mesh", + # checkpoint + "save_checkpoint_dtcp", + "load_checkpoint_dtcp", + # ulysses + "distributed_attention", + "sp_split", + "sp_gather", + "sp_unpad", +] diff --git a/primus/backends/diffusion/distributed/checkpoint.py b/primus/backends/diffusion/distributed/checkpoint.py new file mode 100644 index 000000000..5ca2f32ae --- /dev/null +++ b/primus/backends/diffusion/distributed/checkpoint.py @@ -0,0 +1,114 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Distributed Tensor Checkpointing (DTCP) save / load utilities. + +Handles FSDP2 sharded checkpoints without gathering to rank 0. +""" + +from __future__ import annotations + +import json +import os +from typing import Any, Dict, Optional + +import torch +import torch.distributed as dist +from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter, load, save +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + set_model_state_dict, + set_optimizer_state_dict, +) + +from primus.backends.diffusion.utils.log import logger + +from .mesh import _ensure_process_group + +_META_FILENAME = "meta.json" + + +def save_checkpoint_dtcp( + model: torch.nn.Module, + optimizer: Optional[torch.optim.Optimizer], + path: str, + epoch: int, + step: int, + additional_data: Dict[str, Any] = None, + *, + model_state_options: Optional[StateDictOptions] = None, + optim_state_options: Optional[StateDictOptions] = None, +): + """Save checkpoint using DTCP (sharded, no gather to rank 0).""" + os.makedirs(path, exist_ok=True) + _ensure_process_group(backend="nccl" if torch.cuda.is_available() else "gloo") + + model_state_options = model_state_options or StateDictOptions(full_state_dict=False) + optim_state_options = optim_state_options or StateDictOptions(full_state_dict=False) + + model_state = get_model_state_dict(model, options=model_state_options) + optim_state = None + if optimizer is not None: + optim_state = get_optimizer_state_dict(model, optimizer, options=optim_state_options) + + meta = {"epoch": epoch, "step": step, **(additional_data or {})} + state_dict = { + "model": model_state, + **({"optimizer": optim_state} if optim_state is not None else {}), + "meta": meta, + } + + save(state_dict, FileSystemWriter(path)) + if (not dist.is_initialized()) or dist.get_rank() == 0: + with open(os.path.join(path, _META_FILENAME), "w", encoding="utf-8") as f: + json.dump(meta, f, indent=2, sort_keys=True) + logger.info(f"Saved DTCP checkpoint to {path}") + + +def load_checkpoint_dtcp( + model: torch.nn.Module, + optimizer: Optional[torch.optim.Optimizer], + path: str, + *, + model_state_options: Optional[StateDictOptions] = None, + optim_state_options: Optional[StateDictOptions] = None, +) -> Dict[str, Any]: + """Load checkpoint using DTCP. Updates model/optimizer in-place. Returns metadata.""" + if not os.path.exists(path): + raise FileNotFoundError(f"Checkpoint not found at {path}") + + _ensure_process_group(backend="nccl" if torch.cuda.is_available() else "gloo") + + model_state_options = model_state_options or StateDictOptions(full_state_dict=False) + optim_state_options = optim_state_options or StateDictOptions(full_state_dict=False) + + model_state = get_model_state_dict(model, options=model_state_options) + optim_state = None + if optimizer is not None: + optim_state = get_optimizer_state_dict(model, optimizer, options=optim_state_options) + + state_dict = { + "model": model_state, + **({"optimizer": optim_state} if optim_state is not None else {}), + "meta": {}, + } + + load(state_dict, FileSystemReader(path)) + + set_model_state_dict(model, model_state, options=model_state_options) + if optimizer is not None and optim_state is not None: + set_optimizer_state_dict(model, optimizer, optim_state, options=optim_state_options) + + meta = state_dict.get("meta", {}) + if not meta: + meta_path = os.path.join(path, _META_FILENAME) + if os.path.exists(meta_path): + with open(meta_path, encoding="utf-8") as f: + meta = json.load(f) + return meta diff --git a/primus/backends/diffusion/distributed/mesh.py b/primus/backends/diffusion/distributed/mesh.py new file mode 100644 index 000000000..b9ee3b7b4 --- /dev/null +++ b/primus/backends/diffusion/distributed/mesh.py @@ -0,0 +1,126 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Device mesh creation and distributed setup. + +Supports: + - Pure FSDP2 (1D mesh: dp_shard) + - FSDP2 + HSDP (2D mesh: dp_replicate × dp_shard) + - FSDP2 + Ulysses SP (mesh: dp_shard × ulysses, flattened into dp_shard_sp) + - FSDP2 + HSDP + Ulysses SP (mesh: dp_replicate × dp_shard × ulysses) + +SP ranks are included in the FSDP sharding mesh so that parameters are +sharded across both DP and SP ranks — this reduces per-rank memory. +""" + +from __future__ import annotations + +import os +from datetime import timedelta +from typing import Optional + +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh + +from primus.backends.diffusion.utils.log import logger + + +def _ensure_process_group(*, backend: str) -> None: + """Best-effort init_process_group (lazy, avoids single-GPU hangs).""" + if dist.is_initialized(): + return + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "12345") + dist.init_process_group(backend, timeout=timedelta(minutes=60)) + + +def setup_distributed() -> tuple[int, int, int]: + """ + Initialize distributed process group. + + Returns: (rank, world_size, local_rank) + """ + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + + if torch.cuda.is_available(): + torch.cuda.set_device(local_rank) + + if world_size > 1 and not dist.is_initialized(): + _ensure_process_group(backend="nccl") + + return rank, world_size, local_rank + + +def create_device_mesh( + world_size: int, + sp_size: int = 1, + dp_replicate: int = 1, +) -> Optional[DeviceMesh]: + """ + Create a DeviceMesh for FSDP2 + optional Ulysses Sequence Parallel. + + Mesh layout (innermost → outermost): + [dp_replicate?] × [dp_shard] × [ulysses?] + + Flattened sub-meshes created automatically: + - "dp_shard_sp": dp_shard × ulysses (used for FSDP2 fully_shard) + - "dp": dp_replicate × dp_shard (used for DistributedSampler) + + Args: + world_size: total number of ranks + sp_size: Ulysses sequence parallel size (must divide world_size) + dp_replicate: HSDP replicate dimension (1 = no HSDP) + + Returns: + DeviceMesh, or None if world_size <= 1 + """ + if world_size <= 1: + return None + + _ensure_process_group(backend="nccl") + + dp_shard = world_size // (sp_size * dp_replicate) + if dp_shard * sp_size * dp_replicate != world_size: + raise ValueError( + f"world_size={world_size} is not divisible by " f"sp_size={sp_size} * dp_replicate={dp_replicate}" + ) + + # Build mesh dimensions + dims = [] + names = [] + if dp_replicate > 1: + dims.append(dp_replicate) + names.append("dp_replicate") + dims.append(dp_shard) + names.append("dp_shard") + if sp_size > 1: + dims.append(sp_size) + names.append("ulysses") + + mesh = init_device_mesh("cuda", tuple(dims), mesh_dim_names=tuple(names)) + + # Flatten composite sub-meshes for convenient access + if sp_size > 1: + mesh["dp_shard", "ulysses"]._flatten("dp_shard_sp") + + if dp_replicate > 1: + mesh["dp_replicate", "dp_shard"]._flatten("dp") + + rank = dist.get_rank() if dist.is_initialized() else 0 + if rank == 0: + logger.info( + f"DeviceMesh created: dims={dict(zip(names, dims))}, " + f"sp_size={sp_size}, dp_replicate={dp_replicate}" + ) + + return mesh diff --git a/primus/backends/diffusion/distributed/ulysses.py b/primus/backends/diffusion/distributed/ulysses.py new file mode 100644 index 000000000..b61ae80e2 --- /dev/null +++ b/primus/backends/diffusion/distributed/ulysses.py @@ -0,0 +1,211 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Ulysses Sequence Parallel primitives. + +Follows the Wan2.2 official design: a self-contained ``distributed_attention`` +function that wraps all-to-all + attention + all-to-all, plus ``sp_split`` / +``sp_gather`` / ``sp_unpad`` for slicing model inputs and gathering outputs. + +The model code only needs to: + 1. ``sp_split`` inputs before the blocks + 2. call ``distributed_attention`` instead of ``attention`` in self-attn + 3. ``sp_gather`` + ``sp_unpad`` the output after the head + +Gradient scaling (÷ sp_size in sp_slice backward, × sp_size in sp_gather +backward) ensures correctness with FSDP2 gradient averaging across the +combined DP+SP mesh. +""" + +from __future__ import annotations + +from typing import Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +# --------------------------------------------------------------------------- +# Core autograd primitives +# --------------------------------------------------------------------------- + + +def _require_divisible_dim(x: torch.Tensor, dim: int, divisor: int, op_name: str) -> None: + size = x.shape[dim] + if size % divisor != 0: + raise ValueError( + f"{op_name} requires tensor dimension {dim} (size={size}) to be divisible by " + f"sequence parallel size {divisor}. Got shape={tuple(x.shape)}." + ) + + +class _SeqAllToAll(torch.autograd.Function): + """All-to-all with autograd. Backward is the inverse (swap dims).""" + + @staticmethod + def forward(ctx, group, x, scatter_dim, gather_dim): + ctx.group = group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + sp_size = dist.get_world_size(group) + _require_divisible_dim(x, scatter_dim, sp_size, "_SeqAllToAll") + input_list = [t.contiguous() for t in x.tensor_split(sp_size, scatter_dim)] + output_list = [torch.empty_like(input_list[0]) for _ in range(sp_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + @staticmethod + def backward(ctx, grad_output): + return ( + None, + _SeqAllToAll.apply(ctx.group, grad_output, ctx.gather_dim, ctx.scatter_dim), + None, + None, + ) + + +class _SliceWithGather(torch.autograd.Function): + """Forward: slice. Backward: all-gather ÷ sp_size (FSDP2 compat).""" + + @staticmethod + def forward(ctx, x, dim, group): + ctx.dim = dim + ctx.group = group + sp_size = dist.get_world_size(group) + sp_rank = dist.get_rank(group) + ctx.sp_size = sp_size + _require_divisible_dim(x, dim, sp_size, "_SliceWithGather") + chunk_size = x.shape[dim] // sp_size + return x.narrow(dim, sp_rank * chunk_size, chunk_size).contiguous() + + @staticmethod + def backward(ctx, grad_output): + gathered = [torch.empty_like(grad_output) for _ in range(ctx.sp_size)] + dist.all_gather(gathered, grad_output.contiguous(), group=ctx.group) + return torch.cat(gathered, dim=ctx.dim) / ctx.sp_size, None, None + + +class _GatherWithSlice(torch.autograd.Function): + """Forward: all-gather. Backward: slice × sp_size (FSDP2 compat).""" + + @staticmethod + def forward(ctx, x, dim, group): + ctx.dim = dim + ctx.group = group + sp_size = dist.get_world_size(group) + sp_rank = dist.get_rank(group) + ctx.sp_size = sp_size + ctx.sp_rank = sp_rank + ctx.chunk_size = x.shape[dim] + gathered = [torch.empty_like(x) for _ in range(sp_size)] + dist.all_gather(gathered, x.contiguous(), group=group) + return torch.cat(gathered, dim=dim) + + @staticmethod + def backward(ctx, grad_output): + chunk = grad_output.narrow(ctx.dim, ctx.sp_rank * ctx.chunk_size, ctx.chunk_size).contiguous() + return chunk * ctx.sp_size, None, None + + +# --------------------------------------------------------------------------- +# Public API — high-level +# --------------------------------------------------------------------------- + + +def distributed_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + group: dist.ProcessGroup, + attention_fn: Callable, + **attention_kwargs, +) -> torch.Tensor: + """ + Ulysses distributed attention (DeepSpeed Ulysses, arXiv:2309.14509). + + Wraps ``attention_fn`` with all-to-all communication so that each rank + computes attention on the **full sequence** with a **subset of heads**. + + Input / output shapes: ``[B, S/P, H, D]`` (sharded-seq, full-heads). + + Args: + q, k, v: query / key / value ``[B, S/P, H, D]`` + group: Ulysses SP process group + attention_fn: e.g. ``primus.backends.diffusion.attention.attention`` + **attention_kwargs: forwarded to ``attention_fn(q=, k=, v=, ...)`` + """ + # scatter heads, gather seq: [B, S/P, H, D] → [B, S, H/P, D] + q = _SeqAllToAll.apply(group, q, 2, 1) + k = _SeqAllToAll.apply(group, k, 2, 1) + v = _SeqAllToAll.apply(group, v, 2, 1) + # standard attention on full sequence with partial heads + out = attention_fn(q=q, k=k, v=v, **attention_kwargs) + # scatter seq, gather heads: [B, S, H/P, D] → [B, S/P, H, D] + return _SeqAllToAll.apply(group, out, 1, 2) + + +def sp_split( + tensors: List[torch.Tensor], + dim: int, + group: dist.ProcessGroup, +) -> Tuple[List[torch.Tensor], int]: + """ + Pad + slice multiple tensors for sequence parallelism. + + Returns ``(sliced_tensors, original_size)`` where *original_size* is + the size along *dim* before padding (needed by :func:`sp_unpad`). + """ + sp_size = dist.get_world_size(group) + original_size: Optional[int] = None + results: List[torch.Tensor] = [] + for t in tensors: + t, orig = _sp_pad(t, dim, sp_size) + if original_size is None: + original_size = orig + results.append(_SliceWithGather.apply(t, dim, group)) + assert original_size is not None + return results, original_size + + +# --------------------------------------------------------------------------- +# Public API — low-level (used directly for gathering the output) +# --------------------------------------------------------------------------- + + +def sp_slice(x: torch.Tensor, dim: int, group: dist.ProcessGroup) -> torch.Tensor: + """Slice ``x`` along *dim* for this SP rank (with autograd). Low-level.""" + return _SliceWithGather.apply(x, dim, group) + + +def sp_gather(x: torch.Tensor, dim: int, group: dist.ProcessGroup) -> torch.Tensor: + """All-gather ``x`` along *dim* from all SP ranks (with autograd).""" + return _GatherWithSlice.apply(x, dim, group) + + +def sp_unpad(x: torch.Tensor, dim: int, original_size: int) -> torch.Tensor: + """Remove padding added by :func:`sp_split`.""" + if x.shape[dim] == original_size: + return x + return x.narrow(dim, 0, original_size) + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _sp_pad(x: torch.Tensor, dim: int, sp_size: int) -> Tuple[torch.Tensor, int]: + """Pad *x* along *dim* so its size is divisible by *sp_size*.""" + original_size = x.shape[dim] + remainder = original_size % sp_size + if remainder == 0: + return x, original_size + pad_amount = sp_size - remainder + pad_config = [0] * (2 * x.dim()) + pos = 2 * (x.dim() - 1 - dim) + pad_config[pos + 1] = pad_amount + return F.pad(x, pad_config), original_size diff --git a/primus/backends/diffusion/models/__init__.py b/primus/backends/diffusion/models/__init__.py new file mode 100644 index 000000000..2798dc3fb --- /dev/null +++ b/primus/backends/diffusion/models/__init__.py @@ -0,0 +1,13 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Wan model exports for the Primus Wan backend.""" + +from .wan import WanForTraining + +__all__ = [ + "WanForTraining", +] diff --git a/primus/backends/diffusion/models/interface.py b/primus/backends/diffusion/models/interface.py new file mode 100644 index 000000000..1355c3658 --- /dev/null +++ b/primus/backends/diffusion/models/interface.py @@ -0,0 +1,48 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from abc import ABC, abstractmethod +from typing import Any, Dict + +import torch +import torch.nn as nn + + +class GenAIModel(nn.Module, ABC): + """ + Unified interface for Generative Models (DiT, AR, Hybrid). + Wraps the Backbone (DiT/Transformer), VAE, and TextEncoder. + """ + + @abstractmethod + def forward_train(self, batch: Dict[str, Any], scheduler: Any = None) -> Dict[str, torch.Tensor]: + """ + The single entry point for training. + + Responsibilities: + 1. Process raw batch (Tokenization, VAE Encoding if needed). + 2. Apply Training Recipe (e.g., Add Noise for DiT, Shift Tokens for AR). + 3. Forward pass through Backbone. + 4. Calculate Loss. + + Args: + batch: Raw batch from data loader. + scheduler: Optional scheduler (e.g., FlowMatchScheduler) passed from Trainer. + + Returns: + { + "loss": torch.Tensor, # Main optimization objective + "log_metrics": Dict, # Metrics for WandB (MSE, Accuracy, etc.) + } + """ + + @abstractmethod + def forward_inference(self, batch: Dict[str, Any], **kwargs): + """ + Entry point for validation/inference. + For DiT: Runs the diffusion sampling loop. + For AR: Runs autoregressive generation. + """ diff --git a/primus/backends/diffusion/models/registrations/__init__.py b/primus/backends/diffusion/models/registrations/__init__.py new file mode 100644 index 000000000..98d014572 --- /dev/null +++ b/primus/backends/diffusion/models/registrations/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Model registrations.""" diff --git a/primus/backends/diffusion/models/registrations/wan.py b/primus/backends/diffusion/models/registrations/wan.py new file mode 100644 index 000000000..491284ff7 --- /dev/null +++ b/primus/backends/diffusion/models/registrations/wan.py @@ -0,0 +1,226 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Register the Wan model builder.""" + +from __future__ import annotations + +import glob +import os +from typing import Any + +import torch +from safetensors.torch import load_file as safe_load_file + +from primus.backends.diffusion.models.wan.adapter import WanForTraining +from primus.backends.diffusion.models.wan.components import WanComponents +from primus.backends.diffusion.models.wan.configuration_wanvideo import WanVideoConfig +from primus.backends.diffusion.models.wan.t5 import umt5_xxl_encoder_from_checkpoint +from primus.backends.diffusion.models.wan.train_pipeline import ( + WanFlowMatchTrainPipeline, +) +from primus.backends.diffusion.models.wan.vae2_1 import Wan2_1_VAE +from primus.backends.diffusion.models.wan.vae2_2 import Wan2_2_VAE +from primus.backends.diffusion.models.wan.wan_dit import WanModel as WanDiT +from primus.backends.diffusion.utils.log import logger +from primus.backends.diffusion.utils.train_utils import count_parameters + + +def _strip_module_prefix( + state_dict: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + for k, v in state_dict.items(): + if k.startswith("module."): + out[k[len("module.") :]] = v + else: + out[k] = v + return out + + +def _load_state_dict(path: str) -> dict[str, torch.Tensor]: + if path.endswith(".safetensors"): + return dict(safe_load_file(path)) + obj = torch.load(path, map_location="cpu") + if isinstance(obj, dict) and "model" in obj and isinstance(obj["model"], dict): + obj = obj["model"] + if not isinstance(obj, dict): + raise ValueError(f"Unsupported checkpoint format at {path}") + return obj + + +def _load_dit_weights_into_module(dit: torch.nn.Module, pretrained_path: str): + """ + Support two common cases: + 1) Wan training export: `dit_model.safetensors` with keys like `blocks.0...` + 2) Official-ish export: `*model*.safetensors|bin` possibly with keys like `dit.blocks.0...` or `blocks.0...` + """ + # Direct file + if os.path.isfile(pretrained_path): + state = _strip_module_prefix(_load_state_dict(pretrained_path)) + # Accept either `dit.*` or plain keys. + if any(k.startswith("dit.") for k in state): + state = {k[len("dit.") :]: v for k, v in state.items() if k.startswith("dit.")} + result = dit.load_state_dict(state, strict=False) + logger.info( + f"Loaded DiT from file. missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}" + ) + return + + # Directory: prefer explicit Wan trainer file name, else fall back to pattern search + candidates: list[str] = [] + for fname in ( + "dit_model.safetensors", + "diffusion_pytorch_model.safetensors", + "model.safetensors", + ): + p = os.path.join(pretrained_path, fname) + if os.path.exists(p): + candidates.append(p) + if not candidates: + candidates = sorted(glob.glob(os.path.join(pretrained_path, "*model*.safetensors"))) + if not candidates: + candidates = sorted(glob.glob(os.path.join(pretrained_path, "*model*.bin"))) + if not candidates: + raise FileNotFoundError(f"No DiT weights found under {pretrained_path}") + + merged: dict[str, torch.Tensor] = {} + for ckpt in candidates: + part = _strip_module_prefix(_load_state_dict(ckpt)) + merged.update(part) + + if any(k.startswith("dit.") for k in merged): + merged = {k[len("dit.") :]: v for k, v in merged.items() if k.startswith("dit.")} + + result = dit.load_state_dict(merged, strict=False) + logger.info( + f"Loaded DiT from dir. files={len(candidates)} missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}" + ) + + +def build_wan_model(model_config: dict): + """ + YAML compatibility: + model_config: + name: wan + load_from_pretrained_path: /path/to/Wan2.2-TI2V-5B (optional) + config: {...} (optional overrides) + encoder: + t5_encoder: ... + autoencoder: ... + """ + encoder_cfg = model_config.get("encoder", {}) if isinstance(model_config.get("encoder"), dict) else {} + cfg_dict: dict[str, Any] = dict(model_config.get("config", {}) or {}) + if encoder_cfg: + cfg_dict.setdefault("encoder", encoder_cfg) + + # 1. Try to load config.json from pretrained path if available + pretrained_path = model_config.get("load_from_pretrained_path") + if pretrained_path: + import json + + cfg_path = os.path.join(pretrained_path, "config.json") + if os.path.exists(cfg_path): + logger.info(f"Loading config from {cfg_path}") + with open(cfg_path) as f: + loaded_cfg = json.load(f) + + # Map WanModel checkpoint config keys to WanVideoConfig keys. + mapping = { + "dim": "dit_hidden_size", + "num_layers": "dit_num_layers", + "num_heads": "dit_num_heads", + "ffn_dim": "dit_intermediate_size", + "in_dim": "dit_in_channels", + "out_dim": "dit_out_channels", + "freq_dim": "dit_freq_dim", + "text_len": "text_len", + } + + for k, v in loaded_cfg.items(): + target_key = mapping.get(k) + # Only infer values the user did not set explicitly in YAML. + if target_key is not None and target_key not in cfg_dict: + cfg_dict[target_key] = v + + # Build a WanVideoConfig-compatible object to maximize reuse of existing fields. + model_cfg = WanVideoConfig(**cfg_dict) + + # Build DiT (close to official Wan2.2 modules/model.py signature). + # model_config.config.model_type: one of {"t2v","i2v","ti2v","s2v"} + dit_task_type = getattr(model_cfg, "model_type", None) + if dit_task_type not in ("t2v", "i2v", "ti2v", "s2v"): + raise ValueError( + "wan requires explicit `model_config.config.model_type` in YAML " + "(one of: t2v / i2v / ti2v / s2v). " + f"Got: {dit_task_type!r}" + ) + dit = WanDiT( + model_type=dit_task_type, + patch_size=tuple(model_cfg.dit_patch_size), + text_len=int(getattr(model_cfg, "text_len", 512)), + in_dim=int(model_cfg.dit_in_channels), + dim=int(model_cfg.dit_hidden_size), + ffn_dim=int(model_cfg.dit_intermediate_size), + freq_dim=int(model_cfg.dit_freq_dim), + text_dim=int(model_cfg.dit_text_dim), + out_dim=int(model_cfg.dit_out_channels), + num_heads=int(model_cfg.dit_num_heads), + num_layers=int(model_cfg.dit_num_layers), + window_size=tuple(getattr(model_cfg, "dit_window_size", (-1, -1))), + qk_norm=bool(getattr(model_cfg, "dit_qk_norm", True)), + cross_attn_norm=bool(getattr(model_cfg, "dit_cross_attn_norm", True)), + eps=float(model_cfg.dit_eps), + ) + + # Build VAE + vae_type = getattr(model_cfg, "vae_type", "wan_video_vae_38") + vae_ckpt = encoder_cfg.get("vae_checkpoint") or encoder_cfg.get("autoencoder") + if vae_type in ("wan2.2", "wan_video_vae_38"): + if not vae_ckpt: + raise ValueError("wan requires `model_config.encoder.autoencoder` (Wan2.2 VAE checkpoint path)") + vae = Wan2_2_VAE(z_dim=48, vae_pth=vae_ckpt) + vae.upsampling_factor = 16 + else: + if not vae_ckpt: + raise ValueError("wan requires `model_config.encoder.autoencoder` (Wan2.1 VAE checkpoint path)") + vae = Wan2_1_VAE(z_dim=16, vae_pth=vae_ckpt) + vae.upsampling_factor = 8 + + # Build text encoder (UMT5-XXL, encoder-only). Tokenization is handled by the dataset processor. + t5_ckpt = encoder_cfg.get("t5_encoder") + if not t5_ckpt: + raise ValueError("wan requires `model_config.encoder.t5_encoder` (UMT5 encoder checkpoint path)") + text_encoder = umt5_xxl_encoder_from_checkpoint(t5_ckpt, dtype=torch.bfloat16, device="cpu") + + # Optionally load pretrained DiT + pretrained_path = model_config.get("load_from_pretrained_path") + if pretrained_path: + logger.info(f"Loading DiT weights from {pretrained_path}") + _load_dit_weights_into_module(dit, pretrained_path) + + components = WanComponents(dit=dit, vae=vae, text_encoder=text_encoder, image_encoder=None) + pipeline = WanFlowMatchTrainPipeline() + + model = WanForTraining( + components=components, + train_pipeline=pipeline, + model_config=model_cfg, + raw_config={ + "model_config": model_config, + "resolved_model_cfg": getattr(model_cfg, "__dict__", cfg_dict), + }, + trainable_modules=getattr(model_cfg, "trainable_modules", None), + ) + + total_params, trainable_params = count_parameters(model) + logger.info(f"wan parameters: total={total_params/1e9:.3f}B trainable={trainable_params/1e9:.3f}B") + if hasattr(model, "freeze_except"): + model.freeze_except() + total_params, trainable_params = count_parameters(model) + logger.info(f"wan after freeze: total={total_params/1e9:.3f}B trainable={trainable_params/1e9:.3f}B") + + return model diff --git a/primus/backends/diffusion/models/wan/__init__.py b/primus/backends/diffusion/models/wan/__init__.py new file mode 100644 index 000000000..86c01a615 --- /dev/null +++ b/primus/backends/diffusion/models/wan/__init__.py @@ -0,0 +1,20 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +wan: Components + TaskPipeline design for Wan PyTorch/FSDP training. + +Goal: +- Keep modeling close to official Wan2.2 (pure torch modules). +- Keep trainer generic: trainer calls model(batch, scheduler) -> {"loss": ...}. +- Decouple training/inference workflow (pipeline) from modeling (modules). +""" + +from .adapter import WanForTraining + +__all__ = [ + "WanForTraining", +] diff --git a/primus/backends/diffusion/models/wan/adapter.py b/primus/backends/diffusion/models/wan/adapter.py new file mode 100644 index 000000000..995177d6e --- /dev/null +++ b/primus/backends/diffusion/models/wan/adapter.py @@ -0,0 +1,154 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch +import torch.nn as nn + +from primus.backends.diffusion.models.interface import GenAIModel + +from .components import WanComponents +from .train_pipeline import WanFlowMatchTrainPipeline + + +@dataclass +class WanConfigShim: + """ + A tiny config shim to keep trainer saving behavior happy. + We intentionally do not implement a full HF/Diffusers config system. + """ + + raw: dict + + def save_pretrained(self, save_directory: str): + # Best-effort: keep minimal JSON for debugging/reproducibility. + import json + import os + + os.makedirs(save_directory, exist_ok=True) + path = os.path.join(save_directory, "wan_config.json") + with open(path, "w") as f: + json.dump(self.raw, f, indent=2, sort_keys=True) + + def to_dict(self): + return self.raw + + +class WanForTraining(GenAIModel, nn.Module): + """ + A thin adapter that exposes the call pattern expected by Wan trainers. + + Trainers call: + outputs = model(batch, scheduler) -> {"loss": ...} + + Internally we delegate workflow to the pipeline. + """ + + def __init__( + self, + *, + components: WanComponents, + train_pipeline: WanFlowMatchTrainPipeline, + model_config: Any, + raw_config: Optional[dict] = None, + trainable_modules: Optional[str] = None, + ): + super().__init__() + self.components = components + self.train_pipeline = train_pipeline + self.model_config = model_config + self.trainable_modules = trainable_modules + + # Expose common attribute names expected by trainers/FSDP ignore regex. + self.dit = components.dit + self.vae = components.vae + self.text_encoder = components.text_encoder + self.image_encoder = components.image_encoder + + self.config = WanConfigShim(raw=raw_config or {}) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def to(self, *args, **kwargs): + """ + Override to propagate .to() to non-nn.Module components. + + Wan VAE wrappers (Wan2_1_VAE, Wan2_2_VAE) are plain Python classes + with a custom .to() method, not nn.Module subclasses. PyTorch's + nn.Module.to() only recurses into registered submodules, so the VAE + would be silently skipped — causing device mismatches on multi-GPU. + """ + result = super().to(*args, **kwargs) + for component in (self.vae, self.image_encoder): + if component is not None and not isinstance(component, nn.Module) and hasattr(component, "to"): + component.to(*args, **kwargs) + return result + + def freeze_except(self): + """ + Keep the Wan training behavior: freeze non-trainable modules. + Default: train only DiT. + """ + mode = ( + self.trainable_modules or getattr(self.model_config, "trainable_modules", None) or "dit" + ).lower() + + def freeze(m: nn.Module): + for p in m.parameters(): + p.requires_grad_(False) + + def unfreeze(m: nn.Module): + for p in m.parameters(): + p.requires_grad_(True) + + # Freeze everything first + freeze(self) + + # Unfreeze requested parts + if mode in ("dit", "diffusion", "backbone"): + unfreeze(self.dit) + elif mode in ("all",): + unfreeze(self) + else: + # Conservative default: DiT only + unfreeze(self.dit) + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + """ + Activated by HF Trainer when `gradient_checkpointing=True`. + """ + if self.dit and hasattr(self.dit, "gradient_checkpointing"): + self.dit.gradient_checkpointing = True + + def forward(self, *args, **kwargs): + # Match existing trainer call convention: model(batch, scheduler) + if len(args) >= 1 and isinstance(args[0], dict): + scheduler = args[1] if len(args) > 1 else kwargs.get("scheduler", None) + return self.forward_train(args[0], scheduler=scheduler) + raise TypeError("WanForTraining.forward expects (batch_dict, scheduler)") + + def forward_train(self, batch: Dict[str, Any], scheduler: Any = None) -> Dict[str, torch.Tensor]: + if scheduler is None: + raise ValueError("scheduler must be provided by trainer") + return self.train_pipeline.compute_loss( + components=self.components, + batch=batch, + scheduler=scheduler, + model_config=self.model_config, + ) + + def forward_inference(self, batch: Dict[str, Any], **kwargs): + raise NotImplementedError("wan inference pipeline not wired yet") diff --git a/primus/backends/diffusion/models/wan/attention_backend.py b/primus/backends/diffusion/models/wan/attention_backend.py new file mode 100644 index 000000000..306f8afe4 --- /dev/null +++ b/primus/backends/diffusion/models/wan/attention_backend.py @@ -0,0 +1,228 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +# +# This file is vendored from `Wan2.2/wan/modules/attention.py` with minimal changes: +# - keep PyTorch implementation +# - allow CPU fallback when flash-attn is unavailable or tensors are on CPU + +import warnings + +import torch + +try: + import flash_attn_interface + + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +__all__ = [ + "flash_attention", + "attention", +] + + +def flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + version=None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == "cuda" and q.size(-1) <= 256 + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # Fast path (fixed-length, no padding mask): + # Prefer flash_attn_func over flash_attn_varlen_func to reduce overhead. + if q_lens is None and k_lens is None: + qh = half(q) + kh = half(k) + vh = half(v) + qh = qh.to(vh.dtype) + kh = kh.to(vh.dtype) + if q_scale is not None: + qh = qh * q_scale + + # apply attention + if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + x = flash_attn_interface.flash_attn_func(qh, kh, vh) + if isinstance(x, tuple): + x = x[0] + else: + assert FLASH_ATTN_2_AVAILABLE + # flash_attn_func signature differs across versions; be conservative + try: + x = flash_attn.flash_attn_func( + qh, + kh, + vh, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + ) + except TypeError: + x = flash_attn.flash_attn_func(qh, kh, vh) + return x.type(out_dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + + if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn("Flash attention 3 is not available, use flash attention 2 instead.") + + # apply attention + if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p, window_size are not supported in FA3 now. + x = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + seqused_q=None, + seqused_k=None, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic, + )[0].unflatten(0, (b, lq)) + else: + assert FLASH_ATTN_2_AVAILABLE + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]) + .cumsum(0, dtype=torch.int32) + .to(q.device, non_blocking=True), + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + ).unflatten(0, (b, lq)) + + # output + return x.type(out_dtype) + + +def attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0.0, + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + fa_version=None, +): + # Minimal change vs upstream: allow CPU fallback. + if (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE) and q.device.type == "cuda": + return flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=fa_version, + ) + + if q_lens is not None or k_lens is not None: + warnings.warn( + "Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance." + ) + attn_mask = None + + out_dtype = q.dtype + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p + ) + + out = out.transpose(1, 2).contiguous().to(out_dtype) + return out diff --git a/primus/backends/diffusion/models/wan/components.py b/primus/backends/diffusion/models/wan/components.py new file mode 100644 index 000000000..3e5a854d2 --- /dev/null +++ b/primus/backends/diffusion/models/wan/components.py @@ -0,0 +1,26 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch.nn as nn + + +@dataclass +class WanComponents: + """ + A thin container for model components. + + Keep this intentionally simple: it's just a bundle of modules that a pipeline can use. + """ + + dit: nn.Module + vae: nn.Module + text_encoder: nn.Module + image_encoder: Optional[nn.Module] = None diff --git a/primus/backends/diffusion/models/wan/configuration_wanvideo.py b/primus/backends/diffusion/models/wan/configuration_wanvideo.py new file mode 100644 index 000000000..0df88d0ec --- /dev/null +++ b/primus/backends/diffusion/models/wan/configuration_wanvideo.py @@ -0,0 +1,84 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +# coding=utf-8 +# Copyright 2024 WanVideo team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict + + +class WanVideoConfig: + model_type = "wanvideo" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + dit_hidden_size: int = 3072, + dit_num_layers: int = 30, + dit_num_heads: int = 24, + dit_intermediate_size: int = 14336, + dit_patch_size: tuple = (1, 2, 2), + dit_in_channels: int = 48, + dit_out_channels: int = 48, + dit_freq_dim: int = 256, + dit_text_dim: int = 4096, + dit_eps: float = 1e-6, + dit_has_image_input: bool = False, + dit_has_image_pos_emb: bool = False, + dit_has_ref_conv: bool = False, + trainable_modules=None, + separated_timestep: bool = True, + require_clip_embedding: bool = False, + require_vae_embedding: bool = False, + fuse_vae_embedding_in_latents: bool = True, + tie_word_embeddings: bool = False, + vae_type: str = "wan_video_vae_38", + **kwargs, + ): + # DiT configuration + self.vae_type = vae_type + self.dit_hidden_size = dit_hidden_size + self.dit_num_layers = dit_num_layers + self.dit_num_heads = dit_num_heads + self.dit_intermediate_size = dit_intermediate_size + self.dit_patch_size = dit_patch_size + self.dit_in_channels = dit_in_channels + self.dit_out_channels = dit_out_channels + self.dit_freq_dim = dit_freq_dim + self.dit_text_dim = dit_text_dim + self.dit_eps = dit_eps + self.dit_has_image_input = dit_has_image_input + + self.dit_has_image_pos_emb = dit_has_image_pos_emb + self.dit_has_ref_conv = dit_has_ref_conv + + self.separated_timestep = bool(separated_timestep) + self.require_clip_embedding = require_clip_embedding + self.require_vae_embedding = require_vae_embedding + self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents + + self.trainable_modules = trainable_modules + + for k, v in kwargs.items(): + setattr(self, k, v) + + def to_dict(self) -> Dict[str, Any]: + """Convert configuration to dictionary""" + output = {k: v for k, v in self.__dict__.items() if not k.startswith("_")} + output["model_type"] = self.model_type + return output diff --git a/primus/backends/diffusion/models/wan/t5.py b/primus/backends/diffusion/models/wan/t5.py new file mode 100644 index 000000000..d9cdfab71 --- /dev/null +++ b/primus/backends/diffusion/models/wan/t5.py @@ -0,0 +1,299 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +# +# Notes for Wan training: +# - We keep this file focused on the PyTorch encoder only. +# - We do NOT vendor the tokenizer wrapper (ftfy/regex deps). Training uses +# `input_ids/attention_mask` produced by the existing processor instead. + +from __future__ import annotations + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = [ + "T5Encoder", + "umt5_xxl_encoder", + "umt5_xxl_encoder_from_checkpoint", +] + + +def fp16_clamp(x): + if x.dtype == torch.float16 and torch.isinf(x).any(): + clamp = torch.finfo(x.dtype).max - 1000 + x = torch.clamp(x, min=-clamp, max=clamp) + return x + + +class GELU(nn.Module): + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + + +class T5LayerNorm(nn.Module): + def __init__(self, dim, eps=1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) + if self.weight.dtype in (torch.float16, torch.bfloat16): + x = x.type_as(self.weight) + return self.weight * x + + +class T5Attention(nn.Module): + def __init__(self, dim, dim_attn, num_heads, dropout=0.1): + assert dim_attn % num_heads == 0 + super().__init__() + self.dim = dim + self.dim_attn = dim_attn + self.num_heads = num_heads + self.head_dim = dim_attn // num_heads + + self.q = nn.Linear(dim, dim_attn, bias=False) + self.k = nn.Linear(dim, dim_attn, bias=False) + self.v = nn.Linear(dim, dim_attn, bias=False) + self.o = nn.Linear(dim_attn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, context=None, mask=None, pos_bias=None): + """ + x: [B, L1, C] + context: [B, L2, C] or None + mask: [B, L2] or [B, L1, L2] or None + """ + context = x if context is None else context + b, n, c = x.size(0), self.num_heads, self.head_dim + + q = self.q(x).view(b, -1, n, c) + k = self.k(context).view(b, -1, n, c) + v = self.v(context).view(b, -1, n, c) + + attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) + if pos_bias is not None: + attn_bias += pos_bias + if mask is not None: + assert mask.ndim in (2, 3) + mask = mask.view(b, 1, 1, -1) if mask.ndim == 2 else mask.unsqueeze(1) + attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min) + + # T5 does not use scaling + attn = torch.einsum("binc,bjnc->bnij", q, k) + attn_bias + attn = F.softmax(attn.float(), dim=-1).type_as(attn) + x = torch.einsum("bnij,bjnc->binc", attn, v) + + x = x.reshape(b, -1, n * c) + x = self.o(x) + x = self.dropout(x) + return x + + +class T5FeedForward(nn.Module): + def __init__(self, dim, dim_ffn, dropout=0.1): + super().__init__() + self.dim = dim + self.dim_ffn = dim_ffn + self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU()) + self.fc1 = nn.Linear(dim, dim_ffn, bias=False) + self.fc2 = nn.Linear(dim_ffn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + x = self.fc1(x) * self.gate(x) + x = self.dropout(x) + x = self.fc2(x) + x = self.dropout(x) + return x + + +class T5RelativeEmbedding(nn.Module): + def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): + super().__init__() + self.num_buckets = num_buckets + self.num_heads = num_heads + self.bidirectional = bidirectional + self.max_dist = max_dist + self.embedding = nn.Embedding(num_buckets, num_heads) + + def forward(self, lq, lk): + device = self.embedding.weight.device + rel_pos = torch.arange(lk, device=device).unsqueeze(0) - torch.arange(lq, device=device).unsqueeze(1) + rel_pos = self._relative_position_bucket(rel_pos) + rel_pos_embeds = self.embedding(rel_pos) + rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(0) # [1, N, Lq, Lk] + return rel_pos_embeds.contiguous() + + def _relative_position_bucket(self, rel_pos): + if self.bidirectional: + num_buckets = self.num_buckets // 2 + rel_buckets = (rel_pos > 0).long() * num_buckets + rel_pos = torch.abs(rel_pos) + else: + num_buckets = self.num_buckets + rel_buckets = 0 + rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) + + max_exact = num_buckets // 2 + rel_pos_large = ( + max_exact + + ( + torch.log(rel_pos.float() / max_exact) + / math.log(self.max_dist / max_exact) + * (num_buckets - max_exact) + ).long() + ) + rel_pos_large = torch.min(rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) + rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) + return rel_buckets + + +class T5SelfAttention(nn.Module): + def __init__(self, dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos=True, dropout=0.1): + super().__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + self.norm1 = T5LayerNorm(dim) + self.attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm2 = T5LayerNorm(dim) + self.ffn = T5FeedForward(dim, dim_ffn, dropout) + self.pos_embedding = ( + None if shared_pos else T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True) + ) + + def forward(self, x, mask=None, pos_bias=None): + e = pos_bias if self.shared_pos else self.pos_embedding(x.size(1), x.size(1)) + x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) + x = fp16_clamp(x + self.ffn(self.norm2(x))) + return x + + +def init_weights(m): + if isinstance(m, T5LayerNorm): + nn.init.ones_(m.weight) + elif isinstance(m, T5FeedForward): + nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5) + nn.init.normal_(m.fc1.weight, std=m.dim**-0.5) + nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5) + elif isinstance(m, T5Attention): + nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn) ** -0.5) + nn.init.normal_(m.k.weight, std=m.dim**-0.5) + nn.init.normal_(m.v.weight, std=m.dim**-0.5) + nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn) ** -0.5) + elif isinstance(m, T5RelativeEmbedding): + nn.init.normal_(m.embedding.weight, std=(2 * m.num_buckets * m.num_heads) ** -0.5) + + +class T5Encoder(nn.Module): + def __init__( + self, vocab, dim, dim_attn, dim_ffn, num_heads, num_layers, num_buckets, shared_pos=True, dropout=0.1 + ): + super().__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_layers = num_layers + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + self.token_embedding = vocab if isinstance(vocab, nn.Embedding) else nn.Embedding(vocab, dim) + self.pos_embedding = ( + T5RelativeEmbedding(num_buckets, num_heads, bidirectional=True) if shared_pos else None + ) + self.dropout = nn.Dropout(dropout) + self.blocks = nn.ModuleList( + [ + T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, shared_pos, dropout) + for _ in range(num_layers) + ] + ) + self.norm = T5LayerNorm(dim) + self.apply(init_weights) + + def forward(self, ids, mask=None): + x = self.token_embedding(ids) + x = self.dropout(x) + e = self.pos_embedding(x.size(1), x.size(1)) if self.shared_pos else None + for block in self.blocks: + x = block(x, mask, pos_bias=e) + x = self.norm(x) + x = self.dropout(x) + return x + + +def umt5_xxl_encoder(**kwargs) -> T5Encoder: + """ + Official UMT5-XXL encoder-only config. + """ + cfg = dict( + vocab=256384, + dim=4096, + dim_attn=4096, + dim_ffn=10240, + num_heads=64, + num_layers=24, + num_buckets=32, + shared_pos=False, + dropout=0.1, + ) + cfg.update(**kwargs) + model = T5Encoder(**cfg) + model.eval().requires_grad_(False) + return model + + +def umt5_xxl_encoder_from_checkpoint( + checkpoint_path: str, + *, + dtype: torch.dtype = torch.bfloat16, + device: str | torch.device = "cpu", +) -> T5Encoder: + """ + Build UMT5-XXL encoder-only model and load weights without allocating + full parameter tensors upfront. + + This mirrors the official pattern of meta-init + `assign=True` loading. + """ + if not checkpoint_path: + raise ValueError("checkpoint_path is required") + + # Build on meta to avoid allocating a ~1B parameter embedding upfront. + with torch.device("meta"): + model = umt5_xxl_encoder() + + state = torch.load(checkpoint_path, map_location="cpu") + if isinstance(state, dict) and "model" in state and isinstance(state["model"], dict): + state = state["model"] + if not isinstance(state, dict): + raise ValueError(f"Unsupported checkpoint format at {checkpoint_path}") + + # Strip potential DDP prefix. + if any(k.startswith("module.") for k in state.keys()): + state = {k[len("module.") :]: v for k, v in state.items()} + + try: + model.load_state_dict(state, strict=False, assign=True) + except TypeError as exc: + raise RuntimeError( + "Your PyTorch build does not support `assign=True` for loading into meta-initialized modules. " + "Please upgrade torch (>=2.0) or adjust loading strategy." + ) from exc + + model = model.to(device=device, dtype=dtype).eval().requires_grad_(False) + return model diff --git a/primus/backends/diffusion/models/wan/train_pipeline.py b/primus/backends/diffusion/models/wan/train_pipeline.py new file mode 100644 index 000000000..3f8a4d1b8 --- /dev/null +++ b/primus/backends/diffusion/models/wan/train_pipeline.py @@ -0,0 +1,314 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F + +from .components import WanComponents + + +@dataclass +class WanFlowMatchTrainPipelineConfig: + """ + Minimal training-pipeline config. + + We intentionally keep this small and derive most behavior from the model config + to maximize YAML reuse. + """ + + # For Wan VAEs, temporal downsample is typically (4, 1): 4n+1 frames. + time_division_factor: int = 4 + time_division_remainder: int = 1 + + +class WanFlowMatchTrainPipeline: + """ + Flow-Matching training pipeline for Wan-style DiT models. + + Contract: + compute_loss(components, batch, scheduler) -> {"loss": Tensor, ...} + """ + + def __init__(self, cfg: Optional[WanFlowMatchTrainPipelineConfig] = None): + self.cfg = cfg or WanFlowMatchTrainPipelineConfig() + + @staticmethod + def _encode_prompt( + text_encoder: torch.nn.Module, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + ): + # Match existing wan_new behavior: zero-out padded embeddings explicitly. + seq_lens = attention_mask.gt(0).sum(dim=1).long() + prompt_emb = text_encoder(input_ids, attention_mask) + for i, v in enumerate(seq_lens): + prompt_emb[i, v:] = 0 + return prompt_emb + + @staticmethod + def _get_seed_from_env_or_batch(batch: Dict[str, Any]) -> Optional[int]: + # Keep parity with existing scripts/wan_new behavior. + import os + + if os.environ.get("FIXED_SEED"): + try: + return int(os.environ["FIXED_SEED"]) + except ValueError as exc: + raise ValueError(f"Invalid FIXED_SEED value: {os.environ['FIXED_SEED']}") from exc + seed = batch.get("seed", None) + if seed is None: + return None + if isinstance(seed, torch.Tensor): + flat_seed = seed.detach().reshape(-1).cpu() + if flat_seed.numel() == 0: + return None + if flat_seed.numel() > 1 and not bool((flat_seed == flat_seed[0]).all()): + raise ValueError( + "Per-sample dataset seeds are not supported; got different seeds in one batch." + ) + seed = flat_seed[0].item() + elif isinstance(seed, (list, tuple)): + if not seed: + return None + if len(seed) > 1 and any(value != seed[0] for value in seed): + raise ValueError( + "Per-sample dataset seeds are not supported; got different seeds in one batch." + ) + seed = seed[0] + return int(seed) + + @staticmethod + def _randn_like_on_cpu_then_to( + x: torch.Tensor, + *, + seed: Optional[int], + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + generator = None + if seed is not None: + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + noise = torch.randn(x.shape, generator=generator, device="cpu", dtype=torch.float32) + return noise.to(device=device, dtype=dtype) + + @staticmethod + def _vae_encode(vae: torch.nn.Module, videos_bcthw: torch.Tensor) -> torch.Tensor: + # Wan VAEs in this repo use list-of-tensors interface: List[[C,T,H,W]] -> List[[C,F,H',W']] + videos_list = [videos_bcthw[i] for i in range(videos_bcthw.shape[0])] + latents_list = vae.encode(videos_list) + return torch.stack(latents_list) + + @staticmethod + def _get_dit_patch_size(model_config: Any) -> tuple[int, int, int]: + patch_size = tuple(getattr(model_config, "dit_patch_size", (1, 2, 2))) + if len(patch_size) != 3: + raise ValueError(f"Expected 3D dit_patch_size, got {patch_size!r}") + return patch_size + + @staticmethod + def _pad_latents_for_dit( + latents: torch.Tensor, *, patch_size: tuple[int, int, int] + ) -> tuple[torch.Tensor, tuple[int, int]]: + _, _, _, height, width = latents.shape + _, patch_h, patch_w = patch_size + pad_h = (-height) % patch_h + pad_w = (-width) % patch_w + if pad_h == 0 and pad_w == 0: + return latents, (height, width) + + # Pad only the latent-space bottom/right edges so DiT patchify works on any + # VAE output shape. We crop predictions back before computing loss. + latents = F.pad(latents, (0, pad_w, 0, pad_h)) + return latents, (height, width) + + @staticmethod + def _crop_latents(latents: torch.Tensor, *, spatial_size: tuple[int, int]) -> torch.Tensor: + height, width = spatial_size + return latents[..., :height, :width] + + @staticmethod + def _select_timestep(scheduler: Any, device: torch.device) -> torch.Tensor: + """ + Match `wan_new` timestep selection: + - If FIXED_TIMESTEP is set, use that discrete index into [0, num_train_timesteps). + - Else uniform randint over [0, num_train_timesteps). + """ + import os + + if os.environ.get("FIXED_TIMESTEP"): + try: + fixed_step = int(os.environ["FIXED_TIMESTEP"]) + except ValueError as exc: + raise ValueError(f"Invalid FIXED_TIMESTEP value: {os.environ['FIXED_TIMESTEP']}") from exc + max_step = int(scheduler.num_train_timesteps) - 1 + fixed_step = max(0, min(fixed_step, max_step)) + timestep_id = torch.tensor([fixed_step], device=device) + else: + timestep_id = torch.randint(0, int(scheduler.num_train_timesteps), (1,), device=device) + + # scheduler.timesteps live on CPU in this repo; `wan_new` indexes with cpu tensor. + timestep = scheduler.timesteps[timestep_id.cpu()].float() + return timestep.to(device=device) + + @staticmethod + def _maybe_expand_separated_timestep( + *, + timestep: torch.Tensor, + x_list: list[torch.Tensor], + patch_size: tuple[int, int, int], + enabled: bool, + ) -> torch.Tensor: + """ + Match `wan_new.forward_dit` separated-timestep behavior (only used when enabled). + For each sample, build per-token timestep [L] and set first-frame patches to 0. + Returns: + - t: [B] if not enabled + - t: [B, L] if enabled + """ + if not enabled: + if timestep.ndim == 0: + return timestep.unsqueeze(0).repeat(len(x_list)) + if timestep.ndim == 1 and timestep.numel() == 1 and len(x_list) > 1: + return timestep.repeat(len(x_list)) + if timestep.ndim == 1 and timestep.shape[0] != len(x_list): + return timestep.repeat(len(x_list)) + return timestep + + d_f, d_h, d_w = patch_size + if timestep.ndim == 0: + timestep_b = timestep.unsqueeze(0).repeat(len(x_list)) + elif timestep.ndim == 1 and timestep.shape[0] != len(x_list): + timestep_b = timestep.repeat(len(x_list)) + else: + timestep_b = timestep + + t_expand_list: list[torch.Tensor] = [] + for i, x in enumerate(x_list): + # x: [C, F, H, W] in latent space + f, h, w = x.shape[1], x.shape[2], x.shape[3] + seq_len = (f // d_f) * (h // d_h) * (w // d_w) + t_seq = torch.full( + (seq_len,), + timestep_b[i], + device=timestep_b.device, + dtype=timestep_b.dtype, + ) + spatial_patches = (h // d_h) * (w // d_w) + t_seq[:spatial_patches] = 0 + t_expand_list.append(t_seq) + return torch.stack(t_expand_list) + + def compute_loss( + self, + *, + components: WanComponents, + batch: Dict[str, Any], + scheduler: Any, + model_config: Any, + ) -> Dict[str, torch.Tensor]: + """ + batch requirements (from current dataset/processor): + - video: Tensor [B, C, T, H, W] + - input_ids: Tensor [B, L] + - attention_mask: Tensor [B, L] + """ + video = batch.get("video") + if video is None: + raise ValueError("Batch must contain 'video'") + if not isinstance(video, torch.Tensor) or video.ndim != 5: + raise ValueError( + f"Expected batch['video'] as 5D tensor [B,C,T,H,W], got {type(video)} shape={getattr(video, 'shape', None)}" + ) + + input_ids = batch.get("input_ids") + attention_mask = batch.get("attention_mask") + if input_ids is None or attention_mask is None: + raise ValueError("Batch must contain 'input_ids' and 'attention_mask'") + + device = next(components.dit.parameters()).device + dtype = next(components.dit.parameters()).dtype + + video = video.to(device=device, dtype=dtype, non_blocking=True) + input_ids = input_ids.to(device=device, non_blocking=True) + attention_mask = attention_mask.to(device=device, non_blocking=True) + + # 1) Encode to latents + # Keep VAE/text encoder in eval() by default (consistent with wan_new practice). + components.text_encoder.eval() + try: + components.vae.eval() + except AttributeError: + # Some VAE compatibility wrappers may not expose eval(); real eval + # failures should still surface. + pass + + with torch.no_grad(): + input_latents = self._vae_encode(components.vae, video).to(device=device, dtype=dtype) + patch_size = self._get_dit_patch_size(model_config) + input_latents, original_latent_spatial_size = self._pad_latents_for_dit( + input_latents, patch_size=patch_size + ) + + # 2) Sample timestep + noise (match `wan_new`: noise on CPU, timestep from scheduler.timesteps) + timestep = self._select_timestep(scheduler, device=device) # [1] (float) + seed = self._get_seed_from_env_or_batch(batch) + noise = self._randn_like_on_cpu_then_to(input_latents, seed=seed, dtype=dtype, device=device) + + # 3) Diffusion target + noise injection (match `wan_new`) + # Note: in this repo's FlowMatchScheduler, `training_target(sample, noise, t)` uses `noise - sample`. + # `wan_new` passes `input_latents` (not noisy latents). + target = scheduler.training_target(input_latents, noise, timestep) + noisy_latents = scheduler.add_noise(input_latents, noise, timestep=timestep) + + # 4) Text embeddings + with torch.no_grad(): + context = self._encode_prompt(components.text_encoder, input_ids, attention_mask) + + # 5) DiT forward (official interface: List[Tensor] per-sample) + x_list = [noisy_latents[i] for i in range(noisy_latents.shape[0])] + context_list = [context[i] for i in range(context.shape[0])] + + # seq_len matches DiT Conv3d patchification over [F, H, W]. + d_f, d_h, d_w = patch_size + max_seq_len = 0 + for x in x_list: + seq_len = (x.shape[1] // d_f) * (x.shape[2] // d_h) * (x.shape[3] // d_w) + max_seq_len = max(max_seq_len, seq_len) + + separated = bool(getattr(model_config, "separated_timestep", False)) + fuse_flag = bool(getattr(model_config, "fuse_vae_embedding_in_latents", False)) + t = self._maybe_expand_separated_timestep( + timestep=timestep, + x_list=x_list, + patch_size=(d_f, d_h, d_w), + enabled=bool(separated and fuse_flag), + ) + + sp_group = batch.get("sp_group", None) + noise_pred_list = components.dit( + x=x_list, + t=t, + context=context_list, + seq_len=max_seq_len, + y=None, + sp_group=sp_group, + ) + noise_pred = torch.stack(noise_pred_list) + noise_pred = self._crop_latents(noise_pred, spatial_size=original_latent_spatial_size) + target = self._crop_latents(target, spatial_size=original_latent_spatial_size) + + # 6) Loss + loss = F.mse_loss(noise_pred.float(), target.float(), reduction="mean") + # `wan_new` uses scalar `timestep` for weighting (not expanded per-token) + weight = scheduler.training_weight(timestep) + loss = loss * weight + return {"loss": loss} diff --git a/primus/backends/diffusion/models/wan/vae2_1.py b/primus/backends/diffusion/models/wan/vae2_1.py new file mode 100644 index 000000000..ac6b44760 --- /dev/null +++ b/primus/backends/diffusion/models/wan/vae2_1.py @@ -0,0 +1,646 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Wan2.1 VAE (vendored from official Wan2.2 repo). + +Source: `Wan2.2/wan/modules/vae2_1.py` + +Minimal changes: +- add small `to()/eval()/train()` helpers on the wrapper to fit Wan trainer usage + (HF trainer sometimes calls `model.vae.to(dtype=...)`). +""" + +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import logging +from typing import List + +import torch +import torch.cuda.amp as amp +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +__all__ = [ + "Wan2_1_VAE", + "WanVAE", +] + +CACHE_T = 2 + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = ( + self.padding[2], + self.padding[2], + self.padding[1], + self.padding[1], + 2 * self.padding[0], + 0, + ) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + return super().forward(x) + + +class RMS_norm(nn.Module): + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + def __init__(self, dim, mode): + assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d") + super().__init__() + self.dim = dim + self.mode = mode + + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim // 2, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim // 2, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1, c2) + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[: c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2 :, :, -1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + +class ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), + nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1), + ) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + h = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm(x) + q, k, v = ( + self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1) + ) + x = F.scaled_dot_product_attention(q, k, v) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + x = self.proj(x) + x = rearrange(x, "(b t) c h w-> b c t h w", t=t) + return x + identity + + +class Encoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + if i != len(dim_mult) - 1: + mode = "downsample3d" if temperal_downsample[i] else "downsample2d" + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Decoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + self.middle = nn.Sequential( + ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout), + ) + + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + if i == 1 or i == 2 or i == 3: + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + if i != len(dim_mult) - 1: + mode = "upsample3d" if temperal_upsample[i] else "upsample2d" + upsamples.append(Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, 3, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class WanVAE_(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + self.encoder = Encoder3d( + dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout + ) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d( + dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout + ) + + def forward(self, x): + mu, log_var = self.encode(x) + z = self.reparameterize(mu, log_var) + x_recon = self.decode(z) + return x_recon, mu, log_var + + def encode(self, x, scale): + self.clear_cache() + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder( + x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx + ) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1) + else: + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + def decode(self, z, scale): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1) + else: + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx + ) + else: + out_ = self.decoder( + x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx + ) + out = torch.cat([out, out_], 2) + self.clear_cache() + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +def _video_vae(pretrained_path=None, z_dim=None, device="cpu", **kwargs): + cfg = dict( + dim=96, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0, + ) + cfg.update(**kwargs) + + with torch.device("meta"): + model = WanVAE_(**cfg) + + logging.info(f"loading {pretrained_path}") + model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True) + return model + + +class Wan2_1_VAE: + def __init__(self, z_dim=16, vae_pth="cache/vae_step_411000.pth", dtype=torch.float, device="cuda"): + self.dtype = dtype + self.device = device + + mean = [ + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, + ] + std = [ + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, + ] + self.mean = torch.tensor(mean, dtype=dtype, device=device) + self.std = torch.tensor(std, dtype=dtype, device=device) + self.scale = [self.mean, 1.0 / self.std] + + if vae_pth is None: + raise ValueError("Wan2_1_VAE requires `vae_pth` (checkpoint path).") + self.model = _video_vae(pretrained_path=vae_pth, z_dim=z_dim).eval().requires_grad_(False).to(device) + + # --- Wan trainer compatibility helpers (non-upstream) --- + def to(self, *args, **kwargs): + self.model.to(*args, **kwargs) + if "dtype" in kwargs and kwargs["dtype"] is not None: + self.dtype = kwargs["dtype"] + device = kwargs.get("device", None) + if device is None and len(args) > 0: + device = args[0] + if device is not None: + self.device = device + self.mean = self.mean.to(device) + self.std = self.std.to(device) + self.scale = [self.mean, 1.0 / self.std] + return self + + def eval(self): + self.model.eval() + return self + + def train(self, mode: bool = True): + self.model.train(mode) + return self + + def encode(self, videos: List[torch.Tensor]): + """ + videos: A list of videos each with shape [C, T, H, W]. + """ + with amp.autocast(dtype=self.dtype): + # Avoid forcing fp32 outputs: it increases memory and slows training. + return [self.model.encode(u.unsqueeze(0), self.scale).squeeze(0) for u in videos] + + def decode(self, zs: List[torch.Tensor]): + with amp.autocast(dtype=self.dtype): + return [self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0) for u in zs] + + +# Backward-compat alias (older Wan naming) +WanVAE = Wan2_1_VAE diff --git a/primus/backends/diffusion/models/wan/vae2_2.py b/primus/backends/diffusion/models/wan/vae2_2.py new file mode 100644 index 000000000..05998db45 --- /dev/null +++ b/primus/backends/diffusion/models/wan/vae2_2.py @@ -0,0 +1,873 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Wan2.2 VAE (vendored from official Wan2.2 repo). + +Source: `Wan2.2/wan/modules/vae2_2.py` + +Minimal changes: +- add small `to()/eval()/train()` helpers on the wrapper to fit Wan trainer usage + (HF trainer sometimes calls `model.vae.to(dtype=...)`). +""" + +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import logging +from typing import List + +import torch +import torch.cuda.amp as amp +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +__all__ = [ + "Wan2_2_VAE", +] + +CACHE_T = 2 + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = ( + self.padding[2], + self.padding[2], + self.padding[1], + self.padding[1], + 2 * self.padding[0], + 0, + ) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + return super().forward(x) + + +class RMS_norm(nn.Module): + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + def __init__(self, dim, mode): + assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d") + super().__init__() + self.dim = dim + self.mode = mode + + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight.detach().clone() + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1, c2) + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight = nn.Parameter(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data.detach().clone() + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[: c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2 :, :, -1, 0, 0] = init_matrix + conv.weight = nn.Parameter(conv_weight) + nn.init.zeros_(conv.bias.data) + + +class ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), + nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1), + ) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + h = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm(x) + q, k, v = ( + self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1) + ) + x = F.scaled_dot_product_attention(q, k, v) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + x = self.proj(x) + x = rearrange(x, "(b t) c h w-> b c t h w", t=t) + return x + identity + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + return x + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size) + return x + + +class AvgDown3D(nn.Module): + def __init__(self, in_channels, out_channels, factor_t, factor_s=1): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + b, c, t, h, w = x.shape + x = x.view( + b, + c, + t // self.factor_t, + self.factor_t, + h // self.factor_s, + self.factor_s, + w // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(b, c * self.factor, t // self.factor_t, h // self.factor_s, w // self.factor_s) + x = x.view( + b, self.out_channels, self.group_size, t // self.factor_t, h // self.factor_s, w // self.factor_s + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + def __init__(self, in_channels: int, out_channels: int, factor_t, factor_s=1): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False): + super().__init__() + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample(out_dim, mode=mode)) + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + x_copy = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + return x + self.avg_shortcut(x_copy) + + +class Up_ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False): + super().__init__() + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2 if up_flag else 1, + ) + else: + self.avg_shortcut = None + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + if up_flag: + mode = "upsample3d" if temperal_upsample else "upsample2d" + upsamples.append(Resample(out_dim, mode=mode)) + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False): + if feat_idx is None: + feat_idx = [0] + x_main = x.clone() + for module in self.upsamples: + x_main = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + x_shortcut = self.avg_shortcut(x, first_chunk) + return x_main + x_shortcut + return x_main + + +class Encoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + dims = [dim * u for u in [1] + dim_mult] + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = temperal_downsample[i] if i < len(temperal_downsample) else False + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(dim_mult) - 1, + ) + ) + self.downsamples = nn.Sequential(*downsamples) + + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=None): + if feat_idx is None: + feat_idx = [0] + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Decoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0, + ): + super().__init__() + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + self.middle = nn.Sequential( + ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout), + ) + + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False + upsamples.append( + Up_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temperal_upsample=t_up_flag, + up_flag=i != len(dim_mult) - 1, + ) + ) + self.upsamples = nn.Sequential(*upsamples) + + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=None, first_chunk=False): + if feat_idx is None: + feat_idx = [0] + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2 + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class WanVAE_(nn.Module): + # NOTE: Keep this identical to upstream for numerical parity. + def __init__( + self, + dim=160, + dec_dim=256, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + self.encoder = Encoder3d( + dim, z_dim * 2, dim_mult, num_res_blocks, attn_scales, self.temperal_downsample, dropout + ) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d( + dec_dim, z_dim, dim_mult, num_res_blocks, attn_scales, self.temperal_upsample, dropout + ) + + def forward(self, x, scale=[0, 1]): + mu = self.encode(x, scale) + x_recon = self.decode(mu, scale) + return x_recon, mu + + def encode(self, x, scale): + self.clear_cache() + x = patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder( + x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx + ) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1) + else: + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu + + def decode(self, z, scale): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1) + else: + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i : i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx, + first_chunk=True, + ) + else: + out_ = self.decoder( + x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx + ) + out = torch.cat([out, out_], 2) + out = unpatchify(out, patch_size=2) + self.clear_cache() + return out + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +def _video_vae(pretrained_path=None, z_dim=16, dim=160, device="cpu", **kwargs): + cfg = dict( + dim=dim, + z_dim=z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, True], + dropout=0.0, + ) + cfg.update(**kwargs) + with torch.device("meta"): + model = WanVAE_(**cfg) + logging.info(f"loading {pretrained_path}") + model.load_state_dict(torch.load(pretrained_path, map_location=device), assign=True) + return model + + +class Wan2_2_VAE: + def __init__( + self, + z_dim=48, + c_dim=160, + vae_pth=None, + dim_mult=[1, 2, 4, 4], + temperal_downsample=[False, True, True], + dtype=torch.bfloat16, + device="cuda", + ): + self.dtype = dtype + self.device = device + mean = torch.tensor( + [ + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.1570, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.1230, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.0520, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, + ], + dtype=dtype, + device=device, + ) + std = torch.tensor( + [ + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.4990, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.0600, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, + ], + dtype=dtype, + device=device, + ) + self.scale = [mean, 1.0 / std] + + if vae_pth is None: + raise ValueError("Wan2_2_VAE requires `vae_pth` (checkpoint path).") + self.model = ( + _video_vae( + pretrained_path=vae_pth, + z_dim=z_dim, + dim=c_dim, + dim_mult=dim_mult, + temperal_downsample=temperal_downsample, + ) + .eval() + .requires_grad_(False) + .to(device) + ) + + # --- Wan trainer compatibility helpers (non-upstream) --- + def to(self, *args, **kwargs): + dtype = kwargs.get("dtype", None) + if dtype is None: + dtype = self.dtype + else: + self.dtype = dtype + device = kwargs.get("device", None) + if device is None and len(args) > 0: + device = args[0] + if device is not None: + self.model.to(device=device, dtype=dtype) + else: + self.model.to(*args, **kwargs) + if device is not None: + self.device = device + dev = next(self.model.parameters()).device + self.scale = [t.to(device=dev, dtype=dtype) for t in self.scale] + return self + + def eval(self): + self.model.eval() + return self + + def train(self, mode: bool = True): + self.model.train(mode) + return self + + def encode(self, videos: List[torch.Tensor]): + try: + if not isinstance(videos, list): + raise TypeError("videos should be a list") + with amp.autocast(dtype=self.dtype): + # Avoid forcing fp32 outputs: it increases memory and slows training. + # Keep outputs in the model/autocast dtype (typically bf16). + return [self.model.encode(u.unsqueeze(0), self.scale).squeeze(0) for u in videos] + except TypeError as e: + logging.info(e) + return None + + def decode(self, zs: List[torch.Tensor]): + try: + if not isinstance(zs, list): + raise TypeError("zs should be a list") + with amp.autocast(dtype=self.dtype): + return [self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0) for u in zs] + except TypeError as e: + logging.info(e) + return None diff --git a/primus/backends/diffusion/models/wan/wan_dit.py b/primus/backends/diffusion/models/wan/wan_dit.py new file mode 100644 index 000000000..0d975b250 --- /dev/null +++ b/primus/backends/diffusion/models/wan/wan_dit.py @@ -0,0 +1,580 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +# +# Vendored from `Wan2.2/wan/modules/model.py` with minimal changes: +# - remove diffusers dependencies (ModelMixin/ConfigMixin/register_to_config) +# - keep pure PyTorch modeling as close as possible +# - provide a `DiTBlock` alias class so existing Wan configs can keep +# `fsdp_transformer_layer_cls_to_wrap: "DiTBlock"` + +from __future__ import annotations + +import math +from typing import Optional + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint + +from primus.backends.diffusion.attention import attention +from primus.backends.diffusion.distributed.ulysses import ( + distributed_attention, + sp_gather, + sp_split, + sp_unpad, +) + +__all__ = ["WanModel", "DiTBlock"] + + +def sinusoidal_embedding_1d(dim, position): + # preprocess + assert dim % 2 == 0 + half = dim // 2 + position = position.type(torch.float64) + + # calculation + sinusoid = torch.outer(position, torch.pow(10000, -torch.arange(half).to(position).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +@torch.amp.autocast("cuda", enabled=False) +def rope_params(max_seq_len, dim, theta=10000): + assert dim % 2 == 0 + freqs = torch.outer( + torch.arange(max_seq_len), + 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)), + ) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + +@torch.amp.autocast("cuda", enabled=False) +def rope_apply(x, grid_sizes, freqs, sp_group=None): + """ + Apply 3-D RoPE. When *sp_group* is given the input is assumed to hold + only this rank's local token chunk (S/P) and the correct positional + frequencies are selected automatically (following Wan2.2 official). + """ + n, c = x.size(2), x.size(3) // 2 + s = x.size(1) # local token count (= full seq when no SP) + + if sp_group is not None: + sp_size = dist.get_world_size(sp_group) + sp_rank = dist.get_rank(sp_group) + else: + sp_size, sp_rank = 1, 0 + + # split freqs + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + seq_len = f * h * w + + # tokens to process: local chunk when SP, valid tokens when single-GPU + t = s if sp_size > 1 else seq_len + + # precompute multipliers + x_i = torch.view_as_complex(x[i, :t].to(torch.float64).reshape(t, n, -1, 2)) + freqs_i = torch.cat( + [ + freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + + # SP: pad freqs to padded-full-length, select this rank's range + if sp_size > 1: + full_len = s * sp_size + if seq_len < full_len: + freqs_i = torch.cat( + [ + freqs_i, + torch.ones( + full_len - seq_len, + 1, + freqs_i.size(-1), + dtype=freqs_i.dtype, + device=freqs_i.device, + ), + ], + dim=0, + ) + freqs_i = freqs_i[sp_rank * s : (sp_rank + 1) * s] + + # apply rotary embedding + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, t:]]) + + # append to collection + output.append(x_i) + return torch.stack(output).to(x.dtype) + + +class WanRMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class WanLayerNorm(nn.LayerNorm): + def __init__(self, dim, eps=1e-6, elementwise_affine=False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + # Keep numerical stability by normalizing in fp32, but do NOT rely on + # module parameters being fp32 (FSDP/bf16 can cast weights/bias). + weight = self.weight.float() if self.weight is not None else None + bias = self.bias.float() if self.bias is not None else None + y = F.layer_norm(x.float(), self.normalized_shape, weight, bias, self.eps) + return y.type_as(x) + + +class WanSelfAttention(nn.Module): + def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, seq_lens, grid_sizes, freqs, sp_group: Optional[dist.ProcessGroup] = None): + r""" + Args: + x(Tensor): Shape [B, L, C] (L = S/P when SP enabled) + seq_lens(Tensor): Shape [B] — valid lengths in the *full* sequence + grid_sizes(Tensor): Shape [B, 3], (F, H, W) — full spatial grid + freqs(Tensor): Rope freqs + sp_group: Ulysses SP process group (None = no SP) + """ + b, s = x.shape[:2] + n = self.num_heads + d = self.head_dim + + x_ = x.to(dtype=self.q.weight.dtype) + q = self.norm_q(self.q(x_)).view(b, s, n, d) + k = self.norm_k(self.k(x_)).view(b, s, n, d) + v = self.v(x_).view(b, s, n, d) + + # RoPE on local tokens (SP-aware: uses rank-offset freqs when sp_group is set) + q = rope_apply(q, grid_sizes, freqs, sp_group=sp_group) + k = rope_apply(k, grid_sizes, freqs, sp_group=sp_group) + + attn_kwargs = { + "q_lens": seq_lens, + "k_lens": seq_lens, + "window_size": self.window_size, + "dtype": self.q.weight.dtype, + } + if sp_group is not None: + x = distributed_attention(q, k, v, group=sp_group, attention_fn=attention, **attn_kwargs) + else: + x = attention(q=q, k=k, v=v, **attn_kwargs) + + x = x.to(dtype=self.o.weight.dtype) + x = x.flatten(2) + x = self.o(x) + return x + + +class WanCrossAttention(WanSelfAttention): + def forward( + self, + x, + seq_lens=None, + grid_sizes=None, + freqs=None, + sp_group: Optional[dist.ProcessGroup] = None, + context=None, + context_lens=None, + ): + r""" + Args: + x(Tensor): Shape [B, L1, C] + seq_lens/grid_sizes/freqs/sp_group: accepted for WanSelfAttention + signature compatibility; cross-attention uses text context only. + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + """ + if context is None: + # Backward-compatible positional form: forward(x, context, context_lens). + context = seq_lens + context_lens = grid_sizes if context_lens is None else context_lens + + b, n, d = x.size(0), self.num_heads, self.head_dim + + x = x.to(dtype=self.q.weight.dtype) + context = context.to(dtype=self.k.weight.dtype) + + q = self.norm_q(self.q(x)).view(b, -1, n, d) + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + + x = attention(q, k, v, k_lens=context_lens, dtype=self.q.weight.dtype) + + x = x.to(dtype=self.o.weight.dtype) + x = x.flatten(2) + x = self.o(x) + return x + + +class WanAttentionBlock(nn.Module): + def __init__( + self, + dim, + ffn_dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + ): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps) + self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim) + ) + + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + e, + seq_lens, + grid_sizes, + freqs, + context, + context_lens, + sp_group: Optional[dist.ProcessGroup] = None, + ): + # Memory-critical: + # `e` can be very large ([B, seq_len, 6, dim]). Keeping it in fp32 and + # doing broadcast add in fp32 easily OOMs for Wan2.1 (more tokens). + # Align to x.dtype (bf16/fp16) like `wan_new` implementation. + e = e.to(dtype=x.dtype) + modulation = self.modulation.to(dtype=x.dtype) + e = (modulation.unsqueeze(0) + e).chunk(6, dim=2) + + # self-attention (Ulysses all-to-all happens inside self_attn) + y = self.self_attn( + self.norm1(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2), + seq_lens, + grid_sizes, + freqs, + sp_group=sp_group, + ) + x = x + y * e[2].squeeze(2) + + # cross-attention + FFN (no SP communication needed: + # each rank's visual queries attend to full text keys independently) + def cross_attn_ffn(x_, context_, context_lens_, e_): + x_ = x_ + self.cross_attn( + self.norm3(x_), + context=context_, + context_lens=context_lens_, + ) + ffn_in = self.norm2(x_) * (1 + e_[4].squeeze(2)) + e_[3].squeeze(2) + y_ = self.ffn(ffn_in) + x_ = x_ + y_ * e_[5].squeeze(2) + return x_ + + x = cross_attn_ffn(x, context, context_lens, e) + return x + + +# Wan config compatibility: FSDP auto-wrap often uses "DiTBlock". +class DiTBlock(WanAttentionBlock): + pass + + +class Head(nn.Module): + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + # Align modulation embedding dtype to x (saves memory vs fp32 broadcast). + e = e.to(dtype=x.dtype) + modulation = self.modulation.to(dtype=x.dtype) + e = (modulation.unsqueeze(0) + e.unsqueeze(2)).chunk(2, dim=2) + x = self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2)) + return x + + +class WanModel(nn.Module): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + def __init__( + self, + model_type="t2v", + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6, + ): + super().__init__() + + assert model_type in ["t2v", "i2v", "ti2v", "s2v"] + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + # Used by the Wan trainer: it sets `model.dit.gradient_checkpointing = True` + # when `trainer_args.gradient_checkpointing: true`. + self.gradient_checkpointing = False + + # embeddings + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim) + ) + self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks (use DiTBlock for YAML compatibility) + self.blocks = nn.ModuleList( + [ + DiTBlock(dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps) + for _ in range(num_layers) + ] + ) + + # head + self.head = Head(dim, out_dim, patch_size, eps) + + # rope buffers (avoid register_buffer to keep dtype stable under .to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + self.freqs = torch.cat( + [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + ], + dim=1, + ) + + self.init_weights() + + def forward(self, x, t, context, seq_len, y=None, sp_group: Optional[dist.ProcessGroup] = None): + if self.model_type == "i2v": + assert y is not None + + device = self.patch_embedding.weight.device + if self.freqs.device != device: + self.freqs = self.freqs.to(device) + + if y is not None: + x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + assert seq_lens.max() <= seq_len + x = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x]) + + # time embeddings + if t.dim() == 1: + # [B] -> [B, seq_len]: per-sample timestep broadcast across the sequence. + # Must unsqueeze before expand, otherwise the batch dim (B) is wrongly + # aligned with seq_len and breaks for batch_size > 1. + t = t.unsqueeze(1).expand(t.size(0), seq_len) + with torch.amp.autocast("cuda", dtype=torch.float32): + bt = t.size(0) + t_flat = t.flatten() + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, t_flat).unflatten(0, (bt, seq_len)).float() + ) + e0 = self.time_projection(e).unflatten(2, (6, self.dim)) + assert e.dtype == torch.float32 and e0.dtype == torch.float32 + + # Cast large modulation tensors back to model dtype early to reduce peak memory. + e = e.to(dtype=x.dtype) + e0 = e0.to(dtype=x.dtype) + + # --- Ulysses SP: split sequence-dim tensors across SP ranks --- + original_seq_len = None + if sp_group is not None: + (x, e, e0), original_seq_len = sp_split([x, e, e0], dim=1, group=sp_group) + + # context (text embeddings — NOT sliced, full text on every SP rank) + context_lens = None + context_in = torch.stack( + [torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context] + ) + context_in = context_in.to(dtype=self.text_embedding[0].weight.dtype) + context = self.text_embedding(context_in) + + freqs = self.freqs + seq_lens_ = seq_lens + grid_sizes_ = grid_sizes + e0_ = e0 + context_ = context + context_lens_ = context_lens + + for block in self.blocks: + if self.training and getattr(self, "gradient_checkpointing", False): + + def create_custom_forward(module, _sp_group): + def custom_forward(x_in, e_in, ctx_in): + return module( + x_in, + e=e_in, + seq_lens=seq_lens_, + grid_sizes=grid_sizes_, + freqs=freqs, + context=ctx_in, + context_lens=context_lens_, + sp_group=_sp_group, + ) + + return custom_forward + + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block, sp_group), + x, + e0_, + context_, + use_reentrant=False, + ) + else: + x = block( + x, + e=e0_, + seq_lens=seq_lens_, + grid_sizes=grid_sizes_, + freqs=freqs, + context=context_, + context_lens=context_lens_, + sp_group=sp_group, + ) + + # head operates on local tokens (like Wan2.2 official) + x = self.head(x, e) + + # --- Ulysses SP: gather output back to full sequence --- + if sp_group is not None: + x = sp_unpad(sp_gather(x, dim=1, group=sp_group), dim=1, original_size=original_seq_len) + + x = self.unpatchify(x, grid_sizes) + # IMPORTANT: + # Returning fp32 here dramatically increases activation/grad memory and slows + # training (DiffSynth returns bf16 here). Keep dtype consistent with model. + return x + + def unpatchify(self, x, grid_sizes): + c = self.out_dim + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[: math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum("fhwpqrc->cfphqwr", u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + def init_weights(self): + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=0.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/primus/backends/diffusion/optim/adamw_fp32_state.py b/primus/backends/diffusion/optim/adamw_fp32_state.py new file mode 100644 index 000000000..1403043f2 --- /dev/null +++ b/primus/backends/diffusion/optim/adamw_fp32_state.py @@ -0,0 +1,111 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +AdamW that keeps optimizer state in FP32 even when parameters are BF16/FP16. + +Why: +- DeepSpeed bf16 commonly keeps FP32 master weights / FP32 optimizer states. +- Vanilla torch.optim.AdamW will create exp_avg/exp_avg_sq with the same dtype as the parameter, + so bf16 params -> bf16 states, which changes early-step behavior and can destabilize. + +This optimizer: +- Stores exp_avg/exp_avg_sq in FP32 +- Applies the AdamW update in FP32 on a FP32 view of the param +- Writes the updated value back to the original parameter dtype + +Scope: +- Intended for this repo's bf16 diffusion training (single GPU or FSDP world_size=1). +- Supports common AdamW args: lr, betas, eps, weight_decay. +""" + +from __future__ import annotations + +from typing import Iterable, Optional + +import torch + + +class AdamWFP32State(torch.optim.Optimizer): + def __init__( + self, + params: Iterable[torch.nn.Parameter], + lr: float = 1e-3, + betas: tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0.0, + ): + if lr < 0.0: + raise ValueError(f"Invalid lr: {lr}") + if eps < 0.0: + raise ValueError(f"Invalid eps: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta1: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta2: {betas[1]}") + if weight_decay < 0.0: + raise ValueError(f"Invalid weight_decay: {weight_decay}") + + defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) + super().__init__(params, defaults) + + @torch.no_grad() + def step(self, closure: Optional[callable] = None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + lr: float = group["lr"] + beta1, beta2 = group["betas"] + eps: float = group["eps"] + weight_decay: float = group["weight_decay"] + + for p in group["params"]: + if p.grad is None: + continue + if p.grad.is_sparse: + raise RuntimeError("AdamWFP32State does not support sparse gradients") + + grad = p.grad + state = self.state[p] + + # Initialize master weights in state if not present + if len(state) == 0: + state["step"] = 0 + state["exp_avg"] = torch.zeros_like(grad, dtype=torch.float32) + state["exp_avg_sq"] = torch.zeros_like(grad, dtype=torch.float32) + state["master_param"] = p.detach().clone().float() + + # Always use the persistent master weight for updates + p_fp32 = state["master_param"] + + exp_avg: torch.Tensor = state["exp_avg"] + exp_avg_sq: torch.Tensor = state["exp_avg_sq"] + state["step"] += 1 + step: int = state["step"] + + # Decoupled weight decay (AdamW) + if weight_decay != 0.0: + p_fp32.add_(p_fp32, alpha=-lr * weight_decay) + + # Adam moments + exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) + + bias_correction1 = 1.0 - beta1**step + bias_correction2 = 1.0 - beta2**step + + denom = (exp_avg_sq.sqrt() / (bias_correction2**0.5)).add_(eps) + step_size = lr / bias_correction1 + + p_fp32.addcdiv_(exp_avg, denom, value=-step_size) + + # Write back to original dtype + p.copy_(p_fp32.to(dtype=p.dtype)) + + return loss diff --git a/primus/backends/diffusion/registry.py b/primus/backends/diffusion/registry.py new file mode 100644 index 000000000..b0a728827 --- /dev/null +++ b/primus/backends/diffusion/registry.py @@ -0,0 +1,66 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Explicit factory registry for diffusion model, dataset, and trainer builders.""" + +from __future__ import annotations + +from typing import Any, Callable, Dict, Tuple + + +def _build_wan_model(model_config: dict): + from primus.backends.diffusion.models.registrations.wan import build_wan_model + + return build_wan_model(model_config) + + +def _build_wan_dataset(dataset_config: dict): + from primus.backends.diffusion.data.registrations.wan import build_wan_dataset + + return build_wan_dataset(dataset_config) + + +def _build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict): + from primus.backends.diffusion.trainers.fsdp2 import build_fsdp2_trainer + + return build_fsdp2_trainer( + model=model, + dataset=dataset, + processor=processor, + trainer_args=trainer_args, + ) + + +MODEL_BUILDERS: Dict[str, Callable[[dict], Any]] = { + "wan": _build_wan_model, +} +DATASET_BUILDERS: Dict[str, Callable[[dict], Tuple[Any, Any]]] = { + "wan": _build_wan_dataset, +} +TRAINER_BUILDERS: Dict[str, Callable[..., Any]] = { + "fsdp2": _build_fsdp2_trainer, +} + + +def get_model_builder(name: str) -> Callable[[dict], Any]: + try: + return MODEL_BUILDERS[name] + except KeyError: + raise KeyError(f"Unknown model name: {name}") + + +def get_dataset_builder(name: str) -> Callable[[dict], Tuple[Any, Any]]: + try: + return DATASET_BUILDERS[name] + except KeyError: + raise KeyError(f"Unknown dataset name: {name}") + + +def get_trainer_builder(name: str) -> Callable[..., Any]: + try: + return TRAINER_BUILDERS[name] + except KeyError: + raise KeyError(f"Unknown trainer name: {name}") diff --git a/primus/backends/diffusion/schedulers/flow_match.py b/primus/backends/diffusion/schedulers/flow_match.py new file mode 100644 index 000000000..ec5519574 --- /dev/null +++ b/primus/backends/diffusion/schedulers/flow_match.py @@ -0,0 +1,130 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import math + +import torch + + +class FlowMatchScheduler: + def __init__( + self, + num_inference_steps=100, + num_train_timesteps=1000, + shift=3.0, + sigma_max=1.0, + sigma_min=0.003 / 1.002, + inverse_timesteps=False, + extra_one_step=False, + reverse_sigmas=False, + exponential_shift=False, + exponential_shift_mu=None, + shift_terminal=None, + ): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.exponential_shift = exponential_shift + self.exponential_shift_mu = exponential_shift_mu + self.shift_terminal = shift_terminal + self.set_timesteps(num_inference_steps) + + def set_timesteps( + self, + num_inference_steps=100, + denoising_strength=1.0, + training=False, + shift=None, + dynamic_shift_len=None, + ): + if shift is not None: + self.shift = shift + sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + if self.exponential_shift: + mu = ( + self.calculate_shift(dynamic_shift_len) + if dynamic_shift_len is not None + else self.exponential_shift_mu + ) + self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1)) + else: + self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) + if self.shift_terminal is not None: + one_minus_z = 1 - self.sigmas + scale_factor = one_minus_z[-1] / (1 - self.shift_terminal) + self.sigmas = 1 - (one_minus_z / scale_factor) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + self.training = True + else: + self.training = False + + def step(self, model_output, timestep, sample, to_final=False, **kwargs): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + if to_final or timestep_id + 1 >= len(self.timesteps): + sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 + else: + sigma_ = self.sigmas[timestep_id + 1] + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample + + def return_to_timestep(self, timestep, sample, sample_stablized): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + model_output = (sample - sample_stablized) / sigma + return model_output + + def add_noise(self, original_samples, noise, timestep): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + sample = (1 - sigma) * original_samples + sigma * noise + return sample + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + def training_weight(self, timestep): + timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs()) + weights = self.linear_timesteps_weights[timestep_id] + return weights + + def calculate_shift( + self, + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 8192, + base_shift: float = 0.5, + max_shift: float = 0.9, + ): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu diff --git a/primus/backends/diffusion/trainers/__init__.py b/primus/backends/diffusion/trainers/__init__.py new file mode 100644 index 000000000..d4df6355e --- /dev/null +++ b/primus/backends/diffusion/trainers/__init__.py @@ -0,0 +1,11 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Trainer registrations for the Primus Wan backend.""" + +from .fsdp2 import build_fsdp2_trainer + +__all__ = ["build_fsdp2_trainer"] diff --git a/primus/backends/diffusion/trainers/base.py b/primus/backends/diffusion/trainers/base.py new file mode 100644 index 000000000..c5d5227fd --- /dev/null +++ b/primus/backends/diffusion/trainers/base.py @@ -0,0 +1,551 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Base trainer with shared logic for Wan PyTorch trainers. + +This module holds functionality shared across FSDP-style Wan trainers: +config parsing, optimizer creation, LR scheduling, training loop, logging, +and W&B integration. Concrete trainers (e.g. FSDP2Trainer) subclass it. +""" + +from __future__ import annotations + +import math +import os +import time +from contextlib import contextmanager + +import torch + +from primus.backends.diffusion.optim.adamw_fp32_state import AdamWFP32State +from primus.backends.diffusion.schedulers.flow_match import FlowMatchScheduler +from primus.backends.diffusion.utils.log import logger +from primus.backends.diffusion.utils.train_utils import ( + get_memory, + resolve_dtype, + set_seed, +) + +try: + import wandb +except ImportError: + wandb = None + + +def create_lr_scheduler(optimizer, scheduler_type, warmup_steps, total_steps): + """ + Create LR scheduler with warmup support. + + Supports: constant, constant_with_warmup, linear, cosine, polynomial. + Shared between FSDP and FSDP2 trainers for consistency. + """ + if total_steps <= 0: + return torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: 1.0) + + def linear_warmup(step): + if warmup_steps == 0: + return 1.0 + return min(1.0, float(step) / float(max(1, warmup_steps))) + + def linear_decay(step): + if step <= warmup_steps: + return linear_warmup(step) + progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps)) + return max(0.0, 1.0 - progress) + + def cosine_decay(step): + if step <= warmup_steps: + return linear_warmup(step) + progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps)) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + def constant_with_warmup(step): + return linear_warmup(step) + + def polynomial_decay(step, power=1.0): + if step <= warmup_steps: + return linear_warmup(step) + progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps)) + return max(0.0, (1.0 - progress) ** power) + + scheduler_type = (scheduler_type or "constant").lower() + lambdas = { + "constant": lambda step: 1.0, + "constant_with_warmup": constant_with_warmup, + "linear": linear_decay, + "cosine": cosine_decay, + "cosine_with_restarts": cosine_decay, + "polynomial": polynomial_decay, + } + lr_lambda = lambdas.get(scheduler_type) + if lr_lambda is None: + logger.warning(f"Unknown lr_scheduler_type={scheduler_type}, falling back to constant.") + + def lr_lambda(step): + return 1.0 + + return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) + + +class BaseWanTrainer: + """ + Shared base class for Wan PyTorch FSDP-style trainers. + + Subclasses must implement: + - _apply_parallelism(): set up distributed wrapping / sharding + - save_model(): save final model + + Subclasses may override: + - _grad_sync_context(is_update_step): gradient sync during accumulation + - _clip_grad_norm(): gradient clipping strategy + - _save_checkpoint(): periodic checkpoint saving + """ + + def __init__( + self, + model: torch.nn.Module, + args: dict, + train_dataset, + data_collator, + processing_class, + rank: int, + world_size: int, + local_rank: int, + ): + self.model = model + self.args = args + self.rank = rank + self.world_size = world_size + self.local_rank = local_rank + self.device = torch.device(f"cuda:{local_rank}") + + # --- Config extraction --- + self.output_dir = self.args.get("output_dir", "./output") + self.logging_steps = int(self.args.get("logging_steps", 1)) + self.save_steps = int(self.args.get("save_steps", 0)) + self.max_steps = int(self.args.get("max_steps", -1) if self.args.get("max_steps") is not None else -1) + self.grad_accum_steps = int(self.args.get("gradient_accumulation_steps", 1)) + self.max_grad_norm = float(self.args.get("max_grad_norm", 1.0)) + self.num_train_epochs = int(self.args.get("num_train_epochs", 1)) + + if self.rank == 0: + os.makedirs(self.output_dir, exist_ok=True) + + # --- Seeding --- + seed = self.args.get("seed") + if seed is not None: + set_seed(int(seed)) + if os.environ.get("FIXED_SEED"): + set_seed(int(os.environ["FIXED_SEED"])) + + # --- Gradient Checkpointing --- + if self.args.get("gradient_checkpointing", False): + if hasattr(self.model, "gradient_checkpointing_enable"): + self.model.gradient_checkpointing_enable() + elif hasattr(self.model, "dit") and hasattr(self.model.dit, "gradient_checkpointing"): + self.model.dit.gradient_checkpointing = True + if self.rank == 0: + logger.info("Gradient checkpointing enabled") + + # --- W&B --- + self._setup_wandb() + + # --- Parallelism (subclass hook) --- + # Subclass sets self.sp_group (Ulysses SP group) if SP is enabled. + self.sp_group = None + self._apply_parallelism() + + # --- DataLoader --- + self.train_dataset = train_dataset + self.processing_class = processing_class + self.data_collator = data_collator + + # When SP is enabled, all ranks in the same SP group process the same sample. + # DistributedSampler should use DP-only rank/size so SP peers get identical data. + self.sp_size = 1 + dp_world_size = world_size + dp_rank = rank + if self.sp_group is not None: + import torch.distributed as dist + + self.sp_size = dist.get_world_size(self.sp_group) + dp_world_size = world_size // self.sp_size + dp_rank = rank // self.sp_size + + self.data_parallel_world_size = dp_world_size + self.per_device_train_batch_size = int(self.args.get("per_device_train_batch_size", 1)) + + self.sampler = torch.utils.data.distributed.DistributedSampler( + train_dataset, + num_replicas=dp_world_size, + rank=dp_rank, + shuffle=self.args.get("shuffle", True), + ) + + num_workers = int(self.args.get("dataloader_num_workers", 4) or 0) + self.dataloader = torch.utils.data.DataLoader( + train_dataset, + batch_size=self.per_device_train_batch_size, + sampler=self.sampler, + num_workers=num_workers, + collate_fn=data_collator, + pin_memory=True, + persistent_workers=num_workers > 0, + prefetch_factor=2 if num_workers > 0 else None, + ) + + # --- Optimizer --- + self.optimizer = self._create_optimizer() + + # --- LR Scheduler --- + steps_per_epoch = math.ceil(len(self.dataloader) / max(1, self.grad_accum_steps)) + self.total_steps = self.max_steps if self.max_steps > 0 else self.num_train_epochs * steps_per_epoch + self.lr_scheduler = create_lr_scheduler( + self.optimizer, + self.args.get("lr_scheduler_type", "constant"), + int(self.args.get("warmup_steps", 0)), + self.total_steps, + ) + + # --- Diffusion Scheduler (configurable from YAML) --- + scheduler_cfg = self.args.get("flow_match_scheduler", {}) or {} + self.scheduler = FlowMatchScheduler( + shift=float(scheduler_cfg.get("shift", 5)), + sigma_min=float(scheduler_cfg.get("sigma_min", 0.0)), + extra_one_step=bool(scheduler_cfg.get("extra_one_step", True)), + ) + self.scheduler.set_timesteps( + int(scheduler_cfg.get("num_train_timesteps", 1000)), + training=True, + ) + + self.global_step = 0 + + # ------------------------------------------------------------------ # + # Subclass hooks # + # ------------------------------------------------------------------ # + + def _apply_parallelism(self): + """Set up distributed parallelism. Called during __init__.""" + raise NotImplementedError + + @contextmanager + def _grad_sync_context(self, is_update_step: bool): + """Context manager for gradient sync control during accumulation.""" + yield + + def _clip_grad_norm(self) -> float: + """Clip gradient norm. Returns the total norm value.""" + if self.max_grad_norm > 0: + norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm) + return norm.item() if isinstance(norm, torch.Tensor) else float(norm) + return 0.0 + + def _save_checkpoint(self): + """Save checkpoint at save_steps intervals. Override for custom strategies.""" + + # ------------------------------------------------------------------ # + # Common methods # + # ------------------------------------------------------------------ # + + def _setup_wandb(self): + self.use_wandb = False + if self.rank != 0: + return + if str(self.args.get("report_to", "")).lower() != "wandb": + return + if self.args.get("use_wandb") is False: + return + if wandb is None: + logger.warning("W&B requested but wandb is not installed.") + return + + project = self.args.get("wandb_project") or os.environ.get("WANDB_PROJECT", "primus-diffusion") + run_name = self.args.get("wandb_name") or self.args.get("run_name") + wandb_dir = self.args.get("wandb_dir") or os.environ.get("WANDB_DIR") + wandb.init(project=project, name=run_name, dir=wandb_dir, config=self.args) + self.use_wandb = True + + def _resolve_dtype(self) -> torch.dtype: + return resolve_dtype(self.args) + + def _create_optimizer(self): + lr = float(self.args.get("learning_rate", 1e-4)) + wd = float(self.args.get("weight_decay", 0.01)) + betas = ( + float(self.args.get("adam_beta1", 0.9)), + float(self.args.get("adam_beta2", 0.999)), + ) + eps = float(self.args.get("adam_epsilon", 1e-8)) + + params = [p for p in self.model.parameters() if p.requires_grad] + optimizer_kwargs = { + "params": params, + "lr": lr, + "betas": betas, + "eps": eps, + "weight_decay": wd, + } + + optimizer = None + try: + optimizer = torch.optim.AdamW(**optimizer_kwargs, fused=True) + if self.rank == 0: + logger.info("Optimizer: torch.optim.AdamW(fused=True)") + except (TypeError, RuntimeError): + try: + optimizer = torch.optim.AdamW(**optimizer_kwargs, foreach=True) + if self.rank == 0: + logger.info("Optimizer: torch.optim.AdamW(foreach=True)") + except (TypeError, RuntimeError): + optimizer = torch.optim.AdamW(**optimizer_kwargs) + if self.rank == 0: + logger.info("Optimizer: torch.optim.AdamW(default)") + + if (self.args.get("bf16", False) or self.args.get("fp16", False)) and os.getenv( + "FP32_MASTER_WEIGHTS", "0" + ) == "1": + if self.rank == 0: + logger.info( + "FP32_MASTER_WEIGHTS=1: using AdamWFP32State (fp32 master weights + fp32 moments)." + ) + optimizer = AdamWFP32State( + optimizer.param_groups, + lr=lr, + betas=betas, + eps=eps, + weight_decay=wd, + ) + + return optimizer + + def compute_loss(self, batch): + """Prepare batch and compute training loss.""" + prepare_batch = getattr(self.processing_class, "prepare_batch", None) + if callable(prepare_batch): + batch = prepare_batch( + batch=batch, + device=self.device, + dtype=self._resolve_dtype(), + ) + # Ensure all tensors are on the correct device + for k, v in batch.items(): + if isinstance(v, torch.Tensor): + batch[k] = v.to(self.device, non_blocking=True) + + # Pass SP group so model can shard sequences across SP ranks + if self.sp_group is not None: + batch["sp_group"] = self.sp_group + + # Use explicit training entry point if available (GenAIModel interface) + forward_train = getattr(self.model, "forward_train", None) + if callable(forward_train): + outputs = forward_train(batch, scheduler=self.scheduler) + else: + outputs = self.model(batch, self.scheduler) + return outputs["loss"] + + def _infer_batch_size_from_tensors(self, value) -> int | None: + if isinstance(value, torch.Tensor): + return int(value.shape[0]) if value.ndim > 0 else 1 + if isinstance(value, dict): + for item in value.values(): + batch_size = self._infer_batch_size_from_tensors(item) + if batch_size is not None: + return batch_size + if isinstance(value, (list, tuple)): + for item in value: + batch_size = self._infer_batch_size_from_tensors(item) + if batch_size is not None: + return batch_size + return None + + def _infer_batch_size_from_sequences(self, value) -> int | None: + if isinstance(value, dict): + for item in value.values(): + batch_size = self._infer_batch_size_from_sequences(item) + if batch_size is not None: + return batch_size + return None + if isinstance(value, (list, tuple)): + if not value: + return 0 + first = value[0] + if isinstance(first, (dict, list, tuple, torch.Tensor)): + for item in value: + batch_size = self._infer_batch_size_from_sequences(item) + if batch_size is not None: + return batch_size + return None + return len(value) + return None + + def _infer_local_batch_size(self, batch) -> int: + tensor_batch_size = self._infer_batch_size_from_tensors(batch) + if tensor_batch_size is not None: + return tensor_batch_size + + sequence_batch_size = self._infer_batch_size_from_sequences(batch) + if sequence_batch_size is not None: + return sequence_batch_size + + return self.per_device_train_batch_size + + def _compute_samples_per_gpu_per_second( + self, + local_samples: int, + interval_seconds: float | None, + ) -> float | None: + if interval_seconds is None or interval_seconds <= 0 or local_samples <= 0 or self.world_size <= 0: + return None + + global_samples = float(local_samples) * float(self.data_parallel_world_size) + return global_samples / float(self.world_size) / float(interval_seconds) + + def _log_step( + self, + loss_value: float, + grad_norm: float = 0.0, + step_time: float | None = None, + elapsed: float | None = None, + eta_seconds: float | None = None, + throughput_samples_per_gpu_s: float | None = None, + ): + """Log training metrics. Format matches test regex expectations.""" + if self.rank != 0: + return + alloc, res, max_mem = get_memory() + lr = self.optimizer.param_groups[0]["lr"] + + # NOTE: The "step=... loss=... mem=.../...GB" line format is relied on by + # downstream log parsers; keep it stable when editing. + msg = ( + f"step={self.global_step} loss={loss_value:.4f} " + f"mem={alloc:.2f}/{res:.2f}GB peak_mem={max_mem:.2f}GB " + f"gnorm={grad_norm:.4f}" + ) + if step_time is not None: + msg += f" step_time={step_time:.2f}s" + if throughput_samples_per_gpu_s is not None: + msg += f" throughput={throughput_samples_per_gpu_s:.4f}samples/gpu/s" + if elapsed is not None: + msg += f" elapsed={elapsed / 60:.2f}m" + if eta_seconds is not None: + msg += f" eta={eta_seconds / 60:.2f}m" + logger.info(msg) + + if self.use_wandb: + payload = { + "train/loss": loss_value, + "train/step": self.global_step, + "train/grad_norm": grad_norm, + "train/lr": lr, + "mem/allocated_gb": alloc, + "mem/reserved_gb": res, + "mem/max_alloc_gb": max_mem, + } + if step_time is not None: + payload["time/step_s"] = step_time + if throughput_samples_per_gpu_s is not None: + payload["perf/samples_per_gpu_s"] = throughput_samples_per_gpu_s + if elapsed is not None: + payload["time/elapsed_s"] = elapsed + if eta_seconds is not None: + payload["time/eta_s"] = eta_seconds + wandb.log(payload, step=self.global_step) + + def train(self): + if self.rank == 0: + logger.info("Starting training...") + + # Ensure frozen state (idempotent) + core = getattr(self.model, "module", self.model) + if hasattr(core, "freeze_except"): + core.freeze_except() + + self.model.train() + self.optimizer.zero_grad(set_to_none=True) + torch.cuda.reset_peak_memory_stats() + + start_time = time.time() + last_log_time = start_time + local_samples_in_update = 0 + local_samples_since_log = 0 + update_steps_since_log = 0 + update_loss_sum = 0.0 + update_loss_count = 0 + + for epoch in range(self.num_train_epochs): + self.sampler.set_epoch(epoch) + + for batch_idx, batch in enumerate(self.dataloader): + is_update_step = ((batch_idx + 1) % max(1, self.grad_accum_steps)) == 0 + local_samples_in_update += self._infer_local_batch_size(batch) + + with self._grad_sync_context(is_update_step): + raw_loss = self.compute_loss(batch) + update_loss_sum += raw_loss.detach().float().item() + update_loss_count += 1 + loss = raw_loss / max(1, self.grad_accum_steps) + loss.backward() + + if is_update_step: + loss_val = update_loss_sum / max(1, update_loss_count) + update_loss_sum = 0.0 + update_loss_count = 0 + grad_norm = self._clip_grad_norm() + + self.optimizer.step() + self.lr_scheduler.step() + self.optimizer.zero_grad(set_to_none=True) + self.global_step += 1 + update_steps_since_log += 1 + local_samples_since_log += local_samples_in_update + local_samples_in_update = 0 + + # Logging + if self.global_step % self.logging_steps == 0: + now = time.time() + log_interval = now - last_log_time + step_time = log_interval / max(1, update_steps_since_log) + last_log_time = now + elapsed = now - start_time + steps_left = max(0, self.total_steps - self.global_step) + eta_seconds = step_time * steps_left + throughput_samples_per_gpu_s = self._compute_samples_per_gpu_per_second( + local_samples=local_samples_since_log, + interval_seconds=log_interval, + ) + self._log_step( + loss_val, + grad_norm=grad_norm, + step_time=step_time, + elapsed=elapsed, + eta_seconds=eta_seconds, + throughput_samples_per_gpu_s=throughput_samples_per_gpu_s, + ) + local_samples_since_log = 0 + update_steps_since_log = 0 + + # Periodic save + if self.save_steps > 0 and self.global_step % self.save_steps == 0: + self._save_checkpoint() + + # Early termination + if self.max_steps > 0 and self.global_step >= self.max_steps: + return + + if self.max_steps > 0 and self.global_step >= self.max_steps: + break + + if self.rank == 0: + elapsed = time.time() - start_time + logger.info(f"Training finished in {elapsed / 60:.2f} min") + + def save_model(self): + """Save final model. Override in subclass.""" + raise NotImplementedError diff --git a/primus/backends/diffusion/trainers/fsdp2.py b/primus/backends/diffusion/trainers/fsdp2.py new file mode 100644 index 000000000..1fa79b177 --- /dev/null +++ b/primus/backends/diffusion/trainers/fsdp2.py @@ -0,0 +1,395 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Wan PyTorch FSDP2 trainer +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager +from importlib import import_module + +import torch +from safetensors.torch import load_file as safe_load_file +from safetensors.torch import save_file as safe_save_file +from torch.distributed._composable.fsdp import MixedPrecisionPolicy, fully_shard +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, +) + +from primus.backends.diffusion.distributed import ( + create_device_mesh, + load_checkpoint_dtcp, + save_checkpoint_dtcp, + setup_distributed, +) +from primus.backends.diffusion.utils.log import logger + +from .base import BaseWanTrainer + + +class FSDP2Trainer(BaseWanTrainer): + def __init__( + self, + model: torch.nn.Module, + args: dict, + train_dataset, + data_collator, + processing_class, + rank: int, + world_size: int, + local_rank: int, + ): + super().__init__( + model=model, + args=args, + train_dataset=train_dataset, + data_collator=data_collator, + processing_class=processing_class, + rank=rank, + world_size=world_size, + local_rank=local_rank, + ) + + # FSDP2-specific: checkpoint strategy + # - "dit_only": save `dit_model.safetensors` (default for T2V training) + # - "dtcp_full": save model + optimizer via DTCP + # - "dtcp_model_only": save only model via DTCP + # - "dtcp_trainable": DTCP save only trainable params + optimizer + self.save_strategy = str(self.args.get("save_strategy", "dit_only")).lower() + + # Checkpoint loading (after optimizer & scheduler are fully initialized) + resume_from = self.args.get("resume_from_checkpoint") + if resume_from: + self._load_checkpoint(resume_from) + + # ------------------------------------------------------------------ # + # Parallelism # + # ------------------------------------------------------------------ # + + def _apply_parallelism(self): + """Set up FSDP2 composable sharding, optionally with Ulysses SP.""" + sp_size = int(self.args.get("sp_size", 1)) + dp_replicate = int(self.args.get("dp_replicate", 1)) + + self.mesh = create_device_mesh(self.world_size, sp_size=sp_size, dp_replicate=dp_replicate) + self.sp_group = self.mesh.get_group("ulysses") if (self.mesh is not None and sp_size > 1) else None + self.model.to(self.device) + + # Freeze non-trainable params BEFORE FSDP and optimizer creation + if hasattr(self.model, "freeze_except"): + self.model.freeze_except() + if self.rank == 0: + logger.info("FSDP2: Applied freeze_except (frozen non-trainable params)") + + self._apply_fsdp2() + + def _apply_fsdp2(self): + """Apply torch.distributed._composable.fsdp.fully_shard to the model.""" + mp_dtype = self._resolve_dtype() + + # ---- Mixed-precision policy (aligned with DeepSpeed bf16) ---- + # 1. Pre-cast params to bf16 BEFORE FSDP wrapping (bf16 storage) + # 2. reduce_dtype = bf16 (gradient reduce matches DeepSpeed bf16 all-reduce) + # 3. AdamWFP32State maintains fp32 master weights and writes back to bf16 + if mp_dtype != torch.float32: + mp_policy = MixedPrecisionPolicy( + param_dtype=mp_dtype, + reduce_dtype=mp_dtype, # bf16 reduce (matches DeepSpeed bf16) + ) + else: + mp_policy = None + + # When SP is enabled, FSDP2 shards across dp_shard_sp (DP + SP combined). + # This reduces per-rank parameter memory. + fsdp_mesh = self.mesh + if self.mesh is not None and self.sp_group is not None: + try: + fsdp_mesh = self.mesh["dp_shard_sp"] + except KeyError: + pass # fallback to full mesh + + wrap_target = str(self.args.get("fsdp2_wrap_target", "") or "").strip() + wrap_root = self._get_module_by_path(self.model, wrap_target) if wrap_target else self.model + + # Pre-cast to bf16 before FSDP wrapping so parameters are stored in bf16, + # matching DiffSynth/DeepSpeed bf16 behavior. + if mp_dtype != torch.float32: + wrap_root.to(dtype=mp_dtype) + if self.rank == 0: + logger.info(f"FSDP2: pre-cast '{wrap_target or ''}' to {mp_dtype} (DiffSynth-aligned)") + + if self.world_size == 1: + if self.rank == 0: + logger.info("FSDP2: world_size=1; skipping composable FSDP wrapping.") + return + + reshard_after_forward = bool(self.args.get("fsdp2_reshard_after_forward", True)) + + # Wrap transformer blocks first for optimal memory management + layer_cls_spec = self.args.get("fsdp_transformer_layer_cls_to_wrap") + if layer_cls_spec: + if isinstance(layer_cls_spec, str): + layer_cls_items = [x.strip() for x in layer_cls_spec.split(",") if x.strip()] + else: + layer_cls_items = [str(x).strip() for x in layer_cls_spec if str(x).strip()] + + cls_objs = set() + cls_names = set() + for item in layer_cls_items: + if "." in item: + mod_path, cls_name = item.rsplit(".", 1) + cls_objs.add(getattr(import_module(mod_path), cls_name)) + else: + cls_names.add(item) + + wrapped_count = 0 + seen = set() + for _, module in wrap_root.named_modules(): + if id(module) in seen: + continue + seen.add(id(module)) + if module is wrap_root: + continue + if (cls_objs and isinstance(module, tuple(cls_objs))) or ( + cls_names and module.__class__.__name__ in cls_names + ): + fully_shard( + module, + mesh=fsdp_mesh, + reshard_after_forward=reshard_after_forward, + mp_policy=mp_policy, + ) + wrapped_count += 1 + + if self.rank == 0: + logger.info(f"FSDP2: wrapped {wrapped_count} submodules under '{wrap_target or ''}'") + + fully_shard( + wrap_root, + mesh=fsdp_mesh, + reshard_after_forward=reshard_after_forward, + mp_policy=mp_policy, + ) + if self.rank == 0: + logger.info(f"FSDP2: applied fully_shard to '{wrap_target or ''}' with mp={mp_dtype}") + + @staticmethod + def _get_module_by_path(root: torch.nn.Module, path: str) -> torch.nn.Module: + """Resolve a dot-separated attribute path on a module.""" + cur = root + if not path: + return cur + for part in path.split("."): + if not hasattr(cur, part): + raise ValueError(f"fsdp2_wrap_target='{path}' is invalid: missing attribute '{part}'") + cur = getattr(cur, part) + if not isinstance(cur, torch.nn.Module): + raise ValueError(f"fsdp2_wrap_target='{path}' did not resolve to a torch.nn.Module") + return cur + + # ------------------------------------------------------------------ # + # Gradient sync # + # ------------------------------------------------------------------ # + + def _set_requires_gradient_sync(self, enabled: bool) -> None: + """ + Best-effort gradient sync control for composable FSDP2. + `fully_shard()` turns modules into FSDPM which exposes + `set_requires_gradient_sync`. For grad accumulation, we disable + sync on non-update micro-steps. + """ + seen = set() + for _, m in self.model.named_modules(): + if id(m) in seen: + continue + seen.add(id(m)) + setter = getattr(m, "set_requires_gradient_sync", None) + if callable(setter): + setter(bool(enabled)) + + @contextmanager + def _grad_sync_context(self, is_update_step: bool): + if self.world_size > 1 and self.grad_accum_steps > 1: + self._set_requires_gradient_sync(is_update_step) + yield + if is_update_step and self.world_size > 1 and self.grad_accum_steps > 1: + self._set_requires_gradient_sync(True) + + # ------------------------------------------------------------------ # + # Checkpointing # + # ------------------------------------------------------------------ # + + def _dtcp_save_args(self): + """Compute DTCP save parameters based on save_strategy.""" + if self.save_strategy == "dtcp_model_only": + return None, None + if self.save_strategy == "dtcp_trainable": + opts = StateDictOptions(full_state_dict=False, ignore_frozen_params=True) + return self.optimizer, opts + # dtcp_full (default non-dit strategy) + return self.optimizer, None + + def _save_dtcp(self, path): + """Save checkpoint via Distributed Tensor Checkpointing.""" + optimizer, opts = self._dtcp_save_args() + kwargs = {} + if opts is not None: + kwargs["model_state_options"] = opts + kwargs["optim_state_options"] = opts + save_checkpoint_dtcp( + self.model, + optimizer, + path, + epoch=0, + step=self.global_step, + additional_data={ + "lr_scheduler": (self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None) + }, + **kwargs, + ) + + def _save_checkpoint(self): + path = os.path.join(self.output_dir, f"checkpoint-{self.global_step}") + if self.save_strategy == "dit_only": + self._save_dit(os.path.join(path, "dit_model.safetensors")) + else: + self._save_dtcp(path) + + def _load_checkpoint(self, path): + if self.rank == 0: + logger.info(f"Loading checkpoint from {path}") + + if self.save_strategy == "dit_only": + core_model = self.model + if not hasattr(core_model, "dit"): + raise ValueError("save_strategy=dit_only requires model.dit to exist for checkpoint loading") + candidate = path + if os.path.isdir(path): + candidate = os.path.join(path, "dit_model.safetensors") + if not os.path.exists(candidate): + raise FileNotFoundError(f"dit_only checkpoint not found at: {candidate}") + state = safe_load_file(candidate) + missing, unexpected = core_model.dit.load_state_dict(state, strict=False) + if self.rank == 0: + logger.info( + f"Loaded DiT weights from {candidate}. " + f"Missing keys: {len(missing)}, " + f"Unexpected keys: {len(unexpected)}" + ) + return + + if self.save_strategy == "dtcp_trainable": + opts = StateDictOptions(full_state_dict=False, ignore_frozen_params=True) + meta = load_checkpoint_dtcp( + self.model, + self.optimizer, + path, + model_state_options=opts, + optim_state_options=opts, + ) + elif self.save_strategy == "dtcp_model_only": + meta = load_checkpoint_dtcp(self.model, None, path) + else: + meta = load_checkpoint_dtcp(self.model, self.optimizer, path) + + if "step" in meta: + self.global_step = meta["step"] + if self.rank == 0: + logger.info(f"Resumed from step {self.global_step}") + if self.lr_scheduler is not None and isinstance(meta, dict) and meta.get("lr_scheduler") is not None: + try: + self.lr_scheduler.load_state_dict(meta["lr_scheduler"]) + if self.rank == 0: + logger.info("Resumed lr_scheduler state from checkpoint meta") + except Exception as exc: + if self.rank == 0: + logger.warning(f"Failed to restore lr_scheduler state: {exc}") + + def _save_dit(self, save_path: str) -> None: + """ + Save `dit` weights as a single safetensors file. + + Notes: + - On composable FSDP (world_size>1), ``get_model_state_dict`` with + ``full_state_dict=True`` returns the full dict **only on rank 0**; + other ranks receive an empty dict. This is expected – do NOT + early-return on the empty-dict check, or non-rank-0 processes will + race ahead and desynchronize from rank 0 (which still needs to + write to disk). + - A ``dist.barrier()`` at the end keeps all ranks in lockstep so + that back-to-back saves (e.g. periodic checkpoint + final save) + never overlap. + """ + import torch.distributed as dist + + core_model = self.model + if not hasattr(core_model, "dit"): + logger.warning("save_model: model has no `dit` attribute; skipping save.") + return + + if self.world_size > 1: + # FSDP-wrapped: use get_model_state_dict to unshard DTensors + full_state = get_model_state_dict( + core_model, + options=StateDictOptions(full_state_dict=True, cpu_offload=True), + ) + dit_state_dict = {k[len("dit.") :]: v for k, v in full_state.items() if k.startswith("dit.")} + else: + # Single GPU (no FSDP wrapping): direct state_dict + dit_state_dict = core_model.dit.state_dict() + if not dit_state_dict: + logger.warning("save_model: DiT state dict is empty; skipping save.") + return + + # Only rank 0 has the full dict in multi-GPU; write from rank 0 only. + if self.rank == 0: + if not dit_state_dict: + logger.warning("save_model: DiT state dict is empty on rank 0; skipping save.") + else: + os.makedirs(os.path.dirname(save_path), exist_ok=True) + safe_save_file(dit_state_dict, save_path) + logger.info(f"Saved DiT weights to {save_path}") + + if hasattr(core_model, "config"): + try: + core_model.config.save_pretrained(os.path.dirname(save_path)) + except Exception as exc: + logger.warning(f"save_model: failed to save config: {exc}") + + # Barrier: ensure all ranks wait for rank 0 to finish writing before + # any rank proceeds to the next operation (e.g. another save or exit). + if self.world_size > 1: + dist.barrier() + + def save_model(self): + """Save final model using the configured strategy.""" + if self.save_strategy == "dit_only": + save_path = os.path.join(self.output_dir, "dit_model.safetensors") + self._save_dit(save_path) + return + + path = os.path.join(self.output_dir, "checkpoint-final") + if self.rank == 0: + logger.info(f"Saving final checkpoint to {path} (strategy={self.save_strategy})") + self._save_dtcp(path) + + +def build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict): + rank, world_size, local_rank = setup_distributed() + return FSDP2Trainer( + model=model, + args=trainer_args, + train_dataset=dataset, + data_collator=dataset.get_collator(), + processing_class=processor, + rank=rank, + world_size=world_size, + local_rank=local_rank, + ) diff --git a/primus/backends/diffusion/utils/__init__.py b/primus/backends/diffusion/utils/__init__.py new file mode 100644 index 000000000..3d7aa5fe7 --- /dev/null +++ b/primus/backends/diffusion/utils/__init__.py @@ -0,0 +1,10 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Utility package for the Primus Wan backend. + +Heavy vision helpers are imported lazily by the qwen_vl_utils video path. +""" diff --git a/primus/backends/diffusion/utils/data_utils.py b/primus/backends/diffusion/utils/data_utils.py new file mode 100644 index 000000000..4e78671da --- /dev/null +++ b/primus/backends/diffusion/utils/data_utils.py @@ -0,0 +1,56 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + + +import math + +FRAME_FACTOR = 2 +FPS = 2.0 +FPS_MIN_FRAMES = 4 +FPS_MAX_FRAMES = 768 + + +def smart_nframes( + total_frames: int, + video_fps: int | float, + fps: int | float, +) -> int: + """calculate the number of frames for video used for model inputs. + + Args: + total_frames (int): the original total number of frames of the video. + video_fps (int | float): the original fps of the video. + fps (int | float): the target sampling fps for model inputs. + + Raises: + ValueError: nframes should in interval [FRAME_FACTOR, total_frames]. + + Returns: + int: the number of frames for video used for model inputs. + """ + min_frames = ceil_by_factor(FPS_MIN_FRAMES, FRAME_FACTOR) + max_frames = floor_by_factor(min(FPS_MAX_FRAMES, total_frames), FRAME_FACTOR) + nframes = total_frames / video_fps * fps + nframes = min(min(max(nframes, min_frames), max_frames), total_frames) + nframes = floor_by_factor(nframes, FRAME_FACTOR) + if not (FRAME_FACTOR <= nframes and nframes <= total_frames): + raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.") + return nframes + + +def round_by_factor(number: int, factor: int) -> int: + """Returns the closest integer to 'number' that is divisible by 'factor'.""" + return round(number / factor) * factor + + +def ceil_by_factor(number: int, factor: int) -> int: + """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'.""" + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int, factor: int) -> int: + """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'.""" + return math.floor(number / factor) * factor diff --git a/primus/backends/diffusion/utils/log.py b/primus/backends/diffusion/utils/log.py new file mode 100644 index 000000000..2193ddebc --- /dev/null +++ b/primus/backends/diffusion/utils/log.py @@ -0,0 +1,42 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Logger proxy for the diffusion backend. + +Primus configures loguru sinks whose format references bound ``extra`` fields +(``rank``, ``world_size``, ``user``, ``team``, ``module_name``, ``node_ip``). +Those fields are only present on the Primus-bound logger created in +``primus.core.utils.logger.setup_logger``. Emitting from the raw, unbound +global loguru logger (``from loguru import logger``) produces records without +those extras, which makes every sink raise "Logging error in Handler". + +This module exposes a ``logger`` proxy that forwards attribute access to the +live Primus-bound logger, so diffusion code keeps the familiar +``logger.info(...)`` style while inheriting the bound extras. Before the Primus +logger is initialized (e.g. standalone tooling), it falls back to the raw +loguru logger, which at that point still uses loguru's default sink. +""" + +from __future__ import annotations + +from primus.core.utils import logger as _primus_logger_module + + +class _LoggerProxy: + @staticmethod + def _resolve(): + bound = getattr(_primus_logger_module, "_logger", None) + if bound is not None: + return bound + from loguru import logger as _raw_logger + + return _raw_logger + + def __getattr__(self, name): + return getattr(self._resolve(), name) + + +logger = _LoggerProxy() diff --git a/primus/backends/diffusion/utils/train_utils.py b/primus/backends/diffusion/utils/train_utils.py new file mode 100644 index 000000000..315ee9e67 --- /dev/null +++ b/primus/backends/diffusion/utils/train_utils.py @@ -0,0 +1,201 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + + +""" +utils for training +""" + +import hashlib +import os +import random +from contextlib import contextmanager + +import numpy as np +import torch +from safetensors import safe_open + + +def resolve_dtype(config_or_args) -> torch.dtype: + """ + Resolve mixed-precision dtype from either: + - a dict-like config (e.g. trainer_args) + - an object with attributes (e.g. HF TrainingArguments) + + Priority: + bf16 -> fp16 -> fp32 + """ + # Dict-like + if isinstance(config_or_args, dict): + if config_or_args.get("bf16", False): + return torch.bfloat16 + if config_or_args.get("fp16", False): + return torch.float16 + return torch.float32 + + # Attribute-like (HF TrainingArguments etc.) + bf16 = bool(getattr(config_or_args, "bf16", False)) + fp16 = bool(getattr(config_or_args, "fp16", False)) + if bf16: + return torch.bfloat16 + if fp16: + return torch.float16 + return torch.float32 + + +def set_seed(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def count_parameters(model): + total = sum(p.numel() for p in model.parameters()) + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + return total, trainable + + +def get_memory(unit=1e9): + torch.cuda.synchronize() + allocated = torch.cuda.memory_allocated() + reserved = torch.cuda.memory_reserved() + max_alloc = torch.cuda.max_memory_allocated() + return allocated / unit, reserved / unit, max_alloc / unit + + +def print_cuda_memory(prefix="", unit=1e9): + allocated = torch.cuda.memory_allocated() / unit + reserved = torch.cuda.memory_reserved() / unit + max_alloc = torch.cuda.max_memory_allocated() / unit + print( + f"{prefix} " + f"allocated={allocated:.2f}GB, " + f"reserved={reserved:.2f}GB, " + f"max_alloc={max_alloc:.2f}GB" + ) + + +@contextmanager +def init_weights_on_device(device=torch.device("meta"), include_buffers: bool = False): + + old_register_parameter = torch.nn.Module.register_parameter + if include_buffers: + old_register_buffer = torch.nn.Module.register_buffer + + def register_empty_parameter(module, name, param): + old_register_parameter(module, name, param) + if param is not None: + param_cls = type(module._parameters[name]) + kwargs = module._parameters[name].__dict__ + kwargs["requires_grad"] = param.requires_grad + module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) + + def register_empty_buffer(module, name, buffer, persistent=True): + old_register_buffer(module, name, buffer, persistent=persistent) + if buffer is not None: + module._buffers[name] = module._buffers[name].to(device) + + def patch_tensor_constructor(fn): + def wrapper(*args, **kwargs): + kwargs["device"] = device + return fn(*args, **kwargs) + + return wrapper + + if include_buffers: + tensor_constructors_to_patch = { + torch_function_name: getattr(torch, torch_function_name) + for torch_function_name in ["empty", "zeros", "ones", "full"] + } + else: + tensor_constructors_to_patch = {} + + try: + torch.nn.Module.register_parameter = register_empty_parameter + if include_buffers: + torch.nn.Module.register_buffer = register_empty_buffer + for torch_function_name in tensor_constructors_to_patch.keys(): + setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name))) + yield + finally: + torch.nn.Module.register_parameter = old_register_parameter + if include_buffers: + torch.nn.Module.register_buffer = old_register_buffer + for torch_function_name, old_torch_function in tensor_constructors_to_patch.items(): + setattr(torch, torch_function_name, old_torch_function) + + +def load_state_dict_from_folder(file_path, torch_dtype=None): + state_dict = {} + for file_name in os.listdir(file_path): + if "." in file_name and file_name.split(".")[-1] in ["safetensors", "bin", "ckpt", "pth", "pt"]: + state_dict.update(load_state_dict(os.path.join(file_path, file_name), torch_dtype=torch_dtype)) + return state_dict + + +def load_state_dict(file_path, torch_dtype=None, device="cpu"): + if file_path.endswith(".safetensors"): + return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device) + else: + return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device) + + +def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"): + state_dict = {} + with safe_open(file_path, framework="pt", device=str(device)) as f: + for k in f.keys(): + state_dict[k] = f.get_tensor(k) + if torch_dtype is not None: + state_dict[k] = state_dict[k].to(torch_dtype) + return state_dict + + +def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"): + state_dict = torch.load(file_path, map_location=device, weights_only=True) + if torch_dtype is not None: + for i in state_dict: + if isinstance(state_dict[i], torch.Tensor): + state_dict[i] = state_dict[i].to(torch_dtype) + return state_dict + + +def convert_state_dict_keys_to_single_str(state_dict, with_shape=True): + keys = [] + for key, value in state_dict.items(): + if isinstance(key, str): + if isinstance(value, torch.Tensor): + if with_shape: + shape = "_".join(map(str, list(value.shape))) + keys.append(key + ":" + shape) + else: + keys.append(key) + elif isinstance(value, dict): + keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape)) + keys.sort() + keys_str = ",".join(keys) + return keys_str + + +def split_state_dict_with_prefix(state_dict): + keys = sorted([key for key in state_dict if isinstance(key, str)]) + prefix_dict = {} + for key in keys: + prefix = key if "." not in key else key.split(".")[0] + if prefix not in prefix_dict: + prefix_dict[prefix] = [] + prefix_dict[prefix].append(key) + state_dicts = [] + for prefix, keys in prefix_dict.items(): + sub_state_dict = {key: state_dict[key] for key in keys} + state_dicts.append(sub_state_dict) + return state_dicts + + +def hash_state_dict_keys(state_dict, with_shape=True): + keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape) + keys_str = keys_str.encode(encoding="UTF-8") + return hashlib.md5(keys_str).hexdigest() diff --git a/primus/backends/diffusion/utils/vision_process.py b/primus/backends/diffusion/utils/vision_process.py new file mode 100644 index 000000000..2058179be --- /dev/null +++ b/primus/backends/diffusion/utils/vision_process.py @@ -0,0 +1,583 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +# Adapted from qwen-vl-utils +""" + +import base64 +import copy +import logging +import math +import os +import sys +import time +import warnings +from concurrent.futures import ThreadPoolExecutor +from functools import lru_cache +from io import BytesIO +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import requests +import torch +import torchvision +from packaging import version +from PIL import Image +from torchvision import io, transforms +from torchvision.transforms import InterpolationMode + +MAX_RATIO = 200 +SPATIAL_MERGE_SIZE = 2 +IMAGE_MIN_TOKEN_NUM = 4 +IMAGE_MAX_TOKEN_NUM = 16384 +VIDEO_MIN_TOKEN_NUM = 128 +VIDEO_MAX_TOKEN_NUM = 768 + +FPS = 2.0 +FRAME_FACTOR = 2 +FPS_MIN_FRAMES = 4 +FPS_MAX_FRAMES = 768 +MAX_NUM_WORKERS_FETCH_VIDEO = 8 +REMOTE_FETCH_TIMEOUT = (5, 30) + +MODEL_SEQ_LEN = int(float(os.environ.get("MODEL_SEQ_LEN", 128000))) +logger = logging.getLogger(__name__) + +VideoMetadata = Dict[str, Any] +VideoReaderOutput = Tuple[torch.Tensor, VideoMetadata, float] +FetchVideoOutput = Union[ + torch.Tensor, + Tuple[torch.Tensor, VideoMetadata], + Tuple[torch.Tensor, float], + Tuple[Tuple[torch.Tensor, VideoMetadata], float], +] + + +def strip_file_uri(path: str) -> str: + """Return a filesystem path for local file URIs.""" + return path[7:] if path.startswith("file://") else path + + +def round_by_factor(number: int, factor: int) -> int: + """Returns the closest integer to 'number' that is divisible by 'factor'.""" + return round(number / factor) * factor + + +def ceil_by_factor(number: int, factor: int) -> int: + """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'.""" + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int, factor: int) -> int: + """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'.""" + return math.floor(number / factor) * factor + + +def smart_resize( + height: int, width: int, factor: int, min_pixels: Optional[int] = None, max_pixels: Optional[int] = None +) -> Tuple[int, int]: + """ + Rescales the image so that the following conditions are met: + + 1. Both dimensions (height and width) are divisible by 'factor'. + 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. + 3. The aspect ratio of the image is maintained as closely as possible. + """ + max_pixels = max_pixels if max_pixels is not None else (IMAGE_MAX_TOKEN_NUM * factor**2) + min_pixels = min_pixels if min_pixels is not None else (IMAGE_MIN_TOKEN_NUM * factor**2) + assert max_pixels >= min_pixels, "The max_pixels of image must be greater than or equal to min_pixels." + if max(height, width) / min(height, width) > MAX_RATIO: + raise ValueError( + f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" + ) + h_bar = max(factor, round_by_factor(height, factor)) + w_bar = max(factor, round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = floor_by_factor(height / beta, factor) + w_bar = floor_by_factor(width / beta, factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = ceil_by_factor(height * beta, factor) + w_bar = ceil_by_factor(width * beta, factor) + return h_bar, w_bar + + +def to_rgb(pil_image: Image.Image) -> Image.Image: + if pil_image.mode == "RGBA": + white_background = Image.new("RGB", pil_image.size, (255, 255, 255)) + white_background.paste(pil_image, mask=pil_image.split()[3]) # Use alpha channel as mask + return white_background + else: + return pil_image.convert("RGB") + + +def fetch_image(ele: Dict[str, Union[str, Image.Image]], image_patch_size: int = 14) -> Image.Image: + if "image" in ele: + image = ele["image"] + else: + image = ele["image_url"] + + image_obj = None + patch_factor = int(image_patch_size * SPATIAL_MERGE_SIZE) + if isinstance(image, Image.Image): + image_obj = image + elif image.startswith("http://") or image.startswith("https://"): + with requests.get(image, stream=True, timeout=REMOTE_FETCH_TIMEOUT) as response: + response.raise_for_status() + with BytesIO(response.content) as bio: + image_obj = copy.deepcopy(Image.open(bio)) + elif image.startswith("file://"): + image_obj = Image.open(image[7:]) + elif image.startswith("data:image"): + if "base64," in image: + _, base64_data = image.split("base64,", 1) + data = base64.b64decode(base64_data) + with BytesIO(data) as bio: + image_obj = copy.deepcopy(Image.open(bio)) + else: + image_obj = Image.open(image) + if image_obj is None: + raise ValueError( + f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}" + ) + image = to_rgb(image_obj) + + ## resize + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=patch_factor, + ) + else: + width, height = image.size + min_pixels = ele.get("min_pixels", IMAGE_MIN_TOKEN_NUM * patch_factor**2) + max_pixels = ele.get("max_pixels", IMAGE_MAX_TOKEN_NUM * patch_factor**2) + resized_height, resized_width = smart_resize( + height, + width, + factor=patch_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + image = image.resize((resized_width, resized_height)) + return image + + +def smart_nframes( + ele: Dict[str, Any], + total_frames: int, + video_fps: Union[int, float], +) -> int: + """calculate the number of frames for video used for model inputs. + + Args: + ele (dict): a dict contains the configuration of video. + support either `fps` or `nframes`: + - nframes: the number of frames to extract for model inputs. + - fps: the fps to extract frames for model inputs. + - min_frames: the minimum number of frames of the video, only used when fps is provided. + - max_frames: the maximum number of frames of the video, only used when fps is provided. + total_frames (int): the original total number of frames of the video. + video_fps (int | float): the original fps of the video. + + Raises: + ValueError: nframes should in interval [FRAME_FACTOR, total_frames]. + + Returns: + int: the number of frames for video used for model inputs. + """ + assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`" + if "nframes" in ele: + nframes = round_by_factor(ele["nframes"], FRAME_FACTOR) + else: + fps = ele.get("fps", FPS) + min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR) + max_frames = floor_by_factor(ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR) + nframes = total_frames / video_fps * fps + if nframes > total_frames: + logger.warning(f"smart_nframes: nframes[{nframes}] > total_frames[{total_frames}]") + nframes = min(min(max(nframes, min_frames), max_frames), total_frames) + nframes = floor_by_factor(nframes, FRAME_FACTOR) + if not (FRAME_FACTOR <= nframes and nframes <= total_frames): + raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.") + return nframes + + +def _read_video_torchvision( + ele: Dict[str, Any], +) -> VideoReaderOutput: + """read video using torchvision.io.read_video + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + Tuple of video tensor (T, C, H, W), metadata, and sampled FPS. + """ + video_path = ele["video"] + if version.parse(torchvision.__version__) < version.parse("0.19.0"): + if "http://" in video_path or "https://" in video_path: + warnings.warn( + "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0." + ) + video_path = strip_file_uri(video_path) + st = time.time() + video, audio, info = io.read_video( + video_path, + start_pts=ele.get("video_start", 0.0), + end_pts=ele.get("video_end", None), + pts_unit="sec", + output_format="TCHW", + ) + total_frames, video_fps = video.size(0), info["video_fps"] + logger.info(f"torchvision: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(0, total_frames - 1, nframes).round().long() + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + video = video[idx] + + video_metadata = dict( + fps=video_fps, + frames_indices=idx, + total_num_frames=total_frames, + video_backend="torchvision", + ) + return video, video_metadata, sample_fps + + +def is_decord_available() -> bool: + import importlib.util + + return importlib.util.find_spec("decord") is not None + + +def calculate_video_frame_range( + ele: Dict[str, Any], + total_frames: int, + video_fps: float, +) -> Tuple[int, int, int]: + """ + Calculate the start and end frame indices based on the given time range. + + Args: + ele (dict): A dictionary containing optional 'video_start' and 'video_end' keys (in seconds). + total_frames (int): Total number of frames in the video. + video_fps (float): Frames per second of the video. + + Returns: + tuple: A tuple containing (start_frame, end_frame, frame_count). + + Raises: + ValueError: If input parameters are invalid or the time range is inconsistent. + """ + # Validate essential parameters + if video_fps <= 0: + raise ValueError("video_fps must be a positive number") + if total_frames <= 0: + raise ValueError("total_frames must be a positive integer") + + # Get start and end time in seconds + video_start = ele.get("video_start", None) + video_end = ele.get("video_end", None) + if video_start is None and video_end is None: + return 0, total_frames - 1, total_frames + + max_duration = total_frames / video_fps + # Process start frame + if video_start is not None: + video_start_clamped = max(0.0, min(video_start, max_duration)) + start_frame = math.ceil(video_start_clamped * video_fps) + else: + start_frame = 0 + # Process end frame + if video_end is not None: + video_end_clamped = max(0.0, min(video_end, max_duration)) + end_frame = math.floor(video_end_clamped * video_fps) + end_frame = min(end_frame, total_frames - 1) + else: + end_frame = total_frames - 1 + + # Validate frame order + if start_frame >= end_frame: + raise ValueError( + f"Invalid time range: Start frame {start_frame} (at {video_start_clamped if video_start is not None else 0}s) " + f"exceeds end frame {end_frame} (at {video_end_clamped if video_end is not None else max_duration}s). " + f"Video duration: {max_duration:.2f}s ({total_frames} frames @ {video_fps}fps)" + ) + + logger.info( + f"calculate video frame range: {start_frame=}, {end_frame=}, {total_frames=} from {video_start=}, {video_end=}, {video_fps=:.3f}" + ) + return start_frame, end_frame, end_frame - start_frame + 1 + + +def _read_video_decord( + ele: Dict[str, Any], +) -> VideoReaderOutput: + """read video using decord.VideoReader + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + Tuple of video tensor (T, C, H, W), metadata, and sampled FPS. + """ + import decord + + video_path = strip_file_uri(ele["video"]) + st = time.time() + vr = decord.VideoReader(video_path) + total_frames, video_fps = len(vr), vr.get_avg_fps() + start_frame, end_frame, total_frames = calculate_video_frame_range( + ele, + total_frames, + video_fps, + ) + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist() + video = vr.get_batch(idx).asnumpy() + video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format + logger.info(f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + + video_metadata = dict( + fps=video_fps, + frames_indices=idx, + total_num_frames=total_frames, + video_backend="decord", + ) + return video, video_metadata, sample_fps + + +def is_torchcodec_available() -> bool: + import importlib.util + + return importlib.util.find_spec("torchcodec") is not None + + +def _read_video_torchcodec( + ele: Dict[str, Any], +) -> VideoReaderOutput: + """read video using torchcodec.decoders.VideoDecoder + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + Tuple of video tensor (T, C, H, W), metadata, and sampled FPS. + """ + from torchcodec.decoders import VideoDecoder + + TORCHCODEC_NUM_THREADS = int(os.environ.get("TORCHCODEC_NUM_THREADS", 8)) + logger.info(f"set TORCHCODEC_NUM_THREADS: {TORCHCODEC_NUM_THREADS}") + video_path = strip_file_uri(ele["video"]) + st = time.time() + decoder = VideoDecoder(video_path, num_ffmpeg_threads=TORCHCODEC_NUM_THREADS) + video_fps = decoder.metadata.average_fps + total_frames = decoder.metadata.num_frames + start_frame, end_frame, total_frames = calculate_video_frame_range( + ele, + total_frames, + video_fps, + ) + nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps) + idx = torch.linspace(start_frame, end_frame, nframes).round().long().tolist() + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + video = decoder.get_frames_at(indices=idx).data + logger.info(f"torchcodec: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s") + + video_metadata = dict( + fps=video_fps, + frames_indices=idx, + total_num_frames=total_frames, + video_backend="torchcodec", + ) + return video, video_metadata, sample_fps + + +VIDEO_READER_BACKENDS = { + "decord": _read_video_decord, + "torchvision": _read_video_torchvision, + "torchcodec": _read_video_torchcodec, +} + +FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None) + + +@lru_cache(maxsize=1) +def get_video_reader_backend() -> str: + if FORCE_QWENVL_VIDEO_READER is not None: + video_reader_backend = FORCE_QWENVL_VIDEO_READER + elif is_torchcodec_available(): + video_reader_backend = "torchcodec" + elif is_decord_available(): + video_reader_backend = "decord" + else: + video_reader_backend = "torchvision" + print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr) + return video_reader_backend + + +def fetch_video( + ele: Dict[str, Any], + image_patch_size: int = 14, + return_video_sample_fps: bool = False, + return_video_metadata: bool = False, +) -> FetchVideoOutput: + image_factor = image_patch_size * SPATIAL_MERGE_SIZE + VIDEO_FRAME_MIN_PIXELS = VIDEO_MIN_TOKEN_NUM * image_factor * image_factor + VIDEO_FRAME_MAX_PIXELS = VIDEO_MAX_TOKEN_NUM * image_factor * image_factor + if isinstance(ele["video"], str): + video_reader_backend = get_video_reader_backend() + try: + video, video_metadata, sample_fps = VIDEO_READER_BACKENDS[video_reader_backend](ele) + except Exception as e: + # logger.warning(f"video_reader_backend {video_reader_backend} error, use torchvision as default, msg: {e}") + logger.warning( + f"video_reader_backend {video_reader_backend} error, for {ele=} use torchvision as default, msg: {e}" + ) + video, video_metadata, sample_fps = VIDEO_READER_BACKENDS["torchvision"](ele) + else: + # The input is a list of frames + assert isinstance(ele["video"], (list, tuple)) + process_info = ele.copy() + process_info.pop("type", None) + process_info.pop("video", None) + # use ThreadPoolExecutor to parallel process frames + max_workers = min(MAX_NUM_WORKERS_FETCH_VIDEO, len(ele["video"])) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit(fetch_image, {"image": video_element, **process_info}, image_factor) + for video_element in ele["video"] + ] + image_list = [future.result() for future in futures] + + nframes = ceil_by_factor(len(image_list), FRAME_FACTOR) + if len(image_list) < nframes: + image_list.extend([image_list[-1]] * (nframes - len(image_list))) + + sample_fps = ele.get("sample_fps", 2.0) + video = torch.stack([torch.from_numpy(np.array(image).transpose(2, 0, 1)) for image in image_list]) + + # fake video metadata + raw_fps = process_info.pop("raw_fps", sample_fps) + video_metadata = dict( + fps=raw_fps, + frames_indices=[i for i in range(len(video))], + total_num_frames=(nframes / sample_fps) * raw_fps, + ) + + nframes, _, height, width = video.shape + min_pixels = ele.get("min_pixels", VIDEO_FRAME_MIN_PIXELS) + total_pixels = ele.get("total_pixels", MODEL_SEQ_LEN * image_factor * image_factor * 0.9) + max_pixels = max( + min(VIDEO_FRAME_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05) + ) + max_pixels_supposed = ele.get("max_pixels", max_pixels) + if max_pixels_supposed > max_pixels: + logger.warning(f"The given max_pixels[{max_pixels_supposed}] exceeds limit[{max_pixels}].") + max_pixels = min(max_pixels_supposed, max_pixels) + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=image_factor, + ) + else: + resized_height, resized_width = smart_resize( + height, + width, + factor=image_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + video = transforms.functional.resize( + video, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ).float() + + final_video = (video, video_metadata) if return_video_metadata else video + if return_video_sample_fps: + return final_video, sample_fps + return final_video + + +def extract_vision_info( + conversations: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]] +) -> List[Dict[str, Any]]: + vision_infos = [] + if isinstance(conversations[0], dict): + conversations = [conversations] + for conversation in conversations: + for message in conversation: + if isinstance(message["content"], list): + for ele in message["content"]: + if ( + "image" in ele + or "image_url" in ele + or "video" in ele + or ele.get("type", "text") in ("image", "image_url", "video") + ): + vision_infos.append(ele) + return vision_infos + + +def process_vision_info( + conversations: Union[List[Dict[str, Any]], List[List[Dict[str, Any]]]], + return_video_kwargs: bool = False, + return_video_metadata: bool = False, + image_patch_size: int = 14, +) -> Tuple[ + Optional[List[Image.Image]], + Optional[List[Union[torch.Tensor, List[Image.Image]]]], + Optional[Dict[str, Any]], +]: + + vision_infos = extract_vision_info(conversations) + ## Read images or videos + image_inputs = [] + video_inputs = [] + video_sample_fps_list = [] + for vision_info in vision_infos: + if "image" in vision_info or "image_url" in vision_info: + image_inputs.append(fetch_image(vision_info, image_patch_size=image_patch_size)) + elif "video" in vision_info: + video_input, video_sample_fps = fetch_video( + vision_info, + return_video_sample_fps=True, + image_patch_size=image_patch_size, + return_video_metadata=return_video_metadata, + ) + video_sample_fps_list.append(video_sample_fps) + video_inputs.append(video_input) + else: + raise ValueError("image, image_url or video should in content.") + if len(image_inputs) == 0: + image_inputs = None + if len(video_inputs) == 0: + video_inputs = None + + video_kwargs = {"do_sample_frames": False} + if not return_video_metadata: # BC for qwen2.5vl + video_kwargs.update({"fps": video_sample_fps_list}) + + if return_video_kwargs: + return image_inputs, video_inputs, video_kwargs + return image_inputs, video_inputs diff --git a/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml b/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml new file mode 100644 index 000000000..023776ecd --- /dev/null +++ b/primus/configs/models/diffusion/wan2.1_t2v_1.3b.yaml @@ -0,0 +1,13 @@ +model: + name: wan + config: + load_from_pretrained_path: ${PRETRAINED_PATH:/models/Wan2.1-T2V-1.3B} + config: + model_type: t2v + separated_timestep: false + fuse_vae_embedding_in_latents: false + trainable_modules: dit + vae_type: wan_video_vae + encoder: + t5_encoder: ${TEXT_ENCODER:/models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth} + autoencoder: ${VAE_CHECKPOINT:/models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth} diff --git a/primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml b/primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml new file mode 100644 index 000000000..d9f089b3a --- /dev/null +++ b/primus/configs/models/diffusion/wan2.1_t2v_1.3b_sft.yaml @@ -0,0 +1,8 @@ +extends: + - wan2.1_t2v_1.3b.yaml + +# SFT reuses the Wan2.1-T2V-1.3B preset and only swaps the DiT +# initialization source to INIT_CHECKPOINT (post-train convention). +model: + config: + load_from_pretrained_path: ${INIT_CHECKPOINT:/models/Wan2.1-T2V-1.3B} diff --git a/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml b/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml new file mode 100644 index 000000000..0f23001e2 --- /dev/null +++ b/primus/configs/models/diffusion/wan2.2_ti2v_5b.yaml @@ -0,0 +1,13 @@ +model: + name: wan + config: + load_from_pretrained_path: ${PRETRAINED_PATH:/models/Wan2.2-TI2V-5B} + config: + model_type: ti2v + separated_timestep: true + fuse_vae_embedding_in_latents: true + trainable_modules: dit + vae_type: wan_video_vae_38 + encoder: + t5_encoder: ${TEXT_ENCODER:/models/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth} + autoencoder: ${VAE_CHECKPOINT:/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth} diff --git a/primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml b/primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml new file mode 100644 index 000000000..b8b1cd8d8 --- /dev/null +++ b/primus/configs/models/diffusion/wan2.2_ti2v_5b_sft.yaml @@ -0,0 +1,8 @@ +extends: + - wan2.2_ti2v_5b.yaml + +# SFT reuses the Wan2.2-TI2V-5B preset and only swaps the DiT +# initialization source to INIT_CHECKPOINT (post-train convention). +model: + config: + load_from_pretrained_path: ${INIT_CHECKPOINT:/models/Wan2.2-TI2V-5B} diff --git a/primus/configs/modules/diffusion/post_trainer.yaml b/primus/configs/modules/diffusion/post_trainer.yaml new file mode 100644 index 000000000..e89d7b62f --- /dev/null +++ b/primus/configs/modules/diffusion/post_trainer.yaml @@ -0,0 +1,5 @@ +extends: + - ../module_base.yaml + +trainable: true +stage: posttrain diff --git a/primus/configs/modules/diffusion/pre_trainer.yaml b/primus/configs/modules/diffusion/pre_trainer.yaml new file mode 100644 index 000000000..3d9d3cbdb --- /dev/null +++ b/primus/configs/modules/diffusion/pre_trainer.yaml @@ -0,0 +1,4 @@ +extends: + - ../module_base.yaml + +stage: pretrain diff --git a/runner/helpers/hooks/train/posttrain/diffusion/prepare.py b/runner/helpers/hooks/train/posttrain/diffusion/prepare.py new file mode 100644 index 000000000..4f011269f --- /dev/null +++ b/runner/helpers/hooks/train/posttrain/diffusion/prepare.py @@ -0,0 +1,24 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import runpy +import sys +from pathlib import Path + + +def main() -> None: + primus_root = Path(__file__).resolve().parents[6] + hook = primus_root / "runner" / "helpers" / "hooks" / "train" / "pretrain" / "diffusion" / "prepare.py" + sys.argv[0] = str(hook) + if "--module_name" not in sys.argv: + sys.argv.extend(["--module_name", "post_trainer"]) + runpy.run_path(str(hook), run_name="__main__") + + +if __name__ == "__main__": + main() diff --git a/runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh b/runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh new file mode 100755 index 000000000..e0c3c97c0 --- /dev/null +++ b/runner/helpers/hooks/train/pretrain/diffusion/00_install_requirements.sh @@ -0,0 +1,40 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +DEFAULT_PRIMUS_ROOT="$(cd "${SCRIPT_DIR}/../../../../../.." && pwd)" +PRIMUS_ROOT="${PRIMUS_PATH:-${DEFAULT_PRIMUS_ROOT}}" + +while [[ $# -gt 0 ]]; do + case "$1" in + --data_path) + DATA_PATH="$2" + shift 2 + ;; + --primus_path) + PRIMUS_ROOT="$2" + shift 2 + ;; + *) + shift + ;; + esac +done + +DATA_PATH="${DATA_PATH:-${PRIMUS_ROOT}/data}" +PIP_CACHE_DIR="${PIP_CACHE_DIR:-${DATA_PATH}/pip_cache}" + +echo "[INFO] Using pip cache: ${PIP_CACHE_DIR}" +mkdir -p "${PIP_CACHE_DIR}" + +REQ_FILE="${SCRIPT_DIR}/requirements-diffusion.txt" +if [[ -f "${REQ_FILE}" ]] && grep -qE '^[[:space:]]*[^#[:space:]]' "${REQ_FILE}"; then + echo "[+] Installing Diffusion dependencies..." + pip install --cache-dir="${PIP_CACHE_DIR}" -r "${REQ_FILE}" + echo "[OK] Diffusion dependencies installed" +fi diff --git a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py new file mode 100644 index 000000000..b23e6a5ae --- /dev/null +++ b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py @@ -0,0 +1,132 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path +from typing import Any + +PRIMUS_ROOT = Path(__file__).resolve().parents[6] +if str(PRIMUS_ROOT) not in sys.path: + sys.path.insert(0, str(PRIMUS_ROOT)) + +from primus.backends.diffusion.diffusion_adapter import DiffusionAdapter +from primus.core.config.primus_config import get_module_config, load_primus_config +from primus.core.utils.yaml_utils import nested_namespace_to_dict + + +def _log(message: str) -> None: + print(f"[INFO] diffusion prepare: {message}", file=sys.stderr) + + +def _fail(message: str) -> None: + print(f"[ERROR] diffusion prepare: {message}", file=sys.stderr) + raise SystemExit(1) + + +def _as_dict(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return value + return nested_namespace_to_dict(value) + + +def _select_module_name(cfg: Any, requested: str | None) -> str: + if requested: + get_module_config(cfg, requested) + return requested + + for module_name in ("pre_trainer", "post_trainer"): + try: + get_module_config(cfg, module_name) + return module_name + except Exception: + continue + + _fail("config must contain either modules.pre_trainer or modules.post_trainer") + raise AssertionError("unreachable") + + +def _is_placeholder(path: str | None) -> bool: + return not path or path.startswith("/path/to/") + + +def _require_path(path: str | None, description: str, *, kind: str = "any") -> None: + if _is_placeholder(path): + _fail(f"{description} is not configured: {path!r}") + + resolved = Path(path).expanduser() + if kind == "file" and not resolved.is_file(): + _fail(f"{description} file not found: {resolved}") + if kind == "dir" and not resolved.is_dir(): + _fail(f"{description} directory not found: {resolved}") + if kind == "any" and not resolved.exists(): + _fail(f"{description} path not found: {resolved}") + + _log(f"{description}: {resolved}") + + +def validate_diffusion_config(config_path: Path, module_name: str | None = None) -> None: + cfg = load_primus_config(config_path) + selected_module = _select_module_name(cfg, module_name) + module_cfg = get_module_config(cfg, selected_module) + + if getattr(module_cfg, "framework", None) != "diffusion": + _log(f"module {selected_module} framework is not diffusion; skipping") + return + + backend_args = DiffusionAdapter().convert_config(module_cfg.params) + model = _as_dict(backend_args.model) + dataset = _as_dict(backend_args.dataset) + trainer = _as_dict(backend_args.trainer) + + dataset_cfg = dataset.get("config", {}) + processor_cfg = dataset_cfg.get("processor_config", {}) + encoder_cfg = model.get("config", {}).get("encoder", {}) or model.get("encoder", {}) + + _require_path(dataset_cfg.get("dataset_path"), "dataset metadata", kind="file") + _require_path(dataset_cfg.get("data_folder"), "dataset media folder", kind="dir") + _require_path(processor_cfg.get("text_tokenizer"), "text tokenizer", kind="dir") + + model_cfg = model.get("config", {}) + _require_path(model_cfg.get("load_from_pretrained_path"), "DiT initialization checkpoint") + _require_path(encoder_cfg.get("t5_encoder"), "text encoder checkpoint", kind="file") + _require_path(encoder_cfg.get("autoencoder"), "VAE checkpoint", kind="file") + + _log( + "validated " + f"module={selected_module} stage={getattr(backend_args, 'stage', None)} " + f"trainer={trainer.get('name')} model={model.get('name')}" + ) + print("env.PREPARED=1") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Prepare Primus diffusion training environment") + parser.add_argument("--config", type=Path, required=True, help="Experiment YAML config") + parser.add_argument("--data_path", type=Path, required=False, help="Reserved for hook API compatibility") + parser.add_argument( + "--primus_path", type=Path, required=False, help="Reserved for hook API compatibility" + ) + parser.add_argument("--patch_args", type=Path, required=False, help="Reserved for hook API compatibility") + parser.add_argument("--backend_path", type=str, default=None, help="Unused; diffusion is in-tree") + parser.add_argument("--module_name", type=str, default=None, help="Override module name to validate") + args, _unknown = parser.parse_known_args() + + if os.environ.get("SKIP_PREPARE") == "1": + _log("SKIP_PREPARE=1; skipping validation") + return + + if args.backend_path: + _fail("diffusion is an in-tree backend and does not support --backend_path") + + validate_diffusion_config(args.config, module_name=args.module_name) + + +if __name__ == "__main__": + main() diff --git a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt new file mode 100644 index 000000000..38fabbeef --- /dev/null +++ b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt @@ -0,0 +1,12 @@ +einops +loguru +safetensors +numpy +pillow +packaging +requests +decord +imageio[ffmpeg] +wandb +pydantic +transformers==4.50.0 diff --git a/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py b/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py new file mode 100644 index 000000000..9d07d4f07 --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_wan_argument_builder.py @@ -0,0 +1,194 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from types import SimpleNamespace + +import pytest + +from primus.backends.diffusion.argument_builder import WanArgBuilder + + +def _minimal_params(): + return { + "model": { + "name": "wan", + "config": { + "model_type": "t2v", + }, + }, + } + + +def test_rejects_legacy_public_dataset_override(): + params = _minimal_params() + params["dataset"] = {"name": "wan"} + + builder = WanArgBuilder() + builder.update(params) + + with pytest.raises(ValueError, match="no longer accepts public `dataset` overrides"): + builder.finalize() + + +def test_rejects_legacy_public_trainer_override(): + params = _minimal_params() + params["trainer"] = {"name": "fsdp2"} + + builder = WanArgBuilder() + builder.update(params) + + with pytest.raises(ValueError, match="no longer accepts public `trainer` overrides"): + builder.finalize() + + +def test_maps_primus_style_sections_to_wan_runtime_config(): + params = { + **_minimal_params(), + "stage": "posttrain", + "primus": {"experiment": "wan-smoke"}, + "training": { + "steps": 7, + "local_batch_size": 2, + "global_batch_size": 16, + "gradient_accumulation_steps": 4, + "output_dir": "/tmp/wan-out", + "save_steps": 3, + "run_name": "wan-test", + "num_train_epochs": 9, + "dataloader_num_workers": 0, + "resume_from_checkpoint": "/tmp/resume", + }, + "data": { + "dataset_path": "/data/meta.jsonl", + "data_folder": "/data/videos", + "frame_num": 17, + "video_backend": "decord", + "text_tokenizer": "/models/umt5", + "height": 256, + "width": 384, + }, + "parallelism": { + "sp_size": 4, + "dp_replicate": 2, + }, + "optimizer": { + "lr": 2.0e-5, + "weight_decay": 0.02, + "adam_beta1": 0.8, + "adam_beta2": 0.95, + "adam_epsilon": 1.0e-7, + "max_grad_norm": 0.5, + }, + "runtime": { + "attention_backend": "sdpa", + "report_to": "none", + "seed": 1234, + "fsdp2_reshard_after_forward": False, + }, + "metrics": { + "log_freq": 5, + "enable_wandb": False, + }, + } + + builder = WanArgBuilder() + builder.update(params) + result = builder.finalize() + + dataset_cfg = result.dataset["config"] + processor_cfg = dataset_cfg["processor_config"] + trainer_args = result.trainer["args"] + + assert result.stage == "posttrain" + assert result.primus == {"experiment": "wan-smoke"} + assert result.model == params["model"] + + assert dataset_cfg["dataset_path"] == "/data/meta.jsonl" + assert dataset_cfg["data_folder"] == "/data/videos" + assert dataset_cfg["frame_num"] == 17 + assert dataset_cfg["video_backend"] == "decord" + assert processor_cfg["text_tokenizer"] == "/models/umt5" + assert processor_cfg["extra_kwargs"]["size"] == {"height": 256, "width": 384} + + assert trainer_args["max_steps"] == 7 + assert trainer_args["per_device_train_batch_size"] == 2 + assert trainer_args["global_batch_size"] == 16 + assert trainer_args["gradient_accumulation_steps"] == 4 + assert trainer_args["output_dir"] == "/tmp/wan-out" + assert trainer_args["save_steps"] == 3 + assert trainer_args["run_name"] == "wan-test" + assert trainer_args["num_train_epochs"] == 9 + assert trainer_args["dataloader_num_workers"] == 0 + assert trainer_args["resume_from_checkpoint"] == "/tmp/resume" + + assert trainer_args["sp_size"] == 4 + assert trainer_args["dp_replicate"] == 2 + assert trainer_args["learning_rate"] == 2.0e-5 + assert trainer_args["weight_decay"] == 0.02 + assert trainer_args["adam_beta1"] == 0.8 + assert trainer_args["adam_beta2"] == 0.95 + assert trainer_args["adam_epsilon"] == 1.0e-7 + assert trainer_args["max_grad_norm"] == 0.5 + assert trainer_args["attention_backend"] == "sdpa" + assert trainer_args["report_to"] == "none" + assert trainer_args["seed"] == 1234 + assert trainer_args["fsdp2_reshard_after_forward"] is False + assert trainer_args["logging_steps"] == 5 + + +def test_defaults_propagate_when_optional_sections_are_omitted(): + builder = WanArgBuilder() + builder.update(_minimal_params()) + result = builder.finalize() + + dataset_cfg = result.dataset["config"] + trainer_args = result.trainer["args"] + + assert result.stage == "pretrain" + assert dataset_cfg["dataset_path"] == "/path/to/meta.jsonl" + assert dataset_cfg["data_folder"] == "/path/to/videos" + assert dataset_cfg["processor_config"]["text_tokenizer"] == "/path/to/umt5-xxl" + assert result.trainer["name"] == "fsdp2" + assert trainer_args["max_steps"] == 100 + assert trainer_args["attention_backend"] == "flash_attn_aiter" + assert trainer_args["sp_size"] == 1 + assert trainer_args["dp_replicate"] == 1 + assert trainer_args["report_to"] == "none" + + +def test_metrics_enable_wandb_sets_report_to_when_runtime_omits_it(): + builder = WanArgBuilder() + builder.update( + { + **_minimal_params(), + "metrics": { + "enable_wandb": True, + }, + } + ) + + result = builder.finalize() + + assert result.trainer["args"]["report_to"] == "wandb" + + +def test_update_accepts_simplenamespace_input(): + builder = WanArgBuilder() + builder.update( + SimpleNamespace( + model=SimpleNamespace( + name="wan", + config=SimpleNamespace(model_type="t2v"), + ), + data=SimpleNamespace(dataset_path="/data/meta.jsonl"), + ) + ) + + result = builder.finalize() + + assert result.model["name"] == "wan" + assert result.model["config"]["model_type"] == "t2v" + assert result.dataset["config"]["dataset_path"] == "/data/meta.jsonl" diff --git a/tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py b/tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py new file mode 100644 index 000000000..6636942e9 --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_wan_trainer_optimizer.py @@ -0,0 +1,61 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import torch + +from primus.backends.diffusion.trainers.base import BaseWanTrainer + + +def _make_trainer(): + trainer = BaseWanTrainer.__new__(BaseWanTrainer) + trainer.model = torch.nn.Linear(2, 2) + trainer.rank = 0 + trainer.args = { + "learning_rate": 1.0e-4, + "weight_decay": 0.01, + "adam_beta1": 0.9, + "adam_beta2": 0.999, + "adam_epsilon": 1.0e-8, + } + return trainer + + +def test_optimizer_falls_back_to_foreach_when_fused_raises_runtime_error(monkeypatch): + calls = [] + + def fake_adamw(**kwargs): + calls.append(kwargs) + if kwargs.get("fused"): + raise RuntimeError("fused AdamW is unsupported") + return kwargs + + monkeypatch.setattr(torch.optim, "AdamW", fake_adamw) + + optimizer = _make_trainer()._create_optimizer() + + assert calls[0]["fused"] is True + assert calls[1]["foreach"] is True + assert optimizer is calls[1] + + +def test_optimizer_falls_back_to_default_when_foreach_raises_runtime_error(monkeypatch): + calls = [] + + def fake_adamw(**kwargs): + calls.append(kwargs) + if kwargs.get("fused") or kwargs.get("foreach"): + raise RuntimeError("optimized AdamW path is unsupported") + return kwargs + + monkeypatch.setattr(torch.optim, "AdamW", fake_adamw) + + optimizer = _make_trainer()._create_optimizer() + + assert calls[0]["fused"] is True + assert calls[1]["foreach"] is True + assert "fused" not in calls[2] + assert "foreach" not in calls[2] + assert optimizer is calls[2] From 03eedade9e0968284f4bec704cb509fd02dd9688 Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Wed, 8 Jul 2026 18:16:20 +0800 Subject: [PATCH 014/127] fix: remove duplicated flag use_turbo_fp4_autocast (#860) # Description This PR removes the redundant `use_turbo_fp4_autocast` flag and simplifies FP4 autocast routing in Megatron. Previously, enabling Primus-Turbo FP4 autocast required both `enable_primus_turbo` and `use_turbo_fp4_autocast` to be set. Actually the Primus-Turbo autocast is compatible with TE autocast. That duplicated control was confusing and easy to misconfigure (e.g. Turbo GEMM/attention enabled while FP4 still fell back to Transformer Engine). With this change, MXFP4 training uses the Primus-Turbo FP4 autocast path whenever `enable_primus_turbo` is enabled, which aligns FP4 behavior with other Primus-Turbo features. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Remove `use_turbo_fp4_autocast` from `primus/configs/modules/megatron/primus_turbo.yaml`. - Update `_primus_turbo_enabled()` in `fp4_utils.py` to gate the Turbo FP4 path on `enable_primus_turbo` only, instead of requiring both `enable_primus_turbo` and `use_turbo_fp4_autocast`. # Checklist: - [x] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --- .../megatron/configs/MI355X/llama3.1_8B-MXFP4-pretrain.yaml | 1 - primus/backends/megatron/core/fp4_utils.py | 3 +-- primus/configs/modules/megatron/primus_turbo.yaml | 2 -- 3 files changed, 1 insertion(+), 5 deletions(-) diff --git a/examples/megatron/configs/MI355X/llama3.1_8B-MXFP4-pretrain.yaml b/examples/megatron/configs/MI355X/llama3.1_8B-MXFP4-pretrain.yaml index f5072cb46..e5e920603 100644 --- a/examples/megatron/configs/MI355X/llama3.1_8B-MXFP4-pretrain.yaml +++ b/examples/megatron/configs/MI355X/llama3.1_8B-MXFP4-pretrain.yaml @@ -110,7 +110,6 @@ modules: # --- Primus Turbo Config --- enable_primus_turbo: true use_turbo_attention: true - use_turbo_fp4_autocast: false # TE mxfp4 recipe should set it to false use_turbo_gemm: false # can't use together with delayed recipe use_turbo_grouped_gemm: false moe_use_fused_router_with_aux_score: false diff --git a/primus/backends/megatron/core/fp4_utils.py b/primus/backends/megatron/core/fp4_utils.py index e141b4f8e..319f18990 100644 --- a/primus/backends/megatron/core/fp4_utils.py +++ b/primus/backends/megatron/core/fp4_utils.py @@ -45,8 +45,7 @@ def _primus_turbo_enabled() -> bool: args = get_args() enable_primus_turbo = bool(getattr(args, "enable_primus_turbo", False)) - use_turbo_fp4_autocast = bool(getattr(args, "use_turbo_fp4_autocast", False)) - return enable_primus_turbo and use_turbo_fp4_autocast + return enable_primus_turbo except Exception: return False diff --git a/primus/configs/modules/megatron/primus_turbo.yaml b/primus/configs/modules/megatron/primus_turbo.yaml index 290024b3a..bad3403d6 100644 --- a/primus/configs/modules/megatron/primus_turbo.yaml +++ b/primus/configs/modules/megatron/primus_turbo.yaml @@ -1,8 +1,6 @@ # ===== Global ===== # main control flag enable_primus_turbo: false -# use primus_turbo_fp4_autocast instead of TE fp8_autocast for FP4 (requires enable_primus_turbo) -use_turbo_fp4_autocast: false # ===== Attention ===== # operator switch From 0d168b3618bddc5082ed30b68f6114ccf00f2634 Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Wed, 8 Jul 2026 19:04:24 +0800 Subject: [PATCH 015/127] refactor: remove primus/modules and migrate still-used code into core/backends (#851) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Remove the legacy `primus/modules/` package entirely. Code still used by the current core-runtime training path is migrated/refactored into its natural home in `primus/core` and `primus/backends/megatron`; the legacy `PRIMUS_TRAIN_RUNTIME=legacy` flow and dead trainer code are deleted. ## What changed - **core**: `module_utils` / `base_module` moved out of `primus/modules` into `primus/core/utils/module_utils.py` and `primus/core/base_module.py`; all importers updated. - **megatron backend**: still-used trainer utilities rehomed — PP-visualization wrappers to `core/pipeline_parallel/pp_visualizer.py`; ROCm arg validation to `patches/args/rocm_arg_validation.py`; batch loader (`DataLoaderStore`) to `data_loader_store.py`; `is_v_schedule_enabled` to `training/utils.py`. - **legacy runtime removed**: dropped the `PRIMUS_TRAIN_RUNTIME=legacy` branch/resolver in the train subcommand and the `load_backend_trainer` / `launch_pretrain_*` entrypoints in `primus/pretrain.py` (kept `setup_backend_path` / `setup_env`, still used by projection/runner/examples). Removed torchtitan/maxtext legacy pretrainers and dead code (`sft_trainer`, torchtitan `parse_utils`). - **projection**: performance/memory layer-benchmark now builds its model via a new training-neutral `PrimusRuntime.setup_model_only()` / `MegatronPretrainTrainer.setup_model_only()` instead of the legacy `MegatronPretrainTrainer`; the entire `primus/backends/megatron/pretrainer/` package is deleted. - **bug fixes surfaced during migration**: restore the FSDP2 availability probe import (autoflake had reduced the guard to `pass`, making `HAVE_FSDP2` unconditionally True), and make `get_batch_func` return a consistent 5-tuple. --------- Co-authored-by: Xiaoming-AMD --- docs/backends/extending-backends.md | 4 +- primus/agents/tuning_agent/evaluator.py | 2 +- .../hummingbirdxt/hummingbirdxt_adapter.py | 2 +- primus/backends/maxtext/maxtext_adapter.py | 2 +- .../maxtext/maxtext_pretrain_trainer.py | 2 +- .../maxtext/patches/logger_patches.py | 2 +- .../backends/maxtext/patches/train_patches.py | 2 +- .../strategies/filesystem_async.py | 2 +- .../core/distributed/finalize_model_grad.py | 2 +- .../torch_fully_sharded_data_parallel.py | 6 +- primus/backends/megatron/core/fp4_utils.py | 2 +- primus/backends/megatron/core/fp8_utils.py | 2 +- .../fsdp2_bf16_master_weight_optimizer.py | 2 +- .../core/optimizer/fsdp2_fp32_optimizer.py | 2 +- .../megatron/core/optimizer/zbpp_optimizer.py | 2 +- .../backends/megatron/core/parallel_state.py | 2 +- .../core/pipeline_parallel/pp_visualizer.py | 193 ++ .../primuspipe/handlers/bwd_handler.py | 4 +- .../handlers/combined_fwd_bkwd_handler.py | 4 +- .../primuspipe/handlers/fwd_handler.py | 4 +- .../primuspipe/pipeline_launcher.py | 6 +- .../pipeline_parallel/zerobubble/offload.py | 2 +- .../pipeline_parallel/zerobubble/runtime.py | 12 +- .../zerobubble/scheduler/basic1f1b.py | 2 +- .../zerobubble/scheduler/communication.py | 2 +- .../scheduler/group_interleaved_1f1b.py | 2 +- .../zerobubble/scheduler/offloading.py | 2 +- .../zerobubble/scheduler/passes.py | 2 +- .../zerobubble/scheduler/v_auto_schedule.py | 2 +- .../zerobubble/scheduler/vpp.py | 2 +- .../zerobubble/scheduler/zb.py | 2 +- .../zerobubble/scheduler/zbv.py | 2 +- .../zerobubble/scheduler/zbv_greedy.py | 2 +- .../zerobubble/zbpp_utils.py | 4 +- .../pipeline_parallel_layer_layout.py | 2 +- .../core/transformer/transformer_layer.py | 2 +- primus/backends/megatron/data_loader_store.py | 132 + primus/backends/megatron/megatron_adapter.py | 3 +- .../megatron/megatron_base_trainer.py | 2 +- .../megatron/megatron_pretrain_trainer.py | 58 +- .../backends/megatron/megatron_sft_trainer.py | 2 +- primus/backends/megatron/patches/__init__.py | 2 +- .../patches/args/checkpoint_path_patches.py | 2 +- .../patches/args/data_path_split_patches.py | 2 +- .../patches/args/hsdp_args_patches.py | 2 +- .../iterations_to_skip_default_patches.py | 2 +- .../patches/args/logging_level_patches.py | 2 +- .../patches/args/mock_data_patches.py | 2 +- .../patches/args/moe_layer_freq_patches.py | 2 +- .../patches/args/rocm_arg_validation.py | 197 ++ .../args/sequence_parallel_tp1_patches.py | 2 +- .../patches/args/tensorboard_path_patches.py | 2 +- .../patches/args/validate_args_patches.py | 6 +- .../patches/args/wandb_config_patches.py | 2 +- .../megatron/patches/args_compat_patches.py | 2 +- .../megatron/patches/build_model_patches.py | 2 +- .../megatron/patches/checkpoint_patches.py | 2 +- .../megatron/patches/dataloader_patch.py | 2 +- .../patches/distributed_init_patches.py | 2 +- .../backends/megatron/patches/env_patches.py | 2 +- .../megatron/patches/evaluate_patches.py | 2 +- .../patches/fsdp2_fp8_cache_patches.py | 2 +- .../gpt_decoder_layer_specs_patches.py | 2 +- .../megatron/patches/mamba_rocm_patches.py | 2 +- .../megatron/patches/megatron_fsdp_patches.py | 2 +- .../backends/megatron/patches/mla_patches.py | 2 +- .../moe_patches/deprecated_layer_patches.py | 2 +- .../moe_patches/permute_fusion_patches.py | 2 +- .../moe_patches/skip_identity_sort_patches.py | 2 +- .../moe_patches/topk_router_patches.py | 2 +- .../megatron/patches/mp_sync_skip_patches.py | 2 +- .../patches/muon_optimizer_patches.py | 2 +- .../megatron/patches/optimizer_patches.py | 2 +- .../parallelism/forward_step_patches.py | 4 +- .../parallelism/linear_grad_split_patches.py | 2 +- .../pipeline_parallel_layout_patches.py | 2 +- .../patches/parallelism/schedule_patches.py | 2 +- .../sdma_param_all_gather_patches.py | 2 +- .../parallelism/te_wgrad_split_patches.py | 2 +- .../patches/parallelism/train_step_patches.py | 2 +- .../patches/parallelism/v_schedule_patches.py | 2 +- .../parallelism/zero_bubble_patches.py | 2 +- .../megatron/patches/pp_dump_data_patches.py | 12 +- .../megatron/patches/pp_warmup_patches.py | 2 +- .../patches/recompute_layer_patches.py | 2 +- .../megatron/patches/runtime_hooks_patches.py | 2 +- .../sdma_symm_mem_collectives_patches.py | 2 +- .../patches/sft_grad_sanitize_patches.py | 2 +- .../patches/te_patches/bshd_layout_patches.py | 2 +- .../te_patches/delayed_scaling_patches.py | 2 +- .../general_gemm_workspace_patches.py | 2 +- .../layernorm_linear_fp8_cache_patches.py | 2 +- .../legacy_grouped_mlp_wgrad_patches.py | 2 +- .../te_patches/linear_fp8_cache_patches.py | 2 +- .../patches/te_patches/tp_overlap_patches.py | 2 +- .../patches/tokenizer_builder_patches.py | 2 +- .../megatron/patches/torch_fsdp2_patches.py | 2 +- .../patches/torch_profiler_patches.py | 4 +- .../training_log/print_rank_last_patches.py | 2 +- .../patches/turbo/aiter_deepbind_patches.py | 2 +- .../megatron/patches/turbo/fp4_patches.py | 2 +- .../megatron/patches/turbo/fp8_patches.py | 2 +- .../turbo/fused_residual_norm_patches.py | 2 +- .../patches/turbo/moe_dispatcher_patches.py | 2 +- .../patches/turbo/rms_norm_patches.py | 2 +- .../patches/turbo/te_spec_provider_patches.py | 2 +- .../backends/megatron/patches/turbo/utils.py | 2 +- .../patches/zebra_llama_flops_patches.py | 2 +- primus/backends/megatron/peft/recompute.py | 2 +- primus/backends/megatron/sft/preprocessing.py | 2 +- primus/backends/megatron/sft/runtime.py | 2 +- .../backends/megatron/training/evaluator.py | 2 +- .../backends/megatron/training/global_vars.py | 2 +- .../megatron/training/mlflow_artifacts.py | 2 +- .../megatron/training/tokenizer/tokenizer.py | 2 +- primus/backends/megatron/training/utils.py | 15 +- .../backends/megatron_bridge/config_utils.py | 2 +- .../megatron_bridge_adapter.py | 2 +- .../megatron_bridge_base_trainer.py | 2 +- .../megatron_bridge_posttrain_trainer.py | 2 +- .../megatron_bridge_pretrain_trainer.py | 2 +- .../bridge_training_log_patches.py | 2 +- primus/backends/torchtitan/config_utils.py | 2 +- .../patches/dcp_consolidate_patches.py | 2 +- .../patches/embedding_amp_patches.py | 2 +- .../patches/flex_attention_patches.py | 2 +- .../torchtitan/patches/logger_patches.py | 2 +- .../patches/mock_dataset_patches.py | 5 +- .../patches/model_override_patches.py | 2 +- .../patches/pipelining_schedule_patches.py | 2 +- .../patches/sdma_symm_mem_collectives.py | 2 +- .../patches/turbo/async_tp_patches.py | 2 +- .../patches/turbo/attention_patches.py | 2 +- .../deepseek_v3_classic_attention_patches.py | 2 +- .../patches/turbo/fp8_linear_patches.py | 2 +- .../patches/turbo/moe_grouped_mm_patches.py | 2 +- .../patches/turbo/mx_linear_patches.py | 2 +- .../torchtitan/patches/wandb_patches.py | 2 +- .../backends/torchtitan/torchtitan_adapter.py | 2 +- .../transformer_engine/pytorch/module/base.py | 2 +- primus/cli/subcommands/train.py | 39 +- primus/core/backend/backend_adapter.py | 4 +- primus/core/backend/backend_registry.py | 2 +- primus/{modules => core}/base_module.py | 3 +- primus/core/patches/patch.py | 2 +- primus/core/patches/patch_runner.py | 2 +- .../handler/wgrad_handler.py | 4 +- .../scheduler/algorithms/base.py | 2 +- .../projection/memory_projection/benchmark.py | 2 +- .../performance_projection/projection.py | 88 +- primus/core/runtime/logging.py | 2 +- primus/core/runtime/train_runtime.py | 81 +- primus/core/trainer/base_trainer.py | 2 +- primus/core/utils/import_utils.py | 2 +- .../{modules => core/utils}/module_utils.py | 0 primus/modules/__init__.py | 0 primus/modules/trainer/__init__.py | 0 primus/modules/trainer/base_trainer.py | 24 - primus/modules/trainer/maxtext/__init__.py | 0 primus/modules/trainer/maxtext/pre_trainer.py | 144 - primus/modules/trainer/megatron/__init__.py | 0 .../trainer/megatron/model_provider.py | 53 - .../modules/trainer/megatron/pre_trainer.py | 307 -- .../modules/trainer/megatron/sft_trainer.py | 22 - primus/modules/trainer/megatron/trainer.py | 2561 ----------------- primus/modules/trainer/megatron/utils.py | 611 ---- primus/modules/trainer/torchtitan/__init__.py | 0 .../modules/trainer/torchtitan/parse_utils.py | 66 - .../modules/trainer/torchtitan/pre_trainer.py | 297 -- primus/pretrain.py | 152 +- skills/backend-gap-report/examples.md | 2 +- .../patches/test_fsdp2_fp32_patches.py | 6 +- .../megatron/test_megatron_adapter.py | 6 +- .../megatron/test_training_log_patches.py | 6 +- .../megatron/test_validate_args_patches.py | 16 +- tests/unit_tests/cli/test_train_subcommand.py | 40 - .../core/trainer/test_base_trainer.py | 2 +- .../megatron/cco/test_tp_overlap.py | 2 +- tests/unit_tests/test_backend_loader.py | 23 +- 179 files changed, 943 insertions(+), 4547 deletions(-) create mode 100644 primus/backends/megatron/core/pipeline_parallel/pp_visualizer.py create mode 100644 primus/backends/megatron/data_loader_store.py create mode 100644 primus/backends/megatron/patches/args/rocm_arg_validation.py rename primus/{modules => core}/base_module.py (98%) rename primus/{modules => core/utils}/module_utils.py (100%) delete mode 100644 primus/modules/__init__.py delete mode 100644 primus/modules/trainer/__init__.py delete mode 100644 primus/modules/trainer/base_trainer.py delete mode 100644 primus/modules/trainer/maxtext/__init__.py delete mode 100644 primus/modules/trainer/maxtext/pre_trainer.py delete mode 100644 primus/modules/trainer/megatron/__init__.py delete mode 100644 primus/modules/trainer/megatron/model_provider.py delete mode 100644 primus/modules/trainer/megatron/pre_trainer.py delete mode 100644 primus/modules/trainer/megatron/sft_trainer.py delete mode 100644 primus/modules/trainer/megatron/trainer.py delete mode 100644 primus/modules/trainer/megatron/utils.py delete mode 100644 primus/modules/trainer/torchtitan/__init__.py delete mode 100644 primus/modules/trainer/torchtitan/parse_utils.py delete mode 100644 primus/modules/trainer/torchtitan/pre_trainer.py delete mode 100644 tests/unit_tests/cli/test_train_subcommand.py diff --git a/docs/backends/extending-backends.md b/docs/backends/extending-backends.md index 129e57a51..e0ffc0bc6 100644 --- a/docs/backends/extending-backends.md +++ b/docs/backends/extending-backends.md @@ -95,7 +95,7 @@ from typing import Any, Dict from primus.core.backend.backend_adapter import BackendAdapter from primus.core.backend.backend_registry import BackendRegistry -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class DummyAdapter(BackendAdapter): @@ -154,7 +154,7 @@ Key points: from typing import Any from primus.core.trainer.base_trainer import BaseTrainer -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class DummyPretrainTrainer(BaseTrainer): diff --git a/primus/agents/tuning_agent/evaluator.py b/primus/agents/tuning_agent/evaluator.py index 4aa2330c9..2966545a4 100644 --- a/primus/agents/tuning_agent/evaluator.py +++ b/primus/agents/tuning_agent/evaluator.py @@ -227,7 +227,7 @@ def write_trial_yaml(arch: ArchitectureRecord, cfg: TrialConfig, out_dir: Path, # enable ``use_turbo_gemm `` inside Primus's projection. Two # things go wrong on the v26.2 container ``primus_turbo==0.2.0``: # * The default ``fp8_recipe: delayed`` is incompatible with that path - # (``primus/modules/trainer/megatron/utils.py:464`` asserts). + # (``primus/backends/megatron/patches/args/rocm_arg_validation.py`` asserts). # * The dense FP8 GEMM op (``primus_turbo.pytorch.ops.gemm_fp8``) on # this version raises ``ValueError: Unsupported FP8 format: HYBRID`` # for ``fp8: hybrid`` (a common DSv3 / Kimi-K2 configuration). diff --git a/primus/backends/hummingbirdxt/hummingbirdxt_adapter.py b/primus/backends/hummingbirdxt/hummingbirdxt_adapter.py index 8b7156e40..5970dd135 100644 --- a/primus/backends/hummingbirdxt/hummingbirdxt_adapter.py +++ b/primus/backends/hummingbirdxt/hummingbirdxt_adapter.py @@ -9,7 +9,7 @@ from primus.backends.hummingbirdxt.argument_builder import HummingbirdXTArgBuilder from primus.core.backend.backend_adapter import BackendAdapter from primus.core.backend.backend_registry import BackendRegistry -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class HummingbirdXTAdapter(BackendAdapter): diff --git a/primus/backends/maxtext/maxtext_adapter.py b/primus/backends/maxtext/maxtext_adapter.py index 7e44f7c49..417070633 100644 --- a/primus/backends/maxtext/maxtext_adapter.py +++ b/primus/backends/maxtext/maxtext_adapter.py @@ -24,7 +24,7 @@ import primus.backends.maxtext.patches # noqa: F401 # Register patches from primus.backends.maxtext.argument_builder import MaxTextConfigBuilder from primus.core.backend.backend_adapter import BackendAdapter -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 class MaxTextAdapter(BackendAdapter): diff --git a/primus/backends/maxtext/maxtext_pretrain_trainer.py b/primus/backends/maxtext/maxtext_pretrain_trainer.py index 8009230ac..17d207ad1 100644 --- a/primus/backends/maxtext/maxtext_pretrain_trainer.py +++ b/primus/backends/maxtext/maxtext_pretrain_trainer.py @@ -28,7 +28,7 @@ from typing import Any, Dict, Optional from primus.core.trainer.base_trainer import BaseTrainer -from primus.modules.module_utils import ( +from primus.core.utils.module_utils import ( error_rank_0, log_rank_0, set_logging_rank, diff --git a/primus/backends/maxtext/patches/logger_patches.py b/primus/backends/maxtext/patches/logger_patches.py index 8d7e50833..4c5cd2bf2 100644 --- a/primus/backends/maxtext/patches/logger_patches.py +++ b/primus/backends/maxtext/patches/logger_patches.py @@ -15,7 +15,7 @@ from primus.core.patches import PatchContext, register_patch from primus.core.utils import checker from primus.core.utils.logger import _logger as primus_logger -from primus.modules.module_utils import error_rank_0, log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import error_rank_0, log_rank_0, warning_rank_0 @register_patch( diff --git a/primus/backends/maxtext/patches/train_patches.py b/primus/backends/maxtext/patches/train_patches.py index ca2ca1d1d..39aa017fe 100644 --- a/primus/backends/maxtext/patches/train_patches.py +++ b/primus/backends/maxtext/patches/train_patches.py @@ -16,7 +16,7 @@ from typing import Any, Sequence from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 @register_patch( diff --git a/primus/backends/megatron/core/dist_checkpointing/strategies/filesystem_async.py b/primus/backends/megatron/core/dist_checkpointing/strategies/filesystem_async.py index 38108e464..f37c7d4dd 100644 --- a/primus/backends/megatron/core/dist_checkpointing/strategies/filesystem_async.py +++ b/primus/backends/megatron/core/dist_checkpointing/strategies/filesystem_async.py @@ -11,7 +11,7 @@ FileSystemWriterAsync, ) -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class PrimusFileSystemWriterAsync(FileSystemWriterAsync): diff --git a/primus/backends/megatron/core/distributed/finalize_model_grad.py b/primus/backends/megatron/core/distributed/finalize_model_grad.py index 02743cf3a..320b48f33 100644 --- a/primus/backends/megatron/core/distributed/finalize_model_grad.py +++ b/primus/backends/megatron/core/distributed/finalize_model_grad.py @@ -17,7 +17,7 @@ ) from megatron.core.utils import get_model_config -from primus.modules.trainer.megatron.utils import is_v_schedule_enabled +from primus.backends.megatron.training.utils import is_v_schedule_enabled def finalize_model_grads(model: List[torch.nn.Module], num_tokens: Optional[torch.Tensor] = None): diff --git a/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py b/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py index 8e7e75118..292dd1e68 100644 --- a/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py +++ b/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py @@ -22,7 +22,7 @@ from megatron.core.transformer.transformer_layer import TransformerLayer from torch.distributed import ProcessGroup -from primus.modules.module_utils import warning_rank_0 +from primus.core.utils.module_utils import warning_rank_0 try: from torch.distributed import DeviceMesh @@ -102,7 +102,7 @@ def __init__( # Build DeviceMesh from Megatron's process groups from megatron.training import get_args - from primus.modules.module_utils import log_rank_0 + from primus.core.utils.module_utils import log_rank_0 args = get_args() replicate_degree = getattr(args, "data_parallel_replicate_degree", 1) @@ -385,7 +385,7 @@ def compile_model(self): the actual model (e.g., Flux) and calls its compile_model() if present. Skipped if enable_torch_compile is False or no compile_model method is found. """ - from primus.modules.module_utils import log_rank_0 + from primus.core.utils.module_utils import log_rank_0 try: from megatron.training import get_args diff --git a/primus/backends/megatron/core/fp4_utils.py b/primus/backends/megatron/core/fp4_utils.py index 319f18990..280949961 100644 --- a/primus/backends/megatron/core/fp4_utils.py +++ b/primus/backends/megatron/core/fp4_utils.py @@ -14,7 +14,7 @@ from megatron.core.utils import is_te_min_version from primus.backends.megatron.core.enums import Fp4Recipe -from primus.modules.module_utils import warning_rank_0 +from primus.core.utils.module_utils import warning_rank_0 # Check if Transformer Engine is installed HAVE_TE = False diff --git a/primus/backends/megatron/core/fp8_utils.py b/primus/backends/megatron/core/fp8_utils.py index 7fa662c28..0e7ce7995 100644 --- a/primus/backends/megatron/core/fp8_utils.py +++ b/primus/backends/megatron/core/fp8_utils.py @@ -12,7 +12,7 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_te_min_version -from primus.modules.module_utils import warning_rank_0 +from primus.core.utils.module_utils import warning_rank_0 # Check if Transformer Engine is installed HAVE_TE = False diff --git a/primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py b/primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py index 6b8058ce6..73714fefc 100644 --- a/primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py +++ b/primus/backends/megatron/core/optimizer/fsdp2_bf16_master_weight_optimizer.py @@ -34,7 +34,7 @@ from megatron.core.optimizer.optimizer import MegatronOptimizer, _zero_grad_group_helper from megatron.core.optimizer.optimizer_config import OptimizerConfig -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 if TYPE_CHECKING: from megatron.core.process_groups_config import ProcessGroupCollection diff --git a/primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py b/primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py index 6c00caa9c..04766ac3d 100644 --- a/primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py +++ b/primus/backends/megatron/core/optimizer/fsdp2_fp32_optimizer.py @@ -33,7 +33,7 @@ from megatron.core.optimizer.optimizer import MegatronOptimizer from megatron.core.optimizer.optimizer_config import OptimizerConfig -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 if TYPE_CHECKING: from megatron.core.process_groups_config import ProcessGroupCollection diff --git a/primus/backends/megatron/core/optimizer/zbpp_optimizer.py b/primus/backends/megatron/core/optimizer/zbpp_optimizer.py index c554d3d28..2f50d2413 100644 --- a/primus/backends/megatron/core/optimizer/zbpp_optimizer.py +++ b/primus/backends/megatron/core/optimizer/zbpp_optimizer.py @@ -18,7 +18,7 @@ multi_tensor_applier, ) -from primus.modules.module_utils import log_rank_all +from primus.core.utils.module_utils import log_rank_all class ZeroBubblePPChainedOptimizer(ChainedOptimizer): diff --git a/primus/backends/megatron/core/parallel_state.py b/primus/backends/megatron/core/parallel_state.py index 71b76d7cc..b7e6b7e4f 100644 --- a/primus/backends/megatron/core/parallel_state.py +++ b/primus/backends/megatron/core/parallel_state.py @@ -17,7 +17,7 @@ ) from megatron.core.utils import get_pg_rank, get_pg_size -from primus.modules.trainer.megatron.utils import is_v_schedule_enabled +from primus.backends.megatron.training.utils import is_v_schedule_enabled def is_pp_first_stage(pp_group: torch.distributed.ProcessGroup): diff --git a/primus/backends/megatron/core/pipeline_parallel/pp_visualizer.py b/primus/backends/megatron/core/pipeline_parallel/pp_visualizer.py new file mode 100644 index 000000000..85387b258 --- /dev/null +++ b/primus/backends/megatron/core/pipeline_parallel/pp_visualizer.py @@ -0,0 +1,193 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Pipeline-parallel visualization helpers. + +Records per-iteration forward/backward/weight-grad CUDA events so the offline +``dump_pp_data`` step can emit a JSON timeline consumed by the PP visualizer. +These wrappers are installed by the ``megatron.pp.dump_pp_data`` patch (and by +the Primus-Pipe / ZeroBubble handlers) when ``--dump_pp_data`` is enabled. +""" + +import json +import os + +import torch +from megatron.core import parallel_state + +_GLOBAL_PP_VIS_EVENTS = [] +_GLOBAL_PP_VIS_EVENTS_PER_ITER = None + + +def schedule_wrapper(func): + def wrapper(*args, **kwargs): + global _GLOBAL_PP_VIS_EVENTS_PER_ITER + _GLOBAL_PP_VIS_EVENTS_PER_ITER = { + "start": None, + "end": None, + "memory": None, + "fwd_start": [], + "fwd_end": [], + "fwd_minibatch": [], + "fwd_chunk": [], + "bwd_start": [], + "bwd_end": [], + "bwd_minibatch": [], + "bwd_chunk": [], + "wgrad_start": [], + "wgrad_end": [], + "wgrad_minibatch": [], + "wgrad_chunk": [], + } + + _GLOBAL_PP_VIS_EVENTS_PER_ITER["start"] = torch.cuda.Event(enable_timing=True) + _GLOBAL_PP_VIS_EVENTS_PER_ITER["start"].record() + res = func(*args, **kwargs) + _GLOBAL_PP_VIS_EVENTS_PER_ITER["end"] = torch.cuda.Event(enable_timing=True) + _GLOBAL_PP_VIS_EVENTS_PER_ITER["end"].record() + + _GLOBAL_PP_VIS_EVENTS_PER_ITER["memory"] = torch.cuda.max_memory_reserved() / 1024**3 + + global _GLOBAL_PP_VIS_EVENTS + _GLOBAL_PP_VIS_EVENTS.append(_GLOBAL_PP_VIS_EVENTS_PER_ITER) + + return res + + return wrapper + + +def fwd_bwd_wrapper(func, mode, minibatch=None, chunk=None): + def wrapper(*args, **kwargs): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + res = func(*args, **kwargs) + end.record() + + global _GLOBAL_PP_VIS_EVENTS_PER_ITER + _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_start"].append(start) + _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_end"].append(end) + + if minibatch is not None: + _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_minibatch"].append(minibatch) + if chunk is not None: + _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_chunk"].append(chunk) + return res + + return wrapper + + +def combined_fwd_bwd_wrapper(func, fwd_minibatch, fwd_chunk, bwd_minibatch, bwd_chunk): + """Record a single combined forward+backward call as both an ``fwd`` event + and a ``bwd`` event sharing the same ``[start, end]`` interval. + + Used by ``megatron_combined_fwd_bkwd_handler`` so that nodes collapsed into + a combined FB group still appear in the dump_pp_data output. Without this + the visualizer's per-rank F/B/W totals are heavily under-counted on ranks + that hit the steady state (the F and B halves are interleaved inside + ``combined_forward_backward_step`` and cannot be timed separately). + """ + + def wrapper(*args, **kwargs): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + res = func(*args, **kwargs) + end.record() + + global _GLOBAL_PP_VIS_EVENTS_PER_ITER + _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_start"].append(start) + _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_end"].append(end) + _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_minibatch"].append(fwd_minibatch) + _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_chunk"].append(fwd_chunk) + _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_start"].append(start) + _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_end"].append(end) + _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_minibatch"].append(bwd_minibatch) + _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_chunk"].append(bwd_chunk) + return res + + return wrapper + + +def set_dump_pp_data_patch(): + from megatron.core.pipeline_parallel import schedules + + schedules.forward_step = fwd_bwd_wrapper(schedules.forward_step, "fwd") + schedules.backward_step = fwd_bwd_wrapper(schedules.backward_step, "bwd") + + +def dump_pp_data(args, num_mbs, pp_data_dir): + torch.cuda.synchronize() + + global _GLOBAL_PP_VIS_EVENTS + all_iter_data = {} + for iter_idx, iter_events in enumerate(_GLOBAL_PP_VIS_EVENTS): + iter_data = { + "total": None, + "memory": None, + "fwd_start": [], + "fwd_end": [], + "fwd_minibatch": [], + "fwd_chunk": [], + "bwd_start": [], + "bwd_end": [], + "bwd_minibatch": [], + "bwd_chunk": [], + "wgrad_start": [], + "wgrad_end": [], + "wgrad_minibatch": [], + "wgrad_chunk": [], + } + iter_data["total"] = iter_events["start"].elapsed_time(iter_events["end"]) + iter_data["memory"] = iter_events["memory"] + + for i in range(len(iter_events["fwd_start"])): + for key in ["fwd_start", "fwd_end", "bwd_start", "bwd_end", "wgrad_start", "wgrad_end"]: + if i >= len(iter_events[key]): + continue + event_time = iter_events["start"].elapsed_time(iter_events[key][i]) + iter_data[key].append(event_time) + for key in [ + "fwd_minibatch", + "fwd_chunk", + "bwd_minibatch", + "bwd_chunk", + "wgrad_minibatch", + "wgrad_chunk", + ]: + if i >= len(iter_events[key]): + continue + iter_data[key].append(iter_events[key][i]) + + all_iter_data[iter_idx + 1] = iter_data + + rank = torch.distributed.get_rank() + dp_rank = parallel_state.get_data_parallel_rank() + pp_rank = parallel_state.get_pipeline_model_parallel_rank() + os.makedirs(pp_data_dir, exist_ok=True) + if dp_rank == 0: + log_path = os.path.join(pp_data_dir, f"pp_rank_{pp_rank}.json") + with open(log_path, "w") as f: + json.dump(all_iter_data, f, indent=2) + + if rank == 0: + vp_size = args.virtual_pipeline_model_parallel_size + vp_size = 1 if vp_size is None else vp_size + config_dict = { + "world_size": args.world_size, + "dp_size": args.data_parallel_size, + "tp_size": args.tensor_model_parallel_size, + "ep_size": args.expert_model_parallel_size, + "pp_size": args.pipeline_model_parallel_size, + "vp_size": vp_size, + "num_mbs": num_mbs, + "train_iters": args.train_iters, + } + log_path = os.path.join(pp_data_dir, f"config.json") + with open(log_path, "w") as f: + json.dump(config_dict, f, indent=2) diff --git a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/bwd_handler.py b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/bwd_handler.py index 628f78dc4..8b6eec8b7 100644 --- a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/bwd_handler.py +++ b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/bwd_handler.py @@ -8,6 +8,9 @@ from megatron.core.pipeline_parallel.schedules import backward_step from megatron.training.global_vars import get_args +from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import ( + fwd_bwd_wrapper, +) from primus.core.pipeline_parallel.handler.offload_handler import OFFLOAD_BUFFER from primus.core.pipeline_parallel.handler.wgrad_handler import WGRAD_RUNNING_CACHE from primus.core.pipeline_parallel.scheduler.scheduler_node import ( @@ -15,7 +18,6 @@ SchedulerNode, ) from primus.core.pipeline_parallel.utils import find_prev_node_with_type -from primus.modules.trainer.megatron.utils import fwd_bwd_wrapper def megatron_check_bwd_node_valid(node: SchedulerNode): diff --git a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/combined_fwd_bkwd_handler.py b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/combined_fwd_bkwd_handler.py index 4a61cfd1e..d3ea91072 100644 --- a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/combined_fwd_bkwd_handler.py +++ b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/combined_fwd_bkwd_handler.py @@ -8,6 +8,9 @@ from megatron.core.pipeline_parallel.schedules import deallocate_output_tensor from megatron.training.global_vars import get_args +from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import ( + combined_fwd_bwd_wrapper, +) from primus.backends.megatron.core.pipeline_parallel.primuspipe.handlers.communication_handler import ( batch_p2p_communication_handler, ) @@ -17,7 +20,6 @@ SchedulerNode, ) from primus.core.pipeline_parallel.utils import find_prev_node_with_type -from primus.modules.trainer.megatron.utils import combined_fwd_bwd_wrapper def megatron_check_combined_fwd_bkwd_node_valid(node: SchedulerNode): diff --git a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/fwd_handler.py b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/fwd_handler.py index 374610fd9..b876119c6 100644 --- a/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/fwd_handler.py +++ b/primus/backends/megatron/core/pipeline_parallel/primuspipe/handlers/fwd_handler.py @@ -11,13 +11,15 @@ ) from megatron.training.global_vars import get_args +from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import ( + fwd_bwd_wrapper, +) from primus.core.pipeline_parallel.handler.offload_handler import OFFLOAD_BUFFER from primus.core.pipeline_parallel.scheduler.scheduler_node import ( FuncType, SchedulerNode, ) from primus.core.pipeline_parallel.utils import find_prev_node_with_type -from primus.modules.trainer.megatron.utils import fwd_bwd_wrapper def megatron_check_fwd_node_valid(node: SchedulerNode): diff --git a/primus/backends/megatron/core/pipeline_parallel/primuspipe/pipeline_launcher.py b/primus/backends/megatron/core/pipeline_parallel/primuspipe/pipeline_launcher.py index d196d8420..a0a09eb85 100644 --- a/primus/backends/megatron/core/pipeline_parallel/primuspipe/pipeline_launcher.py +++ b/primus/backends/megatron/core/pipeline_parallel/primuspipe/pipeline_launcher.py @@ -33,7 +33,7 @@ ) from primus.core.pipeline_parallel.scheduler.scheduler import ScheduleRunner from primus.core.pipeline_parallel.scheduler.scheduler_node import FuncType -from primus.modules.module_utils import warning_rank_0 +from primus.core.utils.module_utils import warning_rank_0 class PrimusPipelineParallelLauncher: @@ -309,7 +309,9 @@ def enable_grad_sync(): node.args["pp_group"] = pg_collection.pp if args.dump_pp_data: - from primus.modules.trainer.megatron.utils import schedule_wrapper + from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import ( + schedule_wrapper, + ) schedule_wrapper(self.schedule_runner.run)(self.schedule_table, self.pp_rank) else: diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/offload.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/offload.py index d730f74ad..eb94068e0 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/offload.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/offload.py @@ -15,7 +15,7 @@ import torch from torch.autograd.graph import saved_tensors_hooks -from primus.modules.module_utils import log_rank_all +from primus.core.utils.module_utils import log_rank_all def checksum(tensor): diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/runtime.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/runtime.py index ff67048cb..296304fc6 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/runtime.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/runtime.py @@ -35,10 +35,12 @@ from megatron.core.utils import get_model_config, get_model_type, get_model_xattn from megatron.training import get_args, print_rank_0 +from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import ( + fwd_bwd_wrapper, +) from primus.backends.megatron.training.training import RollbackDataIteratorWrapper from primus.backends.megatron.training.utils import is_second_last_pipeline_stage -from primus.modules.module_utils import log_rank_0, log_rank_all -from primus.modules.trainer.megatron.utils import fwd_bwd_wrapper +from primus.core.utils.module_utils import log_rank_0, log_rank_all from .offload import ActivationStorePool, FakeActivationStore, partial_recompute from .scheduler import ( @@ -454,7 +456,7 @@ def pre_load_batch(self, idx): cnt = (int(offload_time) + 1) * 3 multi_chunks = get_virtual_pipeline_number() > 1 conf = self.iteration_config - from primus.modules.trainer.megatron.pre_trainer import DataLoaderStore + from primus.backends.megatron.data_loader_store import DataLoaderStore count = len(DataLoaderStore.cache) for i in range(cnt): @@ -478,7 +480,7 @@ def pre_load_batch(self, idx): def load_all_batch(self): conf = self.iteration_config multi_chunks = get_virtual_pipeline_number() > 1 - from primus.modules.trainer.megatron.pre_trainer import DataLoaderStore + from primus.backends.megatron.data_loader_store import DataLoaderStore assert len(DataLoaderStore.cache) == 0 for scheduled_node in conf.schedules: @@ -531,7 +533,7 @@ def schedule_f_impl(self, scheduled_node: ScheduledNode): mem_before = torch.cuda.memory_allocated() set_seq_split_idx(scheduled_node.seq_split_idx) - from primus.modules.trainer.megatron.pre_trainer import DataLoaderStore + from primus.backends.megatron.data_loader_store import DataLoaderStore if len(DataLoaderStore.cache) == 0: DataLoaderStore.push(conf.data_iterator[scheduled_node.chunk], vp_stage=scheduled_node.chunk) diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/basic1f1b.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/basic1f1b.py index a02982606..d405c2b79 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/basic1f1b.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/basic1f1b.py @@ -7,7 +7,7 @@ # See LICENSE for license information. ############################################################################### -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 from .graph import BW, F, GraphConfig, ScheduledNode diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/communication.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/communication.py index c28b679f0..e59ad7c04 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/communication.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/communication.py @@ -14,7 +14,7 @@ from megatron.training.global_vars import get_args -from primus.modules.module_utils import log_rank_all +from primus.core.utils.module_utils import log_rank_all from .graph import BW, B, CommDirection, F, FuncType, GraphConfig, ScheduledNode diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/group_interleaved_1f1b.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/group_interleaved_1f1b.py index 112419319..b66ac5e25 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/group_interleaved_1f1b.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/group_interleaved_1f1b.py @@ -13,7 +13,7 @@ from enum import Enum from typing import List -from primus.modules.module_utils import log_rank_0, log_rank_all +from primus.core.utils.module_utils import log_rank_0, log_rank_all class PassType(Enum): diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/offloading.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/offloading.py index 4758ed9c6..a5049e12f 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/offloading.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/offloading.py @@ -11,7 +11,7 @@ from dataclasses import dataclass from typing import List -from primus.modules.module_utils import log_rank_all +from primus.core.utils.module_utils import log_rank_all from .graph import BW, B, F, FuncType, GraphConfig, NodeKey, ScheduledNode, W diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/passes.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/passes.py index 9edaf49ba..180a216de 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/passes.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/passes.py @@ -12,7 +12,7 @@ from megatron.training.global_vars import get_args -from primus.modules.module_utils import log_rank_all +from primus.core.utils.module_utils import log_rank_all from .communication import ( add_communication_nodes_without_sorting, diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/v_auto_schedule.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/v_auto_schedule.py index e35e9a9b4..c6388cfbe 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/v_auto_schedule.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/v_auto_schedule.py @@ -11,7 +11,7 @@ from collections import Counter, deque from dataclasses import dataclass -from primus.modules.module_utils import log_rank_all +from primus.core.utils.module_utils import log_rank_all @dataclass(eq=True, frozen=True) diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/vpp.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/vpp.py index 2ca912c43..9af46992b 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/vpp.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/vpp.py @@ -7,7 +7,7 @@ # See LICENSE for license information. ############################################################################### -from primus.modules.module_utils import log_rank_all +from primus.core.utils.module_utils import log_rank_all from .graph import BW, F, GraphConfig, ScheduledNode diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zb.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zb.py index b0ea0c47d..2a450220b 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zb.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zb.py @@ -13,7 +13,7 @@ import pulp import torch -from primus.modules.module_utils import log_rank_all +from primus.core.utils.module_utils import log_rank_all from .graph import FuncType, GraphConfig, ScheduledNode diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv.py index 6ebfc15f3..203dae691 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv.py @@ -9,7 +9,7 @@ from collections import deque -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 from .graph import B, F, ScheduledNode, W diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv_greedy.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv_greedy.py index ad68b5e97..aea98f32c 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv_greedy.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/scheduler/zbv_greedy.py @@ -10,7 +10,7 @@ # Implementation of vhalf and vmin schedules of Pipeline Parallelism # with Controllable Memory (https://arxiv.org/abs/2405.15362) # The reordering is based on a greedy algorithm. -from primus.modules.module_utils import log_rank_all +from primus.core.utils.module_utils import log_rank_all from .graph import B, F, ScheduledNode, W diff --git a/primus/backends/megatron/core/pipeline_parallel/zerobubble/zbpp_utils.py b/primus/backends/megatron/core/pipeline_parallel/zerobubble/zbpp_utils.py index 7ef4315e0..6d7a36ddf 100644 --- a/primus/backends/megatron/core/pipeline_parallel/zerobubble/zbpp_utils.py +++ b/primus/backends/megatron/core/pipeline_parallel/zerobubble/zbpp_utils.py @@ -15,7 +15,9 @@ from megatron.core import parallel_state from megatron.training import get_args -from primus.modules.trainer.megatron.utils import fwd_bwd_wrapper +from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import ( + fwd_bwd_wrapper, +) def add_zero_bubble_args(parser): diff --git a/primus/backends/megatron/core/transformer/pipeline_parallel_layer_layout.py b/primus/backends/megatron/core/transformer/pipeline_parallel_layer_layout.py index f1155098e..589ce54f5 100644 --- a/primus/backends/megatron/core/transformer/pipeline_parallel_layer_layout.py +++ b/primus/backends/megatron/core/transformer/pipeline_parallel_layer_layout.py @@ -6,7 +6,7 @@ PipelineParallelLayerLayout, ) -from primus.modules.trainer.megatron.utils import is_v_schedule_enabled +from primus.backends.megatron.training.utils import is_v_schedule_enabled class PrimusPipelineParallelLayerLayout(PipelineParallelLayerLayout): diff --git a/primus/backends/megatron/core/transformer/transformer_layer.py b/primus/backends/megatron/core/transformer/transformer_layer.py index e544845a8..8fe9f4f5c 100644 --- a/primus/backends/megatron/core/transformer/transformer_layer.py +++ b/primus/backends/megatron/core/transformer/transformer_layer.py @@ -13,7 +13,7 @@ from megatron.core.transformer.enums import LayerType from megatron.core.transformer.transformer_config import TransformerConfig -from primus.modules.trainer.megatron.utils import is_v_schedule_enabled +from primus.backends.megatron.training.utils import is_v_schedule_enabled def get_transformer_layer_offset( diff --git a/primus/backends/megatron/data_loader_store.py b/primus/backends/megatron/data_loader_store.py new file mode 100644 index 000000000..02372e946 --- /dev/null +++ b/primus/backends/megatron/data_loader_store.py @@ -0,0 +1,132 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Megatron batch-loading helpers used by the pipeline schedules. + +``get_batch_func`` builds a micro-batch honoring TP/CP sharding (and the +ZeroBubble sequence-split path), and ``DataLoaderStore`` provides the +push/pop cache the ZeroBubble runtime and the ``forward_step`` patch use to +pre-fetch batches (optionally on a dedicated H2D stream). +""" + +import collections + +import torch +from megatron.core.utils import StragglerDetector +from megatron.training import get_args, get_timers +from megatron.training.utils import ( + get_batch_on_this_cp_rank, + get_batch_on_this_tp_rank, + is_first_or_last_pipeline_stage, +) + +stimer = StragglerDetector() + +mb_batch = None + + +def get_batch_func(data_iterator, vp_stage=None): + # TODO: this is pretty hacky, find a better way + if not is_first_or_last_pipeline_stage(vp_stage): + return None, None, None, None, None + + # assert data_iterator is not None, f"data_iterator is None vp_stage: {vp_stage}" + # get batches based on the TP rank you are on + batch = get_batch_on_this_tp_rank(data_iterator) + + # slice batch along sequence dimension for context parallelism + args = get_args() + + if args.patch_zero_bubble: + from primus.backends.megatron.core.pipeline_parallel.zerobubble.zbpp_vars import ( + get_seq_split_idx, + ) + + global mb_batch + # "or 0" to support original 1f1b and interleaved-1f1b in schedules.py + seq_split_idx = get_seq_split_idx() or 0 + if seq_split_idx == 0: + # get batches based on the TP rank you are on + mb_batch = get_batch_on_this_tp_rank(data_iterator) + assert ( + mb_batch["attention_mask"] is None + ), "attention_mask should be None, please enable --no-create-attention-mask-in-dataloader" + batch = {} + for k in mb_batch.keys(): + v = mb_batch[k] + if v is None: + batch[k] = v + continue + + assert v.shape[1] % get_args().num_seq_splits == 0, f"{k} size {v.shape}" + start_idx = seq_split_idx * v.shape[1] // get_args().num_seq_splits + end_idx = (seq_split_idx + 1) * v.shape[1] // get_args().num_seq_splits + if len(v.shape) > 2: + batch[k] = v[:, start_idx:end_idx, :].contiguous() + else: + batch[k] = v[:, start_idx:end_idx].contiguous() + + if args.context_parallel_size > 1 and args.enable_primus_turbo and args.use_turbo_attention: + try: + from primus.backends.megatron.core.utils import ( + produce_attention_sharder, + shard_batch_on_this_cp_rank, + ) + except: + raise ImportError("Module 'primus_turbo' may not installed. Please install it") + sharder = produce_attention_sharder(args.cp_comm_type) + batch = shard_batch_on_this_cp_rank(sharder, batch) + else: + batch = get_batch_on_this_cp_rank(batch) + + # Return a stable, explicitly-ordered 5-tuple so both this path and the + # early not-first/last-stage return path have the same shape/type. The keys + # match megatron's ``get_batch_on_this_tp_rank`` and the unpack order at all + # call sites (tokens, labels, loss_mask, attention_mask, position_ids). + return ( + batch["tokens"], + batch["labels"], + batch["loss_mask"], + batch["attention_mask"], + batch["position_ids"], + ) + + +class DataLoaderStore: + cache = collections.deque() + + @classmethod + def push(cls, data_iterator, h2d_stream=False, vp_stage=None): + timers = get_timers() + # Get the batch. + timers("batch-generator", log_level=2).start() + global stimer + + with stimer(bdata=True): + if h2d_stream: + from primus.backends.megatron.core.pipeline_parallel.zerobubble.offload import ( + get_offload_h2d_stream, + ) + + load_event = torch.cuda.Event() + original_stream = torch.cuda.current_stream() + with torch.cuda.stream(get_offload_h2d_stream()): + data = get_batch_func(data_iterator, vp_stage) + for x in data: + if x is not None: + x.record_stream(original_stream) + load_event.record() + cls.cache.append((data, load_event)) + else: + cls.cache.append((get_batch_func(data_iterator, vp_stage), None)) + timers("batch-generator").stop() + + @classmethod + def pop(cls): + data, load_event = cls.cache.popleft() + if load_event: + load_event.wait() + return data diff --git a/primus/backends/megatron/megatron_adapter.py b/primus/backends/megatron/megatron_adapter.py index 66574ec00..698d95e37 100644 --- a/primus/backends/megatron/megatron_adapter.py +++ b/primus/backends/megatron/megatron_adapter.py @@ -11,7 +11,7 @@ from primus.backends.megatron.argument_builder import MegatronArgBuilder from primus.core.backend.backend_adapter import BackendAdapter from primus.core.backend.backend_registry import BackendRegistry -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class MegatronAdapter(BackendAdapter): @@ -82,7 +82,6 @@ def _load_trainer_class_by_name(self, trainer_class: str): possible_paths = [ f"primus.backends.megatron.{trainer_class.lower()}.{trainer_class}", - f"primus.modules.trainer.megatron.{trainer_class.lower()}.{trainer_class}", f"primus.backends.megatron.{trainer_class}", ] diff --git a/primus/backends/megatron/megatron_base_trainer.py b/primus/backends/megatron/megatron_base_trainer.py index 89cbb3813..de0d17587 100644 --- a/primus/backends/megatron/megatron_base_trainer.py +++ b/primus/backends/megatron/megatron_base_trainer.py @@ -14,7 +14,7 @@ ) from primus.backends.megatron.training.mlflow_setup import upload_mlflow_artifacts from primus.core.trainer.base_trainer import BaseTrainer -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 class MegatronBaseTrainer(BaseTrainer): diff --git a/primus/backends/megatron/megatron_pretrain_trainer.py b/primus/backends/megatron/megatron_pretrain_trainer.py index 4955491d8..b7daff509 100644 --- a/primus/backends/megatron/megatron_pretrain_trainer.py +++ b/primus/backends/megatron/megatron_pretrain_trainer.py @@ -5,12 +5,62 @@ ############################################################################### from primus.backends.megatron.megatron_base_trainer import MegatronBaseTrainer -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class MegatronPretrainTrainer(MegatronBaseTrainer): """Trainer for Megatron-LM pre-training.""" + def setup_model_only(self): + """Initialize Megatron and build the model WITHOUT running the training loop. + + A general, training-neutral capability: mirrors the front of + ``megatron.training.pretrain`` (``initialize_megatron`` followed by + ``setup_model_and_optimizer``) to construct a model identical to the real + training path, but stops before datasets / the train loop. Useful for any + "build the model only" scenario (offline profiling, layer benchmarking, + model inspection); performance/memory projection is the current consumer. + + Prerequisites (handled by the runtime before calling this): ``setup()`` + has patched ``parse_args`` to return ``self.backend_args`` and set the + Primus global vars, and the build_args/setup/before_train patch phases + have been applied. Returns the built model (a list of model chunks, as + produced by megatron's ``get_model``) and also stores it on ``self.model``. + """ + log_rank_0("Setting up Megatron model only (no training loop)...") + + from megatron.core.enums import ModelType + from megatron.training.initialize import initialize_megatron + from megatron.training.training import setup_model_and_optimizer + + from primus.core.utils.import_utils import get_model_provider + + # Determine model type (gpt or mamba) from backend_args + model_type = getattr(self.backend_args, "model_type", "gpt") + log_rank_0(f"-detected model_type: {model_type}") + + # parse_args was patched in setup() to return backend_args, so + # initialize_megatron consumes the Primus-configured arguments (same as + # the front of megatron's pretrain()). + initialize_megatron(args_defaults={"tokenizer_type": "GPT2BPETokenizer"}) + + # Get model provider with correct model_type (reuse the core runtime helper) + if model_type != "gpt": + model_provider = get_model_provider(model_type=model_type) + else: + model_provider = get_model_provider() + log_rank_0(f"-model_provider: {model_provider}") + + model, optimizer, opt_param_scheduler = setup_model_and_optimizer( + model_provider, ModelType.encoder_or_decoder, checkpointing_context={} + ) + self.model = model + self.optimizer = optimizer + self.opt_param_scheduler = opt_param_scheduler + + log_rank_0("Megatron model-only setup completed.") + return model + def get_forward_step(self): """ Return forward step function for training loop. @@ -151,7 +201,7 @@ def patched_get_forward_backward_func(*args, **kwargs): try: m_args = get_megatron_args() if getattr(m_args, "dump_pp_data", False): - from primus.modules.trainer.megatron.utils import ( + from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import ( schedule_wrapper, set_dump_pp_data_patch, ) @@ -194,7 +244,9 @@ def patched_get_forward_backward_func(*args, **kwargs): get_num_microbatches, ) - from primus.modules.trainer.megatron.utils import dump_pp_data + from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import ( + dump_pp_data, + ) pp_data_dir = os.environ.get("DUMP_PP_DIR", "output/pp_data") dump_pp_data(megatron_args, get_num_microbatches(), pp_data_dir) diff --git a/primus/backends/megatron/megatron_sft_trainer.py b/primus/backends/megatron/megatron_sft_trainer.py index 0f95eebc9..46493434f 100644 --- a/primus/backends/megatron/megatron_sft_trainer.py +++ b/primus/backends/megatron/megatron_sft_trainer.py @@ -9,7 +9,7 @@ from typing import Any from primus.backends.megatron.megatron_base_trainer import MegatronBaseTrainer -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class MegatronSFTTrainer(MegatronBaseTrainer): diff --git a/primus/backends/megatron/patches/__init__.py b/primus/backends/megatron/patches/__init__.py index 6ae718c8b..dd51eaf5e 100644 --- a/primus/backends/megatron/patches/__init__.py +++ b/primus/backends/megatron/patches/__init__.py @@ -31,7 +31,7 @@ import pkgutil # from primus.core.patches import run_patches -# from primus.modules.module_utils import log_rank_0 +# from primus.core.utils.module_utils import log_rank_0 def _auto_import_patch_modules() -> None: diff --git a/primus/backends/megatron/patches/args/checkpoint_path_patches.py b/primus/backends/megatron/patches/args/checkpoint_path_patches.py index b0736551f..bdb56eef9 100644 --- a/primus/backends/megatron/patches/args/checkpoint_path_patches.py +++ b/primus/backends/megatron/patches/args/checkpoint_path_patches.py @@ -7,7 +7,7 @@ import os from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/args/data_path_split_patches.py b/primus/backends/megatron/patches/args/data_path_split_patches.py index 4c326206e..a96d40dec 100644 --- a/primus/backends/megatron/patches/args/data_path_split_patches.py +++ b/primus/backends/megatron/patches/args/data_path_split_patches.py @@ -5,7 +5,7 @@ ############################################################################### from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_kv_rank_0 +from primus.core.utils.module_utils import log_kv_rank_0 def _normalize_data_path(path_value): diff --git a/primus/backends/megatron/patches/args/hsdp_args_patches.py b/primus/backends/megatron/patches/args/hsdp_args_patches.py index da1d76bc5..57b796344 100644 --- a/primus/backends/megatron/patches/args/hsdp_args_patches.py +++ b/primus/backends/megatron/patches/args/hsdp_args_patches.py @@ -5,7 +5,7 @@ ############################################################################### from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_kv_rank_0, log_rank_0 +from primus.core.utils.module_utils import log_kv_rank_0, log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/args/iterations_to_skip_default_patches.py b/primus/backends/megatron/patches/args/iterations_to_skip_default_patches.py index 525f66ade..b8651b9aa 100644 --- a/primus/backends/megatron/patches/args/iterations_to_skip_default_patches.py +++ b/primus/backends/megatron/patches/args/iterations_to_skip_default_patches.py @@ -5,7 +5,7 @@ ############################################################################### from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_kv_rank_0 +from primus.core.utils.module_utils import log_kv_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/args/logging_level_patches.py b/primus/backends/megatron/patches/args/logging_level_patches.py index cfb39c613..7477ed236 100644 --- a/primus/backends/megatron/patches/args/logging_level_patches.py +++ b/primus/backends/megatron/patches/args/logging_level_patches.py @@ -5,7 +5,7 @@ ############################################################################### from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/args/mock_data_patches.py b/primus/backends/megatron/patches/args/mock_data_patches.py index 25253d4b9..2089b9b3f 100644 --- a/primus/backends/megatron/patches/args/mock_data_patches.py +++ b/primus/backends/megatron/patches/args/mock_data_patches.py @@ -5,7 +5,7 @@ ############################################################################### from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_kv_rank_0, log_rank_0 +from primus.core.utils.module_utils import log_kv_rank_0, log_rank_0 DEFAULT_NULL_TOKENIZER_VOCAB_SIZE = 131072 diff --git a/primus/backends/megatron/patches/args/moe_layer_freq_patches.py b/primus/backends/megatron/patches/args/moe_layer_freq_patches.py index fa15d0998..7e19daef4 100644 --- a/primus/backends/megatron/patches/args/moe_layer_freq_patches.py +++ b/primus/backends/megatron/patches/args/moe_layer_freq_patches.py @@ -5,7 +5,7 @@ ############################################################################### from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_kv_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_kv_rank_0, warning_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/args/rocm_arg_validation.py b/primus/backends/megatron/patches/args/rocm_arg_validation.py new file mode 100644 index 000000000..e68251f8d --- /dev/null +++ b/primus/backends/megatron/patches/args/rocm_arg_validation.py @@ -0,0 +1,197 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""ROCm-specific Megatron argument validation. + +``validate_args_on_rocm`` is invoked by the ``megatron.validate_args`` patch +after Megatron's own ``validate_args`` runs. It enforces ROCm/Primus-Turbo +constraints (deterministic-mode env vars, Turbo FP8/FP4 recipes, sync-free MoE +auto-config, DeepEP restrictions, ...). +""" + +import inspect +import os + +import torch + +from primus.core.utils import logger + + +def is_last_rank(): + return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1) + + +def print_rank_last(msg): + """If distributed is initialized, print only on last rank.""" + log_func = logger.info_with_caller + + caller = inspect.stack()[1] + caller_frame = caller.frame + function_name = caller_frame.f_code.co_name + module_name = caller_frame.f_globals["__name__"].split(".")[-1] + line = caller.lineno + + if torch.distributed.is_initialized(): + if is_last_rank(): + log_func(msg, module_name, function_name, line) + else: + log_func(msg, module_name, function_name, line) + + +def _get_sync_free_moe_options(args) -> dict: + stage = args.turbo_sync_free_moe_stage + + if stage > 3 or stage < 0: + raise ValueError("turbo_sync_free_moe_stage only support [0-3]") + + sync_free_moe = { + 1: { + "moe_use_fused_router_with_aux_score": True, + "moe_permute_fusion": True, + "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False, + }, + 2: { + "moe_use_fused_router_with_aux_score": True, + "use_turbo_deepep": True, + "moe_permute_fusion": True, + "use_turbo_grouped_gemm": True, + "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False, + }, + 3: { + "moe_use_fused_router_with_aux_score": True, + "use_turbo_deepep": True, + "moe_permute_fusion": True, + "use_turbo_grouped_gemm": True, + "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False, + "use_turbo_fused_act_with_probs": True, + }, + } + + return sync_free_moe[stage] + + +# FSDP2 custom optimizer selection flags. Each one monkeypatches +# get_megatron_optimizer at the same priority (50), so at most one may be set. +_FSDP2_OPTIMIZER_FLAGS = ( + "use_fsdp2_fp32_param_optimizer", + "use_fsdp2_bf16_master_weight_optimizer", +) + + +def validate_fsdp2_optimizer_exclusivity(args) -> None: + """Ensure at most one FSDP2 custom optimizer flag is enabled. + + Enabling more than one would silently let whichever optimizer patch applies + last win the monkeypatch, so raise ValueError to fail loudly at + arg-validation time (before training starts). + """ + enabled = [flag for flag in _FSDP2_OPTIMIZER_FLAGS if getattr(args, flag, False)] + if len(enabled) > 1: + raise ValueError( + "Conflicting FSDP2 optimizer selection: at most one of " + f"{list(_FSDP2_OPTIMIZER_FLAGS)} may be enabled, but got {enabled}. " + "Enable exactly one." + ) + + +def validate_args_on_rocm(args): + # Deterministic mode + if args.deterministic_mode: + # NOTE: Some environment variables affect deterministic mode on ROCm. Need to do extra check. + NON_DETERMINISTIC_ENVS = { + "TORCH_COMPILE_DISABLE": "1", + "ROCBLAS_DEFAULT_ATOMICS_MODE": "0", + "PRIMUS_TURBO_AUTO_TUNE": "0", + "PRIMUS_DETERMINISTIC": "1", + } + # NOTE: Some version triton compile exist potential racing condition issue. + for env, value in NON_DETERMINISTIC_ENVS.items(): + assert ( + os.environ.get(env, None) == value + ), f"{env} must be set to {value} in deterministic mode but got {os.environ.get(env, None)} instead." + + # Set fill_uninitialized_memory to False to avoid calling extra fill kernel in deterministic mode. + torch.utils.deterministic.fill_uninitialized_memory = False + + assert not getattr( + args, "use_turbo_parallel_linear", False + ), "use_turbo_parallel_linear has been removed; please use use_turbo_gemm instead." + + validate_fsdp2_optimizer_exclusivity(args) + + use_turbo_gemm = getattr(args, "use_turbo_gemm", False) + # Turbo FP8 linear check + if args.fp8 and use_turbo_gemm: + support_fp8_recipe = ["tensorwise", "blockwise", "mxfp8"] + assert ( + args.fp8_recipe in support_fp8_recipe + ), f"{args.fp8_recipe} recipe is not support when enable `use_turbo_gemm`." + + # Turbo FP4 linear check + if args.fp4 and use_turbo_gemm: + support_fp4_recipe = ["mxfp4"] + assert ( + args.fp4_recipe in support_fp4_recipe + ), f"{args.fp4_recipe} recipe is not support when enable `use_turbo_gemm`." + + # NOTE: mxfp8 environment variable must be set to 1 to enable mxfp8 recipe on ROCm. + if args.fp8_recipe == "mxfp8": + assert ( + os.getenv("NVTE_ROCM_ENABLE_MXFP8", "0") == "1" + ), "Please set `NVTE_ROCM_ENABLE_MXFP8=1` to enable `mxfp8` recipe." + + # dump pp data + if args.dump_pp_data and args.pipeline_model_parallel_size == 1: + args.dump_pp_data = False + print_rank_last(f"Disable args.dump_pp_data since args.pipeline_model_parallel_size=1") + + # PrimusTurboGroupedMLP no longer depends on legacy GroupedMLP; the two + # flags are mutually exclusive when turbo is enabled. + assert not getattr( + args, "use_turbo_grouped_mlp", False + ), "use_turbo_grouped_mlp has been removed; please use use_turbo_grouped_gemm instead." + use_turbo_grouped_gemm = getattr(args, "use_turbo_grouped_gemm", False) + if use_turbo_grouped_gemm: + if getattr(args, "moe_use_legacy_grouped_gemm", False): + raise ValueError( + "use_turbo_grouped_gemm=True is incompatible with moe_use_legacy_grouped_gemm=True. " + "please set moe_use_legacy_grouped_gemm=False." + ) + + # sync-free MoE + if args.turbo_sync_free_moe_stage > 0: + assert args.enable_primus_turbo, "Please set `enable_primus_turbo=True` to enable sync-free MoE." + + if args.turbo_sync_free_moe_stage > 1 and not use_turbo_grouped_gemm: + raise ValueError( + "Sync-Free MoE stage 2 or 3 require PrimusTurboGroupedLinear, please set `use_turbo_grouped_gemm=True`" + ) + options = _get_sync_free_moe_options(args) + print_rank_last( + f"========== Enable Sync-Free MoE Stage {args.turbo_sync_free_moe_stage} (Auto-Enabled Options) ==========" + ) + for flag, value in options.items(): + dots = "." * (73 - len(flag) - len(str(value))) + print_rank_last(f"{flag}{dots}{value}") + setattr(args, flag, value) + print_rank_last( + f"========== Enable Sync-Free MoE Stage {args.turbo_sync_free_moe_stage} (Auto-Enabled Options) ==========" + ) + + # turbo deepep + if args.use_turbo_deepep: + assert ( + not args.moe_shared_expert_overlap + ), "DeepEP not support moe_shared_expert_overlap, please set `moe_shared_expert_overlap=False`." + assert ( + args.moe_router_dtype == "fp32" + ), "DeepEP only supports float32 probs, please set `moe_router_dtype=fp32`" + if ( + args.expert_model_parallel_size >= 16 + and os.getenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND", "TURBO") == "TURBO" + ): + # Turbo DeepEP is not supported for CUs > 32 when using internode dispatch/combine. + assert args.turbo_deepep_num_cu <= 32, "Set `turbo_deepep_num_cu<=32` when using ep_size >= 16." diff --git a/primus/backends/megatron/patches/args/sequence_parallel_tp1_patches.py b/primus/backends/megatron/patches/args/sequence_parallel_tp1_patches.py index e0d022304..90f732e9f 100644 --- a/primus/backends/megatron/patches/args/sequence_parallel_tp1_patches.py +++ b/primus/backends/megatron/patches/args/sequence_parallel_tp1_patches.py @@ -5,7 +5,7 @@ ############################################################################### from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_kv_rank_0, log_rank_0 +from primus.core.utils.module_utils import log_kv_rank_0, log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/args/tensorboard_path_patches.py b/primus/backends/megatron/patches/args/tensorboard_path_patches.py index f1a3fbb27..98c77c70b 100644 --- a/primus/backends/megatron/patches/args/tensorboard_path_patches.py +++ b/primus/backends/megatron/patches/args/tensorboard_path_patches.py @@ -7,7 +7,7 @@ import os from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/args/validate_args_patches.py b/primus/backends/megatron/patches/args/validate_args_patches.py index b389899f9..2f1e0a622 100644 --- a/primus/backends/megatron/patches/args/validate_args_patches.py +++ b/primus/backends/megatron/patches/args/validate_args_patches.py @@ -14,7 +14,7 @@ import inspect from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 # --------------------------------------------------------------------------- # Base wrapper — always active, must be registered first @@ -34,7 +34,9 @@ def patch_validate_args(ctx: PatchContext): import megatron.training.arguments as megatron_args import megatron.training.initialize as megatron_init - from primus.modules.trainer.megatron.utils import validate_args_on_rocm + from primus.backends.megatron.patches.args.rocm_arg_validation import ( + validate_args_on_rocm, + ) megatron_args._primus_original_validate_args = megatron_args.validate_args diff --git a/primus/backends/megatron/patches/args/wandb_config_patches.py b/primus/backends/megatron/patches/args/wandb_config_patches.py index c26e8307b..c1dc58879 100644 --- a/primus/backends/megatron/patches/args/wandb_config_patches.py +++ b/primus/backends/megatron/patches/args/wandb_config_patches.py @@ -7,7 +7,7 @@ import os from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_kv_rank_0, log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_kv_rank_0, log_rank_0, warning_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/args_compat_patches.py b/primus/backends/megatron/patches/args_compat_patches.py index c443fee69..80e67838c 100644 --- a/primus/backends/megatron/patches/args_compat_patches.py +++ b/primus/backends/megatron/patches/args_compat_patches.py @@ -23,7 +23,7 @@ from typing import List from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 _CANONICAL_ADD_FN_NAMES: List[str] = [ "_add_network_size_args", diff --git a/primus/backends/megatron/patches/build_model_patches.py b/primus/backends/megatron/patches/build_model_patches.py index 601524c8b..d81906d80 100644 --- a/primus/backends/megatron/patches/build_model_patches.py +++ b/primus/backends/megatron/patches/build_model_patches.py @@ -20,7 +20,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/checkpoint_patches.py b/primus/backends/megatron/patches/checkpoint_patches.py index 24755b394..1461eb50a 100644 --- a/primus/backends/megatron/patches/checkpoint_patches.py +++ b/primus/backends/megatron/patches/checkpoint_patches.py @@ -12,7 +12,7 @@ """ from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/dataloader_patch.py b/primus/backends/megatron/patches/dataloader_patch.py index 54735dc94..774cda1d3 100644 --- a/primus/backends/megatron/patches/dataloader_patch.py +++ b/primus/backends/megatron/patches/dataloader_patch.py @@ -23,7 +23,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 # Guard against double-patching (e.g. when the patch system runs twice in tests). _PATCHED_ATTR = "_primus_dataloader_mp_context_patched" diff --git a/primus/backends/megatron/patches/distributed_init_patches.py b/primus/backends/megatron/patches/distributed_init_patches.py index 7d7da8b92..b86a8b3b5 100644 --- a/primus/backends/megatron/patches/distributed_init_patches.py +++ b/primus/backends/megatron/patches/distributed_init_patches.py @@ -23,7 +23,7 @@ import torch from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/env_patches.py b/primus/backends/megatron/patches/env_patches.py index c8c9ae4c0..52e3a5781 100644 --- a/primus/backends/megatron/patches/env_patches.py +++ b/primus/backends/megatron/patches/env_patches.py @@ -13,7 +13,7 @@ import os from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_kv_rank_0 +from primus.core.utils.module_utils import log_kv_rank_0 # ============================================================================ # CUDA Device Configuration diff --git a/primus/backends/megatron/patches/evaluate_patches.py b/primus/backends/megatron/patches/evaluate_patches.py index 2c1932e5e..108c19d1f 100644 --- a/primus/backends/megatron/patches/evaluate_patches.py +++ b/primus/backends/megatron/patches/evaluate_patches.py @@ -13,7 +13,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py b/primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py index dbbbcc056..e1180b9b9 100644 --- a/primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py +++ b/primus/backends/megatron/patches/fsdp2_fp8_cache_patches.py @@ -17,7 +17,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _needs_fp8_cache_update(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/gpt_decoder_layer_specs_patches.py b/primus/backends/megatron/patches/gpt_decoder_layer_specs_patches.py index 3847f0022..378d839cf 100644 --- a/primus/backends/megatron/patches/gpt_decoder_layer_specs_patches.py +++ b/primus/backends/megatron/patches/gpt_decoder_layer_specs_patches.py @@ -11,7 +11,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/mamba_rocm_patches.py b/primus/backends/megatron/patches/mamba_rocm_patches.py index 9d6243247..31caa2c94 100644 --- a/primus/backends/megatron/patches/mamba_rocm_patches.py +++ b/primus/backends/megatron/patches/mamba_rocm_patches.py @@ -15,7 +15,7 @@ import torch from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _is_rocm(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/megatron_fsdp_patches.py b/primus/backends/megatron/patches/megatron_fsdp_patches.py index 5a3b48ea8..a6fb3b921 100644 --- a/primus/backends/megatron/patches/megatron_fsdp_patches.py +++ b/primus/backends/megatron/patches/megatron_fsdp_patches.py @@ -15,7 +15,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/mla_patches.py b/primus/backends/megatron/patches/mla_patches.py index c138112f4..e25ac2d2d 100644 --- a/primus/backends/megatron/patches/mla_patches.py +++ b/primus/backends/megatron/patches/mla_patches.py @@ -12,7 +12,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/moe_patches/deprecated_layer_patches.py b/primus/backends/megatron/patches/moe_patches/deprecated_layer_patches.py index b2259599d..79db6b803 100644 --- a/primus/backends/megatron/patches/moe_patches/deprecated_layer_patches.py +++ b/primus/backends/megatron/patches/moe_patches/deprecated_layer_patches.py @@ -13,7 +13,7 @@ import sys from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/moe_patches/permute_fusion_patches.py b/primus/backends/megatron/patches/moe_patches/permute_fusion_patches.py index d00d2b84b..b889b64c8 100644 --- a/primus/backends/megatron/patches/moe_patches/permute_fusion_patches.py +++ b/primus/backends/megatron/patches/moe_patches/permute_fusion_patches.py @@ -11,7 +11,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py b/primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py index d072e3d02..8d08d4af1 100644 --- a/primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py +++ b/primus/backends/megatron/patches/moe_patches/skip_identity_sort_patches.py @@ -33,7 +33,7 @@ import torch from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 # Cache the identity decision keyed on the index tensor's ``id()``. The topology # index tensors persist for the dispatcher's lifetime, so caching by identity is diff --git a/primus/backends/megatron/patches/moe_patches/topk_router_patches.py b/primus/backends/megatron/patches/moe_patches/topk_router_patches.py index 77f4c7e5c..6fe2e90e6 100644 --- a/primus/backends/megatron/patches/moe_patches/topk_router_patches.py +++ b/primus/backends/megatron/patches/moe_patches/topk_router_patches.py @@ -24,7 +24,7 @@ class MoESubmodules: import sys from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/mp_sync_skip_patches.py b/primus/backends/megatron/patches/mp_sync_skip_patches.py index b7b9170c2..c32ca3861 100644 --- a/primus/backends/megatron/patches/mp_sync_skip_patches.py +++ b/primus/backends/megatron/patches/mp_sync_skip_patches.py @@ -24,7 +24,7 @@ import torch from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _is_pure_dp(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/muon_optimizer_patches.py b/primus/backends/megatron/patches/muon_optimizer_patches.py index 2d96dfa39..7ecd2bc19 100644 --- a/primus/backends/megatron/patches/muon_optimizer_patches.py +++ b/primus/backends/megatron/patches/muon_optimizer_patches.py @@ -18,7 +18,7 @@ import inspect from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/optimizer_patches.py b/primus/backends/megatron/patches/optimizer_patches.py index 5667c7cd9..69e2958a9 100644 --- a/primus/backends/megatron/patches/optimizer_patches.py +++ b/primus/backends/megatron/patches/optimizer_patches.py @@ -12,7 +12,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/parallelism/forward_step_patches.py b/primus/backends/megatron/patches/parallelism/forward_step_patches.py index ec83eb9b0..dddf9c033 100644 --- a/primus/backends/megatron/patches/parallelism/forward_step_patches.py +++ b/primus/backends/megatron/patches/parallelism/forward_step_patches.py @@ -17,7 +17,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _needs_forward_step_patch(ctx: PatchContext) -> bool: @@ -49,7 +49,7 @@ def _patched_forward_step(data_iterator, model, return_schedule_plan=False): if not args.patch_zero_bubble: return _original_forward_step(data_iterator, model, return_schedule_plan) - from primus.modules.trainer.megatron.pre_trainer import DataLoaderStore + from primus.backends.megatron.data_loader_store import DataLoaderStore timers = None try: diff --git a/primus/backends/megatron/patches/parallelism/linear_grad_split_patches.py b/primus/backends/megatron/patches/parallelism/linear_grad_split_patches.py index a65a9cfcb..92e74ac72 100644 --- a/primus/backends/megatron/patches/parallelism/linear_grad_split_patches.py +++ b/primus/backends/megatron/patches/parallelism/linear_grad_split_patches.py @@ -9,7 +9,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/parallelism/pipeline_parallel_layout_patches.py b/primus/backends/megatron/patches/parallelism/pipeline_parallel_layout_patches.py index 65ee1695e..f38dde4aa 100644 --- a/primus/backends/megatron/patches/parallelism/pipeline_parallel_layout_patches.py +++ b/primus/backends/megatron/patches/parallelism/pipeline_parallel_layout_patches.py @@ -9,7 +9,7 @@ """ from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/parallelism/schedule_patches.py b/primus/backends/megatron/patches/parallelism/schedule_patches.py index 138ae4a1d..d0b71bbbd 100644 --- a/primus/backends/megatron/patches/parallelism/schedule_patches.py +++ b/primus/backends/megatron/patches/parallelism/schedule_patches.py @@ -9,7 +9,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( 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 64f9275f8..f44875b78 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 @@ -37,7 +37,7 @@ import torch from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 def _sdma_allgather_enabled(_ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/parallelism/te_wgrad_split_patches.py b/primus/backends/megatron/patches/parallelism/te_wgrad_split_patches.py index 74d49f6ca..e3bfde632 100644 --- a/primus/backends/megatron/patches/parallelism/te_wgrad_split_patches.py +++ b/primus/backends/megatron/patches/parallelism/te_wgrad_split_patches.py @@ -17,7 +17,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/parallelism/train_step_patches.py b/primus/backends/megatron/patches/parallelism/train_step_patches.py index 9b539adaf..8916ed9c7 100644 --- a/primus/backends/megatron/patches/parallelism/train_step_patches.py +++ b/primus/backends/megatron/patches/parallelism/train_step_patches.py @@ -14,7 +14,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _needs_seq_split_adjustment(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/parallelism/v_schedule_patches.py b/primus/backends/megatron/patches/parallelism/v_schedule_patches.py index fece7a507..73a0604a1 100644 --- a/primus/backends/megatron/patches/parallelism/v_schedule_patches.py +++ b/primus/backends/megatron/patches/parallelism/v_schedule_patches.py @@ -9,7 +9,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _is_v_schedule_enabled(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/parallelism/zero_bubble_patches.py b/primus/backends/megatron/patches/parallelism/zero_bubble_patches.py index 4c1bbde29..d2786f5cb 100644 --- a/primus/backends/megatron/patches/parallelism/zero_bubble_patches.py +++ b/primus/backends/megatron/patches/parallelism/zero_bubble_patches.py @@ -9,7 +9,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/pp_dump_data_patches.py b/primus/backends/megatron/patches/pp_dump_data_patches.py index 6e4816deb..72c939c8f 100644 --- a/primus/backends/megatron/patches/pp_dump_data_patches.py +++ b/primus/backends/megatron/patches/pp_dump_data_patches.py @@ -21,7 +21,7 @@ import os from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _make_wrapped_get_forward_backward_func(original_get_forward_backward_func): @@ -30,12 +30,12 @@ def _make_wrapped_get_forward_backward_func(original_get_forward_backward_func): def wrapped_get_forward_backward_func(): from megatron.training import get_args as get_megatron_args - from primus.modules.trainer.megatron.utils import schedule_wrapper + import primus.backends.megatron.core.pipeline_parallel.pp_visualizer as utils func = original_get_forward_backward_func() args = get_megatron_args() if getattr(args, "dump_pp_data", False): - func = schedule_wrapper(func) + func = utils.schedule_wrapper(func) func._pp_dump_schedule_wrapped = True # noqa: B010 return func @@ -81,7 +81,7 @@ def patch_pp_dump_data_before_train(ctx: PatchContext): import megatron.core.pipeline_parallel as pp_module import megatron.training.training as training_module - import primus.modules.trainer.megatron.utils as utils + import primus.backends.megatron.core.pipeline_parallel.pp_visualizer as utils args = get_args(ctx) is_primus_pipeline = getattr(args, "patch_primus_pipeline", False) @@ -131,11 +131,11 @@ def patch_pp_dump_data_after_train(ctx: PatchContext): from megatron.core.num_microbatches_calculator import get_num_microbatches from megatron.training import get_args as get_megatron_args - from primus.modules.trainer.megatron.utils import dump_pp_data + import primus.backends.megatron.core.pipeline_parallel.pp_visualizer as utils args = get_megatron_args() pp_data_dir = os.environ.get("DUMP_PP_DIR", "output/pp_data") - dump_pp_data(args, get_num_microbatches(), pp_data_dir) + utils.dump_pp_data(args, get_num_microbatches(), pp_data_dir) log_rank_0(f"[Patch:megatron.pp.dump_pp_data] pp schedule data dumped to {pp_data_dir}") except Exception as e: log_rank_0(f"[Patch:megatron.pp.dump_pp_data] WARNING: failed to dump pp data: {e}") diff --git a/primus/backends/megatron/patches/pp_warmup_patches.py b/primus/backends/megatron/patches/pp_warmup_patches.py index 6deb8d703..852a93433 100644 --- a/primus/backends/megatron/patches/pp_warmup_patches.py +++ b/primus/backends/megatron/patches/pp_warmup_patches.py @@ -15,7 +15,7 @@ import torch from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def run_pp_warmup(forward_step_func, model, optimizer, config): diff --git a/primus/backends/megatron/patches/recompute_layer_patches.py b/primus/backends/megatron/patches/recompute_layer_patches.py index c63d4008c..e79677395 100644 --- a/primus/backends/megatron/patches/recompute_layer_patches.py +++ b/primus/backends/megatron/patches/recompute_layer_patches.py @@ -38,7 +38,7 @@ from contextlib import nullcontext from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def validate_specified_recompute_layers(config, args): diff --git a/primus/backends/megatron/patches/runtime_hooks_patches.py b/primus/backends/megatron/patches/runtime_hooks_patches.py index 94bb3ac53..6d3dc5307 100644 --- a/primus/backends/megatron/patches/runtime_hooks_patches.py +++ b/primus/backends/megatron/patches/runtime_hooks_patches.py @@ -12,7 +12,7 @@ """ from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py b/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py index 6d8593768..8f6bf5708 100644 --- a/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py +++ b/primus/backends/megatron/patches/sdma_symm_mem_collectives_patches.py @@ -41,7 +41,7 @@ import os from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 def _sdma_all_gather_enabled(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/sft_grad_sanitize_patches.py b/primus/backends/megatron/patches/sft_grad_sanitize_patches.py index 59f847e61..896e9574f 100644 --- a/primus/backends/megatron/patches/sft_grad_sanitize_patches.py +++ b/primus/backends/megatron/patches/sft_grad_sanitize_patches.py @@ -36,7 +36,7 @@ import torch from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 _HAS_LOGGED_FIRST_SANITIZE = False diff --git a/primus/backends/megatron/patches/te_patches/bshd_layout_patches.py b/primus/backends/megatron/patches/te_patches/bshd_layout_patches.py index dba74f6e6..1f6178412 100644 --- a/primus/backends/megatron/patches/te_patches/bshd_layout_patches.py +++ b/primus/backends/megatron/patches/te_patches/bshd_layout_patches.py @@ -30,7 +30,7 @@ import os from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 def _bshd_enabled(_ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/te_patches/delayed_scaling_patches.py b/primus/backends/megatron/patches/te_patches/delayed_scaling_patches.py index 0a99f5a22..bda764962 100644 --- a/primus/backends/megatron/patches/te_patches/delayed_scaling_patches.py +++ b/primus/backends/megatron/patches/te_patches/delayed_scaling_patches.py @@ -15,7 +15,7 @@ make_get_extra_te_kwargs_with_override, ) from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/te_patches/general_gemm_workspace_patches.py b/primus/backends/megatron/patches/te_patches/general_gemm_workspace_patches.py index eceb47243..bd36f8644 100644 --- a/primus/backends/megatron/patches/te_patches/general_gemm_workspace_patches.py +++ b/primus/backends/megatron/patches/te_patches/general_gemm_workspace_patches.py @@ -15,7 +15,7 @@ import inspect from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _te_general_gemm_workspace_mode(): diff --git a/primus/backends/megatron/patches/te_patches/layernorm_linear_fp8_cache_patches.py b/primus/backends/megatron/patches/te_patches/layernorm_linear_fp8_cache_patches.py index 5e0a9ea83..56b1f6caf 100644 --- a/primus/backends/megatron/patches/te_patches/layernorm_linear_fp8_cache_patches.py +++ b/primus/backends/megatron/patches/te_patches/layernorm_linear_fp8_cache_patches.py @@ -14,7 +14,7 @@ make_get_extra_te_kwargs_with_override, ) from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/te_patches/legacy_grouped_mlp_wgrad_patches.py b/primus/backends/megatron/patches/te_patches/legacy_grouped_mlp_wgrad_patches.py index beb38bbaf..536470455 100644 --- a/primus/backends/megatron/patches/te_patches/legacy_grouped_mlp_wgrad_patches.py +++ b/primus/backends/megatron/patches/te_patches/legacy_grouped_mlp_wgrad_patches.py @@ -35,7 +35,7 @@ import torch from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _resolve_grouped_mlp_class(): diff --git a/primus/backends/megatron/patches/te_patches/linear_fp8_cache_patches.py b/primus/backends/megatron/patches/te_patches/linear_fp8_cache_patches.py index af539b82b..7f74dea81 100644 --- a/primus/backends/megatron/patches/te_patches/linear_fp8_cache_patches.py +++ b/primus/backends/megatron/patches/te_patches/linear_fp8_cache_patches.py @@ -15,7 +15,7 @@ make_get_extra_te_kwargs_with_override, ) from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/te_patches/tp_overlap_patches.py b/primus/backends/megatron/patches/te_patches/tp_overlap_patches.py index 89a7e2027..58cc89712 100644 --- a/primus/backends/megatron/patches/te_patches/tp_overlap_patches.py +++ b/primus/backends/megatron/patches/te_patches/tp_overlap_patches.py @@ -18,7 +18,7 @@ is_te_v2_or_above, ) from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _check_tp_overlap_conditions(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/tokenizer_builder_patches.py b/primus/backends/megatron/patches/tokenizer_builder_patches.py index f441af53f..61c819989 100644 --- a/primus/backends/megatron/patches/tokenizer_builder_patches.py +++ b/primus/backends/megatron/patches/tokenizer_builder_patches.py @@ -25,7 +25,7 @@ """ from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/torch_fsdp2_patches.py b/primus/backends/megatron/patches/torch_fsdp2_patches.py index 7eaf39e98..b8ba726e2 100644 --- a/primus/backends/megatron/patches/torch_fsdp2_patches.py +++ b/primus/backends/megatron/patches/torch_fsdp2_patches.py @@ -13,7 +13,7 @@ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/megatron/patches/torch_profiler_patches.py b/primus/backends/megatron/patches/torch_profiler_patches.py index 0149f5aa5..5a039442d 100644 --- a/primus/backends/megatron/patches/torch_profiler_patches.py +++ b/primus/backends/megatron/patches/torch_profiler_patches.py @@ -14,7 +14,7 @@ import inspect from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _is_called_from_training_train() -> bool: @@ -31,8 +31,6 @@ def _is_called_from_training_train() -> bool: def _create_primus_prof(args, exp_name: str, original_profile): """ Create torch profiler with Primus options. - - Logic from primus/modules/trainer/megatron/trainer.py L1277-1298. """ import torch diff --git a/primus/backends/megatron/patches/training_log/print_rank_last_patches.py b/primus/backends/megatron/patches/training_log/print_rank_last_patches.py index a43a83b26..db0cedf04 100644 --- a/primus/backends/megatron/patches/training_log/print_rank_last_patches.py +++ b/primus/backends/megatron/patches/training_log/print_rank_last_patches.py @@ -31,8 +31,8 @@ from primus.core.patches import PatchContext, get_args, register_patch from primus.core.utils import logger as primus_logger +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 from primus.core.utils.rocm_mem_info import get_rocm_smi_mem_info -from primus.modules.module_utils import log_rank_0, warning_rank_0 @dataclass diff --git a/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py b/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py index 11d649857..cf0852cda 100644 --- a/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py +++ b/primus/backends/megatron/patches/turbo/aiter_deepbind_patches.py @@ -36,7 +36,7 @@ from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 _LOG_PREFIX = "[Patch:megatron.turbo.aiter_deepbind]" _ENV_SWITCH = "PRIMUS_AITER_DEEPBIND" diff --git a/primus/backends/megatron/patches/turbo/fp4_patches.py b/primus/backends/megatron/patches/turbo/fp4_patches.py index d95fa0e52..c1990fac1 100644 --- a/primus/backends/megatron/patches/turbo/fp4_patches.py +++ b/primus/backends/megatron/patches/turbo/fp4_patches.py @@ -10,7 +10,7 @@ from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _is_fp4_can_patch(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/turbo/fp8_patches.py b/primus/backends/megatron/patches/turbo/fp8_patches.py index 22a9989b6..7a077f631 100644 --- a/primus/backends/megatron/patches/turbo/fp8_patches.py +++ b/primus/backends/megatron/patches/turbo/fp8_patches.py @@ -13,7 +13,7 @@ from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _is_fp8_can_patch(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py b/primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py index 86ff432e9..0d93be267 100644 --- a/primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py +++ b/primus/backends/megatron/patches/turbo/fused_residual_norm_patches.py @@ -26,7 +26,7 @@ import os from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _env_truthy(name: str) -> bool: diff --git a/primus/backends/megatron/patches/turbo/moe_dispatcher_patches.py b/primus/backends/megatron/patches/turbo/moe_dispatcher_patches.py index 3cbc0dd90..2185fdcaa 100644 --- a/primus/backends/megatron/patches/turbo/moe_dispatcher_patches.py +++ b/primus/backends/megatron/patches/turbo/moe_dispatcher_patches.py @@ -13,7 +13,7 @@ from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _is_turbo_deepep_can_patch(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/turbo/rms_norm_patches.py b/primus/backends/megatron/patches/turbo/rms_norm_patches.py index 0611b465b..cdf45243b 100644 --- a/primus/backends/megatron/patches/turbo/rms_norm_patches.py +++ b/primus/backends/megatron/patches/turbo/rms_norm_patches.py @@ -13,7 +13,7 @@ from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _is_turbo_rms_norm_can_patch(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/turbo/te_spec_provider_patches.py b/primus/backends/megatron/patches/turbo/te_spec_provider_patches.py index b9d424289..192be08bc 100644 --- a/primus/backends/megatron/patches/turbo/te_spec_provider_patches.py +++ b/primus/backends/megatron/patches/turbo/te_spec_provider_patches.py @@ -12,7 +12,7 @@ from primus.backends.megatron.patches.turbo.utils import is_primus_turbo_can_patch from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _use_legacy_grouped_gemm(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/turbo/utils.py b/primus/backends/megatron/patches/turbo/utils.py index 6452b407b..4b1b8ce20 100644 --- a/primus/backends/megatron/patches/turbo/utils.py +++ b/primus/backends/megatron/patches/turbo/utils.py @@ -7,7 +7,7 @@ import importlib.util from primus.core.patches import PatchContext, get_args -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 def _is_primus_turbo_enabled(ctx: PatchContext) -> bool: diff --git a/primus/backends/megatron/patches/zebra_llama_flops_patches.py b/primus/backends/megatron/patches/zebra_llama_flops_patches.py index 1193fabd8..6566d4e35 100644 --- a/primus/backends/megatron/patches/zebra_llama_flops_patches.py +++ b/primus/backends/megatron/patches/zebra_llama_flops_patches.py @@ -13,7 +13,7 @@ """ from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def zebra_llama_flops(args, batch_size): diff --git a/primus/backends/megatron/peft/recompute.py b/primus/backends/megatron/peft/recompute.py index 0137a3b74..3688e1749 100644 --- a/primus/backends/megatron/peft/recompute.py +++ b/primus/backends/megatron/peft/recompute.py @@ -22,7 +22,7 @@ import torch from megatron.core.utils import unwrap_model -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 PEFT_RECOMPUTE_PATCHED: Set[int] = set() diff --git a/primus/backends/megatron/sft/preprocessing.py b/primus/backends/megatron/sft/preprocessing.py index 22aaf08e7..599f769ff 100644 --- a/primus/backends/megatron/sft/preprocessing.py +++ b/primus/backends/megatron/sft/preprocessing.py @@ -16,7 +16,7 @@ from primus.backends.megatron.sft.schema import FormattedSFTSample, SFTSample try: - from primus.modules.module_utils import log_rank_0 as _primus_log_rank_0 + from primus.core.utils.module_utils import log_rank_0 as _primus_log_rank_0 except ImportError: _primus_log_rank_0 = None diff --git a/primus/backends/megatron/sft/runtime.py b/primus/backends/megatron/sft/runtime.py index 68b333f2f..a578e38e4 100644 --- a/primus/backends/megatron/sft/runtime.py +++ b/primus/backends/megatron/sft/runtime.py @@ -11,7 +11,7 @@ from typing import Any, Optional from primus.backends.megatron.sft.dataset import build_train_valid_test_datasets -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _safe_signature(fn: Callable[..., Any]) -> inspect.Signature | None: diff --git a/primus/backends/megatron/training/evaluator.py b/primus/backends/megatron/training/evaluator.py index 8d04bafcc..06071b46f 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -16,7 +16,7 @@ from primus.backends.megatron.training.global_vars import get_train_start_time from primus.backends.megatron.training.utils import is_pipeline_stage_containing_loss -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def primus_evaluate( diff --git a/primus/backends/megatron/training/global_vars.py b/primus/backends/megatron/training/global_vars.py index 4e0a52a86..e0fd44884 100644 --- a/primus/backends/megatron/training/global_vars.py +++ b/primus/backends/megatron/training/global_vars.py @@ -13,7 +13,7 @@ collect_git_metadata, get_env_variables, ) -from primus.modules.module_utils import debug_rank_0 +from primus.core.utils.module_utils import debug_rank_0 _GLOBAL_ARGS = None _GLOBAL_MLFLOW_WRITER = None diff --git a/primus/backends/megatron/training/mlflow_artifacts.py b/primus/backends/megatron/training/mlflow_artifacts.py index 2b6e32310..6e38f8d62 100644 --- a/primus/backends/megatron/training/mlflow_artifacts.py +++ b/primus/backends/megatron/training/mlflow_artifacts.py @@ -40,7 +40,7 @@ import sys from typing import List, Optional -from primus.modules.module_utils import log_rank_0, log_rank_last, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, log_rank_last, warning_rank_0 # Pinned to immutable commit SHA for supply-chain safety (tags can be moved). # This corresponds to tag v0.4.0 in AMD-AGI/TraceLens. diff --git a/primus/backends/megatron/training/tokenizer/tokenizer.py b/primus/backends/megatron/training/tokenizer/tokenizer.py index 376996521..5e97a987a 100644 --- a/primus/backends/megatron/training/tokenizer/tokenizer.py +++ b/primus/backends/megatron/training/tokenizer/tokenizer.py @@ -26,7 +26,7 @@ HuggingFaceTokenizer as _HuggingFaceTokenizer, ) -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 CUSTOM_TOKENIZER_TYPES = { "DeepSeekV2Tokenizer", diff --git a/primus/backends/megatron/training/utils.py b/primus/backends/megatron/training/utils.py index 61a262f2a..4484f5307 100644 --- a/primus/backends/megatron/training/utils.py +++ b/primus/backends/megatron/training/utils.py @@ -10,6 +10,19 @@ """General utilities.""" import torch from megatron.core import mpu, parallel_state +from megatron.training.global_vars import get_args + + +def is_v_schedule_enabled(args=None): + """Return True when a V-shaped pipeline schedule (ZeroBubble-V / Primus-Pipe + zbv/v-half/v-min) is active for the current run.""" + if args is None: + args = get_args() + return ( + args.patch_zero_bubble + and args.enable_zero_bubble + and (args.zero_bubble_v_schedule or args.enable_1f1b_v) + ) or (args.pp_algorithm in ("zbv-formatted", "v-half", "v-min") and args.patch_primus_pipeline) def is_second_last_pipeline_stage(): @@ -31,8 +44,6 @@ def print_second_last_pipeline_stage(message): def is_pipeline_stage_containing_loss(): - from primus.modules.trainer.megatron.utils import is_v_schedule_enabled - if is_v_schedule_enabled(): return mpu.is_pipeline_first_stage(ignore_virtual=True) else: diff --git a/primus/backends/megatron_bridge/config_utils.py b/primus/backends/megatron_bridge/config_utils.py index 987e628d1..31c8b0090 100644 --- a/primus/backends/megatron_bridge/config_utils.py +++ b/primus/backends/megatron_bridge/config_utils.py @@ -21,7 +21,7 @@ from types import SimpleNamespace from typing import Any, Callable, Dict -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 DATASET_PATH_KEYS = ( "data_paths", diff --git a/primus/backends/megatron_bridge/megatron_bridge_adapter.py b/primus/backends/megatron_bridge/megatron_bridge_adapter.py index f0cdeecf4..972c91f6a 100644 --- a/primus/backends/megatron_bridge/megatron_bridge_adapter.py +++ b/primus/backends/megatron_bridge/megatron_bridge_adapter.py @@ -29,7 +29,7 @@ ) from primus.core.backend.backend_adapter import BackendAdapter from primus.core.backend.backend_registry import BackendRegistry -from primus.modules.module_utils import log_dict_aligned, log_rank_0 +from primus.core.utils.module_utils import log_dict_aligned, log_rank_0 def _install_modelopt_stub() -> bool: diff --git a/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py index 2f46dfe7b..f2943e381 100644 --- a/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py +++ b/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py @@ -19,7 +19,7 @@ from primus.backends.megatron.training.global_vars import set_primus_global_variables from primus.core.patches import run_patches from primus.core.trainer.base_trainer import BaseTrainer -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class MegatronBridgeBaseTrainer(BaseTrainer): diff --git a/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py index 4ed170ec4..dd9ed4a9e 100644 --- a/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py +++ b/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py @@ -22,7 +22,7 @@ from primus.backends.megatron_bridge.megatron_bridge_base_trainer import ( MegatronBridgeBaseTrainer, ) -from primus.modules.module_utils import log_dict_aligned, log_rank_0 +from primus.core.utils.module_utils import log_dict_aligned, log_rank_0 class MegatronBridgePosttrainTrainer(MegatronBridgeBaseTrainer): diff --git a/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py index ac73d7362..aa76e8daa 100644 --- a/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py +++ b/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py @@ -20,7 +20,7 @@ from primus.backends.megatron_bridge.megatron_bridge_base_trainer import ( MegatronBridgeBaseTrainer, ) -from primus.modules.module_utils import log_dict_aligned, log_rank_0 +from primus.core.utils.module_utils import log_dict_aligned, log_rank_0 class MegatronBridgePretrainTrainer(MegatronBridgeBaseTrainer): diff --git a/primus/backends/megatron_bridge/patches/training_log/bridge_training_log_patches.py b/primus/backends/megatron_bridge/patches/training_log/bridge_training_log_patches.py index 0e948a533..2e8d79326 100644 --- a/primus/backends/megatron_bridge/patches/training_log/bridge_training_log_patches.py +++ b/primus/backends/megatron_bridge/patches/training_log/bridge_training_log_patches.py @@ -36,7 +36,7 @@ from typing import Optional from primus.core.patches import PatchContext, get_args, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 # Per-step actual sequence length captured from prepare_finetuning_batch. # Accessed from the same training-loop thread; the lock guards against diff --git a/primus/backends/torchtitan/config_utils.py b/primus/backends/torchtitan/config_utils.py index 39f212cc4..695dc8a49 100644 --- a/primus/backends/torchtitan/config_utils.py +++ b/primus/backends/torchtitan/config_utils.py @@ -20,8 +20,8 @@ from types import SimpleNamespace from typing import Any +from primus.core.utils.module_utils import log_rank_0 from primus.core.utils.yaml_utils import dict_to_nested_namespace -from primus.modules.module_utils import log_rank_0 def build_job_config_from_namespace(ns: SimpleNamespace) -> Any: diff --git a/primus/backends/torchtitan/patches/dcp_consolidate_patches.py b/primus/backends/torchtitan/patches/dcp_consolidate_patches.py index 973a9be04..702e981c9 100644 --- a/primus/backends/torchtitan/patches/dcp_consolidate_patches.py +++ b/primus/backends/torchtitan/patches/dcp_consolidate_patches.py @@ -15,7 +15,7 @@ """ from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/embedding_amp_patches.py b/primus/backends/torchtitan/patches/embedding_amp_patches.py index 3daedf44e..df4e153a3 100644 --- a/primus/backends/torchtitan/patches/embedding_amp_patches.py +++ b/primus/backends/torchtitan/patches/embedding_amp_patches.py @@ -21,7 +21,7 @@ """ from primus.core.patches import PatchContext, get_param, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/flex_attention_patches.py b/primus/backends/torchtitan/patches/flex_attention_patches.py index 5acadd41c..c05ed3f18 100644 --- a/primus/backends/torchtitan/patches/flex_attention_patches.py +++ b/primus/backends/torchtitan/patches/flex_attention_patches.py @@ -19,7 +19,7 @@ """ from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/logger_patches.py b/primus/backends/torchtitan/patches/logger_patches.py index 3940c52f6..0e3292bb8 100644 --- a/primus/backends/torchtitan/patches/logger_patches.py +++ b/primus/backends/torchtitan/patches/logger_patches.py @@ -33,7 +33,7 @@ """ from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/mock_dataset_patches.py b/primus/backends/torchtitan/patches/mock_dataset_patches.py index 9703504bf..4ecb49918 100644 --- a/primus/backends/torchtitan/patches/mock_dataset_patches.py +++ b/primus/backends/torchtitan/patches/mock_dataset_patches.py @@ -7,8 +7,7 @@ """ TorchTitan Mock HF Dataset Patch -This patch mirrors ``patch_mock_hf_dataset`` from -``primus.modules.trainer.torchtitan.patch_utils`` using the generic Primus +This patch implements ``patch_mock_hf_dataset`` using the generic Primus patch system, so that HF dataset mocking can be enabled via config without tightly coupling it to the trainer implementation. """ @@ -17,7 +16,7 @@ from datasets.search import np from primus.core.patches import PatchContext, get_param, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _create_mock_text_dataset(num_samples: int = 128) -> Dataset: diff --git a/primus/backends/torchtitan/patches/model_override_patches.py b/primus/backends/torchtitan/patches/model_override_patches.py index 37d27e458..ffeaa9464 100644 --- a/primus/backends/torchtitan/patches/model_override_patches.py +++ b/primus/backends/torchtitan/patches/model_override_patches.py @@ -50,7 +50,7 @@ from typing import Any from primus.core.patches import PatchContext, get_param, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def get_standard_model_config_fields() -> set: diff --git a/primus/backends/torchtitan/patches/pipelining_schedule_patches.py b/primus/backends/torchtitan/patches/pipelining_schedule_patches.py index fb2040c1a..9b9a6958e 100644 --- a/primus/backends/torchtitan/patches/pipelining_schedule_patches.py +++ b/primus/backends/torchtitan/patches/pipelining_schedule_patches.py @@ -14,7 +14,7 @@ """ from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py b/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py index 943d08080..60aa0a639 100644 --- a/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py +++ b/primus/backends/torchtitan/patches/sdma_symm_mem_collectives.py @@ -39,7 +39,7 @@ import os from primus.core.patches import PatchContext, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def _sdma_all_gather_enabled(ctx: PatchContext) -> bool: diff --git a/primus/backends/torchtitan/patches/turbo/async_tp_patches.py b/primus/backends/torchtitan/patches/turbo/async_tp_patches.py index ae260e556..0db17c15d 100644 --- a/primus/backends/torchtitan/patches/turbo/async_tp_patches.py +++ b/primus/backends/torchtitan/patches/turbo/async_tp_patches.py @@ -18,7 +18,7 @@ from typing import Any, Optional from primus.core.patches import PatchContext, get_param, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/turbo/attention_patches.py b/primus/backends/torchtitan/patches/turbo/attention_patches.py index 9146cefb1..8dbf1fe05 100644 --- a/primus/backends/torchtitan/patches/turbo/attention_patches.py +++ b/primus/backends/torchtitan/patches/turbo/attention_patches.py @@ -16,7 +16,7 @@ """ from primus.core.patches import PatchContext, get_param, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/turbo/deepseek_v3_classic_attention_patches.py b/primus/backends/torchtitan/patches/turbo/deepseek_v3_classic_attention_patches.py index 119417188..ca96fe194 100644 --- a/primus/backends/torchtitan/patches/turbo/deepseek_v3_classic_attention_patches.py +++ b/primus/backends/torchtitan/patches/turbo/deepseek_v3_classic_attention_patches.py @@ -17,7 +17,7 @@ """ from primus.core.patches import PatchContext, get_param, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/turbo/fp8_linear_patches.py b/primus/backends/torchtitan/patches/turbo/fp8_linear_patches.py index fbd582bcf..b5df8a9ff 100644 --- a/primus/backends/torchtitan/patches/turbo/fp8_linear_patches.py +++ b/primus/backends/torchtitan/patches/turbo/fp8_linear_patches.py @@ -16,7 +16,7 @@ """ from primus.core.patches import PatchContext, get_param, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py b/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py index f48a1d120..fcd2569bf 100644 --- a/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py +++ b/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py @@ -22,7 +22,7 @@ import functools from primus.core.patches import PatchContext, get_param, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/turbo/mx_linear_patches.py b/primus/backends/torchtitan/patches/turbo/mx_linear_patches.py index 15473ffb2..bb2a23c1f 100644 --- a/primus/backends/torchtitan/patches/turbo/mx_linear_patches.py +++ b/primus/backends/torchtitan/patches/turbo/mx_linear_patches.py @@ -16,7 +16,7 @@ """ from primus.core.patches import PatchContext, get_param, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/patches/wandb_patches.py b/primus/backends/torchtitan/patches/wandb_patches.py index 129427ba7..538ebea2b 100644 --- a/primus/backends/torchtitan/patches/wandb_patches.py +++ b/primus/backends/torchtitan/patches/wandb_patches.py @@ -23,7 +23,7 @@ import os from primus.core.patches import PatchContext, get_param, register_patch -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @register_patch( diff --git a/primus/backends/torchtitan/torchtitan_adapter.py b/primus/backends/torchtitan/torchtitan_adapter.py index 0bace2840..f8fa1e5f9 100644 --- a/primus/backends/torchtitan/torchtitan_adapter.py +++ b/primus/backends/torchtitan/torchtitan_adapter.py @@ -24,7 +24,7 @@ from primus.backends.torchtitan.argument_builder import TorchTitanJobConfigBuilder from primus.core.backend.backend_adapter import BackendAdapter from primus.core.backend.backend_registry import BackendRegistry -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class TorchTitanAdapter(BackendAdapter): diff --git a/primus/backends/transformer_engine/pytorch/module/base.py b/primus/backends/transformer_engine/pytorch/module/base.py index 003e18b45..80be73600 100644 --- a/primus/backends/transformer_engine/pytorch/module/base.py +++ b/primus/backends/transformer_engine/pytorch/module/base.py @@ -13,7 +13,7 @@ from transformer_engine.pytorch.module import base import primus.backends.transformer_engine.transformer_engine_torch as ptex -from primus.modules.module_utils import log_rank_0, warning_rank_0 +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 def get_cublas_workspace_size_bytes() -> None: diff --git a/primus/cli/subcommands/train.py b/primus/cli/subcommands/train.py index ac8d2da97..5cb2262ba 100644 --- a/primus/cli/subcommands/train.py +++ b/primus/cli/subcommands/train.py @@ -6,50 +6,19 @@ from __future__ import annotations -import sys -from os import getenv from typing import List -def _resolve_pretrain_runtime(args) -> str: - """ - Resolve the runtime entry for pretrain. - - Priority: - 1) Explicit env override via PRIMUS_TRAIN_RUNTIME - 2) Framework-based default (MaxText -> legacy, others -> core) - """ - runtime_entry = getenv("PRIMUS_TRAIN_RUNTIME", "").strip().lower() - if runtime_entry in ("legacy", "core"): - return runtime_entry - if runtime_entry: - print( - f"[Primus:Train] Ignoring invalid PRIMUS_TRAIN_RUNTIME='{runtime_entry}'.", - file=sys.stderr, - ) - - # Default: use the new core runtime for all supported frameworks. - return "core" - - def run(args, overrides: List[str]): """ Entry point for the 'train' subcommand. """ if args.suite == "pretrain": - runtime_entry = _resolve_pretrain_runtime(args) - - if runtime_entry == "core": - # New core runtime path: mirror `train_launcher.launch_train`. - from primus.core.runtime.train_runtime import PrimusRuntime - - runtime = PrimusRuntime(args=args) - runtime.run_train_module(module_name="pre_trainer", overrides=overrides or []) - else: - # Legacy pretrain flow. - from primus.pretrain import launch_pretrain_from_cli + # All frameworks train via the core runtime. + from primus.core.runtime.train_runtime import PrimusRuntime - launch_pretrain_from_cli(args, overrides) + runtime = PrimusRuntime(args=args) + runtime.run_train_module(module_name="pre_trainer", overrides=overrides or []) elif args.suite == "posttrain": # Post-training (SFT/alignment) currently runs via the new core runtime. # It expects a training module named "sft_trainer" in the experiment config. diff --git a/primus/core/backend/backend_adapter.py b/primus/core/backend/backend_adapter.py index 738191b5f..d0aeb16bf 100644 --- a/primus/core/backend/backend_adapter.py +++ b/primus/core/backend/backend_adapter.py @@ -40,7 +40,7 @@ def prepare_backend(self, config: Any): beyond `BackendRegistry.run_setup(self.framework)`. """ from primus.core.backend.backend_registry import BackendRegistry - from primus.modules.module_utils import log_rank_0 + from primus.core.utils.module_utils import log_rank_0 BackendRegistry.run_setup(self.framework) log_rank_0(f"[Primus:{self.framework}] Backend prepared") @@ -57,7 +57,7 @@ def setup_backend_path(self, backend_path=None) -> str: import sys from pathlib import Path - from primus.modules.module_utils import log_rank_0 + from primus.core.utils.module_utils import log_rank_0 def _use_path(path: str, error_msg: str) -> str: norm_path = os.path.abspath(os.path.normpath(str(path))) diff --git a/primus/core/backend/backend_registry.py b/primus/core/backend/backend_registry.py index e11001877..d198c7073 100644 --- a/primus/core/backend/backend_registry.py +++ b/primus/core/backend/backend_registry.py @@ -20,7 +20,7 @@ from typing import Callable, Dict, List, Type -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class BackendRegistry: diff --git a/primus/modules/base_module.py b/primus/core/base_module.py similarity index 98% rename from primus/modules/base_module.py rename to primus/core/base_module.py index 71fa19f02..8c5283cb1 100644 --- a/primus/modules/base_module.py +++ b/primus/core/base_module.py @@ -15,8 +15,7 @@ get_target_platform, set_global_variables, ) - -from .module_utils import debug_rank_all, set_logging_rank +from primus.core.utils.module_utils import debug_rank_all, set_logging_rank class BaseModule(ABC): diff --git a/primus/core/patches/patch.py b/primus/core/patches/patch.py index 50522f163..466bfddfc 100644 --- a/primus/core/patches/patch.py +++ b/primus/core/patches/patch.py @@ -15,7 +15,7 @@ from primus.core.patches.context import PatchContext from primus.core.patches.utils import version_in_range -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 # ----------------------------------------------------------------------------- # FunctionPatch diff --git a/primus/core/patches/patch_runner.py b/primus/core/patches/patch_runner.py index 1a7eae356..f0ba5f058 100644 --- a/primus/core/patches/patch_runner.py +++ b/primus/core/patches/patch_runner.py @@ -16,7 +16,7 @@ from primus.core.patches.context import PatchContext from primus.core.patches.patch_registry import PatchRegistry -from primus.modules.module_utils import error_rank_0, log_rank_0 +from primus.core.utils.module_utils import error_rank_0, log_rank_0 # ----------------------------------------------------------------------------- # Parse PRIMUS_PATCHES Environment Variable diff --git a/primus/core/pipeline_parallel/handler/wgrad_handler.py b/primus/core/pipeline_parallel/handler/wgrad_handler.py index 421cde6f1..b5fb723e6 100644 --- a/primus/core/pipeline_parallel/handler/wgrad_handler.py +++ b/primus/core/pipeline_parallel/handler/wgrad_handler.py @@ -8,8 +8,10 @@ from megatron.training.global_vars import get_args +from primus.backends.megatron.core.pipeline_parallel.pp_visualizer import ( + fwd_bwd_wrapper, +) from primus.core.pipeline_parallel.scheduler.scheduler_node import SchedulerNode -from primus.modules.trainer.megatron.utils import fwd_bwd_wrapper class WGradRunningCache: diff --git a/primus/core/pipeline_parallel/scheduler/algorithms/base.py b/primus/core/pipeline_parallel/scheduler/algorithms/base.py index 05d145bb7..d51277236 100644 --- a/primus/core/pipeline_parallel/scheduler/algorithms/base.py +++ b/primus/core/pipeline_parallel/scheduler/algorithms/base.py @@ -11,7 +11,7 @@ FuncType, SchedulerNode, ) -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class PipelineScheduleAlgo(ABC): diff --git a/primus/core/projection/memory_projection/benchmark.py b/primus/core/projection/memory_projection/benchmark.py index bfcafd466..0b3eb106f 100644 --- a/primus/core/projection/memory_projection/benchmark.py +++ b/primus/core/projection/memory_projection/benchmark.py @@ -174,7 +174,7 @@ def _run_bench( if reduction_info.get("benchmark_num_experts") is not None: bench_module.num_experts = reduction_info["benchmark_num_experts"] - profiling_results = _run_layer_benchmark(primus_config_bench, overrides, reduction_info) + profiling_results = _run_layer_benchmark(primus_config_bench, overrides, reduction_info, args=args) # Optionally persist the artifact. We piggy-back on the perf saver so # that the JSON is identical-shape to the perf-side artifact (one diff --git a/primus/core/projection/performance_projection/projection.py b/primus/core/projection/performance_projection/projection.py index 5f9a7d57a..fdcac1433 100644 --- a/primus/core/projection/performance_projection/projection.py +++ b/primus/core/projection/performance_projection/projection.py @@ -39,9 +39,9 @@ convert_primus_config_to_projection_config, ) -# NOTE: MegatronPretrainTrainer is imported lazily inside _run_layer_benchmark() -# to avoid pulling in the megatron dependency when running in pure simulation mode -# (--profiling-mode simulate). +# NOTE: The core runtime (PrimusRuntime) and megatron backend are imported +# lazily inside _run_layer_benchmark() to avoid pulling in the megatron +# dependency when running in pure simulation mode (--profiling-mode simulate). _MAX_EXPERT_PARALLEL_SIZE = 8 _BYTES_PER_GB = 1024**3 @@ -2536,9 +2536,43 @@ def _summarize_bench_training_config(training_config) -> Dict[str, Any]: } -def _run_layer_benchmark(primus_config, unknown_overrides, reduction_info=None): - from primus.modules.trainer.megatron.pre_trainer import MegatronPretrainTrainer +def _build_runtime_primus_config(legacy_primus_config, args, module_name="pre_trainer"): + """Adapt the (already-mutated) legacy PrimusConfig into a runtime-shaped config. + Performance projection loads config via the legacy ``PrimusConfig`` and mutates + the flat ``pre_trainer`` module in place (layer limiting, EP rescale, turbo / + overlap overrides, ...). The core runtime, however, consumes the normalized + ``cfg.modules[*].params`` shape. This helper reuses the runtime loader to build + a correct scaffold (exp_root_path / exp_meta_info / platform / ...), then swaps + in the mutated ``pre_trainer`` module normalized to the runtime shape, so the + benchmark model is built from exactly the mutated config. + """ + from pathlib import Path + + from primus.core.config.primus_config import ( + _normalize_module_for_runtime, + load_primus_config, + ) + + runtime_cfg = load_primus_config(Path(args.config), args) + normalized = _normalize_module_for_runtime( + legacy_primus_config.get_module_config(module_name), module_name + ) + + modules = list(getattr(runtime_cfg, "modules", []) or []) + replaced = False + for i, m in enumerate(modules): + if getattr(m, "name", None) == module_name: + modules[i] = normalized + replaced = True + break + if not replaced: + modules.append(normalized) + runtime_cfg.modules = modules + return runtime_cfg + + +def _run_layer_benchmark(primus_config, unknown_overrides, reduction_info=None, args=None): module_config = primus_config.get_module_config("pre_trainer") _limit_layers_for_projection(module_config) rescale_info = _rescale_expert_parallelism(module_config) @@ -2549,8 +2583,6 @@ def _run_layer_benchmark(primus_config, unknown_overrides, reduction_info=None): if reduction_info is not None: reduction_info["bench_training_config_summary"] = _summarize_bench_training_config(training_config) - master_addr = os.getenv("MASTER_ADDR", "127.0.0.1") - master_port = int(os.getenv("MASTER_PORT", "29500")) rank = int(os.getenv("RANK", "0")) world_size = int(os.getenv("WORLD_SIZE", "1")) @@ -2560,7 +2592,7 @@ def _run_layer_benchmark(primus_config, unknown_overrides, reduction_info=None): mem_recorder = MemoryBenchmarkRecorder(rank=rank) mem_recorder.snapshot("pre_trainer_init") - print("[Primus:Performance Projection] Initializing MegatronPretrainTrainer...") + print("[Primus:Performance Projection] Preparing benchmark model build...") # Disable overlap features and FSDP2 for profiling (they add complexity without benefiting isolated layer benchmarking) # FSDP2 uses DTensor which causes issues with benchmarking inputs cfg = primus_config.get_module_config("pre_trainer") @@ -2599,22 +2631,32 @@ def _run_layer_benchmark(primus_config, unknown_overrides, reduction_info=None): if getattr(cfg, "multi_latent_attention", False): print(f" enable_primus_turbo: {cfg.enable_primus_turbo}") print(f" use_turbo_gemm: {cfg.use_turbo_gemm}") - trainer = MegatronPretrainTrainer( + + # Build the model via the new core runtime instead of the legacy trainer. + # The runtime reuses the exact real-training init pipeline (adapter config + # conversion + build_args/setup/before_train patches + megatron + # initialize_megatron + setup_model_and_optimizer) but stops before the + # training loop, giving us a model identical to real training for the + # per-layer benchmark. `args` provides data_path / backend_path / config. + from primus.core.runtime.train_runtime import PrimusRuntime + + if args is None: + raise ValueError( + "[Primus:Performance Projection] _run_layer_benchmark requires `args` " + "(CLI namespace with config/data_path/backend_path) to build the model " + "via the core runtime." + ) + + runtime_primus_config = _build_runtime_primus_config(primus_config, args, module_name="pre_trainer") + + print("[Primus:Performance Projection] Initializing Megatron and building model...") + runtime = PrimusRuntime(args=args) + trainer = runtime.setup_model_only( module_name="pre_trainer", - primus_config=primus_config, - module_rank=rank, - module_world_size=world_size, - module_master_addr=master_addr, - module_master_port=master_port, - extra_args=unknown_overrides, + overrides=unknown_overrides, + primus_config=runtime_primus_config, ) - - print("[Primus:Performance Projection] Initializing Megatron...") - trainer.init() mem_recorder.snapshot("post_megatron_init") - - print("[Primus:Performance Projection] Setting up model and optimizer...") - trainer.setup() # post_setup captures: params + distributed-optimizer state + grad buffers # at the bench config. This is the "static" memory anchor — what every # rank pays before any forward pass. @@ -3930,7 +3972,7 @@ def launch_projection_from_cli(args, overrides): # downstream pipeline simulation / multinode projection, but print # a side-by-side comparison. sim_results = _run_layer_simulation(copy.deepcopy(primus_config), args) - bench_results = _run_layer_benchmark(primus_config, unknown_overrides, reduction_info) + bench_results = _run_layer_benchmark(primus_config, unknown_overrides, reduction_info, args=args) is_rank_0 = int(os.getenv("RANK", "0")) == 0 if is_rank_0: @@ -3960,7 +4002,7 @@ def launch_projection_from_cli(args, overrides): profiling_results = bench_results else: # Default: actual GPU benchmark - profiling_results = _run_layer_benchmark(primus_config, unknown_overrides, reduction_info) + profiling_results = _run_layer_benchmark(primus_config, unknown_overrides, reduction_info, args=args) # ── Save bench artifact if requested ── # ``--save-benchmark`` (preferred) and ``--save-profiling`` (deprecated) diff --git a/primus/core/runtime/logging.py b/primus/core/runtime/logging.py index 0e56a86f2..bb2cb918c 100644 --- a/primus/core/runtime/logging.py +++ b/primus/core/runtime/logging.py @@ -33,7 +33,7 @@ from primus.core.utils import logger from primus.core.utils.env import get_torchrun_env -from primus.modules.module_utils import debug_rank_all, set_logging_rank +from primus.core.utils.module_utils import debug_rank_all, set_logging_rank # from primus.core.utils.distributed_logging import debug_rank_all, set_logging_rank diff --git a/primus/core/runtime/train_runtime.py b/primus/core/runtime/train_runtime.py index 67d72f364..61c8b7479 100644 --- a/primus/core/runtime/train_runtime.py +++ b/primus/core/runtime/train_runtime.py @@ -26,12 +26,12 @@ from primus.core.runtime.runtime_state import RuntimeState from primus.core.utils.arg_utils import parse_cli_overrides from primus.core.utils.env_setup import setup_training_env +from primus.core.utils.module_utils import log_dict_aligned, log_rank_0, warning_rank_0 from primus.core.utils.yaml_utils import ( dict_to_nested_namespace, merge_namespace, nested_namespace_to_dict, ) -from primus.modules.module_utils import log_dict_aligned, log_rank_0, warning_rank_0 # --------------------------------------------------------------------------- # Context & Hooks @@ -111,6 +111,64 @@ def run_train_module(self, module_name: str, overrides: Optional[List[str]] = No self._safe_cleanup(error=e) raise RuntimeError(f"Training execution failed: {e}") from e + def setup_model_only( + self, + module_name: str = "pre_trainer", + overrides: Optional[List[str]] = None, + primus_config: Any = None, + ) -> Any: + """Run the runtime pipeline but stop after building the model (no training). + + A general, training-neutral entry that mirrors :meth:`run_train_module` + except it replaces the training step with a model-only build. Reuses the + full runtime initialization pipeline (config → environment → distributed → + logging → backend adapter → args conversion → build_args patches → trainer + setup/init → before_train patches) and then asks the trainer to build only + the model via ``setup_model_only``. Useful for any "build the model only" + scenario (offline profiling, layer benchmarking, model inspection); + performance/memory projection is the current consumer. + + Args: + module_name: training module to build (default ``pre_trainer``). + overrides: CLI-style overrides to merge into the module params. + primus_config: optional pre-loaded (runtime-shaped) PrimusConfig. When + provided, config is not re-loaded from ``self.args.config`` — this + lets callers (e.g. projection) apply their own config mutations + before building. + + Returns: + The trainer instance whose ``.model`` has been built. + """ + overrides = overrides or [] + + # 1) Configuration (optionally injected, so callers can pre-mutate it) + self._initialize_configuration(module_name, overrides, primus_config=primus_config) + # 2) Runtime environment (paths, distributed, logging) + self._initialize_runtime_environment() + # 3) Backend adapter + trainer (convert config, build_args patches, instantiate) + self._initialize_adapter() + self._initialize_trainer() + + assert self.ctx is not None and self.ctx.trainer is not None + trainer = self.ctx.trainer + + # 4) Setup + init + before_train patches (mirror _run_trainer_lifecycle up + # to, but excluding, the training step). + self._run_phase_patches(phase="setup", backend_args=self.ctx.backend_args) + trainer.setup() + trainer.init() + self._run_phase_patches(phase="before_train", backend_args=self.ctx.backend_args) + + # 5) Build only the model (no datasets / no train loop). + build_fn = getattr(trainer, "setup_model_only", None) + if build_fn is None: + raise NotImplementedError( + f"Trainer '{type(trainer).__name__}' for framework " + f"'{self.ctx.framework}' does not implement setup_model_only()." + ) + build_fn() + return trainer + # --------------------------- Internal Steps --------------------------- # def _initialize_runtime_environment(self) -> None: @@ -173,11 +231,22 @@ def _initialize_environment(self) -> None: # setup_training_env expects a string path. setup_training_env(str(data_path), setup_hf=True) - def _initialize_configuration(self, module_name: str, overrides: Optional[List[str]] = None) -> None: - cfg_path = Path(self.args.config) - assert cfg_path.exists(), f"[Primus:TrainRuntime] Config file not found: {cfg_path}" - - primus_cfg = load_primus_config(cfg_path, self.args) + def _initialize_configuration( + self, + module_name: str, + overrides: Optional[List[str]] = None, + primus_config: Any = None, + ) -> None: + # Resolve a config path for diagnostics/context (falls back to args.config + # when a pre-loaded config is injected and args may lack a usable path). + cfg_path = Path(getattr(self.args, "config", None) or getattr(primus_config, "config_file", "") or "") + if primus_config is not None: + # Caller supplied a pre-loaded (and possibly mutated) runtime config; + # skip re-loading from disk so those mutations are preserved. + primus_cfg = primus_config + else: + assert cfg_path.exists(), f"[Primus:TrainRuntime] Config file not found: {cfg_path}" + primus_cfg = load_primus_config(cfg_path, self.args) # Reuse the legacy PrimusConfig that load_primus_config already parsed # (exposed as `_legacy`) instead of re-parsing the YAML a second time. diff --git a/primus/core/trainer/base_trainer.py b/primus/core/trainer/base_trainer.py index 86e447166..2b0a75ce0 100644 --- a/primus/core/trainer/base_trainer.py +++ b/primus/core/trainer/base_trainer.py @@ -77,7 +77,7 @@ def __init__(self, backend_args: Any = None, *args, **kwargs): # Cooperative multiple inheritance: pass kwargs to BaseModule if present in MRO, # otherwise call super().__init__() with no args to avoid object.__init__() error. - from primus.modules.base_module import BaseModule + from primus.core.base_module import BaseModule if BaseModule in type(self).__mro__: super().__init__(**kwargs) diff --git a/primus/core/utils/import_utils.py b/primus/core/utils/import_utils.py index 34c2de16d..1fd93c408 100644 --- a/primus/core/utils/import_utils.py +++ b/primus/core/utils/import_utils.py @@ -6,7 +6,7 @@ import importlib from functools import partial -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 def lazy_import(paths, symbol, log_prefix="[Primus]"): diff --git a/primus/modules/module_utils.py b/primus/core/utils/module_utils.py similarity index 100% rename from primus/modules/module_utils.py rename to primus/core/utils/module_utils.py diff --git a/primus/modules/__init__.py b/primus/modules/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/primus/modules/trainer/__init__.py b/primus/modules/trainer/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/primus/modules/trainer/base_trainer.py b/primus/modules/trainer/base_trainer.py deleted file mode 100644 index e555156cb..000000000 --- a/primus/modules/trainer/base_trainer.py +++ /dev/null @@ -1,24 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -from abc import ABC, abstractmethod - -import torch -from megatron.core.models.gpt import GPTModel - - -class BaseTrainer(ABC): - @abstractmethod - def get_batch(self, data_iterator): - pass - - @abstractmethod - def loss_func(self, loss_mask: torch.Tensor, output_tensor: torch.Tensor): - pass - - @abstractmethod - def forward_step(self, data_iterator, model: GPTModel): - pass diff --git a/primus/modules/trainer/maxtext/__init__.py b/primus/modules/trainer/maxtext/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/primus/modules/trainer/maxtext/pre_trainer.py b/primus/modules/trainer/maxtext/pre_trainer.py deleted file mode 100644 index 9c862186d..000000000 --- a/primus/modules/trainer/maxtext/pre_trainer.py +++ /dev/null @@ -1,144 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -from typing import Any, Dict - -from primus.core.utils import checker -from primus.modules.base_module import BaseModule -from primus.modules.module_utils import error_rank_0, log_rank_0, warning_rank_0 - - -class MaxTextPretrainTrainer(BaseModule): - def __init__(self, *args, **kwargs): - extra_args = kwargs.pop("extra_args", None) - super().__init__(*args, **kwargs) - - # important: make sure patch maxtext logger first - self.patch_maxtext_logger() - - self.primus_cfg = kwargs.pop("primus_config", None) - - if self.primus_cfg is None: - raise ValueError("primus_config is required") - self.primus_cfg.export_module_config("pre_trainer") - self.pre_trainer_cfg_path = self.primus_cfg.module_config_path("pre_trainer") - self.override_model_args = self.prepare_model_overrides(extra_args) - - def setup(self): - log_rank_0(f"setup MaxText") - - def init(self, *init_args, **kwargs): - import functools - - from MaxText import pyconfig - from MaxText.train import initialize - - argv = ["MaxText.train", self.pre_trainer_cfg_path] - log_rank_0(f"init MaxText with argv {argv}") - - if self.override_model_args: - _orig = pyconfig.initialize - pyconfig.initialize = functools.partial(_orig, **self.override_model_args) - try: - self.train_config, self.recorder, self.diagnostic_config = initialize(argv) - finally: - pyconfig.initialize = _orig - else: - self.train_config, self.recorder, self.diagnostic_config = initialize(argv) - - self._update_logger_rank() - - def run(self, *args, **kwargs): - log_rank_0(f"MaxText Pre-Trainer: begin training...") - - from MaxText.train import run - - run(self.train_config, self.recorder, self.diagnostic_config) - log_rank_0("MaxText Pre-Trainer: after training is done") - - def prepare_model_overrides(self, override_args: Dict[str, Any]): - """ - Monkey patch maxtext cli args to override model args dynamically. - Supports nested overrides like: - {"override_model": {"num_experts": 16, "base_num_decoder_layers": 4}} - - All override keys MUST be under the "model" key. - """ - - if not override_args: - warning_rank_0("MaxText Pre-Trainer: No override_args provided, skip patch.") - return {} - - warning_rank_0(f"MaxText Pre-Trainer: Applying override_args: {override_args}") - - # --- Step 1. Flatten any nested dict under 'override_model' - flat_overrides = {} - for k, v in override_args.items(): - if k != "override_model": - raise ValueError(f"Only the 'override_model' key is supported for overrides, found: {k}") - if not isinstance(v, dict): - raise ValueError( - f"MaxText Pre-Trainer: The value for 'override_model' must be a dict, got {type(v).__name__}." - ) - for subk, subv in v.items(): - if isinstance(subv, dict): - raise ValueError( - f"MaxText Pre-Trainer: Invalid override key-value detected: {k}.{subk}-{subv}" - ) - flat_overrides[subk] = subv - return flat_overrides - - def _update_logger_rank(self): - """Refresh Primus logger rank/world_size from JAX distributed state. - - The logger is created before ``jax.distributed.initialize()`` runs, - so it defaults to rank=0, world_size=1. After JAX init we have the - real values and can patch them in. - """ - import jax - - rank = jax.process_index() - world_size = jax.process_count() - - from primus.core.utils.logger import update_rank_info - from primus.modules.module_utils import set_logging_rank - - update_rank_info(rank, world_size) - set_logging_rank(rank, world_size) - log_rank_0( - f"JAX distributed ready: rank={rank}, world_size={world_size}, " - f"devices={jax.device_count()}, local_devices={jax.local_device_count()}" - ) - - def patch_maxtext_logger(self): - import logging - - from primus.core.utils.logger import _logger as primus_logger - - try: - import MaxText.max_logging as maxtext_logging - - if hasattr(maxtext_logging, "log"): - maxtext_logging.log = primus_logger.info - warning_rank_0("MaxText Pre-Trainer: patch logger successfully.") - else: - error_rank_0("MaxText Pre-Trainer: logging module does not have a 'log' function.") - except ImportError: - error_rank_0("MaxText Pre-Trainer: failed to import MaxText Pre-Trainer's logging module.") - - level_map = {"DEBUG": 10, "INFO": 20, "WARNING": 30, "ERROR": 40} - - stderr_sink_level = self.module_config.stderr_sink_level - checker.check_true(stderr_sink_level in level_map) - logging_level = level_map[stderr_sink_level] - - jax_loggers = [logging.getLogger("jax"), logging.getLogger("jaxlib")] - for jax_logger in jax_loggers: - jax_logger.setLevel(logging_level) - - warning_rank_0( - f"jax.logging_level is deprecated, set logging_level={logging_level} [stderr_sink_level]" - ) diff --git a/primus/modules/trainer/megatron/__init__.py b/primus/modules/trainer/megatron/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/primus/modules/trainer/megatron/model_provider.py b/primus/modules/trainer/megatron/model_provider.py deleted file mode 100644 index 38de87202..000000000 --- a/primus/modules/trainer/megatron/model_provider.py +++ /dev/null @@ -1,53 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -from types import MethodType -from typing import Callable, Optional, Union - -import torch -from megatron.core.models.gpt import GPTModel -from megatron.core.models.mamba import MambaModel -from megatron.training import get_args - -from primus.backends.megatron.core.extensions.logits_processor import fused_softcap - -import megatron.legacy.model # isort: skip - -g_final_logit_softcapping: Optional[float] = None -original_compute_language_model_loss: Optional[MethodType] = None - - -def wrapped_compute_language_model_loss(self, labels: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: - global g_final_logit_softcapping - assert g_final_logit_softcapping is not None - - logits = logits.float() - fused_softcap(logits, g_final_logit_softcapping) - - global original_compute_language_model_loss - return original_compute_language_model_loss(labels, logits) - - -def primus_model_provider( - model_provider: Callable, pre_process=True, post_process=True, vp_stage: Optional[int] = None -) -> Union[GPTModel, megatron.legacy.model.GPTModel, MambaModel]: - # get model - model = model_provider(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) - - args = get_args() - if args.final_logit_softcapping is not None and args.final_logit_softcapping > 0.0: - - global g_final_logit_softcapping - g_final_logit_softcapping = args.final_logit_softcapping - - # save original func - global original_compute_language_model_loss - original_compute_language_model_loss = model.compute_language_model_loss - - # wrap with logits softcapping - model.compute_language_model_loss = MethodType(wrapped_compute_language_model_loss, model) - - return model diff --git a/primus/modules/trainer/megatron/pre_trainer.py b/primus/modules/trainer/megatron/pre_trainer.py deleted file mode 100644 index 4609f9c1a..000000000 --- a/primus/modules/trainer/megatron/pre_trainer.py +++ /dev/null @@ -1,307 +0,0 @@ -############################################################################### -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -import collections -from functools import partial - -import torch -from megatron.core import mpu -from megatron.core.models.gpt import GPTModel -from megatron.core.rerun_state_machine import get_rerun_state_machine -from megatron.core.utils import StragglerDetector, get_attr_wrapped_model -from megatron.training import get_args, get_timers -from megatron.training.utils import ( - get_batch_on_this_cp_rank, - get_batch_on_this_tp_rank, - is_first_or_last_pipeline_stage, -) - -stimer = StragglerDetector() - -from .trainer import MegatronTrainer - -mb_batch = None - - -def get_batch_func(data_iterator, vp_stage=None): - # TODO: this is pretty hacky, find a better way - if not is_first_or_last_pipeline_stage(vp_stage): - return None, None, None, None, None - - # assert data_iterator is not None, f"data_iterator is None vp_stage: {vp_stage}" - # get batches based on the TP rank you are on - batch = get_batch_on_this_tp_rank(data_iterator) - - # slice batch along sequence dimension for context parallelism - args = get_args() - - if args.patch_zero_bubble: - from primus.backends.megatron.core.pipeline_parallel.zerobubble.zbpp_vars import ( - get_seq_split_idx, - ) - - global mb_batch - # "or 0" to support original 1f1b and interleaved-1f1b in schedules.py - seq_split_idx = get_seq_split_idx() or 0 - if seq_split_idx == 0: - # get batches based on the TP rank you are on - mb_batch = get_batch_on_this_tp_rank(data_iterator) - assert ( - mb_batch["attention_mask"] is None - ), "attention_mask should be None, please enable --no-create-attention-mask-in-dataloader" - batch = {} - for k in mb_batch.keys(): - v = mb_batch[k] - if v is None: - batch[k] = v - continue - - assert v.shape[1] % get_args().num_seq_splits == 0, f"{k} size {v.shape}" - start_idx = seq_split_idx * v.shape[1] // get_args().num_seq_splits - end_idx = (seq_split_idx + 1) * v.shape[1] // get_args().num_seq_splits - if len(v.shape) > 2: - batch[k] = v[:, start_idx:end_idx, :].contiguous() - else: - batch[k] = v[:, start_idx:end_idx].contiguous() - - if args.context_parallel_size > 1 and args.enable_primus_turbo and args.use_turbo_attention: - try: - from primus.backends.megatron.core.utils import ( - produce_attention_sharder, - shard_batch_on_this_cp_rank, - ) - except: - raise ImportError("Module 'primus_turbo' may not installed. Please install it") - sharder = produce_attention_sharder(args.cp_comm_type) - batch = shard_batch_on_this_cp_rank(sharder, batch) - else: - batch = get_batch_on_this_cp_rank(batch) - - return batch.values() - - -class DataLoaderStore: - cache = collections.deque() - - @classmethod - def push(cls, data_iterator, h2d_stream=False, vp_stage=None): - timers = get_timers() - # Get the batch. - timers("batch-generator", log_level=2).start() - global stimer - - with stimer(bdata=True): - if h2d_stream: - from primus.backends.megatron.core.pipeline_parallel.zerobubble.offload import ( - get_offload_h2d_stream, - ) - - load_event = torch.cuda.Event() - original_stream = torch.cuda.current_stream() - with torch.cuda.stream(get_offload_h2d_stream()): - data = get_batch_func(data_iterator, vp_stage) - for x in data: - if x is not None: - x.record_stream(original_stream) - load_event.record() - cls.cache.append((data, load_event)) - else: - cls.cache.append((get_batch_func(data_iterator, vp_stage), None)) - timers("batch-generator").stop() - - @classmethod - def pop(cls): - data, load_event = cls.cache.popleft() - if load_event: - load_event.wait() - return data - - -class MegatronPretrainTrainer(MegatronTrainer): - def __init__(self, *args, **kwargs): - kwargs["module_name"] = "pre_trainer" - - # Explicitly reject unknown extra_args - extra_args = kwargs.pop("extra_args", None) - if extra_args: - raise ValueError( - f"[MegatronPretrainTrainer] Unexpected extra_args detected: {extra_args}. " - f"Megatron backend does not support unregistered config keys." - ) - - try: - super().__init__(*args, **kwargs) - except Exception: - import traceback - - traceback.print_exc() - raise - - def get_batch(self, data_iterator, vp_stage=None): - """Generate a batch.""" - return get_batch_func(data_iterator, vp_stage) - - def loss_func(self, loss_mask: torch.Tensor, output_tensor: torch.Tensor): - """Loss function. - - Args: - loss_mask (torch.Tensor): Used to mask out some portions of the loss - output_tensor (torch.Tensor): The tensor with the losses - - Returns: - the loss scalar for this micro-batch - the number of non-padded tokens in this microbatch - a dict containing reporting metrics on the loss and number of tokens across - the data parallel ranks - """ - args = get_args() - - losses = output_tensor.float() - loss_mask = loss_mask.view(-1).float() - total_tokens = loss_mask.sum() - loss = torch.cat([torch.sum(losses.view(-1) * loss_mask).view(1), total_tokens.view(1)]) - - if args.context_parallel_size > 1: - torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group()) - - # Check individual rank losses are not NaN prior to DP all-reduce. - rerun_state_machine = get_rerun_state_machine() - if args.check_for_nan_in_loss_and_grad: - rerun_state_machine.validate_result( - result=loss[0], - rejection_func=torch.isnan, - message="found NaN in local forward loss calculation", - tolerance=0.0, # forward pass calculations are determinisic - fatal=True, - ) - rerun_state_machine.validate_result( - result=loss[0], - rejection_func=torch.isinf, - message="found Inf in local forward loss calculation", - tolerance=0.0, # forward pass calculations are determinisic - fatal=True, - ) - # Check for spiky loss - if args.check_for_spiky_loss: - rerun_state_machine.validate_result( - result=loss[0], - rejection_func=partial( - rerun_state_machine.is_unexpectedly_large, - threshold=SPIKY_LOSS_FACTOR, - context="loss", - ), - message="Spiky loss", - tolerance=0.0, # forward pass calculations are determinisic - fatal=False, - ) - # Reduce loss for logging. - reporting_loss = loss.clone().detach() - torch.distributed.all_reduce(reporting_loss, group=mpu.get_data_parallel_group()) - - # loss[0] is a view of loss, so it has ._base not None, which triggers assert error - # in core/pipeline_parallel/schedule.py::deallocate_output_tensor, calling .clone() - # on loss[0] fixes this - local_num_tokens = loss[1].clone().detach().to(torch.int) - return ( - loss[0].clone(), - local_num_tokens, - {"lm loss": (reporting_loss[0], reporting_loss[1])}, - ) - - def forward_step(self, data_iterator, model: GPTModel, return_schedule_plan=False): - """Forward training step. - - Args: - data_iterator : Input data iterator - model (GPTModel): The GPT Model - """ - args = get_args() - timers = get_timers() - - # Get the batch. - if not args.patch_zero_bubble: - timers("batch-generator", log_level=2).start() - global stimer - with stimer(bdata=True): - vp_stage = get_attr_wrapped_model(model, "vp_stage") - tokens, labels, loss_mask, attention_mask, position_ids = self.get_batch( - data_iterator, vp_stage - ) - timers("batch-generator").stop() - else: - from collections.abc import Iterable - - vp_stage = get_attr_wrapped_model(model, "vp_stage") - if ( - not isinstance(data_iterator, Iterable) and not data_iterator is None - ): # isinstance(data_iterator, DataLoaderStore): - tokens, labels, loss_mask, attention_mask, position_ids = data_iterator.pop() - else: - DataLoaderStore.push(data_iterator, h2d_stream=False, vp_stage=vp_stage) - tokens, labels, loss_mask, attention_mask, position_ids = DataLoaderStore.pop() - - with stimer: - if return_schedule_plan: - assert ( - args.overlap_moe_expert_parallel_comm - ), "overlap_moe_expert_parallel_comm must be enabled to return the schedule plan" - - # Schedule plan building is only supported for GPT models - # Check if this is a Mamba model - unwrapped_model = model - while hasattr(unwrapped_model, "module"): - unwrapped_model = unwrapped_model.module - model_class_name = unwrapped_model.__class__.__name__ - - if "Mamba" in model_class_name: - raise NotImplementedError( - "Schedule plan building is not supported for Mamba models. " - "Please disable overlap_moe_expert_parallel_comm for Mamba." - ) - - if args.patch_moe_overlap: - assert ( - not args.delay_wgrad_compute - ), "Primus MoE overlap handles wgrad separately from the original Megatron implementation" - from primus.backends.megatron.core.pipeline_parallel.zerobubble.zbpp_utils import ( - WeightGradStore, - ) - - WeightGradStore.enable_split_bw() - assert ( - WeightGradStore.split_bw() - ), "WeightGradStore.split_bw is not supported, please make sure overlap_grad_reduce is disabled and gradient_accumulation_fusion is enabled" - from primus.backends.megatron.core.models.common.model_chunk_schedule_plan import ( - TransformerModelChunkSchedulePlan, - ) - - schedule_plan = TransformerModelChunkSchedulePlan( - model, tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask - ) - else: - schedule_plan = model.build_schedule_plan( - tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask - ) - return schedule_plan, partial(self.loss_func, loss_mask) - else: - # Check if model supports loss_mask parameter - # MambaModel doesn't accept loss_mask, but GPTModel does - # Unwrap the model to get the actual model class - unwrapped_model = model - while hasattr(unwrapped_model, "module"): - unwrapped_model = unwrapped_model.module - model_class_name = unwrapped_model.__class__.__name__ - - if "Mamba" in model_class_name: - # MambaModel doesn't accept loss_mask parameter - output_tensor = model(tokens, position_ids, attention_mask, labels=labels) - else: - # GPTModel and other models accept loss_mask parameter - output_tensor = model( - tokens, position_ids, attention_mask, labels=labels, loss_mask=loss_mask - ) - - return output_tensor, partial(self.loss_func, loss_mask) diff --git a/primus/modules/trainer/megatron/sft_trainer.py b/primus/modules/trainer/megatron/sft_trainer.py deleted file mode 100644 index 6aec95d15..000000000 --- a/primus/modules/trainer/megatron/sft_trainer.py +++ /dev/null @@ -1,22 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -from .trainer import MegatronTrainer - - -class MegatronSFTTrainer(MegatronTrainer): - def __init__(self, *args, **kwargs): - kwargs["module_name"] = "sft_trainer" - super().__init__(*args, **kwargs) - - def get_batch(self, data_iterator): - raise NotImplementedError - - def loss_func(self, loss_mask: torch.Tensor, output_tensor: torch.Tensor): - raise NotImplementedError - - def forward_step(self, data_iterator, model: GPTModel): - raise NotImplementedError diff --git a/primus/modules/trainer/megatron/trainer.py b/primus/modules/trainer/megatron/trainer.py deleted file mode 100644 index 130d0edca..000000000 --- a/primus/modules/trainer/megatron/trainer.py +++ /dev/null @@ -1,2561 +0,0 @@ -############################################################################### -# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -import argparse -import dataclasses -import functools -import gc -import importlib.util -import inspect -import json -import os -import statistics -import sys -import time - -import megatron -import torch -import torch.distributed as dist -from megatron.core import mpu, parallel_state, tensor_parallel -from megatron.core.distributed import DistributedDataParallel as DDP -from megatron.core.distributed import finalize_model_grads -from megatron.core.distributed.distributed_data_parallel_config import ( - DistributedDataParallelConfig, -) -from megatron.core.distributed.torch_fully_sharded_data_parallel import ( - TorchFullyShardedDataParallel as torch_FSDP, -) -from megatron.core.utils import check_param_hashes_across_dp_replicas, get_model_config -from megatron.training.checkpointing import ( - checkpoint_exists, - load_checkpoint, - save_checkpoint, -) -from megatron.training.training import save_checkpoint_and_time - -from primus.backends.megatron.core.optimizer.moun import get_megatron_muon_optimizer -from primus.backends.megatron.core.optimizer.moun_optimizer_config import ( - MounOptimizerConfig, -) -from primus.backends.megatron.training.utils import is_pipeline_stage_containing_loss -from primus.core.utils.import_utils import get_custom_fsdp, get_model_provider - -try: - pass - - HAVE_FSDP2 = True -except ImportError: - HAVE_FSDP2 = False -from megatron.core.datasets.blended_megatron_dataset_builder import ( - BlendedMegatronDatasetBuilder, -) -from megatron.core.datasets.gpt_dataset import ( - GPTDataset, - GPTDatasetConfig, - MockGPTDataset, -) -from megatron.core.enums import ModelType -from megatron.core.num_microbatches_calculator import ( - get_current_global_batch_size, - get_current_running_global_batch_size, - get_num_microbatches, - update_num_microbatches, -) -from megatron.core.optimizer import get_megatron_optimizer, get_mup_config_overrides -from megatron.core.rerun_state_machine import ( - RerunDiagnostic, - RerunErrorInjector, - RerunMode, - get_rerun_state_machine, - initialize_rerun_state_machine, -) -from megatron.core.utils import check_param_hashes_across_dp_replicas, get_model_config -from megatron.training import ( - ft_integration, - get_args, - get_tensorboard_writer, - get_timers, - global_vars, - one_logger_utils, -) -from megatron.training.arguments import validate_args -from megatron.training.async_utils import ( - init_persistent_async_worker, - maybe_finalize_async_save, -) -from megatron.training.checkpointing import ( - checkpoint_exists, - load_args_from_checkpoint, - load_checkpoint, - save_checkpoint, -) -from megatron.training.global_vars import ( - get_args, - get_one_logger, - get_tensorboard_writer, - get_timers, - get_tokenizer, - get_wandb_writer, - set_global_variables, -) -from megatron.training.initialize import ( - _compile_dependencies, - _init_autoresume, - _initialize_distributed, - _initialize_tp_communicators, - _set_random_seed, - set_jit_fusion_options, - setup_logging, - write_args_to_tensorboard, -) -from megatron.training.theoretical_memory_usage import report_theoretical_memory -from megatron.training.training import ( - build_train_valid_test_data_iterators, - checkpoint_and_decide_exit, - disable_forward_pre_hook, - dummy_train_step, - enable_forward_pre_hook, - evaluate_and_print_results, - get_megatron_optimizer_config, - get_model, - get_optimizer_param_scheduler, - num_floating_point_operations, - post_training_step_callbacks, - preprocess_common_state_dict, - print_datetime, - should_disable_forward_pre_hook, -) -from megatron.training.utils import ( - append_to_progress_log, - calc_params_l2_norm, - get_blend_and_blend_per_split, - is_first_or_last_pipeline_stage, - logical_and_across_model_parallel_group, - reduce_max_stat_across_model_parallel_group, - report_memory, - unwrap_model, - update_use_dist_ckpt, -) -from megatron.training.yaml_arguments import validate_yaml - -from primus.backends.megatron.argument_builder import _load_megatron_defaults -from primus.backends.megatron.core.transformer.moe.moe_utils import track_moe_metrics -from primus.backends.megatron.training.global_vars import ( - get_mlflow_writer, - get_train_start_time, - set_primus_global_variables, - set_train_start_time, -) -from primus.backends.megatron.training.mlflow_setup import upload_mlflow_artifacts -from primus.backends.megatron.training.tokenizer.tokenizer import build_tokenizer -from primus.core.utils import checker, file_utils -from primus.core.utils.rocm_mem_info import get_rocm_smi_gpu_util, get_rocm_smi_mem_info -from primus.core.utils.yaml_utils import nested_namespace_to_dict -from primus.modules.base_module import BaseModule -from primus.modules.module_utils import ( - debug_rank_0, - log_kv_rank_0, - log_rank_0, - log_rank_last, - warning_rank_0, -) -from primus.modules.trainer.base_trainer import BaseTrainer -from primus.modules.trainer.megatron.model_provider import primus_model_provider - -from .utils import ( - is_v_schedule_enabled, - schedule_wrapper, - set_wandb_writer_patch, - validate_args_on_rocm, -) - -# The earliest we can measure the start time. -set_train_start_time() - - -def _normalize_data_path_arg(path_value): - """Normalize data path args to list form when paths are passed as strings.""" - if path_value is None: - return None - if isinstance(path_value, str): - return path_value.split() - if isinstance(path_value, (list, tuple)): - return list(path_value) - return path_value - - -class MegatronTrainer(BaseTrainer, BaseModule): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - self.is_v_schedule = is_v_schedule_enabled(self.module_config) - - self.app_metrics = {} - - # Cache for GPU utilization sampling to avoid rocm-smi overhead on every rank - # See: https://github.com/AMD-AGI/Primus/pull/439 (Copilot review) - self._last_rocm_gpu_util = None - self._last_rocm_gpu_util_time = 0.0 - self._rocm_gpu_util_sample_interval = 5.0 # seconds - - # disable all logging handlers - import logging - - for handler in logging.root.handlers[:]: - logging.root.removeHandler(handler) - - if "aiter" in logging.root.manager.loggerDict: - logging.getLogger("aiter").setLevel(logging.ERROR) - - def init(self, *init_args, **kwargs): - allowed_keys = { - "extra_args_provider", - "args_defaults", - "ignore_unknown_args", - "allow_no_cuda", - "skip_mpu_initialization", - "skip_setup", - } - - invalid_keys = set(kwargs.keys()) - allowed_keys - if invalid_keys: - raise TypeError(f"Invalid keyword arguments for MegatronTrainer: {invalid_keys}") - - log_rank_0(f"-run update_primus_config...") - self.update_primus_config( - args=self.module_config, - exp_root_path=self.exp_root_path, - exp_meta_info=self.exp_meta_info, - ) - - # Apply patches during initialization - # These patches need to be applied before model setup - from types import SimpleNamespace - - from primus.core.patches import run_patches - - # Construct a temporary module_config for patch context - # This allows patches to access Megatron args via get_args(ctx) - temp_module_config = SimpleNamespace() - temp_module_config.params = self.module_config - - run_patches( - backend="megatron", - phase="before_train", - backend_version="0.15.0rc8", - extra={ - "module_config": temp_module_config, - }, - ) - - # Initalize and get arguments, timers, and Tensorboard writer. - log_rank_0(f"-run initialize_megatron...") - self.initialize_megatron( - extra_args_provider=kwargs.get("extra_args_provider", None), - args_defaults=kwargs.get("args_defaults", {}), - ignore_unknown_args=kwargs.get("ignore_unknown_args", False), - allow_no_cuda=kwargs.get("allow_no_cuda", False), - skip_mpu_initialization=kwargs.get("skip_mpu_initialization", False), - ) - - args = get_args() - # There are some extra limitation on ROCm need extra validate. - validate_args_on_rocm(args) - - # Enable manually split layers in (interleaved) 1f1b pipeline - # parallelism by monkey patching - if args.decoder_pipeline_manual_split_list is not None: - log_rank_0( - f"-decoder_pipeline_manual_split_list has been deprecated, please use pipeline_model_parallel_layout instead" - ) - from .utils import set_manual_pipeline_split_patch, validate_manual_split - - log_rank_0(f"-monkey patch to enable manual pipeline split...") - if validate_manual_split(args): - set_manual_pipeline_split_patch(args) - - if args.recompute_layer_ids is not None: - from .utils import validate_specified_recompute_layers - - validate_specified_recompute_layers(args) - - if args.log_progress: - append_to_progress_log("Starting job") - - self.log_avg_reset_interval = args.log_avg_reset_interval - self.log_avg_skip_iterations = args.log_avg_skip_iterations - self.recent_tflop_throughputs = [] - self.recent_iteration_times = [] - self.recent_token_throughputs = [] - - # Initialize fault tolerance - # NOTE: ft_integration functions other than `setup` are no-op if the FT is not initialized - if args.enable_ft_package: - ft_integration.setup(args) - ft_integration.maybe_setup_simulated_fault() - - # Set pytorch JIT layer fusion options and warmup JIT functions. - set_jit_fusion_options() - - # Adjust the startup time so it reflects the largest value. - # This will be closer to what scheduler will see (outside of - # image ... launches. - start_time_tensor = torch.tensor([get_train_start_time()], dtype=torch.double, device="cuda") - torch.distributed.all_reduce(start_time_tensor, op=torch.distributed.ReduceOp.MIN) - set_train_start_time(start_time_tensor.item()) - - self.app_metrics["app_start_time"] = round(get_train_start_time() * 1000.0) - self.app_metrics["app_model_init_start_time"] = round(get_train_start_time() * 1000.0) - - log_rank_0( - "time to initialize megatron (seconds): {:.3f}".format(time.time() - get_train_start_time()) - ) - print_datetime("after megatron is initialized") - self.app_metrics["app_model_init_finish_time"] = one_logger_utils.get_timestamp_in_ms() - - # Track E2E metrics on pretrain start - one_logger_utils.on_pretrain_start() - - # Context used for persisting some state between checkpoint saves. - if args.non_persistent_ckpt_type == "local": - try: - from nvidia_resiliency_ext.checkpointing.local.ckpt_managers.local_manager import ( - LocalCheckpointManager, - ) - from nvidia_resiliency_ext.checkpointing.local.replication.strategies import ( - CliqueReplicationStrategy, - ) - except ModuleNotFoundError: - raise RuntimeError( - "The 'nvidia_resiliency_ext' module is required for local " - "checkpointing but was not found. Please ensure it is installed." - ) - - if args.replication: - repl_strategy = CliqueReplicationStrategy.from_replication_params( - args.replication_jump, args.replication_factor - ) - else: - repl_strategy = None - - self.checkpointing_context = { - "local_checkpoint_manager": LocalCheckpointManager( - args.non_persistent_local_ckpt_dir, repl_strategy=repl_strategy - ) - } - else: - self.checkpointing_context = {} - - if not kwargs.get("skip_setup", False): - self.setup() - - def update_primus_config( - self, - args, - exp_meta_info, - exp_root_path, - ): - # rank/world_size - args.rank = self.module_rank - args.world_size = self.module_world_size - args.local_rank = self.module_local_rank - log_kv_rank_0(f"-rank", f"{args.rank}") - log_kv_rank_0(f"-local_rank", f"{args.local_rank}") - log_kv_rank_0(f"-world_size", f"{args.world_size}") - - # cuda - if not args.use_torch_fsdp2 and not args.use_custom_fsdp: - os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" - else: - os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "8" - - # profile - if args.profile: - args.disable_tensorboard = False - - # checkpoint - ckpt_path = os.path.abspath(os.path.join(exp_root_path, "checkpoints")) - if args.save is not None: - warning_rank_0(f" args.save is deprecated, the checkpoint path is: {ckpt_path}") - args.save = ckpt_path - log_kv_rank_0(f"-save", f"{args.save}") - - # auto_continue_train - # Note that if args.auto_continue_train is enabled, check if there are existing save checkpoints - # for the current training experiment. If checkpoints are found, update the training - # configuration by disable finetuning (args.finetune), loading the optimizer state - # (args.no_load_optim), and loading the random number generator state (args.no_load_rng). - log_kv_rank_0(f"-auto_continue_train", f"{args.auto_continue_train}") - if args.auto_continue_train: - latest_file = f"{args.save}/latest_checkpointed_iteration.txt" - ckpt_exist = file_utils.is_file(latest_file) - if ckpt_exist: - with open(latest_file, "r") as file: - iter_str = file.read().strip() - log_rank_0(f"-find '{latest_file}', latest iteration is {iter_str}.") - if args.load != args.save: - warning_rank_0( - f"-set args.load={args.save}, path '{args.load}' is deprecated. [auto_continue_train]" - ) - args.load = args.save - if args.finetune: - args.finetune = False - warning_rank_0(f"-set args.finetune=False [auto_continue_train]") - if args.no_load_optim: - args.no_load_optim = False - warning_rank_0(f"-set args.no_load_optim=False [auto_continue_train]") - if args.no_load_rng: - args.no_load_rng = False - warning_rank_0(f"-set args.no_load_rng=False [auto_continue_train]") - if not args.use_checkpoint_args: - args.use_checkpoint_args = True - warning_rank_0(f"-set args.use_checkpoint_args=True [auto_continue_train]") - else: - log_rank_0(f"-{latest_file} does not exist, skip auto_continue_train.") - - # Auto-enable profiling and tensorboard when traces are needed: for MLflow upload - # (only if MLflow is enabled) or for local TraceLens report generation. - # Without this, generate_tracelens_report=True with profile=False would produce no traces. - needs_profiling = ( - ( - getattr(args, "mlflow_upload_traces", False) - or getattr(args, "mlflow_upload_tracelens_report", False) - ) - and not args.disable_mlflow - ) or getattr(args, "generate_tracelens_report", False) - if needs_profiling: - if not getattr(args, "profile", False): - args.profile = True - debug_rank_0("Auto-enabled profile=True for trace/tracelens (upload or local generation)") - if not getattr(args, "use_pytorch_profiler", False): - args.use_pytorch_profiler = True - debug_rank_0( - "Auto-enabled use_pytorch_profiler=True for trace/tracelens (upload or local generation)" - ) - if getattr(args, "disable_tensorboard", True): - args.disable_tensorboard = False - debug_rank_0("Auto-enabled tensorboard (disable_tensorboard=False) for profiler trace output") - - # tensorboard - if not args.disable_tensorboard: - tb_path = os.path.abspath(os.path.join(exp_root_path, "tensorboard")) - if args.tensorboard_dir is not None: - warning_rank_0(f"args.tensorboard_dir is deprecated, the tensorboard path is: {tb_path}") - args.tensorboard_dir = tb_path - else: - args.tensorboard_dir = None - log_kv_rank_0(f"-disable_tensorboard", f"{args.disable_tensorboard}") - log_kv_rank_0(f" -tensorboard_dir", f"{args.tensorboard_dir}") - - # wandb - if not args.disable_wandb: - wandb_path = exp_root_path - if args.wandb_save_dir is not None: - warning_rank_0(f"args.wandb_save_dir is deprecated, the wandb path is: {wandb_path}/wandb") - if not hasattr(args, "wandb_project") or args.wandb_project is None: - args.wandb_project = f"{exp_meta_info['work_group']}_{exp_meta_info['user_name']}" - debug_rank_0(f" -create new wandb project name: {args.wandb_project}") - if not hasattr(args, "wandb_exp_name") or args.wandb_exp_name is None: - args.wandb_exp_name = exp_meta_info["exp_name"] - debug_rank_0(f" -create new exp name: {args.wandb_exp_name}") - args.wandb_save_dir = wandb_path - elif args.wandb_project is not None: - args.wandb_project = None - debug_rank_0(f"args.wandb_project is disabled, as args.disable_wandb=True.") - log_kv_rank_0(f"-disable_wandb", f"{args.disable_wandb}") - if not args.disable_wandb and "WANDB_API_KEY" not in os.environ: - warning_rank_0( - "The environment variable WANDB_API_KEY is not set. " - "Please set it before proceeding or enable 'disable_wandb' in yaml config" - ) - log_kv_rank_0(f" -wandb_project", f"{args.wandb_project}") - log_kv_rank_0(f" -wandb_exp_name", f"{args.wandb_exp_name}") - log_kv_rank_0(f" -wandb_save_dir", f"{args.wandb_save_dir}") - log_kv_rank_0(f" -wandb_entity", f"{args.wandb_entity}") - - # mlflow - log_kv_rank_0(f"-disable_mlflow", f"{args.disable_mlflow}") - if not args.disable_mlflow: - if not hasattr(args, "mlflow_run_name") or args.mlflow_run_name is None: - args.mlflow_run_name = f"{exp_meta_info['work_group']}_{exp_meta_info['user_name']}" - debug_rank_0(f" -create new mlflow run name: {args.mlflow_run_name}") - elif args.mlflow_run_name is not None: - args.mlflow_run_name = None - args.mlflow_experiment_name = None - debug_rank_0(f"args.mlflow_run_name is disabled, as args.disable_mlflow=True.") - if not args.disable_mlflow and "DATABRICKS_HOST" not in os.environ: - warning_rank_0( - "The environment variable DATABRICKS_HOST is not set. " - "Please set it before proceeding or enable 'disable_mlflow' in yaml config" - ) - log_kv_rank_0(f" -mlflow_run_name", f"{args.mlflow_run_name}") - log_kv_rank_0(f" -mlflow_experiment_name", f"{args.mlflow_experiment_name}") - - # sink_level: logging_level - level_map = {"DEBUG": 10, "INFO": 20, "WARNING": 30, "ERROR": 40} - checker.check_true(args.stderr_sink_level in level_map) - logging_level = level_map[args.stderr_sink_level] - if args.logging_level is not None: - warning_rank_0( - f"-args.logging_level is deprecated, set args.logging_level={logging_level} [stderr_sink_level]" - ) - args.logging_level = logging_level - - # update data path - # "data1 data2 data3" -> ['data1', 'data2', 'data3'] - if args.data_path is not None: - args.data_path = _normalize_data_path_arg(args.data_path) - log_rank_0(f"-data_path: {args.data_path}") - - if args.train_data_path is not None: - args.train_data_path = _normalize_data_path_arg(args.train_data_path) - log_rank_0(f"-train_data_path: {args.train_data_path}") - if args.valid_data_path is not None: - args.valid_data_path = _normalize_data_path_arg(args.valid_data_path) - log_rank_0(f"-valid_data_path: {args.valid_data_path}") - if args.test_data_path is not None: - args.test_data_path = _normalize_data_path_arg(args.test_data_path) - log_rank_0(f"-test_data_path: {args.test_data_path}") - - # update sp - if args.tensor_model_parallel_size == 1: - args.sequence_parallel = False - - if args.iterations_to_skip is None: - args.iterations_to_skip = [] - - # support moe_freq_type - ensure moe_layer_freq has a default value - if not hasattr(args, "moe_layer_freq"): - args.moe_layer_freq = 1 - elif isinstance(args.moe_layer_freq, str): - try: - args.moe_layer_freq = eval(args.moe_layer_freq) - except Exception: - raise ValueError(f"Invalid moe_layer_freq format: {args.moe_layer_freq}") - - if args.mock_data: - args.data_path = None - args.train_data_path = None - args.valid_data_path = None - args.test_data_path = None - - # Determine model type (gpt or mamba) - model_type = getattr(args, "model_type", "gpt") - log_rank_0(f"-detected model_type: {model_type}") - - # Ensure required attributes have safe defaults if missing from config - if not hasattr(args, "final_logit_softcapping"): - args.final_logit_softcapping = None - if not hasattr(args, "router_logit_softcapping"): - args.router_logit_softcapping = None - - # Only pass model_type parameter when it's "mamba" to maintain backward compatibility - # with main branch behavior for "gpt" (default) case - if args.final_logit_softcapping is not None and args.final_logit_softcapping > 0.0: - log_rank_0(f"-enable final_logit_softcapping: {args.final_logit_softcapping}") - if model_type == "mamba": - self.model_provider = functools.partial( - primus_model_provider, get_model_provider(model_type=model_type) - ) - else: - self.model_provider = functools.partial(primus_model_provider, get_model_provider()) - else: - if model_type == "mamba": - log_rank_0(f"-getting model provider for model_type={model_type}") - model_provider = get_model_provider(model_type=model_type) - log_rank_0(f"-model_provider: {model_provider}") - self.model_provider = model_provider - else: - # For "gpt" (default), call without arguments to match main branch behavior - self.model_provider = get_model_provider() - - if args.router_logit_softcapping is not None and args.router_logit_softcapping > 0.0: - log_rank_0(f"-enable router_logit_softcapping: {args.router_logit_softcapping}") - - def vocab_size_with_padding(self, orig_vocab_size, args): - """Pad vocab size so it is divisible by model parallel size and - still having GPU friendly size.""" - - after = orig_vocab_size - multiple = args.make_vocab_size_divisible_by * args.tensor_model_parallel_size - while (after % multiple) != 0: - after += 1 - debug_rank_0( - " -padded vocab (size: {}) with {} dummy tokens " - "(new size: {})".format(orig_vocab_size, after - orig_vocab_size, after) - ) - return after - - def setup(self): - args = get_args() - timers = get_timers() - # Model, optimizer, and learning rate. - timers("model-and-optimizer-setup", log_level=0).start(barrier=True) - self.app_metrics["app_build_optimizer_start_time"] = one_logger_utils.get_timestamp_in_ms() - log_rank_0(f"-setup_model_and_optimizer...") - self.model, self.optimizer, self.opt_param_scheduler = self.setup_model_and_optimizer( - self.model_provider, - ModelType.encoder_or_decoder, - checkpointing_context=self.checkpointing_context, - ) - - timers("model-and-optimizer-setup").stop() - print_datetime("after model, optimizer, and learning rate " "scheduler are built") - self.app_metrics["app_build_optimizer_finish_time"] = one_logger_utils.get_timestamp_in_ms() - self.config = get_model_config(self.model[0]) - - # Data stuff. - self.app_metrics["app_build_dataiters_start_time"] = one_logger_utils.get_timestamp_in_ms() - timers("train/valid/test-data-iterators-setup", log_level=0).start(barrier=True) - - def train_valid_test_datasets_provider_func(train_val_test_num_samples, vp_stage=None): - return self.train_valid_test_datasets_provider(train_val_test_num_samples, vp_stage=vp_stage) - - train_valid_test_datasets_provider_func.is_distributed = True - - if args.virtual_pipeline_model_parallel_size is not None: - self.train_data_iterator = [] - self.valid_data_iterator = [] - self.test_data_iterator = [] - for vp_stage in range(len(self.model)): - dataset_provider_parameters = inspect.signature( - train_valid_test_datasets_provider_func - ).parameters - assert ( - "vp_stage" in dataset_provider_parameters - ), "vp_stage must be a kwarg in train_valid_test_dataset_provider when using virtual pipeline parallelism" - vp_stage_train_valid_test_dataset_provider = functools.partial( - train_valid_test_datasets_provider_func, vp_stage=vp_stage - ) - if getattr(train_valid_test_datasets_provider_func, "is_distributed", False): - vp_stage_train_valid_test_dataset_provider.is_distributed = True - iterators = build_train_valid_test_data_iterators(vp_stage_train_valid_test_dataset_provider) - self.train_data_iterator.append(iterators[0]) - self.valid_data_iterator.append(iterators[1]) - self.test_data_iterator.append(iterators[2]) - else: - ( - self.train_data_iterator, - self.valid_data_iterator, - self.test_data_iterator, - ) = build_train_valid_test_data_iterators(train_valid_test_datasets_provider_func) - timers("train/valid/test-data-iterators-setup").stop() - print_datetime("after dataloaders are built") - self.app_metrics["app_build_dataiters_finish_time"] = one_logger_utils.get_timestamp_in_ms() - - # Track if training is enabled. Can only be done once args.do_train is assigned after dataloader is built. - # The pinned Megatron-LM's ``track_config_flags`` signature varies (6 vs 8 - # positional args across commits); pass only what the installed version accepts. - import inspect as _inspect - - _tcf_args = [ - args.train_iters, - args.skip_train, - args.do_train, - args.do_valid, - args.do_test, - args.dataloader_type, - args.retro_project_dir, - args.retro_cyclic_train_iters, - ] - try: - _tcf_nparams = len(_inspect.signature(one_logger_utils.track_config_flags).parameters) - except (TypeError, ValueError): - _tcf_nparams = len(_tcf_args) - one_logger_utils.track_config_flags(*_tcf_args[:_tcf_nparams]) - - # Print setup timing. - log_rank_0("done with setup ...") - timers.log( - ["model-and-optimizer-setup", "train/valid/test-data-iterators-setup"], - barrier=True, - ) - - one_logger = get_one_logger() - one_logger and one_logger.log_metrics(self.app_metrics) - - def core_gpt_dataset_config_from_args(self, args): - tokenizer = get_tokenizer() - - # Keep legacy trainer aligned with upstream pretrain_gpt dataset argument handling. - blend, blend_per_split = get_blend_and_blend_per_split(args) - - sequences_per_dataset = None - per_dataset_sequences_path = getattr(args, "per_dataset_sequences_path", None) - if per_dataset_sequences_path is not None: - with open(per_dataset_sequences_path, "r") as f: - sequences_per_dataset = json.load(f) - - data_args = { - "random_seed": args.seed, - "sequence_length": args.seq_length, - "blend": blend, - "blend_per_split": blend_per_split, - "split": args.split, - "multiple_validation_sets": getattr(args, "multiple_validation_sets", None), - "full_validation": getattr(args, "full_validation", None), - "num_dataset_builder_threads": args.num_dataset_builder_threads, - "path_to_cache": args.data_cache_path, - "mmap_bin_files": args.mmap_bin_files, - "tokenizer": tokenizer, - "reset_position_ids": args.reset_position_ids, - "reset_attention_mask": args.reset_attention_mask, - "eod_mask_loss": args.eod_mask_loss, - "create_attention_mask": args.create_attention_mask_in_dataloader, - "object_storage_cache_path": getattr(args, "object_storage_cache_path", None), - "mid_level_dataset_surplus": getattr(args, "mid_level_dataset_surplus", 0.005), - "allow_ambiguous_pad_tokens": getattr(args, "allow_ambiguous_pad_tokens", False), - "fast_cache_load": getattr(args, "dataloader_fast_cache_load", False), - "sequences_per_dataset": sequences_per_dataset, - "defer_npy_index_mmap": getattr(args, "dataloader_defer_npy_index_mmap", False), - "context_parallel_size": getattr(args, "context_parallel_size", 1), - "data_parallel_size": getattr(args, "data_parallel_size", 1), - "sequence_parallel_size": getattr(args, "tensor_model_parallel_size", 1) - * getattr(args, "sequence_parallel", False), - "hybrid_context_parallel": getattr(args, "hybrid_context_parallel", False), - } - - return GPTDatasetConfig(**data_args) - - def train_valid_test_datasets_provider(self, train_val_test_num_samples, vp_stage=None): - """Build the train test and validation datasets. - - Args: - train_val_test_num_samples : A list containing the number of samples in train test and validation. - """ - args = get_args() - - config = self.core_gpt_dataset_config_from_args(args) - - if args.mock_data: - dataset_type = MockGPTDataset - else: - dataset_type = GPTDataset - - def is_dataset_built_on_rank(vp_stage=None): - return ( - is_first_or_last_pipeline_stage(vp_stage) - and parallel_state.get_tensor_model_parallel_rank() == 0 - ) - - log_rank_0("> building train, validation, and test datasets for GPT ...") - train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder( - dataset_type, - train_val_test_num_samples, - functools.partial(is_dataset_built_on_rank, vp_stage=vp_stage), - config, - ).build() - - log_rank_0("> finished creating GPT datasets ...") - - return train_ds, valid_ds, test_ds - - def initialize_megatron( - self, - extra_args_provider=None, - args_defaults={}, - ignore_unknown_args=False, - allow_no_cuda=False, - skip_mpu_initialization=False, - get_embedding_ranks=None, - get_position_embedding_ranks=None, - ): - """Set global variables, initialize distributed, and - set autoresume and random seeds. - `allow_no_cuda` should not be set unless using megatron for cpu only - data processing. In general this arg should not be set unless you know - what you are doing. - Returns a function to finalize distributed env initialization - (optionally, only when args.lazy_mpu_init == True) - """ - if not allow_no_cuda: - # Make sure cuda is available. - assert torch.cuda.is_available(), "Megatron requires CUDA." - - # Note: parse_args is deprecated in megatron trainer, use primus yaml config instead. - # Parse arguments - # args = parse_args(extra_args_provider, ignore_unknown_args) - - # Use trainer args from primus - # args = self.module_config - - # Build Megatron arguments by merging Primus config into Megatron defaults. - # 1) Load Megatron defaults (no Primus overrides) - megatron_defaults = _load_megatron_defaults() - - # 2) Convert Primus module_config (nested namespace) to a plain dict - primus_args = nested_namespace_to_dict(self.module_config) - - # 3) Merge: defaults < primus_args - merged_args = megatron_defaults.copy() - merged_args.update(primus_args) - - # 4) Convert to Namespace for compatibility with downstream Megatron utilities - args = argparse.Namespace(**merged_args) - - # Prep for checkpoint conversion. - if args.ckpt_convert_format is not None: - assert args.ckpt_convert_save is not None - assert args.load is not None - args.exit_on_missing_checkpoint = True - - log_kv_rank_0(f"-load", f"{args.load}") - log_kv_rank_0(f"-use_checkpoint_args", f"{args.use_checkpoint_args}") - if args.use_checkpoint_args or args_defaults.get("use_checkpoint_args", False): - checker.check_true(args.load is not None, "--use-checkpoints-args requires --load argument") - log_rank_0(f"-load_args_from_checkpoint...") - assert args.non_persistent_ckpt_type != "local", ( - "--use-checkpoint-args is not supported with --non_persistent_ckpt_type=local. " - "Two-stage checkpoint loading is not implemented, and all arguments must be defined " - "before initializing LocalCheckpointManager." - ) - load_args_from_checkpoint(args) - - if args.async_save and args.use_persistent_ckpt_worker: - init_persistent_async_worker() - - checker.check_true(args.yaml_cfg is None, "Xpipe doesn't support megatron yaml config.") - if args.yaml_cfg is not None: - args = validate_yaml(args, args_defaults) - else: - if args.decoder_pipeline_manual_split_list is not None: - from .utils import validate_args_modified - - ori_code = "if args.decoder_first_pipeline_num_layers is None and args.decoder_last_pipeline_num_layers is None:" - new_code = ( - "if args.decoder_pipeline_manual_split_list is None and " + ori_code.split("if ")[-1] - ) - - validate_args_modified(args, args_defaults, ori_code=ori_code, new_code=new_code) - elif args.fp4 is not None: - # TODO(ruibin): Remove it when ROCm TE upgrade to 2.7.0.dev0 - from .utils import validate_args_modified - - ori_code = """raise ValueError("--fp4-format requires Transformer Engine >= 2.7.0.dev0 for NVFP4BlockScaling support.")""" - new_code = """pass""" - - validate_args_modified(args, args_defaults, ori_code=ori_code, new_code=new_code) - else: - validate_args(args, args_defaults) - - # monkey patch _set_wandb_writer before set_global_variables - log_rank_0(f"-monkey patch megatron.training.global_vars._set_wandb_writer...") - megatron.training.global_vars._set_wandb_writer = set_wandb_writer_patch - - # set global args, build tokenizer, and set adlr-autoresume, - # tensorboard-writer, and timers. - log_rank_0(f"-set_global_variables...") - set_global_variables(args, build_tokenizer=False) - log_rank_0(f"-set_primus_global_variables...") - set_primus_global_variables(args) - args = get_args() - - # set tokenizer - log_rank_0(f"-build_tokenizer...") - global_vars._ensure_var_is_not_initialized(global_vars._GLOBAL_TOKENIZER, "tokenizer") - global_vars._GLOBAL_TOKENIZER = build_tokenizer(args) - - # set logging level - setup_logging() - - # init rerun state - def state_save_func(): - return {"rng_tracker_states": tensor_parallel.get_cuda_rng_tracker().get_states()} - - def state_restore_func(state_dict): - if state_dict["rng_tracker_states"]: - tensor_parallel.get_cuda_rng_tracker().set_states(state_dict["rng_tracker_states"]) - - initialize_rerun_state_machine( - state_save_func=state_save_func, - state_restore_func=state_restore_func, - mode=RerunMode(args.rerun_mode), - error_injector=RerunErrorInjector( - error_injection_rate=args.error_injection_rate, - error_injection_type=RerunDiagnostic(args.error_injection_type), - ), - result_rejected_tracker_filename=args.result_rejected_tracker_filename, - ) - - # torch.distributed initialization - def finish_mpu_init(): - args = get_args() - # Pytorch distributed. - log_rank_0(f"-initialize_distributed...") - _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks, None) - - # Random seeds for reproducibility. - log_kv_rank_0(f"-seeds", f"{args.seed}") - _set_random_seed( - args.seed, - args.data_parallel_random_init, - args.te_rng_tracker, - args.inference_rng_tracker, - use_cudagraphable_rng=args.enable_cuda_graph, - ) - - # Setup MoE aux loss scale value. - if args.num_experts is not None: - from megatron.core.transformer.moe.router import MoEAuxLossAutoScaler - - MoEAuxLossAutoScaler.set_loss_scale(torch.ones(1, device=torch.cuda.current_device())) - - if skip_mpu_initialization: - return None - - args = get_args() - log_kv_rank_0(f"-lazy_mpu_init", f"{args.lazy_mpu_init}") - if args.lazy_mpu_init: - # TODO is this still a necessary option? - args.use_cpu_initialization = True - # delayed initialization of DDP-related stuff - # We only set basic DDP globals - mpu.set_tensor_model_parallel_world_size(args.tensor_model_parallel_size) - # and return function for external DDP manager - # to call when it has DDP initialized - mpu.set_tensor_model_parallel_rank(args.rank) - return finish_mpu_init - else: - # Megatron's MPU is the master. Complete initialization right away. - finish_mpu_init() - - # Autoresume. - _init_autoresume() - - # Compile dependencies. - if not args.disable_compile_dependencies: - log_rank_0(f"-compile_dependencies...") - _compile_dependencies() - - if args.tp_comm_overlap: - # TODO: Should this be activated with just decoder-tp-comm-overlap too? - _initialize_tp_communicators() - - # No continuation function - return None - - def setup_model_and_optimizer( - self, - model_provider_func, - model_type, - no_wd_decay_cond=None, - scale_lr_cond=None, - lr_mult=1.0, - checkpointing_context=None, - ): - """Setup model and optimizer.""" - args = get_args() - timers = get_timers() - one_logger = get_one_logger() - - # Primus Turbo patches are now applied automatically via the new patch system - # in BaseTrainer.run() with phase="before_train" - if importlib.util.find_spec("primus_turbo") is not None: - args = get_args() - if args.tensor_model_parallel_size == 1: - if args.enable_primus_turbo: - log_rank_0(f"use pt backend...") - else: - log_rank_0(f"use te backend...") - elif args.enable_primus_turbo: - log_rank_0(f"primus turbo does not support tp, use te backend...") - else: - log_rank_0(f"use te backend...") - - log_rank_0(f"-run get_model") - log_rank_0(f"-model_provider_func: {model_provider_func}") - log_rank_0(f"-model_type: {model_type}") - model = get_model(model_provider_func, model_type) - log_rank_0(model) - # get_megatron_optimizer will use the ddp_config - if isinstance(model[0], torch_FSDP): - model[0].ddp_config = DistributedDataParallelConfig() - model[0].ddp_config.use_custom_fsdp = False - - unwrapped_model = unwrap_model(model) - - config, config_overrides = get_megatron_optimizer_config(args) - config.timers = timers - if getattr(args, "use_mup", False): - model_config_source = unwrapped_model[0] if isinstance(unwrapped_model, list) else unwrapped_model - model_config = get_model_config(model_config_source) - mup_overrides = get_mup_config_overrides( - config=config, - mup_width_mult=model_config.mup_width_mult, - optimizer_type=config.optimizer, - ) - if mup_overrides: - config_overrides = {**(config_overrides or {}), **mup_overrides} - - if "muon" not in args.optimizer: - optimizer = get_megatron_optimizer( - config, - model, - config_overrides=config_overrides, - use_gloo_process_groups=args.enable_gloo_process_groups, - dump_param_to_param_group_map=getattr(args, "dump_param_to_param_group_map", None), - ) - else: - kwargs = {} - for f in dataclasses.fields(MounOptimizerConfig): - if hasattr(args, f.name): - kwargs[f.name] = getattr(args, f.name) - - config = MounOptimizerConfig(**kwargs) - config.timers = timers - optimizer = get_megatron_muon_optimizer( - config, - model, - config_overrides=config_overrides, - use_gloo_process_groups=args.enable_gloo_process_groups, - layer_wise_distributed_optimizer="dist" in config.optimizer, - dump_param_to_param_group_map=getattr(args, "dump_param_to_param_group_map", None), - ) - - opt_param_scheduler = get_optimizer_param_scheduler(optimizer) - - if args.moe_use_upcycling: - torch.distributed.barrier() - assert not checkpoint_exists(args.save), ( - "The upcycling destination directory already exists. " - "Please check if --moe-use-upcycling is mistakenly enabled. " - "Upcycling should only be set for the first run when converting the dense model. " - "All subsequent runs should remove this flag. " - ) - num_experts = args.num_experts - args.num_experts = None - expert_model_parallel_size = args.expert_model_parallel_size - args.expert_model_parallel_size = 1 - dense_model_for_upcycling = get_model(model_provider_func, model_type) - args.num_experts = num_experts - args.expert_model_parallel_size = expert_model_parallel_size - _, args.num_floating_point_operations_so_far = upcycling_utils.load_and_upcycle_model( - load_checkpoint, - unwrapped_model, - dense_model_for_upcycling, - load_kwargs={ - "model": dense_model_for_upcycling, - "optimizer": None, - "opt_param_scheduler": None, - }, - ) - args.iteration = 1 - save_checkpoint( - args.iteration, - model, - None, - None, - args.num_floating_point_operations_so_far, - ) - torch.distributed.barrier() - del dense_model_for_upcycling - if (args.fp16 or args.bf16) and optimizer is not None: - optimizer.reload_model_params() - log_rank_0(f"Upcycled checkpoint saved to {args.save}") - - if (args.load is not None or args.pretrained_checkpoint is not None) and not args.moe_use_upcycling: - one_logger and one_logger.log_metrics( - {"load_checkpoint_start_time": one_logger_utils.get_timestamp_in_ms()} - ) - timers("load-checkpoint", log_level=0).start(barrier=True) - - log_rank_0(f"-run load_checkpoint") - log_rank_0(f" -args.load={args.load}") - args.iteration, args.num_floating_point_operations_so_far = load_checkpoint( - model, - optimizer, - opt_param_scheduler, - checkpointing_context=self.checkpointing_context, - skip_load_to_model_and_opt=HAVE_FSDP2 and args.use_torch_fsdp2, - ) - if ( - HAVE_FSDP2 - and args.use_torch_fsdp2 - and optimizer is not None - and hasattr(optimizer, "finalize_dist_ckpt_load") - and args.iteration > 0 - ): - optimizer.finalize_dist_ckpt_load(args.iteration) - timers("load-checkpoint").stop(barrier=True) - timers.log(["load-checkpoint"]) - one_logger and one_logger.log_metrics( - { - "load_checkpoint_finish_time": one_logger_utils.get_timestamp_in_ms(), - "load_checkpoint_time": timers("load-checkpoint").active_time(), - } - ) - else: - args.iteration = 0 - args.num_floating_point_operations_so_far = 0 - - # get model without FP16 and/or DDP wrappers - if ( - args.iteration == 0 - and len(unwrapped_model) == 1 - and hasattr(unwrapped_model[0], "init_state_dict_from_bert") - ): - log_rank_0("Initializing ICT from pretrained BERT model") - unwrapped_model[0].init_state_dict_from_bert() - if args.fp16: - optimizer.reload_model_params() - - # Convert checkpoint format. - if args.ckpt_convert_format is not None: - load_ckpt_format = args.ckpt_format - args.ckpt_format = args.ckpt_convert_format - args.save = os.path.join(args.ckpt_convert_save, args.ckpt_convert_format) - update_use_dist_ckpt(args) - - save_checkpoint( - args.iteration, - model, - optimizer, - opt_param_scheduler, - args.num_floating_point_operations_so_far, - preprocess_common_state_dict_fn=preprocess_common_state_dict, - ) - - log_rank_0("> converted checkpoint: %s -> %s." % (load_ckpt_format, args.ckpt_format)) - torch.distributed.barrier() - exit() - - return model, optimizer, opt_param_scheduler - - def run(self, *args, **kwargs): - one_logger = get_one_logger() - args = get_args() - - process_non_loss_data_func = None - non_loss_data_func = None - if not args.skip_train: - log_rank_0("training ...") - - if args.dataloader_type == "cyclic" and args.retro_project_dir: - assert args.retro_cyclic_train_iters is not None - args.train_iters = args.retro_cyclic_train_iters - log_rank_0("retro cyclic train iters : %d" % args.train_iters) - - iteration = 0 - if args.do_train and args.train_iters > 0: - iteration, num_floating_point_operations_so_far = self.train( - self.forward_step, - self.model, - self.optimizer, - self.opt_param_scheduler, - self.train_data_iterator, - self.valid_data_iterator, - process_non_loss_data_func, - self.config, - self.checkpointing_context, - non_loss_data_func, - ) - - print_datetime("after training is done") - - if ( - args.save - and iteration != 0 - and iteration % args.save_interval != 0 - and not args.disable_last_saving - ): - save_checkpoint( - iteration, - self.model, - self.optimizer, - self.opt_param_scheduler, - num_floating_point_operations_so_far, - self.checkpointing_context, - train_data_iterator=self.train_data_iterator, - preprocess_common_state_dict_fn=preprocess_common_state_dict, - ) - - one_logger and one_logger.log_metrics( - {"app_train_loop_finish_time": one_logger_utils.get_timestamp_in_ms()} - ) - - else: - log_rank_0("skipping training (--skip-train is on) ...") - - iteration = args.iteration - - if args.do_valid: - prefix = f"iteration {iteration} on validation set" - evaluate_and_print_results( - prefix, - self.forward_step, - self.valid_data_iterator, - self.model, - iteration, - process_non_loss_data_func, - self.config, - verbose=True, - write_to_tensorboard=not args.skip_train, - non_loss_data_func=non_loss_data_func, - ) - - if args.do_test: - prefix = f"iteration {iteration} on test set" - evaluate_and_print_results( - prefix, - self.forward_step, - self.test_data_iterator, - self.model, - iteration, - process_non_loss_data_func, - self.config, - verbose=True, - write_to_tensorboard=not args.skip_train, - non_loss_data_func=non_loss_data_func, - ) - - wandb_writer = get_wandb_writer() - if wandb_writer: - wandb_writer.finish() - - ft_integration.on_checkpointing_start() - maybe_finalize_async_save(blocking=True, terminate=True) - ft_integration.on_checkpointing_end(is_async_finalization=True) - - mlflow_writer = get_mlflow_writer() - # Barrier to ensure all ranks have finished writing files before upload. - # Must run on ALL ranks to avoid deadlock (only last rank has mlflow_writer). - if dist.is_initialized(): - dist.barrier() - - # Always call: uploads to MLflow when enabled; when MLflow disabled, still runs - # local-only TraceLens report generation if generate_tracelens_report=True. - try: - upload_mlflow_artifacts( - tensorboard_dir=args.tensorboard_dir, - exp_root_path=self.exp_root_path, - upload_traces=getattr(args, "mlflow_upload_traces", False), - upload_logs=getattr(args, "mlflow_upload_logs", False), - generate_tracelens_report=getattr(args, "generate_tracelens_report", False), - upload_tracelens_report=getattr(args, "mlflow_upload_tracelens_report", False), - tracelens_ranks=getattr(args, "mlflow_tracelens_ranks", None), - tracelens_output_format=getattr(args, "mlflow_tracelens_output_format", "xlsx"), - tracelens_cleanup_after_upload=getattr(args, "mlflow_tracelens_cleanup_after_upload", False), - tracelens_auto_install=getattr(args, "mlflow_tracelens_auto_install", True), - ) - except Exception as e: - import logging - - logging.getLogger(__name__).warning("[MLflow] Artifact upload failed: %s", e) - finally: - if mlflow_writer: - mlflow_writer.end_run() - if dist.is_initialized(): - dist.barrier() - - one_logger and one_logger.log_metrics({"app_finish_time": one_logger_utils.get_timestamp_in_ms()}) - - ft_integration.shutdown() - one_logger_utils.finish() - - # clean up torch pg resources on exit - if dist.is_initialized(): - dist.destroy_process_group() - - def train( - self, - forward_step_func, - model, - optimizer, - opt_param_scheduler, - train_data_iterator, - valid_data_iterator, - process_non_loss_data_func, - config, - checkpointing_context, - non_loss_data_func, - ): - """Training function: run train_step desired number of times, run validation, checkpoint.""" - args = get_args() - timers = get_timers() - one_logger = get_one_logger() - - if args.run_workload_inspector_server: - try: - import threading - - from workload_inspector.utils.webserver import run_server - - threading.Thread(target=run_server, daemon=True, args=(torch.distributed.get_rank(),)).start() - except ModuleNotFoundError: - log_rank_0("workload inspector module not found.") - - # Write args to tensorboard - write_args_to_tensorboard() - - # Turn on training mode which enables dropout. - for model_module in model: - model_module.train() - - # Tracking loss. - total_loss_dict = {} - - # Iterations. - iteration = args.iteration - # Make sure rerun_state_machine has the right iteration loaded from checkpoint. - rerun_state_machine = get_rerun_state_machine() - if rerun_state_machine.current_iteration != iteration: - log_rank_0(f"Setting rerun_state_machine.current_iteration to {iteration}...") - rerun_state_machine.current_iteration = iteration - - # Track E2E metrics at the start of training. - one_logger_utils.on_train_start( - iteration=iteration, - consumed_train_samples=args.consumed_train_samples, - train_samples=args.train_samples, - seq_length=args.seq_length, - train_iters=args.train_iters, - save=args.save, - async_save=args.async_save, - log_throughput=args.log_throughput, - num_floating_point_operations_so_far=args.num_floating_point_operations_so_far, - ) - - num_floating_point_operations_so_far = args.num_floating_point_operations_so_far - - # Setup some training config params. - config.grad_scale_func = optimizer.scale_loss - config.timers = timers - - if isinstance(model[0], (get_custom_fsdp(), DDP)) and args.overlap_grad_reduce: - assert config.no_sync_func is None, ( - "When overlap_grad_reduce is True, config.no_sync_func must be None; " - "a custom no_sync_func is not supported when overlapping grad-reduce" - ) - config.no_sync_func = [model_chunk.no_sync for model_chunk in model] - if len(model) == 1: - config.no_sync_func = config.no_sync_func[0] - if args.align_grad_reduce: - config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in model] - if len(model) == 1: - config.grad_sync_func = config.grad_sync_func[0] - if args.overlap_param_gather and args.align_param_gather: - config.param_sync_func = [model_chunk.start_param_sync for model_chunk in model] - if len(model) == 1: - config.param_sync_func = config.param_sync_func[0] - config.finalize_model_grads_func = finalize_model_grads - - timers("interval-time", log_level=0).start(barrier=True) - print_datetime("before the start of training step") - report_memory_flag = True - pre_hook_enabled = False - should_exit = False - exit_code = 0 - - if args.manual_gc: - # Disable the default garbage collector and perform the collection manually. - # This is to align the timing of garbage collection across ranks. - assert ( - args.manual_gc_interval >= 0 - ), "Manual garbage collection interval should be larger than or equal to 0" - gc.disable() - gc.collect() - - # Singleton initialization of straggler detector. - if args.log_straggler: - global stimer - world = torch.distributed.get_world_size() - rank = torch.distributed.get_rank() - mmcnt = args.straggler_minmax_count - stimer.configure( - world, - rank, - mmcnt=mmcnt, - enabled=not args.disable_straggler_on_startup, - port=args.straggler_ctrlr_port, - ) - num_floating_point_operations_since_last_log_event = 0.0 - - num_microbatches = get_num_microbatches() - eval_duration = 0.0 - eval_iterations = 0 - - def get_e2e_base_metrics(): - """Get base metrics values for one-logger to calculate E2E tracking metrics.""" - num_floating_point_operations_since_current_train_start = ( - num_floating_point_operations_so_far - args.num_floating_point_operations_so_far - ) - return { - "iteration": iteration, - "train_duration": timers("interval-time").active_time(), - "eval_duration": eval_duration, - "eval_iterations": eval_iterations, - "total_flops_since_current_train_start": num_floating_point_operations_since_current_train_start, - "num_floating_point_operations_so_far": num_floating_point_operations_so_far, - "consumed_train_samples": args.consumed_train_samples, - "world_size": args.world_size, - "seq_length": args.seq_length, - } - - # Cache into one-logger for callback. - if one_logger: - with one_logger.get_context_manager(): - one_logger.store_set("get_e2e_base_metrics", get_e2e_base_metrics) - - prof = None - if args.profile and torch.distributed.get_rank() in args.profile_ranks and args.use_pytorch_profiler: - activities = [torch.profiler.ProfilerActivity.CUDA] - if not args.disable_profiler_activity_cpu: - activities.append(torch.profiler.ProfilerActivity.CPU) - worker_name = ( - f"primus-megatron-exp[{self.exp_meta_info['exp_name']}]-rank[{torch.distributed.get_rank()}]" - ) - prof = torch.profiler.profile( - activities=activities, - schedule=torch.profiler.schedule( - wait=max(args.profile_step_start - 1, 0), - warmup=1 if args.profile_step_start > 0 else 0, - active=args.profile_step_end - args.profile_step_start, - repeat=1, - ), - on_trace_ready=torch.profiler.tensorboard_trace_handler( - args.tensorboard_dir, - worker_name=worker_name, - use_gzip=args.torch_profiler_use_gzip, - ), - record_shapes=args.torch_profiler_record_shapes, - with_stack=args.torch_profiler_with_stack, - ) - prof.start() - - start_iteration = iteration - # Disable forward pre-hook to start training to ensure that errors in checkpoint loading - # or random initialization don't propagate to all ranks in first all-gather (which is a - # no-op if things work correctly). - if should_disable_forward_pre_hook(args): - disable_forward_pre_hook(model, param_sync=False) - # Also remove param_sync_func temporarily so that sync calls made in - # `forward_backward_func` are no-ops. - param_sync_func = config.param_sync_func - config.param_sync_func = None - pre_hook_enabled = False - # Also, check weight hash across DP replicas to be very pedantic. - if args.check_weight_hash_across_dp_replicas_interval is not None: - assert check_param_hashes_across_dp_replicas( - model, cross_check=True - ), "Parameter hashes not matching across DP replicas" - torch.distributed.barrier() - log_rank_0(f">>> Weight hashes match after {iteration} iterations...") - - if args.dump_pp_data: - from .utils import set_dump_pp_data_patch - - set_dump_pp_data_patch() - log_rank_0(f"dump pp schedule data for visualization") - - # Run training iterations till done. - while iteration < args.train_iters: - if args.profile and torch.distributed.get_rank() in args.profile_ranks: - if args.use_pytorch_profiler: - prof.step() - elif iteration == args.profile_step_start: - torch.cuda.cudart().cudaProfilerStart() - torch.autograd.profiler.emit_nvtx(record_shapes=True).__enter__() - - ft_integration.on_checkpointing_start() - maybe_finalize_async_save(blocking=False) - ft_integration.on_checkpointing_end(is_async_finalization=True) - - # Update number of microbatches first without consistency check to decide if a - # checkpoint should be saved. If the number of microbatches is different - # from the previous iteration, save a checkpoint. Then run consistency check - # to make sure training configuration is still valid. - update_num_microbatches(args.consumed_train_samples, consistency_check=False, verbose=True) - if get_num_microbatches() != num_microbatches and iteration != 0: - assert get_num_microbatches() > num_microbatches, ( - f"Number of microbatches should be increasing due to batch size rampup; " - f"instead going from {num_microbatches} to {get_num_microbatches()}" - ) - if args.save is not None: - save_checkpoint_and_time( - iteration, - model, - optimizer, - opt_param_scheduler, - num_floating_point_operations_so_far, - checkpointing_context, - train_data_iterator=train_data_iterator, - ) - num_microbatches = get_num_microbatches() - update_num_microbatches(args.consumed_train_samples, consistency_check=True, verbose=True) - - # Completely skip iteration if needed. - if iteration in args.iterations_to_skip: - # Dummy train_step to fast forward train_data_iterator. - dummy_train_step(train_data_iterator) - iteration += 1 - batch_size = ( - mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches() - ) - args.consumed_train_samples += batch_size - args.skipped_train_samples += batch_size - continue - - # Run training step. - args.curr_iteration = iteration - ft_integration.on_training_step_start() - ( - loss_dict, - skipped_iter, - should_checkpoint, - should_exit, - exit_code, - grad_norm, - num_zeros_in_grad, - ) = self.train_step( - forward_step_func, - train_data_iterator, - model, - optimizer, - opt_param_scheduler, - config, - ) - ft_integration.on_training_step_end() - if should_checkpoint: - save_checkpoint_and_time( - iteration, - model, - optimizer, - opt_param_scheduler, - num_floating_point_operations_so_far, - checkpointing_context, - train_data_iterator=train_data_iterator, - ) - if should_exit: - break - - # Enable forward pre-hooks after first set of forward and backward passes. - # When running in fp16, skip all NaN iterations until steady-state loss scaling value - # is reached. - if iteration == start_iteration: - if skipped_iter: - # Only enable forward pre-hook after a training step has successfully run. Relevant - # for fp16 codepath where first XX iterations are skipped until steady-state loss - # scale value is reached. - start_iteration = iteration + 1 - else: - # Enable forward pre-hook after training step has successfully run. All subsequent - # forward passes will use the forward pre-hook / `param_sync_func` in - # `forward_backward_func`. - if should_disable_forward_pre_hook(args): - enable_forward_pre_hook(model) - config.param_sync_func = param_sync_func - pre_hook_enabled = True - - iteration += 1 - batch_size = mpu.get_data_parallel_world_size() * args.micro_batch_size * get_num_microbatches() - args.consumed_train_samples += batch_size - num_skipped_samples_in_batch = ( - get_current_global_batch_size() - get_current_running_global_batch_size() - ) - if args.decrease_batch_size_if_needed: - assert num_skipped_samples_in_batch >= 0 - else: - assert num_skipped_samples_in_batch == 0 - args.skipped_train_samples += num_skipped_samples_in_batch - flops_calc = ( - num_floating_point_operations - if not args.multi_latent_attention - else self.num_floating_point_operations_mla_moe - ) - num_floating_point_operations_in_batch = flops_calc(args, batch_size) - num_floating_point_operations_so_far += num_floating_point_operations_in_batch - num_floating_point_operations_since_last_log_event += num_floating_point_operations_in_batch - - # Logging. - if not optimizer.is_stub_optimizer: - loss_scale = optimizer.get_loss_scale().item() - else: - loss_scale = 1.0 - params_norm = None - - if args.log_params_norm: - params_norm = calc_params_l2_norm(model) - learning_rate = None - decoupled_learning_rate = None - for param_group in optimizer.param_groups: - if param_group["is_decoupled_lr"]: - decoupled_learning_rate = param_group["lr"] - else: - learning_rate = param_group["lr"] - report_memory_flag = self.training_log( - loss_dict, - total_loss_dict, - learning_rate, - decoupled_learning_rate, - iteration, - loss_scale, - report_memory_flag, - skipped_iter, - grad_norm, - params_norm, - num_zeros_in_grad, - ) - - # Evaluation. - if args.eval_interval and iteration % args.eval_interval == 0 and args.do_valid: - timers("interval-time").stop() - if should_disable_forward_pre_hook(args): - disable_forward_pre_hook(model) - pre_hook_enabled = False - if args.manual_gc and args.manual_gc_eval: - # Collect all objects. - gc.collect() - prefix = f"iteration {iteration}" - timers("eval-time", log_level=0).start(barrier=True) - evaluate_and_print_results( - prefix, - forward_step_func, - valid_data_iterator, - model, - iteration, - process_non_loss_data_func, - config, - verbose=False, - write_to_tensorboard=True, - non_loss_data_func=non_loss_data_func, - ) - eval_duration += timers("eval-time").elapsed() - eval_iterations += args.eval_iters - timers("eval-time").stop() - one_logger_utils.track_e2e_metrics() - - if args.manual_gc and args.manual_gc_eval: - # Collect only the objects created and used in evaluation. - gc.collect(generation=0) - if should_disable_forward_pre_hook(args): - enable_forward_pre_hook(model) - pre_hook_enabled = True - timers("interval-time", log_level=0).start(barrier=True) - - # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC). - # Some of these only happen at specific iterations. - post_training_step_callbacks( - model, - optimizer, - opt_param_scheduler, - iteration, - prof, - num_floating_point_operations_since_last_log_event, - ) - - # Checkpoint and decide whether to exit. - should_exit = checkpoint_and_decide_exit( - model, - optimizer, - opt_param_scheduler, - iteration, - num_floating_point_operations_so_far, - checkpointing_context, - train_data_iterator, - ) - if should_exit: - break - - one_logger_utils.track_e2e_metrics() - - if args.dump_pp_data: - from .utils import dump_pp_data - - pp_data_dir = os.environ.get("DUMP_PP_DIR", "output/pp_data") - dump_pp_data(args, get_num_microbatches(), pp_data_dir) - log_rank_0(f"pp schedule data dumped to {pp_data_dir}") - - # Flush TensorBoard, WandB writers and one-logger. - writer = get_tensorboard_writer() - if writer: - writer.flush() - - # Close out pre-hooks if using distributed optimizer and overlapped param gather. - if pre_hook_enabled: - disable_forward_pre_hook(model) - - ft_integration.on_checkpointing_start() - # This will finalize all unfinalized async request and terminate - # a persistent async worker if persistent ckpt worker is enabled - maybe_finalize_async_save(blocking=True, terminate=True) - ft_integration.on_checkpointing_end(is_async_finalization=True) - if args.enable_ft_package and ft_integration.get_rank_monitor_client() is not None: - ft_integration.get_rank_monitor_client().shutdown_workload_monitoring() - - # If any exit conditions (signal handler, duration, iterations) have been reached, exit. - if should_exit: - wandb_writer = get_wandb_writer() - if wandb_writer: - wandb_writer.finish() - mlflow_writer = get_mlflow_writer() - # Barrier to ensure all ranks have finished writing files before upload. - # Must run on ALL ranks to avoid deadlock (only last rank has mlflow_writer). - if dist.is_initialized(): - dist.barrier() - try: - upload_mlflow_artifacts( - tensorboard_dir=args.tensorboard_dir, - exp_root_path=self.exp_root_path, - upload_traces=getattr(args, "mlflow_upload_traces", False), - upload_logs=getattr(args, "mlflow_upload_logs", False), - generate_tracelens_report=getattr(args, "generate_tracelens_report", False), - upload_tracelens_report=getattr(args, "mlflow_upload_tracelens_report", False), - tracelens_ranks=getattr(args, "mlflow_tracelens_ranks", None), - tracelens_output_format=getattr(args, "mlflow_tracelens_output_format", "xlsx"), - tracelens_cleanup_after_upload=getattr( - args, "mlflow_tracelens_cleanup_after_upload", False - ), - tracelens_auto_install=getattr(args, "mlflow_tracelens_auto_install", True), - ) - except Exception as e: - import logging - - logging.getLogger(__name__).warning("[MLflow] Artifact upload failed: %s", e) - finally: - if mlflow_writer: - mlflow_writer.end_run() - if dist.is_initialized(): - dist.barrier() - ft_integration.shutdown() - sys.exit(exit_code) - - return iteration, num_floating_point_operations_so_far - - def train_step( - self, - forward_step_func, - data_iterator, - model, - optimizer, - opt_param_scheduler, - config, - no_optimizer_post_validation=False, - ): - """Single training step.""" - args = get_args() - timers = get_timers() - - def run_forward_backward_func(optimizer=None): - """Forward pass. - optimizer is not None for running post validation.""" - from megatron.core.pipeline_parallel import get_forward_backward_func - - forward_backward_func = get_forward_backward_func() - if optimizer is None and args.dump_pp_data: - forward_backward_func = schedule_wrapper(forward_backward_func) - kwargs = {} - if optimizer is not None: - kwargs["optimizer"] = optimizer - return forward_backward_func( - forward_step_func=forward_step_func, - data_iterator=data_iterator, - model=model, - num_microbatches=get_num_microbatches() * args.num_seq_splits, - seq_length=args.seq_length // args.num_seq_splits, - micro_batch_size=args.micro_batch_size, - decoder_seq_length=args.decoder_seq_length, - forward_only=False, - **kwargs, - ) - - rerun_state_machine = get_rerun_state_machine() - while rerun_state_machine.should_run_forward_backward(data_iterator): - # Set grad to zero. - for model_chunk in model: - model_chunk.zero_grad_buffer() - optimizer.zero_grad() - - # Forward pass. - losses_reduced = run_forward_backward_func() - - should_checkpoint, should_exit, exit_code = rerun_state_machine.should_checkpoint_and_exit() - if should_exit: - return {}, True, should_checkpoint, should_exit, exit_code, None, None - - # Empty unused memory. - if args.empty_unused_memory_level >= 1: - torch.cuda.empty_cache() - - # Vision gradients. - if args.vision_pretraining and args.vision_pretraining_type == "dino": - unwrapped_model = unwrap_model(model[0]) - unwrapped_model.cancel_gradients_last_layer(args.curr_iteration) - - # Update parameters. - - timers("optimizer", log_level=1).start(barrier=args.barrier_with_L1_time) - # update_successful, grad_norm, num_zeros_in_grad = optimizer.step() - if get_args().profile: - torch.cuda.nvtx.range_push("Optimizer") - if args.patch_zero_bubble and args.enable_optimizer_post_validation: - if optimizer.post_validation_enabled and not no_optimizer_post_validation: - optimizer.pre_step(args, timers) - if get_args().profile: - torch.cuda.nvtx.range_pop() - if get_args().profile: - torch.cuda.nvtx.range_push("post_validation_phase") - update_successful, grad_norm, num_zeros_in_grad = run_forward_backward_func(optimizer) - if get_args().profile: - torch.cuda.nvtx.range_pop() - # Here num_zeros_in_grad is a fake name, representing for optimizer_rollback - else: - update_successful, grad_norm, num_zeros_in_grad = optimizer.step() - if get_args().profile: - torch.cuda.nvtx.range_pop() - optimizer.record_grad_norm(grad_norm) - else: - update_successful, grad_norm, num_zeros_in_grad = optimizer.step() - if get_args().profile: - torch.cuda.nvtx.range_pop() - - timers("optimizer").stop() - - if getattr(args, "use_fsdp2_fp8_all_gather", False): - from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( - precompute_fp8_scales_for_fsdp, - ) - - precompute_fp8_scales_for_fsdp( - model[0], - stochastic_rounding=getattr(args, "fp8_all_gather_stochastic_rounding", False), - ) - - # FSDP2FP32Optimizer returns grad_norm as a GPU tensor to avoid a - # torch.compile graph break inside the compiled optimizer.step(). - # Materialize to float here, outside the compiled region. - if isinstance(grad_norm, torch.Tensor): - grad_norm = grad_norm.item() - - # when freezing sub-models we may have a mixture of successful and unsucessful ranks, - # so we must gather across mp ranks - update_successful = logical_and_across_model_parallel_group(update_successful) - # grad_norm and num_zeros_in_grad will be None on ranks without trainable params, - # so we must gather across mp ranks - grad_norm = reduce_max_stat_across_model_parallel_group(grad_norm) - if args.log_num_zeros_in_grad: - num_zeros_in_grad = reduce_max_stat_across_model_parallel_group(num_zeros_in_grad) - - # Vision momentum. - if args.vision_pretraining and args.vision_pretraining_type == "dino": - unwrapped_model = unwrap_model(model[0]) - unwrapped_model.update_momentum(args.curr_iteration) - - # Update learning rate. - if update_successful: - increment = get_num_microbatches() * args.micro_batch_size * args.data_parallel_size - opt_param_scheduler.step(increment=increment) - skipped_iter = 0 - else: - skipped_iter = 1 - - # Empty unused memory. - if args.empty_unused_memory_level >= 2: - torch.cuda.empty_cache() - - if is_pipeline_stage_containing_loss(): - # Average loss across microbatches. - loss_reduced = {} - for key in losses_reduced[0].keys(): - numerator = 0 - denominator = 0 - for x in losses_reduced: - val = x[key] - # there is one dict per microbatch. in new reporting, we average - # over the total number of tokens across the global batch. - if isinstance(val, tuple) or isinstance(val, list): - numerator += val[0] - denominator += val[1] - elif isinstance(val, torch.Tensor) and val.numel() == 2: - # Handle 2-element tensor [loss, num_tokens] format - # (upstream Megatron compatibility) - numerator += val[0] - denominator += val[1] - else: - # legacy behavior. we average over the number of microbatches, - # and so the denominator is 1. - numerator += val - denominator += 1 - loss_reduced[key] = numerator / denominator - return ( - loss_reduced, - skipped_iter, - should_checkpoint, - should_exit, - exit_code, - grad_norm, - num_zeros_in_grad, - ) - return ( - {}, - skipped_iter, - should_checkpoint, - should_exit, - exit_code, - grad_norm, - num_zeros_in_grad, - ) - - def training_log( - self, - loss_dict, - total_loss_dict, - learning_rate, - decoupled_learning_rate, - iteration, - loss_scale, - report_memory_flag, - skipped_iter, - grad_norm, - params_norm, - num_zeros_in_grad, - ): - """Log training information such as losses, timing, ....""" - args = get_args() - timers = get_timers() - writer = get_tensorboard_writer() - wandb_writer = get_wandb_writer() - mlflow_writer = get_mlflow_writer() - get_one_logger() - - # Advanced, skipped, and Nan iterations. - advanced_iters_key = "advanced iterations" - skipped_iters_key = "skipped iterations" - nan_iters_key = "nan iterations" - # Advanced iterations. - if not skipped_iter: - total_loss_dict[advanced_iters_key] = total_loss_dict.get(advanced_iters_key, 0) + 1 - else: - if advanced_iters_key not in total_loss_dict: - total_loss_dict[advanced_iters_key] = 0 - # Skipped iterations. - total_loss_dict[skipped_iters_key] = total_loss_dict.get(skipped_iters_key, 0) + skipped_iter - # Update losses and set nan iterations - got_nan = False - for key in loss_dict: - if not skipped_iter: - total_loss_dict[key] = ( - total_loss_dict.get(key, torch.tensor([0.0], dtype=torch.float, device="cuda")) - + loss_dict[key] - ) - else: - value = loss_dict[key].float().sum().item() - is_nan = value == float("inf") or value == -float("inf") or value != value - got_nan = got_nan or is_nan - total_loss_dict[nan_iters_key] = total_loss_dict.get(nan_iters_key, 0) + int(got_nan) - - # Logging. - timers_to_log = [ - "forward-backward", - "forward-compute", - "backward-compute", - "batch-generator", - "forward-recv", - "forward-send", - "backward-recv", - "backward-send", - "forward-send-forward-recv", - "forward-send-backward-recv", - "backward-send-forward-recv", - "backward-send-backward-recv", - "forward-backward-send-forward-backward-recv", - "layernorm-grads-all-reduce", - "embedding-grads-all-reduce", - "all-grads-sync", - "params-all-gather", - "optimizer-copy-to-main-grad", - "optimizer-unscale-and-check-inf", - "optimizer-clip-main-grad", - "optimizer-count-zeros", - "optimizer-inner-step", - "optimizer-copy-main-to-model-params", - "optimizer", - ] - - # Calculate batch size. - batch_size = args.micro_batch_size * args.data_parallel_size * get_num_microbatches() - - # Track app tag & app tag ID - one_logger_utils.track_app_tag(batch_size, args.world_size, args.seq_length) - - total_iterations = total_loss_dict[advanced_iters_key] + total_loss_dict[skipped_iters_key] - - # learning rate will be None on ranks without trainable params, so we must gather across mp ranks - learning_rate = reduce_max_stat_across_model_parallel_group(learning_rate) - # Tensorboard values. - # Timer requires all the ranks to call. - if args.log_timers_to_tensorboard and (iteration % args.tensorboard_log_interval == 0): - timers.write(timers_to_log, writer, iteration, normalizer=total_iterations) - if iteration % args.tensorboard_log_interval == 0: - if wandb_writer: - wandb_writer.log({"samples vs steps": args.consumed_train_samples}, iteration) - if mlflow_writer: - mlflow_writer.log_metric("samples vs steps", args.consumed_train_samples, step=iteration) - if writer: - writer.add_scalar("learning-rate", learning_rate, iteration) - if args.decoupled_lr is not None: - writer.add_scalar("decoupled-learning-rate", decoupled_learning_rate, iteration) - writer.add_scalar( - "learning-rate vs samples", - learning_rate, - args.consumed_train_samples, - ) - if wandb_writer: - wandb_writer.log({"learning-rate": learning_rate}, iteration) - if mlflow_writer: - mlflow_writer.log_metric("learning-rate", learning_rate, step=iteration) - if writer: - writer.add_scalar("batch-size", batch_size, iteration) - writer.add_scalar("batch-size vs samples", batch_size, args.consumed_train_samples) - if mlflow_writer: - mlflow_writer.log_metric("batch-size", batch_size, iteration) - if wandb_writer: - wandb_writer.log({"batch-size": batch_size}, iteration) - for key in loss_dict: - if writer: - writer.add_scalar(key, loss_dict[key], iteration) - writer.add_scalar(key + " vs samples", loss_dict[key], args.consumed_train_samples) - if wandb_writer: - wandb_writer.log({key: loss_dict[key]}, iteration) - if mlflow_writer: - mlflow_writer.log_metric(key, loss_dict[key], step=iteration) - if args.log_loss_scale_to_tensorboard: - if writer: - writer.add_scalar("loss-scale", loss_scale, iteration) - writer.add_scalar("loss-scale vs samples", loss_scale, args.consumed_train_samples) - if wandb_writer: - wandb_writer.log({"loss-scale": loss_scale}, iteration) - if mlflow_writer: - mlflow_writer.log_metric("loss-scale", loss_scale, step=iteration) - if args.log_world_size_to_tensorboard: - if writer: - writer.add_scalar("world-size", args.world_size, iteration) - writer.add_scalar( - "world-size vs samples", - args.world_size, - args.consumed_train_samples, - ) - if wandb_writer: - wandb_writer.log({"world-size": args.world_size}, iteration) - if mlflow_writer: - mlflow_writer.log_metric("world-size", args.world_size, step=iteration) - if grad_norm is not None: - if writer: - writer.add_scalar("grad-norm", grad_norm, iteration) - writer.add_scalar("grad-norm vs samples", grad_norm, args.consumed_train_samples) - if wandb_writer: - wandb_writer.log({"grad-norm": grad_norm}, iteration) - if mlflow_writer: - mlflow_writer.log_metric("grad-norm", grad_norm, step=iteration) - if num_zeros_in_grad is not None: - if writer: - writer.add_scalar("num-zeros", num_zeros_in_grad, iteration) - writer.add_scalar( - "num-zeros vs samples", - num_zeros_in_grad, - args.consumed_train_samples, - ) - if wandb_writer: - wandb_writer.log({"num-zeros": num_zeros_in_grad}, iteration) - if mlflow_writer: - mlflow_writer.log_metric("num-zeros", num_zeros_in_grad, iteration) - if params_norm is not None: - if writer: - writer.add_scalar("params-norm", params_norm, iteration) - writer.add_scalar( - "params-norm vs samples", - params_norm, - args.consumed_train_samples, - ) - if wandb_writer: - wandb_writer.log({"params-norm": params_norm}, iteration) - if mlflow_writer: - mlflow_writer.log_metric("params-norm", params_norm, iteration) - if args.log_memory_to_tensorboard: - mem_stats = torch.cuda.memory_stats() - if writer: - writer.add_scalar( - "mem-reserved-bytes", - mem_stats["reserved_bytes.all.current"], - iteration, - ) - writer.add_scalar( - "mem-allocated-bytes", - mem_stats["allocated_bytes.all.current"], - iteration, - ) - writer.add_scalar( - "mem-max-allocated-bytes", - mem_stats["allocated_bytes.all.peak"], - iteration, - ) - writer.add_scalar( - "mem-allocated-count", - mem_stats["allocation.all.current"], - iteration, - ) - if wandb_writer: - wandb_writer.log( - {"mem-reserved-bytes": mem_stats["reserved_bytes.all.current"]}, - iteration, - ) - wandb_writer.log( - {"mem-allocated-bytes": mem_stats["allocated_bytes.all.current"]}, - iteration, - ) - wandb_writer.log( - {"mem-max-allocated-bytes": mem_stats["allocated_bytes.all.peak"]}, - iteration, - ) - wandb_writer.log( - {"mem-allocated-count": mem_stats["allocation.all.current"]}, - iteration, - ) - if args.num_experts is not None: - moe_loss_scale = 1 / get_num_microbatches() - track_moe_metrics( - loss_scale=moe_loss_scale, - iteration=iteration, - writer=writer, - wandb_writer=wandb_writer, - mlflow_writer=mlflow_writer, - total_loss_dict=total_loss_dict, - per_layer_logging=args.moe_per_layer_logging, - moe_layer_freq=args.moe_layer_freq, - num_layers=args.num_layers, - ) - - if iteration % args.log_interval == 0: - # Note(wenx): If we want to collect rocm-smi memory information for the first two iterations, - # place the collection before the timer to minimize its impact on latency measurements for iterations ≥ 3. - rocm_gpu_util = None - # Enable throughput calculations if log_throughput or mlflow_upload_performance_metrics is set - enable_perf_metrics = args.log_throughput or getattr( - args, "mlflow_upload_performance_metrics", False - ) - if enable_perf_metrics: - if args.use_rocm_mem_info or ( - args.use_rocm_mem_info_iters is not None and iteration in args.use_rocm_mem_info_iters - ): - rocm_total_mem, rocm_used_mem, rocm_free_mem = get_rocm_smi_mem_info( - self.module_local_rank - ) - # Collect GPU utilization for performance metrics - # Every rank samples its own GPU util (rate-limited to avoid rocm-smi - # subprocess overhead) so the all_gather below gets real per-rank values. - if getattr(args, "mlflow_upload_performance_metrics", False): - import time - - now = time.time() - should_sample = ( - self._last_rocm_gpu_util is None - or (now - self._last_rocm_gpu_util_time) >= self._rocm_gpu_util_sample_interval - ) - if should_sample: - try: - self._last_rocm_gpu_util = get_rocm_smi_gpu_util(self.module_local_rank) - self._last_rocm_gpu_util_time = now - except Exception: - pass # Keep previous cached value if any - rocm_gpu_util = self._last_rocm_gpu_util - else: - rocm_gpu_util = None - - elapsed_time = timers("interval-time").elapsed(barrier=True) - elapsed_time_per_iteration = elapsed_time / total_iterations - - flops_calc = ( - num_floating_point_operations - if not args.multi_latent_attention - else self.num_floating_point_operations_mla_moe - ) - throughput = flops_calc(args, batch_size) / ( - elapsed_time_per_iteration * 10**12 * args.world_size - ) - - if args.log_timers_to_tensorboard: - if writer: - writer.add_scalar("iteration-time", elapsed_time_per_iteration, iteration) - if wandb_writer: - wandb_writer.log({"iteration-time": elapsed_time_per_iteration}, iteration) - if mlflow_writer: - mlflow_writer.log_metric("iteration-time", elapsed_time_per_iteration, iteration) - # log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]" - log_string = f"" - if hasattr(self, "episode_count") and self.episode_count is not None: - log_string += f" episode {self.episode_count} |" - log_string += " iteration {:8d}/{:8d} |".format(iteration, args.train_iters) - log_string += " consumed samples: {:12d} |".format(args.consumed_train_samples) - if ( - iteration == self.log_avg_skip_iterations + 1 - or len(self.recent_iteration_times) >= self.log_avg_reset_interval - ): - self.recent_iteration_times.clear() - self.recent_iteration_times.append(elapsed_time_per_iteration * 1000.0) - log_string += " elapsed time per iteration (ms): {:.1f}/{:.1f} |".format( - elapsed_time_per_iteration * 1000.0, - statistics.mean(self.recent_iteration_times), - ) - if enable_perf_metrics: - if ( - iteration == self.log_avg_skip_iterations + 1 - or len(self.recent_tflop_throughputs) >= self.log_avg_reset_interval - ): - self.recent_tflop_throughputs.clear() - self.recent_tflop_throughputs.append(throughput) - - if not args.use_rocm_mem_info: - hip_free_mem, hip_total_mem = torch.cuda.mem_get_info() - hip_used_mem = hip_total_mem - hip_free_mem - hip_mem_usage = hip_used_mem / hip_total_mem - log_string += ( - f" hip mem usage/free/total/usage_ratio: {hip_used_mem/1024/1024/1024:.2f}GiB/" - ) - log_string += f"{hip_free_mem/1024/1024/1024:.2f}GiB/" - log_string += f"{hip_total_mem/1024/1024/1024:.2f}GiB/{hip_mem_usage*100:.2f}% |" - - if args.use_rocm_mem_info or ( - args.use_rocm_mem_info_iters is not None and iteration in args.use_rocm_mem_info_iters - ): - rocm_mem_usage = rocm_used_mem / rocm_total_mem - - # get the max rocm_mem_usage - usage_tensor = torch.tensor([rocm_mem_usage], device="cuda", dtype=torch.float32) - world_size = dist.get_world_size() - gathered_usage = [torch.zeros_like(usage_tensor) for _ in range(world_size)] - dist.all_gather(gathered_usage, usage_tensor) - - rocm_mem_usages = [t.item() for t in gathered_usage] - max_usage = max(rocm_mem_usages) - max_rank = rocm_mem_usages.index(max_usage) - - log_string += ( - f" rocm mem usage/free/total/usage_ratio: {rocm_used_mem/1024/1024/1024:.2f}GiB/" - ) - log_string += f"{rocm_free_mem/1024/1024/1024:.2f}GiB/" - log_string += f"{rocm_total_mem/1024/1024/1024:.2f}GiB/{rocm_mem_usage*100:.2f}% |" - log_string += f" rank-{max_rank} max mem usage/usage_ratio: " - log_string += f"{rocm_total_mem*max_usage/1024/1024/1024:.2f}GiB/{max_usage*100:.2f}% |" - - log_string += ( - f" throughput per GPU (TFLOP/s/GPU): {throughput:.1f}/" - f"{statistics.mean(self.recent_tflop_throughputs):.1f} |" - ) - token_throughput = args.seq_length * batch_size / elapsed_time_per_iteration / args.world_size - if ( - iteration == self.log_avg_skip_iterations + 1 - or len(self.recent_token_throughputs) >= self.log_avg_reset_interval - ): - self.recent_token_throughputs.clear() - self.recent_token_throughputs.append(token_throughput) - log_string += ( - f" tokens per GPU (tokens/s/GPU): {token_throughput:.1f}/" - f"{statistics.mean(self.recent_token_throughputs):.1f} |" - ) - if args.log_timers_to_tensorboard: - if args.use_rocm_mem_info or ( - args.use_rocm_mem_info_iters is not None and iteration in args.use_rocm_mem_info_iters - ): - mem_collector = "rocm" - used_mem, free_mem, total_mem, mem_usage = ( - rocm_used_mem, - rocm_free_mem, - rocm_total_mem, - rocm_mem_usage, - ) - else: - mem_collector = "hip" - used_mem, free_mem, total_mem, mem_usage = ( - hip_used_mem, - hip_free_mem, - hip_total_mem, - hip_mem_usage, - ) - if writer: - writer.add_scalar("throughput(tflops/sec/gpu)", throughput, iteration) - writer.add_scalar( - "token_throughput(tokens/sec/gpu)", - token_throughput, - iteration, - ) - writer.add_scalar( - f"{mem_collector}_used_mem(GiB)", - used_mem / 1024 / 1024 / 1024, - iteration, - ) - writer.add_scalar( - f"{mem_collector}_free_mem(GiB)", - free_mem / 1024 / 1024 / 1024, - iteration, - ) - writer.add_scalar( - f"{mem_collector}_total_mem(GiB)", - total_mem / 1024 / 1024 / 1024, - iteration, - ) - writer.add_scalar(f"{mem_collector}_mem_usage(%)", mem_usage * 100.0, iteration) - if wandb_writer: - wandb_writer.log({"throughput(tflops/sec/gpu)": throughput}, iteration) - wandb_writer.log( - {"token_throughput(tokens/sec/gpu)": token_throughput}, - iteration, - ) - wandb_writer.log( - {f"{mem_collector}_used_mem(GiB)": used_mem / 1024 / 1024 / 1024}, - iteration, - ) - wandb_writer.log( - {f"{mem_collector}_free_mem(GiB)": free_mem / 1024 / 1024 / 1024}, - iteration, - ) - wandb_writer.log( - {f"{mem_collector}_total_mem(GiB)": total_mem / 1024 / 1024 / 1024}, - iteration, - ) - wandb_writer.log({f"{mem_collector}_mem_usage(%)": mem_usage * 100.0}, iteration) - if mlflow_writer: - mlflow_writer.log_metric("throughput_tflops_per_sec_per_gpu", throughput, iteration) - mlflow_writer.log_metric( - "token_throughput_tokens_per_sec_per_gpu", - token_throughput, - iteration, - ) - mlflow_writer.log_metric( - f"{mem_collector}_used_mem_GiB", - used_mem / 1024 / 1024 / 1024, - iteration, - ) - mlflow_writer.log_metric( - f"{mem_collector}_free_mem_GiB", - free_mem / 1024 / 1024 / 1024, - iteration, - ) - mlflow_writer.log_metric( - f"{mem_collector}_total_mem_GiB", - total_mem / 1024 / 1024 / 1024, - iteration, - ) - mlflow_writer.log_metric( - f"{mem_collector}_mem_usage_percent", mem_usage * 100.0, iteration - ) - - # Upload performance metrics to MLflow - # Groups: Performance (throughput, TPS, iteration time), Memory (peak, usage %), System (GPU util) - # NOTE: mlflow_writer only exists on last rank, but all_gather requires all ranks to participate - if getattr(args, "mlflow_upload_performance_metrics", False): - # Ensure memory metrics are available when log_timers_to_tensorboard is False - # (mem_collector, used_mem, mem_usage are otherwise only set inside log_timers_to_tensorboard) - if not args.log_timers_to_tensorboard: - hip_free_mem, hip_total_mem = torch.cuda.mem_get_info() - used_mem = hip_total_mem - hip_free_mem - mem_usage = used_mem / hip_total_mem - mem_collector = "hip" - if args.use_rocm_mem_info or ( - args.use_rocm_mem_info_iters is not None - and iteration in args.use_rocm_mem_info_iters - ): - mem_collector = "rocm" - used_mem = rocm_used_mem - mem_usage = rocm_mem_usage - # System metrics - GPU utilization per rank - # ALL ranks must participate in all_gather, even if they don't have mlflow_writer - # Use -1 as sentinel for unavailable GPU util - util_value = rocm_gpu_util if rocm_gpu_util is not None else -1.0 - util_tensor = torch.tensor([util_value], device="cuda", dtype=torch.float32) - world_size = dist.get_world_size() - gathered_utils = [torch.zeros_like(util_tensor) for _ in range(world_size)] - dist.all_gather(gathered_utils, util_tensor) - - # Only the last rank (which has mlflow_writer) logs the metrics - if mlflow_writer: - # Performance metrics - mlflow_writer.log_metric("perf/throughput_tflops_per_gpu", throughput, iteration) - mlflow_writer.log_metric( - "perf/tps_tokens_per_sec_per_gpu", token_throughput, iteration - ) - mlflow_writer.log_metric( - "perf/iteration_time_ms", - elapsed_time_per_iteration * 1000.0, - iteration, - ) - # Memory metrics - mlflow_writer.log_metric( - f"perf/{mem_collector}_current_mem_gb", - used_mem / 1024 / 1024 / 1024, - iteration, - ) - mlflow_writer.log_metric( - f"perf/{mem_collector}_mem_utilization_pct", - mem_usage * 100.0, - iteration, - ) - # Log GPU utilization from gathered values - valid_utils = [] - for rank, util_val in enumerate(gathered_utils): - util = util_val.item() - if util >= 0: # Filter out sentinel values (-1) - mlflow_writer.log_metric( - f"perf/gpu_utilization_pct_rank{rank}", - util, - iteration, - ) - valid_utils.append(util) - # Also log average GPU utilization (only from valid values) - if valid_utils: - avg_util = sum(valid_utils) / len(valid_utils) - mlflow_writer.log_metric( - "perf/gpu_utilization_pct_avg", - avg_util, - iteration, - ) - - assert learning_rate is not None - # Decoupled_learning_rate should be not None only on first and last pipeline stage. - log_string += " learning rate: {:.6E} |".format(learning_rate) - if args.decoupled_lr is not None and ( - mpu.is_pipeline_first_stage(ignore_virtual=True) - or mpu.is_pipeline_last_stage(ignore_virtual=True) - ): - assert decoupled_learning_rate is not None - log_string += " decoupled learning rate: {:.6E} |".format(decoupled_learning_rate) - else: - assert decoupled_learning_rate is None - log_string += " global batch size: {:5d} |".format(batch_size) - for key in total_loss_dict: - if key not in [advanced_iters_key, skipped_iters_key, nan_iters_key]: - avg = total_loss_dict[key].item() / float(max(1, total_loss_dict[advanced_iters_key])) - if avg > 0.0: - log_string += " {}: {:.6E} |".format(key, avg) - total_loss_dict[key] = torch.tensor([0.0], dtype=torch.float, device="cuda") - log_string += " loss scale: {:.1f} |".format(loss_scale) - if grad_norm is not None: - log_string += " grad norm: {:.3f} |".format(grad_norm) - if num_zeros_in_grad is not None: - log_string += " num zeros: {:.1f} |".format(num_zeros_in_grad) - if params_norm is not None: - log_string += " params norm: {:.3f} |".format(params_norm) - log_string += " number of skipped iterations: {:3d} |".format(total_loss_dict[skipped_iters_key]) - log_string += " number of nan iterations: {:3d} |".format(total_loss_dict[nan_iters_key]) - total_loss_dict[advanced_iters_key] = 0 - total_loss_dict[skipped_iters_key] = 0 - total_loss_dict[nan_iters_key] = 0 - - if self.is_v_schedule: - log_rank_0(log_string) - else: - log_rank_last(log_string) - if report_memory_flag and learning_rate > 0.0: - # Report memory after optimizer state has been initialized. - if torch.distributed.get_rank() == 0: - num_microbatches = get_num_microbatches() - report_theoretical_memory(args, num_microbatches=num_microbatches, verbose=True) - report_memory("(after {} iterations)".format(iteration)) - report_memory_flag = False - - # Removed to avoid global sync in zero bubble schedules. - if not (get_args().zero_bubble_v_schedule or get_args().patch_zero_bubble): - timers.log(timers_to_log, normalizer=args.log_interval) - - return report_memory_flag - - def num_floating_point_operations_mla_moe(self, args, batch_size): - # MoE. - gated_linear_multiplier = 3 / 2 if args.swiglu else 1 - num_experts_routed_to = 1 if args.num_experts is None else args.moe_router_topk - ffn_hidden_size = ( - args.moe_ffn_hidden_size if args.moe_ffn_hidden_size is not None else args.ffn_hidden_size - ) - shared_expert_ffn_hidden_size = ( - 0 - if args.moe_shared_expert_intermediate_size is None - else args.moe_shared_expert_intermediate_size - ) - - # The 12x term below comes from the following factors; for more details, see - # "APPENDIX: FLOATING-POINT OPERATIONS" in https://arxiv.org/abs/2104.04473. - # - 3x: Each GEMM in the model needs to be performed 3 times (forward pass, - # backward wgrad [weight gradient], backward dgrad [data gradient]). - # - 2x: GEMMs of a particular size are stacked twice in the standard Transformer model - # architectures implemented in this codebase (e.g., h->ffn_h GEMM and ffn_h->h GEMM - # in MLP layer). - # - 2x: A GEMM of a m*n tensor with a n*k tensor requires 2mnk floating-point operations. - expansion_factor_fw_bw = 3 - expansion_factor_gemm_stack = 2 - expansion_factor_gemm_ops = 2 - - return ( - expansion_factor_fw_bw - * expansion_factor_gemm_ops - * batch_size - * args.seq_length - * args.num_layers - * args.hidden_size - * args.hidden_size - * ( - # Attention - q. - ( - ( - # Attention - q_proj. - args.num_attention_heads - * (args.qk_head_dim + args.qk_pos_emb_head_dim) - / args.hidden_size - ) - if not args.q_lora_rank - else ( - # Attention - q_down_proj. - (args.q_lora_rank / args.hidden_size) - # Attention - q_up_proj. - + ( - (args.q_lora_rank / args.hidden_size) - * ( - args.num_attention_heads - * (args.qk_head_dim + args.qk_pos_emb_head_dim) - / args.hidden_size - ) - ) - ) - ) - # Attention - kv_down_proj. - + ((args.kv_lora_rank + args.qk_pos_emb_head_dim) / args.hidden_size) - # Attention - kv_up_proj. - + ( - (args.kv_lora_rank / args.hidden_size) - * (args.num_attention_heads * (args.qk_head_dim + args.v_head_dim) / args.hidden_size) - ) - # Attention - Q*K. - + ( - (args.seq_length / args.hidden_size) - * ( - args.num_attention_heads - * (args.qk_head_dim + args.qk_pos_emb_head_dim) - / args.hidden_size - ) - ) - # Attention - A(ttention Score)*V. - + ( - (args.seq_length / args.hidden_size) - * (args.num_attention_heads * args.v_head_dim / args.hidden_size) - ) - # Attention - output proj. - + (args.num_attention_heads * args.v_head_dim / args.hidden_size) - # MLP. - + ( - (ffn_hidden_size / args.hidden_size) - * num_experts_routed_to - * gated_linear_multiplier - * expansion_factor_gemm_stack - ) - # Shared Experts. - + ( - (shared_expert_ffn_hidden_size / args.hidden_size) - * gated_linear_multiplier - * expansion_factor_gemm_stack - ) - # Logit. (untie_embeddings_and_output_weights=True) - + (args.padded_vocab_size / (2 * args.num_layers * args.hidden_size)) - * expansion_factor_gemm_stack - ) - ) diff --git a/primus/modules/trainer/megatron/utils.py b/primus/modules/trainer/megatron/utils.py deleted file mode 100644 index d97842547..000000000 --- a/primus/modules/trainer/megatron/utils.py +++ /dev/null @@ -1,611 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -"""megatron utils""" - -import inspect -import json -import os - -import megatron -import torch -from megatron.core import parallel_state -from megatron.training.global_vars import get_args - -from primus.core.utils import logger - -_GLOBAL_PP_VIS_EVENTS = [] -_GLOBAL_PP_VIS_EVENTS_PER_ITER = None - - -######################################################log after torch distributed initialized - - -def is_v_schedule_enabled(args=None): - if args is None: - args = get_args() - return ( - args.patch_zero_bubble - and args.enable_zero_bubble - and (args.zero_bubble_v_schedule or args.enable_1f1b_v) - ) or (args.pp_algorithm in ("zbv-formatted", "v-half", "v-min") and args.patch_primus_pipeline) - - -def is_last_rank(): - return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1) - - -def print_rank_last(msg): - """If distributed is initialized, print only on last rank.""" - log_func = logger.info_with_caller - - caller = inspect.stack()[1] - caller_frame = caller.frame - function_name = caller_frame.f_code.co_name - module_name = caller_frame.f_globals["__name__"].split(".")[-1] - line = caller.lineno - - if torch.distributed.is_initialized(): - if is_last_rank(): - log_func(msg, module_name, function_name, line) - else: - log_func(msg, module_name, function_name, line) - - -def set_wandb_writer_patch(args): # monkey patch - """ - This function is adapted from the original Megatron implementation, with an additional - wandb argument `entity` be added. - Monkey-patch note: - - The original function will be replaced at runtime by this implementation. - - """ - - megatron.training.global_vars._ensure_var_is_not_initialized( - megatron.training.global_vars._GLOBAL_WANDB_WRITER, "wandb writer" - ) - - if getattr(args, "wandb_project", "") and args.rank == (args.world_size - 1): - if args.wandb_exp_name == "": - raise ValueError("Please specify the wandb experiment name!") - - import wandb - - if args.wandb_save_dir: - save_dir = args.wandb_save_dir - else: - # Defaults to the save dir. - save_dir = os.path.join(args.save, "wandb") - wandb_kwargs = { - "dir": save_dir, - "name": args.wandb_exp_name, - "project": args.wandb_project, - "entity": args.wandb_entity, - "config": vars(args), - } - os.makedirs(wandb_kwargs["dir"], exist_ok=True) - wandb.init(**wandb_kwargs) - megatron.training.global_vars._GLOBAL_WANDB_WRITER = wandb - - -def validate_specified_recompute_layers(args): - if args.recompute_layer_ids is None: - return - - assert isinstance( - args.recompute_layer_ids, list - ), f"recompute_layer_ids={args.recompute_layer_ids} should be a list" - recompute_layer_ids = list(set(args.recompute_layer_ids)) - assert len(recompute_layer_ids) > 0, "recompute layer ids is null" - for layer_id in recompute_layer_ids: - assert ( - layer_id >= 0 and layer_id < args.num_layers - ), f"recompute layer id must be between 0 and {args.num_layers - 1}" - - if args.recompute_granularity != "full": - raise ValueError( - f'When using recompute_layer_ids, recompute_granuarlity: {args.recompute_granularity} must be "full"' - ) - - if args.recompute_method is not None: - raise ValueError( - f"When using recompute_layer_ids, recompute_method: {args.recompute_method} must be None." - ) - - if args.distribute_saved_activations and args.sequence_parallel: - raise ValueError( - f"distribute_saved_activations: {args.distribute_saved_activations} must be " - f"false when sequence parallel is enabled: {args.sequence_parallel}" - ) - - -def validate_manual_split(args): - """ - The use of decoder_pipeline_manual_split_list is to relax the divisibility - restriction of the current (interleaved) 1f1b pipeline schedule. The layer - split or number of each pp rank is - decoder_pipeline_manual_split_list[pp_rank*vp_size:(pp_rank+1)*vp_size] or - decoder_pipeline_manual_split_list[pp_rank] when interleaved pipeline is - used or not. For example, the split list could be "[2,3,2,2,2,2,2,1]" - in layer16-pp4-vpp2 config, where the vpp split of - pp_rank0/pp_rank1/pp_rank2/pp_rank3 is [2,3]/[2,2]/[2,2]/[2,1]. - - if chosen pipeline is v_schedule like zbv/v-half, - the split list will be the actual layer sequence. - For example, layer16-pp4-vpp2 config, the vpp split of - pp_rank0/pp_rank1/pp_rank2/pp_rank3 is [3,2,2,2,2,2,2,1] - indicate the pipeline as follows: - pp_rank0: 3 1 - pp_rank1: 2 2 - pp_rank2: 2 2 - pp_rank3: 2 2 - - """ - - if ( - args.num_layers_per_virtual_pipeline_stage is not None - or args.decoder_first_pipeline_num_layers is not None - or args.decoder_last_pipeline_num_layers is not None - or args.account_for_embedding_in_pipeline_split - or args.account_for_loss_in_pipeline_split - ): - raise ValueError( - "decoder_pipeline_manual_split_list is not compatible " - "with num_layers_per_virtual_pipeline_stage/" - "decoder_first_pipeline_num_layers/" - "decoder_last_pipeline_num_layers/" - "account_for_embedding_in_pipeline_split/" - "account_for_loss_in_pipeline_split yet" - ) - - num_layers = args.num_layers - pp_size = args.pipeline_model_parallel_size - vp_size = args.virtual_pipeline_model_parallel_size - pp_split = args.decoder_pipeline_manual_split_list - - if pp_size <= 1: - raise ValueError( - f"pipeline_model_parallel_size={pp_size} should be larger " - f"than 1 when decoder_pipeline_manual_split_list is used" - ) - - if not isinstance(pp_split, list): - raise ValueError(f"decoder_pipeline_manual_split_list={pp_split} should be a list") - - split_size = pp_size if vp_size is None else pp_size * vp_size - if len(pp_split) != split_size: - raise ValueError( - f"the size of decoder_pipeline_manual_split_list=" - f"{pp_split} should be {split_size} " - f"given pipeline_model_parallel_size={pp_size} and " - f"virtual_pipeline_model_parallel_size={vp_size}" - ) - - if not all(x > 0 for x in pp_split): - raise ValueError( - f"layer numbers in decoder_pipeline_manual_split_list={pp_split} should all be larger than 0" - ) - - if sum(pp_split) != num_layers: - raise ValueError( - f"the sum of decoder_pipeline_manual_split_list=" - f"{pp_split} is {sum(pp_split)} and " - f"should be equal to num_layers={num_layers}" - ) - - return True - - -def validate_args_modified(*args, **kwargs): - def validate_args_modifier(func, modification): - import inspect - - source = inspect.getsource(func) - modified_source = modification(source) - namespace = {} - exec(modified_source, func.__globals__, namespace) - return namespace[func.__name__] - - ori_code = kwargs.pop("ori_code", None) - new_code = kwargs.pop("new_code", None) - - assert ori_code is not None and new_code is not None, "ori_code and new_code must be provided." - - megatron.training.arguments.validate_args = validate_args_modifier( - megatron.training.arguments.validate_args, lambda s: s.replace(ori_code, new_code) - ) - megatron.training.arguments.validate_args(*args, **kwargs) - - -def set_manual_pipeline_split_patch(args): - """ - Monkey-patch note: - - The original function will be replaced at runtime by this implementation. - - """ - - megatron.core.transformer.TransformerConfig.decoder_pipeline_manual_split_list = ( - args.decoder_pipeline_manual_split_list - ) - - # patch get_num_layers_to_build - def get_num_layers_to_build_patch(config, vp_stage, pp_rank=None): - if pp_rank is None: - pp_rank = parallel_state.get_pipeline_model_parallel_rank() - vp_size = config.virtual_pipeline_model_parallel_size - - if not is_v_schedule_enabled(): - pp_idx = pp_rank if vp_size is None else pp_rank * vp_size + vp_stage - num_layers_to_build = config.decoder_pipeline_manual_split_list[pp_idx] - return num_layers_to_build - else: - assert vp_stage is not None and vp_stage in (0, 1) - pp_size = config.pipeline_model_parallel_size - chunk_id = pp_rank if vp_stage == 0 else 2 * pp_size - pp_rank - 1 - num_layers_to_build = config.decoder_pipeline_manual_split_list[chunk_id] - return num_layers_to_build - - megatron.core.transformer.transformer_block.get_num_layers_to_build = get_num_layers_to_build_patch - megatron.core.models.gpt.gpt_layer_specs.get_num_layers_to_build = get_num_layers_to_build_patch - - # patch get_transformer_layer_offset - def get_transformer_layer_offset_patch(config, vp_stage, pp_rank=None): - if pp_rank is None: - pp_rank = parallel_state.get_pipeline_model_parallel_rank() - pp_size = config.pipeline_model_parallel_size - vp_size = config.virtual_pipeline_model_parallel_size - - offset = 0 - - if not is_v_schedule_enabled(): - if vp_stage is not None: - for vp_idx in range(vp_stage): - for pp_idx in range(pp_size): - offset += config.decoder_pipeline_manual_split_list[pp_idx * vp_size + vp_idx] - for pp_idx in range(pp_rank): - offset += config.decoder_pipeline_manual_split_list[pp_idx * vp_size + vp_stage] - else: - offset = sum(config.decoder_pipeline_manual_split_list[:pp_rank]) - else: - assert vp_stage is not None and vp_stage in (0, 1) - chunk_id = pp_rank if vp_stage == 0 else 2 * pp_size - pp_rank - 1 - offset = sum(config.decoder_pipeline_manual_split_list[:chunk_id]) - return offset - - megatron.core.transformer.transformer_layer.get_transformer_layer_offset = ( - get_transformer_layer_offset_patch - ) - megatron.core.transformer.transformer_block.get_transformer_layer_offset = ( - get_transformer_layer_offset_patch - ) - megatron.core.models.gpt.gpt_layer_specs.get_transformer_layer_offset = get_transformer_layer_offset_patch - - -def schedule_wrapper(func): - def wrapper(*args, **kwargs): - global _GLOBAL_PP_VIS_EVENTS_PER_ITER - _GLOBAL_PP_VIS_EVENTS_PER_ITER = { - "start": None, - "end": None, - "memory": None, - "fwd_start": [], - "fwd_end": [], - "fwd_minibatch": [], - "fwd_chunk": [], - "bwd_start": [], - "bwd_end": [], - "bwd_minibatch": [], - "bwd_chunk": [], - "wgrad_start": [], - "wgrad_end": [], - "wgrad_minibatch": [], - "wgrad_chunk": [], - } - - _GLOBAL_PP_VIS_EVENTS_PER_ITER["start"] = torch.cuda.Event(enable_timing=True) - _GLOBAL_PP_VIS_EVENTS_PER_ITER["start"].record() - res = func(*args, **kwargs) - _GLOBAL_PP_VIS_EVENTS_PER_ITER["end"] = torch.cuda.Event(enable_timing=True) - _GLOBAL_PP_VIS_EVENTS_PER_ITER["end"].record() - - _GLOBAL_PP_VIS_EVENTS_PER_ITER["memory"] = torch.cuda.max_memory_reserved() / 1024**3 - - global _GLOBAL_PP_VIS_EVENTS - _GLOBAL_PP_VIS_EVENTS.append(_GLOBAL_PP_VIS_EVENTS_PER_ITER) - - return res - - return wrapper - - -def fwd_bwd_wrapper(func, mode, minibatch=None, chunk=None): - def wrapper(*args, **kwargs): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - - start.record() - res = func(*args, **kwargs) - end.record() - - global _GLOBAL_PP_VIS_EVENTS_PER_ITER - _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_start"].append(start) - _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_end"].append(end) - - if minibatch is not None: - _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_minibatch"].append(minibatch) - if chunk is not None: - _GLOBAL_PP_VIS_EVENTS_PER_ITER[mode + "_chunk"].append(chunk) - return res - - return wrapper - - -def combined_fwd_bwd_wrapper(func, fwd_minibatch, fwd_chunk, bwd_minibatch, bwd_chunk): - """Record a single combined forward+backward call as both an ``fwd`` event - and a ``bwd`` event sharing the same ``[start, end]`` interval. - - Used by ``megatron_combined_fwd_bkwd_handler`` so that nodes collapsed into - a combined FB group still appear in the dump_pp_data output. Without this - the visualizer's per-rank F/B/W totals are heavily under-counted on ranks - that hit the steady state (the F and B halves are interleaved inside - ``combined_forward_backward_step`` and cannot be timed separately). - """ - - def wrapper(*args, **kwargs): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - - start.record() - res = func(*args, **kwargs) - end.record() - - global _GLOBAL_PP_VIS_EVENTS_PER_ITER - _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_start"].append(start) - _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_end"].append(end) - _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_minibatch"].append(fwd_minibatch) - _GLOBAL_PP_VIS_EVENTS_PER_ITER["fwd_chunk"].append(fwd_chunk) - _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_start"].append(start) - _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_end"].append(end) - _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_minibatch"].append(bwd_minibatch) - _GLOBAL_PP_VIS_EVENTS_PER_ITER["bwd_chunk"].append(bwd_chunk) - return res - - return wrapper - - -def set_dump_pp_data_patch(): - from megatron.core.pipeline_parallel import schedules - - schedules.forward_step = fwd_bwd_wrapper(schedules.forward_step, "fwd") - schedules.backward_step = fwd_bwd_wrapper(schedules.backward_step, "bwd") - - -def dump_pp_data(args, num_mbs, pp_data_dir): - torch.cuda.synchronize() - - global _GLOBAL_PP_VIS_EVENTS - all_iter_data = {} - for iter_idx, iter_events in enumerate(_GLOBAL_PP_VIS_EVENTS): - iter_data = { - "total": None, - "memory": None, - "fwd_start": [], - "fwd_end": [], - "fwd_minibatch": [], - "fwd_chunk": [], - "bwd_start": [], - "bwd_end": [], - "bwd_minibatch": [], - "bwd_chunk": [], - "wgrad_start": [], - "wgrad_end": [], - "wgrad_minibatch": [], - "wgrad_chunk": [], - } - iter_data["total"] = iter_events["start"].elapsed_time(iter_events["end"]) - iter_data["memory"] = iter_events["memory"] - - for i in range(len(iter_events["fwd_start"])): - for key in ["fwd_start", "fwd_end", "bwd_start", "bwd_end", "wgrad_start", "wgrad_end"]: - if i >= len(iter_events[key]): - continue - event_time = iter_events["start"].elapsed_time(iter_events[key][i]) - iter_data[key].append(event_time) - for key in [ - "fwd_minibatch", - "fwd_chunk", - "bwd_minibatch", - "bwd_chunk", - "wgrad_minibatch", - "wgrad_chunk", - ]: - if i >= len(iter_events[key]): - continue - iter_data[key].append(iter_events[key][i]) - - all_iter_data[iter_idx + 1] = iter_data - - rank = torch.distributed.get_rank() - dp_rank = parallel_state.get_data_parallel_rank() - pp_rank = parallel_state.get_pipeline_model_parallel_rank() - os.makedirs(pp_data_dir, exist_ok=True) - if dp_rank == 0: - log_path = os.path.join(pp_data_dir, f"pp_rank_{pp_rank}.json") - with open(log_path, "w") as f: - json.dump(all_iter_data, f, indent=2) - - if rank == 0: - vp_size = args.virtual_pipeline_model_parallel_size - vp_size = 1 if vp_size is None else vp_size - config_dict = { - "world_size": args.world_size, - "dp_size": args.data_parallel_size, - "tp_size": args.tensor_model_parallel_size, - "ep_size": args.expert_model_parallel_size, - "pp_size": args.pipeline_model_parallel_size, - "vp_size": vp_size, - "num_mbs": num_mbs, - "train_iters": args.train_iters, - } - log_path = os.path.join(pp_data_dir, f"config.json") - with open(log_path, "w") as f: - json.dump(config_dict, f, indent=2) - - -def _get_sync_free_moe_options(args) -> dict: - stage = args.turbo_sync_free_moe_stage - - if stage > 3 or stage < 0: - raise ValueError("turbo_sync_free_moe_stage only support [0-3]") - - sync_free_moe = { - 1: { - "moe_use_fused_router_with_aux_score": True, - "moe_permute_fusion": True, - "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False, - }, - 2: { - "moe_use_fused_router_with_aux_score": True, - "use_turbo_deepep": True, - "moe_permute_fusion": True, - "use_turbo_grouped_gemm": True, - "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False, - }, - 3: { - "moe_use_fused_router_with_aux_score": True, - "use_turbo_deepep": True, - "moe_permute_fusion": True, - "use_turbo_grouped_gemm": True, - "moe_router_padding_for_quantization": True if args.fp8 or args.fp4 else False, - "use_turbo_fused_act_with_probs": True, - }, - } - - return sync_free_moe[stage] - - -# FSDP2 custom optimizer selection flags. Each one monkeypatches -# get_megatron_optimizer at the same priority (50), so at most one may be set. -_FSDP2_OPTIMIZER_FLAGS = ( - "use_fsdp2_fp32_param_optimizer", - "use_fsdp2_bf16_master_weight_optimizer", -) - - -def validate_fsdp2_optimizer_exclusivity(args) -> None: - """Ensure at most one FSDP2 custom optimizer flag is enabled. - - Enabling more than one would silently let whichever optimizer patch applies - last win the monkeypatch, so raise ValueError to fail loudly at - arg-validation time (before training starts). - """ - enabled = [flag for flag in _FSDP2_OPTIMIZER_FLAGS if getattr(args, flag, False)] - if len(enabled) > 1: - raise ValueError( - "Conflicting FSDP2 optimizer selection: at most one of " - f"{list(_FSDP2_OPTIMIZER_FLAGS)} may be enabled, but got {enabled}. " - "Enable exactly one." - ) - - -def validate_args_on_rocm(args): - # Deterministic mode - if args.deterministic_mode: - # NOTE: Some environment variables affect deterministic mode on ROCm. Need to do extra check. - NON_DETERMINISTIC_ENVS = { - "TORCH_COMPILE_DISABLE": "1", - "ROCBLAS_DEFAULT_ATOMICS_MODE": "0", - "PRIMUS_TURBO_AUTO_TUNE": "0", - "PRIMUS_DETERMINISTIC": "1", - } - # NOTE: Some version triton compile exist potential racing condition issue. - for env, value in NON_DETERMINISTIC_ENVS.items(): - assert ( - os.environ.get(env, None) == value - ), f"{env} must be set to {value} in deterministic mode but got {os.environ.get(env, None)} instead." - - # Set fill_uninitialized_memory to False to avoid calling extra fill kernel in deterministic mode. - torch.utils.deterministic.fill_uninitialized_memory = False - - assert not getattr( - args, "use_turbo_parallel_linear", False - ), "use_turbo_parallel_linear has been removed; please use use_turbo_gemm instead." - - validate_fsdp2_optimizer_exclusivity(args) - - use_turbo_gemm = getattr(args, "use_turbo_gemm", False) - # Turbo FP8 linear check - if args.fp8 and use_turbo_gemm: - support_fp8_recipe = ["tensorwise", "blockwise", "mxfp8"] - assert ( - args.fp8_recipe in support_fp8_recipe - ), f"{args.fp8_recipe} recipe is not support when enable `use_turbo_gemm`." - - # Turbo FP4 linear check - if args.fp4 and use_turbo_gemm: - support_fp4_recipe = ["mxfp4"] - assert ( - args.fp4_recipe in support_fp4_recipe - ), f"{args.fp4_recipe} recipe is not support when enable `use_turbo_gemm`." - - # NOTE: mxfp8 environment variable must be set to 1 to enable mxfp8 recipe on ROCm. - if args.fp8_recipe == "mxfp8": - assert ( - os.getenv("NVTE_ROCM_ENABLE_MXFP8", "0") == "1" - ), "Please set `NVTE_ROCM_ENABLE_MXFP8=1` to enable `mxfp8` recipe." - - # dump pp data - if args.dump_pp_data and args.pipeline_model_parallel_size == 1: - args.dump_pp_data = False - print_rank_last(f"Disable args.dump_pp_data since args.pipeline_model_parallel_size=1") - - # PrimusTurboGroupedMLP no longer depends on legacy GroupedMLP; the two - # flags are mutually exclusive when turbo is enabled. - assert not getattr( - args, "use_turbo_grouped_mlp", False - ), "use_turbo_grouped_mlp has been removed; please use use_turbo_grouped_gemm instead." - use_turbo_grouped_gemm = getattr(args, "use_turbo_grouped_gemm", False) - if use_turbo_grouped_gemm: - if getattr(args, "moe_use_legacy_grouped_gemm", False): - raise ValueError( - "use_turbo_grouped_gemm=True is incompatible with moe_use_legacy_grouped_gemm=True. " - "please set moe_use_legacy_grouped_gemm=False." - ) - - # sync-free MoE - if args.turbo_sync_free_moe_stage > 0: - assert args.enable_primus_turbo, "Please set `enable_primus_turbo=True` to enable sync-free MoE." - - if args.turbo_sync_free_moe_stage > 1 and not use_turbo_grouped_gemm: - raise ValueError( - "Sync-Free MoE stage 2 or 3 require PrimusTurboGroupedLinear, please set `use_turbo_grouped_gemm=True`" - ) - options = _get_sync_free_moe_options(args) - print_rank_last( - f"========== Enable Sync-Free MoE Stage {args.turbo_sync_free_moe_stage} (Auto-Enabled Options) ==========" - ) - for flag, value in options.items(): - dots = "." * (73 - len(flag) - len(str(value))) - print_rank_last(f"{flag}{dots}{value}") - setattr(args, flag, value) - print_rank_last( - f"========== Enable Sync-Free MoE Stage {args.turbo_sync_free_moe_stage} (Auto-Enabled Options) ==========" - ) - - # turbo deepep - if args.use_turbo_deepep: - assert ( - not args.moe_shared_expert_overlap - ), "DeepEP not support moe_shared_expert_overlap, please set `moe_shared_expert_overlap=False`." - assert ( - args.moe_router_dtype == "fp32" - ), "DeepEP only supports float32 probs, please set `moe_router_dtype=fp32`" - if ( - args.expert_model_parallel_size >= 16 - and os.getenv("PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND", "TURBO") == "TURBO" - ): - # Turbo DeepEP is not supported for CUs > 32 when using internode dispatch/combine. - assert args.turbo_deepep_num_cu <= 32, "Set `turbo_deepep_num_cu<=32` when using ep_size >= 16." diff --git a/primus/modules/trainer/torchtitan/__init__.py b/primus/modules/trainer/torchtitan/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/primus/modules/trainer/torchtitan/parse_utils.py b/primus/modules/trainer/torchtitan/parse_utils.py deleted file mode 100644 index 17efaaaa2..000000000 --- a/primus/modules/trainer/torchtitan/parse_utils.py +++ /dev/null @@ -1,66 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -import argparse -from argparse import Namespace -from typing import List - - -def _parse_args(ignore_unknown_args: bool = False) -> tuple[Namespace, List[str]]: - parser = argparse.ArgumentParser(description="Primus Arguments", allow_abbrev=False) - parser.add_argument( - "--exp", - type=str, - required=True, - help="Primus experiment yaml config file.", - ) - return parser.parse_known_args() if ignore_unknown_args else (parser.parse_args(), []) - - -def _merge_args(args: Namespace, unknown_args: List[str]) -> Namespace: - merged_dict = vars(args).copy() - temp_parser = argparse.ArgumentParser() - - i = 0 - while i < len(unknown_args): - key = unknown_args[i] - if key.startswith("--"): - if i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith("--"): - temp_parser.add_argument(key, type=str) - i += 2 - else: - temp_parser.add_argument(key, action="store_true") - i += 1 - else: - i += 1 - - parsed_unknown, _ = temp_parser.parse_known_args(unknown_args) - merged_dict.update(vars(parsed_unknown)) - return Namespace(**merged_dict) - - -def _convert_args_to_cli(args: Namespace) -> List[str]: - cli_args = [] - - if hasattr(args, "exp"): - cli_args += ["--job.config_file", args.exp] - - for k, v in vars(args).items(): - if k in {"exp", "backend"}: - continue - if isinstance(v, bool): - if v: - cli_args.append(f"--{k}") - else: - cli_args += [f"--{k}", str(v)] - - return cli_args - - -def get_torchtitan_config_args() -> List[str]: - args, unknown_args = _parse_args(ignore_unknown_args=True) - merged = _merge_args(args, unknown_args) - return _convert_args_to_cli(merged) diff --git a/primus/modules/trainer/torchtitan/pre_trainer.py b/primus/modules/trainer/torchtitan/pre_trainer.py deleted file mode 100644 index c178d853b..000000000 --- a/primus/modules/trainer/torchtitan/pre_trainer.py +++ /dev/null @@ -1,297 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -from dataclasses import asdict, is_dataclass -from types import SimpleNamespace -from typing import Any, Dict - -# Trigger registration of all TorchTitan patches -import primus.backends.torchtitan.patches # noqa: F401 -from primus.core.patches import run_patches -from primus.core.utils.yaml_utils import nested_namespace_to_dict -from primus.modules.base_module import BaseModule - - -class TorchTitanPretrainTrainer(BaseModule): - def __init__(self, *args, **kwargs): - extra_args = kwargs.pop("extra_args", None) - super().__init__(*args, **kwargs) - - self.primus_cfg = kwargs.pop("primus_config", None) - if self.primus_cfg is None: - raise ValueError("primus_config is required") - - pre_trainer_cfg = self.primus_cfg.get_module_config("pre_trainer") - cfg_dict = nested_namespace_to_dict(pre_trainer_cfg) - - # Run before_train phase patches (e.g., Primus-Turbo patches) - # Construct a temporary module_config for patch context - temp_module_config = SimpleNamespace() - temp_module_config.params = pre_trainer_cfg - temp_module_config.name = "pre_trainer" - temp_module_config.framework = "torchtitan" - temp_module_config.model = getattr(pre_trainer_cfg, "model.name", None) - - # Merge extra_args into temp_module_config.params if provided - if extra_args: - temp_module_config.params = self._merge_extra_args_into_params( - temp_module_config.params, extra_args - ) - - run_patches( - backend="torchtitan", - phase="setup", - backend_version="unknown", - extra={ - "module_config": temp_module_config, - }, - ) - - from torchtitan.config.job_config import JobConfig - from torchtitan.train import Trainer - - self.TrainerClass = Trainer - self.JobConfigClass = JobConfig - - self.titan_config = self.build_job_config(cfg_dict, self.JobConfigClass) - - self.log_config(self.titan_config) - self.trainer = None - - def _merge_extra_args_into_params( - self, params: SimpleNamespace, extra_args: Dict[str, Any] - ) -> SimpleNamespace: - """ - Merge extra_args into params with specific rules: - - 1. Non-model.* parameters: Only merge if they already exist in params - 2. model.* parameters: Extract and merge (can add new fields) - - Args: - params: Original params namespace - extra_args: Additional arguments to merge - - Returns: - Updated params with merged extra_args - """ - from primus.modules.module_utils import log_rank_0 - - # First, flatten extra_args if it contains nested dicts - flattened_extra_args = self._flatten_dict(extra_args) - - # Separate model.* and non-model.* parameters - model_params = {} - non_model_params = {} - - for key, value in flattened_extra_args.items(): - if key.startswith("model."): - # Extract model parameter (remove "model." prefix) - model_key = key[6:] # Remove "model." prefix - model_params[model_key] = value - else: - non_model_params[key] = value - - # Convert params to dict for easier manipulation - params_dict = nested_namespace_to_dict(params) - - # Merge non-model.* parameters (only if they already exist) - for key, value in non_model_params.items(): - if self._key_exists_in_dict(params_dict, key): - self._set_nested_dict_value(params_dict, key, value) - log_rank_0(f"[ExtraArgs] Merged non-model param: {key} = {value}") - else: - log_rank_0(f"[ExtraArgs] Skipped non-model param (not in params): {key} = {value}") - - # Merge model.* parameters (can add new fields) - if model_params: - if "model" not in params_dict: - params_dict["model"] = {} - - for key, value in model_params.items(): - self._set_nested_dict_value(params_dict, f"model.{key}", value) - log_rank_0(f"[ExtraArgs] Merged model param: model.{key} = {value}") - - # Convert back to SimpleNamespace - return self._dict_to_namespace(params_dict) - - def _flatten_dict(self, d: Dict[str, Any], prefix: str = "") -> Dict[str, Any]: - """Flatten nested dictionary using dot notation.""" - result = {} - - for key, value in d.items(): - full_key = f"{prefix}.{key}" if prefix else key - - if isinstance(value, dict): - # Recursively flatten nested dict - result.update(self._flatten_dict(value, full_key)) - else: - # Leaf value - result[full_key] = value - - return result - - def _key_exists_in_dict(self, d: Dict[str, Any], key: str) -> bool: - """Check if a nested key exists in dictionary (supports dot notation).""" - parts = key.split(".") - current = d - - for part in parts: - if not isinstance(current, dict) or part not in current: - return False - current = current[part] - - return True - - def _set_nested_dict_value(self, d: Dict[str, Any], key: str, value: Any) -> None: - """Set a nested dictionary value using dot notation.""" - parts = key.split(".") - current = d - - for part in parts[:-1]: - if part not in current: - current[part] = {} - current = current[part] - - current[parts[-1]] = value - - def _dict_to_namespace(self, d: Dict[str, Any]) -> SimpleNamespace: - """Recursively convert dictionary to SimpleNamespace.""" - if isinstance(d, dict): - return SimpleNamespace(**{k: self._dict_to_namespace(v) for k, v in d.items()}) - elif isinstance(d, list): - return [self._dict_to_namespace(item) for item in d] - else: - return d - - def setup(self): - pass - - def init(self, *init_args, **kwargs): - self.trainer = self.TrainerClass(self.titan_config) - - def run(self, *args, **kwargs): - if self.trainer is None: - raise RuntimeError("Trainer has not been initialized. Call init() first.") - self.trainer.train() - - def flatten_config(self, obj: Any, prefix: str = "") -> Dict[str, Any]: - flat_dict = {} - if is_dataclass(obj): - obj = asdict(obj) - - if isinstance(obj, dict): - for key, value in obj.items(): - full_key = f"{prefix}.{key}" if prefix else key - if is_dataclass(value) or isinstance(value, dict): - flat_dict.update(self.flatten_config(value, full_key)) - else: - flat_dict[full_key] = value - else: - flat_dict[prefix] = obj - - return flat_dict - - def log_config(self, obj: Any, header: str = "TorchTitan Config"): - from torchtitan.tools.logging import logger - - logger.info("========== %s ==========" % header) - flat = self.flatten_config(obj) - max_key_len = max(len(k) for k in flat.keys()) - for key in sorted(flat): - val = flat[key] - formatted_line = f"arguments {key.ljust(max_key_len, '.')} {val}" - logger.info(formatted_line) - - def build_job_config(self, cfg_dict: dict, JobConfigType) -> Any: - import importlib - - from torchtitan.config.job_config import Experimental - from torchtitan.tools.logging import logger - - # Step 1: Parse the experimental section to check for a custom JobConfig extension - experimental_cfg = cfg_dict.get("experimental", {}) - experimental = Experimental(**experimental_cfg) - - # Step 2: If a custom_args_module is defined, import and merge with JobConfig - custom_job_config_cls = JobConfigType - if experimental and getattr(experimental, "custom_args_module", None): - try: - module = importlib.import_module(experimental.custom_args_module) - ExtendedJobConfig = getattr(module, "JobConfig") - custom_job_config_cls = self.merge_configs(JobConfigType, ExtendedJobConfig) - logger.info(f"Loaded and merged custom JobConfig from {experimental.custom_args_module}") - except Exception as e: - logger.warning(f"Failed to load custom_args_module '{experimental.custom_args_module}': {e}") - - # Step 3: Parse config dict (including custom fields) into dataclass recursively - return self._dict_to_dataclass(custom_job_config_cls, cfg_dict) - - @staticmethod - def merge_configs(base_cls, custom_cls): - """ - Merges two dataclass types into one unified dataclass. - - Merge logic: - - If a field exists in both: - - If both fields are dataclasses, recursively merge them. - - Otherwise, the custom field overrides the base. - - Fields only in base or only in custom are included as-is. - """ - from dataclasses import field, fields, make_dataclass - - base_fields = {f.name: f for f in fields(base_cls)} - custom_fields = {f.name: f for f in fields(custom_cls)} - - merged = [] - - # Merge overlapping and base-only fields - for name, base_f in base_fields.items(): - if name in custom_fields: - custom_f = custom_fields[name] - if is_dataclass(base_f.type) and is_dataclass(custom_f.type): - merged_type = TorchTitanPretrainTrainer.merge_configs(base_f.type, custom_f.type) - merged.append((name, merged_type, field(default_factory=merged_type))) - else: - merged.append((name, custom_f.type, custom_f)) - else: - merged.append((name, base_f.type, base_f)) - - # Add custom-only fields - for name, custom_f in custom_fields.items(): - if name not in base_fields: - merged.append((name, custom_f.type, custom_f)) - - return make_dataclass(f"Merged{base_cls.__name__}", merged, bases=(base_cls,)) - - def _dict_to_dataclass(self, cls, data: dict[str, Any]) -> Any: - """Recursively convert dictionary to dataclass, handling nested and custom fields.""" - from dataclasses import fields, is_dataclass - - if not is_dataclass(cls): - return data - - # collect valid field names - field_names = {f.name for f in fields(cls)} - init_values = {} - - # only use known fields for constructor - for f in fields(cls): - if f.name in data: - val = data[f.name] - if is_dataclass(f.type) and isinstance(val, dict): - init_values[f.name] = self._dict_to_dataclass(f.type, val) - else: - init_values[f.name] = val - - # instantiate dataclass - obj = cls(**init_values) - - # attach unknown fields dynamically - for k, v in data.items(): - if k not in field_names: - setattr(obj, k, v) - - return obj diff --git a/primus/pretrain.py b/primus/pretrain.py index 9e8dbb26f..e4b4ca0e8 100644 --- a/primus/pretrain.py +++ b/primus/pretrain.py @@ -4,14 +4,19 @@ # See LICENSE for license information. ############################################################################### -import argparse +"""Backend-path / environment helpers shared by the training and projection +entry points. + +Training is driven entirely by the core runtime +(:mod:`primus.core.runtime.train_runtime`). Only the backend-path resolution +utilities remain here, still used by the projection subcommand, runner hooks +and examples. +""" + import os import sys from pathlib import Path -from primus.core.launcher.config import PrimusConfig -from primus.core.launcher.parser import add_pretrain_parser, load_primus_config - def _info_enabled() -> bool: """True when PRIMUS_LOG_LEVEL permits INFO-level chatter (DEBUG/INFO). @@ -22,65 +27,6 @@ def _info_enabled() -> bool: return os.environ.get("PRIMUS_LOG_LEVEL", "INFO").upper() in ("DEBUG", "INFO") -# Lazy backend loader -def load_backend_trainer(framework: str): - if framework == "megatron": - import megatron.training.training as training - import torch - - _original_build_model = training.get_model - - def _patched_get_model(*args, **kwargs): - """ - Monkey-patched version of build_model that removes the second - DDP construction inside torch.cuda.stream() block. - """ - import inspect - - from megatron.training import training as tr - - inspect.getsource(tr.get_model) - print("[PrimusPatch] Overriding build_model to disable second DDP construction...") - - _orig_stream_ctx = torch.cuda.stream - - def _noop_stream(*args, **kwargs): - class DummyCtx: - def __enter__(self): - return None - - def __exit__(self, *a): - return False - - return DummyCtx() - - torch.cuda.stream = _noop_stream - - try: - return _original_build_model(*args, **kwargs) - finally: - torch.cuda.stream = _orig_stream_ctx - - training.get_model = _patched_get_model - print("[PrimusPatch] Applied Megatron build_model monkey-patch to disable second DDP.") - - from primus.modules.trainer.megatron.pre_trainer import MegatronPretrainTrainer - - return MegatronPretrainTrainer - elif framework == "torchtitan": - from primus.modules.trainer.torchtitan.pre_trainer import ( - TorchTitanPretrainTrainer, - ) - - return TorchTitanPretrainTrainer - elif framework == "maxtext": - from primus.modules.trainer.maxtext.pre_trainer import MaxTextPretrainTrainer - - return MaxTextPretrainTrainer - else: - raise ValueError(f"Unsupported framework: {framework}") - - def setup_backend_path(framework: str, backend_path=None, verbose: bool = True): """ Setup Python path for backend modules. @@ -159,83 +105,3 @@ def setup_env(data_path: str): else: hf_home = os.environ["HF_HOME"] print(f"[Primus CLI] HF_HOME already set: {hf_home}") - - -def launch_pretrain_trainer(primus_cfg: PrimusConfig, extra_args=None): - """ - Launch the training using the Primus trainer. - - Args: - primus_cfg (PrimusConfig): Parsed Primus configuration object. - """ - # Get pre_trainer module configuration - pre_trainer_cfg = primus_cfg.get_module_config("pre_trainer") - framework = pre_trainer_cfg.framework - - # Lazy import backend trainer - TrainerClass = load_backend_trainer(framework) - - master_addr = os.getenv("MASTER_ADDR", "127.0.0.1") - master_port = int(os.getenv("MASTER_PORT", "29500")) - - if framework == "maxtext": - rank = int(os.getenv("NODE_RANK", "0")) - world_size = int(os.getenv("NNODES", "1")) - else: - # envs set by torchrun - rank = int(os.getenv("RANK", "0")) - world_size = int(os.getenv("WORLD_SIZE", "1")) - - # Initialize trainer - trainer = TrainerClass( - module_name="pre_trainer", - primus_config=primus_cfg, - module_rank=rank, - module_world_size=world_size, - module_master_addr=master_addr, - module_master_port=master_port, - extra_args=extra_args, - ) - - # Launch training - trainer.init() - trainer.run() - - -def launch_pretrain_from_cli(args, overrides): - """ - Entry point for the 'train' subcommand. - - Steps: - 1. Load and parse the experiment YAML config - 2. Merge CLI overrides into the config - 3. Optionally export the merged config - 4. Setup backend path - 5. Launch the training - """ - cfg_path = Path(args.config) - if not cfg_path.exists(): - raise FileNotFoundError(f"[Primus:Train] Config file '{cfg_path}' not found.") - - setup_env(data_path=args.data_path) - - primus_cfg, unknown_overrides = load_primus_config(args, overrides) - - # Export merged config if requested - if args.export_config: - primus_cfg.export(export_path=args.export_config) - - # Setup backend path for dynamic import - framework = primus_cfg.get_module_config("pre_trainer").framework - setup_backend_path(framework=framework, backend_path=args.backend_path, verbose=True) - - launch_pretrain_trainer(primus_cfg=primus_cfg, extra_args=unknown_overrides) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="pretrain") - add_pretrain_parser(parser) - - args, unknown_args = parser.parse_known_args() - - launch_pretrain_from_cli(args, unknown_args) diff --git a/skills/backend-gap-report/examples.md b/skills/backend-gap-report/examples.md index 047003287..1def34df7 100644 --- a/skills/backend-gap-report/examples.md +++ b/skills/backend-gap-report/examples.md @@ -44,7 +44,7 @@ Suggested evidence sources: - `third_party/Megatron-LM/requirements*` - `third_party/Megatron-LM/.github/workflows/*` - `primus/backends/megatron/*` -- `primus/modules/trainer/megatron/*` +- `primus/backends/megatron/pretrainer/*` - `docs/backends/megatron/patch-notes.md` Expected outputs: diff --git a/tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py b/tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py index e4951cfe5..7869ce074 100644 --- a/tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py +++ b/tests/unit_tests/backends/megatron/patches/test_fsdp2_fp32_patches.py @@ -453,7 +453,7 @@ def test_two_conflicting_optimizer_flags_raise(self): patch get_megatron_optimizer at the same priority, so enabling more than one is ambiguous. validate_fsdp2_optimizer_exclusivity enforces this at arg-validation time.""" - from primus.modules.trainer.megatron.utils import ( + from primus.backends.megatron.patches.args.rocm_arg_validation import ( validate_fsdp2_optimizer_exclusivity, ) @@ -467,7 +467,7 @@ def test_two_conflicting_optimizer_flags_raise(self): def test_single_optimizer_flag_is_allowed(self): """Exactly one FSDP2 optimizer flag is the supported case (no raise).""" - from primus.modules.trainer.megatron.utils import ( + from primus.backends.megatron.patches.args.rocm_arg_validation import ( validate_fsdp2_optimizer_exclusivity, ) @@ -481,7 +481,7 @@ def test_single_optimizer_flag_is_allowed(self): def test_no_optimizer_flag_is_allowed(self): """No FSDP2 optimizer flag set (default Megatron optimizer) must not raise.""" - from primus.modules.trainer.megatron.utils import ( + from primus.backends.megatron.patches.args.rocm_arg_validation import ( validate_fsdp2_optimizer_exclusivity, ) diff --git a/tests/unit_tests/backends/megatron/test_megatron_adapter.py b/tests/unit_tests/backends/megatron/test_megatron_adapter.py index 74e4a098c..a6c8ba1ef 100644 --- a/tests/unit_tests/backends/megatron/test_megatron_adapter.py +++ b/tests/unit_tests/backends/megatron/test_megatron_adapter.py @@ -188,7 +188,7 @@ def test_load_trainer_class_fallback_tries_multiple_paths(self, mock_log): def side_effect(module_name, *args, **kwargs): import_calls.append(module_name) - if len(import_calls) <= 2: + if len(import_calls) <= 1: raise ImportError(f"Path {len(import_calls)} failed") return mock_module @@ -197,7 +197,7 @@ def side_effect(module_name, *args, **kwargs): result = adapter.load_trainer_class(trainer_class="CustomExperimentalTrainer") assert result == mock_trainer_class - assert len(import_calls) == 3 + assert len(import_calls) == 2 assert "primus.backends.megatron.customexperimentaltrainer" in import_calls[0].lower() @patch("primus.backends.megatron.megatron_adapter.log_rank_0") @@ -283,7 +283,7 @@ def side_effect(module_name, *args, **kwargs): class TestMegatronAdapterIntegration: """Integration tests for complete adapter workflow.""" - @patch("primus.modules.module_utils.log_rank_0") + @patch("primus.core.utils.module_utils.log_rank_0") @patch("primus.backends.megatron.megatron_adapter.MegatronAdapter.detect_backend_version") @patch("primus.backends.megatron.megatron_adapter.MegatronArgBuilder") @patch("primus.backends.megatron.megatron_adapter.log_rank_0") diff --git a/tests/unit_tests/backends/megatron/test_training_log_patches.py b/tests/unit_tests/backends/megatron/test_training_log_patches.py index 2f8f49683..6776c23ff 100644 --- a/tests/unit_tests/backends/megatron/test_training_log_patches.py +++ b/tests/unit_tests/backends/megatron/test_training_log_patches.py @@ -32,9 +32,9 @@ def fake_training_log(*args, **kwargs): return "ok" # Provide a minimal `get_model` stub so that any code which expects - # `megatron.training.training.get_model` (e.g., Primus monkey patches in - # `primus.pretrain.load_backend_trainer`) can safely import and patch this - # fake training module during tests without raising AttributeError. + # `megatron.training.training.get_model` (e.g., Primus monkey patches) + # can safely import and patch this fake training module during tests + # without raising AttributeError. def fake_get_model(*args, **kwargs): return None diff --git a/tests/unit_tests/backends/megatron/test_validate_args_patches.py b/tests/unit_tests/backends/megatron/test_validate_args_patches.py index 790775ffb..bd58796cf 100644 --- a/tests/unit_tests/backends/megatron/test_validate_args_patches.py +++ b/tests/unit_tests/backends/megatron/test_validate_args_patches.py @@ -53,7 +53,7 @@ def _install_fake_megatron(monkeypatch: pytest.MonkeyPatch): """Create fake ``megatron.*`` modules sufficient for validate_args patches. Includes ``megatron.core.parallel_state`` and ``megatron.training.global_vars`` - so that transitive imports from ``primus.modules.trainer.megatron.utils`` succeed. + so that transitive imports from ``primus.backends.megatron.patches.args.rocm_arg_validation`` succeed. """ import sys @@ -149,7 +149,7 @@ def test_wraps_validate_args_on_both_modules(self, monkeypatch): args_mod, init_mod = _install_fake_megatron(monkeypatch) _silence_logging(monkeypatch) monkeypatch.setattr( - "primus.modules.trainer.megatron.utils.validate_args_on_rocm", + "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm", lambda args: setattr(args, "_rocm_validated", True), ) @@ -167,7 +167,7 @@ def test_calls_original_and_rocm_validation(self, monkeypatch): _install_fake_megatron(monkeypatch) _silence_logging(monkeypatch) monkeypatch.setattr( - "primus.modules.trainer.megatron.utils.validate_args_on_rocm", + "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm", lambda args: setattr(args, "_rocm_validated", True), ) @@ -190,7 +190,7 @@ def test_stores_original_on_module(self, monkeypatch): args_mod, _ = _install_fake_megatron(monkeypatch) _silence_logging(monkeypatch) monkeypatch.setattr( - "primus.modules.trainer.megatron.utils.validate_args_on_rocm", + "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm", lambda args: None, ) @@ -212,7 +212,7 @@ def _apply_base_and_split(self, monkeypatch): _install_fake_megatron(monkeypatch) _silence_logging(monkeypatch) monkeypatch.setattr( - "primus.modules.trainer.megatron.utils.validate_args_on_rocm", + "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm", lambda args: None, ) @@ -252,7 +252,7 @@ def _apply_base_and_fp4(self, monkeypatch): _install_fake_megatron(monkeypatch) _silence_logging(monkeypatch) monkeypatch.setattr( - "primus.modules.trainer.megatron.utils.validate_args_on_rocm", + "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm", lambda args: None, ) @@ -288,7 +288,7 @@ def test_without_patch_fp4_raises(self, monkeypatch): _install_fake_megatron(monkeypatch) _silence_logging(monkeypatch) monkeypatch.setattr( - "primus.modules.trainer.megatron.utils.validate_args_on_rocm", + "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm", lambda args: None, ) @@ -314,7 +314,7 @@ def test_pipeline_split_and_rocm_validation(self, monkeypatch): rocm_calls = [] monkeypatch.setattr( - "primus.modules.trainer.megatron.utils.validate_args_on_rocm", + "primus.backends.megatron.patches.args.rocm_arg_validation.validate_args_on_rocm", lambda args: rocm_calls.append(True), ) diff --git a/tests/unit_tests/cli/test_train_subcommand.py b/tests/unit_tests/cli/test_train_subcommand.py deleted file mode 100644 index 04955245a..000000000 --- a/tests/unit_tests/cli/test_train_subcommand.py +++ /dev/null @@ -1,40 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -"""Unit tests for the train subcommand's runtime router (``_resolve_pretrain_runtime``). - -This selects the active core runtime vs the legacy pretrain flow; run()/register -import heavy runtime modules and are out of scope. -""" - -from __future__ import annotations - -import pytest - -pytest.importorskip("primus.cli.subcommands.train") - -from primus.cli.subcommands.train import _resolve_pretrain_runtime # noqa: E402 - - -def test_runtime_defaults_to_core(monkeypatch): - monkeypatch.delenv("PRIMUS_TRAIN_RUNTIME", raising=False) - assert _resolve_pretrain_runtime(None) == "core" - - -def test_runtime_explicit_legacy(monkeypatch): - monkeypatch.setenv("PRIMUS_TRAIN_RUNTIME", "legacy") - assert _resolve_pretrain_runtime(None) == "legacy" - - -def test_runtime_explicit_core_case_insensitive(monkeypatch): - monkeypatch.setenv("PRIMUS_TRAIN_RUNTIME", "CORE") - assert _resolve_pretrain_runtime(None) == "core" - - -def test_runtime_invalid_value_warns_and_falls_back_to_core(monkeypatch, capsys): - monkeypatch.setenv("PRIMUS_TRAIN_RUNTIME", "bogus") - assert _resolve_pretrain_runtime(None) == "core" - assert "Ignoring invalid" in capsys.readouterr().err diff --git a/tests/unit_tests/core/trainer/test_base_trainer.py b/tests/unit_tests/core/trainer/test_base_trainer.py index 9595b71c1..b3ee0beab 100644 --- a/tests/unit_tests/core/trainer/test_base_trainer.py +++ b/tests/unit_tests/core/trainer/test_base_trainer.py @@ -59,7 +59,7 @@ def train(self): def test_init_mro_with_base_module(self, monkeypatch: pytest.MonkeyPatch): """Test MRO handling when BaseModule IS in inheritance chain (legacy pattern).""" - from primus.modules.base_module import BaseModule + from primus.core.base_module import BaseModule # Create a class that inherits from both BaseTrainer and BaseModule class LegacyTrainer(BaseTrainer, BaseModule): diff --git a/tests/unit_tests/megatron/cco/test_tp_overlap.py b/tests/unit_tests/megatron/cco/test_tp_overlap.py index b906a34b3..ceab475e2 100644 --- a/tests/unit_tests/megatron/cco/test_tp_overlap.py +++ b/tests/unit_tests/megatron/cco/test_tp_overlap.py @@ -28,7 +28,7 @@ initialize_ub, ) from primus.core.utils import logger -from primus.modules.module_utils import set_logging_rank +from primus.core.utils.module_utils import set_logging_rank @contextmanager diff --git a/tests/unit_tests/test_backend_loader.py b/tests/unit_tests/test_backend_loader.py index efee65af5..d143aabff 100644 --- a/tests/unit_tests/test_backend_loader.py +++ b/tests/unit_tests/test_backend_loader.py @@ -8,12 +8,11 @@ import shutil import sys import tempfile -import types from pathlib import Path import pytest -from primus.pretrain import load_backend_trainer, setup_backend_path +from primus.pretrain import setup_backend_path @pytest.fixture @@ -64,23 +63,3 @@ def test_setup_backend_path_failure(): """Test that FileNotFoundError is raised when no valid backend path exists.""" with pytest.raises(FileNotFoundError): setup_backend_path(framework="nonexistent_backend") - - -def test_load_backend_trainer_supported(monkeypatch): - """Test that load_backend_trainer returns correct class for supported frameworks.""" - # Mock: Define a dummy class to simulate MegatronPretrainTrainer - dummy_class = type("DummyTrainer", (), {}) - dummy_module = types.ModuleType("primus.modules.trainer.megatron.pre_trainer") - dummy_module.MegatronPretrainTrainer = dummy_class - - # Inject into sys.modules to bypass real import - sys.modules["primus.modules.trainer.megatron.pre_trainer"] = dummy_module - - trainer_cls = load_backend_trainer("megatron") - assert trainer_cls is dummy_class - - -def test_load_backend_trainer_unsupported(): - """Test that unsupported framework raises ValueError.""" - with pytest.raises(ValueError, match="Unsupported framework"): - load_backend_trainer("invalid_framework") From aafbc3f2f55c9d9b709de8f268118d1cde70088d Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Wed, 8 Jul 2026 18:29:43 +0300 Subject: [PATCH 016/127] feat(flux): common diffusion module (embeddings, normalization, DiT block) (#810) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/core` — review after it. Parent of the Flux model PR. ## What this changes The model-agnostic diffusion building blocks shared by all diffusion models: the common diffusion module, diffusion config/embeddings/normalization, and the diffusion transformer (DiT) block. Also lands the shared diffusion test scaffolding (`conftest`/`helpers`/`constants`) the later diffusion PRs reuse. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/core`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "embeddings or normalization"`. Validated locally on an AMD GPU container: 7 passed. ## Files 14 (common diffusion module, embeddings/normalization/config, DiT block, shared diffusion test scaffolding). --------- Co-authored-by: Flux Split Trial Co-authored-by: luiza-amd Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- .../common/diffusion_module/__init__.py | 12 + .../diffusion_module/diffusion_module.py | 176 +++ .../core/models/diffusion/__init__.py | 36 + .../core/models/diffusion/common/__init__.py | 38 + .../core/models/diffusion/common/config.py | 190 +++ .../models/diffusion/common/embeddings.py | 296 ++++ .../models/diffusion/common/normalization.py | 1226 +++++++++++++++++ .../diffusion_transformer_block.py | 352 +++++ .../backends/megatron/diffusion/__init__.py | 2 + .../backends/megatron/diffusion/conftest.py | 47 + .../backends/megatron/diffusion/constants.py | 73 + .../backends/megatron/diffusion/helpers.py | 245 ++++ .../diffusion/test_flux_embeddings.py | 57 + .../diffusion/test_flux_normalization.py | 147 ++ 14 files changed, 2897 insertions(+) create mode 100644 primus/backends/megatron/core/models/common/diffusion_module/__init__.py create mode 100644 primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py create mode 100644 primus/backends/megatron/core/models/diffusion/__init__.py create mode 100644 primus/backends/megatron/core/models/diffusion/common/__init__.py create mode 100644 primus/backends/megatron/core/models/diffusion/common/config.py create mode 100644 primus/backends/megatron/core/models/diffusion/common/embeddings.py create mode 100644 primus/backends/megatron/core/models/diffusion/common/normalization.py create mode 100644 primus/backends/megatron/core/transformer/diffusion_transformer_block.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/__init__.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/conftest.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/constants.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/helpers.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_embeddings.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py diff --git a/primus/backends/megatron/core/models/common/diffusion_module/__init__.py b/primus/backends/megatron/core/models/common/diffusion_module/__init__.py new file mode 100644 index 000000000..399364b6a --- /dev/null +++ b/primus/backends/megatron/core/models/common/diffusion_module/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Base diffusion module with Megatron-Core integration. +""" + +from primus.backends.megatron.core.models.common.diffusion_module.diffusion_module import ( + DiffusionModule, +) + +__all__ = ["DiffusionModule"] diff --git a/primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py b/primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py new file mode 100644 index 000000000..fd5d6cec1 --- /dev/null +++ b/primus/backends/megatron/core/models/common/diffusion_module/diffusion_module.py @@ -0,0 +1,176 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Base diffusion module with Megatron-Core integration. + +This module provides common infrastructure for all diffusion models (Flux, DiT, etc.), +following the architecture pattern established by LanguageModule in Megatron-LM. +""" + +import os +from abc import abstractmethod +from typing import Any, Dict, Optional + +import torch.nn as nn +from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.enums import AttnBackend +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import get_tensor_model_parallel_group_if_none + + +class DiffusionModule(MegatronModule): + """ + Base diffusion module with Megatron-Core integration. + + Provides common infrastructure for all diffusion models (Flux, DiT, MovieGen, etc.): + - Process group management (TP, PP, CP, DP) + - Attention backend configuration + - Distributed checkpointing support + - Encoder management (VAE, T5, CLIP, etc.) + - Common loss computation utilities + + This class follows the architecture pattern of LanguageModule from Megatron-LM, + adapted for diffusion model requirements. + + Args: + config (TransformerConfig): Transformer config with diffusion-specific parameters + pg_collection (Optional[ProcessGroupCollection]): Model communication process groups. + Defaults to None (uses MPU process groups). + encoder_configs (Optional[Dict[str, Any]]): Optional encoder configurations. + Defaults to None. + Format: {'encoder_name': EncoderConfig, ...} + Example: {'vae': VAEConfig(...), 't5': T5XXLConfig(...), 'clip': CLIPLConfig(...)} + + Example: + >>> from primus.backends.megatron.core.models.diffusion.flux import FluxConfig + >>> config = FluxConfig.flux_12b() + >>> model = Flux(config=config) + + Reference: + Megatron-LM LanguageModule: megatron/core/models/common/language_module/language_module.py + """ + + def __init__( + self, + config: TransformerConfig, + pg_collection: Optional[ProcessGroupCollection] = None, + encoder_configs: Optional[Dict[str, Any]] = None, + ) -> None: + super().__init__(config=config) + + # Configure attention backend + self._set_attention_backend() + + # Setup process groups for distributed training + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + self.pg_collection = pg_collection + self.cp_group = pg_collection.cp + self.tp_group = get_tensor_model_parallel_group_if_none(pg_collection.tp) + self.pp_group = pg_collection.pp + + # Note: Diffusion models typically don't need embd_group since they don't share + # embeddings across pipeline stages like language models do + + # Virtual pipeline parallelism + self.vp_stage = None + self.vp_size = self.config.virtual_pipeline_model_parallel_size + + # Encoder management (VAE, T5, CLIP, etc.) + self.encoder_configs = encoder_configs or {} + self.encoders = nn.ModuleDict() + + def _set_attention_backend(self): + """ + Set attention backend for TransformerEngine. + + Configures environment variables to control which attention implementation + TransformerEngine uses. Options: + - flash: Flash Attention (fastest, requires Ampere+ GPUs) + - fused: Fused attention kernels + - unfused: Standard unfused attention + - auto: Let TE choose automatically + - local: Disable TE attention, use PyTorch native + + TransformerEngine works on opt-out basis. By default all three attention + backend flags are set to 1. If the user chooses a particular attention + backend, we set the other two to 0. If the user chooses local, we set + all 3 TE env variables to 0. + + Reference: + LanguageModule._set_attention_backend() in Megatron-LM + """ + + def check_and_set_env_variable( + env_variable_name: str, expected_value: int, attn_type: AttnBackend + ) -> None: + current_value = os.getenv(env_variable_name) + if current_value is not None and current_value != str(expected_value): + raise ValueError( + f"{env_variable_name} set to {current_value}, but expected {expected_value} " + f"for attention backend type {attn_type.name}. Unset NVTE_FLASH_ATTN, " + f"NVTE_FUSED_ATTN and NVTE_UNFUSED_ATTN. Use the --attention-backend argument " + f"if you want to choose between (flash/fused/unfused/auto/local). Default is auto." + ) + os.environ[env_variable_name] = str(expected_value) + + if self.config.attention_backend == AttnBackend.local: + check_and_set_env_variable("NVTE_FLASH_ATTN", 0, AttnBackend.local) + check_and_set_env_variable("NVTE_FUSED_ATTN", 0, AttnBackend.local) + check_and_set_env_variable("NVTE_UNFUSED_ATTN", 0, AttnBackend.local) + elif self.config.attention_backend == AttnBackend.flash: + check_and_set_env_variable("NVTE_FLASH_ATTN", 1, AttnBackend.flash) + check_and_set_env_variable("NVTE_FUSED_ATTN", 0, AttnBackend.flash) + check_and_set_env_variable("NVTE_UNFUSED_ATTN", 0, AttnBackend.flash) + elif self.config.attention_backend == AttnBackend.fused: + check_and_set_env_variable("NVTE_FLASH_ATTN", 0, AttnBackend.fused) + check_and_set_env_variable("NVTE_FUSED_ATTN", 1, AttnBackend.fused) + check_and_set_env_variable("NVTE_UNFUSED_ATTN", 0, AttnBackend.fused) + elif self.config.attention_backend == AttnBackend.unfused: + check_and_set_env_variable("NVTE_FLASH_ATTN", 0, AttnBackend.unfused) + check_and_set_env_variable("NVTE_FUSED_ATTN", 0, AttnBackend.unfused) + check_and_set_env_variable("NVTE_UNFUSED_ATTN", 1, AttnBackend.unfused) + elif self.config.attention_backend == AttnBackend.auto: + check_and_set_env_variable("NVTE_FLASH_ATTN", 1, AttnBackend.auto) + check_and_set_env_variable("NVTE_FUSED_ATTN", 1, AttnBackend.auto) + check_and_set_env_variable("NVTE_UNFUSED_ATTN", 1, AttnBackend.auto) + + @abstractmethod + def forward(self, *args, **kwargs): + """ + Forward pass through diffusion model. + + This method must be implemented by all subclasses to define the + model's forward computation. + + Returns: + Model prediction (noise, velocity, or other target type depending on model) + """ + + def sharded_state_dict( + self, + prefix: str = "", + sharded_offsets: tuple = (), + metadata: Optional[dict] = None, + ) -> ShardedStateDict: + """ + Generate sharded state dictionary for distributed checkpointing. + + Subclasses should override this to handle model-specific checkpointing logic + (e.g., weight tying, custom sharding strategies). + + Args: + prefix: Prefix for state dict keys (e.g., 'module.') + sharded_offsets: Pipeline parallel offsets + metadata: Optional metadata for checkpoint conversion + + Returns: + Dictionary mapping state dict keys to ShardedTensor objects + + Reference: + LanguageModule.sharded_state_dict() in Megatron-LM + """ + return super().sharded_state_dict(prefix, sharded_offsets, metadata) diff --git a/primus/backends/megatron/core/models/diffusion/__init__.py b/primus/backends/megatron/core/models/diffusion/__init__.py new file mode 100644 index 000000000..95a3c358f --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Diffusion models for Primus-Megatron. + +This module contains implementations of various diffusion model architectures +following Megatron-Core conventions. + +Supported models: + - Flux: Flow-based diffusion with MMDiT architecture + - (Future) DiT: Diffusion Transformer + - (Future) MovieGen: Video diffusion +""" + + +# Lazy import for Flux to avoid early dependencies. Import common components and +# other models from their submodules (e.g. ``...diffusion.common``, +# ``...diffusion.flux``) directly. +def __getattr__(name): + """Lazy import for the Flux model and config.""" + if name == "Flux": + from primus.backends.megatron.core.models.diffusion.flux import Flux + + return Flux + elif name == "FluxConfig": + from primus.backends.megatron.core.models.diffusion.flux import FluxConfig + + return FluxConfig + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") + + +__all__ = [ + "Flux", + "FluxConfig", +] diff --git a/primus/backends/megatron/core/models/diffusion/common/__init__.py b/primus/backends/megatron/core/models/diffusion/common/__init__.py new file mode 100644 index 000000000..090f51849 --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/common/__init__.py @@ -0,0 +1,38 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Common components for diffusion models in Primus. + +This module provides shared components used across different diffusion model +architectures (Flux, DiT, MovieGen, etc.). +""" + +from primus.backends.megatron.core.models.diffusion.common.config import ( + BaseDiffusionConfig, +) +from primus.backends.megatron.core.models.diffusion.common.embeddings import ( + MLPEmbedder, + TimeStepEmbedder, + Timesteps, + get_timestep_embedding, +) +from primus.backends.megatron.core.models.diffusion.common.normalization import ( + AdaLN, + AdaLNContinuous, + RMSNorm, +) + +__all__ = [ + # Base classes + "BaseDiffusionConfig", + # Embeddings + "TimeStepEmbedder", + "MLPEmbedder", + "Timesteps", + "get_timestep_embedding", + # Normalization + "RMSNorm", + "AdaLN", + "AdaLNContinuous", +] diff --git a/primus/backends/megatron/core/models/diffusion/common/config.py b/primus/backends/megatron/core/models/diffusion/common/config.py new file mode 100644 index 000000000..de4074744 --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/common/config.py @@ -0,0 +1,190 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Configuration classes for diffusion models. + +This module defines configuration dataclasses that extend Megatron-Core's +TransformerConfig to include diffusion-specific parameters. +""" + +from dataclasses import dataclass +from typing import Optional + +from megatron.core.enums import Fp8Recipe +from megatron.core.transformer.transformer_config import TransformerConfig + + +@dataclass +class BaseDiffusionConfig(TransformerConfig): + """ + Base configuration for all diffusion models in Primus. + + This class extends Megatron-Core's TransformerConfig to add common + diffusion model parameters. Model-specific configurations (FluxConfig, + DiTConfig, etc.) should inherit from this class. + + Attributes: + model_type: Type of diffusion model (e.g., 'flux', 'dit', 'moviegen') + in_channels: Number of input channels in latent space + out_channels: Number of output channels (default: same as in_channels) + patch_size: Patch size for patchification (if applicable) + fp8_scaling_strategy: FP8 scaling strategy for local spec provider (default: 'dynamic') + fp8_force_nt_layout: FP8 backward GEMM layout (default: False) + fp8_reduce_amax: Whether to allreduce amax across ranks (default: False) + mxfp4_backward_precision: MXFP4 backward precision, 'mxfp4' or 'fp8' (default: 'mxfp4') + mxfp4_gradient_stochastic_rounding: Stochastic rounding on gradients (default: False) + sensitive_layers_enabled: Enable sensitive layer configuration (default: False) + sensitive_layers_start: Number of sensitive layers at start (default: 0) + sensitive_layers_end: Number of sensitive layers at end (default: 0) + sensitive_layer_precision: Precision for sensitive layers (default: 'bf16') + + Inherited from TransformerConfig: + hidden_size: Hidden dimension size + num_attention_heads: Number of attention heads + ffn_hidden_size: FFN intermediate dimension + layernorm_epsilon: LayerNorm epsilon value + bf16, fp16, params_dtype: Precision settings + And many more Megatron-Core transformer parameters... + """ + + # Model identification + model_type: str = "base" + + # Input/output dimensions + in_channels: int = 64 + out_channels: Optional[int] = None # Defaults to in_channels if None + + # Patchification + patch_size: int = 1 + + # FP8 scaling strategy for local spec provider + fp8_scaling_strategy: str = "dynamic" + + # FP8 backward GEMM layout for the local spec provider (tensorwise path only). + # False (default) = native layouts (dgrad=NN, wgrad=TN), the validated 0-NaN path + # on hipBLASLt 1.3. True = forced-NT (every GEMM normalized to NT via pre-transposed + # operands); faster on some stacks but NaN-prone on hipBLASLt 1.3 (gfx950). + # Only affects ScalingGranularity.TENSORWISE; rowwise/blockwise ignore it. + fp8_force_nt_layout: bool = False + + # Whether to allreduce amax across DP/TP ranks for delayed FP8 scaling + fp8_reduce_amax: bool = False + + # MXFP4 backward precision: "mxfp4" (pure) or "fp8" (hybrid) + mxfp4_backward_precision: str = "mxfp4" + + # Stochastic rounding on MXFP4 gradients (paper Section 4.4) + mxfp4_gradient_stochastic_rounding: bool = False + + # Sensitive layer configuration (clean naming, maps to Megatron internals) + sensitive_layers_enabled: bool = False + sensitive_layers_start: int = 0 + sensitive_layers_end: int = 0 + sensitive_layer_precision: str = "bf16" # "bf16", "tw_fp8", or "mxfp8" (future) + + def __post_init__(self): + """Post-initialization processing.""" + # Pipeline parallelism is not implemented for diffusion models: the + # forward path runs embeddings/output head on every rank and does not + # relay activations between stages, so PP > 1 would silently + # miscompute. Reject it explicitly (before TransformerConfig validation) + # rather than producing wrong results. + if self.pipeline_model_parallel_size > 1: + raise ValueError( + "Diffusion models do not support pipeline parallelism; " + f"got pipeline_model_parallel_size={self.pipeline_model_parallel_size}. " + "Set pipeline_model_parallel_size=1." + ) + + if self.sensitive_layers_enabled: + if self.num_layers <= 1: + raise ValueError( + "sensitive_layers_enabled=True requires num_layers to be set by the child config " + "BEFORE calling super().__post_init__(). Set self.num_layers in your model config's " + "__post_init__ before the super() call." + ) + if self.sensitive_layers_start + self.sensitive_layers_end <= 0: + raise ValueError("sensitive_layers_enabled=True but both start and end counts are 0") + if self.sensitive_layers_start + self.sensitive_layers_end > self.num_layers: + raise ValueError( + f"sensitive_layers_start ({self.sensitive_layers_start}) + " + f"sensitive_layers_end ({self.sensitive_layers_end}) exceeds " + f"num_layers ({self.num_layers})" + ) + self.first_last_layers_bf16 = True + self.num_layers_at_start_in_bf16 = self.sensitive_layers_start + self.num_layers_at_end_in_bf16 = self.sensitive_layers_end + + if self.sensitive_layers_enabled and self.sensitive_layer_precision == "tw_fp8": + _deferred_fp8 = "e4m3" if self.fp8 is None else None + _deferred_fp8_recipe = ( + Fp8Recipe.tensorwise + if self.fp8_recipe is None or self.fp8_recipe == Fp8Recipe.delayed + else None + ) + else: + _deferred_fp8 = None + _deferred_fp8_recipe = None + + super().__post_init__() + + # Apply deferred FP8 settings for sensitive layers (set after super to + # avoid Megatron's "fp4 and fp8 cannot coexist" validation). + if _deferred_fp8 is not None: + self.fp8 = _deferred_fp8 + if _deferred_fp8_recipe is not None: + self.fp8_recipe = _deferred_fp8_recipe + + # Re-run the FP8 validations that Megatron skipped because self.fp8 was + # None during super().__post_init__() (TransformerConfig lines 988-1017). + if self.fp8 and self.sensitive_layers_enabled: + if self.first_last_layers_bf16 and self.fp8_recipe == Fp8Recipe.delayed: + raise ValueError("Delayed scaling does not support first / last layer in BF16.") + max_bf16 = self.num_layers // self.pipeline_model_parallel_size + if self.first_last_layers_bf16: + if not (0 <= self.num_layers_at_start_in_bf16 <= max_bf16): + raise ValueError( + f"num_layers_at_start_in_bf16 ({self.num_layers_at_start_in_bf16}) " + f"must be between 0 and {max_bf16}." + ) + if not (0 <= self.num_layers_at_end_in_bf16 <= max_bf16): + raise ValueError( + f"num_layers_at_end_in_bf16 ({self.num_layers_at_end_in_bf16}) " + f"must be between 0 and {max_bf16}." + ) + + if self.out_channels is None: + self.out_channels = self.in_channels + + # Run configuration validation on construction. (Subclass fields used by + # validate() are plain dataclass fields, so they are already populated.) + self.validate() + + def validate(self): + """ + Validate configuration parameters. + + Raises: + ValueError: If configuration is invalid + """ + if self.in_channels <= 0: + raise ValueError(f"in_channels must be positive, got {self.in_channels}") + + if self.out_channels <= 0: + raise ValueError(f"out_channels must be positive, got {self.out_channels}") + + if self.patch_size <= 0: + raise ValueError(f"patch_size must be positive, got {self.patch_size}") + + if self.hidden_size <= 0: + raise ValueError(f"hidden_size must be positive, got {self.hidden_size}") + + if self.num_attention_heads <= 0: + raise ValueError(f"num_attention_heads must be positive, got {self.num_attention_heads}") + + if self.hidden_size % self.num_attention_heads != 0: + raise ValueError( + f"hidden_size ({self.hidden_size}) must be divisible by " + f"num_attention_heads ({self.num_attention_heads})" + ) diff --git a/primus/backends/megatron/core/models/diffusion/common/embeddings.py b/primus/backends/megatron/core/models/diffusion/common/embeddings.py new file mode 100644 index 000000000..8ef927730 --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/common/embeddings.py @@ -0,0 +1,296 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Embedding layers for diffusion models. + +Provides timestep and vector conditioning embeddings used across +Flux and other diffusion architectures. + +This module implements: + - TimeStepEmbedder: Sinusoidal timestep embedding with MLP projection + - MLPEmbedder: Simple 2-layer MLP for vector conditioning (e.g., CLIP pooled) + - Timesteps: Helper module for sinusoidal timestep encoding + - get_timestep_embedding: Standalone function for sinusoidal embeddings + +Reference: + - Based on Denoising Diffusion Probabilistic Models (Ho et al., 2020) +""" + +import math + +import torch +import torch.nn as nn +from torch import Tensor + + +class TimeStepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations using sinusoidal encoding. + + This module converts scalar timestep values into high-dimensional embeddings + using sinusoidal position encoding (similar to transformers), followed by + an MLP projection. + + Architecture: + 1. Sinusoidal position encoding (creates periodic features) + 2. Linear projection to hidden_dim + 3. SiLU activation + 4. Linear projection to hidden_dim + + Args: + embedding_dim: Number of channels for sinusoidal encoding (typically 256) + hidden_dim: Hidden dimension for MLP projection (model hidden size, e.g., 3072) + flip_sin_to_cos: Whether to flip sine and cosine order (default: True) + downscale_freq_shift: Controls delta between frequencies (default: 0) + scale: Scaling factor for timesteps (default: 1.0) + max_period: Maximum period for sinusoidal encoding (default: 10000) + + Input: + t: Timestep scalars [B] or [B, 1], typically in range [0, 1000] + + Output: + Timestep embeddings [B, hidden_dim] + + Example: + >>> embedder = TimeStepEmbedder(embedding_dim=256, hidden_dim=3072) + >>> timesteps = torch.randn(4) # Batch of 4 timesteps + >>> t_emb = embedder(timesteps) + >>> assert t_emb.shape == (4, 3072) + + Reference: + - Denoising Diffusion Probabilistic Models (Ho et al., 2020) + - Adapted from NeMo's Flux implementation + """ + + def __init__( + self, + embedding_dim: int, + hidden_dim: int, + flip_sin_to_cos: bool = True, + downscale_freq_shift: float = 0, + scale: float = 1.0, + max_period: int = 10000, + ): + super().__init__() + self.embedding_dim = embedding_dim + self.hidden_dim = hidden_dim + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + self.max_period = max_period + + # MLP for projecting sinusoidal embeddings to hidden_dim + self.time_proj = Timesteps( + embedding_dim=embedding_dim, + flip_sin_to_cos=flip_sin_to_cos, + downscale_freq_shift=downscale_freq_shift, + scale=scale, + max_period=max_period, + ) + + self.time_embedding = MLPEmbedder(in_dim=embedding_dim, hidden_dim=hidden_dim) + + def forward(self, t: Tensor) -> Tensor: + """ + Forward pass: Convert timesteps to embeddings. + + Args: + t: Timesteps [B] or [B, 1], typically in range [0, 1000] + + Returns: + Timestep embeddings [B, hidden_dim] + """ + # Get sinusoidal embeddings + t_emb = self.time_proj(t) # [B, embedding_dim] + + # Project through MLP + t_emb = self.time_embedding(t_emb) # [B, hidden_dim] + + return t_emb + + +class Timesteps(nn.Module): + """ + Converts timesteps to sinusoidal embeddings. + + This is a helper module that creates sinusoidal position encodings + from scalar timestep values. + + Args: + embedding_dim: Dimension of output embeddings + flip_sin_to_cos: Whether to order as [cos, sin] instead of [sin, cos] + downscale_freq_shift: Frequency shift parameter + scale: Scaling factor for embeddings + max_period: Maximum period for sinusoidal functions + + Input: + timesteps: Scalar timesteps [B] + + Output: + Sinusoidal embeddings [B, embedding_dim] + """ + + def __init__( + self, + embedding_dim: int, + flip_sin_to_cos: bool = True, + downscale_freq_shift: float = 0, + scale: float = 1.0, + max_period: int = 10000, + ): + super().__init__() + self.embedding_dim = embedding_dim + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + self.scale = scale + self.max_period = max_period + + def forward(self, timesteps: Tensor) -> Tensor: + """ + Create sinusoidal timestep embeddings. + + Args: + timesteps: Timesteps [B] + + Returns: + Sinusoidal embeddings [B, embedding_dim] + """ + t_emb = get_timestep_embedding( + timesteps, + self.embedding_dim, + flip_sin_to_cos=self.flip_sin_to_cos, + downscale_freq_shift=self.downscale_freq_shift, + scale=self.scale, + max_period=self.max_period, + ) + return t_emb + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + flip_sin_to_cos: bool = True, + downscale_freq_shift: float = 0, + scale: float = 1.0, + max_period: int = 10000, +) -> torch.Tensor: + """ + Create sinusoidal timestep embeddings. + + This matches the implementation in Denoising Diffusion Probabilistic Models. + It creates position encodings using sine and cosine functions at different + frequencies. + + Args: + timesteps: 1-D Tensor of N indices, one per batch element [B] + embedding_dim: Dimension of the output embeddings + flip_sin_to_cos: Whether embedding order should be [cos, sin] (True) or [sin, cos] (False) + downscale_freq_shift: Controls delta between frequencies + scale: Scaling factor applied to embeddings + max_period: Controls maximum frequency of embeddings + + Returns: + Tensor of positional embeddings [B, embedding_dim] + + Reference: + - Denoising Diffusion Probabilistic Models (Ho et al., 2020) + - Adapted from NeMo's Flux implementation + """ + if len(timesteps.shape) != 1: + raise ValueError("Timesteps should be a 1d-array") + + half_dim = embedding_dim // 2 + + # Remember input dtype to preserve it + input_dtype = timesteps.dtype + + # Compute frequencies + exponent = -math.log(max_period) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timesteps.device + ) + exponent = exponent / (half_dim - downscale_freq_shift) + + emb = torch.exp(exponent) + emb = timesteps[:, None].float() * emb[None, :] + + # Scale embeddings + emb = scale * emb + + # Concatenate sine and cosine embeddings + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + + # Flip sine and cosine embeddings if requested + if flip_sin_to_cos: + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + + # Zero pad if embedding_dim is odd + if embedding_dim % 2 == 1: + emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) + + # Convert back to input dtype + emb = emb.to(input_dtype) + + return emb + + +class MLPEmbedder(nn.Module): + """ + Simple 2-layer MLP for vector conditioning (e.g., CLIP pooled embeddings). + + This module provides a learnable projection from an input vector space + to the model's hidden dimension space. It's commonly used for: + - CLIP pooled text embeddings: [B, 768] -> [B, hidden_dim] + - Other vector-level conditioning signals + + Architecture: + 1. Linear projection: in_dim -> hidden_dim + 2. SiLU activation + 3. Linear projection: hidden_dim -> hidden_dim + + Args: + in_dim: Input dimension (e.g., 768 for CLIP-L pooled embeddings) + hidden_dim: Hidden dimension (model hidden size, e.g., 3072) + + Input: + x: Vector conditioning [B, in_dim] + + Output: + Embedded conditioning [B, hidden_dim] + + Example: + >>> # For CLIP-L pooled embeddings + >>> embedder = MLPEmbedder(in_dim=768, hidden_dim=3072) + >>> clip_pooled = torch.randn(4, 768) + >>> embedded = embedder(clip_pooled) + >>> assert embedded.shape == (4, 3072) + + Reference: + - Adapted from NeMo's Flux implementation + """ + + def __init__(self, in_dim: int, hidden_dim: int): + super().__init__() + self.in_dim = in_dim + self.hidden_dim = hidden_dim + + # Two-layer MLP with SiLU activation + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) + + def forward(self, x: Tensor) -> Tensor: + """ + Forward pass: Project input vectors to hidden dimension. + + Args: + x: Input vectors [B, in_dim] + + Returns: + Embedded vectors [B, hidden_dim] + """ + x = self.in_layer(x) + x = self.silu(x) + x = self.out_layer(x) + return x diff --git a/primus/backends/megatron/core/models/diffusion/common/normalization.py b/primus/backends/megatron/core/models/diffusion/common/normalization.py new file mode 100644 index 000000000..c4c14edaf --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/common/normalization.py @@ -0,0 +1,1226 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Normalization layers for diffusion models. + +Provides adaptive normalization layers that condition on timesteps and other +conditioning signals. These are essential for diffusion model architectures +like Flux and DiT. + +This module implements: + - RMSNorm: Root Mean Square Layer Normalization + - AdaLN: Adaptive Layer Normalization for DiT + - AdaLNContinuous: Continuous variant of AdaLN for Flux + +Reference: + - DiT Paper: "Scalable Diffusion Models with Transformers" (Peebles & Xie, 2023) + - Flux Paper: "Flux: A Scalable Diffusion Model" +""" + +from typing import Tuple + +import torch +import torch.nn as nn +import triton +import triton.language as tl +from megatron.core.jit import jit_fuser +from megatron.core.tensor_parallel.layers import ColumnParallelLinear +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.transformer_config import TransformerConfig +from torch import Tensor +from torch.library import triton_op, wrap_triton + +# --------------------------------------------------------------------------- +# Opaque LayerNorm custom op — prevents Inductor from decomposing +# native_layer_norm into a Triton Welford reduction (which uses a different +# FP32 accumulation order than the eager CUDA kernel, causing numerical +# divergence amplified by FP8 quantisation). +# +# Unlike @torch.compiler.disable, this stays inside the compiled graph +# (no graph break) while remaining opaque to fusion/decomposition. +# --------------------------------------------------------------------------- + +_custom_op = torch.library.custom_op + + +# --------------------------------------------------------------------------- +# Opaque modulate custom op — prevents Inductor from fusing the +# x * (1 + scale) + shift modulation with surrounding ops. +# +# Uses hand-written Triton kernels that fuse all pointwise ops into single +# kernel launches. The backward kernel accumulates dscale/dshift in a +# deterministic sequential loop over the sequence dimension (float32 +# registers), avoiding the non-deterministic parallel-reduction order that +# Inductor's auto-generated Triton kernels would use. +# --------------------------------------------------------------------------- + +_MODULATE_BLOCK_H = 1024 + + +@triton.jit +def _modulate_fwd_kernel( + X_ptr, + Scale_ptr, + Shift_ptr, + Out_ptr, + S, + B, + H, + stride_x_s, + stride_x_b, + stride_sc_b, + BLOCK_H: tl.constexpr, +): + h_block = tl.program_id(0) + sb_idx = tl.program_id(1) + s_idx = sb_idx // B + b_idx = sb_idx % B + + offs_h = h_block * BLOCK_H + tl.arange(0, BLOCK_H) + mask = offs_h < H + + x_off = s_idx * stride_x_s + b_idx * stride_x_b + offs_h + sc_off = b_idx * stride_sc_b + offs_h + + x = tl.load(X_ptr + x_off, mask=mask) + sc = tl.load(Scale_ptr + sc_off, mask=mask) + sh = tl.load(Shift_ptr + sc_off, mask=mask) + out = x * (1.0 + sc) + sh + tl.store(Out_ptr + x_off, out, mask=mask) + + +@triton.jit +def _modulate_bwd_kernel( + Grad_ptr, + X_ptr, + Scale_ptr, + DX_ptr, + DScale_ptr, + DShift_ptr, + S, + B, + H, + stride_g_s, + stride_g_b, + stride_x_s, + stride_x_b, + stride_sc_b, + BLOCK_H: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + h_block = tl.program_id(0) + b_idx = tl.program_id(1) + + offs_h = h_block * BLOCK_H + tl.arange(0, BLOCK_H) + mask = offs_h < H + + sc_off = b_idx * stride_sc_b + offs_h + scale_val = tl.load(Scale_ptr + sc_off, mask=mask).to(tl.float32) + one_plus_scale = 1.0 + scale_val + + acc_dscale = tl.zeros([BLOCK_H], dtype=tl.float32) + acc_dshift = tl.zeros([BLOCK_H], dtype=tl.float32) + + for s_idx in range(S): + g_off = s_idx * stride_g_s + b_idx * stride_g_b + offs_h + x_off = s_idx * stride_x_s + b_idx * stride_x_b + offs_h + g = tl.load(Grad_ptr + g_off, mask=mask).to(tl.float32) + x = tl.load(X_ptr + x_off, mask=mask).to(tl.float32) + dx = g * one_plus_scale + tl.store(DX_ptr + x_off, dx.to(OUT_DTYPE), mask=mask) + acc_dscale += g * x + acc_dshift += g + + tl.store(DScale_ptr + sc_off, acc_dscale.to(OUT_DTYPE), mask=mask) + tl.store(DShift_ptr + sc_off, acc_dshift.to(OUT_DTYPE), mask=mask) + + +_TORCH_TO_TRITON_DTYPE = { + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, + torch.float32: tl.float32, +} + + +@_custom_op("primus::modulate", mutates_args=(), device_types="cuda") +def _opaque_modulate( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + if x.dim() != 3 or scale.dim() != 2: + return x * (1 + scale) + shift + x, scale, shift = x.contiguous(), scale.contiguous(), shift.contiguous() + out = torch.empty_like(x) + S, B, H = x.shape + BLOCK_H = min(triton.next_power_of_2(H), _MODULATE_BLOCK_H) + num_h_blocks = triton.cdiv(H, BLOCK_H) + _modulate_fwd_kernel[(num_h_blocks, S * B)]( + x, + scale, + shift, + out, + S, + B, + H, + x.stride(0), + x.stride(1), + scale.stride(0), + BLOCK_H=BLOCK_H, + ) + return out + + +@_opaque_modulate.register_fake +def _opaque_modulate_fake(x, scale, shift): + return torch.empty_like(x) + + +@_custom_op("primus::modulate_backward", mutates_args=(), device_types="cuda") +def _opaque_modulate_backward_op( + grad_output: torch.Tensor, + x: torch.Tensor, + scale: torch.Tensor, + need_dx: bool, + need_dscale: bool, + need_dshift: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if x.dim() != 3 or scale.dim() != 2: + grad_output = grad_output.to(x.dtype) + dx = grad_output * (1 + scale) if need_dx else torch.empty(0, device=x.device, dtype=x.dtype) + reduce_dims = list(range(grad_output.dim() - scale.dim())) + dscale = ( + (grad_output * x).sum(dim=reduce_dims) + if need_dscale + else torch.empty(0, device=x.device, dtype=x.dtype) + ) + dshift = ( + grad_output.sum(dim=reduce_dims) + if need_dshift + else torch.empty(0, device=x.device, dtype=x.dtype) + ) + return dx, dscale, dshift + + grad_output = grad_output.to(x.dtype).contiguous() + x = x.contiguous() + scale = scale.contiguous() + S, B, H = x.shape + dx = torch.empty_like(x) + dscale = torch.empty_like(scale) + dshift = torch.empty_like(scale) + BLOCK_H = min(triton.next_power_of_2(H), _MODULATE_BLOCK_H) + num_h_blocks = triton.cdiv(H, BLOCK_H) + out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype] + _modulate_bwd_kernel[(num_h_blocks, B)]( + grad_output, + x, + scale, + dx, + dscale, + dshift, + S, + B, + H, + grad_output.stride(0), + grad_output.stride(1), + x.stride(0), + x.stride(1), + scale.stride(0), + BLOCK_H=BLOCK_H, + OUT_DTYPE=out_dtype, + ) + if not need_dx: + dx = torch.empty(0, device=x.device, dtype=x.dtype) + if not need_dscale: + dscale = torch.empty(0, device=x.device, dtype=x.dtype) + if not need_dshift: + dshift = torch.empty(0, device=x.device, dtype=x.dtype) + return dx, dscale, dshift + + +@_opaque_modulate_backward_op.register_fake +def _opaque_modulate_backward_fake(grad_output, x, scale, need_dx, need_dscale, need_dshift): + dx = torch.empty_like(x) if need_dx else torch.empty(0, device=x.device, dtype=x.dtype) + dscale = torch.empty_like(scale) if need_dscale else torch.empty(0, device=x.device, dtype=x.dtype) + dshift = torch.empty_like(scale) if need_dshift else torch.empty(0, device=x.device, dtype=x.dtype) + return dx, dscale, dshift + + +def _opaque_modulate_setup_context(ctx, inputs, output): + x, scale, _shift = inputs + ctx.save_for_backward(x, scale) + + +def _opaque_modulate_backward(ctx, grad_output): + x, scale = ctx.saved_tensors + dx, dscale, dshift = _opaque_modulate_backward_op( + grad_output, + x, + scale, + True, + True, + True, + ) + return dx, dscale, dshift + + +_opaque_modulate.register_autograd( + _opaque_modulate_backward, + setup_context=_opaque_modulate_setup_context, +) + + +# --------------------------------------------------------------------------- +# @triton_op versions — transparent to Inductor for cross-op fusion. +# Enabled via config.use_triton_ops=True. +# --------------------------------------------------------------------------- + + +@triton_op("primus::modulate_v2", mutates_args=()) +def _triton_modulate( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, +) -> torch.Tensor: + if x.dim() != 3 or scale.dim() != 2: + return x * (1 + scale) + shift + x, scale, shift = x.contiguous(), scale.contiguous(), shift.contiguous() + out = torch.empty_like(x) + S, B, H = x.shape + BLOCK_H = min(triton.next_power_of_2(H), _MODULATE_BLOCK_H) + num_h_blocks = triton.cdiv(H, BLOCK_H) + wrap_triton(_modulate_fwd_kernel)[(num_h_blocks, S * B)]( + x, + scale, + shift, + out, + S, + B, + H, + x.stride(0), + x.stride(1), + scale.stride(0), + BLOCK_H=BLOCK_H, + ) + return out + + +@_triton_modulate.register_fake +def _triton_modulate_fake(x, scale, shift): + return torch.empty_like(x) + + +@triton_op("primus::modulate_backward_v2", mutates_args=()) +def _triton_modulate_backward_op( + grad_output: torch.Tensor, + x: torch.Tensor, + scale: torch.Tensor, + need_dx: bool, + need_dscale: bool, + need_dshift: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if x.dim() != 3 or scale.dim() != 2: + grad_output = grad_output.to(x.dtype) + dx = grad_output * (1 + scale) if need_dx else torch.empty(0, device=x.device, dtype=x.dtype) + reduce_dims = list(range(grad_output.dim() - scale.dim())) + dscale = ( + (grad_output * x).sum(dim=reduce_dims) + if need_dscale + else torch.empty(0, device=x.device, dtype=x.dtype) + ) + dshift = ( + grad_output.sum(dim=reduce_dims) + if need_dshift + else torch.empty(0, device=x.device, dtype=x.dtype) + ) + return dx, dscale, dshift + + grad_output = grad_output.to(x.dtype).contiguous() + x = x.contiguous() + scale = scale.contiguous() + S, B, H = x.shape + dx = torch.empty_like(x) + dscale = torch.empty_like(scale) + dshift = torch.empty_like(scale) + BLOCK_H = min(triton.next_power_of_2(H), _MODULATE_BLOCK_H) + num_h_blocks = triton.cdiv(H, BLOCK_H) + out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype] + wrap_triton(_modulate_bwd_kernel)[(num_h_blocks, B)]( + grad_output, + x, + scale, + dx, + dscale, + dshift, + S, + B, + H, + grad_output.stride(0), + grad_output.stride(1), + x.stride(0), + x.stride(1), + scale.stride(0), + BLOCK_H=BLOCK_H, + OUT_DTYPE=out_dtype, + ) + if not need_dx: + dx = torch.empty(0, device=x.device, dtype=x.dtype) + if not need_dscale: + dscale = torch.empty(0, device=x.device, dtype=x.dtype) + if not need_dshift: + dshift = torch.empty(0, device=x.device, dtype=x.dtype) + return dx, dscale, dshift + + +@_triton_modulate_backward_op.register_fake +def _triton_modulate_backward_fake(grad_output, x, scale, need_dx, need_dscale, need_dshift): + dx = torch.empty_like(x) if need_dx else torch.empty(0, device=x.device, dtype=x.dtype) + dscale = torch.empty_like(scale) if need_dscale else torch.empty(0, device=x.device, dtype=x.dtype) + dshift = torch.empty_like(scale) if need_dshift else torch.empty(0, device=x.device, dtype=x.dtype) + return dx, dscale, dshift + + +def _triton_modulate_setup_context(ctx, inputs, output): + x, scale, _shift = inputs + ctx.save_for_backward(x, scale) + + +def _triton_modulate_backward(ctx, grad_output): + x, scale = ctx.saved_tensors + dx, dscale, dshift = _triton_modulate_backward_op( + grad_output, + x, + scale, + True, + True, + True, + ) + return dx, dscale, dshift + + +_triton_modulate.register_autograd( + _triton_modulate_backward, + setup_context=_triton_modulate_setup_context, +) + + +# --------------------------------------------------------------------------- +# Fused LayerNorm + Modulate Triton kernel — computes +# (x - mean) / sqrt(var + eps) * (1 + scale) + shift +# in a single kernel launch, eliminating the intermediate ln_out tensor +# from DRAM. Only works for elementwise_affine=False LayerNorm (no +# learnable weight/bias), which is the case for all AdaLN variants in Flux. +# +# The backward kernel fuses the LN backward formula with the modulate +# backward, using deterministic sequential accumulation over S for +# d_scale / d_shift (same pattern as _modulate_bwd_kernel). +# --------------------------------------------------------------------------- + +_FUSED_LN_MOD_MAX_H = 8192 + + +@triton.jit +def _fused_ln_modulate_fwd_kernel( + X_ptr, + Scale_ptr, + Shift_ptr, + Out_ptr, + Mean_ptr, + Rstd_ptr, + S, + B, + H, + eps, + stride_x_sb, + stride_sc_b, + BLOCK_H: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + row = tl.program_id(0) + b_idx = row % B + offs_h = tl.arange(0, BLOCK_H) + mask = offs_h < H + + x_off = row * stride_x_sb + offs_h + sc_off = b_idx * stride_sc_b + offs_h + + x = tl.load(X_ptr + x_off, mask=mask, other=0.0).to(tl.float32) + mean = tl.sum(x, axis=0) / H + + x_centered = tl.where(mask, x - mean, 0.0) + var = tl.sum(x_centered * x_centered, axis=0) / H + rstd = 1.0 / tl.sqrt(var + eps) + x_hat = x_centered * rstd + + sc = tl.load(Scale_ptr + sc_off, mask=mask).to(tl.float32) + sh = tl.load(Shift_ptr + sc_off, mask=mask).to(tl.float32) + out = x_hat * (1.0 + sc) + sh + + tl.store(Out_ptr + x_off, out.to(OUT_DTYPE), mask=mask) + tl.store(Mean_ptr + row, mean) + tl.store(Rstd_ptr + row, rstd) + + +@triton.jit +def _fused_ln_modulate_bwd_dscale_dshift_kernel( + Grad_ptr, + X_ptr, + Mean_ptr, + Rstd_ptr, + DScale_ptr, + DShift_ptr, + S, + B, + H, + stride_g_sb, + stride_x_sb, + XBLOCK: tl.constexpr, + RBLOCK: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK)[:, None] + xmask = xindex < (B * H) + x_bh = xindex + x_b = xindex // H + + r_base = tl.arange(0, RBLOCK)[None, :] + + acc_dscale = tl.full([XBLOCK, RBLOCK], 0, tl.float32) + acc_dshift = tl.full([XBLOCK, RBLOCK], 0, tl.float32) + + for r_off in range(0, S, RBLOCK): + r_idx = r_off + r_base + r_mask = r_idx < S + mask = xmask & r_mask + + g_off = x_bh + (r_idx * B * H) + x_off = x_bh + (r_idx * B * H) + row = x_b + r_idx * B + + g = tl.load(Grad_ptr + g_off, mask=mask, other=0.0).to(tl.float32) + x = tl.load(X_ptr + x_off, mask=mask, other=0.0).to(tl.float32) + mean_val = tl.load(Mean_ptr + row, mask=mask, other=0.0) + rstd_val = tl.load(Rstd_ptr + row, mask=mask, other=0.0) + + x_hat = (x - mean_val) * rstd_val + acc_dscale = tl.where(mask, acc_dscale + g * x_hat, acc_dscale) + acc_dshift = tl.where(mask, acc_dshift + g, acc_dshift) + + dscale_val = tl.sum(acc_dscale, 1)[:, None] + dshift_val = tl.sum(acc_dshift, 1)[:, None] + tl.store(DScale_ptr + x_bh, dscale_val.to(OUT_DTYPE), mask=xmask) + tl.store(DShift_ptr + x_bh, dshift_val.to(OUT_DTYPE), mask=xmask) + + +@triton.jit +def _fused_ln_modulate_bwd_dx_kernel( + Grad_ptr, + X_ptr, + Mean_ptr, + Rstd_ptr, + Scale_ptr, + DX_ptr, + S, + B, + H, + stride_g_sb, + stride_x_sb, + stride_sc_b, + BLOCK_H: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + sb_idx = tl.program_id(0) + b_idx = sb_idx % B + offs_h = tl.arange(0, BLOCK_H) + mask = offs_h < H + + sc_off = b_idx * stride_sc_b + offs_h + scale_val = tl.load(Scale_ptr + sc_off, mask=mask).to(tl.float32) + one_plus_scale = 1.0 + scale_val + + g_off = sb_idx * stride_g_sb + offs_h + x_off = sb_idx * stride_x_sb + offs_h + + g = tl.load(Grad_ptr + g_off, mask=mask).to(tl.float32) + x = tl.load(X_ptr + x_off, mask=mask).to(tl.float32) + mean = tl.load(Mean_ptr + sb_idx) + rstd = tl.load(Rstd_ptr + sb_idx) + + x_hat = (x - mean) * rstd + d_x_hat = g * one_plus_scale + + inv_H = 1.0 / H + d_x_hat_masked = tl.where(mask, d_x_hat, 0.0) + xhat_masked = tl.where(mask, x_hat, 0.0) + c1 = tl.sum(xhat_masked * d_x_hat_masked, axis=0) * inv_H + c2 = tl.sum(d_x_hat_masked, axis=0) * inv_H + d_x = rstd * (d_x_hat - c2 - x_hat * c1) + + tl.store(DX_ptr + x_off, d_x.to(OUT_DTYPE), mask=mask) + + +@_custom_op("primus::fused_ln_modulate", mutates_args=(), device_types="cuda") +def _opaque_fused_ln_modulate( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if x.dim() != 3 or scale.dim() != 2 or x.shape[-1] > _FUSED_LN_MOD_MAX_H: + out, mean, rstd = torch.ops.aten.native_layer_norm( + x, + [x.shape[-1]], + None, + None, + eps, + ) + out = out * (1 + scale) + shift + return out, mean.flatten(), rstd.flatten() + + x = x.contiguous() + scale, shift = scale.contiguous(), shift.contiguous() + S, B, H = x.shape + out = torch.empty_like(x) + M = S * B + mean = torch.empty(M, dtype=torch.float32, device=x.device) + rstd = torch.empty(M, dtype=torch.float32, device=x.device) + BLOCK_H = triton.next_power_of_2(H) + out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype] + _fused_ln_modulate_fwd_kernel[(M,)]( + x, + scale, + shift, + out, + mean, + rstd, + S, + B, + H, + eps, + x.stride(1), + scale.stride(0), + BLOCK_H=BLOCK_H, + OUT_DTYPE=out_dtype, + ) + return out, mean, rstd + + +@_opaque_fused_ln_modulate.register_fake +def _opaque_fused_ln_modulate_fake(x, scale, shift, eps): + out = torch.empty_like(x) + M = x.numel() // x.shape[-1] + mean = torch.empty(M, dtype=torch.float32, device=x.device) + rstd = torch.empty(M, dtype=torch.float32, device=x.device) + return out, mean, rstd + + +@_custom_op("primus::fused_ln_modulate_backward", mutates_args=(), device_types="cuda") +def _opaque_fused_ln_modulate_backward_op( + grad_output: torch.Tensor, + x: torch.Tensor, + mean: torch.Tensor, + rstd: torch.Tensor, + scale: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if x.dim() != 3 or scale.dim() != 2 or x.shape[-1] > _FUSED_LN_MOD_MAX_H: + grad_output = grad_output.to(x.dtype) + mean_r = mean.view(*x.shape[:-1], 1) + rstd_r = rstd.view(*x.shape[:-1], 1) + x_hat = (x.float() - mean_r) * rstd_r + d_x_hat = grad_output.float() * (1 + scale.float()) + c1 = (x_hat * d_x_hat).mean(dim=-1, keepdim=True) + c2 = d_x_hat.mean(dim=-1, keepdim=True) + d_x = (rstd_r * (d_x_hat - c2 - x_hat * c1)).to(x.dtype) + reduce_dims = list(range(grad_output.dim() - scale.dim())) + dscale = (grad_output.float() * x_hat).sum(dim=reduce_dims).to(scale.dtype) + dshift = grad_output.sum(dim=reduce_dims).to(scale.dtype) + return d_x, dscale, dshift + + grad_output = grad_output.to(x.dtype).contiguous() + x = x.contiguous() + scale = scale.contiguous() + S, B, H = x.shape + dx = torch.empty_like(x) + dscale = torch.empty_like(scale) + dshift = torch.empty_like(scale) + BLOCK_H = triton.next_power_of_2(H) + out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype] + XBLOCK, RBLOCK = 256, 8 + _fused_ln_modulate_bwd_dscale_dshift_kernel[(triton.cdiv(B * H, XBLOCK),)]( + grad_output, + x, + mean, + rstd, + dscale, + dshift, + S, + B, + H, + grad_output.stride(1), + x.stride(1), + XBLOCK=XBLOCK, + RBLOCK=RBLOCK, + OUT_DTYPE=out_dtype, + ) + _fused_ln_modulate_bwd_dx_kernel[(S * B,)]( + grad_output, + x, + mean, + rstd, + scale, + dx, + S, + B, + H, + grad_output.stride(1), + x.stride(1), + scale.stride(0), + BLOCK_H=BLOCK_H, + OUT_DTYPE=out_dtype, + ) + return dx, dscale, dshift + + +@_opaque_fused_ln_modulate_backward_op.register_fake +def _opaque_fused_ln_modulate_backward_fake(grad_output, x, mean, rstd, scale): + dx = torch.empty_like(x) + dscale = torch.empty_like(scale) + dshift = torch.empty_like(scale) + return dx, dscale, dshift + + +def _opaque_fused_ln_modulate_setup_context(ctx, inputs, output): + x, scale, _shift, _eps = inputs + _out, mean, rstd = output + ctx.save_for_backward(x, mean, rstd, scale) + + +def _opaque_fused_ln_modulate_backward(ctx, grad_output, _grad_mean, _grad_rstd): + x, mean, rstd, scale = ctx.saved_tensors + dx, dscale, dshift = _opaque_fused_ln_modulate_backward_op( + grad_output, + x, + mean, + rstd, + scale, + ) + return dx, dscale, dshift, None + + +_opaque_fused_ln_modulate.register_autograd( + _opaque_fused_ln_modulate_backward, + setup_context=_opaque_fused_ln_modulate_setup_context, +) + + +# --------------------------------------------------------------------------- +# @triton_op fused LN+modulate — transparent to Inductor. +# --------------------------------------------------------------------------- + + +@triton_op("primus::fused_ln_modulate_v2", mutates_args=()) +def _triton_fused_ln_modulate( + x: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if x.dim() != 3 or scale.dim() != 2 or x.shape[-1] > _FUSED_LN_MOD_MAX_H: + out, mean, rstd = torch.ops.aten.native_layer_norm( + x, + [x.shape[-1]], + None, + None, + eps, + ) + out = out * (1 + scale) + shift + return out, mean.flatten(), rstd.flatten() + + x = x.contiguous() + scale, shift = scale.contiguous(), shift.contiguous() + S, B, H = x.shape + out = torch.empty_like(x) + M = S * B + mean = torch.empty(M, dtype=torch.float32, device=x.device) + rstd = torch.empty(M, dtype=torch.float32, device=x.device) + BLOCK_H = triton.next_power_of_2(H) + out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype] + wrap_triton(_fused_ln_modulate_fwd_kernel)[(M,)]( + x, + scale, + shift, + out, + mean, + rstd, + S, + B, + H, + eps, + x.stride(1), + scale.stride(0), + BLOCK_H=BLOCK_H, + OUT_DTYPE=out_dtype, + ) + return out, mean, rstd + + +@_triton_fused_ln_modulate.register_fake +def _triton_fused_ln_modulate_fake(x, scale, shift, eps): + out = torch.empty_like(x) + M = x.numel() // x.shape[-1] + mean = torch.empty(M, dtype=torch.float32, device=x.device) + rstd = torch.empty(M, dtype=torch.float32, device=x.device) + return out, mean, rstd + + +@triton_op("primus::fused_ln_modulate_backward_v2", mutates_args=()) +def _triton_fused_ln_modulate_backward_op( + grad_output: torch.Tensor, + x: torch.Tensor, + mean: torch.Tensor, + rstd: torch.Tensor, + scale: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if x.dim() != 3 or scale.dim() != 2 or x.shape[-1] > _FUSED_LN_MOD_MAX_H: + grad_output = grad_output.to(x.dtype) + mean_r = mean.view(*x.shape[:-1], 1) + rstd_r = rstd.view(*x.shape[:-1], 1) + x_hat = (x.float() - mean_r) * rstd_r + d_x_hat = grad_output.float() * (1 + scale.float()) + c1 = (x_hat * d_x_hat).mean(dim=-1, keepdim=True) + c2 = d_x_hat.mean(dim=-1, keepdim=True) + d_x = (rstd_r * (d_x_hat - c2 - x_hat * c1)).to(x.dtype) + reduce_dims = list(range(grad_output.dim() - scale.dim())) + dscale = (grad_output.float() * x_hat).sum(dim=reduce_dims).to(scale.dtype) + dshift = grad_output.sum(dim=reduce_dims).to(scale.dtype) + return d_x, dscale, dshift + + grad_output = grad_output.to(x.dtype).contiguous() + x = x.contiguous() + scale = scale.contiguous() + S, B, H = x.shape + dx = torch.empty_like(x) + dscale = torch.empty_like(scale) + dshift = torch.empty_like(scale) + BLOCK_H = triton.next_power_of_2(H) + out_dtype = _TORCH_TO_TRITON_DTYPE[x.dtype] + XBLOCK, RBLOCK = 256, 8 + wrap_triton(_fused_ln_modulate_bwd_dscale_dshift_kernel)[(triton.cdiv(B * H, XBLOCK),)]( + grad_output, + x, + mean, + rstd, + dscale, + dshift, + S, + B, + H, + grad_output.stride(1), + x.stride(1), + XBLOCK=XBLOCK, + RBLOCK=RBLOCK, + OUT_DTYPE=out_dtype, + ) + wrap_triton(_fused_ln_modulate_bwd_dx_kernel)[(S * B,)]( + grad_output, + x, + mean, + rstd, + scale, + dx, + S, + B, + H, + grad_output.stride(1), + x.stride(1), + scale.stride(0), + BLOCK_H=BLOCK_H, + OUT_DTYPE=out_dtype, + ) + return dx, dscale, dshift + + +@_triton_fused_ln_modulate_backward_op.register_fake +def _triton_fused_ln_modulate_backward_fake(grad_output, x, mean, rstd, scale): + dx = torch.empty_like(x) + dscale = torch.empty_like(scale) + dshift = torch.empty_like(scale) + return dx, dscale, dshift + + +def _triton_fused_ln_modulate_setup_context(ctx, inputs, output): + x, scale, _shift, _eps = inputs + _out, mean, rstd = output + ctx.save_for_backward(x, mean, rstd, scale) + + +def _triton_fused_ln_modulate_backward(ctx, grad_output, _grad_mean, _grad_rstd): + x, mean, rstd, scale = ctx.saved_tensors + dx, dscale, dshift = _triton_fused_ln_modulate_backward_op( + grad_output, + x, + mean, + rstd, + scale, + ) + return dx, dscale, dshift, None + + +_triton_fused_ln_modulate.register_autograd( + _triton_fused_ln_modulate_backward, + setup_context=_triton_fused_ln_modulate_setup_context, +) + + +def _fused_eager_ln_modulate( + norm_module: nn.Module, + x: Tensor, + scale: Tensor, + shift: Tensor, + use_triton_ops: bool = False, +) -> Tensor: + """Fused LayerNorm + modulate via a single Triton kernel. + + Computes norm(x) * (1 + scale) + shift in one pass, avoiding the + intermediate ln_out write/read to DRAM. Only valid for + elementwise_affine=False norms (no learnable weight/bias). + """ + _ln_mod_fn = _triton_fused_ln_modulate if use_triton_ops else _opaque_fused_ln_modulate + out, _mean, _rstd = _ln_mod_fn( + x, + scale, + shift, + norm_module.eps, + ) + return out + + +class RMSNorm(nn.Module): + """ + Root Mean Square Layer Normalization. + + Formula: RMSNorm(x) = (x / sqrt(mean(x^2) + eps)) * weight + + Args: + hidden_size: Size of the normalized dimension + config: Transformer configuration (for compatibility, not used) + eps: Small constant for numerical stability (default: 1e-6) + + Reference: + - "Root Mean Square Layer Normalization" (Zhang & Sennrich, 2019) + - Adapted from NeMo's DiT implementation + """ + + def __init__(self, hidden_size: int, config=None, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + + def _norm(self, x: Tensor) -> Tensor: + """ + Compute RMS normalization. + + Args: + x: Input tensor + + Returns: + Normalized tensor (before scaling) + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x: Tensor) -> Tensor: + """ + Forward pass: Normalize input. + + Args: + x: Input tensor [..., hidden_size] + + Returns: + Normalized and scaled tensor [..., hidden_size] + """ + output = self._norm(x.float()).type_as(x) + return output * self.weight + + +class AdaLN(MegatronModule): + """ + Adaptive Layer Normalization for DiT (Diffusion Transformer). + + Conditions layer normalization on timestep embeddings via learned scale, + shift, and gate parameters. Projects conditioning through SiLU + Linear + into n_adaln_chunks modulation parameters. + + Args: + config: Transformer configuration + n_adaln_chunks: Number of modulation chunks (default: 9, i.e. (shift, scale, gate) x 3) + norm: Normalization layer class (default: nn.LayerNorm) + modulation_bias: Whether to use bias in modulation projection (default: False) + use_second_norm: Whether to use a second normalization layer (default: False) + init_method: Initialization function for the modulation projection + weight (default: zero-init via ``nn.init.zeros_``). Pass + ``nn.init.normal_`` only when matching NeMo's exact RNG draw + sequence -- Flux's ``init_weights()`` re-zeroes these weights + after construction, so the normal_ draw only serves to advance + the CUDA RNG by the same amount as NeMo's init for + cross-framework convergence comparison. Other consumers + (test harnesses, downstream DiT models that reuse AdaLN) get + the same observable zero-init either way and should leave + this at the default. + + Reference: + - "Scalable Diffusion Models with Transformers" (Peebles & Xie, 2023) + - Adapted from NeMo's DiT implementation + """ + + def __init__( + self, + config: TransformerConfig, + n_adaln_chunks: int = 9, + norm: type = nn.LayerNorm, + modulation_bias: bool = False, + use_second_norm: bool = False, + init_method=nn.init.zeros_, + ): + super().__init__(config) + + # Layer normalization (without affine parameters - scale/shift come from conditioning) + self.ln = norm(config.hidden_size, elementwise_affine=False, eps=config.layernorm_epsilon) + + self.n_adaln_chunks = n_adaln_chunks + + # Modulation network: conditioning -> (scale, shift, gate, ...) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + ColumnParallelLinear( + config.hidden_size, + n_adaln_chunks * config.hidden_size, + config=config, + init_method=init_method, + bias=modulation_bias, + gather_output=True, + ), + ) + + self.use_second_norm = use_second_norm + if use_second_norm: + self.ln2 = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=1e-6) + + # Mark weight as sequence parallel if needed + setattr(self.adaLN_modulation[-1].weight, "sequence_parallel", config.sequence_parallel) + + self._adaln_plain_ops = getattr(config, "adaln_plain_ops", False) + self._use_triton_ops = getattr(config, "use_triton_ops", False) + + if self._adaln_plain_ops: + self.use_fused_ln_modulate = False + else: + self.use_fused_ln_modulate = True + + adaln_always_jit = getattr(config, "adaln_always_jit_fuser", False) + if adaln_always_jit or not getattr(config, "enable_torch_compile", False): + AdaLN._apply_jit_fuser() + + @classmethod + def _apply_jit_fuser(cls): + """Apply @jit_fuser to class methods once (shared across all instances).""" + if getattr(cls, "_jit_fuser_applied", False): + return + cls.forward = jit_fuser(cls.forward) + cls.modulate = jit_fuser(cls.modulate) + cls.scale_add = jit_fuser(cls.scale_add) + cls.modulated_layernorm = jit_fuser(cls.modulated_layernorm) + cls.scaled_modulated_layernorm = jit_fuser(cls.scaled_modulated_layernorm) + cls._jit_fuser_applied = True + + def forward(self, timestep_emb: Tensor) -> Tuple[Tensor, ...]: + """ + Generate modulation parameters from timestep embedding. + + Args: + timestep_emb: Timestep embeddings [B, hidden_size] + + Returns: + Tuple of n_adaln_chunks tensors, each [B, hidden_size] + """ + output, bias = self.adaLN_modulation(timestep_emb) + if bias is not None: + output = output + bias + return output.chunk(self.n_adaln_chunks, dim=-1) + + def modulate(self, x: Tensor, shift: Tensor, scale: Tensor) -> Tensor: + """ + Apply adaptive modulation: x * (1 + scale) + shift. + + Args: + x: Input tensor [B, ..., hidden_size] + shift: Shift parameter [B, hidden_size] + scale: Scale parameter [B, hidden_size] + + Returns: + Modulated tensor [B, ..., hidden_size] + """ + return x * (1 + scale) + shift + + def scale_add(self, residual: Tensor, x: Tensor, gate: Tensor) -> Tensor: + """ + Gated residual addition: residual + gate * x. + + Args: + residual: Residual connection [B, ..., hidden_size] + x: Input to add [B, ..., hidden_size] + gate: Gate parameter [B, hidden_size] + + Returns: + Combined tensor [B, ..., hidden_size] + """ + return residual + gate * x + + def modulated_layernorm(self, x: Tensor, shift: Tensor, scale: Tensor, layernorm_idx: int = 0) -> Tensor: + """ + Apply layer normalization followed by adaptive modulation. + + Args: + x: Input tensor [B, ..., hidden_size] + shift: Shift parameter [B, hidden_size] + scale: Scale parameter [B, hidden_size] + layernorm_idx: Which layer norm to use (0 or 1, if use_second_norm=True) + + Returns: + Normalized and modulated tensor [B, ..., hidden_size] + """ + # Select appropriate layer norm + if self.use_second_norm and layernorm_idx == 1: + layernorm = self.ln2 + else: + layernorm = self.ln + + if self._adaln_plain_ops: + input_layernorm_output = layernorm(x).type_as(x) + return self.modulate(input_layernorm_output, shift, scale) + + if self.use_fused_ln_modulate: + return _fused_eager_ln_modulate(layernorm, x, scale, shift, self._use_triton_ops) + + input_layernorm_output = layernorm(x).type_as(x) + return self.modulate(input_layernorm_output, shift, scale) + + def scaled_modulated_layernorm( + self, + residual: Tensor, + x: Tensor, + gate: Tensor, + shift: Tensor, + scale: Tensor, + layernorm_idx: int = 0, + ) -> Tuple[Tensor, Tensor]: + """ + Combined operation: gated residual addition + modulated layer norm. + + This is a common pattern in DiT: add gated residual, then normalize and modulate. + + Args: + residual: Residual connection [B, ..., hidden_size] + x: Input to add [B, ..., hidden_size] + gate: Gate parameter [B, hidden_size] + shift: Shift parameter [B, hidden_size] + scale: Scale parameter [B, hidden_size] + layernorm_idx: Which layer norm to use + + Returns: + Tuple of (hidden_states, shifted_pre_mlp_layernorm_output) + """ + # Gated residual addition + hidden_states = self.scale_add(residual, x, gate) + + # Apply modulated layer normalization + shifted_pre_mlp_layernorm_output = self.modulated_layernorm( + hidden_states, shift, scale, layernorm_idx + ) + + return hidden_states, shifted_pre_mlp_layernorm_output + + +class AdaLNContinuous(MegatronModule): + """ + Continuous Adaptive Layer Normalization for Flux. + + Simpler variant of AdaLN that only produces scale and shift (no gating). + Formula: norm(x) * (1 + scale) + shift + + Args: + config: Transformer configuration + conditioning_embedding_dim: Dimension of conditioning embedding input + modulation_bias: Whether to use bias in modulation layers (default: True) + norm_type: 'layer_norm' (default: 'layer_norm') + + Note: + Uses Megatron's sequence-first format [S, B, D]. + + Reference: + - Flux Paper: "Flux: A Scalable Diffusion Model" + - Adapted from NeMo's DiT implementation + """ + + def __init__( + self, + config: TransformerConfig, + conditioning_embedding_dim: int, + modulation_bias: bool = True, + norm_type: str = "layer_norm", + ): + super().__init__(config) + + # Modulation network: conditioning -> (scale, shift) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(conditioning_embedding_dim, config.hidden_size * 2, bias=modulation_bias), + ) + + # Normalization layer + if norm_type == "layer_norm": + self.norm = nn.LayerNorm( + config.hidden_size, elementwise_affine=False, eps=1e-6, bias=modulation_bias + ) + else: + raise ValueError(f"Unknown normalization type: {norm_type}") + + self._adaln_plain_ops = getattr(config, "adaln_plain_ops", False) + + if self._adaln_plain_ops: + self.use_fused_ln_modulate = False + else: + self.use_fused_ln_modulate = True + + self._use_triton_ops = getattr(config, "use_triton_ops", False) + + def forward(self, x: Tensor, conditioning_embedding: Tensor) -> Tensor: + """ + Forward pass: Apply continuous adaptive normalization. + + Args: + x: Input tensor [seq_len, B, hidden_size] (sequence-first format) + conditioning_embedding: Conditioning signal [B, conditioning_embedding_dim] + + Returns: + Normalized and modulated tensor [seq_len, B, hidden_size] + """ + # Generate scale and shift from conditioning + emb = self.adaLN_modulation(conditioning_embedding) # [B, 2 * hidden_size] + # NeMo convention: first half is scale, second half is shift. + # NOTE: Checkpoints trained before this fix had (shift, scale) order and + # need the two halves of norm_out.adaLN_modulation weight/bias swapped. + scale, shift = torch.chunk(emb, 2, dim=1) + + if self._adaln_plain_ops: + x = self.norm(x) * (1 + scale) + shift + return x + + if self.use_fused_ln_modulate: + return _fused_eager_ln_modulate(self.norm, x, scale, shift, self._use_triton_ops) + + ln_out = self.norm(x) + _mod_fn = _triton_modulate if self._use_triton_ops else _opaque_modulate + return _mod_fn(ln_out, scale, shift) diff --git a/primus/backends/megatron/core/transformer/diffusion_transformer_block.py b/primus/backends/megatron/core/transformer/diffusion_transformer_block.py new file mode 100644 index 000000000..9e21033ca --- /dev/null +++ b/primus/backends/megatron/core/transformer/diffusion_transformer_block.py @@ -0,0 +1,352 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +DiffusionTransformerBlock for diffusion models with timestep conditioning. + +This module provides a specialized TransformerBlock for diffusion models that properly +handles timestep embeddings and other conditioning parameters through gradient checkpointing. +""" + +from typing import Optional, Union + +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.transformer.transformer_block import TransformerBlock +from megatron.core.utils import WrappedTensor +from torch import Tensor + + +class DiffusionTransformerBlock(TransformerBlock): + """ + TransformerBlock for diffusion models with timestep conditioning. + + Extends upstream Megatron TransformerBlock with explicit timestep_emb and + guidance_emb parameters, routed through gradient checkpointing in a + thread-safe manner (no instance attribute storage for conditioning). + + Call chain: + Flux.forward() -> transformer(timestep_emb=...) -> + DiffusionTransformerBlock.forward() -> checkpoint(custom_forward, timestep_emb) -> + layer(timestep_emb=...) -> adaln(timestep_emb) + """ + + def forward( + self, + hidden_states: Union[Tensor, WrappedTensor], + attention_mask: Optional[Tensor], + context: Optional[Tensor] = None, + context_mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + rotary_pos_cos: Optional[Tensor] = None, + rotary_pos_sin: Optional[Tensor] = None, + rotary_pos_cos_sin: Optional[Tensor] = None, + attention_bias: Optional[Tensor] = None, + inference_context: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[Tensor] = None, + # Diffusion-specific conditioning parameters + timestep_emb: Optional[Tensor] = None, + guidance_emb: Optional[Tensor] = None, + **kwargs, + ): + """ + Forward pass with explicit diffusion conditioning. + + Timestep and guidance embeddings are passed as tensors and participate in + gradient checkpointing like other forward inputs. + + Args: + hidden_states: Input tensor [seq, batch, hidden] + attention_mask: Attention mask tensor + context: Context tensor for cross-attention (e.g., text embeddings) + context_mask: Mask for context in cross-attention + timestep_emb: Timestep conditioning [batch, hidden] for diffusion models + guidance_emb: Guidance conditioning [batch, hidden] for classifier-free guidance + All other args: Same as TransformerBlock.forward() + + Returns: + Output tensor [seq, batch, hidden], or a tuple (hidden_states, context) + when cross-attention context is present. + + Note: + This implementation is thread-safe because no conditioning parameters are + stored as instance attributes. All conditioning flows through function arguments. + """ + # Validate diffusion parameters if model type requires them + if hasattr(self.config, "model_type") and self.config.model_type in ["flux", "diffusion"]: + if timestep_emb is None: + raise ValueError( + f"DiffusionTransformerBlock requires timestep_emb for model_type={self.config.model_type}" + ) + + # Build conditioning_kwargs for passing to layers + # Only include non-None values to avoid passing unused parameters + conditioning_kwargs = {} + if timestep_emb is not None: + conditioning_kwargs["timestep_emb"] = timestep_emb + if guidance_emb is not None: + conditioning_kwargs["guidance_emb"] = guidance_emb + + # Extract kwargs that need to be passed to parent (excluding conditioning) + parent_kwargs = { + "hidden_states": hidden_states, + "attention_mask": attention_mask, + "context": context, + "context_mask": context_mask, + "rotary_pos_emb": rotary_pos_emb, + "rotary_pos_cos": rotary_pos_cos, + "rotary_pos_sin": rotary_pos_sin, + "rotary_pos_cos_sin": rotary_pos_cos_sin, + "attention_bias": attention_bias, + "inference_context": inference_context, + "packed_seq_params": packed_seq_params, + "sequence_len_offset": sequence_len_offset, + } + # Filter out conditioning parameters from kwargs before updating + filtered_kwargs = {k: v for k, v in kwargs.items() if k not in ("timestep_emb", "guidance_emb")} + parent_kwargs.update(filtered_kwargs) + + # Check if checkpointing is needed (based on parent's logic) + if self.config.recompute_granularity == "full" and self.training: + # Checkpointed path: Need to override _checkpointed_forward to handle conditioning + # Delete the obsolete reference to the initial input tensor if necessary + if isinstance(hidden_states, WrappedTensor): + hidden_states = hidden_states.unwrap() + + if not self.pre_process: + # See set_input_tensor() + hidden_states = self.input_tensor + + # Determine inner quantization context usage (from parent logic) + from megatron.core.enums import Fp8Recipe + + if self.config.fp8: + use_inner_quantization_context = self.config.fp8_recipe != Fp8Recipe.delayed + elif self.config.fp4: + use_inner_quantization_context = True + else: + use_inner_quantization_context = False + + # Call custom checkpointed forward with conditioning + return self._checkpointed_forward_with_conditioning( + hidden_states=hidden_states, + attention_mask=attention_mask, + context=context, + context_mask=context_mask, + rotary_pos_emb=rotary_pos_emb, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + use_inner_quantization_context=use_inner_quantization_context, + **conditioning_kwargs, # Pass conditioning to checkpoint + ) + else: + # Normal forward path - manually iterate layers with conditioning + # Parent TransformerBlock doesn't accept timestep_emb/guidance_emb + # So we need to manually process layers when conditioning is present + if conditioning_kwargs: + # Process layers manually with conditioning + current_hidden = hidden_states + current_context = context + + for layer in self.layers: + layer_kwargs = { + "hidden_states": current_hidden, + "attention_mask": attention_mask, + "context": current_context, + "context_mask": context_mask, + "rotary_pos_emb": rotary_pos_emb, + "attention_bias": attention_bias, + "inference_context": inference_context, + "packed_seq_params": packed_seq_params, + } + # Add conditioning + layer_kwargs.update(conditioning_kwargs) + + layer_output = layer(**layer_kwargs) + if isinstance(layer_output, tuple): + current_hidden, current_context = layer_output + else: + current_hidden = layer_output + current_context = None + + return current_hidden if current_context is None else (current_hidden, current_context) + else: + # No conditioning - use parent's forward + return super().forward(**parent_kwargs) + + def _checkpointed_forward_with_conditioning( + self, + hidden_states: Tensor, + attention_mask: Tensor, + context: Tensor, + context_mask: Tensor, + rotary_pos_emb: Tensor, + attention_bias: Tensor, + packed_seq_params: PackedSeqParams, + use_inner_quantization_context: bool, + timestep_emb: Optional[Tensor] = None, + guidance_emb: Optional[Tensor] = None, + ): + """ + Checkpointed forward with conditioning support. + + This method extends the parent's _checkpointed_forward to handle + diffusion conditioning parameters (timestep_emb, guidance_emb). + + Note: Unlike the non-checkpointed path in forward(), this method + only returns hidden_states and discards context. + """ + from contextlib import nullcontext + + from megatron.core import tensor_parallel + from megatron.core.fp4_utils import get_fp4_context + from megatron.core.fp8_utils import get_fp8_context + from megatron.core.transformer.transformer_layer import ( + get_transformer_layer_offset, + ) + + try: + from megatron.core.extensions.transformer_engine import te_checkpoint + except ImportError: + te_checkpoint = None + + def custom(start: int, end: int): + def custom_forward( + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb, + timestep_emb=None, + guidance_emb=None, + ): + for index in range(start, end): + layer = self._get_layer(index) + + # Get appropriate inner quantization context + if use_inner_quantization_context: + if self.config.fp8: + inner_quantization_context = get_fp8_context(self.config, layer.layer_number - 1) + elif self.config.fp4: + inner_quantization_context = get_fp4_context(self.config, layer.layer_number - 1) + else: + inner_quantization_context = nullcontext() + else: + inner_quantization_context = nullcontext() + + with inner_quantization_context: + # Build layer kwargs with conditioning + layer_kwargs = { + "hidden_states": hidden_states, + "attention_mask": attention_mask, + "context": context, + "context_mask": context_mask, + "rotary_pos_emb": rotary_pos_emb, + "attention_bias": attention_bias, + "inference_context": None, + "packed_seq_params": packed_seq_params, + } + # Add conditioning if present + if timestep_emb is not None: + layer_kwargs["timestep_emb"] = timestep_emb + if guidance_emb is not None: + layer_kwargs["guidance_emb"] = guidance_emb + + hidden_states, context = layer(**layer_kwargs) + return hidden_states, context + + return custom_forward + + def checkpoint_handler(forward_func): + """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`""" + if self.config.fp8 or self.config.fp4: + return te_checkpoint( + forward_func, + self.config.distribute_saved_activations, + tensor_parallel.random.get_cuda_rng_tracker, + self.pg_collection.tp, + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb, + ( + timestep_emb if timestep_emb is not None else context + ), # Use context as placeholder if None + ( + guidance_emb if guidance_emb is not None else context + ), # Use context as placeholder if None + ) + else: + return tensor_parallel.checkpoint( + forward_func, + self.config.distribute_saved_activations, + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb, + ( + timestep_emb if timestep_emb is not None else context + ), # Use context as placeholder if None + ( + guidance_emb if guidance_emb is not None else context + ), # Use context as placeholder if None + ) + + recompute_layer_ids = getattr(self.config, "recompute_layer_ids", None) + + if self.config.recompute_method == "uniform": + layer_idx = 0 + while layer_idx < self.num_layers_per_pipeline_rank: + hidden_states, context = checkpoint_handler( + custom(layer_idx, layer_idx + self.config.recompute_num_layers) + ) + layer_idx += self.config.recompute_num_layers + + elif self.config.recompute_method == "block": + recompute_skip_num_layers = 0 + for layer_idx in range(self.num_layers_per_pipeline_rank): + if (self.config.fp8 or self.config.fp4) and not hidden_states.requires_grad: + recompute_skip_num_layers += 1 + if ( + layer_idx >= recompute_skip_num_layers + and layer_idx < self.config.recompute_num_layers + recompute_skip_num_layers + ): + hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1)) + else: + # Build args for non-checkpointed path + forward_args = ( + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb, + timestep_emb, + guidance_emb, + ) + hidden_states, context = custom(layer_idx, layer_idx + 1)(*forward_args) + + elif recompute_layer_ids is not None: + for block_layer_idx in range(self.num_layers_per_pipeline_rank): + layer_idx = block_layer_idx + get_transformer_layer_offset(self.config, self.vp_stage) + if layer_idx not in recompute_layer_ids or ( + (self.config.fp8 or self.config.fp4) and not hidden_states.requires_grad + ): + forward_args = ( + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb, + timestep_emb, + guidance_emb, + ) + hidden_states, context = custom(block_layer_idx, block_layer_idx + 1)(*forward_args) + else: + hidden_states, context = checkpoint_handler(custom(block_layer_idx, block_layer_idx + 1)) + else: + raise ValueError("Invalid activation recompute method.") + + return hidden_states diff --git a/tests/unit_tests/backends/megatron/diffusion/__init__.py b/tests/unit_tests/backends/megatron/diffusion/__init__.py new file mode 100644 index 000000000..89778402a --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. diff --git a/tests/unit_tests/backends/megatron/diffusion/conftest.py b/tests/unit_tests/backends/megatron/diffusion/conftest.py new file mode 100644 index 000000000..8fe29f243 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/conftest.py @@ -0,0 +1,47 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Pytest fixtures for Flux diffusion model testing. + +Re-exports shared fixtures from the parent megatron conftest. +""" + +import pytest + +from tests.utils import install_aiter_deepbind_hook + +# Install Primus' production aiter RTLD_DEEPBIND import hook before any diffusion +# test imports the aiter mha kernels. On gfx942/gfx950, transformer_engine's +# stale vendored libmha interposes the global ``aiter::mha_bwd`` over Turbo's +# pinned aiter, crashing attention backward (ROCm/aiter#1332). Real training +# installs this via the ``megatron.turbo.aiter_deepbind`` before_train patch; +# unit tests call flash_attn_func directly and never hit that phase, so we wire +# up the same hook here. No-op without a GPU. +install_aiter_deepbind_hook() + +from tests.unit_tests.backends.megatron.conftest import ( # noqa: F401,E402 + init_parallel_state, +) + + +@pytest.fixture(autouse=True) +def _unset_nvte_attention_env(monkeypatch): + """Clear the TE attention-backend env vars for diffusion tests. + + Some container images bake ``NVTE_FLASH_ATTN=0`` (they target the fused/CK + attention path). Flux's ``DiffusionModule._set_attention_backend()`` defaults + to the ``auto`` backend, which validates that ``NVTE_FLASH_ATTN``/ + ``NVTE_FUSED_ATTN``/``NVTE_UNFUSED_ATTN`` are unset-or-1, so the baked ``0`` + makes every Flux model construction fail. Mirror Megatron's own harness + (``Utils.initialize_distributed`` in ``tests/unit_tests/test_utilities.py``), + which pops these three vars so a baked/leaked value cannot poison the auto + backend check. Use ``monkeypatch.delenv`` (not ``os.environ.pop``): + ``_set_attention_backend`` writes ``os.environ[var] = 1`` after its check, so + monkeypatch's teardown is what restores each var to its pre-test state and + contains that write-back leak. Scoped to the diffusion suite only -- + non-diffusion Primus megatron tests never construct a model that runs this + assertion. + """ + for var in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"): + monkeypatch.delenv(var, raising=False) diff --git a/tests/unit_tests/backends/megatron/diffusion/constants.py b/tests/unit_tests/backends/megatron/diffusion/constants.py new file mode 100644 index 000000000..2689f67eb --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/constants.py @@ -0,0 +1,73 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Constants for Flux diffusion model testing. + +This module centralizes common dimensional and architectural constants +used across diffusion tests to reduce magic numbers and improve maintainability. +""" + +# Model Architecture Dimensions +HIDDEN_DIM_FLUX = 3072 # Flux model hidden dimension +NUM_ATTENTION_HEADS_FLUX = 24 # Number of attention heads in Flux +HEAD_DIM_FLUX = 128 # Hidden dimension per attention head (3072 / 24) + +# Encoder Dimensions +T5_XXL_EMBEDDING_DIM = 4096 # T5-XXL text encoder output dimension +CLIP_L_EMBEDDING_DIM = 768 # CLIP-L pooled embedding dimension +VAE_LATENT_CHANNELS = 16 # VAE encoder output channels for Flux + +# Position Encoding +ROPE_THETA_DEFAULT = 10000 # Default theta value for RoPE (Rotary Position Embedding) +FLUX_AXES_DIM = (16, 56, 56) # Default axes dimensions for Flux 3D RoPE + +# Embedding Dimensions +TIMESTEP_EMBEDDING_DIM = 256 # Standard timestep embedding dimension + +# Batch Sizes (count-based naming) +BATCH_SIZE_SINGLE = 1 # Single sample tests +BATCH_SIZE_PAIR = 2 # Paired sample tests +BATCH_SIZE_QUAD = 4 # Quad sample tests (most common) +BATCH_SIZE_OCTO = 8 # Octuple sample tests +BATCH_SIZE_HEX = 16 # Hexadecuple sample tests + +# Image Dimensions (size-based naming) +IMG_SIZE_MICRO = 4 # Micro size for unit tests +IMG_SIZE_MINI = 8 # Mini size for quick tests +IMG_SIZE_TINY = 16 # Tiny size for fast tests +IMG_SIZE_SMALL = 32 # Small size for standard tests +IMG_SIZE_MEDIUM = 64 # Medium size +IMG_SIZE_LARGE = 128 # Large size + +# Text Sequence Lengths +TEXT_SEQ_LEN_SHORT = 77 # Standard CLIP text length +TEXT_SEQ_LEN_MEDIUM = 128 # Medium sequence length +TEXT_SEQ_LEN_LONG = 256 # Long sequence length +TEXT_SEQ_LEN_XLARGE = 512 # Extra long sequence length + +# Additional Sequence Lengths +SEQ_LEN_TINY = 100 # Small sequence for basic tests +ATTENTION_SEQ_LEN = 256 # Standard attention sequence length + +# Position Encoding Dimensions +POS_GRID_SMALL = 4 # Small grid for position tests (4x3, 4x4) +CHANNEL_GROUPS = 16 # Channel groups for position IDs + +# Training Hyperparameters (common test values) +DEFAULT_LEARNING_RATE = 1e-4 +DEFAULT_WEIGHT_DECAY = 0.01 +DEFAULT_GRADIENT_CLIP_NORM = 1.0 +DEFAULT_GUIDANCE_SCALE = 3.5 # Typical classifier-free guidance scale + +# Scheduler Parameters +DEFAULT_NUM_TRAIN_TIMESTEPS = 1000 +DEFAULT_SHIFT = 1.0 # Default shift for flow matching scheduler + +# Tensor Dimensions +TENSOR_CHANNELS_RGB = 3 # RGB channels for scheduler tests + +# Iteration Counts +TRAINING_STEPS_FEW = 5 # Few training steps +TRAINING_STEPS_MODERATE = 10 # Moderate training steps +ACCUMULATION_STEPS = 4 # Gradient accumulation steps diff --git a/tests/unit_tests/backends/megatron/diffusion/helpers.py b/tests/unit_tests/backends/megatron/diffusion/helpers.py new file mode 100644 index 000000000..e279b3a10 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/helpers.py @@ -0,0 +1,245 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Test helper utilities for diffusion model testing. + +Provides utility functions for tensor validation, mock data creation, +and common test operations. +""" + +from typing import Optional, Tuple + +import torch + +from tests.unit_tests.backends.megatron.diffusion.constants import ( + CLIP_L_EMBEDDING_DIM, + T5_XXL_EMBEDDING_DIM, + TEXT_SEQ_LEN_LONG, + VAE_LATENT_CHANNELS, +) + + +def assert_tensor_shape(tensor: torch.Tensor, expected_shape: Tuple[int, ...], name: Optional[str] = None): + """ + Assert that a tensor has the expected shape. + + Args: + tensor: Tensor to check + expected_shape: Expected shape tuple + name: Optional tensor name for error messages + + Raises: + AssertionError if shapes don't match + """ + name_str = f"{name} " if name else "" + assert ( + tensor.shape == expected_shape + ), f"{name_str}shape mismatch: expected {expected_shape}, got {tensor.shape}" + + +def assert_tensor_dtype(tensor: torch.Tensor, expected_dtype: torch.dtype, name: Optional[str] = None): + """ + Assert that a tensor has the expected dtype. + + Args: + tensor: Tensor to check + expected_dtype: Expected dtype + name: Optional tensor name for error messages + + Raises: + AssertionError if dtypes don't match + """ + name_str = f"{name} " if name else "" + assert ( + tensor.dtype == expected_dtype + ), f"{name_str}dtype mismatch: expected {expected_dtype}, got {tensor.dtype}" + + +def assert_no_nan_inf(tensor: torch.Tensor, name: Optional[str] = None): + """ + Assert that a tensor contains no NaN or Inf values. + + Args: + tensor: Tensor to check + name: Optional tensor name for error messages + + Raises: + AssertionError if NaN or Inf found + """ + name_str = f"{name} " if name else "" + assert not torch.isnan(tensor).any(), f"{name_str}contains NaN values" + assert not torch.isinf(tensor).any(), f"{name_str}contains Inf values" + + +def assert_tensor_close( + tensor1: torch.Tensor, + tensor2: torch.Tensor, + rtol: float = 1e-5, + atol: float = 1e-8, + name: Optional[str] = None, +): + """ + Assert that two tensors are close within tolerance. + + Args: + tensor1: First tensor + tensor2: Second tensor + rtol: Relative tolerance + atol: Absolute tolerance + name: Optional tensor name for error messages + + Raises: + AssertionError if tensors are not close + """ + name_str = f"{name} " if name else "" + assert torch.allclose( + tensor1, tensor2, rtol=rtol, atol=atol + ), f"{name_str}tensors are not close within rtol={rtol}, atol={atol}" + + +def create_mock_latents( + batch_size: int, + height: int, + width: int, + channels: int = VAE_LATENT_CHANNELS, # VAE output channels for Flux + device: str = "cpu", + dtype: torch.dtype = torch.float32, + seed: Optional[int] = None, +) -> torch.Tensor: + """ + Create mock latent tensors for testing. + + Args: + batch_size: Batch size + height: Latent height + width: Latent width + channels: Number of channels (default: 16 for Flux) + device: Device to create tensors on + dtype: Tensor dtype + seed: Optional random seed for reproducibility + + Returns: + Mock latents of shape (batch_size, channels, height, width) + """ + if seed is not None: + generator = torch.Generator(device=device).manual_seed(seed) + return torch.randn( + batch_size, channels, height, width, device=device, dtype=dtype, generator=generator + ) + else: + return torch.randn(batch_size, channels, height, width, device=device, dtype=dtype) + + +def create_mock_text_embeddings( + batch_size: int, + seq_len: int, + hidden_dim: int = T5_XXL_EMBEDDING_DIM, # T5-XXL dimension + device: str = "cpu", + dtype: torch.dtype = torch.float32, + seed: Optional[int] = None, +) -> torch.Tensor: + """ + Create mock T5 text embeddings for testing. + + Args: + batch_size: Batch size + seq_len: Sequence length + hidden_dim: Hidden dimension (default: 4096 for T5-XXL) + device: Device to create tensors on + dtype: Tensor dtype + seed: Optional random seed for reproducibility + + Returns: + Mock text embeddings of shape (batch_size, seq_len, hidden_dim) + """ + if seed is not None: + generator = torch.Generator(device=device).manual_seed(seed) + return torch.randn(batch_size, seq_len, hidden_dim, device=device, dtype=dtype, generator=generator) + else: + return torch.randn(batch_size, seq_len, hidden_dim, device=device, dtype=dtype) + + +def create_mock_clip_embeddings( + batch_size: int, + hidden_dim: int = CLIP_L_EMBEDDING_DIM, # CLIP-L dimension + device: str = "cpu", + dtype: torch.dtype = torch.float32, + seed: Optional[int] = None, +) -> torch.Tensor: + """ + Create mock CLIP pooled embeddings for testing. + + Args: + batch_size: Batch size + hidden_dim: Hidden dimension (default: 768 for CLIP-L) + device: Device to create tensors on + dtype: Tensor dtype + seed: Optional random seed for reproducibility + + Returns: + Mock CLIP embeddings of shape (batch_size, hidden_dim) + """ + if seed is not None: + generator = torch.Generator(device=device).manual_seed(seed) + return torch.randn(batch_size, hidden_dim, device=device, dtype=dtype, generator=generator) + else: + return torch.randn(batch_size, hidden_dim, device=device, dtype=dtype) + + +def create_mock_position_ids_3d( + batch_size: int, + height: int, + width: int, + text_seq_len: int = TEXT_SEQ_LEN_LONG, # Default text sequence length + device: str = "cpu", + dtype: torch.dtype = torch.float32, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Create mock 3D position IDs for Flux (RoPE format). + + Args: + batch_size: Batch size + height: Image height (in latent space / 2) + width: Image width (in latent space / 2) + text_seq_len: Text sequence length + device: Device to create tensors on + dtype: Tensor dtype + + Returns: + Tuple of (img_ids, txt_ids) where: + - img_ids: (batch_size, height*width, 3) + - txt_ids: (batch_size, text_seq_len, 3) + """ + # Image position IDs + img_ids = torch.zeros(batch_size, height * width, 3, device=device, dtype=dtype) + + # Generate 2D grid for spatial positions + h_coords = torch.arange(height, device=device, dtype=dtype) + w_coords = torch.arange(width, device=device, dtype=dtype) + h_grid, w_grid = torch.meshgrid(h_coords, w_coords, indexing="ij") + + # Flatten and repeat for batch + h_flat = h_grid.reshape(-1) + w_flat = w_grid.reshape(-1) + + for b in range(batch_size): + img_ids[b, :, 1] = h_flat + img_ids[b, :, 2] = w_flat + + # Text position IDs (all zeros for text) + txt_ids = torch.zeros(batch_size, text_seq_len, 3, device=device, dtype=dtype) + + return img_ids, txt_ids + + +__all__ = [ + "assert_tensor_shape", + "assert_tensor_dtype", + "assert_no_nan_inf", + "assert_tensor_close", + "create_mock_latents", + "create_mock_text_embeddings", + "create_mock_clip_embeddings", + "create_mock_position_ids_3d", +] diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_embeddings.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_embeddings.py new file mode 100644 index 000000000..edf5be8b1 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_embeddings.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for diffusion model embedding layers. + +Tests TimeStepEmbedder and MLPEmbedder to ensure correct shapes and behavior. +""" + +import pytest +import torch + +from primus.backends.megatron.core.models.diffusion.common.embeddings import ( + get_timestep_embedding, +) +from tests.unit_tests.backends.megatron.diffusion.constants import ( + TIMESTEP_EMBEDDING_DIM, +) +from tests.utils import PrimusUT + + +class TestGetTimestepEmbedding(PrimusUT): + """Tests for get_timestep_embedding function (Primus implementation of sinusoidal embeddings).""" + + def test_odd_embedding_dim(self): + """Test that odd embedding dimensions are handled correctly by Primus implementation.""" + timesteps = torch.randn(4) + embedding_dim = 257 # Odd number + + emb = get_timestep_embedding(timesteps, embedding_dim) + + assert emb.shape == (4, embedding_dim), f"Expected shape (4, {embedding_dim}), got {emb.shape}" + + def test_sinusoidal_properties(self): + """Test that Primus sinusoidal embeddings have correct mathematical properties.""" + timesteps = torch.linspace(0, 100, 10) + embedding_dim = TIMESTEP_EMBEDDING_DIM + + emb = get_timestep_embedding(timesteps, embedding_dim) + + assert emb.abs().max() < 100, "Embeddings should be in reasonable range" + assert (emb > 0).any() and (emb < 0).any(), "Should have both positive and negative values" + + # For timestep=0: cos(0)=1 for all frequencies (first half), + # sin(0)=0 for all frequencies (second half) + emb_zero = get_timestep_embedding(torch.tensor([0.0]), embedding_dim) + half = embedding_dim // 2 + assert torch.allclose( + emb_zero[0, :half], torch.ones(half), atol=1e-6 + ), "cos(0) should be 1 for all frequencies" + assert torch.allclose( + emb_zero[0, half : 2 * half], torch.zeros(half), atol=1e-6 + ), "sin(0) should be 0 for all frequencies" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py new file mode 100644 index 000000000..01298962f --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_normalization.py @@ -0,0 +1,147 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for diffusion model normalization layers. + +Tests RMSNorm, AdaLN, and AdaLNContinuous. + +NOTE: These tests require CUDA and Megatron parallel state initialization. +""" + +import pytest +import torch +import torch.nn as nn +from megatron.core.transformer.transformer_config import TransformerConfig + +from primus.backends.megatron.core.models.diffusion.common.normalization import ( + AdaLN, + AdaLNContinuous, +) +from tests.unit_tests.backends.megatron.diffusion.constants import ( + ATTENTION_SEQ_LEN, + BATCH_SIZE_QUAD, + HIDDEN_DIM_FLUX, + NUM_ATTENTION_HEADS_FLUX, +) +from tests.utils import PrimusUT + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA (uses ColumnParallelLinear)") +class TestAdaLN(PrimusUT): + """Tests for AdaLN class.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """ + Initialize parallel state for AdaLN tests. + + AdaLN uses ColumnParallelLinear which requires Megatron's RNG tracker. + The init_parallel_state fixture handles initialization and cleanup. + """ + + def test_forward_output_chunks(self): + """Test that AdaLN produces correct number of chunks and numeric behavior.""" + config = TransformerConfig( + hidden_size=HIDDEN_DIM_FLUX, + num_attention_heads=NUM_ATTENTION_HEADS_FLUX, + num_layers=1, + ) + n_chunks = 9 + adaln = AdaLN(config, n_adaln_chunks=n_chunks).cuda() + + timestep_emb = torch.randn(BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda() + chunks = adaln(timestep_emb) + + assert len(chunks) == n_chunks, f"Expected {n_chunks} chunks, got {len(chunks)}" + for chunk in chunks: + assert chunk.shape == (BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX) + + # Numeric: gate=0 should zero out the contribution in scale_add + x = torch.randn(16, BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda() + residual = torch.randn(16, BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda() + zero_gate = torch.zeros(BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda() + result = adaln.scale_add(residual, x, zero_gate) + assert torch.allclose(result, residual, atol=1e-6), "gate=0 should leave residual unchanged" + + def test_default_init_method_zeros_modulation_weight(self): + """The default init_method is nn.init.zeros_; modulation weight comes up zero. + + Guards the post-Flux-PR contract that AdaLN's default init is + observably zero, so downstream callers don't accidentally pick up + the NeMo-aligned normal_ RNG draw without opting in. + """ + config = TransformerConfig( + hidden_size=HIDDEN_DIM_FLUX, + num_attention_heads=NUM_ATTENTION_HEADS_FLUX, + num_layers=1, + ) + adaln = AdaLN(config, n_adaln_chunks=6).cuda() + weight = adaln.adaLN_modulation[-1].weight + assert torch.equal( + weight, torch.zeros_like(weight) + ), "Default AdaLN init_method should produce a zero modulation weight" + + def test_normal_init_method_produces_nonzero_modulation_weight(self): + """Passing nn.init.normal_ produces a nonzero pre-init_weights() draw. + + This is the call sites used by Flux's layer_spec.py to match NeMo's + RNG sequence. Flux's init_weights() immediately re-zeroes these + weights, so the only observable effect is RNG advancement. + """ + config = TransformerConfig( + hidden_size=HIDDEN_DIM_FLUX, + num_attention_heads=NUM_ATTENTION_HEADS_FLUX, + num_layers=1, + ) + torch.manual_seed(0) + adaln = AdaLN(config, n_adaln_chunks=6, init_method=nn.init.normal_).cuda() + weight = adaln.adaLN_modulation[-1].weight + assert not torch.equal( + weight, torch.zeros_like(weight) + ), "init_method=normal_ should draw a nonzero modulation weight" + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="AdaLNContinuous forward dispatches to primus::fused_ln_modulate, " + "which has no CPU kernel registered.", +) +class TestAdaLNContinuous(PrimusUT): + """Tests for AdaLNContinuous class.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """AdaLNContinuous uses RowParallelLinear which requires Megatron's + RNG tracker / parallel state. Mirror TestAdaLN.""" + + def test_output_shape(self): + """Test that AdaLNContinuous produces correct output shape.""" + config = TransformerConfig( + hidden_size=HIDDEN_DIM_FLUX, + num_attention_heads=NUM_ATTENTION_HEADS_FLUX, + num_layers=1, + ) + adaln = AdaLNContinuous(config, conditioning_embedding_dim=HIDDEN_DIM_FLUX).cuda() + + # Use sequence-first format: [seq_len, batch, hidden] + x = torch.randn(ATTENTION_SEQ_LEN, BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda() + cond = torch.randn(BATCH_SIZE_QUAD, HIDDEN_DIM_FLUX).cuda() + + output = adaln(x, cond) + assert output.shape == x.shape + + def test_invalid_norm_type(self): + """Test that invalid norm type raises error (Primus validation).""" + config = TransformerConfig( + hidden_size=HIDDEN_DIM_FLUX, + num_attention_heads=NUM_ATTENTION_HEADS_FLUX, + num_layers=1, + ) + + with pytest.raises(ValueError, match="Unknown normalization type"): + AdaLNContinuous(config, conditioning_embedding_dim=HIDDEN_DIM_FLUX, norm_type="invalid") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 6094046e7650e412fcb318dc50626b73cc6099a2 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Thu, 9 Jul 2026 03:39:29 +0300 Subject: [PATCH 017/127] feat(flux): Primus-Turbo float8 + local-spec extensions (#809) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/opt` — review after it. Parent of the fp8/mxfp4/compile layers. ## What this changes The Primus-Turbo integration layer: the float8 "local" extension, the turbo local-spec layer wiring, the Triton fp8-cast kernels, the native fp8 layout, and fp8 utilities. ## Why it's stacked here The float8 extension lazily imports the FSDP2 fp8 all-gather added in `feat/flux/opt`, and a turbo test exercises that path — so it bases on `feat/flux/opt`, not `feat/flux/core`. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); its float8/fp8 unit tests are green on the current CI pin (no turbo-bump dependency). Builds on `feat/flux/opt`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "turbo or native_fp8"`. Validated locally on an AMD GPU container: 22 passed. ## Files 8 (turbo float8 + local-spec extensions, Triton fp8-cast kernels, fp8 utils + tests). --------- Co-authored-by: Flux Split Trial Co-authored-by: luiza-amd Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- .../extensions/fp8_cast_kernels_triton.py | 343 ++++ .../megatron/core/extensions/primus_turbo.py | 20 +- .../extensions/primus_turbo_float8_local.py | 1703 +++++++++++++++++ .../extensions/primus_turbo_local_spec.py | 318 +++ primus/backends/megatron/core/fp8_utils.py | 11 +- primus/backends/megatron/core/utils.py | 135 +- .../megatron/test_native_fp8_layout.py | 184 ++ .../test_primus_turbo_float8_local.py | 431 +++++ 8 files changed, 3129 insertions(+), 16 deletions(-) create mode 100644 primus/backends/megatron/core/extensions/fp8_cast_kernels_triton.py create mode 100644 primus/backends/megatron/core/extensions/primus_turbo_float8_local.py create mode 100644 primus/backends/megatron/core/extensions/primus_turbo_local_spec.py create mode 100644 tests/unit_tests/backends/megatron/test_native_fp8_layout.py create mode 100644 tests/unit_tests/backends/megatron/test_primus_turbo_float8_local.py diff --git a/primus/backends/megatron/core/extensions/fp8_cast_kernels_triton.py b/primus/backends/megatron/core/extensions/fp8_cast_kernels_triton.py new file mode 100644 index 000000000..0d3b883bc --- /dev/null +++ b/primus/backends/megatron/core/extensions/fp8_cast_kernels_triton.py @@ -0,0 +1,343 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Fused FP8 cast Triton kernels (vendored into Primus). + +Two self-contained @triton_op kernels used by the FP8 local-spec path: + + - ``cast_fp8_triton`` : fused FP8 cast + amax (no transpose), for the + native NN/TN arm. Single quantize + amax HBM + pass, no [N, M] transpose write. + - ``cast_transpose_fp8_triton`` : fused FP8 cast + transpose + amax, for the + forced-NT / delayed-scaling arm. One kernel + replaces a quantize kernel + .t().contiguous() + copy. + +Both follow TE's _cast_transpose_triton pattern with 2D grouped tiling. + +Vendored from Primus-Turbo ``integration/fp8-native-main`` @ 022f2b2b +(``primus_turbo/pytorch/kernels/quantization/cast_fp8.py`` and +``cast_transpose_fp8.py``). They are kept in Primus so the FP8 local-spec path +builds against stock public Primus-Turbo ``main``, which never carried these +kernels. The op namespaces were renamed ``primus_turbo::`` -> ``primus::`` to +match the other Primus ``@triton_op``s. + +Registered as @triton_ops so that: + - torch.compile / Inductor can see through the op + - Output tensors are standard torch.empty (no triton.reinterpret metadata) + - register_fake provides correct shapes/strides for compile-time validation +""" + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl +from torch.library import triton_op, wrap_triton + + +def _fp8_max(dtype: torch.dtype) -> float: + return torch.finfo(dtype).max + + +# --------------------------------------------------------------------------- +# Cast + transpose + amax +# --------------------------------------------------------------------------- + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "GROUP_M": 1}, num_warps=4), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "GROUP_M": 8}, num_warps=4), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "GROUP_M": 8}, num_warps=8), + ], + key=["M", "N"], +) +@triton.jit +def _cast_transpose_amax_kernel( + X_ptr, + C_ptr, + T_ptr, + stride_xm, + stride_xn, + stride_cm, + stride_cn, + stride_tm, + stride_tn, + M, + N, + scale_ptr, + amax_ptr, + scale_inv_ptr, + max_fp8: tl.constexpr, + COMPUTE_AMAX: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + GROUP_M: tl.constexpr, +): + pid = tl.program_id(0) + scale = tl.load(scale_ptr) + + grid_m = tl.cdiv(M, BLOCK_M) + grid_n = tl.cdiv(N, BLOCK_N) + + width = GROUP_M * grid_n + group_id = pid // width + group_size = tl.minimum(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid % group_size) + pid_n = (pid % width) // group_size + + rm = pid_m.to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n.to(tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N) + mask = (rm < M)[:, None] & (rn < N)[None, :] + + x_ptrs = X_ptr + rm[:, None] * stride_xm + rn[None, :] * stride_xn + a = tl.load(x_ptrs, mask=mask) + val = a.to(tl.float32) + + scaled = val * scale + scaled = tl.clamp(scaled, -max_fp8, max_fp8) + fp8_val = scaled.to(C_ptr.dtype.element_ty) + + c_ptrs = C_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn + tl.store(c_ptrs, fp8_val, mask=mask) + + # Transpose the tile so the store to T is coalesced (stride-1 in the fast dim). + # Without this, the transpose store scatters with stride=M which kills + # write throughput on MI355X under memory pressure. + fp8_val_t = tl.trans(fp8_val) + rn2 = pid_n.to(tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N) + rm2 = pid_m.to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M) + mask_t = (rn2 < N)[:, None] & (rm2 < M)[None, :] + t_ptrs = T_ptr + rn2[:, None] * stride_tm + rm2[None, :] * stride_tn + tl.store(t_ptrs, fp8_val_t, mask=mask_t) + + if COMPUTE_AMAX: + tile_amax = tl.max(tl.abs(val)) + tl.atomic_max(amax_ptr, tile_amax, sem="relaxed") + + if pid == 0: + tl.store(scale_inv_ptr, tl.fdiv(1.0, scale)) + + +@triton_op("primus::cast_transpose_fp8_triton", mutates_args=()) +def cast_transpose_fp8_triton( + x: torch.Tensor, + fp8_dtype: torch.dtype, + scale: torch.Tensor, + amax_out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused FP8 cast + transpose + optional amax. + + Args: + x: 2D input tensor [M, N] (bf16 or f16), must be contiguous. + fp8_dtype: Target FP8 dtype (e.g. torch.float8_e4m3fn). + scale: Scalar float32 tensor with the quantization scale. + amax_out: Optional scalar float32 tensor. If provided, the kernel + writes the abs-max of x into it (atomically reduced). + + Returns: + (cast_out, transpose_out, scale_inv) where: + cast_out: [M, N] FP8 tensor (same layout as input) + transpose_out: [N, M] FP8 tensor (contiguous transpose) + scale_inv: scalar float32, 1/scale + """ + if x.ndim != 2: + raise ValueError(f"Expected 2D input, got {x.ndim}D") + if not x.is_contiguous(): + x = x.contiguous() + + M, N = x.shape + max_fp8 = _fp8_max(fp8_dtype) + + cast_out = torch.empty((M, N), dtype=fp8_dtype, device=x.device) + transpose_out = torch.empty((N, M), dtype=fp8_dtype, device=x.device) + scale_inv = torch.empty((), dtype=torch.float32, device=x.device) + + compute_amax = amax_out is not None + if not compute_amax: + amax_out = torch.empty((), dtype=torch.float32, device=x.device) + else: + amax_out.zero_() + + grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) + + wrap_triton(_cast_transpose_amax_kernel)[grid]( + x, + cast_out, + transpose_out, + x.stride(0), + x.stride(1), + cast_out.stride(0), + cast_out.stride(1), + transpose_out.stride(0), + transpose_out.stride(1), + M, + N, + scale, + amax_out, + scale_inv, + max_fp8, + compute_amax, + ) + + return cast_out, transpose_out, scale_inv + + +@cast_transpose_fp8_triton.register_fake +def _cast_transpose_fp8_triton_meta( + x: torch.Tensor, + fp8_dtype: torch.dtype, + scale: torch.Tensor, + amax_out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + M, N = x.shape + return ( + torch.empty((M, N), dtype=fp8_dtype, device=x.device), + torch.empty((N, M), dtype=fp8_dtype, device=x.device), + torch.empty((), dtype=torch.float32, device=x.device), + ) + + +# --------------------------------------------------------------------------- +# Cast + amax (no transpose) +# --------------------------------------------------------------------------- + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "GROUP_M": 1}, num_warps=4), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "GROUP_M": 8}, num_warps=4), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "GROUP_M": 8}, num_warps=8), + ], + key=["M", "N"], +) +@triton.jit +def _cast_amax_kernel( + X_ptr, + C_ptr, + stride_xm, + stride_xn, + stride_cm, + stride_cn, + M, + N, + scale_ptr, + amax_ptr, + scale_inv_ptr, + max_fp8: tl.constexpr, + COMPUTE_AMAX: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + GROUP_M: tl.constexpr, +): + pid = tl.program_id(0) + scale = tl.load(scale_ptr) + + grid_m = tl.cdiv(M, BLOCK_M) + grid_n = tl.cdiv(N, BLOCK_N) + + width = GROUP_M * grid_n + group_id = pid // width + group_size = tl.minimum(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid % group_size) + pid_n = (pid % width) // group_size + + rm = pid_m.to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M) + rn = pid_n.to(tl.int64) * BLOCK_N + tl.arange(0, BLOCK_N) + mask = (rm < M)[:, None] & (rn < N)[None, :] + + x_ptrs = X_ptr + rm[:, None] * stride_xm + rn[None, :] * stride_xn + a = tl.load(x_ptrs, mask=mask) + val = a.to(tl.float32) + + scaled = val * scale + scaled = tl.clamp(scaled, -max_fp8, max_fp8) + fp8_val = scaled.to(C_ptr.dtype.element_ty) + + c_ptrs = C_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn + tl.store(c_ptrs, fp8_val, mask=mask) + + if COMPUTE_AMAX: + tile_amax = tl.max(tl.abs(val)) + tl.atomic_max(amax_ptr, tile_amax, sem="relaxed") + + if pid == 0: + tl.store(scale_inv_ptr, tl.fdiv(1.0, scale)) + + +@triton_op("primus::cast_fp8_triton", mutates_args=()) +def cast_fp8_triton( + x: torch.Tensor, + fp8_dtype: torch.dtype, + scale: torch.Tensor, + amax_out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fused FP8 cast + optional amax (no transpose). + + Args: + x: 2D input tensor [M, N] (bf16 or f16), must be contiguous. + fp8_dtype: Target FP8 dtype (e.g. torch.float8_e4m3fn). + scale: Scalar float32 tensor with the quantization scale. + amax_out: Optional scalar float32 tensor. If provided, the kernel + writes the abs-max of the *unscaled* x into it (atomically). + + Returns: + (cast_out, scale_inv) where: + cast_out: [M, N] FP8 tensor (same layout as input) + scale_inv: scalar float32, 1 / scale + """ + if x.ndim != 2: + raise ValueError(f"Expected 2D input, got {x.ndim}D") + if not x.is_contiguous(): + x = x.contiguous() + + M, N = x.shape + max_fp8 = _fp8_max(fp8_dtype) + + cast_out = torch.empty((M, N), dtype=fp8_dtype, device=x.device) + scale_inv = torch.empty((), dtype=torch.float32, device=x.device) + + compute_amax = amax_out is not None + if not compute_amax: + amax_out = torch.empty((), dtype=torch.float32, device=x.device) + else: + amax_out.zero_() + + grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) + + wrap_triton(_cast_amax_kernel)[grid]( + x, + cast_out, + x.stride(0), + x.stride(1), + cast_out.stride(0), + cast_out.stride(1), + M, + N, + scale, + amax_out, + scale_inv, + max_fp8, + compute_amax, + ) + + return cast_out, scale_inv + + +@cast_fp8_triton.register_fake +def _cast_fp8_triton_meta( + x: torch.Tensor, + fp8_dtype: torch.dtype, + scale: torch.Tensor, + amax_out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + M, N = x.shape + return ( + torch.empty((M, N), dtype=fp8_dtype, device=x.device), + torch.empty((), dtype=torch.float32, device=x.device), + ) diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index b565bf01d..4d7a99c36 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -77,8 +77,11 @@ float8_e4m3, ) from torch import Tensor + +# Imported from .constants (not .fp8) for TransformerEngine >= 2.12 compat; +# the symbol moved out of transformer_engine.pytorch.fp8 in that release. from transformer_engine.pytorch.constants import dist_group_type -from transformer_engine.pytorch.fp8 import DelayedScaling, FP8GlobalStateManager, Recipe +from transformer_engine.pytorch.fp8 import FP8GlobalStateManager, Recipe from primus.core.pipeline_parallel.handler.offload_handler import OFFLOAD_BUFFER @@ -310,6 +313,7 @@ def __init__( strategy: ScalingStrategy = ScalingStrategy.DYNAMIC, scale_dtype: ScaleDtype = ScaleDtype.FP32, block_size: int = None, + use_gradient_sr: bool = True, ): self._is_fp4 = False self._is_fp8 = False @@ -321,6 +325,7 @@ def __init__( strategy=strategy, scale_dtype=scale_dtype, block_size=block_size, + use_gradient_sr=use_gradient_sr, ) self._is_fp4 = True else: @@ -440,7 +445,7 @@ def get_turbo_quant_config(cls) -> PrimusTurboQuantConfig: @classmethod def get_fp8_autocast_state( cls, - ) -> Tuple[bool, bool, Recipe, dist_group_type, bool, bool, PrimusTurboQuantConfig]: + ) -> Tuple[bool, bool, Recipe, dist_group_type, bool, bool, bool, bool, PrimusTurboQuantConfig]: """FP8 autocast state getter""" return ( FP8GlobalStateManager.FP8_ENABLED, @@ -457,7 +462,7 @@ def get_fp8_autocast_state( @classmethod def set_fp8_autocast_state( cls, - fp8_state: Tuple[bool, bool, DelayedScaling, dist_group_type, bool, bool, PrimusTurboQuantConfig], + fp8_state: Tuple[bool, bool, Recipe, dist_group_type, bool, bool, bool, bool, PrimusTurboQuantConfig], ) -> None: """FP8 autocast state setter""" ( @@ -1828,7 +1833,7 @@ def __init__( pg_collection: Optional[ProcessGroupCollection] = None, ): """ - Initialize the Flex token dispatcher. + Initialize the DeepEP token dispatcher. Args: num_local_experts (int): Number of local experts on the current device. @@ -1838,13 +1843,14 @@ def __init__( """ super().__init__(config=config, pg_collection=pg_collection) - assert self.tp_size * self.ep_size > 1, "Flex token dispatcher requires TPxEP > 1" + if self.tp_size * self.ep_size <= 1: + raise ValueError("DeepEP token dispatcher requires TPxEP > 1") assert ( self.config.moe_enable_deepep ), "DeepEP is not enabled. Please set --moe-enable-deepep to use DeepEP backend." assert ( self.config.moe_pad_expert_input_to_capacity is False - ), "Flex token dispatcher does not support --moe-pad-expert-input-to-capacity" + ), "DeepEP token dispatcher does not support --moe-pad-expert-input-to-capacity" args = get_args() diff --git a/primus/backends/megatron/core/extensions/primus_turbo_float8_local.py b/primus/backends/megatron/core/extensions/primus_turbo_float8_local.py new file mode 100644 index 000000000..4632a54ee --- /dev/null +++ b/primus/backends/megatron/core/extensions/primus_turbo_float8_local.py @@ -0,0 +1,1703 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Compile-friendly FP8 linear layers for Megatron local spec. + +Self-contained autograd Functions that call Primus Turbo's low-level building +blocks (quantize_fp8 + gemm_fp8_impl) directly, bypassing the higher-level +pt.ops.gemm_fp8 wrappers entirely. + +Key properties: +- Tensorwise uses the setup_context pattern with primitive-only args so + torch.compile can trace through without graph breaks. FP8 weight data is + extracted by the caller (Float8*ParallelLinear._forward_impl) to avoid + tensor subclass tracing inside the autograd.Function. +- Rowwise and blockwise still use @allow_in_graph (separate feasibility work). +- gemm_fp8_impl is already a torch.library.custom_op with register_fake +- Zero TransformerEngine dependencies +- Builds against stock public Primus-Turbo main: the only kernels not on main + (the fused FP8 cast / cast+transpose Triton ops) are vendored locally in + fp8_cast_kernels_triton.py; everything else still comes from Turbo main. +- Requires tensor_model_parallel_size=1, no GAF, no sequence_parallel +""" + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl +from megatron.core.enums import Fp8Recipe +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from primus_turbo.pytorch.core.backend import BackendType +from primus_turbo.pytorch.core.low_precision import ( + Float8QuantConfig, + Format, + ScaleDtype, + ScalingGranularity, + float8_e4m3, + float8_e5m2, +) +from primus_turbo.pytorch.kernels.gemm.gemm_fp8_impl import gemm_fp8_impl +from primus_turbo.pytorch.kernels.quantization.quantization_impl import ( + quant_fp8_blockwise_for_weight_impl, + quant_fp8_blockwise_impl, +) +from primus_turbo.pytorch.ops.quantization import quantize_fp8 + +from primus.backends.megatron.core.fp8_utils import ( + MXFP8_SCALING_BLOCK_SIZE, + SCALING_BLOCK_SIZE, +) + +# --------------------------------------------------------------------------- +# torch.compile-friendly wrappers for C++ quantization ops. +# +# The raw C++ ops (primus_turbo_cpp_extension::quantize_fp8_tensorwise etc.) +# lack an Autograd dispatch key, which triggers warnings and can break +# torch.compile graph tracing. Wrapping them with @torch.library.custom_op +# (the same pattern used by gemm_fp8_impl) automatically handles Autograd +# dispatch and provides register_fake for shape inference during tracing. +# --------------------------------------------------------------------------- + +_custom_op = torch.library.custom_op + + +@_custom_op("primus::quantize_fp8_tensorwise", mutates_args=(), device_types="cuda") +def _quantize_fp8_tensorwise_op( + x: torch.Tensor, out_dtype: torch.dtype, scale: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor, torch.Tensor]: + return torch.ops.primus_turbo_cpp_extension.quantize_fp8_tensorwise(x, out_dtype, scale) + + +@_quantize_fp8_tensorwise_op.register_fake +def _quantize_fp8_tensorwise_fake( + x: torch.Tensor, out_dtype: torch.dtype, scale: Optional[torch.Tensor] = None +) -> Tuple[torch.Tensor, torch.Tensor]: + x_fp8 = torch.empty_like(x, dtype=out_dtype) + scale_inv = torch.empty((), dtype=torch.float32, device=x.device) + return x_fp8, scale_inv + + +def _quantize_fp8_tensorwise_setup_context(ctx, inputs, output): + pass + + +def _quantize_fp8_tensorwise_backward(ctx, grad_x_fp8, grad_scale_inv): + return None, None, None + + +_quantize_fp8_tensorwise_op.register_autograd( + _quantize_fp8_tensorwise_backward, + setup_context=_quantize_fp8_tensorwise_setup_context, +) + + +def _cast_transpose_fp8_fused_with_amax( + x: torch.Tensor, + out_dtype: torch.dtype, + scale: torch.Tensor, + amax_out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused FP8 quantize + transpose + optional amax via the Triton @triton_op. + + Returns (fp8_out, fp8_transpose, scale_inv). Inductor-transparent. + + Uses the cast+transpose Triton kernel vendored in Primus + (fp8_cast_kernels_triton), so this arm builds against stock Turbo main. + """ + from primus.backends.megatron.core.extensions.fp8_cast_kernels_triton import ( + cast_transpose_fp8_triton, + ) + + return cast_transpose_fp8_triton(x, out_dtype, scale, amax_out) + + +def _cast_fp8_fused_with_amax( + x: torch.Tensor, + out_dtype: torch.dtype, + scale: torch.Tensor, + amax_out: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Native-layout fused FP8 quantize (no transpose) + optional amax. + + Analog of _cast_transpose_fp8_fused_with_amax for the natural NN/TN arm: + applies the delayed scale and captures the current abs-max in a single HBM + pass, with no [N, M] transpose write. Returns (fp8_out, scale_inv). + + Uses the no-transpose cast Triton kernel vendored in Primus + (fp8_cast_kernels_triton), so this arm builds against stock Turbo main. + """ + from primus.backends.megatron.core.extensions.fp8_cast_kernels_triton import ( + cast_fp8_triton, + ) + + return cast_fp8_triton(x, out_dtype, scale, amax_out) + + +def _quantize_fp8_tw( + x: torch.Tensor, + out_dtype: torch.dtype, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Dispatch tensorwise FP8 quantize via the standard C++ @custom_op.""" + return _quantize_fp8_tensorwise_op(x, out_dtype) + + +def _get_fp8_dtype(format: Format, is_fwd: bool): + if format == Format.E4M3: + return float8_e4m3 + elif format == Format.E5M2: + return float8_e5m2 + elif format == Format.HYBRID: + return float8_e4m3 if is_fwd else float8_e5m2 + else: + raise ValueError(f"Unsupported FP8 format: {format}") + + +def _build_fp8_config(config): + """Build Float8QuantConfig from TransformerConfig without TE dependency.""" + FORMAT_MAP = {"e4m3": Format.E4M3, "hybrid": Format.HYBRID} + fmt = FORMAT_MAP[config.fp8] + + if config.fp8_recipe in (Fp8Recipe.tensorwise, Fp8Recipe.delayed): + # Fp8Recipe.delayed is accepted as a TE-compatible alias: it implies + # tensorwise granularity (the only granularity that makes sense with + # delayed/amax-history scaling). The actual delayed-vs-dynamic + # strategy is resolved separately via fp8_scaling_strategy or the + # recipe itself (see Float8*ParallelLinear._use_delayed_scaling). + return Float8QuantConfig(format=fmt, granularity=ScalingGranularity.TENSORWISE) + elif config.fp8_recipe == Fp8Recipe.blockwise: + return Float8QuantConfig( + format=fmt, + granularity=ScalingGranularity.BLOCKWISE, + block_size=SCALING_BLOCK_SIZE, + ) + elif config.fp8_recipe == Fp8Recipe.mxfp8: + return Float8QuantConfig( + format=fmt, + granularity=ScalingGranularity.MX_BLOCKWISE, + block_size=MXFP8_SCALING_BLOCK_SIZE, + scale_dtype=ScaleDtype.E8M0, + ) + else: + raise ValueError( + f"Float8 local spec does not support fp8_recipe={config.fp8_recipe}. " + f"Supported: tensorwise, delayed, blockwise, mxfp8." + ) + + +# --------------------------------------------------------------------------- +# Decomposed FP8 quantize — native aten ops for Inductor fusion +# --------------------------------------------------------------------------- + + +def _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max): + """Tensorwise FP8 quantize using native aten ops (traceable by Inductor). + + Scale is computed in FP32 for precision, then narrowed to x.dtype (BF16) so + the pointwise chain (mul, clamp, to_fp8) stays in BF16 registers inside the + fused Triton kernel. scale_inv is derived from the FP32 scale before the + narrowing to preserve hipBLASLt's float32 precision requirement. + """ + amax = x.abs().amax().float() + scale_f32 = fp8_max / amax.clamp(min=1e-12) + scale_inv = 1.0 / scale_f32 + scale = scale_f32.clamp(max=torch.finfo(x.dtype).max).to(x.dtype) + x_fp8 = (x * scale).clamp(-fp8_max, fp8_max).to(fp8_dtype) + return x_fp8, scale_inv + + +# --------------------------------------------------------------------------- +# Autograd Functions — one per FP8 scaling granularity. +# --------------------------------------------------------------------------- + + +class DecomposedFP8LinearTensorwiseFunction(torch.autograd.Function): + """FP8 linear (Y = X @ W^T) with tensorwise scaling, using decomposed + quantize for forward-pass Inductor fusion. + + Experimental: designed to enable Inductor fusion of quantize with + surrounding kernels. Currently slower than OpaqueFP8LinearTensorwiseFunction + due to Inductor's internal FP32 promotion and amax reduction barriers. + Kept for future work (delayed scaling, Compiled Autograd). + + Uses setup_context pattern with primitive-only arguments (no config objects) + for clean Dynamo tracing. + """ + + @staticmethod + def forward( + input, weight, fp8_fwd_dtype, fp8_bwd_dtype, fp8_fwd_max, fp8_bwd_max, gran_value, backend_value + ): + out_dtype = input.dtype + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + a_fp8, a_scale_inv = _quantize_fp8_tensorwise(input_2d, fp8_fwd_dtype, fp8_fwd_max) + b_fp8, b_scale_inv = _quantize_fp8_tensorwise(weight, fp8_fwd_dtype, fp8_fwd_max) + + output = gemm_fp8_impl( + a_fp8, + a_scale_inv, + False, + b_fp8, + b_scale_inv, + True, + out_dtype, + False, + granularity=gran_value, + default_backend=backend_value, + ) + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + + return output, a_fp8, a_scale_inv, b_fp8, b_scale_inv + + @staticmethod + def setup_context(ctx, inputs, output): + _, _, fp8_fwd_dtype, fp8_bwd_dtype, fp8_fwd_max, fp8_bwd_max, gran_value, backend_value = inputs + output_val, a_fp8, a_scale_inv, b_fp8, b_scale_inv = output + + ctx.save_for_backward(a_fp8, a_scale_inv, b_fp8, b_scale_inv) + ctx.mark_non_differentiable(a_fp8, a_scale_inv, b_fp8, b_scale_inv) + ctx.out_dtype = inputs[0].dtype + ctx.orig_shape = inputs[0].shape + ctx.fp8_bwd_dtype = fp8_bwd_dtype + ctx.fp8_bwd_max = fp8_bwd_max + ctx.gran_value = gran_value + ctx.backend_value = backend_value + + @staticmethod + def backward(ctx, grad_output, *_): + a_fp8, a_scale_inv, b_fp8, b_scale_inv = ctx.saved_tensors + + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + if not grad_2d.is_contiguous(): + grad_2d = grad_2d.contiguous() + + grad_fp8, grad_scale_inv = _quantize_fp8_tensorwise(grad_2d, ctx.fp8_bwd_dtype, ctx.fp8_bwd_max) + + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + b_fp8, + b_scale_inv, + False, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_weight = gemm_fp8_impl( + a_fp8, + a_scale_inv, + True, + grad_fp8, + grad_scale_inv, + False, + ctx.out_dtype, + True, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + + return grad_input, grad_weight, None, None, None, None, None, None + + +class OpaqueFP8LinearTensorwiseFunction(torch.autograd.Function): + """FP8 linear (Y = X @ W^T) with tensorwise scaling using Primus Turbo's + opaque C++ quantize_fp8 kernel. + + Uses the setup_context pattern with primitive-only saved state so that + torch.compile / TorchDynamo can trace through without graph breaks. + The caller must pre-extract FP8 weight data (b_fp8, b_scale_inv) from + FP8UnshardedWeightTensor *before* calling .apply(). + + Default path for tensorwise -- faster than the decomposed variant due to + optimized C++ quantize kernels. + + Backward GEMMs are normalized to NT layout (transA=F, transB=T, transC=F) + so that hipBLASLt always selects the fast TN Tensile kernel on MI355X. + Pre-transposed FP8 copies of input and weight are saved for this purpose. + """ + + @staticmethod + def forward( + input, + weight, + weight_fp8, + weight_scale_inv, + fp8_fwd_dtype, + fp8_bwd_dtype, + gran_value, + backend_value, + force_nt=True, + ): + out_dtype = input.dtype + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + a_fp8, a_scale_inv = _quantize_fp8_tw(input_2d, fp8_fwd_dtype) + + output = gemm_fp8_impl( + a_fp8, + a_scale_inv, + False, + weight_fp8, + weight_scale_inv, + True, + out_dtype, + False, + granularity=gran_value, + default_backend=backend_value, + ) + + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + + if force_nt: + # Pre-transpose so backward dgrad/wgrad both run as NT GEMMs. + a_t_fp8 = a_fp8.t().contiguous() + w_t_fp8 = weight_fp8.t().contiguous() + return (output, a_t_fp8, a_scale_inv, w_t_fp8) + + # Native: keep operands in their natural layout (no transpose). weight_fp8 + # is not returned -- setup_context saves it straight from inputs. + return (output, a_fp8, a_scale_inv) + + @staticmethod + def setup_context(ctx, inputs, output): + # Tolerant unpack: legacy 8-arg callers default to forced NT; only the + # 9th positional (force_nt) toggles the native arm. + force_nt = inputs[8] if len(inputs) > 8 else True + ( + input, + weight, + weight_fp8, + weight_scale_inv, + fp8_fwd_dtype, + fp8_bwd_dtype, + gran_value, + backend_value, + ) = inputs[:8] + + if force_nt: + output_val, a_t_fp8, a_scale_inv, w_t_fp8 = output + ctx.save_for_backward(a_t_fp8, a_scale_inv, w_t_fp8, weight_scale_inv) + ctx.mark_non_differentiable(a_t_fp8, a_scale_inv, w_t_fp8) + else: + output_val, a_fp8, a_scale_inv = output + ctx.save_for_backward(a_fp8, a_scale_inv, weight_fp8, weight_scale_inv) + ctx.mark_non_differentiable(a_fp8, a_scale_inv) + + ctx.force_nt = force_nt + ctx.out_dtype = input.dtype + ctx.orig_shape = input.shape + ctx.fp8_bwd_dtype = fp8_bwd_dtype + ctx.gran_value = gran_value + ctx.backend_value = backend_value + + @staticmethod + def backward(ctx, grad_output, *_): + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + grad_fp8, grad_scale_inv = _quantize_fp8_tw(grad_2d, ctx.fp8_bwd_dtype) + + if ctx.force_nt: + a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv = ctx.saved_tensors + + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + w_t_fp8, + w_scale_inv, + True, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_t_fp8 = grad_fp8.t().contiguous() + grad_weight = gemm_fp8_impl( + grad_t_fp8, + grad_scale_inv, + False, + a_t_fp8, + a_scale_inv, + True, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + else: + a_fp8, a_scale_inv, w_fp8, w_scale_inv = ctx.saved_tensors + + # dgrad NN: grad @ weight (transA=F, transB=F) + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + w_fp8, + w_scale_inv, + False, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + # wgrad TN: grad^T @ a (transA=T, transB=F) + grad_weight = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + True, + a_fp8, + a_scale_inv, + False, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + + return grad_input, grad_weight, None, None, None, None, None, None, None + + +class DelayedFP8LinearTensorwiseFunction(torch.autograd.Function): + """FP8 linear with delayed tensorwise scaling -- 3 independent scale tracks. + + Weight is quantized inline during the forward pass (inside the compiled + graph where Inductor eliminates CPU dispatcher overhead). Input and + gradient quantization also happen inline with fused amax capture. + Backward GEMM layout follows ``force_nt`` (set from ``fp8_force_nt_layout``): + native by default (dgrad=NN, wgrad=TN, no transpose), or forced-NT (operands + pre-transposed so both backward GEMMs run as NT) when opted in. + """ + + @staticmethod + def forward( + input, + weight, + scale_input, + scale_weight, + scale_grad, + staged_input_amax, + staged_weight_amax, + staged_grad_amax, + fp8_fwd_dtype, + fp8_bwd_dtype, + gran_value, + backend_value, + force_nt=True, + ): + out_dtype = input.dtype + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + if force_nt: + a_fp8, a_t_fp8, a_scale_inv = _cast_transpose_fp8_fused_with_amax( + input_2d, fp8_fwd_dtype, scale_input, staged_input_amax + ) + torch._assert(a_t_fp8.is_contiguous(), "cast_transpose must return contiguous transpose") + + w_fp8, w_t_fp8, w_scale_inv = _cast_transpose_fp8_fused_with_amax( + weight, fp8_fwd_dtype, scale_weight, staged_weight_amax + ) + torch._assert(w_t_fp8.is_contiguous(), "cast_transpose must return contiguous transpose") + + output = gemm_fp8_impl( + a_fp8, + a_scale_inv, + False, + w_fp8, + w_scale_inv, + True, + out_dtype, + False, + granularity=gran_value, + default_backend=backend_value, + ) + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + + return (output, a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv) + + # Native: apply the delayed scale and capture the current amax in a + # single fused pass via the no-transpose Triton kernel. No transpose; + # backward runs dgrad=NN / wgrad=TN (see backward). + a_fp8, a_scale_inv = _cast_fp8_fused_with_amax( + input_2d, fp8_fwd_dtype, scale_input, staged_input_amax + ) + + w_fp8, w_scale_inv = _cast_fp8_fused_with_amax( + weight, fp8_fwd_dtype, scale_weight, staged_weight_amax + ) + + output = gemm_fp8_impl( + a_fp8, + a_scale_inv, + False, + w_fp8, + w_scale_inv, + True, + out_dtype, + False, + granularity=gran_value, + default_backend=backend_value, + ) + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + + return (output, a_fp8, a_scale_inv, w_fp8, w_scale_inv) + + @staticmethod + def setup_context(ctx, inputs, output): + # Tolerant unpack: legacy 12-arg callers default to forced NT; only the + # 13th positional (force_nt) toggles the native arm. + force_nt = inputs[12] if len(inputs) > 12 else True + ( + input, + weight, + scale_input, + scale_weight, + scale_grad, + staged_input_amax, + staged_weight_amax, + staged_grad_amax, + fp8_fwd_dtype, + fp8_bwd_dtype, + gran_value, + backend_value, + ) = inputs[:12] + + if force_nt: + (output_val, a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv) = output + ctx.save_for_backward(a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, scale_grad, staged_grad_amax) + ctx.mark_non_differentiable(a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv) + else: + (output_val, a_fp8, a_scale_inv, w_fp8, w_scale_inv) = output + ctx.save_for_backward(a_fp8, a_scale_inv, w_fp8, w_scale_inv, scale_grad, staged_grad_amax) + ctx.mark_non_differentiable(a_fp8, a_scale_inv, w_fp8, w_scale_inv) + + ctx.force_nt = force_nt + ctx.out_dtype = input.dtype + ctx.orig_shape = input.shape + ctx.fp8_bwd_dtype = fp8_bwd_dtype + ctx.gran_value = gran_value + ctx.backend_value = backend_value + + @staticmethod + def backward(ctx, grad_output, *_): + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + op_a, a_scale_inv, op_w, w_scale_inv, scale_grad, staged_grad_amax = ctx.saved_tensors + + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + + if ctx.force_nt: + # op_a = a_t_fp8, op_w = w_t_fp8 + grad_fp8, grad_t_fp8, grad_scale_inv = _cast_transpose_fp8_fused_with_amax( + grad_2d, ctx.fp8_bwd_dtype, scale_grad, staged_grad_amax + ) + torch._assert(grad_t_fp8.is_contiguous(), "cast_transpose must return contiguous transpose") + + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + op_w, + w_scale_inv, + True, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_weight = gemm_fp8_impl( + grad_t_fp8, + grad_scale_inv, + False, + op_a, + a_scale_inv, + True, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + else: + # Native: op_a = a_fp8, op_w = w_fp8. Fused quantize + amax, no transpose. + grad_fp8, grad_scale_inv = _cast_fp8_fused_with_amax( + grad_2d, ctx.fp8_bwd_dtype, scale_grad, staged_grad_amax + ) + + # dgrad NN: grad @ weight (transA=F, transB=F) + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + op_w, + w_scale_inv, + False, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + # wgrad TN: grad^T @ a (transA=T, transB=F) + grad_weight = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + True, + op_a, + a_scale_inv, + False, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + + # 13 inputs to forward -> 13 grads returned (2 real + 11 None). + return ( + grad_input, + grad_weight, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +class DualFP8LinearTensorwiseFunction(torch.autograd.Function): + """Two independent FP8 linears in a single autograd node. + + Y_a = X_a @ W_a^T, Y_b = X_b @ W_b^T + + Uses the setup_context pattern with primitive-only saved state so that + torch.compile can trace through without graph breaks. The caller must + pre-extract FP8 weight data before calling .apply(). + + Backward GEMMs are normalized to NT layout (transA=F, transB=T, transC=F) + so that hipBLASLt always selects the fast TN Tensile kernel on MI355X. + """ + + @staticmethod + def forward( + input_a, + weight_a, + weight_fp8_a, + weight_scale_a, + input_b, + weight_b, + weight_fp8_b, + weight_scale_b, + fp8_fwd_dtype, + fp8_bwd_dtype, + gran_value, + backend_value, + ): + out_dtype = input_a.dtype + + orig_shape_a = input_a.shape + orig_shape_b = input_b.shape + input_a_2d = input_a.reshape(-1, input_a.shape[-1]) + input_b_2d = input_b.reshape(-1, input_b.shape[-1]) + + a_fp8_a, a_scale_a = _quantize_fp8_tw(input_a_2d, fp8_fwd_dtype) + a_fp8_b, a_scale_b = _quantize_fp8_tw(input_b_2d, fp8_fwd_dtype) + + output_a = gemm_fp8_impl( + a_fp8_a, + a_scale_a, + False, + weight_fp8_a, + weight_scale_a, + True, + out_dtype, + False, + granularity=gran_value, + default_backend=backend_value, + ) + output_b = gemm_fp8_impl( + a_fp8_b, + a_scale_b, + False, + weight_fp8_b, + weight_scale_b, + True, + out_dtype, + False, + granularity=gran_value, + default_backend=backend_value, + ) + + output_a = output_a.reshape(*orig_shape_a[:-1], output_a.shape[-1]) + output_b = output_b.reshape(*orig_shape_b[:-1], output_b.shape[-1]) + + a_t_fp8_a = a_fp8_a.t().contiguous() + w_t_fp8_a = weight_fp8_a.t().contiguous() + a_t_fp8_b = a_fp8_b.t().contiguous() + w_t_fp8_b = weight_fp8_b.t().contiguous() + + return (output_a, output_b, a_t_fp8_a, a_scale_a, w_t_fp8_a, a_t_fp8_b, a_scale_b, w_t_fp8_b) + + @staticmethod + def setup_context(ctx, inputs, output): + ( + input_a, + weight_a, + weight_fp8_a, + weight_scale_a, + input_b, + weight_b, + weight_fp8_b, + weight_scale_b, + fp8_fwd_dtype, + fp8_bwd_dtype, + gran_value, + backend_value, + ) = inputs + (output_a, output_b, a_t_fp8_a, a_scale_a, w_t_fp8_a, a_t_fp8_b, a_scale_b, w_t_fp8_b) = output + + ctx.save_for_backward( + a_t_fp8_a, + a_scale_a, + w_t_fp8_a, + weight_scale_a, + a_t_fp8_b, + a_scale_b, + w_t_fp8_b, + weight_scale_b, + ) + ctx.mark_non_differentiable( + a_t_fp8_a, + a_scale_a, + w_t_fp8_a, + a_t_fp8_b, + a_scale_b, + w_t_fp8_b, + ) + ctx.out_dtype = input_a.dtype + ctx.orig_shape_a = input_a.shape + ctx.orig_shape_b = input_b.shape + ctx.fp8_bwd_dtype = fp8_bwd_dtype + ctx.gran_value = gran_value + ctx.backend_value = backend_value + + @staticmethod + def backward(ctx, grad_output_a, grad_output_b, *_): + if not grad_output_a.is_contiguous(): + grad_output_a = grad_output_a.contiguous() + if not grad_output_b.is_contiguous(): + grad_output_b = grad_output_b.contiguous() + + (a_t_fp8_a, a_scale_a, w_t_fp8_a, w_scale_a, a_t_fp8_b, a_scale_b, w_t_fp8_b, w_scale_b) = ( + ctx.saved_tensors + ) + + # --- Stream A --- + grad_a_2d = grad_output_a.reshape(-1, grad_output_a.shape[-1]) + grad_fp8_a, grad_scale_a = _quantize_fp8_tw(grad_a_2d, ctx.fp8_bwd_dtype) + + grad_input_a = gemm_fp8_impl( + grad_fp8_a, + grad_scale_a, + False, + w_t_fp8_a, + w_scale_a, + True, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + grad_input_a = grad_input_a.reshape(ctx.orig_shape_a) + + grad_t_fp8_a = grad_fp8_a.t().contiguous() + grad_weight_a = gemm_fp8_impl( + grad_t_fp8_a, + grad_scale_a, + False, + a_t_fp8_a, + a_scale_a, + True, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + + # --- Stream B --- + grad_b_2d = grad_output_b.reshape(-1, grad_output_b.shape[-1]) + grad_fp8_b, grad_scale_b = _quantize_fp8_tw(grad_b_2d, ctx.fp8_bwd_dtype) + + grad_input_b = gemm_fp8_impl( + grad_fp8_b, + grad_scale_b, + False, + w_t_fp8_b, + w_scale_b, + True, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + grad_input_b = grad_input_b.reshape(ctx.orig_shape_b) + + grad_t_fp8_b = grad_fp8_b.t().contiguous() + grad_weight_b = gemm_fp8_impl( + grad_t_fp8_b, + grad_scale_b, + False, + a_t_fp8_b, + a_scale_b, + True, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + + return ( + grad_input_a, + grad_weight_a, + None, + None, + grad_input_b, + grad_weight_b, + None, + None, + None, + None, + None, + None, + ) + + +@torch._dynamo.allow_in_graph +class FP8LinearRowwiseFunction(torch.autograd.Function): + """FP8 linear (Y = X @ W^T) with rowwise scaling. Compile-friendly.""" + + @staticmethod + def forward(ctx, input, weight, config): + a_dtype = _get_fp8_dtype(config.format, is_fwd=True) + b_dtype = _get_fp8_dtype(config.format, is_fwd=True) + out_dtype = input.dtype + gran = config.granularity + + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + a_fp8_row, a_scale_inv_row = quantize_fp8(input_2d, a_dtype, gran, axis=-1) + b_fp8_row, b_scale_inv_row = quantize_fp8(weight, b_dtype, gran, axis=-1) + + output = gemm_fp8_impl( + a_fp8_row, + a_scale_inv_row, + False, + b_fp8_row, + b_scale_inv_row, + True, + out_dtype, + False, + granularity=gran.value, + default_backend=BackendType.CK.value, + ) + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + + a_fp8_col, a_scale_inv_col = quantize_fp8(input_2d, a_dtype, gran, axis=-2) + b_fp8_col, b_scale_inv_col = quantize_fp8(weight, b_dtype, gran, axis=-2) + + ctx.save_for_backward(a_fp8_col, a_scale_inv_col, b_fp8_col, b_scale_inv_col) + ctx.out_dtype = out_dtype + ctx.config = config + ctx.orig_shape = orig_shape + return output + + @staticmethod + def backward(ctx, grad_output): + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + a_fp8_col, a_scale_inv_col, b_fp8_col, b_scale_inv_col = ctx.saved_tensors + grad_dtype = _get_fp8_dtype(ctx.config.format, is_fwd=False) + gran = ctx.config.granularity + + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + + grad_fp8_row, grad_scale_inv_row = quantize_fp8(grad_2d, grad_dtype, gran, axis=-1) + + grad_input = gemm_fp8_impl( + grad_fp8_row, + grad_scale_inv_row, + False, + b_fp8_col, + b_scale_inv_col, + False, + ctx.out_dtype, + False, + granularity=gran.value, + default_backend=BackendType.CK.value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_fp8_col, grad_scale_inv_col = quantize_fp8(grad_2d, grad_dtype, gran, axis=-2) + + grad_weight = gemm_fp8_impl( + a_fp8_col, + a_scale_inv_col, + True, + grad_fp8_col, + grad_scale_inv_col, + False, + ctx.out_dtype, + True, + granularity=gran.value, + default_backend=BackendType.CK.value, + ) + + return grad_input, grad_weight, None + + +@torch._dynamo.allow_in_graph +class FP8LinearBlockwiseFunction(torch.autograd.Function): + """FP8 linear (Y = X @ W^T) with blockwise scaling. Compile-friendly. + + Note: saves BF16 input AND BF16 weight for backward (must re-quantize + with different axis/dtype), so no activation memory savings compared to + BF16 baseline. + + Forward uses CK backend (NT layout). Backward uses Triton backend for + NN (grad_input) and TN (grad_weight) layouts because CK's blockwise + GEMM produces NaN on these layouts. To satisfy Triton's same-dtype + constraint, the weight is re-quantized to the backward dtype (E5M2 for + hybrid format) from the saved BF16 copy. + """ + + @staticmethod + def forward(ctx, input, weight, config): + a_dtype = _get_fp8_dtype(config.format, is_fwd=True) + b_dtype = _get_fp8_dtype(config.format, is_fwd=True) + out_dtype = input.dtype + gran = config.granularity + bs = config.block_size + + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + a_fp8_row, a_scale_inv_row = quant_fp8_blockwise_impl(input_2d, a_dtype, axis=1, block_size=bs) + b_fp8, b_scale_inv = quant_fp8_blockwise_for_weight_impl(weight, b_dtype, block_size=bs) + + output = gemm_fp8_impl( + a_fp8_row, + a_scale_inv_row, + False, + b_fp8, + b_scale_inv, + True, + out_dtype, + False, + granularity=gran.value, + default_backend=BackendType.CK.value, + ) + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + + ctx.save_for_backward(input_2d, weight) + ctx.out_dtype = out_dtype + ctx.config = config + ctx.orig_shape = orig_shape + return output + + @staticmethod + def backward(ctx, grad_output): + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + input_saved, weight_saved = ctx.saved_tensors + bwd_dtype = _get_fp8_dtype(ctx.config.format, is_fwd=False) + gran = ctx.config.granularity + bs = ctx.config.block_size + + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + + grad_fp8_row, grad_scale_inv_row = quant_fp8_blockwise_impl( + grad_2d, bwd_dtype, axis=-1, block_size=bs + ) + grad_fp8_col, grad_scale_inv_col = quant_fp8_blockwise_impl( + grad_2d, bwd_dtype, axis=-2, block_size=bs + ) + a_fp8_col, a_scale_inv_col = quant_fp8_blockwise_impl(input_saved, bwd_dtype, axis=0, block_size=bs) + b_fp8_bwd, b_scale_inv_bwd = quant_fp8_blockwise_for_weight_impl( + weight_saved, bwd_dtype, block_size=bs + ) + + grad_input = gemm_fp8_impl( + grad_fp8_row, + grad_scale_inv_row, + False, + b_fp8_bwd, + b_scale_inv_bwd, + False, + ctx.out_dtype, + False, + granularity=gran.value, + default_backend=BackendType.TRITON.value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_weight = gemm_fp8_impl( + a_fp8_col, + a_scale_inv_col, + True, + grad_fp8_col, + grad_scale_inv_col, + False, + ctx.out_dtype, + True, + granularity=gran.value, + default_backend=BackendType.TRITON.value, + ) + + return grad_input, grad_weight, None + + +# --------------------------------------------------------------------------- +# Delayed FP8 scaling helpers — shared by Float8ColumnParallelLinear and +# Float8RowParallelLinear. +# --------------------------------------------------------------------------- + + +def _update_fp8_scale(scale_buf, scale_inv_buf, amax, fp8_max): + """TE-aligned scale computation with edge case handling.""" + sf = fp8_max / amax.clamp(min=1e-12) + sf = torch.where(amax > 0.0, sf, scale_buf) + sf = torch.where(torch.isfinite(amax), sf, scale_buf) + sf = torch.where( + torch.isinf(sf), + torch.tensor(torch.finfo(torch.float32).max, device=sf.device), + sf, + ) + scale_buf.fill_(sf) + if scale_inv_buf is not None: + scale_inv_buf.fill_(1.0 / sf) + + +def _init_delayed_scaling_state(module): + """Initialize delayed FP8 scaling buffers on a Float8*ParallelLinear.""" + config = module.config + history_len = getattr(config, "fp8_amax_history_len", 1) + fp8_fwd_max = torch.finfo(module._fp8_fwd_dtype).max + fp8_bwd_max = torch.finfo(module._fp8_bwd_dtype).max + + for name in ["amax_history_input", "amax_history_weight", "amax_history_grad"]: + module.register_buffer( + name, + torch.zeros(history_len), + persistent=False, + ) + + for name in ["scale_input", "scale_weight", "scale_grad"]: + module.register_buffer( + name, + torch.tensor(1.0, dtype=torch.float32), + persistent=False, + ) + + for name in ["staged_input_amax", "staged_grad_amax", "staged_weight_amax"]: + module.register_buffer( + name, + torch.tensor(0.0), + persistent=False, + ) + + module._fp8_fwd_max = fp8_fwd_max + module._fp8_bwd_max = fp8_bwd_max + module._amax_compute_algo = getattr( + config, + "fp8_amax_compute_algo", + "most_recent", + ) + module._history_idx = 0 + module._first_delayed_step = True + + +def _batch_compute_scales(amaxes, fp8_max, old_scales): + """Vectorized scale computation across all modules at once.""" + sf = fp8_max / amaxes.clamp(min=1e-12) + sf = torch.where(amaxes > 0.0, sf, old_scales) + sf = torch.where(torch.isfinite(amaxes), sf, old_scales) + sf = torch.where(torch.isinf(sf), torch.finfo(torch.float32).max, sf) + return sf + + +# --------------------------------------------------------------------------- +# Fused Triton kernel: replaces Python-loop-based _batched_update per step +# with a single GPU dispatch over all (module, track) pairs. +# --------------------------------------------------------------------------- + + +@triton.jit +def _fused_delayed_scale_update_kernel( + amax_history_ptr, + staged_amaxes_ptr, + scales_ptr, + fp8_maxes_ptr, + history_idx, + N: tl.constexpr, + H: tl.constexpr, + use_max_algo: tl.constexpr, + BLOCK_H: tl.constexpr, + FP32_MAX: tl.constexpr, + FILTER_ZEROS: tl.constexpr = False, +): + track = tl.program_id(1) + mod = tl.program_id(0) + + base = track * N * H + mod * H + + new_amax = tl.load(staged_amaxes_ptr + track * N + mod) + tl.store(amax_history_ptr + base + history_idx, new_amax) + + if use_max_algo: + amax = tl.zeros([], dtype=tl.float32) + has_nonzero = tl.zeros([], dtype=tl.int32) + for off in range(0, H, BLOCK_H): + h_idx = off + tl.arange(0, BLOCK_H) + mask = h_idx < H + vals = tl.load(amax_history_ptr + base + h_idx, mask=mask, other=0.0) + if FILTER_ZEROS: + nonzero_mask = vals > 0.0 + has_nonzero = has_nonzero | tl.sum(nonzero_mask.to(tl.int32), axis=0) + filtered = tl.where(nonzero_mask, vals, 0.0) + amax = tl.maximum(amax, tl.max(filtered, axis=0)) + else: + amax = tl.maximum(amax, tl.max(vals, axis=0)) + else: + amax = new_amax + has_nonzero = tl.where(new_amax > 0.0, 1, 0).to(tl.int32) + + fp8_max = tl.load(fp8_maxes_ptr + track) + old_scale = tl.load(scales_ptr + track * N + mod) + sf = fp8_max / tl.maximum(amax, 1e-12) + if FILTER_ZEROS: + sf = tl.where(has_nonzero > 0, sf, old_scale) + else: + sf = tl.where(amax > 0.0, sf, old_scale) + sf = tl.where(amax == amax, sf, old_scale) + sf = tl.minimum(sf, FP32_MAX) + tl.store(scales_ptr + track * N + mod, sf) + + +# --------------------------------------------------------------------------- +# Global registry: replaces per-module scalar buffers with views into +# contiguous (N,) tensors to eliminate Python iteration in the preamble. +# --------------------------------------------------------------------------- + + +class _DelayedScalingRegistry: + __slots__ = ( + "n", + "modules", + "fwd_max", + "bwd_max", + "algo", + "history_len", + "_first_step", + "amax_history", + "_history_idx", + "fp8_maxes", + "staged_amaxes_3n", + "scales_3n", + "reduce_amax", + "amax_reduce_group", + "skip_bootstrap", + "filter_zeros", + ) + + def __init__(self, modules): + self.n = len(modules) + self.modules = modules + m0 = modules[0] + self.fwd_max = m0._fp8_fwd_max + self.bwd_max = m0._fp8_bwd_max + self.algo = m0._amax_compute_algo + self.history_len = m0.amax_history_input.shape[0] + device = m0.weight.device + + self._first_step = True + + H = self.history_len + N = self.n + + self.scales_3n = torch.ones(3, N, dtype=torch.float32, device=device) + self.staged_amaxes_3n = torch.zeros(3, N, dtype=torch.float32, device=device) + + self.amax_history = torch.zeros(3, N, H, dtype=torch.float32, device=device) + self._history_idx = 0 + self.fp8_maxes = torch.tensor( + [self.fwd_max, self.fwd_max, self.bwd_max], + dtype=torch.float32, + device=device, + ) + + _config = getattr(m0, "config", None) + self.reduce_amax = getattr(_config, "fp8_reduce_amax", False) if _config else False + self.skip_bootstrap = getattr(_config, "fp8_skip_first_step_bootstrap", False) if _config else False + self.filter_zeros = getattr(_config, "fp8_filter_zeros_in_history", False) if _config else False + self.amax_reduce_group = None + if self.reduce_amax: + from megatron.core import parallel_state + + if parallel_state.model_parallel_is_initialized(): + self.amax_reduce_group = parallel_state.get_amax_reduction_group( + with_context_parallel=True, + tp_only_amax_red=getattr(_config, "tp_only_amax_red", False), + ) + + for i, m in enumerate(modules): + m._buffers["scale_input"] = torch.ones((), dtype=torch.float32, device=device) + m._buffers["scale_weight"] = torch.ones((), dtype=torch.float32, device=device) + m._buffers["scale_grad"] = torch.ones((), dtype=torch.float32, device=device) + m._buffers["staged_input_amax"] = torch.zeros((), dtype=torch.float32, device=device) + m._buffers["staged_weight_amax"] = torch.zeros((), dtype=torch.float32, device=device) + m._buffers["staged_grad_amax"] = torch.zeros((), dtype=torch.float32, device=device) + m._buffers["amax_history_input"] = self.amax_history[0, i, :] + m._buffers["amax_history_weight"] = self.amax_history[1, i, :] + m._buffers["amax_history_grad"] = self.amax_history[2, i, :] + + +@torch.no_grad() +def _fast_update_scales(registry): + """Ultra-fast scale update for most_recent + history_len=1. + + Gathers per-module scalar amaxes into batched tensors, computes scales, + then scatters back. Each module has independent storage for its scale + and amax buffers so that torch.compile version tracking is correct. + """ + r = registry + + # Detect device migration AND per-module buffer pointer breakage (e.g. + # after _reset_fp8_local_spec during MLPerf warmup teardown). Mirrors the + # same check in _fast_update_scales_with_history so the warmup-reset + # docstring contract holds for history_len==1 configs too: a fresh + # registry.__init__(modules) call sets _first_step=True, which bootstraps + # weight amaxes from the restored (post-warmup) weights below. + if ( + r.scales_3n.device != r.modules[0].weight.device + or r.amax_history.untyped_storage().data_ptr() + != r.modules[0].amax_history_input.untyped_storage().data_ptr() + ): + r.__init__(r.modules) + + if r._first_step and not r.skip_bootstrap: + weights = [m.weight.data for m in r.modules] + try: + weight_amaxes = torch.stack(torch._foreach_norm(weights, ord=float("inf"))).float() + except (AttributeError, RuntimeError): + weight_amaxes = torch.stack([w.abs().amax().float() for w in weights]) + for i, m in enumerate(r.modules): + m.staged_weight_amax.fill_(weight_amaxes[i].item()) + if r._first_step: + r._first_step = False + + staged_input = torch.stack([m.staged_input_amax for m in r.modules]) + staged_weight = torch.stack([m.staged_weight_amax for m in r.modules]) + staged_grad = torch.stack([m.staged_grad_amax for m in r.modules]) + + if r.reduce_amax and r.amax_reduce_group is not None: + amaxes_3n = torch.stack([staged_input, staged_weight, staged_grad]) + torch.distributed.all_reduce( + amaxes_3n, + op=torch.distributed.ReduceOp.MAX, + group=r.amax_reduce_group, + ) + staged_input, staged_weight, staged_grad = amaxes_3n[0], amaxes_3n[1], amaxes_3n[2] + + scale_input = torch.stack([m.scale_input for m in r.modules]) + scale_weight = torch.stack([m.scale_weight for m in r.modules]) + scale_grad = torch.stack([m.scale_grad for m in r.modules]) + + new_input = _batch_compute_scales(staged_input, r.fwd_max, scale_input) + new_weight = _batch_compute_scales(staged_weight, r.fwd_max, scale_weight) + new_grad = _batch_compute_scales(staged_grad, r.bwd_max, scale_grad) + + # Batched device-only scatter: unbinding a (N,) float32 tensor into N + # 0-d views and copying with torch._foreach_copy_ keeps the per-module + # scale_* buffers independent (required for torch.compile version + # tracking) without the 3*N .item() host syncs the loop version + # incurred. + # + # NOTE (perf): the per-module scalar layout was chosen so that + # torch.compile sees each module's scale_* as an independent buffer + # version. A shared (N,) buffer with module-level views would let us + # skip the unbind+foreach_copy entirely, at the cost of more frequent + # compile cache invalidation. Worth benchmarking once the compile + # convergence work settles. + torch._foreach_copy_( + [m.scale_input for m in r.modules], + list(new_input.unbind()), + ) + torch._foreach_copy_( + [m.scale_weight for m in r.modules], + list(new_weight.unbind()), + ) + torch._foreach_copy_( + [m.scale_grad for m in r.modules], + list(new_grad.unbind()), + ) + + +@torch.no_grad() +def _fast_update_scales_with_history(registry): + """Fused scale update for delayed scaling with history_len > 1. + + Single Triton kernel dispatch replaces the prior Python-loop-based + per-module scale update (~200+ tiny GPU ops) with one launch. + """ + r = registry + + if ( + r.amax_history.device != r.modules[0].weight.device + or r.amax_history.untyped_storage().data_ptr() + != r.modules[0].amax_history_input.untyped_storage().data_ptr() + ): + r.__init__(r.modules) + + if r._first_step and not r.skip_bootstrap: + weights = [m.weight.data for m in r.modules] + try: + weight_amaxes = torch.stack(torch._foreach_norm(weights, ord=float("inf"))).float() + except (AttributeError, RuntimeError): + weight_amaxes = torch.stack([w.abs().amax().float() for w in weights]) + for i, m in enumerate(r.modules): + m.staged_weight_amax.fill_(weight_amaxes[i].item()) + if r._first_step: + r._first_step = False + + r.staged_amaxes_3n[0].copy_(torch.stack([m.staged_input_amax for m in r.modules])) + r.staged_amaxes_3n[1].copy_(torch.stack([m.staged_weight_amax for m in r.modules])) + r.staged_amaxes_3n[2].copy_(torch.stack([m.staged_grad_amax for m in r.modules])) + + if r.reduce_amax and r.amax_reduce_group is not None: + torch.distributed.all_reduce( + r.staged_amaxes_3n, + op=torch.distributed.ReduceOp.MAX, + group=r.amax_reduce_group, + ) + + N = r.n + H = r.history_len + BLOCK_H = triton.next_power_of_2(H) if H <= 1024 else 1024 + + _fused_delayed_scale_update_kernel[(N, 3)]( + r.amax_history, + r.staged_amaxes_3n, + r.scales_3n, + r.fp8_maxes, + r._history_idx, + N=N, + H=H, + use_max_algo=(r.algo == "max"), + BLOCK_H=BLOCK_H, + FP32_MAX=torch.finfo(torch.float32).max, + FILTER_ZEROS=r.filter_zeros, + ) + + # See _fast_update_scales for the rationale behind _foreach_copy_ + + # unbind here. r.scales_3n is the (3, N) registry-batched scale tensor + # produced by the Triton kernel; rows 0/1/2 = input/weight/grad. + torch._foreach_copy_( + [m.scale_input for m in r.modules], + list(r.scales_3n[0].unbind()), + ) + torch._foreach_copy_( + [m.scale_weight for m in r.modules], + list(r.scales_3n[1].unbind()), + ) + torch._foreach_copy_( + [m.scale_grad for m in r.modules], + list(r.scales_3n[2].unbind()), + ) + + r._history_idx = (r._history_idx + 1) % H + for m in r.modules: + m._history_idx = r._history_idx + + +# --------------------------------------------------------------------------- +# Async allreduce split: stage+launch / wait+compute +# --------------------------------------------------------------------------- + + +@torch.no_grad() +def _stage_and_launch_async_allreduce(registry): + """Stage per-module amaxes into r.staged_amaxes_3n and launch async allreduce. + + Returns the async work handle (or None if allreduce is not needed). + Must be called after backward completes so staged_*_amax buffers are + populated. + """ + r = registry + r.staged_amaxes_3n[0].copy_(torch.stack([m.staged_input_amax for m in r.modules])) + r.staged_amaxes_3n[1].copy_(torch.stack([m.staged_weight_amax for m in r.modules])) + r.staged_amaxes_3n[2].copy_(torch.stack([m.staged_grad_amax for m in r.modules])) + + if r.reduce_amax and r.amax_reduce_group is not None: + handle = torch.distributed.all_reduce( + r.staged_amaxes_3n, + op=torch.distributed.ReduceOp.MAX, + group=r.amax_reduce_group, + async_op=True, + ) + return handle + return None + + +@torch.no_grad() +def _wait_and_compute_scales(registry, handle): + """Wait on async allreduce handle and compute new scales from reduced amaxes. + + For most_recent + history_len=1, computes scales directly from + staged_amaxes_3n. For history-based, dispatches the fused Triton kernel. + Handles weight caching if enabled. + """ + r = registry + + if handle is not None: + handle.wait() + + if r.algo == "most_recent" and r.history_len == 1: + staged_input = r.staged_amaxes_3n[0] + staged_weight = r.staged_amaxes_3n[1] + staged_grad = r.staged_amaxes_3n[2] + + scale_input = torch.stack([m.scale_input for m in r.modules]) + scale_weight = torch.stack([m.scale_weight for m in r.modules]) + scale_grad = torch.stack([m.scale_grad for m in r.modules]) + + new_input = _batch_compute_scales(staged_input, r.fwd_max, scale_input) + new_weight = _batch_compute_scales(staged_weight, r.fwd_max, scale_weight) + new_grad = _batch_compute_scales(staged_grad, r.bwd_max, scale_grad) + + for i, m in enumerate(r.modules): + m.scale_input.fill_(new_input[i].item()) + m.scale_weight.fill_(new_weight[i].item()) + m.scale_grad.fill_(new_grad[i].item()) + else: + N = r.n + H = r.history_len + BLOCK_H = triton.next_power_of_2(H) if H <= 1024 else 1024 + + _fused_delayed_scale_update_kernel[(N, 3)]( + r.amax_history, + r.staged_amaxes_3n, + r.scales_3n, + r.fp8_maxes, + r._history_idx, + N=N, + H=H, + use_max_algo=(r.algo == "max"), + BLOCK_H=BLOCK_H, + FP32_MAX=torch.finfo(torch.float32).max, + FILTER_ZEROS=r.filter_zeros, + ) + + for i, m in enumerate(r.modules): + m.scale_input.fill_(r.scales_3n[0, i].item()) + m.scale_weight.fill_(r.scales_3n[1, i].item()) + m.scale_grad.fill_(r.scales_3n[2, i].item()) + + r._history_idx = (r._history_idx + 1) % H + for m in r.modules: + m._history_idx = r._history_idx + + +# --------------------------------------------------------------------------- +# Weight extraction helper — called outside autograd.Function to avoid +# tensor-subclass isinstance checks inside the compiled graph. +# --------------------------------------------------------------------------- + + +def _extract_fp8_weight(weight, fp8_dtype): + """Return (fp8_data, scale_inv) from weight, handling FP8UnshardedWeightTensor.""" + from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + FP8UnshardedWeightTensor, + ) + + if isinstance(weight, FP8UnshardedWeightTensor): + return weight.get_fp8_data_and_scale_inv() + return _quantize_fp8_tw(weight, fp8_dtype) + + +# --------------------------------------------------------------------------- +# Granularity -> Function dispatch map +# --------------------------------------------------------------------------- + +_FP8_FN_MAP = { + ScalingGranularity.TENSORWISE: OpaqueFP8LinearTensorwiseFunction, + ScalingGranularity.ROWWISE: FP8LinearRowwiseFunction, + ScalingGranularity.BLOCKWISE: FP8LinearBlockwiseFunction, +} + +# Public alias +FP8LinearTensorwiseFunction = OpaqueFP8LinearTensorwiseFunction + + +# --------------------------------------------------------------------------- +# FP8-aware parallel linear layers +# --------------------------------------------------------------------------- + + +class _Float8LinearMixin: + """Shared per-module FP8 behavior for Float8{Column,Row}ParallelLinear. + + This is a mixin (not an nn.Module). Each concrete class also inherits a + Megatron ``*ParallelLinear`` and MUST list this mixin FIRST in its bases so + the ``_apply`` / ``_forward_impl`` overrides below take precedence over the + Megatron base (otherwise FP8 would be silently disabled). The concrete + ``__init__`` calls ``self._init_fp8_state()`` after ``super().__init__()``. + + Requires: tensor_model_parallel_size=1, gradient_accumulation_fusion=False, + sequence_parallel=False. Fails fast at construction if violated. + + Tensorwise uses the setup_context autograd pattern with pre-extracted FP8 + weight data for graph-break-free torch.compile tracing. + Rowwise and blockwise still use @allow_in_graph. + """ + + def _init_fp8_state(self): + cls = type(self).__name__ + if self.config.tensor_model_parallel_size != 1: + raise ValueError( + f"{cls} requires tensor_model_parallel_size=1. " + f"Got {self.config.tensor_model_parallel_size}." + ) + if self.gradient_accumulation_fusion: + raise ValueError(f"{cls} requires gradient_accumulation_fusion=False.") + if self.sequence_parallel: + raise ValueError(f"{cls} requires sequence_parallel=False.") + if self.config.fp8 is None: + raise ValueError(f"{cls} requires config.fp8 to be set (e.g. 'e4m3').") + + self._fp8_config = _build_fp8_config(self.config) + self._fp8_fn = _FP8_FN_MAP[self._fp8_config.granularity] + + if self._fp8_config.granularity == ScalingGranularity.TENSORWISE: + self._fp8_fwd_dtype = _get_fp8_dtype(self._fp8_config.format, is_fwd=True) + self._fp8_bwd_dtype = _get_fp8_dtype(self._fp8_config.format, is_fwd=False) + self._fp8_gran_value = ScalingGranularity.TENSORWISE.value + self._fp8_backend_value = BackendType.HIPBLASLT.value + # Native layout (NN/TN) by default; forced-NT is opt-in per config. + # Read once at layer init so it stays a compile-time constant. + self._force_nt = getattr(self.config, "fp8_force_nt_layout", False) + self._use_delayed_scaling = ( + getattr(self.config, "fp8_scaling_strategy", "dynamic") == "delayed" + or self.config.fp8_recipe == Fp8Recipe.delayed + ) + if self._use_delayed_scaling: + _init_delayed_scaling_state(self) + + def _apply(self, fn, recurse=True): + result = super()._apply(fn, recurse) + if getattr(self, "_use_delayed_scaling", False): + for name in [ + "scale_input", + "scale_weight", + "scale_grad", + "amax_history_input", + "amax_history_weight", + "amax_history_grad", + "staged_input_amax", + "staged_grad_amax", + "staged_weight_amax", + ]: + buf = getattr(self, name, None) + if buf is not None and buf.dtype != torch.float32: + self._buffers[name] = buf.data.float() + return result + + def _forward_impl(self, input, weight, *args, **kwargs): + bias = kwargs.get("bias", None) + + if self._fp8_config.granularity == ScalingGranularity.TENSORWISE: + if self._use_delayed_scaling: + result = DelayedFP8LinearTensorwiseFunction.apply( + input, + weight, + self.scale_input, + self.scale_weight, + self.scale_grad, + self.staged_input_amax, + self.staged_weight_amax, + self.staged_grad_amax, + self._fp8_fwd_dtype, + self._fp8_bwd_dtype, + self._fp8_gran_value, + self._fp8_backend_value, + self._force_nt, + ) + output = result[0] + else: + weight_fp8, weight_scale_inv = _extract_fp8_weight( + weight, + self._fp8_fwd_dtype, + ) + result = self._fp8_fn.apply( + input, + weight, + weight_fp8, + weight_scale_inv, + self._fp8_fwd_dtype, + self._fp8_bwd_dtype, + self._fp8_gran_value, + self._fp8_backend_value, + self._force_nt, + ) + output = result[0] + else: + output = self._fp8_fn.apply(input, weight, self._fp8_config) + + if bias is not None: + output = output + bias + return output + + +class Float8ColumnParallelLinear(_Float8LinearMixin, ColumnParallelLinear): + """ColumnParallelLinear with per-module FP8. torch.compile friendly. + + Shared FP8 behavior lives in :class:`_Float8LinearMixin`, which is listed + first so its ``_apply`` / ``_forward_impl`` overrides take effect. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._init_fp8_state() + + +class Float8RowParallelLinear(_Float8LinearMixin, RowParallelLinear): + """RowParallelLinear with per-module FP8. torch.compile friendly. + + Shared FP8 behavior lives in :class:`_Float8LinearMixin`, which is listed + first so its ``_apply`` / ``_forward_impl`` overrides take effect. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._init_fp8_state() diff --git a/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py b/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py new file mode 100644 index 000000000..c800ac54b --- /dev/null +++ b/primus/backends/megatron/core/extensions/primus_turbo_local_spec.py @@ -0,0 +1,318 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Primus Turbo Local Spec Provider + +Clean implementation using native Megatron modules with Primus Turbo +optimizations. NO TransformerEngine dependencies. + +Key Features: +- PrimusTurboLocalAttention: New class, inherits from MegatronModule (not TE) +- Native Megatron linear layers (ColumnParallelLinear, RowParallelLinear) +- Native LayerNorm (WrappedTorchNorm or FusedLayerNorm) +- Maximum torch.compile compatibility +- 52ms attention advantage preserved via pt.ops.flash_attn_func + +Performance Target (vs Inductor baseline): +- Attention: Same (52ms advantage over Inductor's flash_attn) +- GEMM: -7ms (native PyTorch proven within 0.8% of TE) +- Elementwise: -157ms (torch.compile fusion) +- Reduce: -17ms (torch.compile fusion) +- Framework: -107ms (no TE overhead) +- Total: ~1,131ms vs Inductor's 1,415ms = 20% faster +""" + +import math +import os +from typing import Optional + +import primus_turbo.pytorch as pt +import torch +from primus_turbo.pytorch.ops.attention.flash_attn_interface import AiterFlashAttnFunc + +torch._dynamo.allow_in_graph(AiterFlashAttnFunc) + + +@torch._dynamo.disable +def _advance_model_parallel_rng(batch_size: int, num_heads: int) -> None: + """Simulate TE's CUDA RNG consumption inside + PrimusTurboLocalAttention. + + TEDotProductAttention is instantiated with + get_rng_state_tracker=get_cuda_rng_tracker, so TE's fused_attn_fwd kernel + call runs inside tracker.fork() (the 'model-parallel-rng' generator state). + The ASM kernel pulls a fresh philox state per call regardless of dropout, + advancing model-parallel-rng by counter_offset = B*H*warp_size each + attention call. Aiter's _ndropout path does not touch RNG, so PT and + aiter attention see different downstream RNG sequences and diverge. + + This helper replays the same RNG advance. Decorated with @dynamo.disable + so AOTAutograd never tries to trace through the generator manipulation. + """ + warp_size = 64 + counter_offset = batch_size * num_heads * warp_size + dev_idx = torch.cuda.current_device() + gens = torch.cuda.default_generators + if dev_idx >= len(gens): + return + gen = gens[dev_idx] + + def _advance_with(g): + try: + g.set_offset(g.get_offset() + counter_offset) + except AttributeError: + torch.empty(counter_offset, device="cuda", dtype=torch.int32).random_(generator=g) + + try: + from megatron.core.tensor_parallel.random import get_cuda_rng_tracker + + tracker = get_cuda_rng_tracker() + except Exception: + tracker = None + + if tracker is not None and getattr(tracker, "is_initialized", lambda: False)(): + with tracker.fork(): + _advance_with(gen) + else: + _advance_with(gen) + + +from megatron.core.models.backends import LocalSpecProvider +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.parallel_state import ( + get_context_parallel_group, + get_tensor_model_parallel_group, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.training.global_vars import get_args +from torch import Tensor + + +class PrimusTurboLocalAttention(MegatronModule): + """ + Primus Turbo Flash Attention - Native Megatron Implementation. + + This is a COMPLETELY NEW class with ZERO TransformerEngine dependencies. + + Key differences from PrimusTurboAttention (in primus_turbo.py): + - Inherits from MegatronModule (NOT te.pytorch.DotProductAttention) + - No TE infrastructure or overhead + - torch.compile friendly (no @no_torch_dynamo decorator) + - Same 52ms performance advantage via pt.ops.flash_attn_func + + Used exclusively by: PrimusTurboLocalSpecProvider + + Performance: + - 52ms faster than Inductor's flash_attn implementation + - Same as existing PrimusTurboAttention but without TE overhead + - torch.compile compatible for surrounding operations + """ + + def __init__( + self, + config: TransformerConfig, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + attention_dropout: Optional[float] = None, + softmax_scale: Optional[float] = None, + k_channels: Optional[int] = None, + v_channels: Optional[int] = None, + cp_comm_type: str = "p2p", + pg_collection: Optional[ProcessGroupCollection] = None, + ): + super().__init__(config=config) + + self.config = config + self.layer_number = layer_number + self.attn_mask_type = attn_mask_type + self.attention_type = attention_type + + # Calculate softmax scale + kv_channels = k_channels if k_channels is not None else config.kv_channels + self.softmax_scale = softmax_scale or (1.0 / math.sqrt(kv_channels)) + + # Setup process groups + if pg_collection is None: + pg_collection = ProcessGroupCollection( + tp=get_tensor_model_parallel_group(check_initialized=False), + cp=get_context_parallel_group(check_initialized=False), + ) + + # Select Primus Turbo flash attention variant + args = get_args() + if args.enable_turbo_attention_float8: + self.attn_func = ( + pt.ops.flash_attn_fp8_usp_func + if config.context_parallel_size > 1 + else pt.ops.flash_attn_fp8_func + ) + else: + self.attn_func = ( + pt.ops.flash_attn_usp_func if config.context_parallel_size > 1 else pt.ops.flash_attn_func + ) + + # The transpose in forward() produces a non-contiguous, SBHD-strided view. + # aiter's v3 flash-attention backward mishandles that strided layout on + # gfx942 (MI300X), corrupting gradients and causing grad-norm divergence + # ~step 31; gfx950 (MI355X) handles it correctly. Feed contiguous BSHD on + # gfx942 only, so gfx950 keeps the faster zero-copy strided path unchanged. + # Computed here (not in forward) so torch.compile sees a constant guard + # rather than a device query in the traced region. + self.force_contiguous_qkv = torch.cuda.get_device_capability() < (9, 5) + + # Setup context parallel arguments + self.attn_kwargs = {} + if config.context_parallel_size > 1: + self.attn_kwargs["ulysses_group"] = pg_collection.cp + + # Validate configuration + if config.window_size is not None: + raise ValueError("PrimusTurboLocalAttention does not support sliding window attention") + + def forward( + self, + query: Tensor, + key: Tensor, + value: Tensor, + attention_mask: Tensor, + attn_mask_type: AttnMaskType, + attention_bias: Optional[Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ) -> Tensor: + """ + Forward pass using Primus Turbo flash attention. + + Args: + query: Query tensor [seq_len, batch, num_heads, head_dim] (sbhd) + key: Key tensor [seq_len, batch, num_heads, head_dim] (sbhd) + value: Value tensor [seq_len, batch, num_heads, head_dim] (sbhd) + attention_mask: Attention mask (not used by flash attention) + attn_mask_type: Type of attention mask (causal, no_mask, etc.) + attention_bias: Attention bias (not used in this implementation) + packed_seq_params: Packed sequence parameters (optional) + + Returns: + Attention output [seq_len, batch, num_heads * head_dim] (merged heads) + """ + query, key, value = [x.transpose(0, 1) for x in (query, key, value)] + + # gfx942: avoid aiter's broken strided-sbhd backward (see __init__). + if self.force_contiguous_qkv: + query, key, value = query.contiguous(), key.contiguous(), value.contiguous() + + causal = attn_mask_type == AttnMaskType.causal + + if os.environ.get("PRIMUS_PT_MIMIC_TE_RNG", "0") == "1": + B = query.size(0) + H = query.size(2) + _advance_model_parallel_rng(B, H) + + output = self.attn_func( + query, + key, + value, + dropout_p=0.0, + softmax_scale=self.softmax_scale, + causal=causal, + window_size=(-1, -1), + bias=None, + alibi_slopes=None, + deterministic=False, + return_lse=False, + return_attn_probs=False, + **self.attn_kwargs, + ) + + # Transpose back to Megatron format (bshd -> sbhd) and merge heads + output = output.transpose(0, 1) + output = output.reshape(output.shape[0], output.shape[1], -1) + + return output + + +class PrimusTurboMXFP4LocalSpecProvider(LocalSpecProvider): + """ + Compile-friendly MXFP4 spec: Primus Turbo attention + MXFP4 linear layers. + NO TransformerEngine. NO global FP4 state in forward path. + Requires tensor_model_parallel_size=1. + """ + + def column_parallel_linear(self) -> type: + from .primus_turbo_mxfp4_local import MXFP4ColumnParallelLinear + + return MXFP4ColumnParallelLinear + + def row_parallel_linear(self) -> type: + from .primus_turbo_mxfp4_local import MXFP4RowParallelLinear + + return MXFP4RowParallelLinear + + def core_attention(self) -> type: + return PrimusTurboLocalAttention + + +class PrimusTurboFloat8LocalSpecProvider(LocalSpecProvider): + """ + Compile-friendly FP8 spec: Primus Turbo attention + FP8 linear layers. + NO TransformerEngine. NO global FP8 state in forward path. + NO changes to Primus-Turbo repo. Requires tensor_model_parallel_size=1. + """ + + def column_parallel_linear(self) -> type: + from .primus_turbo_float8_local import Float8ColumnParallelLinear + + return Float8ColumnParallelLinear + + def row_parallel_linear(self) -> type: + from .primus_turbo_float8_local import Float8RowParallelLinear + + return Float8RowParallelLinear + + def core_attention(self) -> type: + return PrimusTurboLocalAttention + + +class PrimusTurboLocalSpecProvider(LocalSpecProvider): + """ + Spec provider that extends LocalSpecProvider with Primus Turbo attention. + + Uses native Megatron modules everywhere except attention, where it uses + PrimusTurboLocalAttention for the 52ms performance advantage. + + Key features: + - NO TransformerEngine dependencies + - Maximum torch.compile compatibility + - Minimal custom code (~120 lines total) + - Native PyTorch GEMM (proven within 0.8% of TE) + - Primus Turbo flash attention (52ms advantage) + + Configuration: + config.use_primus_turbo_local_spec = True + + Performance Target: + - Attention: -52ms (Primus Turbo advantage over Inductor) + - Operator fusion: -180ms (torch.compile) + - Framework overhead: -107ms (no TE) + - GEMM: ~same (native PyTorch excellent) + - Total: ~1,131ms vs Inductor's 1,415ms (20% faster) + + vs Current TE implementation: + - Total: ~1,131ms vs 1,588ms (28.8% faster) + """ + + def core_attention(self) -> type: + """Return Primus Turbo local attention (NO TE dependency)""" + return PrimusTurboLocalAttention + + # Everything else inherited from LocalSpecProvider: + # - column_parallel_linear() -> ColumnParallelLinear (native Megatron) + # - row_parallel_linear() -> RowParallelLinear (native Megatron) + # - layer_norm() -> WrappedTorchNorm/FusedLayerNorm (compile-friendly) + # - fuse_layernorm_and_linear() -> False (no fusion overhead) + # - grouped_mlp_modules() -> SequentialMLP (native modules) + # - activation_func() -> None (standard activation) diff --git a/primus/backends/megatron/core/fp8_utils.py b/primus/backends/megatron/core/fp8_utils.py index 0e7ce7995..18076c578 100644 --- a/primus/backends/megatron/core/fp8_utils.py +++ b/primus/backends/megatron/core/fp8_utils.py @@ -73,14 +73,19 @@ def te_fp8_format_mapping(te_format): return format_mapping[te_format] def get_fp8_recipe(config: TransformerConfig): - """Return fp8 recipe. + """Return fp8 recipe (Primus Turbo + TE variant). Arguments: config (TransformerConfig): Configuration object. Returns: - FP8 recipe: Transformer Engine FP8 recipe. - FP8 None reason: reason why the fp8 recipe is None. + Tuple of (fp8_recipe, fp8_recipe_none_reason): + - fp8_recipe: Transformer Engine FP8 recipe, or None. + - fp8_recipe_none_reason: reason why the fp8 recipe is None. + + Note: + This variant (HAVE_TE + HAVE_TURBO) returns a tuple. The TE-only + variant returns a single recipe value. """ if config.fp8 == "e4m3": fp8_format = transformer_engine.common.recipe.Format.E4M3 diff --git a/primus/backends/megatron/core/utils.py b/primus/backends/megatron/core/utils.py index a5855ad58..51d1897ec 100644 --- a/primus/backends/megatron/core/utils.py +++ b/primus/backends/megatron/core/utils.py @@ -1,26 +1,34 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### from functools import lru_cache +from typing import Any, List +import torch from megatron.core import parallel_state -from primus_turbo.pytorch.ops.attention.attention_utils import ( - All2AllAttentionSharder, - AttentionSharder, -) + +from primus.modules.module_utils import log_rank_0 @lru_cache def produce_attention_sharder(cp_comm_type: str): + # Import only when needed to avoid import errors if primus_turbo is not available + try: + from primus_turbo.pytorch.ops.attention.attention_utils import ( + All2AllAttentionSharder, + ) + except ImportError: + raise ImportError("All2AllAttentionSharder not available. Ensure primus_turbo is properly installed.") + if cp_comm_type == "a2a": return All2AllAttentionSharder() else: raise ValueError(f"Unsupported cp_comm_type: {cp_comm_type}") -def shard_batch_on_this_cp_rank(sharder: AttentionSharder, batch): +def shard_batch_on_this_cp_rank(sharder, batch): cp_size = parallel_state.get_context_parallel_world_size() cp_group = parallel_state.get_context_parallel_group() if cp_size > 1: @@ -29,3 +37,118 @@ def shard_batch_on_this_cp_rank(sharder: AttentionSharder, batch): seq_dim = 1 if key != "attention_mask" else 2 batch[key] = sharder.shard_cp_input([val], cp_group, seq_dim)[0] return batch + + +def apply_torch_compile_if_enabled(model: List[Any], args: Any) -> None: + """ + Apply torch.compile to model AFTER distributed wrapping and config setup. + + This must be called: + 1. AFTER get_model() completes (which does FSDP/DDP wrapping internally) + 2. AFTER ddp_config is set on the wrapped model + 3. BEFORE optimizer creation + + Args: + model: List of model modules (already wrapped in DDP/FSDP) + args: Megatron args object + + Raises: + Exception: If compilation fails, exception is raised (not caught) + """ + from megatron.training.utils import unwrap_model + + # Check if compilation is enabled + if not getattr(args, "enable_torch_compile", False): + return + + log_rank_0("=" * 80) + log_rank_0("Applying torch.compile AFTER distributed wrapping...") + log_rank_0("=" * 80) + + # Import torch_FSDP for type checking + try: + from primus.backends.megatron.core.distributed.torch_fully_sharded_data_parallel import ( + PrimusTorchFullyShardedDataParallel as torch_FSDP, + ) + except ImportError: + torch_FSDP = None + + for idx, model_module in enumerate(model): + # Check what type of wrapper we have + wrapper_type = type(model_module).__name__ + log_rank_0(f" Model [{idx}] wrapper: {wrapper_type}") + + # Check if wrapped with FSDP and has ddp_config + if torch_FSDP and isinstance(model_module, torch_FSDP): + has_config = hasattr(model_module, "ddp_config") + log_rank_0(f" FSDP detected - has ddp_config: {has_config}") + + # Check if FSDP wrapper has compile_model method + if hasattr(model_module, "compile_model"): + log_rank_0(f" Calling compile_model on FSDP wrapper...") + model_module.compile_model() + log_rank_0(f" ✓ FSDP wrapper compilation complete") + continue + + # Unwrap to get the actual model (Flux, GPT, etc.) + unwrapped_models = unwrap_model([model_module]) + + for unwrapped in unwrapped_models: + model_type = type(unwrapped).__name__ + + # Check if the model has a compile_model method + if hasattr(unwrapped, "compile_model"): + log_rank_0(f" Compiling {model_type}...") + + # Apply compilation + unwrapped.compile_model() + + log_rank_0(f" ✓ {model_type} compilation complete") + else: + log_rank_0(f" ℹ {model_type} does not support torch.compile " f"(no compile_model method)") + + log_rank_0("torch.compile application complete") + log_rank_0("=" * 80) + + +def apply_torch_compile_to_optimizer_if_enabled(optimizer: Any, args: Any) -> None: + """ + Apply torch.compile to the optimizer's step() when enable_torch_compile is True. + Replaces optimizer.step with a compiled version so the trainer needs no changes. + Uses fullgraph=False for the optimizer step (recommended; allows timers/conditionals). + """ + if optimizer is None: + return + if not getattr(args, "enable_torch_compile", False): + return + if not getattr(args, "torch_compile_optimizer", False): + log_rank_0(" Optimizer compilation disabled (torch_compile_optimizer=False)") + return + + backend = getattr(args, "torch_compile_backend", "inductor") + mode = getattr(args, "torch_compile_mode", "default") + fullgraph = False # Optimizer step: allow graph breaks (timers, conditionals) + + compile_kwargs = { + "backend": backend, + "mode": mode, + "fullgraph": fullgraph, + } + + scope = getattr(args, "torch_compile_optimizer_scope", "full") + log_rank_0(f"Applying torch.compile to optimizer step (scope={scope})...") + log_rank_0(f" backend={backend}, mode={mode}, fullgraph={fullgraph}") + + if scope == "inner_only": + inner_opt = getattr(optimizer, "optimizer", None) + if inner_opt is None: + log_rank_0(" ⚠ No inner optimizer found, falling back to full scope") + scope = "full" + else: + inner_opt.step = torch.compile(inner_opt.step, **compile_kwargs) + log_rank_0(" ✓ Inner optimizer step compiled (clip_grad_norm/DTensor ops excluded)") + + if scope == "full": + original_step = optimizer.step + optimizer.step = torch.compile(original_step, **compile_kwargs) + log_rank_0(" ✓ Optimizer step compiled (full scope)") diff --git a/tests/unit_tests/backends/megatron/test_native_fp8_layout.py b/tests/unit_tests/backends/megatron/test_native_fp8_layout.py new file mode 100644 index 000000000..655071b59 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_native_fp8_layout.py @@ -0,0 +1,184 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""GPU-only tests for the native (NN/TN) FP8 backward layout. + +FP8 GEMMs require an AMD GPU (gfx950); the fused no-transpose cast+amax kernel +is vendored in Primus (fp8_cast_kernels_triton), so every test here is skipped +when CUDA is unavailable. Run on an AMD GPU (gfx950) with the FP8 toolchain available. + +Coverage: + - native (``fp8_force_nt_layout=False``) vs legacy forced-NT numerical + equivalence for both forward and backward, + - finiteness of the weight gradient (the wgrad-NaN regression that motivated + the native layout), + - backward gradient flow / return-count validation for both + ``OpaqueFP8LinearTensorwiseFunction`` (9 inputs) and the full delayed-scaling + layer (``DelayedFP8LinearTensorwiseFunction``, 13 inputs); a count mismatch + would make ``autograd`` raise, so a successful backward validates the arity. +""" + +import functools +import os +from types import SimpleNamespace + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from megatron.core.transformer.transformer_config import TransformerConfig + +from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + Float8ColumnParallelLinear, + OpaqueFP8LinearTensorwiseFunction, +) +from tests.utils import PrimusUT + +requires_gpu = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="Native FP8 layout exercises FP8 GEMMs, which require an AMD GPU", +) + + +def _init_method(): + return functools.partial(torch.nn.init.xavier_uniform_) + + +def _make_config(force_nt, recipe="tensorwise"): + """TransformerConfig for an FP8 linear, with the diffusion-only + fp8_force_nt_layout attribute attached (the layer reads it via getattr).""" + config = TransformerConfig( + hidden_size=64, + num_attention_heads=8, + num_layers=1, + params_dtype=torch.bfloat16, + fp8="e4m3", + fp8_recipe=recipe, + ) + config.fp8_force_nt_layout = force_nt + return config + + +def _build_linear(force_nt, recipe="tensorwise"): + return Float8ColumnParallelLinear( + input_size=64, + output_size=128, + config=_make_config(force_nt, recipe=recipe), + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + + +def _forward_only(layer, x): + out = layer(x) + return out[0] if isinstance(out, tuple) else out + + +class _GpuLinearBase(PrimusUT): + """Shared setup: parallel state + a minimal global-args stub for the layers.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, monkeypatch): + dummy_args = SimpleNamespace( + rank=0, + world_size=1, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + offload=False, + offload_ops=[], + patch_primus_pipeline=False, + pp_algorithm=None, + patch_zero_bubble=False, + enable_zero_bubble=False, + rampup_batch_size=None, + global_batch_size=1, + micro_batch_size=1, + data_parallel_size=1, + decrease_batch_size_if_needed=False, + ) + import megatron.training.global_vars as gvars + + monkeypatch.setattr(gvars, "_GLOBAL_ARGS", dummy_args) + + # Some deploy images bake PRIMUS_TURBO_GEMM_BACKEND="" which makes + # GlobalBackendManager raise KeyError('') on the native path. Unset it + # here, exactly as the docs require for native-layout production runs. + if os.environ.get("PRIMUS_TURBO_GEMM_BACKEND", None) == "": + monkeypatch.delenv("PRIMUS_TURBO_GEMM_BACKEND", raising=False) + + +class TestNativeVsForcedNtLayer(_GpuLinearBase): + """End-to-end equivalence of the native and forced-NT layouts via the layer.""" + + @requires_gpu + def test_forward_backward_match_and_finite(self): + # NOTE: looped instead of @pytest.mark.parametrize because PrimusUT is a + # unittest.TestCase, where pytest does not inject parametrized args. + for recipe in ("tensorwise", "delayed"): + with self.subTest(recipe=recipe): + torch.manual_seed(0) + + native = _build_linear(force_nt=False, recipe=recipe).cuda() + forced = _build_linear(force_nt=True, recipe=recipe).cuda() + # Identical weights so the only difference is the backward GEMM layout. + forced.load_state_dict(native.state_dict()) + + assert native._force_nt is False + assert forced._force_nt is True + + x = torch.randn(8, 64, dtype=torch.bfloat16, device="cuda") + x_native = x.clone().requires_grad_(True) + x_forced = x.clone().requires_grad_(True) + + out_native = _forward_only(native, x_native) + out_forced = _forward_only(forced, x_forced) + + assert torch.isfinite(out_native).all(), "native forward produced non-finite output" + torch.testing.assert_close(out_native, out_forced, atol=2e-2, rtol=2e-2) + + out_native.sum().backward() + out_forced.sum().backward() + + # wgrad NaN regression guard + native/forced equivalence. + assert torch.isfinite(native.weight.grad).all(), "native wgrad is non-finite (NaN regression)" + assert torch.isfinite(x_native.grad).all(), "native dgrad is non-finite" + torch.testing.assert_close(native.weight.grad, forced.weight.grad, atol=3e-2, rtol=3e-2) + torch.testing.assert_close(x_native.grad, x_forced.grad, atol=3e-2, rtol=3e-2) + + +class TestOpaqueFunctionBackward(_GpuLinearBase): + """Direct autograd-Function checks for OpaqueFP8LinearTensorwiseFunction.""" + + @requires_gpu + def test_native_backward_returns_grads_for_input_and_weight(self): + from primus_turbo.pytorch.core.backend import BackendType + from primus_turbo.pytorch.core.low_precision import ( + ScalingGranularity, + float8_e4m3, + float8_e5m2, + ) + from primus_turbo.pytorch.ops.quantization import quantize_fp8 + + gran = ScalingGranularity.TENSORWISE.value + backend = BackendType.HIPBLASLT.value + + x = torch.randn(8, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = torch.randn(128, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w_fp8, w_scale = quantize_fp8(w, float8_e4m3, ScalingGranularity.TENSORWISE) + + # force_nt=False exercises the native arm; backward must return exactly + # 9 grads (one per forward input) or autograd raises here. + out = OpaqueFP8LinearTensorwiseFunction.apply( + x, w, w_fp8, w_scale, float8_e4m3, float8_e5m2, gran, backend, False + ) + output = out[0] if isinstance(out, tuple) else out + output.sum().backward() + + assert x.grad is not None and torch.isfinite(x.grad).all() + assert w.grad is not None and torch.isfinite(w.grad).all() diff --git a/tests/unit_tests/backends/megatron/test_primus_turbo_float8_local.py b/tests/unit_tests/backends/megatron/test_primus_turbo_float8_local.py new file mode 100644 index 000000000..0ad1c316c --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_primus_turbo_float8_local.py @@ -0,0 +1,431 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for compile-friendly FP8 linear layers (primus_turbo_float8_local). + +Tests _build_fp8_config mapping, Float8ColumnParallelLinear/Float8RowParallelLinear +construction guards, init-time dispatch to the correct autograd Function, +decomposed tensorwise quantize numerical equivalence, and torch.compile +graph-break validation. +""" + +import functools +from types import SimpleNamespace + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from megatron.core.enums import Fp8Recipe +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from megatron.core.transformer.transformer_config import TransformerConfig +from primus_turbo.pytorch.core.low_precision import ( + ScalingGranularity, + float8_e4m3, + float8_e5m2, +) + +from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + DecomposedFP8LinearTensorwiseFunction, + Float8ColumnParallelLinear, + Float8RowParallelLinear, + FP8LinearBlockwiseFunction, + OpaqueFP8LinearTensorwiseFunction, + _build_fp8_config, + _Float8LinearMixin, + _quantize_fp8_tensorwise, +) +from tests.utils import PrimusUT + + +class TestFloat8MixinStructure: + """Structural guards for the shared _Float8LinearMixin (no construction). + + These do not build modules, so they exercise the class layout rather than + FP8 numerics. They catch the load-bearing footgun where the mixin is not + listed first in the bases: in that case the Megatron base's ``_apply`` / + ``_forward_impl`` would win and FP8 would be silently disabled. The + init-time tests below would NOT catch that, because ``__init__`` calls + ``self._init_fp8_state()`` explicitly regardless of MRO order. + + Note: this module imports primus_turbo (CUDA-at-import), so it is still + gated by the module-level skip_if_no_cuda() and runs in the GPU lane. + """ + + @pytest.mark.parametrize( + "cls, base", + [ + (Float8ColumnParallelLinear, ColumnParallelLinear), + (Float8RowParallelLinear, RowParallelLinear), + ], + ) + def test_mixin_precedes_base_in_mro(self, cls, base): + mro = cls.__mro__ + assert _Float8LinearMixin in mro + assert base in mro + assert mro.index(_Float8LinearMixin) < mro.index(base), ( + f"{cls.__name__}: _Float8LinearMixin must precede {base.__name__} in the MRO " + "or FP8 _apply/_forward_impl would be silently overridden by the base." + ) + + @pytest.mark.parametrize("cls", [Float8ColumnParallelLinear, Float8RowParallelLinear]) + def test_fp8_overrides_resolve_to_mixin(self, cls): + assert cls._forward_impl is _Float8LinearMixin._forward_impl + assert cls._apply is _Float8LinearMixin._apply + + +class TestBuildFP8Config: + """Tests for _build_fp8_config() — no GPU, no parallel state.""" + + def _make_config(self, fp8="e4m3", fp8_recipe=Fp8Recipe.tensorwise): + return SimpleNamespace(fp8=fp8, fp8_recipe=fp8_recipe) + + def test_invalid_recipe_raises(self): + with pytest.raises(ValueError, match="does not support"): + _build_fp8_config(self._make_config(fp8_recipe=Fp8Recipe.custom)) + + +def _init_method(): + return functools.partial(torch.nn.init.xavier_uniform_) + + +def _make_fp8_transformer_config(**overrides): + """Build a TransformerConfig suitable for Float8 linear layers.""" + defaults = dict( + hidden_size=64, + num_attention_heads=8, + num_layers=1, + params_dtype=torch.bfloat16, + fp8="e4m3", + fp8_recipe="tensorwise", + ) + defaults.update(overrides) + return TransformerConfig(**defaults) + + +class TestFloat8LinearGuards(PrimusUT): + """Tests that Float8 linear layers reject invalid configurations at init.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, monkeypatch): + dummy_args = SimpleNamespace( + rank=0, + world_size=1, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + offload=False, + offload_ops=[], + patch_primus_pipeline=False, + pp_algorithm=None, + patch_zero_bubble=False, + enable_zero_bubble=False, + rampup_batch_size=None, + global_batch_size=1, + micro_batch_size=1, + data_parallel_size=1, + decrease_batch_size_if_needed=False, + ) + import megatron.training.global_vars as gvars + + monkeypatch.setattr(gvars, "_GLOBAL_ARGS", dummy_args) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_column_parallel_rejects_tp_gt_1(self): + config = _make_fp8_transformer_config(tensor_model_parallel_size=2) + with pytest.raises(ValueError, match="tensor_model_parallel_size=1"): + Float8ColumnParallelLinear( + input_size=64, + output_size=128, + config=config, + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_column_parallel_rejects_gaf(self): + config = _make_fp8_transformer_config(gradient_accumulation_fusion=True) + with pytest.raises(ValueError, match="gradient_accumulation_fusion"): + Float8ColumnParallelLinear( + input_size=64, + output_size=128, + config=config, + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_column_parallel_rejects_sequence_parallel(self): + # SP requires TP>1 at TransformerConfig level, so we set TP=2 to + # pass config validation and hit our Float8 guard instead. + config = _make_fp8_transformer_config( + sequence_parallel=True, + tensor_model_parallel_size=2, + ) + with pytest.raises(ValueError, match="tensor_model_parallel_size=1"): + Float8ColumnParallelLinear( + input_size=64, + output_size=128, + config=config, + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_column_parallel_rejects_fp8_none(self): + config = _make_fp8_transformer_config(fp8=None) + with pytest.raises(ValueError, match="config.fp8"): + Float8ColumnParallelLinear( + input_size=64, + output_size=128, + config=config, + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_row_parallel_rejects_fp8_none(self): + config = _make_fp8_transformer_config(fp8=None) + with pytest.raises(ValueError, match="config.fp8"): + Float8RowParallelLinear( + input_size=64, + output_size=128, + config=config, + init_method=_init_method(), + bias=False, + input_is_parallel=True, + skip_bias_add=False, + is_expert=False, + ) + + +class TestFloat8LinearInit(PrimusUT): + """Tests successful construction and init-time dispatch.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, monkeypatch): + dummy_args = SimpleNamespace( + rank=0, + world_size=1, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + offload=False, + offload_ops=[], + patch_primus_pipeline=False, + pp_algorithm=None, + patch_zero_bubble=False, + enable_zero_bubble=False, + rampup_batch_size=None, + global_batch_size=1, + micro_batch_size=1, + data_parallel_size=1, + decrease_batch_size_if_needed=False, + ) + import megatron.training.global_vars as gvars + + monkeypatch.setattr(gvars, "_GLOBAL_ARGS", dummy_args) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_column_parallel_dispatch_tensorwise_uses_opaque(self): + config = _make_fp8_transformer_config(fp8="e4m3", fp8_recipe="tensorwise") + layer = Float8ColumnParallelLinear( + input_size=64, + output_size=128, + config=config, + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + assert layer._fp8_fn is OpaqueFP8LinearTensorwiseFunction + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_column_parallel_dispatch_blockwise(self): + config = _make_fp8_transformer_config(fp8="e4m3", fp8_recipe="blockwise") + layer = Float8ColumnParallelLinear( + input_size=64, + output_size=128, + config=config, + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) + assert layer._fp8_fn is FP8LinearBlockwiseFunction + + +class TestQuantizeFP8Tensorwise(PrimusUT): + """Verify _quantize_fp8_tensorwise correctness and properties.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def _reference_quantize(self, x, fp8_dtype, fp8_max): + """Pure-PyTorch reference (known correct): compute in FP32.""" + x_f32 = x.float() + amax = x_f32.abs().amax() + scale = fp8_max / amax.clamp(min=1e-12) + x_fp8 = (x_f32 * scale).clamp(-fp8_max, fp8_max).to(fp8_dtype) + scale_inv = 1.0 / scale + return x_fp8, scale_inv + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU") + def test_e4m3_correctness(self): + """BF16-scale quantize should match FP32 reference within a few percent + of elements (boundary rounding from BF16 scale precision).""" + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + fp8_dtype = float8_e4m3 + fp8_max = torch.finfo(fp8_dtype).max + + decomp_fp8, decomp_sinv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max) + ref_fp8, ref_sinv = self._reference_quantize(x, fp8_dtype, fp8_max) + + mismatch = (decomp_fp8.to(torch.float32) != ref_fp8.to(torch.float32)).float().mean() + assert mismatch < 0.05, f"More than 5% of elements differ: {mismatch:.4f}" + assert decomp_sinv.dtype == torch.float32 + torch.testing.assert_close(decomp_sinv, ref_sinv, atol=0, rtol=0) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU") + def test_e5m2_correctness(self): + """BF16-scale quantize should match FP32 reference within 1 FP8 ULP.""" + x = torch.randn(64, 512, dtype=torch.bfloat16, device="cuda") + fp8_dtype = float8_e5m2 + fp8_max = torch.finfo(fp8_dtype).max + + decomp_fp8, decomp_sinv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max) + ref_fp8, ref_sinv = self._reference_quantize(x, fp8_dtype, fp8_max) + + mismatch = (decomp_fp8.to(torch.float32) != ref_fp8.to(torch.float32)).float().mean() + assert mismatch < 0.02, f"More than 2% of elements differ: {mismatch:.4f}" + assert decomp_sinv.dtype == torch.float32 + torch.testing.assert_close(decomp_sinv, ref_sinv, atol=0, rtol=0) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU") + def test_dequant_roundtrip(self): + """Verify quantize -> dequantize preserves values within FP8 precision.""" + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + fp8_dtype = float8_e4m3 + fp8_max = torch.finfo(fp8_dtype).max + + x_fp8, scale_inv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max) + x_recon = x_fp8.float() * scale_inv + + rel_err = ((x_recon - x.float()).abs() / x.float().abs().clamp(min=1e-12)).mean() + assert rel_err < 0.1, f"Mean relative error {rel_err:.4f} too high for e4m3" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU") + def test_zero_tensor_no_nan(self): + x = torch.zeros(32, 64, dtype=torch.bfloat16, device="cuda") + fp8_dtype = float8_e4m3 + fp8_max = torch.finfo(fp8_dtype).max + + fp8_out, sinv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max) + assert not torch.isnan(sinv).any(), "scale_inv should not be NaN for zero input" + assert not torch.isinf(sinv).any(), "scale_inv should not be Inf for zero input" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU") + def test_near_zero_no_nan(self): + """Tensors with very small amax should not produce NaN or Inf in output.""" + x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") * 1e-6 + fp8_dtype = float8_e4m3 + fp8_max = torch.finfo(fp8_dtype).max + + x_fp8, scale_inv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max) + assert not torch.isnan(x_fp8.float()).any(), "FP8 output has NaN for near-zero input" + assert not torch.isinf(x_fp8.float()).any(), "FP8 output has Inf for near-zero input" + assert not torch.isnan(scale_inv).any(), "scale_inv has NaN for near-zero input" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU") + def test_scale_bf16_overflow_guard(self): + """When amax is tiny, FP32 scale would overflow BF16. Clamp must prevent it.""" + x = torch.full((4, 4), 1e-8, dtype=torch.bfloat16, device="cuda") + fp8_dtype = float8_e4m3 + fp8_max = torch.finfo(fp8_dtype).max + + x_fp8, scale_inv = _quantize_fp8_tensorwise(x, fp8_dtype, fp8_max) + assert not torch.isnan(x_fp8.float()).any(), "BF16 overflow guard failed — NaN in output" + assert torch.isfinite(scale_inv), "scale_inv should be finite" + + +class TestDecomposedTensorwiseCompile(PrimusUT): + """Verify DecomposedFP8LinearTensorwiseFunction works with torch.compile.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU") + def test_no_graph_break(self): + """torch._dynamo.explain should report zero graph breaks for the + decomposed quantize + gemm_fp8_impl forward.""" + from primus_turbo.pytorch.core.backend import BackendType + + fp8_fwd_dtype = float8_e4m3 + fp8_bwd_dtype = float8_e5m2 + fp8_fwd_max = torch.finfo(fp8_fwd_dtype).max + fp8_bwd_max = torch.finfo(fp8_bwd_dtype).max + gran_value = ScalingGranularity.TENSORWISE.value + backend_value = BackendType.HIPBLASLT.value + + x = torch.randn(4, 64, dtype=torch.bfloat16, device="cuda") + w = torch.randn(128, 64, dtype=torch.bfloat16, device="cuda") + + explanation = torch._dynamo.explain( + DecomposedFP8LinearTensorwiseFunction.apply, + )(x, w, fp8_fwd_dtype, fp8_bwd_dtype, fp8_fwd_max, fp8_bwd_max, gran_value, backend_value) + assert explanation.graph_break_count == 0, ( + f"Expected 0 graph breaks, got {explanation.graph_break_count}. " + f"Reasons: {explanation.break_reasons}" + ) + + +class TestOpaqueTensorwiseCompile(PrimusUT): + """Verify the refactored OpaqueFP8LinearTensorwiseFunction has zero graph breaks.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires GPU") + def test_no_graph_break(self): + """torch._dynamo.explain should report zero graph breaks for the + setup_context-based OpaqueFP8LinearTensorwiseFunction forward.""" + from primus_turbo.pytorch.core.backend import BackendType + from primus_turbo.pytorch.ops.quantization import quantize_fp8 + + torch._dynamo.reset() + + fp8_fwd_dtype = float8_e4m3 + fp8_bwd_dtype = float8_e5m2 + gran_value = ScalingGranularity.TENSORWISE.value + backend_value = BackendType.HIPBLASLT.value + + x = torch.randn(4, 64, dtype=torch.bfloat16, device="cuda") + w = torch.randn(128, 64, dtype=torch.bfloat16, device="cuda") + w_fp8, w_scale = quantize_fp8(w, fp8_fwd_dtype, ScalingGranularity.TENSORWISE) + + explanation = torch._dynamo.explain( + OpaqueFP8LinearTensorwiseFunction.apply, + )(x, w, w_fp8, w_scale, fp8_fwd_dtype, fp8_bwd_dtype, gran_value, backend_value) + assert explanation.graph_break_count == 0, ( + f"Expected 0 graph breaks, got {explanation.graph_break_count}. " + f"Reasons: {explanation.break_reasons}" + ) From aef188ab01c968c3c61f3c70393fa8d6b6748016 Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Thu, 9 Jul 2026 16:48:56 +0800 Subject: [PATCH 018/127] fix(diffusion): repoint module_utils import to primus.core (#867) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Fix `ModuleNotFoundError: No module named 'primus.modules'` that currently breaks test collection / `import primus` on `main`. An earlier refactor removed `primus/modules` and migrated the logging helpers to `primus.core.utils.module_utils`. Independently-merged diffusion and flux changes still imported `log_rank_0` from the old `primus.modules.module_utils` path, which now no longer exists — so importing the diffusion backend (e.g. `tests/unit_tests/backends/diffusion/test_wan_argument_builder.py`) fails at collection. ## Changes Repoint the three remaining offenders to the migrated location (`primus.modules.module_utils` → `primus.core.utils.module_utils`): - `primus/backends/diffusion/diffusion_adapter.py` - `primus/backends/diffusion/diffusion_pretrain_trainer.py` - `primus/backends/megatron/core/utils.py` The repo is now free of `primus.modules` references (grep-clean). ## Verification - `pre-commit run --all-files`: all hooks pass. - Import smoke: `import primus.backends.diffusion.argument_builder`, `primus.cli.main`, `train_runtime`, megatron/torchtitan adapters all OK. - `pytest tests/unit_tests/backends/diffusion/`: 8 passed (including the previously-failing `test_wan_argument_builder.py`). - Core unit subset (adapter/runtime/backend/config/base_trainer): 48 passed, no new failures. --- primus/backends/diffusion/diffusion_adapter.py | 2 +- primus/backends/diffusion/diffusion_pretrain_trainer.py | 2 +- primus/backends/megatron/core/utils.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/primus/backends/diffusion/diffusion_adapter.py b/primus/backends/diffusion/diffusion_adapter.py index 202c3914a..344c293b3 100644 --- a/primus/backends/diffusion/diffusion_adapter.py +++ b/primus/backends/diffusion/diffusion_adapter.py @@ -11,7 +11,7 @@ from primus.backends.diffusion.argument_builder import WanArgBuilder from primus.core.backend.backend_adapter import BackendAdapter -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 class DiffusionAdapter(BackendAdapter): diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index ddbb70d49..1b51b6a4c 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -10,8 +10,8 @@ from typing import Any from primus.core.trainer.base_trainer import BaseTrainer +from primus.core.utils.module_utils import log_rank_0 from primus.core.utils.yaml_utils import nested_namespace_to_dict -from primus.modules.module_utils import log_rank_0 class DiffusionPretrainTrainer(BaseTrainer): diff --git a/primus/backends/megatron/core/utils.py b/primus/backends/megatron/core/utils.py index 51d1897ec..be2ca21af 100644 --- a/primus/backends/megatron/core/utils.py +++ b/primus/backends/megatron/core/utils.py @@ -9,7 +9,7 @@ import torch from megatron.core import parallel_state -from primus.modules.module_utils import log_rank_0 +from primus.core.utils.module_utils import log_rank_0 @lru_cache From 608a4bdba9d80326a391729409efb42d34fbf893 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Thu, 9 Jul 2026 16:37:02 +0300 Subject: [PATCH 019/127] feat(flux): Flux DiT model, layers, attention, checkpoint converter (#811) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/model-common` — review after it. This is the central node several later PRs branch from. ## What this changes The Flux model itself: config, layer-spec, layers, attention, the model module, utils, and the checkpoint converter. The converter ships here because the package's `__init__` eagerly imports it (the package will not import without it). ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/model-common`. The data, training-primitives, compile, checkpoint-tools, and trainer PRs all descend from this one. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k flux`. Validated locally on an AMD GPU container: 29 passed. ## Files 15 (Flux config/layers/attention/model/utils, checkpoint converter + tests). --------- Co-authored-by: Flux Split Trial Co-authored-by: Luiza Sayfullina Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- .../core/models/diffusion/flux/__init__.py | 96 ++ .../core/models/diffusion/flux/attention.py | 597 ++++++++++++ .../diffusion/flux/checkpoint_converter.py | 881 ++++++++++++++++++ .../core/models/diffusion/flux/config.py | 315 +++++++ .../core/models/diffusion/flux/layer_spec.py | 594 ++++++++++++ .../core/models/diffusion/flux/layers.py | 130 +++ .../core/models/diffusion/flux/model.py | 857 +++++++++++++++++ .../core/models/diffusion/flux/utils.py | 197 ++++ .../diffusion/test_flux_chimera_init.py | 84 ++ .../megatron/diffusion/test_flux_config.py | 94 ++ .../diffusion/test_flux_fp8_context.py | 51 + .../diffusion/test_flux_init_weights.py | 172 ++++ .../megatron/diffusion/test_flux_layers.py | 54 ++ .../megatron/diffusion/test_flux_model.py | 84 ++ .../megatron/diffusion/test_flux_utils.py | 107 +++ 15 files changed, 4313 insertions(+) create mode 100644 primus/backends/megatron/core/models/diffusion/flux/__init__.py create mode 100644 primus/backends/megatron/core/models/diffusion/flux/attention.py create mode 100644 primus/backends/megatron/core/models/diffusion/flux/checkpoint_converter.py create mode 100644 primus/backends/megatron/core/models/diffusion/flux/config.py create mode 100644 primus/backends/megatron/core/models/diffusion/flux/layer_spec.py create mode 100644 primus/backends/megatron/core/models/diffusion/flux/layers.py create mode 100644 primus/backends/megatron/core/models/diffusion/flux/model.py create mode 100644 primus/backends/megatron/core/models/diffusion/flux/utils.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_chimera_init.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_config.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_init_weights.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_layers.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_model.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_utils.py diff --git a/primus/backends/megatron/core/models/diffusion/flux/__init__.py b/primus/backends/megatron/core/models/diffusion/flux/__init__.py new file mode 100644 index 000000000..5ea3e2236 --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/flux/__init__.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Flux diffusion model components. + +This module provides all components needed for the Flux architecture: +- Model: Flux (main model class) +- Configuration: FluxConfig +- Layers: EmbedND (3D RoPE position embedding) +- Attention: JointSelfAttention, FluxSingleAttention +- Layer specs: MMDiTLayer, FluxSingleTransformerBlock + +Quick Start - Model Creation: + >>> from primus.backends.megatron.core.models.diffusion.flux import Flux, FluxConfig + >>> + >>> # Create and configure model + >>> config = FluxConfig.flux_12b() + >>> model = Flux(config=config) + >>> + >>> # Load checkpoint (native Primus format) + >>> model.load_checkpoint("flux_12b.safetensors") + +Checkpoint Conversion: + >>> from primus.backends.megatron.core.models.diffusion.flux import convert_hf_checkpoint, FluxConfig + >>> + >>> # Convert HuggingFace checkpoint to Primus format + >>> config = FluxConfig.flux_12b() + >>> primus_sd = convert_hf_checkpoint( + ... "black-forest-labs/FLUX.1-dev/transformer", + ... flux_config=config, + ... save_to="primus_flux_12b.safetensors" + ... ) +""" + +from primus.backends.megatron.core.models.diffusion.flux.attention import ( + FluxSingleAttention, + JointSelfAttention, + JointSelfAttentionSubmodules, +) +from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import ( + convert_hf_checkpoint, +) +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.layer_spec import ( + FluxSingleTransformerBlock, + MMDiTLayer, + get_flux_double_transformer_spec_for_backend, + get_flux_layer_spec, + get_flux_single_transformer_spec_for_backend, +) +from primus.backends.megatron.core.models.diffusion.flux.layers import EmbedND, rope +from primus.backends.megatron.core.models.diffusion.flux.utils import ( + generate_image_position_ids, + generate_text_position_ids, + pack_latents, + unpack_latents, +) + + +# LAZY IMPORT: Don't import Flux here to avoid early TransformerBlock import +def __getattr__(name): + """Lazy import for Flux to avoid early TransformerBlock import.""" + if name == "Flux": + from primus.backends.megatron.core.models.diffusion.flux.model import Flux + + return Flux + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") + + +__all__ = [ + # Model + "Flux", + # Configuration + "FluxConfig", + # Layers + "EmbedND", + "rope", + # Attention + "JointSelfAttention", + "FluxSingleAttention", + "JointSelfAttentionSubmodules", + # Layer specs + "MMDiTLayer", + "FluxSingleTransformerBlock", + "get_flux_double_transformer_spec_for_backend", + "get_flux_single_transformer_spec_for_backend", + "get_flux_layer_spec", + # Utils + "pack_latents", + "unpack_latents", + "generate_image_position_ids", + "generate_text_position_ids", + # Checkpoint conversion + "convert_hf_checkpoint", +] diff --git a/primus/backends/megatron/core/models/diffusion/flux/attention.py b/primus/backends/megatron/core/models/diffusion/flux/attention.py new file mode 100644 index 000000000..b444b527c --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/flux/attention.py @@ -0,0 +1,597 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Flux attention mechanisms. + +This module implements specialized attention for Flux's MMDiT architecture: + - JointSelfAttention: Processes concatenated image + text tokens + - FluxSingleAttention: Processes image tokens only + - JointSelfAttentionSubmodules: Configuration for joint attention + +These implementations follow Megatron-Core's attention patterns with +customizations for diffusion model conditioning. + +Reference: + - MMDiT: "Scaling Rectified Flow Transformers" + - Megatron-Core: megatron.core.transformer.attention +""" + +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +from megatron.core.models.common.embeddings.rotary_pos_embedding import ( + apply_rotary_pos_emb, +) +from megatron.core.transformer.attention import ( + Attention, + SelfAttention, + SelfAttentionSubmodules, +) +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig +from torch import Tensor + +try: + from megatron.core.transformer.custom_layers.transformer_engine import SplitAlongDim +except ImportError: + SplitAlongDim = None + + +@dataclass +class JointSelfAttentionSubmodules: + """ + Submodules configuration for Joint Self-Attention layer. + + Joint attention processes both image and text tokens together (MMDiT architecture). + It requires separate QKV projections for image and text (context) streams. + + Attributes: + linear_qkv: QKV projection for main stream (image tokens) + added_linear_qkv: QKV projection for added stream (text/context tokens) + core_attention: Core attention computation module + linear_proj: Output projection for main stream + q_layernorm: Optional layer norm for queries (main stream) + k_layernorm: Optional layer norm for keys (main stream) + added_q_layernorm: Optional layer norm for queries (added stream) + added_k_layernorm: Optional layer norm for keys (added stream) + + Note: + Flux uses RMSNorm for Q/K normalization to improve training stability. + + Reference: + - Paper: "Scaling Rectified Flow Transformers" + """ + + linear_qkv: Union[ModuleSpec, type] = None + added_linear_qkv: Union[ModuleSpec, type] = None + core_attention: Union[ModuleSpec, type] = None + linear_proj: Union[ModuleSpec, type] = None + q_layernorm: Union[ModuleSpec, type] = None + k_layernorm: Union[ModuleSpec, type] = None + added_q_layernorm: Union[ModuleSpec, type] = None + added_k_layernorm: Union[ModuleSpec, type] = None + + +class JointSelfAttention(Attention): + """ + Joint Self-Attention for MMDiT (Multimodal Diffusion Transformer). + + Processes two token streams jointly -- main (image) and added (text) -- + by projecting each through separate QKV layers, concatenating, computing + joint attention, then splitting back. This enables cross-modal interaction + in Flux's "double blocks". + + Args: + config: Transformer configuration + submodules: JointSelfAttentionSubmodules with layer specifications + layer_number: Layer index in the model + attn_mask_type: Type of attention mask (default: padding) + context_pre_only: If True, only compute Q/K/V for context (default: False) + + Input/Output: + hidden_states [seq_main, B, H] + additional_hidden_states [seq_added, B, H] + -> (main_output, added_output) with same shapes + + Reference: + - "Scaling Rectified Flow Transformers" + """ + + def __init__( + self, + config: TransformerConfig, + submodules: JointSelfAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType = AttnMaskType.padding, + context_pre_only: bool = False, + **kwargs, + ): + # Use RMSNorm for Q/K normalization (improves stability) + config.normalization = "RMSNorm" + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + attention_type="self", + **kwargs, + ) + + # QKV projection for main stream (image tokens) + self.linear_qkv = build_module( + submodules.linear_qkv, + self.config.hidden_size, + self.query_projection_size + 2 * self.kv_projection_size, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=self.config.add_bias_linear or self.config.add_qkv_bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="qkv", + ) + + # QKV projection for added stream (text tokens) + if submodules.added_linear_qkv is not None: + self.added_linear_qkv = build_module( + submodules.added_linear_qkv, + self.config.hidden_size, + self.query_projection_size + 2 * self.kv_projection_size, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=self.config.add_qkv_bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="qkv", + ) + + # Output projection for added stream (text tokens) + if not context_pre_only: + self.added_linear_proj = build_module( + submodules.linear_proj, + self.query_projection_size, + self.config.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=self.config.add_bias_linear, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name="proj", + ) + + if ( + not context_pre_only + and getattr(self.config, "use_dual_fp8_output_projection", False) + and hasattr(self.linear_proj, "_fp8_config") + ): + from primus_turbo.pytorch.core.low_precision import ScalingGranularity + + if self.linear_proj._fp8_config.granularity == ScalingGranularity.TENSORWISE: + from primus_turbo.pytorch.core.backend import BackendType + + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + DualFP8LinearTensorwiseFunction, + _get_fp8_dtype, + ) + + self._dual_fp8_fn = DualFP8LinearTensorwiseFunction + cfg = self.linear_proj._fp8_config + self._dual_fp8_fwd_dtype = _get_fp8_dtype(cfg.format, is_fwd=True) + self._dual_fp8_bwd_dtype = _get_fp8_dtype(cfg.format, is_fwd=False) + self._dual_fp8_gran_value = ScalingGranularity.TENSORWISE.value + self._dual_fp8_backend_value = BackendType.HIPBLASLT.value + + # Optional Q/K layer normalization for main stream + if submodules.q_layernorm is not None: + self.q_layernorm = build_module( + submodules.q_layernorm, + hidden_size=self.hidden_size_per_attention_head, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + else: + self.q_layernorm = None + + if submodules.k_layernorm is not None: + self.k_layernorm = build_module( + submodules.k_layernorm, + hidden_size=self.hidden_size_per_attention_head, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + else: + self.k_layernorm = None + + # Optional Q/K layer normalization for added stream + if submodules.added_q_layernorm is not None: + self.added_q_layernorm = build_module( + submodules.added_q_layernorm, + hidden_size=self.hidden_size_per_attention_head, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + else: + self.added_q_layernorm = None + + if submodules.added_k_layernorm is not None: + self.added_k_layernorm = build_module( + submodules.added_k_layernorm, + hidden_size=self.hidden_size_per_attention_head, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + else: + self.added_k_layernorm = None + + def _split_qkv(self, mixed_qkv: Tensor) -> Tuple[Tensor, Tensor, Tensor]: + """ + Split mixed QKV tensor into separate Q, K, V tensors. + + Args: + mixed_qkv: Combined QKV tensor [seq, batch, hidden] + + Returns: + Tuple of (query, key, value) tensors + """ + # Reshape: [sq, b, hp] --> [sq, b, ng, (np/ng + 2) * hn] + new_tensor_shape = mixed_qkv.size()[:-1] + ( + self.num_query_groups_per_partition, + ( + (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2) + * self.hidden_size_per_attention_head + ), + ) + mixed_qkv = mixed_qkv.view(*new_tensor_shape) + + # Define split sizes for Q, K, V + split_arg_list = [ + ( + self.num_attention_heads_per_partition + // self.num_query_groups_per_partition + * self.hidden_size_per_attention_head + ), + self.hidden_size_per_attention_head, + self.hidden_size_per_attention_head, + ] + + # Split tensor + if SplitAlongDim is not None: + # Use Transformer Engine's optimized split if available + (query, key, value) = SplitAlongDim(mixed_qkv, 3, split_arg_list) + else: + # Fallback to PyTorch split + (query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3) + + # Reshape query: [sq, b, ng, np/ng * hn] -> [sq, b, np, hn] + query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head) + + return query, key, value + + def get_query_key_value_tensors( + self, hidden_states: Tensor, key_value_states: Optional[Tensor] = None + ) -> Tuple[Tensor, Tensor, Tensor]: + """ + Derive Q, K, V tensors from main stream hidden states. + + Args: + hidden_states: Main stream tokens [seq, batch, hidden] + key_value_states: Not used for self-attention + + Returns: + Tuple of (query, key, value) tensors + """ + # Project to QKV: [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn)] + mixed_qkv, _ = self.linear_qkv(hidden_states) + + # Split into Q, K, V + query, key, value = self._split_qkv(mixed_qkv) + + # Apply optional Q/K normalization + if self.q_layernorm is not None: + query = self.q_layernorm(query) + + if self.k_layernorm is not None: + key = self.k_layernorm(key) + + return query, key, value + + def get_added_query_key_value_tensors( + self, added_hidden_states: Tensor, key_value_states: Optional[Tensor] = None + ) -> Tuple[Tensor, Tensor, Tensor]: + """ + Derive Q, K, V tensors from added stream (text) hidden states. + + Args: + added_hidden_states: Added stream tokens [seq, batch, hidden] + key_value_states: Not used for self-attention + + Returns: + Tuple of (query, key, value) tensors + """ + # Project to QKV + mixed_qkv, _ = self.added_linear_qkv(added_hidden_states) + + # Split into Q, K, V + query, key, value = self._split_qkv(mixed_qkv) + + # Apply optional Q/K normalization + if self.added_q_layernorm is not None: + query = self.added_q_layernorm(query) + + if self.added_k_layernorm is not None: + key = self.added_k_layernorm(key) + + return query, key, value + + def forward( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor], + key_value_states: Optional[Tensor] = None, + inference_params=None, + rotary_pos_emb=None, + packed_seq_params=None, + additional_hidden_states: Optional[Tensor] = None, + ) -> Tuple[Tensor, Tensor]: + """ + Forward pass: Joint attention over image and text tokens. + + Args: + hidden_states: Main stream (image) tokens [seq_main, batch, hidden] + attention_mask: Attention mask + key_value_states: Not used for self-attention + inference_params: Parameters for inference (e.g., KV cache) + rotary_pos_emb: RoPE position embeddings + packed_seq_params: Parameters for packed sequences + additional_hidden_states: Added stream (text) tokens [seq_added, batch, hidden] + + Returns: + Tuple of (main_output, added_output): + - main_output: Processed main stream [seq_main, batch, hidden] + - added_output: Processed added stream [seq_added, batch, hidden] + """ + # Ensure rotary_pos_emb is a tuple for Q and K + if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple): + rotary_pos_emb = (rotary_pos_emb,) * 2 + + # Get Q, K, V for both streams + query, key, value = self.get_query_key_value_tensors(hidden_states) + added_query, added_key, added_value = self.get_added_query_key_value_tensors(additional_hidden_states) + + # Concatenate streams: [added; main] + query = torch.cat([added_query, query], dim=0) + key = torch.cat([added_key, key], dim=0) + value = torch.cat([added_value, value], dim=0) + + # Adjust for inference (KV caching, etc.) + query, key, value, rotary_pos_emb, attn_mask_type, *_ = self._adjust_key_value_for_inference( + inference_params, query, key, value, rotary_pos_emb + ) + + # Handle packed sequences + if packed_seq_params is not None: + query = query.squeeze(1) + key = key.squeeze(1) + value = value.squeeze(1) + + # Apply RoPE position embeddings + if rotary_pos_emb is not None: + q_pos_emb, k_pos_emb = rotary_pos_emb + + cu_seqlens_q = packed_seq_params.cu_seqlens_q if packed_seq_params is not None else None + cu_seqlens_kv = packed_seq_params.cu_seqlens_kv if packed_seq_params is not None else None + + query = apply_rotary_pos_emb(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q) + key = apply_rotary_pos_emb(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv) + + # Core attention computation + if self.checkpoint_core_attention and self.training: + core_attn_out = self._checkpointed_attention_forward( + query, + key, + value, + attention_mask, + attn_mask_type=attn_mask_type, + packed_seq_params=packed_seq_params, + ) + else: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + attn_mask_type=attn_mask_type, + packed_seq_params=packed_seq_params, + ) + + # Handle packed sequences output + if packed_seq_params is not None: + # Reshape: (t, np, hn) -> (t, b=1, h=np*hn) + core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1) + + # Split output back into added and main streams + encoder_attention_output = core_attn_out[: additional_hidden_states.shape[0], :, :] + attention_output = core_attn_out[additional_hidden_states.shape[0] :, :, :] + + # Project outputs + if hasattr(self, "_dual_fp8_fn"): + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _extract_fp8_weight, + ) + + w_fp8_a, w_scale_a = _extract_fp8_weight( + self.linear_proj.weight, + self._dual_fp8_fwd_dtype, + ) + w_fp8_b, w_scale_b = _extract_fp8_weight( + self.added_linear_proj.weight, + self._dual_fp8_fwd_dtype, + ) + result = self._dual_fp8_fn.apply( + attention_output, + self.linear_proj.weight, + w_fp8_a, + w_scale_a, + encoder_attention_output, + self.added_linear_proj.weight, + w_fp8_b, + w_scale_b, + self._dual_fp8_fwd_dtype, + self._dual_fp8_bwd_dtype, + self._dual_fp8_gran_value, + self._dual_fp8_backend_value, + ) + output, encoder_output = result[0], result[1] + if self.linear_proj.bias is not None: + output = output + self.linear_proj.bias + if self.added_linear_proj.bias is not None: + encoder_output = encoder_output + self.added_linear_proj.bias + else: + output, bias = self.linear_proj(attention_output) + encoder_output, encoder_bias = self.added_linear_proj(encoder_attention_output) + output = output + bias + encoder_output = encoder_output + encoder_bias + + return output, encoder_output + + +class FluxSingleAttention(SelfAttention): + """ + Single-stream Self-Attention for Flux (image tokens only). + + Standard self-attention without cross-modal interaction. Used in Flux's + "single blocks" after the joint MMDiT blocks. + + Args: + config: Transformer configuration + submodules: SelfAttentionSubmodules with layer specifications + layer_number: Layer index in the model + attn_mask_type: Type of attention mask (default: padding) + cp_comm_type: Communication type for context parallelism + """ + + def __init__( + self, + config: TransformerConfig, + submodules: SelfAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType = AttnMaskType.padding, + cp_comm_type: Optional[str] = None, + **kwargs, + ): + # Use RMSNorm for Q/K normalization + config.normalization = "RMSNorm" + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + cp_comm_type=cp_comm_type, + **kwargs, + ) + + # The original Flux proj_out (Diffusers) / linear2 (TorchTitan) is a single fused + # projection with one bias. Megatron splits it into linear_proj + linear_fc2, so the + # bias only needs to be on one path (linear_fc2) to preserve mathematical equivalence. + self.linear_proj = build_module( + submodules.linear_proj, + self.query_projection_size, + self.config.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=False, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name="proj", + ) + + def forward( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor], + key_value_states: Optional[Tensor] = None, + inference_params=None, + rotary_pos_emb=None, + packed_seq_params=None, + ) -> Tuple[Tensor, Optional[Tensor]]: + """ + Forward pass: Self-attention on image tokens. + + Args: + hidden_states: Image tokens [seq, batch, hidden] + attention_mask: Attention mask + key_value_states: Not used for self-attention + inference_params: Parameters for inference (e.g., KV cache) + rotary_pos_emb: RoPE position embeddings + packed_seq_params: Parameters for packed sequences + + Returns: + Tuple of (output, bias): + - output: Projected attention output [seq, batch, hidden] + - bias: Projection bias (None when linear_proj has bias=False) + """ + # Ensure rotary_pos_emb is a tuple for Q and K + if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple): + rotary_pos_emb = (rotary_pos_emb,) * 2 + + # Get Q, K, V + query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states) + + # Adjust for inference + query, key, value, rotary_pos_emb, attn_mask_type, *_ = self._adjust_key_value_for_inference( + inference_params, query, key, value, rotary_pos_emb + ) + + # Handle packed sequences + if packed_seq_params is not None: + query = query.squeeze(1) + key = key.squeeze(1) + value = value.squeeze(1) + + # Apply RoPE position embeddings + if rotary_pos_emb is not None: + q_pos_emb, k_pos_emb = rotary_pos_emb + + cu_seqlens_q = packed_seq_params.cu_seqlens_q if packed_seq_params is not None else None + cu_seqlens_kv = packed_seq_params.cu_seqlens_kv if packed_seq_params is not None else None + + query = apply_rotary_pos_emb(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q) + key = apply_rotary_pos_emb(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv) + + # Core attention computation + if self.checkpoint_core_attention and self.training: + core_attn_out = self._checkpointed_attention_forward( + query, + key, + value, + attention_mask, + attn_mask_type=attn_mask_type, + packed_seq_params=packed_seq_params, + ) + else: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + attn_mask_type=attn_mask_type, + packed_seq_params=packed_seq_params, + ) + + # Handle packed sequences output + if packed_seq_params is not None: + # Reshape: (t, np, hn) -> (t, b=1, h=np*hn) + core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1) + + # Project output (return both output and bias for skip_bias_add pattern) + output, bias = self.linear_proj(core_attn_out) + + return output, bias diff --git a/primus/backends/megatron/core/models/diffusion/flux/checkpoint_converter.py b/primus/backends/megatron/core/models/diffusion/flux/checkpoint_converter.py new file mode 100644 index 000000000..a2715cb7b --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/flux/checkpoint_converter.py @@ -0,0 +1,881 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. +""" +HuggingFace to Primus Flux checkpoint converter. + +Converts HuggingFace Diffusers Flux transformer checkpoints to +Primus/Megatron-Core compatible format. + +Key Conversion: + - HuggingFace: Separate double_blocks, single_blocks + - Primus: Unified transformer.layers.{0-N} with TransformerBlock + + This reflects Primus's architectural enhancement using heterogeneous + layer specifications in a single TransformerBlock container. + +Usage: + from primus.backends.megatron.core.models.diffusion.flux import convert_hf_checkpoint + + primus_state_dict = convert_hf_checkpoint( + checkpoint_path="black-forest-labs/FLUX.1-dev", + flux_config=config, + save_to="primus_flux_12b.safetensors" + ) + +Reference: + - HuggingFace Diffusers checkpoint format + - Megatron-Core TransformerBlock architecture +""" + +import logging +import os +from pathlib import Path +from typing import Dict, Optional, Union + +import torch +from safetensors.torch import load_file as load_safetensors +from safetensors.torch import save_file as save_safetensors + +logger = logging.getLogger(__name__) + + +def _fuse_qkv_weights( + config, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + v_weight: torch.Tensor, +) -> torch.Tensor: + """ + Fuse separate Q, K, V weight matrices into Megatron's fused QKV format. + + Megatron-Core uses grouped-query attention (GQA) format where Q, K, V + are interleaved per attention group: + [Q_group0, K_group0, V_group0, Q_group1, K_group1, V_group1, ...] + + Args: + config: FluxConfig with num_attention_heads, num_query_groups + q_weight: Query weights [hidden_size, hidden_size] + k_weight: Key weights [hidden_size, hidden_size] + v_weight: Value weights [hidden_size, hidden_size] + + Returns: + Fused QKV weights [(heads_per_group + 2) * num_query_groups * head_size, hidden_size] + in GQA interleaved format + + Reference: + - Megatron-Core: Grouped Query Attention patterns + """ + head_num = config.num_attention_heads + num_query_groups = getattr(config, "num_query_groups", head_num) + heads_per_group = head_num // num_query_groups + hidden_size = config.hidden_size + head_size = hidden_size // head_num + + # Get input shape from q_weight + old_tensor_shape = q_weight.size() + + # Reshape to [num_heads, head_size, ...] + new_q_tensor_shape = (head_num, head_size) + old_tensor_shape[1:] + new_kv_tensor_shape = (num_query_groups, head_size) + old_tensor_shape[1:] + + q = q_weight.view(*new_q_tensor_shape) + k = k_weight.view(*new_kv_tensor_shape) + v = v_weight.view(*new_kv_tensor_shape) + + # Interleave by group: [Q_heads_for_group, K_group, V_group, ...] + qkv_list = [] + for i in range(num_query_groups): + qkv_list.append(q[i * heads_per_group : (i + 1) * heads_per_group, :, :]) + qkv_list.append(k[i : i + 1, :, :]) + qkv_list.append(v[i : i + 1, :, :]) + + qkv = torch.cat(qkv_list, dim=0) + + # Validate shape + if qkv.ndim != 3: + raise ValueError(f"Expected 3D QKV tensor, got shape {qkv.shape}") + expected_dim0 = (heads_per_group + 2) * num_query_groups + if qkv.shape[0] != expected_dim0: + raise ValueError(f"Expected QKV dim0 {expected_dim0}, got shape {qkv.shape}") + if qkv.shape[1] != head_size: + raise ValueError(f"Expected QKV head_size {head_size}, got shape {qkv.shape}") + if qkv.shape[2] != old_tensor_shape[1]: + raise ValueError(f"Expected QKV hidden dim {old_tensor_shape[1]}, got shape {qkv.shape}") + + # Reshape to [fused_dim, hidden_size] + qkv = qkv.reshape(head_size * (head_num + 2 * num_query_groups), hidden_size) + + return qkv + + +def _fuse_qkv_bias( + config, + q_bias: torch.Tensor, + k_bias: torch.Tensor, + v_bias: torch.Tensor, +) -> torch.Tensor: + """ + Fuse Q, K, V bias terms into Megatron's fused format. + + Args: + config: FluxConfig with num_attention_heads, num_query_groups + q_bias: Query bias [hidden_size] + k_bias: Key bias [hidden_size] + v_bias: Value bias [hidden_size] + + Returns: + Fused QKV bias [(heads_per_group + 2) * num_query_groups * head_size] + in GQA interleaved format + + Reference: + - Megatron-Core: Grouped Query Attention patterns + """ + head_num = config.num_attention_heads + num_query_groups = getattr(config, "num_query_groups", head_num) + heads_per_group = head_num // num_query_groups + hidden_size = config.hidden_size + head_size = hidden_size // head_num + + # Reshape to [num_heads, head_size] + new_q_bias_shape = (head_num, head_size) + new_kv_bias_shape = (num_query_groups, head_size) + + q = q_bias.view(*new_q_bias_shape) + k = k_bias.view(*new_kv_bias_shape) + v = v_bias.view(*new_kv_bias_shape) + + # Interleave by group + qkv_bias_list = [] + for i in range(num_query_groups): + qkv_bias_list.append(q[i * heads_per_group : (i + 1) * heads_per_group, :]) + qkv_bias_list.append(k[i : i + 1, :]) + qkv_bias_list.append(v[i : i + 1, :]) + + qkv_bias = torch.cat(qkv_bias_list, dim=0) + qkv_bias = qkv_bias.reshape(head_size * (head_num + 2 * num_query_groups)) + + return qkv_bias + + +# Key mapping from HuggingFace to Primus +# After TransformerBlock refactor: double_blocks and single_blocks are now transformer.layers[0-56] +FLUX_KEY_MAPPING = { + "double_blocks": { + "norm1.linear.weight": "adaln.adaLN_modulation.1.weight", + "norm1.linear.bias": "adaln.adaLN_modulation.1.bias", + "norm1_context.linear.weight": "adaln_context.adaLN_modulation.1.weight", + "norm1_context.linear.bias": "adaln_context.adaLN_modulation.1.bias", + "attn.norm_q.weight": "self_attention.q_layernorm.weight", + "attn.norm_k.weight": "self_attention.k_layernorm.weight", + "attn.norm_added_q.weight": "self_attention.added_q_layernorm.weight", + "attn.norm_added_k.weight": "self_attention.added_k_layernorm.weight", + "attn.to_out.0.weight": "self_attention.linear_proj.weight", + "attn.to_out.0.bias": "self_attention.linear_proj.bias", + "attn.to_add_out.weight": "self_attention.added_linear_proj.weight", + "attn.to_add_out.bias": "self_attention.added_linear_proj.bias", + "ff.net.0.proj.weight": "mlp.linear_fc1.weight", + "ff.net.0.proj.bias": "mlp.linear_fc1.bias", + "ff.net.2.weight": "mlp.linear_fc2.weight", + "ff.net.2.bias": "mlp.linear_fc2.bias", + "ff_context.net.0.proj.weight": "context_mlp.linear_fc1.weight", + "ff_context.net.0.proj.bias": "context_mlp.linear_fc1.bias", + "ff_context.net.2.weight": "context_mlp.linear_fc2.weight", + "ff_context.net.2.bias": "context_mlp.linear_fc2.bias", + }, + "single_blocks": { + "norm.linear.weight": "adaln.adaLN_modulation.1.weight", + "norm.linear.bias": "adaln.adaLN_modulation.1.bias", + "proj_mlp.weight": "mlp.linear_fc1.weight", + "proj_mlp.bias": "mlp.linear_fc1.bias", + "attn.norm_q.weight": "self_attention.q_layernorm.weight", + "attn.norm_k.weight": "self_attention.k_layernorm.weight", + }, + # Root-level mappings + "norm_out.linear.bias": "norm_out.adaLN_modulation.1.bias", + "norm_out.linear.weight": "norm_out.adaLN_modulation.1.weight", + "proj_out.bias": "proj_out.bias", + "proj_out.weight": "proj_out.weight", + "time_text_embed.guidance_embedder.linear_1.bias": "guidance_embedding.in_layer.bias", + "time_text_embed.guidance_embedder.linear_1.weight": "guidance_embedding.in_layer.weight", + "time_text_embed.guidance_embedder.linear_2.bias": "guidance_embedding.out_layer.bias", + "time_text_embed.guidance_embedder.linear_2.weight": "guidance_embedding.out_layer.weight", + "x_embedder.bias": "img_embed.bias", + "x_embedder.weight": "img_embed.weight", + "time_text_embed.timestep_embedder.linear_1.bias": "timestep_embedding.time_embedding.in_layer.bias", + "time_text_embed.timestep_embedder.linear_1.weight": "timestep_embedding.time_embedding.in_layer.weight", + "time_text_embed.timestep_embedder.linear_2.bias": "timestep_embedding.time_embedding.out_layer.bias", + "time_text_embed.timestep_embedder.linear_2.weight": "timestep_embedding.time_embedding.out_layer.weight", + "context_embedder.bias": "txt_embed.bias", + "context_embedder.weight": "txt_embed.weight", + "time_text_embed.text_embedder.linear_1.bias": "vector_embedding.in_layer.bias", + "time_text_embed.text_embedder.linear_1.weight": "vector_embedding.in_layer.weight", + "time_text_embed.text_embedder.linear_2.bias": "vector_embedding.out_layer.bias", + "time_text_embed.text_embedder.linear_2.weight": "vector_embedding.out_layer.weight", +} + + +# Key mapping from Black Forest Labs native format to Primus format +# Reference: diffusers Flux checkpoint conversion +# This is the format used in official BFL releases (e.g., flux1-dev.safetensors, flux1-schnell.sft) +BFL_KEY_MAPPING = { + # Root-level embeddings + "img_in.weight": "img_embed.weight", + "img_in.bias": "img_embed.bias", + "txt_in.weight": "txt_embed.weight", + "txt_in.bias": "txt_embed.bias", + # Timestep embedding (time_in → timestep_embedding.time_embedding) + "time_in.in_layer.weight": "timestep_embedding.time_embedding.in_layer.weight", + "time_in.in_layer.bias": "timestep_embedding.time_embedding.in_layer.bias", + "time_in.out_layer.weight": "timestep_embedding.time_embedding.out_layer.weight", + "time_in.out_layer.bias": "timestep_embedding.time_embedding.out_layer.bias", + # Vector embedding (vector_in → vector_embedding) + "vector_in.in_layer.weight": "vector_embedding.in_layer.weight", + "vector_in.in_layer.bias": "vector_embedding.in_layer.bias", + "vector_in.out_layer.weight": "vector_embedding.out_layer.weight", + "vector_in.out_layer.bias": "vector_embedding.out_layer.bias", + # Guidance embedding (guidance_in → guidance_embedding) - optional + "guidance_in.in_layer.weight": "guidance_embedding.in_layer.weight", + "guidance_in.in_layer.bias": "guidance_embedding.in_layer.bias", + "guidance_in.out_layer.weight": "guidance_embedding.out_layer.weight", + "guidance_in.out_layer.bias": "guidance_embedding.out_layer.bias", + # Final layer + "final_layer.linear.weight": "proj_out.weight", + "final_layer.linear.bias": "proj_out.bias", + "final_layer.adaLN_modulation.1.weight": "norm_out.adaLN_modulation.1.weight", + "final_layer.adaLN_modulation.1.bias": "norm_out.adaLN_modulation.1.bias", +} + +# Block-level mappings for BFL double_blocks +BFL_DOUBLE_BLOCK_MAPPING = { + # Note: In BFL format, QKV are already FUSED (img_attn.qkv.weight contains Q+K+V concatenated) + # We need to UNFUSE them first, then REFUSE in Primus/Megatron GQA format + "img_attn.qkv.weight": "self_attention.linear_qkv.weight", # Will need special handling + "img_attn.qkv.bias": "self_attention.linear_qkv.bias", + "img_attn.proj.weight": "self_attention.linear_proj.weight", + "img_attn.proj.bias": "self_attention.linear_proj.bias", + "txt_attn.qkv.weight": "self_attention.added_linear_qkv.weight", # Will need special handling + "txt_attn.qkv.bias": "self_attention.added_linear_qkv.bias", + "txt_attn.proj.weight": "self_attention.added_linear_proj.weight", + "txt_attn.proj.bias": "self_attention.added_linear_proj.bias", + # QK LayerNorms + "img_attn.norm.query_norm.scale": "self_attention.q_layernorm.weight", + "img_attn.norm.key_norm.scale": "self_attention.k_layernorm.weight", + "txt_attn.norm.query_norm.scale": "self_attention.added_q_layernorm.weight", + "txt_attn.norm.key_norm.scale": "self_attention.added_k_layernorm.weight", + # Image MLPs + "img_mlp.0.weight": "mlp.linear_fc1.weight", + "img_mlp.0.bias": "mlp.linear_fc1.bias", + "img_mlp.2.weight": "mlp.linear_fc2.weight", + "img_mlp.2.bias": "mlp.linear_fc2.bias", + # Text MLPs + "txt_mlp.0.weight": "context_mlp.linear_fc1.weight", + "txt_mlp.0.bias": "context_mlp.linear_fc1.bias", + "txt_mlp.2.weight": "context_mlp.linear_fc2.weight", + "txt_mlp.2.bias": "context_mlp.linear_fc2.bias", + # Modulation (AdaLN) + "img_mod.lin.weight": "adaln.adaLN_modulation.1.weight", + "img_mod.lin.bias": "adaln.adaLN_modulation.1.bias", + "txt_mod.lin.weight": "adaln_context.adaLN_modulation.1.weight", + "txt_mod.lin.bias": "adaln_context.adaLN_modulation.1.bias", +} + +# Block-level mappings for BFL single_blocks +BFL_SINGLE_BLOCK_MAPPING = { + # Note: linear1 in BFL is FUSED [Q, K, V, MLP] - need to split + "linear1.weight": None, # Special handling: split into QKV + MLP + "linear1.bias": None, + # Note: linear2 in BFL is FUSED [MLP_out, proj_out] - need to split + "linear2.weight": None, # Special handling: split into proj_out only (simplified in some versions) + "linear2.bias": None, + # Modulation + "modulation.lin.weight": "adaln.adaLN_modulation.1.weight", + "modulation.lin.bias": "adaln.adaLN_modulation.1.bias", + # QK LayerNorms (renamed 'norm' in BFL single blocks) + "norm.query_norm.scale": "self_attention.q_layernorm.weight", + "norm.key_norm.scale": "self_attention.k_layernorm.weight", +} + + +def detect_checkpoint_format(state_dict: Dict[str, torch.Tensor]) -> str: + """ + Detect checkpoint format by inspecting keys. + + Args: + state_dict: Loaded checkpoint state dictionary + + Returns: + 'bfl_native': Black Forest Labs native format (img_in, txt_in, time_in) + 'hf_diffusers': HuggingFace Diffusers format (x_embedder, context_embedder, time_text_embed) + + Reference: diffusers Flux checkpoint conversion + The BFL native format is what Black Forest Labs releases directly. + The HF Diffusers format is what you get after saving from diffusers library. + """ + sample_keys = list(state_dict.keys()) + + # Check for BFL native format markers + # BFL uses: img_in, txt_in, time_in, vector_in, guidance_in + if any(k.startswith("img_in.") for k in sample_keys): + return "bfl_native" + if any(k.startswith("txt_in.") for k in sample_keys): + return "bfl_native" + if any("time_in.in_layer" in k for k in sample_keys): + return "bfl_native" + + # Check for HF Diffusers format markers + # Diffusers uses: x_embedder, context_embedder, time_text_embed + if any(k.startswith("x_embedder.") for k in sample_keys): + return "hf_diffusers" + if any("time_text_embed.timestep_embedder" in k for k in sample_keys): + return "hf_diffusers" + + # Default to HF Diffusers (original behavior) + return "hf_diffusers" + + +def _get_hf_token(token_file: Optional[str] = None) -> Optional[str]: + """ + Get HuggingFace token with fallback options. + + Thin wrapper around the shared + :func:`...preprocessing.auth.setup_hf_authentication` so the token-resolution + priority chain (file with permission check -> HF_TOKEN env -> HF CLI login -> + None) lives in one place. Returns None instead of raising so the converter + can fall back to public-only access. + + Args: + token_file: Optional path to token file + + Returns: + Token string if found, None otherwise + """ + from primus.backends.megatron.data.diffusion.preprocessing.auth import ( + HFAuthError, + setup_hf_authentication, + ) + + try: + return setup_hf_authentication(token_file=token_file, use_env=True) + except HFAuthError: + # Preserve fallback-to-public behavior for the converter rather than + # hard-failing on a bad/insecure token file. + return None + + +def convert_hf_checkpoint( + checkpoint_path: Union[str, Path], + flux_config, + save_to: Optional[Union[str, Path]] = None, +) -> Dict[str, torch.Tensor]: + """ + Convert HuggingFace Flux checkpoint to Primus format. + + Supports both local paths and HuggingFace repo IDs. If a repo ID is provided + (e.g., "black-forest-labs/FLUX.1-dev"), the checkpoint will be + automatically downloaded from HuggingFace Hub. + + Args: + checkpoint_path: Path to HF checkpoint OR HuggingFace repo ID + (e.g., "black-forest-labs/FLUX.1-dev/transformer") + flux_config: FluxConfig instance with model architecture info + save_to: Optional path to save converted checkpoint + + Returns: + Dictionary of converted state dict (Primus format) + + Example: + >>> from primus.backends.megatron.core.models.diffusion.flux import FluxConfig + >>> config = FluxConfig.flux_12b() + >>> # Auto-download from HuggingFace + >>> primus_sd = convert_hf_checkpoint( + ... "black-forest-labs/FLUX.1-dev/transformer", + ... flux_config=config, + ... save_to="primus_flux_12b.safetensors" + ... ) + """ + checkpoint_path_str = str(checkpoint_path) + checkpoint_path_obj = Path(checkpoint_path) + + # Check if path exists locally first + if checkpoint_path_obj.exists(): + # Local file/directory exists, use it directly + pass + else: + # Path doesn't exist - check if it looks like a HuggingFace repo ID + # HF repo IDs: don't start with /, ./, ../, and don't contain .. + looks_like_hf_repo = ( + "/" in checkpoint_path_str + and not checkpoint_path_str.startswith("/") + and not checkpoint_path_str.startswith("./") + and not checkpoint_path_str.startswith("../") + and ".." not in checkpoint_path_str + and not checkpoint_path_str.startswith(".") # Avoid hidden files/dirs + ) + + if looks_like_hf_repo: + logger.info("Detected HuggingFace repo ID: %s", checkpoint_path_str) + logger.info("Downloading from HuggingFace Hub...") + + try: + from huggingface_hub import snapshot_download + except ImportError: + raise ImportError( + "huggingface_hub is required for downloading checkpoints. " + "Install with: pip install huggingface_hub" + ) + + # Setup authentication: .hf_token file → HF_TOKEN env → HF CLI login + token_file = Path(__file__).parents[6] / ".hf_token" + hf_token = None + if token_file.exists(): + logger.info("Using HuggingFace token from project root: %s", token_file) + hf_token = _get_hf_token(token_file=str(token_file)) + else: + # Falls back to HF_TOKEN env var and HF CLI login + hf_token = _get_hf_token(token_file=None) + + # Parse repo_id and subfolder + parts = checkpoint_path_str.split("/") + if len(parts) >= 2: + repo_id = "/".join(parts[:2]) # e.g., "black-forest-labs/FLUX.1-dev" + subfolder = "/".join(parts[2:]) if len(parts) > 2 else None # e.g., "transformer" + else: + repo_id = checkpoint_path_str + subfolder = None + + # Download to cache + cache_dir = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface/hub")) + + try: + local_dir = snapshot_download( + repo_id=repo_id, + allow_patterns=[f"{subfolder}/*"] if subfolder else None, + cache_dir=cache_dir, + token=hf_token, + resume_download=True, + ) + + # Construct path to actual checkpoint + checkpoint_path = Path(local_dir) / subfolder if subfolder else Path(local_dir) + logger.info("Downloaded to: %s", checkpoint_path) + + except Exception as e: + logger.error("Download failed: %s", e) + if "401" in str(e) or "403" in str(e): + raise RuntimeError( + f"Authentication failed. Please set HuggingFace token using one of:\n" + f" 1. Create .hf_token file: echo 'your_token' > .hf_token && chmod 600 .hf_token\n" + f" 2. Set environment variable: export HF_TOKEN=your_token\n" + f" 3. Login via CLI: huggingface-cli login\n" + f"Get your token from: https://huggingface.co/settings/tokens\n" + f"Accept the model license at: https://huggingface.co/{repo_id}" + ) from e + raise + else: + # Looks like a local path that doesn't exist - raise FileNotFoundError + raise FileNotFoundError( + f"Checkpoint file or directory not found: {checkpoint_path_str}\n" + f"If this is a HuggingFace repo ID, ensure it follows the format 'org/repo' or 'org/repo/subfolder'" + ) + + logger.info("Loading HuggingFace checkpoint from: %s", checkpoint_path) + + # Load HF checkpoint + hf_state_dict = {} + checkpoint_path = Path(checkpoint_path) + + if checkpoint_path.is_dir(): + # Load all .safetensors files in directory + safetensor_files = list(checkpoint_path.glob("*.safetensors")) + if not safetensor_files: + raise FileNotFoundError(f"No .safetensors files found in {checkpoint_path}") + + for file in safetensor_files: + logger.info(" Loading %s...", file.name) + hf_state_dict.update(load_safetensors(str(file))) + elif checkpoint_path.is_file(): + hf_state_dict = load_safetensors(str(checkpoint_path)) + else: + raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") + + logger.info("Loaded %d keys from HuggingFace checkpoint", len(hf_state_dict)) + + # Detect checkpoint format + checkpoint_format = detect_checkpoint_format(hf_state_dict) + logger.info("Detected checkpoint format: %s", checkpoint_format) + + # Branch based on format + if checkpoint_format == "bfl_native": + return _convert_bfl_checkpoint(hf_state_dict, flux_config, save_to) + else: + return _convert_hf_diffusers_checkpoint(hf_state_dict, flux_config, save_to) + + +def _convert_hf_diffusers_checkpoint( + hf_state_dict: Dict[str, torch.Tensor], + flux_config, + save_to: Optional[Union[str, Path]] = None, +) -> Dict[str, torch.Tensor]: + """Convert HuggingFace Diffusers format checkpoint to Primus format.""" + # Convert to Primus format + primus_state_dict = {} + num_double_blocks = -1 + num_single_blocks = -1 + + # First pass: Convert simple key mappings (ONLY double blocks and root keys) + logger.info("Converting simple key mappings (double blocks and root keys)...") + for hf_key, value in hf_state_dict.items(): + # Skip QKV weights - will handle separately + if any( + x in hf_key + for x in [ + "attn.to_q", + "attn.to_k", + "attn.to_v", + "attn.add_q_proj", + "attn.add_k_proj", + "attn.add_v_proj", + ] + ): + continue + + # Skip ALL single block keys - will handle in separate pass AFTER determining num_double_blocks + if hf_key.startswith("single_transformer_blocks"): + continue + + # Map double blocks -> transformer.layers[0-18] + if hf_key.startswith("transformer_blocks"): + parts = hf_key.split(".") + idx = int(parts[1]) + sub_key = ".".join(parts[2:]) + num_double_blocks = max(idx, num_double_blocks) + + if sub_key in FLUX_KEY_MAPPING["double_blocks"]: + # New key format: transformer.layers.{idx} instead of double_blocks.{idx} + primus_key = f"transformer.layers.{idx}.{FLUX_KEY_MAPPING['double_blocks'][sub_key]}" + primus_state_dict[primus_key] = value + + # Map root-level keys + elif hf_key in FLUX_KEY_MAPPING: + # Special handling for norm_out: swap scale/shift halves + # HF Diffusers stores [SCALE; SHIFT], but Primus/BFL native expects [SHIFT; SCALE] + if hf_key == "norm_out.linear.weight": + half_size = value.shape[0] // 2 + scale_half = value[:half_size, :] # HF first half = SCALE + shift_half = value[half_size:, :] # HF second half = SHIFT + # Swap to BFL native order: [SHIFT; SCALE] + value = torch.cat([shift_half, scale_half], dim=0) + elif hf_key == "norm_out.linear.bias": + half_size = value.shape[0] // 2 + scale_half = value[:half_size] # HF first half = SCALE + shift_half = value[half_size:] # HF second half = SHIFT + # Swap to BFL native order: [SHIFT; SCALE] + value = torch.cat([shift_half, scale_half], dim=0) + + primus_state_dict[FLUX_KEY_MAPPING[hf_key]] = value + + # Detect number of single blocks + for hf_key in hf_state_dict.keys(): + if hf_key.startswith("single_transformer_blocks"): + parts = hf_key.split(".") + idx = int(parts[1]) + num_single_blocks = max(idx, num_single_blocks) + + logger.info("Found %d double blocks, %d single blocks", num_double_blocks + 1, num_single_blocks + 1) + + # Second pass: Convert single block simple keys (NOW num_double_blocks is known!) + logger.info("Converting single block simple keys...") + for hf_key, value in hf_state_dict.items(): + if not hf_key.startswith("single_transformer_blocks"): + continue + + # Skip QKV and proj_out - will handle separately in later passes + if any(x in hf_key for x in ["attn.to_q", "attn.to_k", "attn.to_v", "proj_out"]): + continue + + parts = hf_key.split(".") + idx = int(parts[1]) + sub_key = ".".join(parts[2:]) + + if sub_key in FLUX_KEY_MAPPING["single_blocks"]: + layer_idx = num_double_blocks + 1 + idx # NOW num_double_blocks is determined! + primus_key = f"transformer.layers.{layer_idx}.{FLUX_KEY_MAPPING['single_blocks'][sub_key]}" + primus_state_dict[primus_key] = value + + # Third pass: Fuse QKV weights for double blocks (now transformer.layers[0-18]) + logger.info("Fusing QKV weights for double blocks...") + for i in range(num_double_blocks + 1): + # Main attention QKV + q_key = f"transformer_blocks.{i}.attn.to_q.weight" + k_key = f"transformer_blocks.{i}.attn.to_k.weight" + v_key = f"transformer_blocks.{i}.attn.to_v.weight" + + fused_qkv = _fuse_qkv_weights( + flux_config, hf_state_dict[q_key], hf_state_dict[k_key], hf_state_dict[v_key] + ) + # New key format: transformer.layers.{i} instead of double_blocks.{i} + primus_state_dict[f"transformer.layers.{i}.self_attention.linear_qkv.weight"] = fused_qkv + + # QKV bias + q_bias_key = f"transformer_blocks.{i}.attn.to_q.bias" + k_bias_key = f"transformer_blocks.{i}.attn.to_k.bias" + v_bias_key = f"transformer_blocks.{i}.attn.to_v.bias" + + fused_qkv_bias = _fuse_qkv_bias( + flux_config, hf_state_dict[q_bias_key], hf_state_dict[k_bias_key], hf_state_dict[v_bias_key] + ) + primus_state_dict[f"transformer.layers.{i}.self_attention.linear_qkv.bias"] = fused_qkv_bias + + # Context (added) attention QKV + add_q_key = f"transformer_blocks.{i}.attn.add_q_proj.weight" + add_k_key = f"transformer_blocks.{i}.attn.add_k_proj.weight" + add_v_key = f"transformer_blocks.{i}.attn.add_v_proj.weight" + + fused_add_qkv = _fuse_qkv_weights( + flux_config, hf_state_dict[add_q_key], hf_state_dict[add_k_key], hf_state_dict[add_v_key] + ) + primus_state_dict[f"transformer.layers.{i}.self_attention.added_linear_qkv.weight"] = fused_add_qkv + + # Added QKV bias + add_q_bias_key = f"transformer_blocks.{i}.attn.add_q_proj.bias" + add_k_bias_key = f"transformer_blocks.{i}.attn.add_k_proj.bias" + add_v_bias_key = f"transformer_blocks.{i}.attn.add_v_proj.bias" + + fused_add_qkv_bias = _fuse_qkv_bias( + flux_config, + hf_state_dict[add_q_bias_key], + hf_state_dict[add_k_bias_key], + hf_state_dict[add_v_bias_key], + ) + primus_state_dict[f"transformer.layers.{i}.self_attention.added_linear_qkv.bias"] = fused_add_qkv_bias + + # Fourth pass: Fuse QKV and split proj_out for single blocks (now transformer.layers[19+]) + logger.info("Fusing QKV and splitting proj_out for single blocks...") + for i in range(num_single_blocks + 1): + # Calculate layer index with offset + layer_idx = num_double_blocks + 1 + i + + # QKV + q_key = f"single_transformer_blocks.{i}.attn.to_q.weight" + k_key = f"single_transformer_blocks.{i}.attn.to_k.weight" + v_key = f"single_transformer_blocks.{i}.attn.to_v.weight" + + fused_qkv = _fuse_qkv_weights( + flux_config, hf_state_dict[q_key], hf_state_dict[k_key], hf_state_dict[v_key] + ) + # New key format with offset + primus_state_dict[f"transformer.layers.{layer_idx}.self_attention.linear_qkv.weight"] = fused_qkv + + # QKV bias + q_bias_key = f"single_transformer_blocks.{i}.attn.to_q.bias" + k_bias_key = f"single_transformer_blocks.{i}.attn.to_k.bias" + v_bias_key = f"single_transformer_blocks.{i}.attn.to_v.bias" + + fused_qkv_bias = _fuse_qkv_bias( + flux_config, hf_state_dict[q_bias_key], hf_state_dict[k_bias_key], hf_state_dict[v_bias_key] + ) + primus_state_dict[f"transformer.layers.{layer_idx}.self_attention.linear_qkv.bias"] = fused_qkv_bias + + # Split proj_out (combined MLP+Attention output in HF) + # HF format: [out_dim, hidden*2] where [:, :3072] is attention, [:, 3072:] is MLP + proj_out_weight = hf_state_dict[f"single_transformer_blocks.{i}.proj_out.weight"] + proj_out_bias = hf_state_dict[f"single_transformer_blocks.{i}.proj_out.bias"] + + # Split weight at hidden_size (3072 for Flux) + hidden_size = flux_config.hidden_size + attn_proj = proj_out_weight[:, :hidden_size].clone() + mlp_proj = proj_out_weight[:, hidden_size:].clone() + + primus_state_dict[f"transformer.layers.{layer_idx}.self_attention.linear_proj.weight"] = attn_proj + primus_state_dict[f"transformer.layers.{layer_idx}.mlp.linear_fc2.weight"] = mlp_proj + + primus_state_dict[f"transformer.layers.{layer_idx}.mlp.linear_fc2.bias"] = proj_out_bias.clone() + + logger.info("Conversion complete! Primus state dict has %d keys", len(primus_state_dict)) + + # Save if requested + if save_to: + save_to = Path(save_to) + save_to.parent.mkdir(parents=True, exist_ok=True) + logger.info("Saving to: %s", save_to) + save_safetensors(primus_state_dict, str(save_to)) + logger.info("Saved successfully!") + + return primus_state_dict + + +def _convert_bfl_checkpoint( + bfl_state_dict: Dict[str, torch.Tensor], + flux_config, + save_to: Optional[Union[str, Path]] = None, +) -> Dict[str, torch.Tensor]: + """ + Convert Black Forest Labs native format checkpoint to Primus format. + + Reference: diffusers Flux checkpoint conversion + + Key differences from HF Diffusers format: + - QKV are already FUSED in BFL (img_attn.qkv.weight contains Q+K+V concatenated) + - Single blocks have fused linear1 (Q, K, V, MLP) and linear2 (proj_out) + """ + logger.info("Converting BFL native format to Primus...") + primus_state_dict = {} + num_double_blocks = -1 + num_single_blocks = -1 + + # First pass: Convert simple root-level mappings + logger.info("Converting root-level keys...") + for bfl_key, value in bfl_state_dict.items(): + # Skip block-level keys - will handle in second pass + if bfl_key.startswith("double_blocks.") or bfl_key.startswith("single_blocks."): + continue + + # Map root-level keys + if bfl_key in BFL_KEY_MAPPING: + primus_key = BFL_KEY_MAPPING[bfl_key] + primus_state_dict[primus_key] = value + + # Detect number of blocks + for key in bfl_state_dict.keys(): + if key.startswith("double_blocks."): + idx = int(key.split(".")[1]) + num_double_blocks = max(idx, num_double_blocks) + elif key.startswith("single_blocks."): + idx = int(key.split(".")[1]) + num_single_blocks = max(idx, num_single_blocks) + + logger.info("Found %d double blocks, %d single blocks", num_double_blocks + 1, num_single_blocks + 1) + + # Second pass: Convert double blocks -> transformer.layers[0-18] + logger.info("Converting double blocks (unfusing and refusing QKV)...") + for i in range(num_double_blocks + 1): + block_prefix_bfl = f"double_blocks.{i}" + block_prefix_primus = f"transformer.layers.{i}" # New key format + + # Convert simple mappings first + for bfl_suffix, primus_suffix in BFL_DOUBLE_BLOCK_MAPPING.items(): + if primus_suffix is None: + continue # Skip special handling keys + + bfl_key = f"{block_prefix_bfl}.{bfl_suffix}" + if bfl_key in bfl_state_dict: + primus_key = f"{block_prefix_primus}.{primus_suffix}" + + # Special handling for QKV (need to unfuse BFL QKV, then refuse in Primus format) + if "qkv" in bfl_suffix and ("weight" in bfl_suffix or "bias" in bfl_suffix): + continue # Handle separately below + else: + primus_state_dict[primus_key] = bfl_state_dict[bfl_key] + + # Handle img_attn QKV (unfuse from BFL concat format, then refuse in Primus GQA format) + # BFL format: [Q; K; V] concatenated along dim 0 + # Reference: diffusers Flux checkpoint conversion + img_qkv_weight = bfl_state_dict[f"{block_prefix_bfl}.img_attn.qkv.weight"] + img_qkv_bias = bfl_state_dict[f"{block_prefix_bfl}.img_attn.qkv.bias"] + + # Unfuse: split into Q, K, V + img_q, img_k, img_v = torch.chunk(img_qkv_weight, 3, dim=0) + img_q_bias, img_k_bias, img_v_bias = torch.chunk(img_qkv_bias, 3, dim=0) + + # Refuse in Primus GQA format + fused_img_qkv = _fuse_qkv_weights(flux_config, img_q, img_k, img_v) + fused_img_qkv_bias = _fuse_qkv_bias(flux_config, img_q_bias, img_k_bias, img_v_bias) + + primus_state_dict[f"{block_prefix_primus}.self_attention.linear_qkv.weight"] = fused_img_qkv + primus_state_dict[f"{block_prefix_primus}.self_attention.linear_qkv.bias"] = fused_img_qkv_bias + + # Handle txt_attn QKV (same process) + txt_qkv_weight = bfl_state_dict[f"{block_prefix_bfl}.txt_attn.qkv.weight"] + txt_qkv_bias = bfl_state_dict[f"{block_prefix_bfl}.txt_attn.qkv.bias"] + + txt_q, txt_k, txt_v = torch.chunk(txt_qkv_weight, 3, dim=0) + txt_q_bias, txt_k_bias, txt_v_bias = torch.chunk(txt_qkv_bias, 3, dim=0) + + fused_txt_qkv = _fuse_qkv_weights(flux_config, txt_q, txt_k, txt_v) + fused_txt_qkv_bias = _fuse_qkv_bias(flux_config, txt_q_bias, txt_k_bias, txt_v_bias) + + primus_state_dict[f"{block_prefix_primus}.self_attention.added_linear_qkv.weight"] = fused_txt_qkv + primus_state_dict[f"{block_prefix_primus}.self_attention.added_linear_qkv.bias"] = fused_txt_qkv_bias + + # Third pass: Convert single blocks -> transformer.layers[19+] + # BFL single blocks have: + # - linear1: fused [Q, K, V, MLP] - need to split and handle separately + # - linear2: just proj_out (simpler than HF Diffusers which has MLP+proj) + # Reference: diffusers Flux checkpoint conversion + logger.info("Converting single blocks (splitting linear1)...") + for i in range(num_single_blocks + 1): + block_prefix_bfl = f"single_blocks.{i}" + # Calculate layer index with offset (num_double_blocks + 1 + i) + layer_idx = num_double_blocks + 1 + i + block_prefix_primus = f"transformer.layers.{layer_idx}" # New key format + + # Convert simple mappings (modulation, norms) + for bfl_suffix, primus_suffix in BFL_SINGLE_BLOCK_MAPPING.items(): + if primus_suffix is None: + continue # Skip special handling + + bfl_key = f"{block_prefix_bfl}.{bfl_suffix}" + if bfl_key in bfl_state_dict: + primus_key = f"{block_prefix_primus}.{primus_suffix}" + primus_state_dict[primus_key] = bfl_state_dict[bfl_key] + + # Handle linear1: fused [Q, K, V, MLP] + linear1_weight = bfl_state_dict[f"{block_prefix_bfl}.linear1.weight"] + linear1_bias = bfl_state_dict[f"{block_prefix_bfl}.linear1.bias"] + + # Split along dim 0: [Q, K, V, MLP] + hidden_size = flux_config.hidden_size + mlp_hidden_dim = int(hidden_size * 4.0) # mlp_ratio = 4.0 + split_sizes = (hidden_size, hidden_size, hidden_size, mlp_hidden_dim) + + q, k, v, mlp = torch.split(linear1_weight, split_sizes, dim=0) + q_bias, k_bias, v_bias, mlp_bias = torch.split(linear1_bias, split_sizes, dim=0) + + # Fuse Q, K, V in Primus GQA format + fused_qkv = _fuse_qkv_weights(flux_config, q, k, v) + fused_qkv_bias = _fuse_qkv_bias(flux_config, q_bias, k_bias, v_bias) + + primus_state_dict[f"{block_prefix_primus}.self_attention.linear_qkv.weight"] = fused_qkv + primus_state_dict[f"{block_prefix_primus}.self_attention.linear_qkv.bias"] = fused_qkv_bias + + # MLP goes to linear_fc1 + primus_state_dict[f"{block_prefix_primus}.mlp.linear_fc1.weight"] = mlp + primus_state_dict[f"{block_prefix_primus}.mlp.linear_fc1.bias"] = mlp_bias + + # Handle linear2: In BFL, this is just proj_out (not fused with MLP like in HF Diffusers) + linear2_weight = bfl_state_dict[f"{block_prefix_bfl}.linear2.weight"] + linear2_bias = bfl_state_dict[f"{block_prefix_bfl}.linear2.bias"] + + # In BFL format, linear2 is [out_channels, hidden_size*2] where first half is attn proj, second is mlp proj + # Split at hidden_size + attn_proj = linear2_weight[:, :hidden_size].clone() + mlp_proj = linear2_weight[:, hidden_size:].clone() + + primus_state_dict[f"{block_prefix_primus}.self_attention.linear_proj.weight"] = attn_proj + primus_state_dict[f"{block_prefix_primus}.mlp.linear_fc2.weight"] = mlp_proj + + primus_state_dict[f"{block_prefix_primus}.mlp.linear_fc2.bias"] = linear2_bias.clone() + + logger.info("Conversion complete! Primus state dict has %d keys", len(primus_state_dict)) + + # Save if requested + if save_to: + save_to = Path(save_to) + save_to.parent.mkdir(parents=True, exist_ok=True) + logger.info("Saving to: %s", save_to) + save_safetensors(primus_state_dict, str(save_to)) + logger.info("Saved successfully!") + + return primus_state_dict + + +__all__ = [ + "convert_hf_checkpoint", + "FLUX_KEY_MAPPING", + "BFL_KEY_MAPPING", + "BFL_DOUBLE_BLOCK_MAPPING", + "BFL_SINGLE_BLOCK_MAPPING", + "detect_checkpoint_format", + "_fuse_qkv_weights", + "_fuse_qkv_bias", +] diff --git a/primus/backends/megatron/core/models/diffusion/flux/config.py b/primus/backends/megatron/core/models/diffusion/flux/config.py new file mode 100644 index 000000000..4d920de51 --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/flux/config.py @@ -0,0 +1,315 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Configuration for Flux diffusion model. + +Flux is a flow-based diffusion model with MMDiT (Multimodal Diffusion Transformer) +architecture that uses separate "joint" and "single" transformer blocks. + +Reference: + - https://github.com/black-forest-labs/flux + - NeMo's Flux implementation +""" + +from dataclasses import dataclass +from typing import Callable, Optional + +import torch +import torch.nn as nn + +from ..common.config import BaseDiffusionConfig + + +# Custom non-JIT compiled openai_gelu to avoid ROCm bugs +def openai_gelu_no_jit(x): + """ + OpenAI's GELU implementation without JIT compilation (for ROCm compatibility). + + This is the tanh-based approximation of GELU used in the original Flux model. + We use a non-JIT version to avoid ROCm compilation bugs that cause NaN values. + """ + return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x))) + + +@dataclass +class FluxConfig(BaseDiffusionConfig): + """ + Flux-specific configuration. + + Configuration for Flux flow-based diffusion model. Flux uses a dual-stream + architecture with joint (multimodal) and single (image-only) transformer layers. + + Standard Configurations: + - Flux.1 [dev/schnell]: 12B parameters (19 joint + 38 single layers) + - Flux 535M: Minimal config for testing (1 joint + 1 single layer) + + Key Configuration Groups: + - Architecture: num_joint_layers, num_single_layers, hidden_size + - Context: context_dim (T5), vec_in_dim (CLIP), model_channels (timestep) + - Position Encoding: theta, axes_dim, rotary_interleaved + - Guidance: guidance_embed, guidance_scale + + For architecture details, see Flux class documentation. + + Attributes: + num_joint_layers: Number of joint (multimodal) transformer layers + num_single_layers: Number of single (image-only) transformer layers + context_dim: Dimension of text context embeddings (T5-XXL: 4096) + vec_in_dim: Dimension of pooled text embeddings (CLIP-L: 768) + model_channels: Channels for timestep embedding (default: 256) + guidance_embed: Whether to use guidance embedding for CFG + guidance_scale: Guidance scale for classifier-free guidance (default: 3.5) + theta: Base for RoPE position embeddings (default: 10000) + axes_dim: Dimension for each axis in 3D RoPE (default: [16, 56, 56]) + rotary_interleaved: Whether RoPE dimensions are interleaved (default: True) + apply_rope_fusion: Whether to apply RoPE fusion optimization (default: False) + add_qkv_bias: Whether to add bias to QKV projections (default: True) + single_block_bias: Whether to add bias to single block linear layers (default: True) + activation_func: Activation function (default: openai_gelu_no_jit) + use_te_rng_tracker: Whether to use Transformer Engine RNG tracker (default: False) + """ + + # Model identification + model_type: str = "flux" + + # Dummy layer for compatibility (actual layers defined below) + num_layers: int = 1 # Not used in Flux, kept for compatibility + + # Architecture: Number of layers + num_joint_layers: int = 19 # Default: Flux 12B + num_single_layers: int = 38 # Default: Flux 12B + + # Architecture: Dimensions (Flux standard: 3072) + hidden_size: int = 3072 + num_attention_heads: int = 24 + + # Input dimensions + in_channels: int = 64 # Packed latent channels (16 VAE channels * 4 from 2x2 patch packing) + + # Context dimensions + context_dim: int = 4096 # T5-XXL hidden dimension + vec_in_dim: int = 768 # CLIP-L pooled dimension + model_channels: int = 256 # Channels for timestep embedding + + # Guidance (for classifier-free guidance) + guidance_embed: bool = False # Set True for CFG support + guidance_scale: float = 3.5 # CFG guidance scale + cfg_dropout_prob: float = ( + 0.0 # Probability of replacing text embeddings with empty encodings (MLPerf: 0.1) + ) + + # Training: timestep sampling strategy + # Options: "logit_normal" (SD3 default), "direct_uniform" (NVIDIA MLPerf), "uniform", "mode" + timestep_sampling_strategy: str = "logit_normal" + + # Position embeddings (RoPE) + theta: int = 10000 # Base for RoPE frequencies + axes_dim: tuple = (16, 56, 56) # 3D RoPE: (channels, height, width) + rotary_interleaved: bool = True # Whether RoPE dimensions are interleaved + apply_rope_fusion: bool = False # Whether to apply RoPE fusion optimization + + # Patchification (Flux uses 1x1 patches by default) + patch_size: int = 1 + + # Attention configuration + add_qkv_bias: bool = True # Whether to add bias to QKV projections + + # Single block configuration + single_block_bias: bool = True # Whether to add bias to single block linear layers + + # Default: non-JIT tanh GELU (ROCm-safe). Override via YAML activation_func: + # "openai_gelu" -> fused F.gelu(approximate="tanh"); "erf_gelu" -> erf-based GELU. + activation_func: Callable = openai_gelu_no_jit + + # Dropout (Flux typically uses 0) + hidden_dropout: float = 0.0 + attention_dropout: float = 0.0 + + # Normalization + layernorm_epsilon: float = 1e-6 + + # Initialization + use_cpu_initialization: bool = True + + # Optimization flags + gradient_accumulation_fusion: bool = False + use_dual_fp8_output_projection: bool = False + use_triton_ops: bool = False + adaln_plain_ops: bool = False + adaln_always_jit_fuser: bool = False + + # FSDP2 prefetch depth: number of layers to prefetch ahead for all-gather overlap. + # 1 = prefetch next layer (default), 2 = prefetch next 2 layers, 0 = no prefetch. + fsdp_prefetch_depth: int = 1 + + # FP8 all-gather data cache: batch-quantize all weights upfront and cache + # the FP8 data so fsdp_pre_all_gather skips per-layer quantization. + # When false, only scales are precomputed; quantization happens on-demand. + fp8_precompute_data_cache: bool = True + + # Optimizer foreach batching: use _foreach_copy_ for batched grad/weight + # copies instead of per-parameter loops. + optimizer_foreach: bool = True + + # Overlap grad norm with reduce-scatter: accumulate squared norms + # incrementally inside FSDP2's post-reduce stream via + # register_post_accumulate_grad_hook, replacing the full recompute + # in clip_grad_norm with a single all-reduce + sqrt + clip. + overlap_grad_norm: bool = False + + # Use the C++ quantize_fp8 kernel from primus_turbo for tensorwise FP8 + # quantization instead of the default inline implementation. + use_cpp_fp8_quantize: bool = False + + # CUDA graph support + enable_cuda_graph: bool = False + cuda_graph_scope: Optional[str] = None # Options: "full", "full_iteration" + cuda_graph_warmup_steps: int = 2 + + # Transformer Engine + use_te_rng_tracker: bool = False + + # Torch compile configuration (applied AFTER distributed setup) + enable_torch_compile: bool = False # Enable compilation (applied after FSDP/DDP wrapping) + torch_compile_backend: str = "inductor" # Compilation backend + torch_compile_mode: str = "default" # "default", "reduce-overhead", "max-autotune" + torch_compile_fullgraph: bool = False # Allow graph breaks (recommended for distributed) + torch_compile_optimizer: bool = False # Whether to compile optimizer step + torch_compile_optimizer_scope: str = ( + "full" # "full" = compile FSDP2FP32Optimizer.step(), "inner_only" = compile only inner AdamW.step() + ) + + # Selective stack compilation (TE spec only — local spec falls back to per_block) + torch_compile_strategy: str = "per_block" + # "per_block": compile each transformer layer individually (default, recommended) + # "whole_model": compile entire forward (incompatible with overlap_param_gather) + # "double_stack": compile only the double (joint) block loop + # "single_stack": compile only the single block loop + # "stack": compile both double and single block loops + # "full_dit": compile the entire DiT (double + cat + single + output) as one region + torch_compile_replace_qk_rmsnorm: bool = False + torch_compile_disable_inductor_cudagraphs: bool = True + torch_compile_emulate_precision_casts: bool = True # Preserve eager BF16 precision in Triton kernels + torch_compile_fused_ln_modulate: bool = True # Use fused LN+modulate Triton kernel in AdaLN + + def __post_init__(self): + """Post-initialization processing.""" + # MUST set num_layers BEFORE super().__post_init__() -- convention for BaseDiffusionConfig + # Required because BaseDiffusionConfig maps sensitive_layer_* fields using num_layers, + # and TransformerConfig validates num_layers_at_start_in_bf16 against num_layers. + self.num_layers = self.num_joint_layers + self.num_single_layers + + super().__post_init__() # BaseDiffusionConfig (mapping) -> TransformerConfig (validation) + + # Xavier uniform for Megatron parallel linear layers (common Flux reference default). + self.init_method = nn.init.xavier_uniform_ + self.output_layer_init_method = nn.init.xavier_uniform_ + + # Ensure axes_dim is a tuple + if isinstance(self.axes_dim, list): + self.axes_dim = tuple(self.axes_dim) + + def validate(self): + """ + Validate Flux-specific configuration. + + Raises: + ValueError: If configuration is invalid + """ + # Call parent validation + super().validate() + + # Flux-specific validations + if self.num_joint_layers <= 0: + raise ValueError(f"num_joint_layers must be positive, got {self.num_joint_layers}") + + if self.num_single_layers <= 0: + raise ValueError(f"num_single_layers must be positive, got {self.num_single_layers}") + + if self.context_dim <= 0: + raise ValueError(f"context_dim must be positive, got {self.context_dim}") + + if self.vec_in_dim <= 0: + raise ValueError(f"vec_in_dim must be positive, got {self.vec_in_dim}") + + if self.theta <= 0: + raise ValueError(f"theta must be positive, got {self.theta}") + + if len(self.axes_dim) != 3: + raise ValueError(f"axes_dim must have 3 elements for 3D RoPE, got {len(self.axes_dim)}") + + if any(d <= 0 for d in self.axes_dim): + raise ValueError(f"All axes_dim values must be positive, got {self.axes_dim}") + + # Validate RoPE fusion constraints + if self.apply_rope_fusion: + import warnings + + warnings.warn( + "\n" + "=" * 80 + "\n" + "RoPE Fusion Enabled: Same-Resolution Batch Requirement\n" + "=" * 80 + "\n" + "RoPE fusion optimization requires ALL images in each training batch to have\n" + "the SAME resolution (height and width). Variable-resolution batches will\n" + "produce incorrect positional encodings and degrade model quality.\n" + "\n" + "Recommended data pipeline configurations:\n" + " 1. Fixed resolution: All images resized to same size (e.g., 512x512)\n" + " 2. Resolution bucketing: Group images by resolution in separate batches\n" + " 3. Aspect ratio bucketing: Use bucketing with consistent dimensions\n" + "\n" + "=" * 80, + UserWarning, + stacklevel=2, + ) + + def get_num_layers(self): + """ + Get total number of transformer layers. + + Returns: + Total number of layers (joint + single) + """ + return self.num_joint_layers + self.num_single_layers + + @classmethod + def flux_535m(cls, **kwargs): + """ + Create configuration for Flux 535M (minimal model for testing). + + Args: + **kwargs: Override default parameters + + Returns: + FluxConfig instance + """ + defaults = { + "num_joint_layers": 1, + "num_single_layers": 1, + "hidden_size": 3072, + "num_attention_heads": 24, + } + defaults.update(kwargs) + return cls(**defaults) + + @classmethod + def flux_12b(cls, **kwargs): + """ + Create configuration for Flux 12B (standard model). + + Args: + **kwargs: Override default parameters + + Returns: + FluxConfig instance + """ + defaults = { + "num_joint_layers": 19, + "num_single_layers": 38, + "hidden_size": 3072, + "num_attention_heads": 24, + } + defaults.update(kwargs) + return cls(**defaults) diff --git a/primus/backends/megatron/core/models/diffusion/flux/layer_spec.py b/primus/backends/megatron/core/models/diffusion/flux/layer_spec.py new file mode 100644 index 000000000..9a1932b09 --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/flux/layer_spec.py @@ -0,0 +1,594 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Flux layer specifications for Megatron-Core integration. + +This module implements transformer layers for Flux's MMDiT architecture: + - MMDiTLayer: Joint image-text transformer block + - FluxSingleTransformerBlock: Image-only transformer block + - Factory functions for creating layer specs with Transformer Engine + +Key Innovation: Heterogeneous layer support via TransformerBlock + Instead of separate ModuleLists, Primus uses a unified TransformerBlock + with heterogeneous layer specifications, enabling better pipeline + parallelism and cleaner checkpoint management. + +Reference: + - Flux Paper: "Flux: A Scalable Diffusion Model" + - MMDiT Paper: "Scaling Rectified Flow Transformers" + - Megatron-Core: Heterogeneous layer patterns +""" + +import copy +from typing import Optional, Tuple + +import torch +import torch.nn as nn +from megatron.core.models.backends import BackendSpecProvider +from megatron.core.transformer.cuda_graphs import CudaGraphManager +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.mlp import MLP, MLPSubmodules +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import ( + TransformerLayer, + TransformerLayerSubmodules, +) +from megatron.core.utils import make_viewless_tensor +from torch import Tensor + +from primus.backends.megatron.core.models.diffusion.common.normalization import ( + AdaLN, + AdaLNContinuous, +) +from primus.backends.megatron.core.models.diffusion.flux.attention import ( + FluxSingleAttention, + JointSelfAttention, + JointSelfAttentionSubmodules, +) + +try: + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + from megatron.core.transformer.attention import SelfAttentionSubmodules + + HAVE_TE_SPEC_PROVIDER = True +except ImportError: + HAVE_TE_SPEC_PROVIDER = False + SelfAttentionSubmodules = None + +# Try to import PrimusTurboLocalSpecProvider (may not exist in all versions) +try: + from primus.backends.megatron.core.extensions.primus_turbo_local_spec import ( + PrimusTurboFloat8LocalSpecProvider, + PrimusTurboLocalSpecProvider, + ) + + HAVE_PRIMUS_TURBO_LOCAL = True +except ImportError: + HAVE_PRIMUS_TURBO_LOCAL = False + PrimusTurboLocalSpecProvider = None + PrimusTurboFloat8LocalSpecProvider = None + +# MXFP4 provider in a separate guard so a missing primus_turbo_mxfp4_local +# doesn't break FP8 imports. +try: + from primus.backends.megatron.core.extensions.primus_turbo_local_spec import ( + PrimusTurboMXFP4LocalSpecProvider, + ) +except ImportError: + PrimusTurboMXFP4LocalSpecProvider = None + + +class MMDiTLayer(TransformerLayer): + """ + Multimodal Diffusion Transformer (MMDiT) Layer. + + Processes image and text streams jointly via AdaLN-conditioned attention + and separate MLPs with gated residual connections. Used in Flux's "double + blocks" for cross-modal image-text interaction. + + Args: + config: Transformer configuration + submodules: Submodule specifications for attention, MLP, etc. + layer_number: Layer index in the model (default: 1) + context_pre_only: If True, context stream only computes pre-attention norm + + Input/Output: + (hidden_states [S_img, B, H], context [S_txt, B, H], timestep_emb [B, H]) + -> (hidden_states, context) with same shapes + + Reference: + - "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis" + """ + + def __init__( + self, + config: TransformerConfig, + submodules: TransformerLayerSubmodules, + layer_number: int = 1, + context_pre_only: bool = False, + **kwargs, # Accept additional kwargs from TransformerBlock (e.g., pg_collection) + ): + hidden_size = config.hidden_size + super().__init__(config=config, submodules=submodules, layer_number=layer_number, **kwargs) + + # Enable per-layer CUDA graph if configured + if config.enable_cuda_graph and config.cuda_graph_scope != "full_iteration": + self.cudagraph_manager = CudaGraphManager(config, share_cudagraph_io_buffers=False) + + # Adaptive layer normalization for main stream (image). + # init_method=nn.init.normal_ draws RNG matching NeMo's DiT init + # sequence so cross-framework convergence comparisons line up. + # init_weights() (model.py) immediately re-zeroes these weights, + # so the only observable effect is RNG-sequence alignment. + self.adaln = AdaLN( + config, + modulation_bias=True, + n_adaln_chunks=6, + use_second_norm=True, + init_method=nn.init.normal_, + ) + + # Adaptive layer normalization for context stream (text) + self.context_pre_only = context_pre_only + context_norm_type = "ada_norm_continuous" if context_pre_only else "ada_norm_zero" + + if context_norm_type == "ada_norm_continuous": + # Continuous AdaLN for context (simpler, used when context is pre-only) + self.adaln_context = AdaLNContinuous( + config, hidden_size, modulation_bias=True, norm_type="layer_norm" + ) + elif context_norm_type == "ada_norm_zero": + # Full AdaLN for context (used when context has full processing). + # See note on self.adaln above re: init_method choice. + self.adaln_context = AdaLN( + config, + modulation_bias=True, + n_adaln_chunks=6, + use_second_norm=True, + init_method=nn.init.normal_, + ) + else: + raise ValueError( + f"Unknown context_norm_type: {context_norm_type}, " + f"currently only support `ada_norm_continuous`, `ada_norm_zero`" + ) + + # Context MLP (only if not pre-only) + if not context_pre_only: + # Disable context parallelism for context MLP + cp_override_config = copy.deepcopy(config) + cp_override_config.context_parallel_size = 1 + cp_override_config.tp_comm_overlap = False + + from megatron.core.transformer.spec_utils import build_module + + self.context_mlp = build_module(submodules.mlp, config=cp_override_config) + else: + self.context_mlp = None + + def forward( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor] = None, + context: Optional[Tensor] = None, + context_mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + timestep_emb: Optional[Tensor] = None, + packed_seq_params=None, + **kwargs, + ) -> Tuple[Tensor, Tensor]: + """ + Forward pass: Joint processing of image and text tokens. + + Args: + hidden_states: Image tokens [seq_img, batch, hidden] + context: Text tokens [seq_txt, batch, hidden] + timestep_emb: Timestep conditioning [batch, hidden] (required) + + Returns: + Tuple of (hidden_states, context) - both updated + """ + # Map TransformerBlock's 'context' to internal 'encoder_hidden_states' + encoder_hidden_states = context + + if timestep_emb is None: + raise ValueError("MMDiTLayer requires timestep_emb for AdaLN conditioning.") + + emb = timestep_emb + + # Get modulation parameters for main stream (image) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln(emb) + + # Apply modulated layer norm to image tokens + norm_hidden_states = self.adaln.modulated_layernorm( + hidden_states, shift=shift_msa, scale=scale_msa, layernorm_idx=0 + ) + + # Apply modulated layer norm to text tokens + if self.context_pre_only: + # Continuous AdaLN (simpler) + norm_encoder_hidden_states = self.adaln_context(encoder_hidden_states, emb) + else: + # Full AdaLN with modulation parameters + c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.adaln_context( + emb + ) + norm_encoder_hidden_states = self.adaln_context.modulated_layernorm( + encoder_hidden_states, shift=c_shift_msa, scale=c_scale_msa, layernorm_idx=0 + ) + + # Joint self-attention + attn_output, context_attn_output = self.self_attention( + norm_hidden_states, + attention_mask=attention_mask, + rotary_pos_emb=rotary_pos_emb, + additional_hidden_states=norm_encoder_hidden_states, + ) + + # MLP for main stream (image) + # Fused operation: gated residual + modulated layernorm for MLP input + hidden_states, norm_hidden_states = self.adaln.scaled_modulated_layernorm( + residual=hidden_states, + x=attn_output, + gate=gate_msa, + shift=shift_mlp, + scale=scale_mlp, + layernorm_idx=1, + ) + mlp_output, mlp_bias = self.mlp(norm_hidden_states) + hidden_states = self.adaln.scale_add(hidden_states, x=(mlp_output + mlp_bias), gate=gate_mlp) + + # MLP for context stream (text) - only if not pre-only + if not self.context_pre_only: + # Fused operation: gated residual + modulated layernorm for context MLP input + encoder_hidden_states, norm_encoder_hidden_states = self.adaln_context.scaled_modulated_layernorm( + residual=encoder_hidden_states, + x=context_attn_output, + gate=c_gate_msa, + shift=c_shift_mlp, + scale=c_scale_mlp, + layernorm_idx=1, + ) + context_mlp_output, context_mlp_bias = self.context_mlp(norm_encoder_hidden_states) + encoder_hidden_states = self.adaln_context.scale_add( + encoder_hidden_states, x=(context_mlp_output + context_mlp_bias), gate=c_gate_mlp + ) + + # Make output viewless for MPU checkpoint compatibility + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True + ) + encoder_hidden_states = make_viewless_tensor( + inp=encoder_hidden_states, requires_grad=encoder_hidden_states.requires_grad, keep_graph=True + ) + + return hidden_states, encoder_hidden_states + + def __call__(self, *args, **kwargs): + """Override call to support CUDA graphs.""" + if hasattr(self, "cudagraph_manager"): + return self.cudagraph_manager(self, args, kwargs) + return super(MegatronModule, self).__call__(*args, **kwargs) + + +class FluxSingleTransformerBlock(TransformerLayer): + """ + Flux Single Transformer Block (image-only processing). + + Processes image tokens with parallel attention + MLP paths, AdaLN + conditioning, and gated residual connections. Used in Flux's "single + blocks" after the joint MMDiT layers. + + Args: + config: Transformer configuration + submodules: Submodule specifications for attention, MLP, etc. + layer_number: Layer index in the model (default: 1) + mlp_ratio: MLP hidden size ratio (default: 4) + n_adaln_chunks: AdaLN modulation chunks (default: 3 for shift, scale, gate) + modulation_bias: Whether to use bias in modulation layers (default: True) + + Input/Output: + (hidden_states [S, B, H], timestep_emb [B, H]) -> (hidden_states, None) + """ + + def __init__( + self, + config: TransformerConfig, + submodules: TransformerLayerSubmodules, + layer_number: int = 1, + mlp_ratio: int = 4, + n_adaln_chunks: int = 3, + modulation_bias: bool = True, + **kwargs, # Accept additional kwargs from TransformerBlock (e.g., pg_collection) + ): + # Override add_bias_linear with single_block_bias for this block + # This allows independent control of bias in single blocks vs joint blocks + original_add_bias_linear = config.add_bias_linear + if hasattr(config, "single_block_bias"): + config.add_bias_linear = config.single_block_bias + + super().__init__(config=config, submodules=submodules, layer_number=layer_number, **kwargs) + + # Restore original value + config.add_bias_linear = original_add_bias_linear + + # Enable per-layer CUDA graph if configured + if config.enable_cuda_graph and config.cuda_graph_scope != "full_iteration": + self.cudagraph_manager = CudaGraphManager(config, share_cudagraph_io_buffers=False) + + # Adaptive layer normalization + # n_adaln_chunks=3 for (shift, scale, gate) + # init_method=nn.init.normal_: see note on FluxTransformerBlock's + # self.adaln above -- NeMo-aligned RNG draw, re-zeroed in init_weights(). + self.adaln = AdaLN( + config=config, + n_adaln_chunks=n_adaln_chunks, + modulation_bias=modulation_bias, + use_second_norm=False, # Single block only uses one norm + init_method=nn.init.normal_, + ) + + def forward( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor] = None, + context: Optional[Tensor] = None, + context_mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + timestep_emb: Optional[Tensor] = None, + packed_seq_params=None, + **kwargs, + ) -> Tuple[Tensor, None]: + """ + Single-stream block: first single layer may receive separate image and text streams + (concatenates them); later layers see one concatenated stream with context=None. + + Args: + hidden_states: Image tokens [seq_img, B, H] or concatenated [seq_total, B, H] + context: Text tokens [seq_txt, B, H] for the first single block only + timestep_emb: Timestep conditioning [B, H] (required) + + Returns: + Tuple of (hidden_states, None) + """ + if timestep_emb is None: + raise ValueError("FluxSingleTransformerBlock requires timestep_emb for AdaLN conditioning.") + + emb = timestep_emb + + # TRANSITION HANDLING: Concatenate if this is first single block + if context is not None: + # This is layer 19 - concatenate text and image + hidden_states = torch.cat([context, hidden_states], dim=0) + + residual = hidden_states + + # Get modulation parameters (shift, scale, gate) + shift, scale, gate = self.adaln(emb) + + # Apply modulated layer normalization + norm_hidden_states = self.adaln.modulated_layernorm(hidden_states, shift=shift, scale=scale) + + # MLP path + mlp_hidden_states, mlp_bias = self.mlp(norm_hidden_states) + + # Attention path + attention_output, attention_bias = self.self_attention( + norm_hidden_states, + attention_mask=attention_mask, + rotary_pos_emb=rotary_pos_emb, + ) + + # Combine MLP and attention (parallel paths) + hidden_states = mlp_hidden_states + mlp_bias + attention_output + if attention_bias is not None: + hidden_states = hidden_states + attention_bias + + # Gated residual connection + hidden_states = self.adaln.scale_add(residual, x=hidden_states, gate=gate) + + return hidden_states, None + + def __call__(self, *args, **kwargs): + """Override call to support CUDA graphs.""" + if hasattr(self, "cudagraph_manager"): + return self.cudagraph_manager(self, args, kwargs) + return super(MegatronModule, self).__call__(*args, **kwargs) + + +def get_flux_single_transformer_spec_for_backend( + backend: BackendSpecProvider, +) -> ModuleSpec: + """ + Get ModuleSpec for Flux single transformer block with backend support. + + This factory function creates the specification for a Flux single block + using the provided backend spec provider for layer implementations. + + Args: + backend: BackendSpecProvider (e.g., TESpecProvider, PrimusTurboSpecProvider) + + Returns: + ModuleSpec for FluxSingleTransformerBlock with backend submodules + """ + return ModuleSpec( + module=FluxSingleTransformerBlock, + submodules=TransformerLayerSubmodules( + self_attention=ModuleSpec( + module=FluxSingleAttention, + params={"attn_mask_type": AttnMaskType.no_mask}, + submodules=SelfAttentionSubmodules( + linear_qkv=backend.column_parallel_linear(), + core_attention=backend.core_attention(), + q_layernorm=backend.layer_norm(rms_norm=True, for_qk=True), + k_layernorm=backend.layer_norm(rms_norm=True, for_qk=True), + linear_proj=backend.row_parallel_linear(), + ), + ), + mlp=ModuleSpec( + module=MLP, + submodules=MLPSubmodules( + linear_fc1=backend.column_parallel_linear(), + linear_fc2=backend.row_parallel_linear(), + ), + ), + ), + ) + + +def get_flux_double_transformer_spec_for_backend( + backend: BackendSpecProvider, +) -> ModuleSpec: + """ + Get ModuleSpec for Flux double (joint) transformer block with backend support. + + This factory function creates the specification for a Flux MMDiT layer + using the provided backend spec provider for layer implementations. + + Args: + backend: BackendSpecProvider (e.g., TESpecProvider, PrimusTurboSpecProvider) + + Returns: + ModuleSpec for MMDiTLayer with backend submodules + """ + return ModuleSpec( + module=MMDiTLayer, + submodules=TransformerLayerSubmodules( + self_attention=ModuleSpec( + module=JointSelfAttention, + params={"attn_mask_type": AttnMaskType.no_mask}, + submodules=JointSelfAttentionSubmodules( + q_layernorm=backend.layer_norm(rms_norm=True, for_qk=True), + k_layernorm=backend.layer_norm(rms_norm=True, for_qk=True), + added_q_layernorm=backend.layer_norm(rms_norm=True, for_qk=True), + added_k_layernorm=backend.layer_norm(rms_norm=True, for_qk=True), + linear_qkv=backend.column_parallel_linear(), + added_linear_qkv=backend.column_parallel_linear(), + core_attention=backend.core_attention(), + linear_proj=backend.row_parallel_linear(), + ), + ), + mlp=ModuleSpec( + module=MLP, + submodules=MLPSubmodules( + linear_fc1=backend.column_parallel_linear(), + linear_fc2=backend.row_parallel_linear(), + ), + ), + ), + ) + + +def get_flux_layer_spec( + config, + backend: Optional[BackendSpecProvider] = None, + vp_stage: Optional[int] = None, + pp_rank: Optional[int] = None, +): + """ + Create heterogeneous layer specifications for Flux TransformerBlock. + + Builds a list of ModuleSpec objects for a unified TransformerBlock: + - Layers 0 to num_joint_layers-1: MMDiTLayer specs (joint image-text processing) + - Layers num_joint_layers to (num_joint_layers + num_single_layers - 1): + FluxSingleTransformerBlock specs (concatenated processing) + + The layer specs are automatically sliced for pipeline parallelism based on + pp_rank and vp_stage, following the pattern from get_gpt_heterogeneous_layer_spec. + + Args: + config: FluxConfig with num_joint_layers, num_single_layers, hidden_size, etc. + backend: BackendSpecProvider (auto-selected based on config.transformer_impl + and available providers if None) + vp_stage: Virtual pipeline stage number (for interleaved PP, optional) + pp_rank: Pipeline parallel rank (optional, auto-detected if None) + + Returns: + TransformerBlockSubmodules containing layer_specs and layer_norm + + Example: + >>> from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + >>> config = FluxConfig.flux_12b() + >>> backend = TESpecProvider() + >>> spec = get_flux_layer_spec(config, backend=backend) + >>> transformer = TransformerBlock(config=config, spec=spec) + + Reference: + Similar to megatron/core/models/gpt/heterogeneous/heterogeneous_layer_specs.py + """ + from megatron.core import parallel_state + from megatron.core.transformer.transformer_block import ( + TransformerBlockSubmodules, + get_num_layers_to_build, + ) + from megatron.core.transformer.transformer_layer import get_transformer_layer_offset + + # Default backend selection based on config + sensitive_backend = None + + if backend is None: + if config.transformer_impl == "local": + if config.fp4 is not None and PrimusTurboMXFP4LocalSpecProvider is not None: + backend = PrimusTurboMXFP4LocalSpecProvider() + + # Resolve sensitive layer backend + sensitive_precision = getattr(config, "sensitive_layer_precision", "bf16") + if sensitive_precision == "tw_fp8": + sensitive_backend = PrimusTurboFloat8LocalSpecProvider() + elif sensitive_precision == "bf16": + sensitive_backend = PrimusTurboLocalSpecProvider() + elif ( + config.fp8 is not None + and HAVE_PRIMUS_TURBO_LOCAL + and PrimusTurboFloat8LocalSpecProvider is not None + ): + backend = PrimusTurboFloat8LocalSpecProvider() + elif HAVE_PRIMUS_TURBO_LOCAL and PrimusTurboLocalSpecProvider is not None: + backend = PrimusTurboLocalSpecProvider() + else: + from megatron.core.models.backends import LocalSpecProvider + + backend = LocalSpecProvider() + elif HAVE_TE_SPEC_PROVIDER: + # Use TransformerEngine (current default) + backend = TESpecProvider() + else: + # Fallback to pure native Megatron + from megatron.core.models.backends import LocalSpecProvider + + backend = LocalSpecProvider() + + # Build per-layer specs with optional sensitive-layer heterogeneity + sensitive_enabled = getattr(config, "sensitive_layers_enabled", False) + num_start = getattr(config, "sensitive_layers_start", 0) if sensitive_enabled else 0 + num_end = getattr(config, "sensitive_layers_end", 0) if sensitive_enabled else 0 + total = config.num_joint_layers + config.num_single_layers + + layer_specs = [] + for i in range(total): + is_sensitive = sensitive_backend is not None and ((i < num_start) or (i >= total - num_end)) + layer_backend = sensitive_backend if is_sensitive else backend + + if i < config.num_joint_layers: + layer_specs.append(get_flux_double_transformer_spec_for_backend(layer_backend)) + else: + layer_specs.append(get_flux_single_transformer_spec_for_backend(layer_backend)) + + # Slice for pipeline parallelism (only if parallel state is initialized) + try: + if parallel_state.model_parallel_is_initialized(): + offset = get_transformer_layer_offset(config, vp_stage=vp_stage, pp_rank=pp_rank) + num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage, pp_rank=pp_rank) + layer_specs = layer_specs[offset : offset + num_layers_to_build] + except (AssertionError, RuntimeError): + # Parallel state not initialized - use all layers (for single-device testing) + pass + + # Get layer norm from backend + layer_norm = backend.layer_norm(rms_norm=False, for_qk=False) + + return TransformerBlockSubmodules(layer_specs=layer_specs, layer_norm=layer_norm) diff --git a/primus/backends/megatron/core/models/diffusion/flux/layers.py b/primus/backends/megatron/core/models/diffusion/flux/layers.py new file mode 100644 index 000000000..62ae3eb7c --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/flux/layers.py @@ -0,0 +1,130 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Flux-specific layers and components. + +This module implements components unique to the Flux architecture: + - EmbedND: Multi-dimensional Rotary Position Embedding (3D RoPE) + - Helper functions for position encoding + +Reference: + - Flux Paper: "Flux: A Scalable Diffusion Model" + - RoPE Paper: "RoFormer: Enhanced Transformer with Rotary Position Embedding" + - Adapted from NeMo's Flux layers +""" + +from typing import List + +import torch +import torch.nn as nn +from torch import Tensor + + +def rope(pos: Tensor, dim: int, theta: int) -> Tensor: + """ + Generate RoPE (Rotary Position Embedding) frequencies for given positions. + + This is adapted for Megatron-Core's attention implementation, which + calculates sin/cos internally. We only generate the frequency matrix here. + + Args: + pos: Position indices [..., n] where n is the number of positions + dim: Dimension of the embedding (must be even) + theta: Base for frequency computation (typically 10000) + + Returns: + Frequency matrix [..., n, dim/2] for RoPE computation + + Note: + This differs from standard RoPE implementations because Megatron + attention applies sin/cos internally, so we only provide frequencies. + + Reference: + - RoPE: "RoFormer: Enhanced Transformer with Rotary Position Embedding" + - Adapted from NeMo's Flux layers + """ + if dim % 2 != 0: + raise ValueError("The dimension must be even for RoPE.") + + # Compute scaling factors for each dimension pair + # scale: [0, 2, 4, ..., dim-2] / dim = [0, 1/dim, 2/dim, ..., (dim-2)/dim] + scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim + + # Compute base frequencies: 1 / (theta ^ scale) + omega = 1.0 / (theta**scale) + + # Outer product of positions and frequencies + # pos: [..., n], omega: [dim/2] -> out: [..., n, dim/2] + out = torch.einsum("...n,d->...nd", pos, omega) + + return out.float() + + +class EmbedND(nn.Module): + """ + Multi-dimensional Rotary Position Embedding (RoPE) for images. + + Flux uses 3D RoPE with three axes: + - Axis 0: Always 0 for 2D images (reserved for video/temporal) + - Axis 1: Height positions (y-coordinates) + - Axis 2: Width positions (x-coordinates) + + Each axis gets independent sinusoidal frequencies, concatenated to form + the complete position embedding. + + Args: + dim: Model hidden dimension (stored for reference, not used in forward) + theta: Base for frequency computation (default: 10000) + axes_dim: Dimensions per axis, e.g., [16, 56, 56] + + Input: + ids: Position IDs [B, seq, num_axes] -- for images: [B, H*W, 3] + + Output: + RoPE frequency matrix [seq, B, 1, dim] for Megatron attention + + Reference: + - Flux: 3D RoPE for spatial + channel position encoding + - RoPE: "RoFormer: Enhanced Transformer with Rotary Position Embedding" + - Adapted from NeMo's Flux layers + """ + + def __init__(self, dim: int, theta: int, axes_dim: List[int]): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: Tensor) -> Tensor: + """ + Compute 3D RoPE frequencies from position IDs. + + Args: + ids: Position IDs [B, seq, num_axes] where num_axes = len(axes_dim) + + Returns: + RoPE frequency matrix [seq, B, 1, dim] for Megatron attention + """ + n_axes = ids.shape[-1] + + # Generate RoPE frequencies for each axis and concatenate + # For each axis i: + # - ids[..., i]: positions for that axis [B, seq] + # - rope(): generates frequencies [B, seq, axes_dim[i]/2] + # Concatenate along last dimension to get [B, seq, dim/2] + emb = torch.cat( + [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], + dim=-1, + ) + + # Reshape for Megatron attention format: + # 1. Add sequence dimension: [B, seq, dim/2] -> [B, 1, seq, dim/2] + # 2. Permute to [seq, B, 1, dim/2] + emb = emb.unsqueeze(1).permute(2, 0, 1, 3) + + # Stack [cos, sin] pairs and reshape to final format + # torch.stack([emb, emb], dim=-1): [seq, B, 1, dim/2, 2] + # reshape: [seq, B, 1, dim] + return torch.stack([emb, emb], dim=-1).reshape(*emb.shape[:-1], -1) diff --git a/primus/backends/megatron/core/models/diffusion/flux/model.py b/primus/backends/megatron/core/models/diffusion/flux/model.py new file mode 100644 index 000000000..b451045fa --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/flux/model.py @@ -0,0 +1,857 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Portions copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Flux diffusion model implementation. + +This module implements the complete Flux model, integrating all components: +- Input embeddings and position encodings +- MMDiT (joint image-text) transformer blocks +- Single (image-only) transformer blocks +- Output projection + +API Note: + Primus's Flux model uses Megatron-Core's TransformerBlock for unified + layer management, enabling heterogeneous layer types (MMDiT + Single) + in a single container. + + Architecture: + - Input: Packed latents [S, B, C*4] (pre-processed) + - Processing: Unified TransformerBlock with heterogeneous specs + - Output: Packed predictions [S, B, C*4] + - S = H*W/4 (2x2 spatial patches grouped into sequence tokens) + + Key Enhancement (vs traditional ModuleList approach): + - Unified TransformerBlock with layer_specs for heterogeneous layers + - Automatic pipeline parallelism splitting + - Cleaner distributed checkpoint format + +Reference: + - Flux Paper: "Flux: A Scalable Diffusion Model for High-Resolution Image Synthesis" + - Megatron-Core: Heterogeneous TransformerBlock patterns +""" + +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.utils import sharded_state_dict_default +from torch import Tensor + +from primus.backends.megatron.core.models.common.diffusion_module.diffusion_module import ( + DiffusionModule, +) +from primus.backends.megatron.core.models.diffusion.common.embeddings import ( + MLPEmbedder, + TimeStepEmbedder, +) +from primus.backends.megatron.core.models.diffusion.common.normalization import ( + AdaLNContinuous, +) +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.layer_spec import ( + get_flux_layer_spec, +) +from primus.backends.megatron.core.models.diffusion.flux.layers import EmbedND +from primus.backends.megatron.core.transformer.diffusion_transformer_block import ( + DiffusionTransformerBlock, +) + +_QK_RMSNORM_PARAM_ATTRS = ( + "sequence_parallel", + "shared", + "allreduce", + "tensor_model_parallel", + "partition_dim", + "partition_stride", +) + + +class _TorchRMSNorm(nn.Module): + """Compile-friendly RMSNorm replacement for TE's QK LayerNorm. + + TE's TENorm is decorated with @no_torch_dynamo, causing graph breaks + inside compiled regions. This pure-PyTorch version is fully traceable. + Matches NeMo's _TorchRMSNorm (custom_flux.py). + """ + + def __init__(self, hidden_size: int, eps: float = 1e-6, *, device=None, dtype=None): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size, device=device, dtype=dtype)) + self.bias = None + + def reset_parameters(self): + nn.init.ones_(self.weight) + + def forward(self, x: Tensor) -> Tensor: + mean_sq = x.square().mean(dim=-1, keepdim=True, dtype=torch.float32) + inv_rms = torch.rsqrt(mean_sq + self.eps).to(dtype=x.dtype) + return x * inv_rms * self.weight + + +def _build_local_qk_rmsnorm(norm: nn.Module) -> "_TorchRMSNorm": + """Build a _TorchRMSNorm from an existing TE norm, preserving weights and attrs.""" + weight = norm.weight + local_norm = _TorchRMSNorm( + weight.numel(), + eps=getattr(norm, "eps", 1e-6), + device=weight.device, + dtype=weight.dtype, + ) + local_norm.weight.data.copy_(weight.data) + local_norm.weight.requires_grad_(weight.requires_grad) + for attr_name in _QK_RMSNORM_PARAM_ATTRS: + if hasattr(weight, attr_name): + setattr(local_norm.weight, attr_name, getattr(weight, attr_name)) + return local_norm + + +class Flux(DiffusionModule): + """ + Flux: Flow-based diffusion model with MMDiT architecture. + + Uses Megatron-Core's TransformerBlock with heterogeneous layer specs, + combining MMDiT "double blocks" (joint image-text) with "single blocks" + (image-only) for text-to-image generation. + + Model Variants: + - Flux 535M: 1 joint + 1 single layer (~535M params, testing) + - Flux 12B: 19 joint + 38 single layers (~12B params, production) + + Args: + config: FluxConfig with all model parameters + encoder_configs: Optional encoder configurations (VAE, T5, CLIP) + pg_collection: ProcessGroupCollection for distributed training + backend: Optional BackendSpecProvider for layer implementations (default: auto-selected) + + Forward args: + img: Packed image latents [S_img, B, C*4] from VAE + txt: Text embeddings [S_txt, B, D_txt] from T5-XXL + y: CLIP pooled embeddings [B, D_pool] + timesteps: Diffusion timesteps [B] in [0, 1] + img_ids, txt_ids: Position IDs [B, S, 3] + guidance: Optional guidance scale [B] + + Returns: + Predicted velocity [S_img, B, C*4] in packed format + + torch.compile: + Set enable_torch_compile=True in config. Compilation is applied AFTER + FSDP/DDP wrapping by the trainer -- do NOT call compile_model() manually. + + Reference: + - "Flux: A Scalable Diffusion Model for High-Resolution Image Synthesis" + - Megatron-Core TransformerBlock + """ + + def __init__( + self, + config: FluxConfig, + encoder_configs: Optional[Dict[str, Any]] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + backend: Optional["BackendSpecProvider"] = None, + ): + super().__init__(config, pg_collection=pg_collection, encoder_configs=encoder_configs) + + self.out_channels = config.in_channels + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.patch_size = config.patch_size + self.in_channels = config.in_channels + self.guidance_embed = config.guidance_embed + + # Position embedding (3D RoPE for image patches) + # axes_dim=[16, 56, 56] for 1024x1024 images with 64 channels + self.pos_embed = EmbedND( + dim=self.hidden_size, + theta=config.theta, + axes_dim=list(config.axes_dim), + ) + + # Input embeddings + # img_embed accepts packed format with in_channels already set to C*4 in config + # Note: config.in_channels = 64 = 16 (VAE channels) * 4 (packing factor) + self.img_embed = nn.Linear(config.in_channels, self.hidden_size) + self.txt_embed = nn.Linear(config.context_dim, self.hidden_size) + + # Conditioning embeddings + self.timestep_embedding = TimeStepEmbedder( + embedding_dim=config.model_channels, + hidden_dim=self.hidden_size, + ) + self.vector_embedding = MLPEmbedder( + in_dim=config.vec_in_dim, + hidden_dim=self.hidden_size, + ) + + # Optional guidance embedding for classifier-free guidance + if config.guidance_embed: + self.guidance_embedding = MLPEmbedder( + in_dim=config.model_channels, + hidden_dim=self.hidden_size, + ) + else: + self.guidance_embedding = nn.Identity() + + # Create unified DiffusionTransformerBlock with heterogeneous layers + # Layers 0-18: MMDiTLayer (joint processing) + # Layers 19-56: FluxSingleTransformerBlock (single processing) + self.transformer = DiffusionTransformerBlock( + config=config, + spec=get_flux_layer_spec(config, backend=backend), + post_layer_norm=False, # Flux uses AdaLN, not standard final layernorm + pre_process=True, + post_process=True, + ) + + # Output layers + self.norm_out = AdaLNContinuous( + config=config, + conditioning_embedding_dim=self.hidden_size, + ) + self.proj_out = nn.Linear( + self.hidden_size, + self.patch_size * self.patch_size * self.out_channels, + bias=True, + ) + + # NOTE: torch.compile is NOT applied here. It must be applied AFTER + # distributed wrapping (FSDP/DDP). The training framework will call + # compile_model() at the appropriate time. + + # Stack runner references (replaced with compiled versions by compile_model) + self._double_block_stack_runner = self._run_double_block_stack + self._single_block_stack_runner = self._run_single_block_stack + self._output_head_runner = self._run_output_head + self._full_dit_runner = self._run_full_dit + + # Replace TE QK RMSNorm BEFORE DDP wrapping so that DDP's forward + # pre-hook registration sees the final module tree. + if config.torch_compile_replace_qk_rmsnorm and config.transformer_impl == "transformer_engine": + try: + from primus.core.utils.module_utils import log_rank_0 as _log + except Exception: + _log = print + self._replace_qk_rmsnorm(_log) + + self.init_weights() + + def init_weights(self): + """ + Custom weight initialization matching NVIDIA MLPerf Training v5.1. + + - img_embed, txt_embed: Xavier uniform + - timestep_embedding, vector_embedding: Normal(std=0.02) for MLP layers + - Per-block AdaLN modulations: re-zeroed after construction (construction + now uses normal_ init_method to match NeMo's RNG sequence) + - norm_out.adaLN_modulation: zero-init (last linear layer) + - norm_out.norm: reset to default + - proj_out: zero-init weight and bias + """ + from primus.core.utils.module_utils import log_rank_0 + + # Embedders: Xavier uniform + nn.init.xavier_uniform_(self.img_embed.weight) + nn.init.constant_(self.img_embed.bias, 0) + nn.init.xavier_uniform_(self.txt_embed.weight) + nn.init.constant_(self.txt_embed.bias, 0) + + # Timestep embedder MLP: Normal(std=0.02) + self._init_mlpembedder(self.timestep_embedding.time_embedding) + self._init_mlpembedder(self.vector_embedding) + + if self.guidance_embed and not isinstance(self.guidance_embedding, nn.Identity): + self._init_mlpembedder(self.guidance_embedding) + + # Per-block AdaLN: zero modulation weights and reset LayerNorms. + # Order matches NeMo: single blocks first, then double blocks. + # These are deterministic ops (no RNG consumption). + for layer in self.transformer.layers[self.config.num_joint_layers :]: + nn.init.constant_(layer.adaln.adaLN_modulation[-1].weight, 0) + nn.init.constant_(layer.adaln.adaLN_modulation[-1].bias, 0) + layer.adaln.ln.reset_parameters() + + for layer in self.transformer.layers[: self.config.num_joint_layers]: + nn.init.constant_(layer.adaln.adaLN_modulation[-1].weight, 0) + nn.init.constant_(layer.adaln.adaLN_modulation[-1].bias, 0) + layer.adaln.ln.reset_parameters() + layer.adaln.ln2.reset_parameters() + nn.init.constant_(layer.adaln_context.adaLN_modulation[-1].weight, 0) + nn.init.constant_(layer.adaln_context.adaLN_modulation[-1].bias, 0) + layer.adaln_context.ln.reset_parameters() + layer.adaln_context.ln2.reset_parameters() + + # Output layers: zero-init for stable training start + nn.init.constant_(self.proj_out.weight, 0) + nn.init.constant_(self.proj_out.bias, 0) + + # norm_out (AdaLNContinuous): zero-init modulation, reset norm + nn.init.constant_(self.norm_out.adaLN_modulation[-1].weight, 0) + nn.init.constant_(self.norm_out.adaLN_modulation[-1].bias, 0) + self.norm_out.norm.reset_parameters() + + log_rank_0("Applied custom weight initialization (MLPerf v5.1 aligned, NeMo RNG-matched)") + + @staticmethod + def _init_mlpembedder(module, init_std: float = 0.02): + """Initialize MLPEmbedder with Normal(std) for both linear layers.""" + nn.init.normal_(module.in_layer.weight, std=init_std) + nn.init.constant_(module.in_layer.bias, 0) + nn.init.normal_(module.out_layer.weight, std=init_std) + nn.init.constant_(module.out_layer.bias, 0) + + def compile_model(self): + """ + Apply torch.compile to the model using the configured strategy. + + CRITICAL: This must be called AFTER distributed wrapping (FSDP/DDP) is complete + AND after ddp_config is set. Do NOT call this in __init__. The training + framework will call this at the appropriate time. + + Strategies: + per_block: compile each transformer layer individually (default, recommended) + whole_model: compile entire forward (incompatible with overlap_param_gather) + double_stack: compile the double (joint) block loop + single_stack: compile the single block loop + stack: compile both double and single block loops + full_dit: compile entire DiT as one region + """ + import os + + if not self.config.enable_torch_compile: + return + + try: + from primus.core.utils.module_utils import log_rank_0 + + log = log_rank_0 + except Exception: + log = print + + strategy = self.config.torch_compile_strategy + valid_strategies = {"whole_model", "per_block", "double_stack", "single_stack", "stack", "full_dit"} + if strategy not in valid_strategies: + raise ValueError( + f"Invalid torch_compile_strategy='{strategy}'. " f"Must be one of: {sorted(valid_strategies)}" + ) + + is_te = self.config.transformer_impl == "transformer_engine" + + if not is_te and strategy not in ("whole_model", "per_block"): + log( + f"WARNING: torch_compile_strategy='{strategy}' is designed for TE spec. " + f"Local spec modules are already torch.compile-friendly; " + f"falling back to 'per_block' strategy." + ) + strategy = "per_block" + + if ( + strategy not in ("whole_model", "per_block") + and getattr(self.config, "recompute_granularity", None) == "full" + ): + raise ValueError( + f"Activation checkpointing (recompute_granularity='full') is incompatible " + f"with torch_compile_strategy='{strategy}'. Stack runners bypass the " + f"checkpointed forward path." + ) + + if strategy == "whole_model": + from megatron.training import get_args + + args = get_args() + if getattr(args, "overlap_param_gather", False): + raise ValueError( + "whole_model compile strategy is incompatible with overlap_param_gather. " + "DDP hooks are traced inside the compiled graph, causing ~20% convergence " + "degradation. Use 'per_block' strategy instead, or disable overlap_param_gather." + ) + + if strategy == "per_block" and getattr(self.config, "enable_cuda_graph", False): + log( + "WARNING: per_block compile strategy may conflict with per-layer " + "CUDA graph __call__ overrides on MMDiTLayer/FluxSingleTransformerBlock." + ) + + compile_kwargs = { + "backend": self.config.torch_compile_backend, + "mode": self.config.torch_compile_mode, + "fullgraph": self.config.torch_compile_fullgraph, + } + + if is_te and self.config.torch_compile_disable_inductor_cudagraphs: + torch._inductor.config.triton.cudagraphs = False + torch._inductor.config.triton.cudagraph_trees = False + os.environ["TORCHINDUCTOR_CUDAGRAPHS"] = "0" + log(" Disabled Inductor CUDA graphs (TE FP8 compatibility)") + + if self.config.torch_compile_emulate_precision_casts: + torch._inductor.config.emulate_precision_casts = True + log( + " Enabled emulate_precision_casts " + "(preserves eager BF16 precision semantics in fused Triton kernels)" + ) + + if not self.config.torch_compile_fused_ln_modulate: + for m in self.modules(): + if hasattr(m, "use_fused_ln_modulate") and not getattr(m, "_adaln_plain_ops", False): + m.use_fused_ln_modulate = False + log(" Disabled fused LN+modulate Triton kernels (using separate opaque ops)") + + if strategy == "whole_model": + self._compile_whole_model(compile_kwargs, log) + elif strategy == "per_block": + log("=" * 80) + log("Applying PER-BLOCK torch.compile on Flux...") + for i, layer in enumerate(self.transformer.layers): + layer.forward = torch.compile(layer.forward, **compile_kwargs) + log(f" Compiled {len(self.transformer.layers)} layers individually") + log("=" * 80) + elif strategy == "double_stack": + log("=" * 80) + log("Applying DOUBLE-STACK torch.compile on Flux...") + self._double_block_stack_runner = torch.compile(self._run_double_block_stack, **compile_kwargs) + log(f" Compiled double block stack ({self.config.num_joint_layers} layers)") + log("=" * 80) + elif strategy == "single_stack": + log("=" * 80) + log("Applying SINGLE-STACK torch.compile on Flux...") + self._single_block_stack_runner = torch.compile(self._run_single_block_stack, **compile_kwargs) + log(f" Compiled single block stack ({self.config.num_single_layers} layers)") + log("=" * 80) + elif strategy == "stack": + log("=" * 80) + log("Applying STACK torch.compile on Flux (double + single)...") + self._double_block_stack_runner = torch.compile(self._run_double_block_stack, **compile_kwargs) + self._single_block_stack_runner = torch.compile(self._run_single_block_stack, **compile_kwargs) + log( + f" Compiled double ({self.config.num_joint_layers}) " + f"+ single ({self.config.num_single_layers}) block stacks" + ) + log("=" * 80) + elif strategy == "full_dit": + log("=" * 80) + log("Applying FULL-DIT torch.compile on Flux...") + self._full_dit_runner = torch.compile(self._run_full_dit, **compile_kwargs) + log(" Compiled entire DiT as one region") + log("=" * 80) + + def _compile_whole_model(self, compile_kwargs, log): + """Apply a single torch.compile on self.forward (whole-model strategy).""" + import torch + + log("=" * 80) + log("Applying WHOLE-MODEL torch.compile on Flux...") + log(f" Backend: {compile_kwargs['backend']}") + log(f" Mode: {compile_kwargs['mode']}") + log(f" Fullgraph: {compile_kwargs['fullgraph']}") + log("=" * 80) + self.forward = torch.compile(self.forward, **compile_kwargs) + log(" Done (actual compilation deferred to first forward pass)") + log("=" * 80) + + def _replace_qk_rmsnorm(self, log): + """Replace TE QK RMSNorm modules with compile-friendly _TorchRMSNorm. + + Only meaningful for transformer_impl="transformer_engine" where QK norms + are TENorm instances (opaque to torch.compile). Local spec already uses + torch.nn.RMSNorm which is fully traceable. + """ + count = 0 + for layer in self.transformer.layers: + attn = getattr(layer, "self_attention", None) + if attn is None: + continue + for attr in ("q_layernorm", "k_layernorm", "added_q_layernorm", "added_k_layernorm"): + old_norm = getattr(attn, attr, None) + if old_norm is not None and not isinstance(old_norm, _TorchRMSNorm): + setattr(attn, attr, _build_local_qk_rmsnorm(old_norm)) + count += 1 + log(f" Replaced {count} TE QK RMSNorm modules with compile-friendly _TorchRMSNorm") + + # ------------------------------------------------------------------ + # Embedding computation: compilable helper that converts raw inputs + # into hidden states, conditioning vectors, and RoPE frequencies. + # ------------------------------------------------------------------ + + def _compute_embeddings(self, img, txt, timesteps, img_ids, txt_ids, guidance, y): + """Compute all embeddings and RoPE in a single compilable region.""" + hidden_states = self.img_embed(img) + encoder_hidden_states = self.txt_embed(txt) + + if hidden_states.dim() == 3 and hidden_states.shape[0] < hidden_states.shape[1]: + hidden_states = hidden_states.transpose(0, 1) + if ( + encoder_hidden_states.dim() == 3 + and encoder_hidden_states.shape[0] < encoder_hidden_states.shape[1] + ): + encoder_hidden_states = encoder_hidden_states.transpose(0, 1) + + txt_seq_len = encoder_hidden_states.shape[0] + + timesteps = timesteps.to(hidden_states.dtype) * 1000.0 + vec_emb = self.timestep_embedding(timesteps) + + if guidance is not None: + guidance_emb = self.guidance_embedding(self.timestep_embedding.time_proj(guidance * 1000.0)) + vec_emb = vec_emb + guidance_emb + + vec_emb = vec_emb + self.vector_embedding(y) + + ids = torch.cat((txt_ids, img_ids), dim=1) + rotary_pos_emb = self.pos_embed(ids) + + return hidden_states, encoder_hidden_states, vec_emb, rotary_pos_emb, txt_seq_len + + # ------------------------------------------------------------------ + # Stack runners: self-contained compilable callables that iterate + # transformer layers directly (bypassing DiffusionTransformerBlock). + # ------------------------------------------------------------------ + + def _run_double_block_stack(self, hidden_states, encoder_hidden_states, rotary_pos_emb, vec_emb): + """Run all double (joint MMDiT) blocks as one compilable region.""" + from megatron.core.utils import make_viewless_tensor + + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True + ) + for layer in self.transformer.layers[: self.config.num_joint_layers]: + hidden_states, encoder_hidden_states = layer( + hidden_states=hidden_states, + attention_mask=None, + context=encoder_hidden_states, + context_mask=None, + rotary_pos_emb=rotary_pos_emb, + timestep_emb=vec_emb, + ) + return hidden_states, encoder_hidden_states + + def _run_single_block_stack(self, hidden_states, rotary_pos_emb, vec_emb): + """Run all single blocks as one compilable region. + + The caller must concatenate [context, hidden] before entering this runner. + All single blocks receive context=None (concat already done). + """ + for layer in self.transformer.layers[self.config.num_joint_layers :]: + hidden_states, _ = layer( + hidden_states=hidden_states, + attention_mask=None, + context=None, + context_mask=None, + rotary_pos_emb=rotary_pos_emb, + timestep_emb=vec_emb, + ) + return hidden_states + + def _run_output_head(self, hidden_states, txt_seq_len, vec_emb): + """Run the output head (slice + norm_out + proj_out).""" + hidden_states = hidden_states[txt_seq_len:, ...] + hidden_states = self.norm_out(hidden_states, vec_emb) + return self.proj_out(hidden_states) + + def _run_full_dit(self, hidden_states, encoder_hidden_states, rotary_pos_emb, vec_emb, txt_seq_len): + """Run the entire DiT (double + cat + single + output) as one compiled region.""" + hidden_states, encoder_hidden_states = self._run_double_block_stack( + hidden_states, encoder_hidden_states, rotary_pos_emb, vec_emb + ) + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=0) + hidden_states = self._run_single_block_stack(hidden_states, rotary_pos_emb, vec_emb) + if self.transformer.final_layernorm is not None: + hidden_states = self.transformer.final_layernorm(hidden_states) + return self._run_output_head(hidden_states, txt_seq_len, vec_emb) + + def set_input_tensor(self, input_tensor): + """ + Set the input tensor on the underlying TransformerBlock. + + Required by the Megatron model interface (the scheduler calls this on + every model). Pipeline parallelism is not supported for diffusion + models (rejected at config construction; PP > 1 raises), so only the + single-stage case is handled here. + + Args: + input_tensor: Union[Tensor, List[Tensor], None] + """ + if input_tensor is None: + return + + # Handle list or single tensor + if not isinstance(input_tensor, list): + input_tensor = [input_tensor] + + # Pass image hidden states to TransformerBlock + self.transformer.set_input_tensor(input_tensor[0]) + + def sharded_state_dict( + self, + prefix: str = "", + sharded_offsets: tuple = (), + metadata: Optional[dict] = None, + ) -> Dict[str, Any]: + """ + Generate sharded state dictionary for distributed checkpointing. + + With TransformerBlock, this delegates to the transformer which handles + all layer numbering automatically for heterogeneous layers. + + Flux is architecturally heterogeneous: joint MMDiT blocks carry params + (e.g. ``self_attention.added_linear_qkv.*``, ``context_mlp.*``) that the + single DiT blocks intentionally omit (FluxSingleAttention rebuilds + ``linear_proj`` with ``bias=False``). Megatron's heterogeneous, per-layer + indexed sharded keys (``transformer.layers..*``) are required here: + the homogeneous layer-stacked path would leave unclaimed slots for the + params absent in single blocks and raise a CheckpointingException at + save time. TransformerBlock.sharded_state_dict provides this + automatically, so no config toggle is needed. + + Args: + prefix: Prefix for state dict keys (e.g., 'module.') + sharded_offsets: Pipeline parallel offsets + metadata: Optional metadata for checkpoint conversion + + Returns: + Dictionary mapping state dict keys to ShardedTensor objects + """ + sharded_state_dict = {} + + # Delegate transformer layers to TransformerBlock + # Handles heterogeneous layers automatically + sharded_state_dict.update( + self.transformer.sharded_state_dict(f"{prefix}transformer.", sharded_offsets, metadata) + ) + + # Handle other submodules (embeddings, projections, norms) + for name, module in self.named_children(): + if module is not self.transformer: + sharded_state_dict.update( + sharded_state_dict_default(module, f"{prefix}{name}.", sharded_offsets, metadata) + ) + + return sharded_state_dict + + def get_fp8_context(self): + """ + Return FP8 context manager for FP8 training. + + When transformer_impl="local", FP8 is handled per-module at init time + (inside Float8ColumnParallelLinear / Float8RowParallelLinear), so no + global FP8 context manager is needed. This avoids mutable global state + that would cause torch.compile graph breaks. + + Returns: + Context manager for FP8 operations (or nullcontext if FP8 disabled + or using local spec) + """ + from contextlib import nullcontext + + if self.config.transformer_impl == "local": + return nullcontext() + + from megatron.core.fp8_utils import get_fp8_context as get_fp8_context_util + + return get_fp8_context_util(self.config) + + def forward( + self, + img: Tensor, + txt: Tensor, + y: Tensor, + timesteps: Tensor, + img_ids: Tensor, + txt_ids: Tensor, + guidance: Optional[Tensor] = None, + ) -> Tensor: + """ + Forward pass through Flux model. + + Args: + img: Image latents [S_img, B, C*4] or [B, S_img, C*4] - PACKED format from prepare_flux_latents + where S_img = H*W/4 (2x2 patches grouped) + txt: Text embeddings [S_txt, B, D_txt] or [B, S_txt, D_txt] from T5-XXL + y: CLIP pooled embeddings [B, D_pool] from CLIP-L + timesteps: Diffusion timesteps [B] in range [0, 1] + img_ids: Position IDs for image patches [B, S_img, 3] + txt_ids: Position IDs for text tokens [B, S_txt, 3] + guidance: Optional guidance scale [B] for classifier-free guidance + + Returns: + Predicted velocity [S_img, B, C*4] or [B, S_img, C*4] in packed format (matches input format) + + Note: + The model accepts pre-packed + latents from the wrapper/training code. The packing groups 2x2 spatial patches, + reducing spatial dimensions by 4x and increasing channels by 4x. + """ + hidden_states, encoder_hidden_states, vec_emb, rotary_pos_emb, txt_seq_len = self._compute_embeddings( + img, txt, timesteps, img_ids, txt_ids, guidance, y + ) + + strategy = self.config.torch_compile_strategy + + if strategy in ("whole_model", "per_block"): + # Original path: delegate to DiffusionTransformerBlock + with self.get_fp8_context(): + hidden_states = self.transformer( + hidden_states=hidden_states, + attention_mask=None, + context=encoder_hidden_states, + context_mask=None, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=None, + rotary_pos_sin=None, + attention_bias=None, + timestep_emb=vec_emb, + packed_seq_params=None, + ) + + hidden_states = hidden_states[txt_seq_len:, ...] + hidden_states = self.norm_out(hidden_states, vec_emb) + output = self.proj_out(hidden_states) + + elif strategy == "full_dit": + # One large compiled region covering double + cat + single + output + with self.get_fp8_context(): + output = self._full_dit_runner( + hidden_states, encoder_hidden_states, rotary_pos_emb, vec_emb, txt_seq_len + ) + + else: + # stack / double_stack / single_stack: compiled stack runners + # with eager transitions between them + with self.get_fp8_context(): + hidden_states, encoder_hidden_states = self._double_block_stack_runner( + hidden_states, encoder_hidden_states, rotary_pos_emb, vec_emb + ) + + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=0) + + with self.get_fp8_context(): + hidden_states = self._single_block_stack_runner(hidden_states, rotary_pos_emb, vec_emb) + + if self.transformer.final_layernorm is not None: + hidden_states = self.transformer.final_layernorm(hidden_states) + + output = self._output_head_runner(hidden_states, txt_seq_len, vec_emb) + + return output + + def load_checkpoint( + self, + checkpoint_path: Union[str, Path], + convert_from_hf: bool = False, + save_converted_to: Optional[Union[str, Path]] = None, + strict: bool = False, + ) -> Tuple[List[str], List[str]]: + """ + Load checkpoint into Flux model. + + Supports both Primus native checkpoints and HuggingFace checkpoints + with automatic conversion. + + Note: HuggingFace checkpoint conversion is handled by the checkpoint_converter module. + + Args: + checkpoint_path: Path to checkpoint file or directory + convert_from_hf: If True, convert from HuggingFace format + save_converted_to: Optional path to save converted checkpoint + strict: Whether to use strict state dict loading + + Returns: + Tuple of (missing_keys, unexpected_keys) + + Example: + >>> config = FluxConfig.flux_12b() + >>> model = Flux(config=config) + >>> # Load native Primus checkpoint + >>> missing, unexpected = model.load_checkpoint("primus_flux_12b.safetensors") + """ + from safetensors.torch import load_file as load_safetensors + + from primus.core.utils.module_utils import ( + error_rank_0, + log_rank_0, + warning_rank_0, + ) + + if convert_from_hf: + try: + from .checkpoint_converter import convert_hf_checkpoint + + log_rank_0(f"Converting HuggingFace checkpoint: {checkpoint_path}") + state_dict = convert_hf_checkpoint( + checkpoint_path, + flux_config=self.config, + save_to=save_converted_to, + ) + except ImportError: + raise NotImplementedError( + "HuggingFace checkpoint conversion requires the checkpoint_converter module. " + "Please use a pre-converted Primus checkpoint." + ) + else: + # Load native Primus checkpoint + log_rank_0(f"Loading Primus checkpoint: {checkpoint_path}") + checkpoint_path = Path(checkpoint_path) + + if checkpoint_path.is_dir(): + # Load all .safetensors files + safetensor_files = list(checkpoint_path.glob("*.safetensors")) + if not safetensor_files: + raise FileNotFoundError(f"No .safetensors files in {checkpoint_path}") + + state_dict = {} + for file in safetensor_files: + state_dict.update(load_safetensors(str(file))) + else: + state_dict = load_safetensors(str(checkpoint_path)) + + # Load state dict into model + missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=strict) + + # Filter out _extra_state keys (Megatron-specific metadata, expected to be missing) + missing_keys = [k for k in missing_keys if not k.endswith("_extra_state")] + + # Check for critical missing keys + critical_patterns = ["timestep_embedding", "img_embed", "txt_embed", "vector_embedding"] + critical_missing = [k for k in missing_keys if any(pattern in k for pattern in critical_patterns)] + + if critical_missing: + variant = "flux_12b" if self.config.num_joint_layers == 19 else "flux_535m" + err_lines = [ + "\n" + "=" * 80, + "CRITICAL ERROR: Key model weights are missing!", + "=" * 80, + f"Missing {len(critical_missing)} critical keys:", + ] + err_lines += [f" - {key}" for key in critical_missing[:10]] + if len(critical_missing) > 10: + err_lines.append(f" ... and {len(critical_missing) - 10} more") + err_lines += [ + "\nThis usually means:", + " 1. Checkpoint format mismatch (HuggingFace vs Primus)", + " 2. Checkpoint was converted with an older version", + " 3. Checkpoint file is corrupted or incomplete", + "\nTo fix, reconvert the checkpoint:", + " python tools/checkpoint_conversion/convert_flux_hf_to_primus.py \\", + " --input black-forest-labs/FLUX.1-dev/transformer \\", + f" --output {checkpoint_path} \\", + f" --variant {variant}", + "=" * 80, + ] + error_rank_0("\n".join(err_lines)) + + raise RuntimeError( + f"Critical model weights missing ({len(critical_missing)} keys). " + f"Checkpoint may be corrupted or incompatible. See details above." + ) + + if missing_keys: + warning_rank_0(f"Missing keys ({len(missing_keys)}): {missing_keys[:5]}...") + if unexpected_keys: + warning_rank_0(f"Unexpected keys ({len(unexpected_keys)}): {unexpected_keys[:5]}...") + + log_rank_0("Checkpoint loaded successfully!") + return missing_keys, unexpected_keys diff --git a/primus/backends/megatron/core/models/diffusion/flux/utils.py b/primus/backends/megatron/core/models/diffusion/flux/utils.py new file mode 100644 index 000000000..604b7d810 --- /dev/null +++ b/primus/backends/megatron/core/models/diffusion/flux/utils.py @@ -0,0 +1,197 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Core utilities shared between Flux training and inference. + +This module provides fundamental operations used by both training and inference: +- Latent packing/unpacking (2x2 spatial grouping) +- Position ID generation (3D RoPE for Flux transformer) + +Design: Pure functions with no side effects, fully type-annotated. +""" + + +import torch +from torch import Tensor + + +def pack_latents(latents: Tensor) -> Tensor: + """ + Pack latents from (B, C, H, W) to (B, H*W/4, C*4) format. + + Groups 2x2 spatial patches into sequence tokens for transformer processing. + Used by both training and inference pipelines. + + Args: + latents: Input tensor of shape (B, C, H, W) + + Returns: + Packed tensor of shape (B, H*W/4, C*4) + + Example: + >>> latents = torch.randn(2, 16, 128, 128) + >>> packed = pack_latents(latents) + >>> packed.shape + torch.Size([2, 4096, 64]) # 128*128/4=4096, 16*4=64 + """ + batch_size, num_channels, height, width = latents.shape + + # Reshape to group 2x2 patches: (B, C, H, W) -> (B, C, H//2, 2, W//2, 2) + latents = latents.view(batch_size, num_channels, height // 2, 2, width // 2, 2) + + # Permute to bring patch dimensions together: (B, H//2, W//2, C, 2, 2) + latents = latents.permute(0, 2, 4, 1, 3, 5) + + # Flatten patches: (B, H//2*W//2, C*4) + latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels * 4) + + return latents + + +def unpack_latents( + latents: Tensor, + height: int, + width: int, + vae_scale_factor: int = 1, +) -> Tensor: + """ + Unpack latents from (B, N, C*4) to (B, C, H, W) format. + + Reverses pack_latents operation. Supports optional VAE scaling for inference. + + Args: + latents: Packed tensor of shape (B, N, C*4) + height: Target spatial height + width: Target spatial width + vae_scale_factor: Downsampling factor (default: 1 for training, 8 for inference) + + Returns: + Unpacked tensor of shape (B, C, H, W) + + Example: + >>> packed = torch.randn(2, 4096, 64) + >>> unpacked = unpack_latents(packed, 1024, 1024, vae_scale_factor=8) + >>> unpacked.shape + torch.Size([2, 16, 128, 128]) # 1024/8 = 128 (VAE downscaling) + """ + batch_size, num_patches, channels = latents.shape + + # Apply VAE downsampling if specified (inference path) + if vae_scale_factor > 1: + height = height // vae_scale_factor + width = width // vae_scale_factor + + # For packed latents: height and width are the dimensions BEFORE packing + # So we need to use height//2 and width//2 for the reshaped view + h_packed = height // 2 + w_packed = width // 2 + + # Reshape to restore patch structure: (B, H//2, W//2, C//4, 2, 2) + latents = latents.view(batch_size, h_packed, w_packed, channels // 4, 2, 2) + + # Permute to restore spatial dimensions: (B, C//4, H//2, 2, W//2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + + # Flatten spatial dimensions: (B, C//4, H, W) + latents = latents.reshape(batch_size, channels // 4, height, width) + + return latents + + +def generate_image_position_ids( + batch_size: int, + height: int, + width: int, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> Tensor: + """ + Generate 3D position IDs for image patches (RoPE). + + Creates position encodings for Flux's 3D RoPE where: + - Dimension 0: Always 0 (reserved for future temporal/video models) + - Dimension 1: Row index (y-coordinate, 0 to height//2-1 after packing) + - Dimension 2: Column index (x-coordinate, 0 to width//2-1 after packing) + + Follows NeMo conventions for Flux position encoding. For video models, dimension 0 + would encode frame indices instead of being 0. + + NOTE FOR ROPE FUSION: When using fused RoPE kernels (apply_rope_fusion=True), + this function should be called with batch_size=1 to satisfy Transformer Engine's + dimension constraints (freqs tensor must have shape [S, 1, 1, D]). The resulting + [1, H*W/4, 3] position IDs will broadcast across the actual batch dimension + during attention computation. This requires all images in the batch to have the + same resolution (same height and width). + + Args: + batch_size: Batch size (use 1 for RoPE fusion, actual batch size otherwise) + height: Latent height in unpacked format (before 2x2 packing) + width: Latent width in unpacked format (before 2x2 packing) + device: Target device + dtype: Target dtype + + Returns: + Position IDs tensor of shape (batch_size, H*W/4, 3) + + Example: + >>> ids = generate_image_position_ids(2, 128, 128, torch.device('cpu')) + >>> ids.shape + torch.Size([2, 4096, 3]) # 128*128/4 = 4096 packed positions + """ + # Generate for packed latents (2x2 grouping) + h_packed = height // 2 + w_packed = width // 2 + + # Create position grid + img_ids = torch.zeros(h_packed, w_packed, 3, device=device, dtype=dtype) + + # Fill row dimension (dim 1) + img_ids[..., 1] = torch.arange(h_packed, device=device, dtype=dtype)[:, None] + + # Fill column dimension (dim 2) + img_ids[..., 2] = torch.arange(w_packed, device=device, dtype=dtype)[None, :] + + # Flatten spatial dimensions and expand for batch + img_ids = img_ids.reshape(h_packed * w_packed, 3) + img_ids = img_ids.unsqueeze(0).expand(batch_size, -1, -1) + + return img_ids + + +def generate_text_position_ids( + batch_size: int, + seq_len: int, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> Tensor: + """ + Generate position IDs for text tokens. + + For Flux, text position IDs are zeros (no explicit positional encoding). + + Args: + batch_size: Batch size + seq_len: Text sequence length + device: Target device + dtype: Target dtype + + Returns: + Position IDs tensor of shape (B, seq_len, 3) filled with zeros + + Example: + >>> ids = generate_text_position_ids(2, 512, torch.device('cpu')) + >>> ids.shape + torch.Size([2, 512, 3]) + >>> ids.sum().item() + 0.0 + """ + return torch.zeros(batch_size, seq_len, 3, device=device, dtype=dtype) + + +__all__ = [ + "pack_latents", + "unpack_latents", + "generate_image_position_ids", + "generate_text_position_ids", +] diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_chimera_init.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_chimera_init.py new file mode 100644 index 000000000..37e41725a --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_chimera_init.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for NeMo chimera initialization behavior. + +Validates that contaminating the default CUDA generator with per-DP-rank seeds +causes non-parallel layers (img_embed, txt_embed, MLPEmbedder) to diverge +across ranks, while ColumnParallelLinear layers (using model-parallel tracker) +remain identical. +""" + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from tests.utils import PrimusUT + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +class TestFluxChimeraInit(PrimusUT): + """Tests that chimera init creates intended per-rank diversity.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state for model tests.""" + + def _build_model_with_rank_seed(self, base_seed, dp_rank): + """Build a Flux model simulating chimera init for a given DP rank.""" + from megatron.core.tensor_parallel import random as tp_random + + per_rank_seed = base_seed + 100 * dp_rank + torch.manual_seed(per_rank_seed) + torch.cuda.manual_seed(per_rank_seed) + # Model-parallel tracker uses the SAME seed across all DP ranks + tp_random.model_parallel_cuda_manual_seed(42) + + config = FluxConfig.flux_535m() + return Flux(config) + + def test_img_embed_differs_across_ranks(self): + """img_embed (nn.Linear) should differ when default generator seed differs.""" + model_rank0 = self._build_model_with_rank_seed(1234, dp_rank=0) + model_rank1 = self._build_model_with_rank_seed(1234, dp_rank=1) + + assert not torch.equal( + model_rank0.img_embed.weight, model_rank1.img_embed.weight + ), "img_embed.weight should differ between DP ranks" + + def test_txt_embed_differs_across_ranks(self): + """txt_embed (nn.Linear) should differ when default generator seed differs.""" + model_rank0 = self._build_model_with_rank_seed(1234, dp_rank=0) + model_rank1 = self._build_model_with_rank_seed(1234, dp_rank=1) + + assert not torch.equal( + model_rank0.txt_embed.weight, model_rank1.txt_embed.weight + ), "txt_embed.weight should differ between DP ranks" + + def test_timestep_embedding_differs_across_ranks(self): + """MLPEmbedder layers (nn.Linear) should differ across DP ranks.""" + model_rank0 = self._build_model_with_rank_seed(1234, dp_rank=0) + model_rank1 = self._build_model_with_rank_seed(1234, dp_rank=1) + + w0 = model_rank0.timestep_embedding.time_embedding.in_layer.weight + w1 = model_rank1.timestep_embedding.time_embedding.in_layer.weight + assert not torch.equal(w0, w1), "timestep_embedding in_layer should differ between DP ranks" + + def test_adaln_zero_init_identical_regardless_of_seed(self): + """AdaLN zero-init is deterministic — identical on both ranks.""" + model_rank0 = self._build_model_with_rank_seed(1234, dp_rank=0) + model_rank1 = self._build_model_with_rank_seed(1234, dp_rank=1) + + for i, (layer0, layer1) in enumerate( + zip(model_rank0.transformer.layers, model_rank1.transformer.layers) + ): + w0 = layer0.adaln.adaLN_modulation[-1].weight + w1 = layer1.adaln.adaLN_modulation[-1].weight + assert torch.equal(w0, w1), f"Layer {i} adaLN weight should be identical (zero)" + assert torch.all(w0 == 0), f"Layer {i} adaLN weight should be zero" diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py new file mode 100644 index 000000000..2a4348832 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_config.py @@ -0,0 +1,94 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for Flux and base diffusion configurations. + +Tests FluxConfig and BaseDiffusionConfig validation, preset configurations. +""" + +import pytest + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.common import BaseDiffusionConfig +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from tests.utils import PrimusUT + +# ======================================================================== +# Base Diffusion Configuration Tests +# ======================================================================== + + +class TestBaseDiffusionConfig(PrimusUT): + """Tests for BaseDiffusionConfig.""" + + def test_base_config_validation_invalid_channels(self): + """Test validation catches invalid channel counts.""" + with self.assertRaises(ValueError) as cm: + config = BaseDiffusionConfig( + in_channels=0, + num_attention_heads=8, + num_layers=1, + ) + config.validate() + self.assertIn("in_channels must be positive", str(cm.exception)) + + +class TestFluxConfig(PrimusUT): + """Tests for FluxConfig class.""" + + def test_validation_positive_joint_layers(self): + """Test validation fails for non-positive num_joint_layers.""" + with self.assertRaises(ValueError) as cm: + config = FluxConfig(num_joint_layers=0) + config.validate() + self.assertIn("num_joint_layers must be positive", str(cm.exception)) + + def test_validation_positive_single_layers(self): + """Test validation fails for non-positive num_single_layers.""" + with self.assertRaises(ValueError) as cm: + config = FluxConfig(num_single_layers=-1) + config.validate() + self.assertIn("num_single_layers must be positive", str(cm.exception)) + + def test_validation_positive_context_dim(self): + """Test validation fails for non-positive context_dim.""" + with self.assertRaises(ValueError) as cm: + config = FluxConfig(context_dim=0) + config.validate() + self.assertIn("context_dim must be positive", str(cm.exception)) + + def test_validation_positive_vec_in_dim(self): + """Test validation fails for non-positive vec_in_dim.""" + with self.assertRaises(ValueError) as cm: + config = FluxConfig(vec_in_dim=-768) + config.validate() + self.assertIn("vec_in_dim must be positive", str(cm.exception)) + + def test_validation_positive_theta(self): + """Test validation fails for non-positive theta.""" + with self.assertRaises(ValueError) as cm: + config = FluxConfig(theta=0) + config.validate() + self.assertIn("theta must be positive", str(cm.exception)) + + def test_validation_axes_dim_length(self): + """Test validation fails for axes_dim with wrong length.""" + with self.assertRaises(ValueError) as cm: + config = FluxConfig(axes_dim=(16, 56)) # Only 2 elements + config.validate() + self.assertIn("axes_dim must have 3 elements", str(cm.exception)) + + def test_validation_axes_dim_positive_values(self): + """Test validation fails for non-positive axes_dim values.""" + with self.assertRaises(ValueError) as cm: + config = FluxConfig(axes_dim=(16, 0, 56)) + config.validate() + self.assertIn("All axes_dim values must be positive", str(cm.exception)) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py new file mode 100644 index 000000000..2bd83cef0 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_fp8_context.py @@ -0,0 +1,51 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for FluxModel.get_fp8_context() behavior. + +Verifies that the local spec FP8 path returns nullcontext (no global state), +while the non-local path delegates to Megatron's get_fp8_context utility. +""" + +import contextlib +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.model import Flux + + +class TestFluxFP8Context: + """Tests for Flux.get_fp8_context() without constructing a full model.""" + + def test_get_fp8_context_local_returns_nullcontext(self): + """Local spec should return nullcontext to avoid global FP8 state.""" + mock_model = MagicMock(spec=Flux) + mock_model.config = SimpleNamespace( + transformer_impl="local", + fp8="e4m3", + ) + + ctx = Flux.get_fp8_context(mock_model) + assert isinstance(ctx, contextlib.nullcontext) + + def test_get_fp8_context_non_local_delegates(self): + """Non-local spec should delegate to megatron.core.fp8_utils.get_fp8_context.""" + mock_model = MagicMock(spec=Flux) + mock_model.config = SimpleNamespace( + transformer_impl="transformer_engine", + fp8="e4m3", + ) + + sentinel = contextlib.nullcontext() + with patch( + "megatron.core.fp8_utils.get_fp8_context", + return_value=sentinel, + ) as mock_get: + ctx = Flux.get_fp8_context(mock_model) + mock_get.assert_called_once_with(mock_model.config) + assert ctx is sentinel diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_init_weights.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_init_weights.py new file mode 100644 index 000000000..687350544 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_init_weights.py @@ -0,0 +1,172 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for Flux model weight initialization. + +Validates the NeMo-aligned init_weights() implementation: +- AdaLN modulation layers are zero-initialized +- Embeddings receive Xavier/Normal initialization +- proj_out and norm_out are zero-initialized +- Initialization is deterministic given fixed RNG state +""" + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from tests.utils import PrimusUT + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +class TestFluxInitWeights(PrimusUT): + """Tests for Flux init_weights() correctness.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state for model tests.""" + + def _create_model(self): + config = FluxConfig.flux_535m() + return Flux(config) + + def test_adaln_modulation_zero_init_single_blocks(self): + """Single-block AdaLN modulation last layer must be zero.""" + model = self._create_model() + num_joint = model.config.num_joint_layers + + for i, layer in enumerate(model.transformer.layers[num_joint:]): + w = layer.adaln.adaLN_modulation[-1].weight + b = layer.adaln.adaLN_modulation[-1].bias + assert torch.all(w == 0), f"Single block {i}: adaLN weight not zero" + assert torch.all(b == 0), f"Single block {i}: adaLN bias not zero" + + def test_adaln_modulation_zero_init_joint_blocks(self): + """Joint-block AdaLN modulation last layers must be zero.""" + model = self._create_model() + num_joint = model.config.num_joint_layers + + for i, layer in enumerate(model.transformer.layers[:num_joint]): + w = layer.adaln.adaLN_modulation[-1].weight + b = layer.adaln.adaLN_modulation[-1].bias + assert torch.all(w == 0), f"Joint block {i}: adaln weight not zero" + assert torch.all(b == 0), f"Joint block {i}: adaln bias not zero" + + w_ctx = layer.adaln_context.adaLN_modulation[-1].weight + b_ctx = layer.adaln_context.adaLN_modulation[-1].bias + assert torch.all(w_ctx == 0), f"Joint block {i}: adaln_context weight not zero" + assert torch.all(b_ctx == 0), f"Joint block {i}: adaln_context bias not zero" + + def test_proj_out_zero_init(self): + """proj_out weight and bias must be zero.""" + model = self._create_model() + assert torch.all(model.proj_out.weight == 0) + assert torch.all(model.proj_out.bias == 0) + + def test_norm_out_adaln_zero_init(self): + """norm_out.adaLN_modulation last layer must be zero.""" + model = self._create_model() + w = model.norm_out.adaLN_modulation[-1].weight + b = model.norm_out.adaLN_modulation[-1].bias + assert torch.all(w == 0) + assert torch.all(b == 0) + + def test_img_embed_xavier_distribution(self): + """img_embed should match Xavier-uniform statistics. + + xavier_uniform_ produces U(-bound, bound) with bound = sqrt(6/(in+out)); + std of that distribution is bound/sqrt(3) = sqrt(2/(in+out)). + """ + import math + + model = self._create_model() + w = model.img_embed.weight + out_dim, in_dim = w.shape + expected_std = math.sqrt(2.0 / (in_dim + out_dim)) + + actual_std = w.float().std().item() + actual_mean = w.float().mean().item() + + rel_err = abs(actual_std - expected_std) / expected_std + assert rel_err < 0.15, ( + f"img_embed std={actual_std:.4f} deviates from Xavier expected " + f"{expected_std:.4f} by {rel_err:.1%} (limit 15%)" + ) + assert abs(actual_mean) < 0.005, f"img_embed mean={actual_mean:.4f} should be ~0 for Xavier" + + def test_txt_embed_xavier_distribution(self): + """txt_embed should match Xavier-uniform statistics.""" + import math + + model = self._create_model() + w = model.txt_embed.weight + out_dim, in_dim = w.shape + expected_std = math.sqrt(2.0 / (in_dim + out_dim)) + + actual_std = w.float().std().item() + actual_mean = w.float().mean().item() + + rel_err = abs(actual_std - expected_std) / expected_std + assert rel_err < 0.15, ( + f"txt_embed std={actual_std:.4f} deviates from Xavier expected " + f"{expected_std:.4f} by {rel_err:.1%} (limit 15%)" + ) + assert abs(actual_mean) < 0.005, f"txt_embed mean={actual_mean:.4f} should be ~0 for Xavier" + + def test_timestep_embedding_normal_init(self): + """Timestep embedding MLP should match Normal(std=0.02) statistics.""" + model = self._create_model() + w = model.timestep_embedding.time_embedding.in_layer.weight + expected_std = 0.02 + + actual_std = w.float().std().item() + actual_mean = w.float().mean().item() + + rel_err = abs(actual_std - expected_std) / expected_std + assert rel_err < 0.15, ( + f"timestep_embedding std={actual_std:.4f} deviates from Normal(0,0.02) " + f"expected std {expected_std} by {rel_err:.1%} (limit 15%)" + ) + assert ( + abs(actual_mean) < 0.005 + ), f"timestep_embedding mean={actual_mean:.4f} should be ~0 for Normal init" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +class TestFluxInitWeightsDeterminism(PrimusUT): + """Tests that init_weights produces deterministic results given fixed RNG.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state for model tests.""" + + def test_deterministic_construction(self): + """Two models built with the same RNG state must be bitwise identical.""" + from megatron.core.tensor_parallel import random as tp_random + + config = FluxConfig.flux_535m() + seed = 123 + + # Build model 1 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + tp_random.model_parallel_cuda_manual_seed(seed) + model1 = Flux(config) + + # Build model 2 with identical RNG state + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + tp_random.model_parallel_cuda_manual_seed(seed) + model2 = Flux(config) + + for (name1, p1), (name2, p2) in zip(model1.named_parameters(), model2.named_parameters()): + assert name1 == name2 + assert torch.equal(p1, p2), ( + f"Parameter {name1} differs between constructions " + f"(max diff: {(p1 - p2).abs().max().item():.2e})" + ) diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_layers.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_layers.py new file mode 100644 index 000000000..15a42efd3 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_layers.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for Flux-specific layers. + +Tests EmbedND (3D RoPE) and related functions. +""" + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.layers import EmbedND +from tests.unit_tests.backends.megatron.diffusion.constants import ( + BATCH_SIZE_PAIR, + HIDDEN_DIM_FLUX, + ROPE_THETA_DEFAULT, + SEQ_LEN_TINY, +) +from tests.utils import PrimusUT + + +class TestEmbedND(PrimusUT): + """Tests for EmbedND class.""" + + def test_forward_output_shape(self): + """Test that forward produces correct output shape. + + Output size is determined by sum(axes_dim), not dim parameter. + """ + dim = HIDDEN_DIM_FLUX # This is just stored, not used for output size + axes_dim = [32, 48, 48] # Sum = 128 + batch_size = BATCH_SIZE_PAIR # Paired sample tests + seq_len = SEQ_LEN_TINY # Small sequence for basic tests + num_axes = 3 + + embed_nd = EmbedND(dim=dim, theta=ROPE_THETA_DEFAULT, axes_dim=axes_dim) + ids = torch.randn(batch_size, seq_len, num_axes) + + output = embed_nd(ids) + + # Output size is based on sum(axes_dim) = 128, not dim = 3072 + # Shape: [seq, B, 1, sum(axes_dim)] after permute and reshape + output_dim = sum(axes_dim) + expected_shape = (seq_len, batch_size, 1, output_dim) + assert output.shape == expected_shape, f"Expected shape {expected_shape}, got {output.shape}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_model.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_model.py new file mode 100644 index 000000000..b8d16c6d3 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_model.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Basic unit tests for Flux model. + +Tests the Flux model initialization and basic forward pass. +""" + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from primus.backends.megatron.core.models.diffusion.flux.utils import ( + generate_image_position_ids, + pack_latents, + unpack_latents, +) +from tests.unit_tests.backends.megatron.diffusion.constants import ( + CLIP_L_EMBEDDING_DIM, + T5_XXL_EMBEDDING_DIM, + TEXT_SEQ_LEN_SHORT, + VAE_LATENT_CHANNELS, +) +from tests.utils import PrimusUT + + +class TestFluxModel(PrimusUT): + """Core tests for Flux model initialization and basic operations.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state for model tests.""" + + def test_forward_pass_small(self): + """Test forward pass with small inputs.""" + if not torch.cuda.is_available(): + self.skipTest("CUDA not available") + + config = FluxConfig.flux_535m() + model = Flux(config).cuda() + model.eval() + + batch_size = 2 + height, width = 16, 16 + channels = VAE_LATENT_CHANNELS + txt_seq_len = TEXT_SEQ_LEN_SHORT + + # Prepare inputs + img = torch.randn(batch_size, channels, height, width).cuda() + txt = torch.randn(batch_size, txt_seq_len, T5_XXL_EMBEDDING_DIM).cuda() + y = torch.randn(batch_size, CLIP_L_EMBEDDING_DIM).cuda() + timesteps = torch.rand(batch_size).cuda() + + # Pack latents + packed_img = pack_latents(img) + packed_img = packed_img.transpose(0, 1) + txt_t = txt.transpose(0, 1) + + # Generate position IDs + img_ids = generate_image_position_ids(batch_size, height, width, device="cuda") + txt_ids = torch.zeros(batch_size, txt_seq_len, 3).cuda() + + # Forward pass + with torch.no_grad(): + output = model(packed_img, txt_t, y, timesteps, img_ids, txt_ids) + + # Unpack output + output = output.transpose(0, 1) + output = unpack_latents(output, height, width, vae_scale_factor=1) + + # Check output shape + assert output.shape == img.shape + assert not torch.isnan(output).any() + assert not torch.isinf(output).any() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_utils.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_utils.py new file mode 100644 index 000000000..ffa2ef010 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_utils.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for shared Flux utilities (pack/unpack, position IDs). + +Tests canonical implementations in flux/utils.py that are shared between +training and inference code paths. +""" + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.utils import ( + generate_image_position_ids, + pack_latents, + unpack_latents, +) +from tests.utils import PrimusUT + + +class TestPackUnpackLatents(PrimusUT): + """Tests for latent packing/unpacking operations.""" + + def test_pack_unpack_roundtrip(self): + """Test that pack -> unpack is reversible.""" + batch_size, channels, height, width = 2, 16, 128, 128 + original = torch.randn(batch_size, channels, height, width) + + # Pack and unpack + packed = pack_latents(original) + unpacked = unpack_latents(packed, height, width, vae_scale_factor=1) + + # Should match original + assert unpacked.shape == original.shape + torch.testing.assert_close(unpacked, original) + + +class TestPositionIDs(PrimusUT): + """Tests for position ID generation.""" + + def test_generate_image_position_ids_values(self): + """Test image position IDs have correct structure.""" + batch_size, height, width = 2, 8, 8 # Small for inspection + device = torch.device("cpu") + + img_ids = generate_image_position_ids(batch_size, height, width, device) + + # Dimension 0 should be all zeros + assert (img_ids[:, :, 0] == 0).all() + + # Dimension 1 (row) should range from 0 to height//2-1 + max_row = img_ids[:, :, 1].max() + assert max_row == height // 2 - 1 + + # Dimension 2 (col) should range from 0 to width//2-1 + max_col = img_ids[:, :, 2].max() + assert max_col == width // 2 - 1 + + +class TestEdgeCases(PrimusUT): + """Tests for edge cases and error conditions.""" + + def test_pack_latents_minimum_size(self): + """Test pack_latents with minimum size (2x2).""" + batch_size, channels = 1, 16 + latents = torch.randn(batch_size, channels, 2, 2) + + packed = pack_latents(latents) + + # 2x2 packed is 1 token + assert packed.shape == (batch_size, 1, channels * 4) + + def test_unpack_latents_minimum_size(self): + """Test unpack_latents with minimum size.""" + batch_size, channels = 1, 64 + packed = torch.randn(batch_size, 1, channels) # 1 packed token + + # Unpack to 2x2 latent (after VAE scale of 1) + unpacked = unpack_latents(packed, 2, 2, vae_scale_factor=1) + + # Output should be (B, C//4, 2, 2) + assert unpacked.shape == (batch_size, channels // 4, 2, 2) + + +class TestCompatibility(PrimusUT): + """Tests for backward compatibility with existing code.""" + + def test_unpack_latents_signature_compatibility(self): + """Test unpack_latents works with both training and inference signatures.""" + packed = torch.randn(2, 4096, 64) + + # Training path (no VAE scaling) + unpacked_train = unpack_latents(packed, 128, 128, vae_scale_factor=1) + assert unpacked_train.shape == (2, 16, 128, 128) + + # Inference path (VAE scaling) + unpacked_infer = unpack_latents(packed, 1024, 1024, vae_scale_factor=8) + assert unpacked_infer.shape == (2, 16, 128, 128) # Same result: 1024/8 = 128 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 9a6ff6a24ab5f35b6e0a6902782cc97bcc3e170f Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Fri, 10 Jul 2026 10:53:43 +0300 Subject: [PATCH 020/127] feat(flux): mxfp4 local-spec extension + fp4 utils/enums (#814) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/turbo` — review after it. ## What this changes The mxfp4 (4-bit) local-spec turbo extension plus the supporting fp4 utils and enums. ## Dependencies Builds on the CI-pins PR (`feat/flux/ci-env`) — this is the path that hard-needs the bumped Primus-Turbo (its head carries that pin; that PR merges first): on the old pin it fails with `gemm_fp4_impl(...)` "expected at most 10 args but received 11" (the concrete motivation for the CI-pins PR). Also builds on `feat/flux/turbo`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "mxfp4 or fp4_utils"`. Validated locally on an AMD GPU container: 7 passed. ## Files 5 (mxfp4 local-spec extension, fp4 utils, enums + tests). Co-authored-by: Flux Split Trial Co-authored-by: luiza-amd --- primus/backends/megatron/core/enums.py | 2 +- .../extensions/primus_turbo_mxfp4_local.py | 586 ++++++++++++++++++ primus/backends/megatron/core/fp4_utils.py | 191 +++--- .../backends/megatron/test_fp4_utils_mxfp4.py | 39 ++ .../megatron/test_primus_turbo_mxfp4_local.py | 340 ++++++++++ 5 files changed, 1075 insertions(+), 83 deletions(-) create mode 100644 primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py create mode 100644 tests/unit_tests/backends/megatron/test_fp4_utils_mxfp4.py create mode 100644 tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py diff --git a/primus/backends/megatron/core/enums.py b/primus/backends/megatron/core/enums.py index 1dd3113a8..c88650fb9 100644 --- a/primus/backends/megatron/core/enums.py +++ b/primus/backends/megatron/core/enums.py @@ -8,7 +8,7 @@ class Fp4Recipe(str, enum.Enum): - """FP4 recipe names: nvfp4.""" + """FP4 recipe names: nvfp4, mxfp4.""" nvfp4 = "nvfp4" mxfp4 = "mxfp4" diff --git a/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py b/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py new file mode 100644 index 000000000..73f5ed68b --- /dev/null +++ b/primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py @@ -0,0 +1,586 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Compile-friendly MXFP4 linear layers for Megatron local spec. + +Self-contained autograd Functions that call Primus Turbo's low-level +quantize_mxfp4_dual / quantize_mxfp4 C++ ops and gemm_fp4_impl directly, +bypassing the higher-level wrappers that construct ScalingRecipe objects +and call check_mxfp4_support() global state. + +Key properties: +- Uses setup_context pattern with primitive-only args so torch.compile + can trace through without graph breaks. +- Two backward modes: pure MXFP4 or hybrid (FP4 fwd / FP8 bwd). +- gemm_fp4_impl and gemm_fp8_impl are torch.library.custom_op with register_fake. +- Zero TransformerEngine dependencies. +- Requires tensor_model_parallel_size=1, no GAF, no sequence_parallel. +""" + +from typing import Tuple + +import torch +from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from primus_turbo.pytorch.core.backend import ( + BackendType, + GlobalBackendManager, + PrecisionType, +) +from primus_turbo.pytorch.core.low_precision import ( + ScalingGranularity, + check_mxfp4_support, +) +from primus_turbo.pytorch.kernels.gemm.gemm_fp4_impl import gemm_fp4_impl +from primus_turbo.pytorch.kernels.gemm.gemm_fp8_impl import gemm_fp8_impl + +from .primus_turbo_float8_local import _quantize_fp8_tw + +# The AITER MXFP4 preshuffle fast path used to be implemented here as a +# monkey patch on GEMMFP4AITERBackend.execute. It now lives in Primus-Turbo +# (commit 683a7de, "feat(gemm): add AITER MXFP4 preshuffle fast path") and +# is opted into per-call via the gemm_fp4_impl(preshuffled=...) kwarg. +# +# Primus-Turbo PR #383 ("refactor preshuffle ...") removed the public +# enable_preshuffle() helper and moved per-call preshuffle control onto +# Float4QuantConfig.use_preshuffle. MXFP4LinearFunction passes a plain bool +# into its custom ops (not a Float4QuantConfig), so we keep the original +# runtime probe here as a module-local helper, _enable_preshuffle(), which +# reproduces the removed upstream logic verbatim. +# +# We resolve _enable_preshuffle() once per MXFP4 linear at __init__ time and +# cache the result on self._preshuffle. If the dispatcher state changes +# after model construction (env unset, set_gemm_backend(None), autotune +# flipped on) the cached flag is stale and the forward path will +# mis-dispatch -- the same caching pattern as before, just now made loud by +# the module-init contract check below (_assert_preshuffle_contract, which +# raises RuntimeError). Resolving it at __init__ (not inside forward) +# also keeps it out of the torch.compile traced region. + + +def _enable_preshuffle() -> bool: + """True iff the AITER FP4 preshuffle fast path is safe. + + Requires: (1) the FP4 GEMM backend is explicitly pinned to AITER (the only + backend that understands the shuffled layout), and (2) autotune is disabled + (AITER opts out of tuning, so the tuner cannot select a backend for + preshuffled inputs). Reproduces the removed Primus-Turbo enable_preshuffle() + (pre Primus-Turbo PR #383) so MXFP4LinearFunction can keep passing a bool + into its custom ops. + """ + return ( + GlobalBackendManager.get_gemm_backend(PrecisionType.FP4) == BackendType.AITER + and not GlobalBackendManager.auto_tune_enabled() + ) + + +def _assert_preshuffle_contract(config, preshuffle: bool) -> None: + """Fail loudly when MXFP4 was explicitly requested but the upstream + AITER preshuffle fast path is unavailable. + + Pre-cleanup, the local monkey patch silently consumed preshuffled + inputs on the AITER backend regardless of dispatcher state, masking + misconfigurations as either a ~5% step-time regression (preshuffle + inputs landed on AITER but every GEMM paid the shuffle cost) or, on + dispatch paths 2/3/4, silent numerical corruption (HipBLASLt received + preshuffled bytes and ran with them). Post-cleanup, both modes become + a hard init-time failure with an actionable message. + + Gated on ``config.fp4 == "mxfp4"`` so any speculative-instantiation + flow that constructs the class without actually opting into MXFP4 is + unaffected (mirrors the precedent of the other ``__init__`` asserts + in this file). + """ + if getattr(config, "fp4", None) != "mxfp4": + return + if not preshuffle: + raise RuntimeError( + "MXFP4 linear requested (config.fp4='mxfp4') but the AITER " + "preshuffle fast path is unavailable. Set " + "PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER (or call " + "GlobalBackendManager.set_gemm_backend(BackendType.AITER, " + "PrecisionType.FP4)) and leave PRIMUS_TURBO_AUTO_TUNE unset. " + "Without this, every FP4 GEMM pays ~3 extra shuffle kernel " + "launches (~5% Flux 12B step-time regression vs the tuned " + "baseline) -- and on non-AITER dispatch the call now raises." + ) + + +# --------------------------------------------------------------------------- +# torch.compile-friendly wrappers for MXFP4 C++ quantization ops. +# +# The raw C++ ops lack an Autograd dispatch key. Wrapping them with +# @torch.library.custom_op provides Autograd dispatch and register_fake +# for shape inference during tracing. +# --------------------------------------------------------------------------- + +_custom_op = torch.library.custom_op + +MXFP4_BLOCK_SIZE = 32 +# Bumped from 16 -> 128 to match Primus-Turbo PR #335 ("feat: add quantized +# tensor support"), which made `padding_align_size` an explicit positional arg +# of quantize_mxfp4{_dual} and asserts it equals MXFP4_PADDING_ALIGN_SIZE (=128) +# in csrc/include/primus_turbo/quantization.h. +MXFP4_PADDING_ALIGN_SIZE = 128 + + +def _cdiv(a, b): + return (a + b - 1) // b + + +@_custom_op("primus::quantize_mxfp4_dual", mutates_args=(), device_types="cuda") +def _quantize_mxfp4_dual_op( + x: torch.Tensor, + out_dtype: torch.dtype, + padding_align_size: int, + rowwise_use_2d_block: bool, + rowwise_use_sr: bool, + rowwise_use_rht: bool, + colwise_use_2d_block: bool, + colwise_use_sr: bool, + colwise_use_rht: bool, + shuffle_rowwise_scale: bool, + shuffle_rowwise: bool, + shuffle_colwise_scale: bool, + shuffle_colwise: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.ops.primus_turbo_cpp_extension.quantize_mxfp4_dual( + x, + out_dtype, + padding_align_size, + rowwise_use_2d_block, + rowwise_use_sr, + rowwise_use_rht, + colwise_use_2d_block, + colwise_use_sr, + colwise_use_rht, + shuffle_rowwise_scale, + shuffle_rowwise, + shuffle_colwise_scale, + shuffle_colwise, + ) + + +@_quantize_mxfp4_dual_op.register_fake +def _quantize_mxfp4_dual_fake( + x: torch.Tensor, + out_dtype: torch.dtype, + padding_align_size: int, + rowwise_use_2d_block: bool, + rowwise_use_sr: bool, + rowwise_use_rht: bool, + colwise_use_2d_block: bool, + colwise_use_sr: bool, + colwise_use_rht: bool, + shuffle_rowwise_scale: bool, + shuffle_rowwise: bool, + shuffle_colwise_scale: bool, + shuffle_colwise: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + M, N = x.shape + M_pad = _cdiv(M, MXFP4_PADDING_ALIGN_SIZE) * MXFP4_PADDING_ALIGN_SIZE + N_pad = _cdiv(N, MXFP4_PADDING_ALIGN_SIZE) * MXFP4_PADDING_ALIGN_SIZE + + if shuffle_rowwise_scale: + rs_M = _cdiv(M, 256) * 256 + rs_N = _cdiv(_cdiv(N_pad, MXFP4_BLOCK_SIZE), 8) * 8 + rowwise_scale = torch.empty(rs_M, rs_N, dtype=torch.uint8, device=x.device) + else: + rowwise_scale = torch.empty(M, _cdiv(N_pad, MXFP4_BLOCK_SIZE), dtype=torch.uint8, device=x.device) + + rowwise_output = torch.empty(M, N_pad // 2, dtype=torch.uint8, device=x.device) + + if shuffle_colwise_scale: + cs_M = _cdiv(N, 256) * 256 + cs_N = _cdiv(_cdiv(M_pad, MXFP4_BLOCK_SIZE), 8) * 8 + colwise_scale = torch.empty(cs_M, cs_N, dtype=torch.uint8, device=x.device) + else: + colwise_scale = torch.empty(N, _cdiv(M_pad, MXFP4_BLOCK_SIZE), dtype=torch.uint8, device=x.device) + + colwise_output = torch.empty(N, M_pad // 2, dtype=torch.uint8, device=x.device) + + return ( + rowwise_output.view(torch.float4_e2m1fn_x2), + rowwise_scale.view(torch.float8_e8m0fnu), + colwise_output.view(torch.float4_e2m1fn_x2), + colwise_scale.view(torch.float8_e8m0fnu), + ) + + +def _quantize_mxfp4_dual_setup_context(ctx, inputs, output): + pass + + +def _quantize_mxfp4_dual_backward(ctx, *grad_outputs): + return (None,) * 13 + + +_quantize_mxfp4_dual_op.register_autograd( + _quantize_mxfp4_dual_backward, + setup_context=_quantize_mxfp4_dual_setup_context, +) + + +# --------------------------------------------------------------------------- +# MXFP4 Autograd Function with setup_context (torch.compile friendly) +# --------------------------------------------------------------------------- + + +_FP4_DTYPE = torch.float4_e2m1fn_x2 +_GRAN_VALUE = ScalingGranularity.MX_BLOCKWISE.value +_DEFAULT_BACKEND = BackendType.HIPBLASLT.value + + +def _quantize_input_dual(input_2d, preshuffle): + """Quantize input (activation) with dual rowwise + colwise.""" + return _quantize_mxfp4_dual_op( + input_2d, + _FP4_DTYPE, + MXFP4_PADDING_ALIGN_SIZE, + False, + False, + False, # rowwise: no 2d_block, no sr, no rht + False, + False, + True, # colwise: no 2d_block, no sr, yes rht + preshuffle, + False, # shuffle_rowwise_scale, shuffle_rowwise + preshuffle, + preshuffle, # shuffle_colwise_scale, shuffle_colwise + ) + + +def _quantize_weight_dual(weight, preshuffle): + """Quantize weight with dual rowwise + colwise.""" + return _quantize_mxfp4_dual_op( + weight, + _FP4_DTYPE, + MXFP4_PADDING_ALIGN_SIZE, + True, + False, + False, # rowwise: 2d_block, no sr, no rht + True, + False, + False, # colwise: 2d_block, no sr, no rht + preshuffle, + preshuffle, # shuffle_rowwise_scale, shuffle_rowwise + preshuffle, + preshuffle, # shuffle_colwise_scale, shuffle_colwise + ) + + +def _quantize_grad_dual(grad_2d, preshuffle, use_sr=True): + """Quantize gradient (activation recipe) with dual rowwise + colwise.""" + return _quantize_mxfp4_dual_op( + grad_2d, + _FP4_DTYPE, + MXFP4_PADDING_ALIGN_SIZE, + False, + use_sr, + False, # rowwise: no 2d_block, SR configurable, no rht + False, + use_sr, + True, # colwise: no 2d_block, SR configurable, yes rht + preshuffle, + False, # shuffle_rowwise_scale, shuffle_rowwise + preshuffle, + False, # shuffle_colwise_scale, shuffle_colwise + ) + + +class MXFP4LinearFunction(torch.autograd.Function): + """MXFP4 linear (Y = X @ W^T) with MX block-of-32 scaling. + + Two modes via backward_is_fp8 bool primitive: + - Pure MXFP4: forward + backward both use FP4 quantization + gemm_fp4_impl + - Hybrid: forward uses FP4, backward re-quantizes saved BF16 to FP8 tensorwise + + Uses setup_context pattern with primitive-only args for torch.compile. + """ + + @staticmethod + def forward( + input, + weight, + preshuffle, + backward_is_fp8, + fp8_bwd_dtype, + fp8_gran_value, + fp8_backend_value, + use_gradient_sr, + ): + out_dtype = input.dtype + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + a_fp4, a_scale, a_t_fp4, a_t_scale = _quantize_input_dual(input_2d, preshuffle) + b_fp4, b_scale, b_t_fp4, b_t_scale = _quantize_weight_dual(weight, preshuffle) + + output = gemm_fp4_impl( + a_fp4, + a_scale, + False, + b_fp4, + b_scale, + True, + out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_DEFAULT_BACKEND, + preshuffled=preshuffle, + ) + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + + if backward_is_fp8: + return output, input_2d.view_as(input_2d), weight.view_as(weight) + else: + # Return FP4/FP8 tensors as uint8 views to avoid + # "fill_cuda not implemented for Float4_e2m1fn_x2" when + # the autograd engine creates zero gradients for + # non-differentiable outputs. + return ( + output, + a_t_fp4.view(torch.uint8), + a_t_scale.view(torch.uint8), + b_t_fp4.view(torch.uint8), + b_t_scale.view(torch.uint8), + ) + + @staticmethod + def setup_context(ctx, inputs, output): + ( + _, + _, + preshuffle, + backward_is_fp8, + fp8_bwd_dtype, + fp8_gran_value, + fp8_backend_value, + use_gradient_sr, + ) = inputs + + ctx.preshuffle = preshuffle + ctx.backward_is_fp8 = backward_is_fp8 + ctx.use_gradient_sr = use_gradient_sr + ctx.out_dtype = inputs[0].dtype + ctx.orig_shape = inputs[0].shape + + if backward_is_fp8: + _, input_2d_saved, weight_saved = output + ctx.save_for_backward(input_2d_saved, weight_saved) + ctx.fp8_bwd_dtype = fp8_bwd_dtype + ctx.fp8_gran_value = fp8_gran_value + ctx.fp8_backend_value = fp8_backend_value + else: + _, a_t_u8, as_u8, b_t_u8, bs_u8 = output + ctx.save_for_backward( + a_t_u8.view(_FP4_DTYPE), + as_u8.view(torch.float8_e8m0fnu), + b_t_u8.view(_FP4_DTYPE), + bs_u8.view(torch.float8_e8m0fnu), + ) + ctx.mark_non_differentiable(a_t_u8, as_u8, b_t_u8, bs_u8) + + @staticmethod + def backward(ctx, grad_output, *_): + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + + if ctx.backward_is_fp8: + input_2d, weight = ctx.saved_tensors + + grad_fp8, grad_scale_inv = _quantize_fp8_tw(grad_2d, ctx.fp8_bwd_dtype) + a_fp8, a_scale_inv = _quantize_fp8_tw(input_2d, ctx.fp8_bwd_dtype) + b_fp8, b_scale_inv = _quantize_fp8_tw(weight, ctx.fp8_bwd_dtype) + + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + b_fp8, + b_scale_inv, + False, + ctx.out_dtype, + False, + granularity=ctx.fp8_gran_value, + default_backend=ctx.fp8_backend_value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_weight = gemm_fp8_impl( + a_fp8, + a_scale_inv, + True, + grad_fp8, + grad_scale_inv, + False, + ctx.out_dtype, + True, + granularity=ctx.fp8_gran_value, + default_backend=ctx.fp8_backend_value, + ) + else: + a_t_fp4, a_t_scale, b_t_fp4, b_t_scale = ctx.saved_tensors + preshuffle = ctx.preshuffle + + g_fp4, g_scale, g_t_fp4, g_t_scale = _quantize_grad_dual( + grad_2d, preshuffle, use_sr=ctx.use_gradient_sr + ) + + grad_input = gemm_fp4_impl( + g_fp4, + g_scale, + False, + b_t_fp4, + b_t_scale, + True, + ctx.out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_DEFAULT_BACKEND, + preshuffled=preshuffle, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_weight = gemm_fp4_impl( + g_t_fp4, + g_t_scale, + False, + a_t_fp4, + a_t_scale, + True, + ctx.out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_DEFAULT_BACKEND, + preshuffled=preshuffle, + ) + + return grad_input, grad_weight, None, None, None, None, None, None + + +# --------------------------------------------------------------------------- +# MXFP4-aware parallel linear layers +# --------------------------------------------------------------------------- + + +class MXFP4ColumnParallelLinear(ColumnParallelLinear): + """ColumnParallelLinear with per-module MXFP4. torch.compile friendly. + + Requires: tensor_model_parallel_size=1, gradient_accumulation_fusion=False, + sequence_parallel=False. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + if self.config.tensor_model_parallel_size != 1: + raise ValueError( + "MXFP4ColumnParallelLinear requires tensor_model_parallel_size=1. " + f"Got {self.config.tensor_model_parallel_size}." + ) + if self.gradient_accumulation_fusion: + raise ValueError("MXFP4ColumnParallelLinear requires gradient_accumulation_fusion=False.") + if self.sequence_parallel: + raise ValueError("MXFP4ColumnParallelLinear requires sequence_parallel=False.") + + supported, reason = check_mxfp4_support() + if not supported: + raise RuntimeError(f"MXFP4 not supported on this device: {reason}") + + self._preshuffle = _enable_preshuffle() + _assert_preshuffle_contract(self.config, self._preshuffle) + self._backward_is_fp8 = getattr(self.config, "mxfp4_backward_precision", "mxfp4") == "fp8" + self._use_gradient_sr = getattr(self.config, "mxfp4_gradient_stochastic_rounding", False) + + if self._backward_is_fp8: + from primus_turbo.pytorch.core.low_precision import float8_e5m2 + + self._fp8_bwd_dtype = float8_e5m2 + self._fp8_gran_value = ScalingGranularity.TENSORWISE.value + self._fp8_backend_value = BackendType.HIPBLASLT.value + else: + self._fp8_bwd_dtype = None + self._fp8_gran_value = 0 + self._fp8_backend_value = 0 + + def _forward_impl(self, input, weight, *args, **kwargs): + bias = kwargs.get("bias", None) + + result = MXFP4LinearFunction.apply( + input, + weight, + self._preshuffle, + self._backward_is_fp8, + self._fp8_bwd_dtype, + self._fp8_gran_value, + self._fp8_backend_value, + self._use_gradient_sr, + ) + output = result[0] + + if bias is not None: + output = output + bias + return output + + +class MXFP4RowParallelLinear(RowParallelLinear): + """RowParallelLinear with per-module MXFP4. torch.compile friendly. + + Requires: tensor_model_parallel_size=1, gradient_accumulation_fusion=False, + sequence_parallel=False. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + if self.config.tensor_model_parallel_size != 1: + raise ValueError( + "MXFP4RowParallelLinear requires tensor_model_parallel_size=1. " + f"Got {self.config.tensor_model_parallel_size}." + ) + if self.gradient_accumulation_fusion: + raise ValueError("MXFP4RowParallelLinear requires gradient_accumulation_fusion=False.") + if self.sequence_parallel: + raise ValueError("MXFP4RowParallelLinear requires sequence_parallel=False.") + + supported, reason = check_mxfp4_support() + if not supported: + raise RuntimeError(f"MXFP4 not supported on this device: {reason}") + + self._preshuffle = _enable_preshuffle() + _assert_preshuffle_contract(self.config, self._preshuffle) + self._backward_is_fp8 = getattr(self.config, "mxfp4_backward_precision", "mxfp4") == "fp8" + self._use_gradient_sr = getattr(self.config, "mxfp4_gradient_stochastic_rounding", False) + + if self._backward_is_fp8: + from primus_turbo.pytorch.core.low_precision import float8_e5m2 + + self._fp8_bwd_dtype = float8_e5m2 + self._fp8_gran_value = ScalingGranularity.TENSORWISE.value + self._fp8_backend_value = BackendType.HIPBLASLT.value + else: + self._fp8_bwd_dtype = None + self._fp8_gran_value = 0 + self._fp8_backend_value = 0 + + def _forward_impl(self, input, weight, *args, **kwargs): + bias = kwargs.get("bias", None) + + result = MXFP4LinearFunction.apply( + input, + weight, + self._preshuffle, + self._backward_is_fp8, + self._fp8_bwd_dtype, + self._fp8_gran_value, + self._fp8_backend_value, + self._use_gradient_sr, + ) + output = result[0] + + if bias is not None: + output = output + bias + return output diff --git a/primus/backends/megatron/core/fp4_utils.py b/primus/backends/megatron/core/fp4_utils.py index 280949961..b747d9b60 100644 --- a/primus/backends/megatron/core/fp4_utils.py +++ b/primus/backends/megatron/core/fp4_utils.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -67,7 +67,8 @@ def get_fp4_recipe(config: TransformerConfig): fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling() except AttributeError: fp4_recipe_none_reason = ( - "NVFP4BlockScaling recipe is not available in this version of Transformer Engine." + "NVFP4BlockScaling recipe is not available in this version of " + "Transformer Engine. Please make sure you are using TE version >= 2.7.0.dev0." ) elif config.fp4_recipe == Fp4Recipe.mxfp4: try: @@ -77,10 +78,14 @@ def get_fp4_recipe(config: TransformerConfig): fp4_recipe.use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1" except AttributeError: fp4_recipe_none_reason = ( - "MXFP4BlockScaling recipe is not available in this version of Transformer Engine." + "MXFP4BlockScaling recipe is not available in this version of " + "Transformer Engine. MXFP4 requires ROCm TE with AITER support." ) else: - fp4_recipe_none_reason = f"Unsupported FP4 recipe: {config.fp4_recipe}." + fp4_recipe_none_reason = ( + f"Unsupported fp4_recipe '{config.fp4_recipe}'. " + "Supported recipes: 'nvfp4' (NVIDIA), 'mxfp4' (AMD ROCm with AITER)." + ) else: fp4_recipe_none_reason = "FP4 support requires TransformerEngine version >= 2.7.0.dev0." @@ -115,6 +120,7 @@ def get_fp4_quant_config(config: TransformerConfig): format=Format.E2M1_X2, block_size=MXFP4_SCALING_BLOCK_SIZE, scale_dtype=ScaleDtype.E8M0, + use_gradient_sr=getattr(config, "mxfp4_gradient_stochastic_rounding", False), ) return fp4_quant_config, "" @@ -132,62 +138,67 @@ def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool elif layer_no >= 0 and config.first_last_layers_bf16 and (is_first_layer or is_last_layer): fp4_context = nullcontext() else: - fp4_recipe, fp4_recipe_none_reason = get_fp4_recipe(config) - turbo_enabled = _primus_turbo_enabled() - - global WARN_ONCE - if WARN_ONCE: - if fp4_recipe is None: - warning_rank_0( - f"TransformerEngine FP4 {config.fp4_recipe} not work since {fp4_recipe_none_reason}" - ) - if is_init: - warning_rank_0( - f"Primus-Turbo FP4 {config.fp4_recipe} not work since Primus-Turbo not support fp4 model init." - ) - WARN_ONCE = False - - fp4_group = None - if parallel_state.model_parallel_is_initialized(): - fp4_group = parallel_state.get_amax_reduction_group( - with_context_parallel=True, tp_only_amax_red=config.tp_only_amax_red - ) + # Local spec bypasses TE autocast -- quantization is handled + # internally by PrimusTurboMXFP4LocalSpecProvider's custom ops. + if getattr(config, "transformer_impl", "transformer_engine") == "local": + fp4_context = nullcontext() + else: + fp4_recipe, fp4_recipe_none_reason = get_fp4_recipe(config) + turbo_enabled = _primus_turbo_enabled() - if not is_init: - # Only touch the Primus-Turbo extension when the Turbo FP4 - # autocast path is explicitly enabled; otherwise use TE directly. - if turbo_enabled: - fp4_quant_config, fp4_quant_config_none_reason = get_fp4_quant_config(config) - if WARN_ONCE and fp4_quant_config is None: + global WARN_ONCE + if WARN_ONCE: + if fp4_recipe is None: + warning_rank_0( + f"TransformerEngine FP4 {config.fp4_recipe} not work since {fp4_recipe_none_reason}" + ) + if is_init: warning_rank_0( - f"Primus-Turbo FP4 {config.fp4_recipe} not work since {fp4_quant_config_none_reason}" + f"Primus-Turbo FP4 {config.fp4_recipe} not work since Primus-Turbo not support fp4 model init." ) + WARN_ONCE = False - from primus.backends.megatron.core.extensions.primus_turbo import ( - primus_turbo_fp4_autocast, + fp4_group = None + if parallel_state.model_parallel_is_initialized(): + fp4_group = parallel_state.get_amax_reduction_group( + with_context_parallel=True, tp_only_amax_red=config.tp_only_amax_red ) - fp4_context = primus_turbo_fp4_autocast( - enabled=True if fp4_recipe is not None else False, - fp4_recipe=fp4_recipe, - fp4_group=fp4_group, - enabled_turbo=True if fp4_quant_config is not None else False, - turbo_quant_config=fp4_quant_config, - ) + if not is_init: + # Only touch the Primus-Turbo extension when the Turbo FP4 + # autocast path is explicitly enabled; otherwise use TE directly. + if turbo_enabled: + fp4_quant_config, fp4_quant_config_none_reason = get_fp4_quant_config(config) + if WARN_ONCE and fp4_quant_config is None: + warning_rank_0( + f"Primus-Turbo FP4 {config.fp4_recipe} not work since {fp4_quant_config_none_reason}" + ) + + from primus.backends.megatron.core.extensions.primus_turbo import ( + primus_turbo_fp4_autocast, + ) + + fp4_context = primus_turbo_fp4_autocast( + enabled=True if fp4_recipe is not None else False, + fp4_recipe=fp4_recipe, + fp4_group=fp4_group, + enabled_turbo=True if fp4_quant_config is not None else False, + turbo_quant_config=fp4_quant_config, + ) + else: + # TE currently uses fp8_autocast for fp8 and fp4 quantization. + fp4_context = transformer_engine.pytorch.fp8_autocast( + enabled=True if fp4_recipe is not None else False, + fp8_recipe=fp4_recipe, + fp8_group=fp4_group, + ) else: - # TE currently uses fp8_autocast for fp8 and fp4 quantization. - fp4_context = transformer_engine.pytorch.fp8_autocast( - enabled=True if fp4_recipe is not None else False, - fp8_recipe=fp4_recipe, - fp8_group=fp4_group, - ) - else: - import inspect + import inspect - context_args = {"enabled": True} - if "recipe" in inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters: - context_args["recipe"] = fp4_recipe - fp4_context = transformer_engine.pytorch.fp8_model_init(**context_args) + context_args = {"enabled": True} + if "recipe" in inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters: + context_args["recipe"] = fp4_recipe + fp4_context = transformer_engine.pytorch.fp8_model_init(**context_args) return fp4_context @@ -195,22 +206,34 @@ def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool def get_fp4_recipe(config: TransformerConfig): """Return fp4 recipe.""" - if config.fp4_recipe == Fp4Recipe.nvfp4: - if not is_te_min_version("2.7.0.dev0"): - raise ValueError("NVFP4BlockScaling requires TransformerEngine >= 2.7.0.dev0.") - fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling() - elif config.fp4_recipe == Fp4Recipe.mxfp4: - try: - import os - - fp4_recipe = transformer_engine.common.recipe.MXFP4BlockScaling() - fp4_recipe.use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1" - except AttributeError: + if is_te_min_version("2.7.0.dev0"): + if config.fp4_recipe == Fp4Recipe.nvfp4: + try: + fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling() + except AttributeError: + raise ValueError( + "NVFP4BlockScaling recipe is not available in this version of " + "Transformer Engine. Please make sure you are using TE version " + ">= 2.7.0.dev0." + ) + elif config.fp4_recipe == Fp4Recipe.mxfp4: + try: + import os + + fp4_recipe = transformer_engine.common.recipe.MXFP4BlockScaling() + fp4_recipe.use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1" + except AttributeError: + raise ValueError( + "MXFP4BlockScaling recipe is not available in this version of " + "Transformer Engine. MXFP4 requires ROCm TE with AITER support." + ) + else: raise ValueError( - "MXFP4BlockScaling recipe is not available in this version of Transformer Engine." + f"Unsupported fp4_recipe '{config.fp4_recipe}'. " + "Supported recipes: 'nvfp4' (NVIDIA), 'mxfp4' (AMD ROCm with AITER)." ) else: - raise ValueError(f"Unsupported FP4 recipe: {config.fp4_recipe}. " "Supported: nvfp4, mxfp4.") + raise ValueError("FP4 support requires TransformerEngine version >= 2.7.0.dev0.") return fp4_recipe def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool = False): @@ -227,25 +250,29 @@ def get_fp4_context(config: TransformerConfig, layer_no: int = -1, is_init: bool elif layer_no >= 0 and config.first_last_layers_bf16 and (is_first_layer or is_last_layer): fp4_context = nullcontext() else: - fp4_recipe = get_fp4_recipe(config) - fp4_group = None - if parallel_state.model_parallel_is_initialized(): - fp4_group = parallel_state.get_amax_reduction_group( - with_context_parallel=True, tp_only_amax_red=config.tp_only_amax_red - ) - - if not is_init: - # TE currently uses fp8_autocast for fp8 and fp4 quantization. - fp4_context = transformer_engine.pytorch.fp8_autocast( - enabled=True, fp8_recipe=fp4_recipe, fp8_group=fp4_group - ) + # Local spec bypasses TE autocast -- quantization is handled + # internally by PrimusTurboMXFP4LocalSpecProvider's custom ops. + if getattr(config, "transformer_impl", "transformer_engine") == "local": + fp4_context = nullcontext() else: - import inspect + fp4_recipe = get_fp4_recipe(config) + fp4_group = None + if parallel_state.model_parallel_is_initialized(): + fp4_group = parallel_state.get_amax_reduction_group( + with_context_parallel=True, tp_only_amax_red=config.tp_only_amax_red + ) + + if not is_init: + fp4_context = transformer_engine.pytorch.fp8_autocast( + enabled=True, fp8_recipe=fp4_recipe, fp8_group=fp4_group + ) + else: + import inspect - context_args = {"enabled": True} - if "recipe" in inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters: - context_args["recipe"] = fp4_recipe - fp4_context = transformer_engine.pytorch.fp8_model_init(**context_args) + context_args = {"enabled": True} + if "recipe" in inspect.signature(transformer_engine.pytorch.fp8_model_init).parameters: + context_args["recipe"] = fp4_recipe + fp4_context = transformer_engine.pytorch.fp8_model_init(**context_args) return fp4_context diff --git a/tests/unit_tests/backends/megatron/test_fp4_utils_mxfp4.py b/tests/unit_tests/backends/megatron/test_fp4_utils_mxfp4.py new file mode 100644 index 000000000..17097796b --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_fp4_utils_mxfp4.py @@ -0,0 +1,39 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for fp4_utils.py MXFP4 recipe and context manager changes. + +Tests get_fp4_recipe error handling for an unsupported recipe. +""" + +from types import SimpleNamespace + +import pytest + +from tests.utils import PrimusUT + + +class TestGetFp4RecipeMXFP4(PrimusUT): + """Verify get_fp4_recipe returns correct recipe objects for MXFP4.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def test_unsupported_recipe_produces_error(self): + pytest.importorskip("transformer_engine") + + from primus.backends.megatron.core.fp4_utils import get_fp4_recipe + + config = SimpleNamespace(fp4_recipe="nonexistent_recipe") + result = get_fp4_recipe(config) + + if isinstance(result, tuple): + recipe, reason = result + assert recipe is None, "Unsupported recipe should return None" + assert ( + "Unsupported" in reason or "unsupported" in reason.lower() + ), f"Expected 'Unsupported' in reason, got: {reason}" + else: + pytest.fail("HAVE_TE-only branch should raise ValueError for unsupported recipe") diff --git a/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py b/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py new file mode 100644 index 000000000..a0fc099b9 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_primus_turbo_mxfp4_local.py @@ -0,0 +1,340 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for compile-friendly MXFP4 linear layers (primus_turbo_mxfp4_local). + +Tests cross-validation against Primus-Turbo's FP4GemmMXFunction reference, +torch.compile graph-break validation, Megatron linear backward flow, +2-step training loop, hybrid (FP4 fwd / FP8 bwd) mode, and init guards. +""" + +import functools +import os +from types import SimpleNamespace + +import pytest +import torch +from megatron.core.transformer.transformer_config import TransformerConfig + +from tests.unit_tests.backends.megatron.conftest import requires_mxfp4 +from tests.utils import PrimusUT + + +def _init_method(): + return functools.partial(torch.nn.init.xavier_uniform_) + + +def _make_mxfp4_transformer_config(**overrides): + defaults = dict( + hidden_size=256, + num_attention_heads=8, + num_layers=1, + params_dtype=torch.bfloat16, + fp4="mxfp4", + fp4_recipe="mxfp4", + ) + defaults.update(overrides) + return TransformerConfig(**defaults) + + +def _pin_fp4_aiter(monkeypatch): + """Pin the FP4 GEMM backend to AITER with autotune off for MXFP4 module tests. + + MXFP4 module __init__ runs _assert_preshuffle_contract, which requires the + FP4 GEMM backend pinned to AITER with autotune off (the only config under + which _enable_preshuffle() is True). Pinning it in-code lets the + module-instantiation tests reach the real path instead of failing the + contract; monkeypatch auto-restores on teardown so the .apply-direct tests + keep their default (preshuffle=False) dispatch. Also clears any baked-empty + PRIMUS_TURBO_GEMM_BACKEND so the in-code pin is authoritative (mirrors + test_native_fp8_layout.py). + """ + from primus_turbo.pytorch.core.backend import ( + BackendType, + GlobalBackendManager, + PrecisionType, + ) + + if os.environ.get("PRIMUS_TURBO_GEMM_BACKEND", None) == "": + monkeypatch.delenv("PRIMUS_TURBO_GEMM_BACKEND", raising=False) + monkeypatch.setattr(GlobalBackendManager, "_gemm_backend", {PrecisionType.FP4: BackendType.AITER}) + monkeypatch.setattr(GlobalBackendManager, "_auto_tune", False) + + +# --------------------------------------------------------------------------- +# Cross-validation against Primus-Turbo's FP4GemmMXFunction reference +# --------------------------------------------------------------------------- + + +class TestMXFP4CrossValidation(PrimusUT): + """Verify MXFP4LinearFunction produces bit-identical results to FP4GemmMXFunction. + + Catches wrong boolean flags in _quantize_input_dual / _quantize_weight_dual / + _quantize_grad_dual. A single wrong flag produces silently incorrect numerics + that may still pass SNR thresholds vs BF16 but diverges from the canonical path. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + @requires_mxfp4 + def test_forward_matches_reference_fp4gemm(self): + from primus_turbo.pytorch.core.low_precision import Float4QuantConfig + from primus_turbo.pytorch.ops.gemm_fp4 import FP4GemmMXFunction + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4LinearFunction, + _enable_preshuffle, + ) + + torch.manual_seed(42) + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda") + preshuffle = _enable_preshuffle() + + result = MXFP4LinearFunction.apply( + x, + w, + preshuffle, + False, + None, + 0, + 0, + False, + ) + our_output = result[0] + + config = Float4QuantConfig(use_preshuffle=preshuffle) + ref_output = FP4GemmMXFunction.apply( + x.clone(), + w.clone(), + None, + None, + False, + True, + x.dtype, + config, + ) + + assert torch.equal(our_output, ref_output), ( + f"Forward outputs differ. Max abs diff: " f"{(our_output - ref_output).abs().max().item():.6e}" + ) + + @requires_mxfp4 + def test_backward_matches_reference_fp4gemm(self): + from primus_turbo.pytorch.core.low_precision import Float4QuantConfig + from primus_turbo.pytorch.ops.gemm_fp4 import FP4GemmMXFunction + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4LinearFunction, + _enable_preshuffle, + ) + + torch.manual_seed(42) + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True) + w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True) + x_ref = x.detach().clone().requires_grad_(True) + w_ref = w.detach().clone().requires_grad_(True) + preshuffle = _enable_preshuffle() + + result = MXFP4LinearFunction.apply(x, w, preshuffle, False, None, 0, 0, False) + our_output = result[0] + grad_out = torch.ones_like(our_output) + our_output.backward(grad_out) + + config = Float4QuantConfig(use_preshuffle=preshuffle) + ref_output = FP4GemmMXFunction.apply( + x_ref, + w_ref, + None, + None, + False, + True, + x_ref.dtype, + config, + ) + ref_output.backward(torch.ones_like(ref_output)) + + # grad_weight stays bit-identical to the reference: our grad_weight + # GEMM pair (g_t colwise + a_t colwise) and the reference's both use the + # RHT recipe, so the quantization is identical. Keep torch.equal here -- + # it still catches a wrong flag in _quantize_input_dual (colwise a_t) or + # the g_t branch of _quantize_grad_dual. + assert torch.equal(w.grad, w_ref.grad), ( + f"grad_weight differs. Max abs diff: " f"{(w.grad - w_ref.grad).abs().max().item():.6e}" + ) + + # grad_input is NOT bit-identical, and that is expected post Primus-Turbo + # PR #383. The grad_input GEMM pair is (grad rowwise) x (weight colwise b_t). + # Our production deliberately quantizes this pair without RHT + # (_quantize_grad_dual rowwise use_rht=False + _quantize_weight_dual + # colwise use_rht=False -- an internally consistent no-RHT pair), whereas + # Primus-Turbo PR #383's FP4GemmMXFunction.backward unconditionally quantizes the grad + # with use_rht=True and derives b_t with use_rht=True. Both compute a + # valid grad_input (RHT cancels within each consistent pair); they only + # differ in MXFP4 quantization noise. Measured against the true BF16 + # gradient, our no-RHT grad_input is actually marginally more accurate + # than the reference's RHT grad_input (~18.4 dB vs ~17.7 dB SNR), so this + # is a recipe choice, not a regression. Assert SNR vs the BF16 truth + # (same >10 dB bar as the forward SNR test) instead of bit-identity. + bf16_grad_input = grad_out.float() @ w.detach().float() + signal = (bf16_grad_input**2).mean() + noise = ((x.grad.float() - bf16_grad_input) ** 2).mean() + snr_db = 10 * torch.log10(signal / noise).item() + assert snr_db > 10, ( + f"grad_input SNR {snr_db:.1f} dB vs BF16 is below the 10 dB threshold " + f"(max abs diff vs Primus-Turbo PR #383 reference: {(x.grad - x_ref.grad).abs().max().item():.6e})" + ) + + +# --------------------------------------------------------------------------- +# torch.compile graph-break validation +# --------------------------------------------------------------------------- + + +class TestMXFP4Compile(PrimusUT): + """Verify MXFP4LinearFunction has zero graph breaks under torch.compile.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + @requires_mxfp4 + def test_no_graph_break_pure_mxfp4(self): + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4LinearFunction, + _enable_preshuffle, + ) + + torch._dynamo.reset() + + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda") + preshuffle = _enable_preshuffle() + + explanation = torch._dynamo.explain( + MXFP4LinearFunction.apply, + )(x, w, preshuffle, False, None, 0, 0, False) + + assert explanation.graph_break_count == 0, ( + f"Expected 0 graph breaks, got {explanation.graph_break_count}. " + f"Reasons: {explanation.break_reasons}" + ) + + @requires_mxfp4 + def test_no_graph_break_hybrid(self): + from primus_turbo.pytorch.core.backend import BackendType + from primus_turbo.pytorch.core.low_precision import ( + ScalingGranularity, + float8_e5m2, + ) + + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4LinearFunction, + _enable_preshuffle, + ) + + torch._dynamo.reset() + + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda") + preshuffle = _enable_preshuffle() + + explanation = torch._dynamo.explain( + MXFP4LinearFunction.apply, + )( + x, + w, + preshuffle, + True, + float8_e5m2, + ScalingGranularity.TENSORWISE.value, + BackendType.HIPBLASLT.value, + False, + ) + + assert explanation.graph_break_count == 0, ( + f"Expected 0 graph breaks, got {explanation.graph_break_count}. " + f"Reasons: {explanation.break_reasons}" + ) + + @requires_mxfp4 + def test_compiled_forward_matches_eager(self): + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4LinearFunction, + _enable_preshuffle, + ) + + torch._dynamo.reset() + torch.manual_seed(42) + + x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda") + w = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda") + preshuffle = _enable_preshuffle() + + eager_result = MXFP4LinearFunction.apply(x, w, preshuffle, False, None, 0, 0, False) + eager_out = eager_result[0] + + compiled_fn = torch.compile(MXFP4LinearFunction.apply) + compiled_result = compiled_fn(x, w, preshuffle, False, None, 0, 0, False) + compiled_out = compiled_result[0] + + assert torch.equal(eager_out, compiled_out), ( + f"Compiled output differs from eager. Max abs diff: " + f"{(eager_out - compiled_out).abs().max().item():.6e}" + ) + + +# --------------------------------------------------------------------------- +# Init guard (TP > 1 rejection) +# --------------------------------------------------------------------------- + + +class TestMXFP4LinearGuard(PrimusUT): + """Test that MXFP4 linear layers reject invalid configurations.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state, monkeypatch): + dummy_args = SimpleNamespace( + rank=0, + world_size=1, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + offload=False, + offload_ops=[], + patch_primus_pipeline=False, + pp_algorithm=None, + patch_zero_bubble=False, + enable_zero_bubble=False, + rampup_batch_size=None, + global_batch_size=1, + micro_batch_size=1, + data_parallel_size=1, + decrease_batch_size_if_needed=False, + ) + import megatron.training.global_vars as gvars + + monkeypatch.setattr(gvars, "_GLOBAL_ARGS", dummy_args) + + _pin_fp4_aiter(monkeypatch) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_column_parallel_rejects_tp_gt_1(self): + from primus.backends.megatron.core.extensions.primus_turbo_mxfp4_local import ( + MXFP4ColumnParallelLinear, + ) + + config = _make_mxfp4_transformer_config(tensor_model_parallel_size=2) + with pytest.raises(ValueError, match="tensor_model_parallel_size=1"): + MXFP4ColumnParallelLinear( + input_size=256, + output_size=512, + config=config, + init_method=_init_method(), + bias=False, + gather_output=False, + skip_bias_add=False, + is_expert=False, + ) From 214c203bf5b28bb0bb02e461c3972b8670ca832b Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Mon, 13 Jul 2026 10:41:53 +0300 Subject: [PATCH 021/127] feat(flux): delayed fp8 scaling + TE DPA prologue patches (#813) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/turbo` — review after it. One of the parents of the trainers PR. ## What this changes The delayed-fp8-scaling patch set plus the consolidated Transformer-Engine dot-product-attention (DPA) prologue patch. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); its fp8 unit tests are green on the current CI pin (no turbo-bump dependency). Builds on `feat/flux/turbo`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "delayed_fp8 or fused_delayed"`. Validated locally on an AMD GPU container: 24 passed. ## Files 5 (delayed-fp8 patches, TE DPA prologue patch + tests). --------- Co-authored-by: Flux Split Trial Co-authored-by: luiza-amd Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- .../patches/delayed_fp8_scaling_patches.py | 342 +++++++++++++++++ .../dpa_consolidated_prologue_patches.py | 305 ++++++++++++++++ .../megatron/patches/turbo/fp8_patches.py | 6 +- .../diffusion/test_delayed_fp8_triton_op.py | 237 ++++++++++++ .../test_fused_delayed_scale_update.py | 343 ++++++++++++++++++ 5 files changed, 1231 insertions(+), 2 deletions(-) create mode 100644 primus/backends/megatron/patches/delayed_fp8_scaling_patches.py create mode 100644 primus/backends/megatron/patches/te_patches/dpa_consolidated_prologue_patches.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py diff --git a/primus/backends/megatron/patches/delayed_fp8_scaling_patches.py b/primus/backends/megatron/patches/delayed_fp8_scaling_patches.py new file mode 100644 index 000000000..e230ad2de --- /dev/null +++ b/primus/backends/megatron/patches/delayed_fp8_scaling_patches.py @@ -0,0 +1,342 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Megatron train_step patch for delayed FP8 scaling updates and preamble +optimization. + +Patch 1 -- Delayed FP8 scale update: + Wraps train_step to call the scale update preamble on every FP8 delayed + module before the original train_step runs. For most_recent + + history_len=1 (production config), a fast path stages per-module amaxes + into the registry-batched ``staged_amaxes_3n`` / ``scales_3n`` tensors, + computes new scales with a few fused ops, and scatters them back to the + per-module scalar ``m.scale_*`` buffers. Per-module scalars (rather + than views into a shared ``(N,)`` tensor) are required so that each + module's buffer storage is independent for torch.compile version + tracking. + +Patch 2 -- Grad-zero stream overlap + data HtoD prefetch: + Dispatches DDP grad-buffer zeroing (grad_data.zero_()) on a secondary + CUDA stream so it overlaps with HtoD data transfer, saving ~4.6 ms/iter. + Additionally, prefetches the next data batch to GPU on a dedicated HtoD + stream. The prefetch iterator is injected at the forward_backward_func + level (not by replacing data_iterator in train_step) to preserve + compatibility with Megatron's RerunDataIterator type assertion. +""" + +import torch + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + +# Module-level handle so the MLPerf warmup hook (or any external code) can +# reach the CudaPrefetchIterator state owned by +# ``patch_grad_zero_and_data_prefetch`` (which would otherwise live solely +# in that patch's closure and be unreachable from outside). +# +# The warmup hook needs this so it can invalidate the cached prefetch +# iterator at the end of the warmup epilogue. Without that invalidation, +# the prefetcher stays bound to the synthetic-data iterator that was passed +# in for warmup step 1, and every subsequent real training step silently +# reads from the cycling synthetic dataset (because +# ``MegatronDataloaderWrapper`` is cyclic and never raises ``StopIteration``). +_PREFETCH_HANDLE: dict = {"state": None} + +# Shared state for async amax allreduce between Patch 1 (scale update) and +# Patch 2 (fwd/bwd wrapper). Patch 2 launches the async allreduce after +# forward_backward_func returns; Patch 1 waits on the handle at the start +# of the next train_step before computing new scales. +# +# Realistic overlap window: the time between forward_backward_func returning +# (in step N) and the Patch 1 wait/compute at the start of step N+1. In +# practice that's roughly ``optimizer.step + grad_zero`` -- not the full +# train_step. The overlap is still useful but smaller than the all-reduce +# cost in most configs, so this is best-effort latency hiding rather than a +# free win. +_ASYNC_AMAX_HANDLE: dict = {"handle": None, "registry": None} + + +def _reset_async_amax_state(): + """Drop any pending async amax allreduce handle. + + Idempotent best-effort cleanup invoked from both call sites in this + module (Patch 1's wait path and Patch 2's launch path). Leaving a + stale handle in ``_ASYNC_AMAX_HANDLE["handle"]`` would cause the next + train_step to try to wait on an already-consumed or never-launched + work object and either deadlock or raise. + """ + _ASYNC_AMAX_HANDLE["handle"] = None + + +def get_prefetch_state(): + """Return the closure-shared ``_prefetch_state`` dict, or ``None`` if + ``patch_grad_zero_and_data_prefetch`` has not been installed yet.""" + return _PREFETCH_HANDLE.get("state") + + +def reset_prefetch_state(): + """Drop the cached ``CudaPrefetchIterator`` so the next ``train_step`` + rebuilds it around its current ``data_iterator`` argument. + + Required after MLPerf warmup, because warmup step 1 is the first call + into ``_patched_train_step``, so the prefetcher gets bound to the + synthetic iterator. Without invalidation, every real training step + afterwards substitutes the cached prefetcher (still wrapping synthetic + data) for the real ``data_iterator`` argument inside the + ``_synced_prefetch_fwd_bwd`` wrapper. + + Returns the evicted iterator (or ``None`` if no cache existed) so + callers can log what was dropped. + """ + state = _PREFETCH_HANDLE.get("state") + if state is None: + return None + return state.pop("iter", None) + + +def _needs_delayed_scaling(ctx: PatchContext) -> bool: + # Imported lazily (not at module top-level) so that importing this patch + # module never depends on ``megatron`` being importable. A top-level + # ``from megatron...`` here can crash the patches package auto-import if + # megatron is momentarily unavailable (e.g. mid circular-import), which + # then leaves the package in a half-initialized state. + from megatron.core.enums import Fp8Recipe + + args = get_args(ctx) + if args is None or not bool(getattr(args, "fp8", False)): + return False + return ( + getattr(args, "fp8_scaling_strategy", "dynamic") == "delayed" + or getattr(args, "fp8_recipe", None) == Fp8Recipe.delayed + ) and not getattr(args, "disable_delayed_scaling_patches", False) + + +@register_patch( + "megatron.fp8.delayed_scaling_update", + backend="megatron", + phase="before_train", + description="Wrap train_step to update delayed FP8 scales before each step.", + priority=40, + condition=_needs_delayed_scaling, +) +def patch_delayed_fp8_update(ctx: PatchContext): + import megatron.training.training as megatron_training + + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _DelayedScalingRegistry, + _fast_update_scales, + _fast_update_scales_with_history, + _wait_and_compute_scales, + ) + from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched + + _PATCH_KEY = "megatron.fp8.delayed_scaling_update" + if is_patched(megatron_training, _PATCH_KEY): + log_rank_0("[Patch:delayed_scaling_update] Already applied; skipping re-wrap.") + return + + _original_train_step = megatron_training.train_step + _cached_delayed_modules = [] + _registry = None + + def _patched_train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=None, + ): + nonlocal _registry + if not _cached_delayed_modules: + _cached_delayed_modules.extend( + m + for model_chunk in model + for m in model_chunk.modules() + if getattr(m, "_use_delayed_scaling", False) + ) + if _cached_delayed_modules: + if _registry is None: + _registry = _DelayedScalingRegistry(_cached_delayed_modules) + _ASYNC_AMAX_HANDLE["registry"] = _registry + + pending_handle = _ASYNC_AMAX_HANDLE.get("handle") + if pending_handle is not None: + try: + _wait_and_compute_scales(_registry, pending_handle) + finally: + # Always drop the handle, even if wait / compute raised: + # leaving a consumed handle here would cause the next + # train_step to try to wait on it again. + _reset_async_amax_state() + else: + if _registry.algo == "most_recent" and _registry.history_len == 1: + _fast_update_scales(_registry) + else: + _fast_update_scales_with_history(_registry) + return _original_train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, + ) + + megatron_training.train_step = _patched_train_step + mark_patched(megatron_training, _PATCH_KEY) + log_rank_0( + "[Patch:delayed_scaling_update] " + "Wrapped train_step with async amax allreduce support. " + "First step uses synchronous fallback; subsequent steps wait on " + "async handle launched by post-fwd_bwd hook." + ) + + +# --------------------------------------------------------------------------- +# Patch 2: Grad-zero stream overlap + data HtoD prefetch +# --------------------------------------------------------------------------- + + +@register_patch( + "megatron.grad_zero_and_data_prefetch", + backend="megatron", + phase="before_train", + description="Overlap grad buffer zeroing and data HtoD transfer via secondary CUDA streams.", + priority=41, + condition=_needs_delayed_scaling, +) +def patch_grad_zero_and_data_prefetch(ctx: PatchContext): + args = get_args(ctx) + if getattr(args, "reuse_grad_buf_for_mxfp8_param_ag", False): + log_rank_0( + "[Patch:grad_zero_and_data_prefetch] SKIPPED — " + "reuse_grad_buf_for_mxfp8_param_ag is set (shared param/grad buffer)." + ) + return + + import megatron.training.training as megatron_training + from megatron.core.distributed import DistributedDataParallel + + from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched + + _PATCH_KEY = "megatron.grad_zero_and_data_prefetch" + if is_patched(megatron_training, _PATCH_KEY): + log_rank_0("[Patch:grad_zero_and_data_prefetch] Already applied; skipping re-wrap.") + return + + tp_size = getattr(args, "tensor_model_parallel_size", 1) + + _original_train_step = megatron_training.train_step + _zero_stream = torch.cuda.Stream() + _prefetch_state: dict = {} + # Expose the closure-local prefetch state so the MLPerf warmup hook can + # invalidate the cached prefetch iterator at the end of warmup; see + # ``reset_prefetch_state`` above for the rationale. + _PREFETCH_HANDLE["state"] = _prefetch_state + + def _stream_zero_grad_buffer(self): + """CPU metadata on main thread; GPU grad_data.zero_() on secondary stream.""" + if getattr(self.config, "cuda_graph_impl", "none") != "transformer_engine": + for param in self.params_with_grad: + param.grad_added_to_main_grad = False + with torch.cuda.stream(_zero_stream): + for buffer in self.buffers + self.expert_parallel_buffers: + buffer.reset() + for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups: + bucket_group.reset() + + DistributedDataParallel.zero_grad_buffer = _stream_zero_grad_buffer + + def _patched_train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=None, + ): + if "iter" not in _prefetch_state and tp_size == 1: + if isinstance(data_iterator, (list, tuple)): + # Virtual pipeline parallel passes a list of per-chunk iterators, + # which the single-stream prefetcher cannot wrap. Skip prefetch + # (a pure HtoD-overlap optimization) and let the original iterator + # flow through unchanged. + if not _prefetch_state.get("vpp_skip_logged"): + log_rank_0( + "[Patch:grad_zero_and_data_prefetch] " + "Skipping CudaPrefetchIterator: data_iterator is a list " + "(virtual pipeline parallel)." + ) + _prefetch_state["vpp_skip_logged"] = True + else: + from primus.backends.megatron.data.cuda_prefetch import ( + CudaPrefetchIterator, + ) + + compute_dtype = torch.bfloat16 if getattr(args, "bf16", False) else torch.float16 + _prefetch_state["iter"] = CudaPrefetchIterator( + data_iterator, + compute_dtype=compute_dtype, + ) + log_rank_0( + "[Patch:grad_zero_and_data_prefetch] " + f"Created CudaPrefetchIterator (dtype={compute_dtype})." + ) + + _pf = _prefetch_state.get("iter") + + def _synced_prefetch_fwd_bwd(*fwd_args, **fwd_kwargs): + torch.cuda.current_stream().wait_stream(_zero_stream) + if _pf is not None: + fwd_kwargs["data_iterator"] = _pf + result = forward_backward_func(*fwd_args, **fwd_kwargs) + + registry = _ASYNC_AMAX_HANDLE.get("registry") + if registry is not None: + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _stage_and_launch_async_allreduce, + ) + + try: + _ASYNC_AMAX_HANDLE["handle"] = _stage_and_launch_async_allreduce(registry) + except Exception: + # Launch failed: clear so the next train_step takes the + # synchronous fallback path instead of trying to wait on + # a non-existent handle. + _reset_async_amax_state() + raise + + return result + + return _original_train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + _synced_prefetch_fwd_bwd, + iteration=iteration, + ) + + megatron_training.train_step = _patched_train_step + mark_patched(megatron_training, _PATCH_KEY) + log_rank_0( + "[Patch:grad_zero_and_data_prefetch] " + "DDP grad_data.zero_() on secondary stream; " + "data HtoD prefetch on secondary stream; " + "secondary streams synced before forward_backward_func; " + "async amax allreduce launched after forward_backward_func " + "(awaited at the start of the next train_step)." + ) diff --git a/primus/backends/megatron/patches/te_patches/dpa_consolidated_prologue_patches.py b/primus/backends/megatron/patches/te_patches/dpa_consolidated_prologue_patches.py new file mode 100644 index 000000000..168457a0a --- /dev/null +++ b/primus/backends/megatron/patches/te_patches/dpa_consolidated_prologue_patches.py @@ -0,0 +1,305 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +DPA Consolidated Prologue Patch + +Separates the eager setup (FP8 metadata, backend selection, cu_seqlens, etc.) +from the FusedAttnFunc kernel call in DotProductAttention, so that the kernel +call can live inside a torch.compile compiled graph rather than running eagerly. + +Graph-break analysis (per DPA call): + + Before: + [compiled A] --break--> [DPA runs fully eager, kernel included] --break--> [compiled B] + + After (attention_dropout == 0, the Flux default - fast path): + [compiled A] --break--> [eager prologue] --break--> [compiled post-ops (kernel inside)] + + After (attention_dropout > 0 - safe path): + [compiled A] --break--> [eager prologue] --break--> [eager fork + kernel] --break--> [compiled post-ops] + +Break counts: 2 for the fast path, 3 for the safe path. The safe path runs +the FusedAttn kernel inside the model-parallel-rng fork context so the +kernel's philox state targets the correct generator. In the fast path the +kernel doesn't consume RNG (dropout is 0), so the fork wrapper is omitted +and the kernel stays inside the compiled graph alongside the post-attention +ops (output reshape, projection, skip connections). + +Mechanism: + FusedAttnFunc.forward is monkey-patched with a thread-local "capture mode". + When capture mode is active the patched forward stores its arguments in + thread-local storage and raises ``_CaptureComplete``, aborting the DPA + forward after all setup is done. The captured args are then passed to the + *real* ``FusedAttnFunc.apply`` inside the compiled graph. +""" + +import threading +from typing import Optional + +import torch +from torch import Tensor + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + +# Positional index of dropout_p inside FusedAttnFunc.forward's *args (i.e. +# inside our captured _capture_tls.captured_args tuple, which excludes ctx). +# Verified against transformer_engine.pytorch.attention.dot_product_attention +# .backends.FusedAttnFunc.forward as of TE 2.12. If a future TE upgrade +# reshuffles the positional layout, _new_te_dpa_forward falls back to the +# always-fork (safe) path rather than silently routing to the dropout=0 fast +# path with the wrong field. +_DROPOUT_P_IDX = 14 + +# --------------------------------------------------------------------------- +# Thread-local capture state +# --------------------------------------------------------------------------- +_capture_tls = threading.local() + + +class _CaptureComplete(Exception): + """Raised inside the patched FusedAttnFunc.forward to abort after capture.""" + + +# --------------------------------------------------------------------------- +# Patch registration +# --------------------------------------------------------------------------- +@register_patch( + "megatron.te.dpa_consolidated_prologue", + backend="megatron", + phase="before_train", + priority=55, + description=( + "Consolidated DPA prologue: separates eager attention setup from " + "FusedAttnFunc kernel call to reduce torch.compile graph breaks" + ), + condition=lambda ctx: ( + ( + ( + getattr(get_args(ctx), "torch_compile", None) is not None + and getattr(get_args(ctx).torch_compile, "enable", False) + ) + # Align with torch_compile_patches.py, which also honors the flat + # enable_torch_compile flag; otherwise this prologue optimization is + # silently skipped for configs that compile via enable_torch_compile. + or getattr(get_args(ctx), "enable_torch_compile", False) + ) + and not getattr(get_args(ctx), "disable_dpa_prologue_patch", False) + ), +) +def patch_dpa_consolidated_prologue(ctx: PatchContext): + """Replace TEDotProductAttention.forward with a two-phase version. + + Phase 1 – eager prologue (``@torch._dynamo.disable``): + Runs the full DPA + FusedAttention setup pipeline to compute all + arguments for ``FusedAttnFunc.apply``, without running the kernel. + + Phase 2 – compiled kernel call: + ``FusedAttnFunc.apply`` (marked ``allow_in_graph``) runs inside the + torch.compile compiled graph, enabling fusion with pre/post-attention + operations. + """ + from megatron.core.extensions.transformer_engine import TEDotProductAttention + from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + FusedAttnFunc, + ) + + # ---- Idempotency guard ----------------------------------------------- + # Re-running this patch would capture the already-patched forward as + # "_orig_*", double-wrapping the capture logic. Guard against it. + if getattr(TEDotProductAttention, "_primus_dpa_prologue_patched", False): + log_rank_0("[Patch:megatron.te.dpa_consolidated_prologue] already applied; skipping re-patch") + return + + # ---- TE version check ------------------------------------------------ + # _DROPOUT_P_IDX (and the captured-args layout) is pinned to TE 2.12. On a + # different TE version the layout may drift; the fast-path detection + # already falls back to the always-fork safe path on mismatch, but warn so + # the perf regression is visible rather than silent. + try: + import transformer_engine + + _te_version = getattr(transformer_engine, "__version__", None) + except Exception: + _te_version = None + if _te_version is not None and not str(_te_version).startswith("2.12"): + log_rank_0( + "[Patch:megatron.te.dpa_consolidated_prologue] WARNING: " + f"_DROPOUT_P_IDX={_DROPOUT_P_IDX} was verified against TE 2.12 but " + f"detected transformer_engine {_te_version}. The dropout=0 fast path " + "will conservatively fall back to the safe rng-fork path if the " + "captured-args layout differs." + ) + + # ---- Mark FusedAttnFunc as graph-safe -------------------------------- + torch._dynamo.allow_in_graph(FusedAttnFunc) + + # ---- Patch FusedAttnFunc.forward for capture mode -------------------- + _orig_fused_attn_func_fwd = FusedAttnFunc.forward + + @staticmethod + def _capturing_fused_attn_func_fwd(ctx, *args): # noqa: N805 – autograd ctx + if getattr(_capture_tls, "capture_mode", False): + _capture_tls.captured_args = args + raise _CaptureComplete() + return _orig_fused_attn_func_fwd(ctx, *args) + + FusedAttnFunc.forward = _capturing_fused_attn_func_fwd + + # ---- Save original TEDotProductAttention.forward --------------------- + _orig_te_dpa_forward = TEDotProductAttention.forward + + # ---- Eager prologue -------------------------------------------------- + @torch._dynamo.disable + def _dpa_eager_prologue( + te_dpa, + query, + key, + value, + attention_mask, + attn_mask_type, + attention_bias, + packed_seq_params, + num_splits, + ): + """Run the full DPA + FusedAttention setup eagerly. + + Returns + ------- + captured_args : tuple | None + Arguments for ``FusedAttnFunc.apply`` if fused backend was chosen. + fallback_result : Tensor | None + Final output tensor if a non-fused backend ran to completion. + """ + _capture_tls.capture_mode = True + _capture_tls.captured_args = None + fallback_result = None + + try: + fallback_result = _orig_te_dpa_forward( + te_dpa, + query, + key, + value, + attention_mask, + attn_mask_type=attn_mask_type, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + num_splits=num_splits, + ) + except _CaptureComplete: + pass # expected – fused attention args captured + finally: + _capture_tls.capture_mode = False + + return _capture_tls.captured_args, fallback_result + + # ---- RNG-correct kernel execution ------------------------------------ + from megatron.core.tensor_parallel.random import get_cuda_rng_tracker + + @torch._dynamo.disable + def _fused_attn_with_rng_fork(captured_args): + """Run FusedAttnFunc.apply inside the model-parallel-rng fork. + + The eager prologue's _CaptureComplete exception unwinds TE's + internal fork context before the kernel runs, so rng_gen=None + causes the default CUDA generator to be advanced instead of + model-parallel-rng. This wrapper re-enters the fork so the + kernel's philox_cuda_state call targets the correct generator. + """ + tracker = get_cuda_rng_tracker() + if tracker.is_initialized(): + states = tracker.get_states() + if "model-parallel-rng" in states: + with tracker.fork("model-parallel-rng"): + return FusedAttnFunc.apply(*captured_args) + return FusedAttnFunc.apply(*captured_args) + + # ---- Max-logit bookkeeping (qk_clip) --------------------------------- + @torch._dynamo.disable + def _update_max_logit_stats(te_dpa, batch_max_logit): + if hasattr(te_dpa, "current_max_attn_logits"): + if te_dpa.current_max_attn_logits is None: + te_dpa.current_max_attn_logits = batch_max_logit + else: + te_dpa.current_max_attn_logits = torch.max(te_dpa.current_max_attn_logits, batch_max_logit) + + # ---- Replacement forward --------------------------------------------- + def _new_te_dpa_forward( + self, + query: Tensor, + key: Tensor, + value: Tensor, + attention_mask: Optional[Tensor], + attn_mask_type=None, + attention_bias: Optional[Tensor] = None, + packed_seq_params=None, + num_splits: Optional[int] = None, + ) -> Tensor: + """TEDotProductAttention.forward with consolidated prologue. + + The eager prologue handles all DPA setup (FP8 metadata, backend + selection, cu_seqlens computation, etc.). Then ``FusedAttnFunc.apply`` + runs inside the compiled graph. + + For non-fused backends (flash / unfused), the original forward runs + fully inside the prologue and the result is returned directly. + """ + captured_args, fallback_result = _dpa_eager_prologue( + self, + query, + key, + value, + attention_mask, + attn_mask_type, + attention_bias, + packed_seq_params, + num_splits, + ) + + if captured_args is not None: + # Fast path: when dropout_p == 0 the FusedAttn kernel doesn't + # consume RNG, so we can skip the eager rng-fork wrapper and + # call FusedAttnFunc.apply directly inside the compiled graph + # (it's already marked allow_in_graph above). This drops a + # graph break in the common Flux case (attention_dropout=0). + # + # Safe path: anything that isn't a clean numeric zero (including + # IndexError if a TE upgrade changes captured_args' layout, or + # an unexpected tensor / object at the dropout slot) falls back + # to the always-fork path so correctness is never compromised + # by a layout drift. + try: + _dropout_p = captured_args[_DROPOUT_P_IDX] + _use_fast_path = isinstance(_dropout_p, (int, float)) and _dropout_p == 0 + except (IndexError, TypeError): + _use_fast_path = False + + if _use_fast_path: + result = FusedAttnFunc.apply(*captured_args) + else: + result = _fused_attn_with_rng_fork(captured_args) + return_max_logit = captured_args[-1] + + if return_max_logit: + core_attn_out = result[0].view(*result[0].shape[:-2], -1) + _update_max_logit_stats(self, result[1]) + return core_attn_out + + return result.view(*result.shape[:-2], -1) + + return fallback_result + + # ---- Apply the monkey-patch ------------------------------------------ + TEDotProductAttention.forward = _new_te_dpa_forward + TEDotProductAttention._primus_dpa_prologue_patched = True + + log_rank_0("[Patch:megatron.te.dpa_consolidated_prologue] " "Applied allow_in_graph to FusedAttnFunc") + log_rank_0( + "[Patch:megatron.te.dpa_consolidated_prologue] " + "Replaced TEDotProductAttention.forward with consolidated prologue" + ) diff --git a/primus/backends/megatron/patches/turbo/fp8_patches.py b/primus/backends/megatron/patches/turbo/fp8_patches.py index 7a077f631..4cbb97c3b 100644 --- a/primus/backends/megatron/patches/turbo/fp8_patches.py +++ b/primus/backends/megatron/patches/turbo/fp8_patches.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -28,7 +28,9 @@ def _is_fp8_can_patch(ctx: PatchContext) -> bool: backend="megatron", phase="before_train", description="Override Megatron get_fp8_context to use Primus implementation when fp8 is enabled", - condition=_is_fp8_can_patch, + condition=lambda ctx: ( + _is_fp8_can_patch(ctx) and not getattr(get_args(ctx), "disable_fp8_context_patches", False) + ), ) def patch_fp8_context(ctx: PatchContext): """ diff --git a/tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py b/tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py new file mode 100644 index 000000000..ba815e97f --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py @@ -0,0 +1,237 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Validate the @triton_op cast_transpose_fp8_triton kernel. + +Tests: + 1. Correctness: FP8 cast + transpose + amax match reference + 2. GEMM interaction: output feeds into torch._scaled_mm without regression + 3. torch.compile: triton_op is traceable (no graph breaks) + 4. Autograd: gradients flow through a minimal FP8 linear autograd.Function + +Run: + python -m pytest tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py -v +or standalone: + python tests/unit_tests/backends/megatron/diffusion/test_delayed_fp8_triton_op.py +""" + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus_turbo.pytorch.core.backend import BackendType +from primus_turbo.pytorch.core.low_precision import ( + ScalingGranularity, + float8_e4m3, + float8_e5m2, +) + +from primus.backends.megatron.core.extensions.fp8_cast_kernels_triton import ( + cast_transpose_fp8_triton, +) +from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + DelayedFP8LinearTensorwiseFunction, +) + +DEVICE = "cuda:0" +DTYPE = torch.bfloat16 +FP8_DTYPE = float8_e4m3 +FP8_MAX = torch.finfo(FP8_DTYPE).max +FP8_BWD_DTYPE = float8_e5m2 +FP8_BWD_MAX = torch.finfo(FP8_BWD_DTYPE).max + + +# ----------------------------------------------------------------------- +# Test 1: Correctness +# ----------------------------------------------------------------------- +@pytest.mark.parametrize( + "shape", + [ + (4096, 3072), + (16384, 3072), + (8192, 8192), + ], +) +def test_correctness(shape): + M, N = shape + x = torch.randn(M, N, dtype=DTYPE, device=DEVICE) + scale = torch.tensor(FP8_MAX / x.abs().amax().item(), dtype=torch.float32, device=DEVICE) + amax_buf = torch.zeros((), dtype=torch.float32, device=DEVICE) + + cast_out, trans_out, scale_inv = cast_transpose_fp8_triton(x, FP8_DTYPE, scale, amax_buf) + + assert cast_out.shape == (M, N), f"Cast shape mismatch: {cast_out.shape}" + assert cast_out.dtype == FP8_DTYPE + assert trans_out.shape == (N, M), f"Trans shape mismatch: {trans_out.shape}" + assert trans_out.dtype == FP8_DTYPE + assert trans_out.is_contiguous(), "Transpose output must be contiguous" + assert scale_inv.shape == (), f"Scale inv shape mismatch: {scale_inv.shape}" + + ref_scaled = (x.float() * scale.item()).clamp(-FP8_MAX, FP8_MAX) + ref_fp8 = ref_scaled.to(FP8_DTYPE) + assert torch.equal(cast_out, ref_fp8), f"Cast output mismatch for {shape}" + + ref_trans = ref_fp8.t().contiguous() + assert torch.equal(trans_out, ref_trans), f"Transpose output mismatch for {shape}" + + expected_scale_inv = 1.0 / scale.item() + assert ( + abs(scale_inv.item() - expected_scale_inv) < 1e-5 + ), f"scale_inv {scale_inv.item()} != {expected_scale_inv}" + + expected_amax = x.float().abs().amax().item() + assert ( + abs(amax_buf.item() - expected_amax) / max(expected_amax, 1e-8) < 1e-3 + ), f"amax {amax_buf.item()} != {expected_amax}" + + +# ----------------------------------------------------------------------- +# Test 2: GEMM interaction (the critical regression test) +# ----------------------------------------------------------------------- +@pytest.mark.parametrize( + "shape", + [ + (16384, 3072), + (4096, 12288), + ], +) +def test_gemm_interaction(shape): + M, N = shape + x = torch.randn(M, N, dtype=DTYPE, device=DEVICE) + scale = torch.tensor(FP8_MAX / x.abs().amax().item(), dtype=torch.float32, device=DEVICE) + amax_buf = torch.zeros((), dtype=torch.float32, device=DEVICE) + + cast_out, _, scale_inv = cast_transpose_fp8_triton(x, FP8_DTYPE, scale, amax_buf) + + w = torch.randn(N, N, dtype=DTYPE, device=DEVICE).to(FP8_DTYPE).t() + w_scale = torch.tensor(1.0, dtype=torch.float32, device=DEVICE) + + result = torch._scaled_mm(cast_out, w, out_dtype=DTYPE, scale_a=scale_inv, scale_b=w_scale) + assert result.shape == (M, N) + assert result.dtype == DTYPE + assert not result.isnan().any(), "GEMM produced NaN" + assert not result.isinf().any(), "GEMM produced Inf" + + +# ----------------------------------------------------------------------- +# Test 3: torch.compile traceability +# ----------------------------------------------------------------------- +def test_torch_compile(): + M, N = 4096, 3072 + x = torch.randn(M, N, dtype=DTYPE, device=DEVICE) + scale = torch.tensor(FP8_MAX / x.abs().amax().item(), dtype=torch.float32, device=DEVICE) + amax_buf = torch.zeros((), dtype=torch.float32, device=DEVICE) + + @torch.compile(fullgraph=True) + def fn(x, scale, amax_buf): + return cast_transpose_fp8_triton(x, FP8_DTYPE, scale, amax_buf) + + cast_out, trans_out, scale_inv = fn(x, scale, amax_buf) + + assert cast_out.shape == (M, N) + assert trans_out.shape == (N, M) + assert cast_out.dtype == FP8_DTYPE + + # Verify match with eager + amax_eager = torch.zeros((), dtype=torch.float32, device=DEVICE) + cast_eager, trans_eager, si_eager = cast_transpose_fp8_triton(x, FP8_DTYPE, scale, amax_eager) + assert torch.equal(cast_out, cast_eager), "Compiled output != eager output" + assert torch.equal(trans_out, trans_eager), "Compiled transpose != eager transpose" + + +# ----------------------------------------------------------------------- +# Test 4: Direct test of the *production* delayed-tensorwise FP8 autograd +# Function (DelayedFP8LinearTensorwiseFunction), exercising both the native +# (force_nt=False -> dgrad=NN/wgrad=TN) and forced-NT (force_nt=True) arms. +# +# This replaces the previous in-file ``_MinimalFP8Linear`` replica: it calls the +# shipped Function directly so a regression in its quantization, GEMM layout, or +# fused amax capture is actually caught. +# ----------------------------------------------------------------------- +def _rel_err(actual: torch.Tensor, ref: torch.Tensor) -> float: + ref = ref.float() + return ((actual.float() - ref).norm() / ref.norm().clamp_min(1e-12)).item() + + +@pytest.mark.parametrize("force_nt", [True, False]) +def test_delayed_fp8_linear_function(force_nt): + M, K, N = 256, 512, 256 + + torch.manual_seed(0) + inp = torch.randn(M, K, dtype=DTYPE, device=DEVICE, requires_grad=True) + weight = torch.randn(N, K, dtype=DTYPE, device=DEVICE, requires_grad=True) + grad_output = torch.randn(M, N, dtype=DTYPE, device=DEVICE) + + # Delayed tensorwise scales (one scalar per track), chosen to fill the FP8 + # range the way the production warmup staging would. + scale_input = torch.tensor(FP8_MAX / inp.detach().abs().amax().item(), dtype=torch.float32, device=DEVICE) + scale_weight = torch.tensor( + FP8_MAX / weight.detach().abs().amax().item(), dtype=torch.float32, device=DEVICE + ) + scale_grad = torch.tensor( + FP8_BWD_MAX / grad_output.abs().amax().item(), dtype=torch.float32, device=DEVICE + ) + + staged_input_amax = torch.zeros((), dtype=torch.float32, device=DEVICE) + staged_weight_amax = torch.zeros((), dtype=torch.float32, device=DEVICE) + staged_grad_amax = torch.zeros((), dtype=torch.float32, device=DEVICE) + + gran_value = ScalingGranularity.TENSORWISE.value + backend_value = BackendType.HIPBLASLT.value + + result = DelayedFP8LinearTensorwiseFunction.apply( + inp, + weight, + scale_input, + scale_weight, + scale_grad, + staged_input_amax, + staged_weight_amax, + staged_grad_amax, + FP8_DTYPE, + FP8_BWD_DTYPE, + gran_value, + backend_value, + force_nt, + ) + output = result[0] + + # Forward: output ~= input @ weight.T (standard nn.Linear), FP8-tensorwise. + ref_out = inp.detach().float() @ weight.detach().float().t() + assert output.shape == (M, N) + assert torch.isfinite(output).all(), "FP8 forward produced non-finite values" + out_rel = _rel_err(output, ref_out) + assert out_rel < 0.1, f"FP8 forward too far from bf16 reference: rel_err={out_rel:.4f}" + + # Fused amax capture (the delayed-scaling contract): the staged buffers must + # be written with the *current* tensor amaxes during forward. + assert staged_input_amax.item() > 0, "staged_input_amax not captured in forward" + assert staged_weight_amax.item() > 0, "staged_weight_amax not captured in forward" + in_amax_rel = abs(staged_input_amax.item() - inp.detach().float().abs().amax().item()) / max( + inp.detach().float().abs().amax().item(), 1e-8 + ) + w_amax_rel = abs(staged_weight_amax.item() - weight.detach().float().abs().amax().item()) / max( + weight.detach().float().abs().amax().item(), 1e-8 + ) + assert in_amax_rel < 1e-2, f"staged input amax mismatch: rel={in_amax_rel:.4f}" + assert w_amax_rel < 1e-2, f"staged weight amax mismatch: rel={w_amax_rel:.4f}" + + # Backward through the production Function. + output.backward(grad_output) + assert inp.grad is not None and weight.grad is not None + assert torch.isfinite(inp.grad).all(), "grad_input non-finite" + assert torch.isfinite(weight.grad).all(), "grad_weight non-finite" + + ref_grad_input = grad_output.float() @ weight.detach().float() + ref_grad_weight = grad_output.float().t() @ inp.detach().float() + # e5m2 backward has only 2 mantissa bits, so use a looser bound than forward. + gi_rel = _rel_err(inp.grad, ref_grad_input) + gw_rel = _rel_err(weight.grad, ref_grad_weight) + assert gi_rel < 0.2, f"grad_input too far from reference: rel_err={gi_rel:.4f}" + assert gw_rel < 0.2, f"grad_weight too far from reference: rel_err={gw_rel:.4f}" + + assert staged_grad_amax.item() > 0, "staged_grad_amax not captured in backward" diff --git a/tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py b/tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py new file mode 100644 index 000000000..b7f14eda7 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py @@ -0,0 +1,343 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Validate the fused Triton kernel for delayed FP8 scale update. + +Tests: + 1. Correctness: fused kernel matches reference Python-loop implementation + 2. Edge cases: zero amaxes, NaN amaxes, very small amaxes + 3. History rollover: circular index wraps correctly + 4. Algorithm variants: 'max' vs 'most_recent' + +Run: + python -m pytest tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py -v +or standalone: + python tests/unit_tests/backends/megatron/diffusion/test_fused_delayed_scale_update.py +""" + +import pytest +import torch + +triton = pytest.importorskip("triton") +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU") + +DEVICE = "cuda:0" +FP8_FWD_MAX = torch.finfo(torch.float8_e4m3fn).max # 448.0 +FP8_BWD_MAX = torch.finfo(torch.float8_e5m2).max # 57344.0 +FP32_MAX = torch.finfo(torch.float32).max + + +def _reference_update(amax_history, staged_amaxes, scales, fp8_maxes, history_idx, use_max_algo): + """Pure-PyTorch reference matching the fused delayed-scale update logic.""" + T, N, H = amax_history.shape + for t in range(T): + for m in range(N): + new_amax = staged_amaxes[t, m].item() + amax_history[t, m, history_idx] = new_amax + + if use_max_algo: + amax = amax_history[t, m, :].max().item() + else: + amax = new_amax + + fp8_max = fp8_maxes[t].item() + old_scale = scales[t, m].item() + + if amax > 0.0 and amax == amax: # positive and not NaN + sf = fp8_max / max(amax, 1e-12) + sf = min(sf, FP32_MAX) + else: + sf = old_scale + + scales[t, m] = sf + + +def _run_fused_kernel(amax_history, staged_amaxes, scales, fp8_maxes, history_idx, use_max_algo): + """Call the actual fused Triton kernel.""" + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _fused_delayed_scale_update_kernel, + ) + + _, N, H = amax_history.shape + BLOCK_H = triton.next_power_of_2(H) if H <= 1024 else 1024 + + _fused_delayed_scale_update_kernel[(N, 3)]( + amax_history, + staged_amaxes, + scales, + fp8_maxes, + history_idx, + N=N, + H=H, + use_max_algo=use_max_algo, + BLOCK_H=BLOCK_H, + FP32_MAX=FP32_MAX, + ) + + +@pytest.mark.parametrize( + "N,H", + [ + (1, 1), + (4, 16), + (57, 1024), + (10, 128), + ], +) +@pytest.mark.parametrize("algo", ["max", "most_recent"]) +def test_correctness(N, H, algo): + """Fused kernel must produce identical scales and history as reference.""" + use_max_algo = algo == "max" + fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE) + + torch.manual_seed(42) + amax_history_ref = torch.rand(3, N, H, dtype=torch.float32, device=DEVICE) * 10.0 + staged_amaxes = torch.rand(3, N, dtype=torch.float32, device=DEVICE) * 5.0 + scales_ref = torch.ones(3, N, dtype=torch.float32, device=DEVICE) + history_idx = H // 3 + + amax_history_fused = amax_history_ref.clone() + scales_fused = scales_ref.clone() + + _reference_update(amax_history_ref, staged_amaxes, scales_ref, fp8_maxes, history_idx, use_max_algo) + _run_fused_kernel(amax_history_fused, staged_amaxes, scales_fused, fp8_maxes, history_idx, use_max_algo) + + torch.testing.assert_close( + amax_history_fused, + amax_history_ref, + atol=0, + rtol=0, + msg=f"amax_history mismatch for N={N}, H={H}, algo={algo}", + ) + torch.testing.assert_close( + scales_fused, scales_ref, atol=1e-5, rtol=1e-5, msg=f"scales mismatch for N={N}, H={H}, algo={algo}" + ) + + +def test_zero_amaxes(): + """When all staged amaxes are zero and history is zero, scales should stay at old_scale.""" + N, H = 8, 16 + fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE) + amax_history = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE) + staged_amaxes = torch.zeros(3, N, dtype=torch.float32, device=DEVICE) + old_scales = torch.full((3, N), 42.0, dtype=torch.float32, device=DEVICE) + scales = old_scales.clone() + + _run_fused_kernel(amax_history, staged_amaxes, scales, fp8_maxes, history_idx=0, use_max_algo=True) + + torch.testing.assert_close( + scales, old_scales, atol=0, rtol=0, msg="Scales should be unchanged when all amaxes are zero" + ) + + +def test_nan_amaxes(): + """NaN amaxes should leave scales unchanged (NaN guard).""" + N, H = 4, 8 + fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE) + amax_history = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE) + staged_amaxes = torch.full((3, N), float("nan"), dtype=torch.float32, device=DEVICE) + old_scales = torch.full((3, N), 7.0, dtype=torch.float32, device=DEVICE) + scales = old_scales.clone() + + _run_fused_kernel(amax_history, staged_amaxes, scales, fp8_maxes, history_idx=0, use_max_algo=True) + + torch.testing.assert_close( + scales, old_scales, atol=0, rtol=0, msg="Scales should be unchanged when amaxes are NaN" + ) + + +def test_very_small_amaxes(): + """Very small amaxes should produce large scales clamped to FP32_MAX.""" + N, H = 4, 8 + fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE) + amax_history = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE) + staged_amaxes = torch.full((3, N), 1e-40, dtype=torch.float32, device=DEVICE) + scales = torch.ones(3, N, dtype=torch.float32, device=DEVICE) + + _run_fused_kernel(amax_history, staged_amaxes, scales, fp8_maxes, history_idx=0, use_max_algo=False) + + assert not scales.isinf().any(), "Scales should not be inf (should be clamped)" + assert not scales.isnan().any(), "Scales should not be NaN" + assert (scales <= FP32_MAX).all(), "Scales should be <= FP32_MAX" + + +@pytest.mark.parametrize("H", [4, 16, 64]) +def test_history_rollover(H): + """Circular index should write to correct position across multiple steps.""" + N = 3 + fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE) + amax_history = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE) + scales = torch.ones(3, N, dtype=torch.float32, device=DEVICE) + + for step in range(H + 5): + idx = step % H + staged = torch.full((3, N), float(step + 1), dtype=torch.float32, device=DEVICE) + _run_fused_kernel(amax_history, staged, scales, fp8_maxes, history_idx=idx, use_max_algo=True) + + assert amax_history[0, 0, idx].item() == float( + step + 1 + ), f"History not written at idx={idx} on step={step}" + + +def test_multi_step_vs_reference(): + """Run 50 steps of both fused and reference, verify they stay in sync.""" + N, H = 10, 32 + use_max_algo = True + fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE) + + torch.manual_seed(123) + amax_history_ref = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE) + amax_history_fused = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE) + scales_ref = torch.ones(3, N, dtype=torch.float32, device=DEVICE) + scales_fused = torch.ones(3, N, dtype=torch.float32, device=DEVICE) + + for step in range(50): + idx = step % H + staged = torch.rand(3, N, dtype=torch.float32, device=DEVICE) * (10.0 + step) + + _reference_update(amax_history_ref, staged, scales_ref, fp8_maxes, idx, use_max_algo) + _run_fused_kernel(amax_history_fused, staged, scales_fused, fp8_maxes, idx, use_max_algo) + + torch.testing.assert_close( + amax_history_fused, amax_history_ref, atol=0, rtol=0, msg="amax_history diverged over 50 steps" + ) + torch.testing.assert_close( + scales_fused, scales_ref, atol=1e-5, rtol=1e-5, msg="scales diverged over 50 steps" + ) + + +def test_registry_integration(): + """Validate _fast_update_scales_with_history via _DelayedScalingRegistry.""" + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _DelayedScalingRegistry, + _fast_update_scales_with_history, + ) + + N, H = 8, 16 + + class _FakeModule(torch.nn.Module): + """Minimal stand-in for Float8*ParallelLinear with real tensors.""" + + def __init__(self, weight_data): + super().__init__() + self._fp8_fwd_max = FP8_FWD_MAX + self._fp8_bwd_max = FP8_BWD_MAX + self._amax_compute_algo = "max" + self._use_delayed_scaling = True + self._first_delayed_step = True + self._history_idx = 0 + self.weight = torch.nn.Parameter(weight_data) + self.register_buffer("amax_history_input", torch.zeros(H, dtype=torch.float32, device=DEVICE)) + self.register_buffer("amax_history_weight", torch.zeros(H, dtype=torch.float32, device=DEVICE)) + self.register_buffer("amax_history_grad", torch.zeros(H, dtype=torch.float32, device=DEVICE)) + self.register_buffer("scale_input", torch.tensor(1.0, dtype=torch.float32, device=DEVICE)) + self.register_buffer("scale_weight", torch.tensor(1.0, dtype=torch.float32, device=DEVICE)) + self.register_buffer("scale_grad", torch.tensor(1.0, dtype=torch.float32, device=DEVICE)) + self.register_buffer("staged_input_amax", torch.tensor(0.0, dtype=torch.float32, device=DEVICE)) + self.register_buffer("staged_weight_amax", torch.tensor(0.0, dtype=torch.float32, device=DEVICE)) + self.register_buffer("staged_grad_amax", torch.tensor(0.0, dtype=torch.float32, device=DEVICE)) + + torch.manual_seed(99) + weights = [torch.randn(64, 64, dtype=torch.bfloat16, device=DEVICE) for _ in range(N)] + modules = [_FakeModule(weights[i]) for i in range(N)] + + registry = _DelayedScalingRegistry(modules) + + # Staging is now done per-module: each m owns scalar buffers m.staged_input_amax, + # m.staged_grad_amax. _fast_update_scales_with_history stacks them into + # registry.staged_amaxes_3n[0/2] before launching the Triton kernel; scales + # land in registry.scales_3n and are scattered back to m.scale_*. + # Parallel pure-Python reference, driven by the EXACT amaxes the kernel used. + fp8_maxes = torch.tensor([FP8_FWD_MAX, FP8_FWD_MAX, FP8_BWD_MAX], dtype=torch.float32, device=DEVICE) + amax_history_ref = torch.zeros(3, N, H, dtype=torch.float32, device=DEVICE) + scales_ref = torch.ones(3, N, dtype=torch.float32, device=DEVICE) + history_idx_ref = 0 + use_max_algo = True # _FakeModule sets _amax_compute_algo = "max" + + torch.manual_seed(77) + for step in range(5): + for i in range(N): + modules[i].staged_input_amax.fill_(torch.rand(1, device=DEVICE).item() * 10) + modules[i].staged_grad_amax.fill_(torch.rand(1, device=DEVICE).item() * 5) + + _fast_update_scales_with_history(registry) + + # The fused kernel only reads registry.staged_amaxes_3n (rows input/weight/grad) + # and writes amax_history/scales, so reading it back yields the exact amaxes the + # kernel consumed -- including the first-step weight bootstrap from ||weight||_inf + # that is staged inside the call. Drive the reference with the same inputs and the + # same history index it used this step. + staged_used = registry.staged_amaxes_3n.clone() + _reference_update(amax_history_ref, staged_used, scales_ref, fp8_maxes, history_idx_ref, use_max_algo) + history_idx_ref = (history_idx_ref + 1) % H + + # Registry results must match the reference: history bit-for-bit, scales within fp32 tol. + torch.testing.assert_close( + registry.amax_history, + amax_history_ref, + atol=0, + rtol=0, + msg="registry amax_history diverged from reference loop", + ) + torch.testing.assert_close( + registry.scales_3n, + scales_ref, + atol=1e-5, + rtol=1e-5, + msg="registry scales diverged from reference loop", + ) + + # Per-module buffers must reflect the registry-batched results (rows input/weight/grad). + for i, m in enumerate(modules): + torch.testing.assert_close(m.scale_input, scales_ref[0, i], atol=1e-5, rtol=1e-5) + torch.testing.assert_close(m.scale_weight, scales_ref[1, i], atol=1e-5, rtol=1e-5) + torch.testing.assert_close(m.scale_grad, scales_ref[2, i], atol=1e-5, rtol=1e-5) + torch.testing.assert_close(m.amax_history_input, amax_history_ref[0, i], atol=0, rtol=0) + torch.testing.assert_close(m.amax_history_weight, amax_history_ref[1, i], atol=0, rtol=0) + torch.testing.assert_close(m.amax_history_grad, amax_history_ref[2, i], atol=0, rtol=0) + + # Sanity: scales finite and positive. + assert not registry.scales_3n.isnan().any(), "scales have NaN" + assert not registry.scales_3n.isinf().any(), "scales have Inf" + assert (registry.scales_3n[0] > 0).all(), "scale_input should be positive" + assert (registry.scales_3n[1] > 0).all(), "scale_weight should be positive" + + assert registry._history_idx == 5 % H, f"history_idx should be {5 % H}, got {registry._history_idx}" + for m in modules: + assert m._history_idx == registry._history_idx, "Per-module _history_idx not synced with registry" + + +if __name__ == "__main__": + print("=== Test 1: Correctness (max, various sizes) ===") + for N, H in [(1, 1), (4, 16), (57, 1024), (10, 128)]: + for algo in ["max", "most_recent"]: + test_correctness(N, H, algo) + print(f" N={N}, H={H}, algo={algo}: PASS") + + print("\n=== Test 2: Zero amaxes ===") + test_zero_amaxes() + print(" PASS") + + print("\n=== Test 3: NaN amaxes ===") + test_nan_amaxes() + print(" PASS") + + print("\n=== Test 4: Very small amaxes ===") + test_very_small_amaxes() + print(" PASS") + + print("\n=== Test 5: History rollover ===") + for H in [4, 16, 64]: + test_history_rollover(H) + print(f" H={H}: PASS") + + print("\n=== Test 6: Multi-step vs reference ===") + test_multi_step_vs_reference() + print(" PASS") + + print("\n=== Test 7: Registry integration ===") + test_registry_integration() + print(" PASS") + + print("\nAll tests passed!") From 68236e74b19a79f62590972e879078d307d3db78 Mon Sep 17 00:00:00 2001 From: Kailash Gogineni Date: Mon, 13 Jul 2026 17:46:45 -0700 Subject: [PATCH 022/127] Auto benchmark tool refinement (#872) * Refined the auto-benchmark tool to support Rock images. * Consolidated the source code into a single `metrics.py` file. * Unified the implementation so it works with both the Megatron and TorchTitan backends via Primus --- tools/auto_benchmark/metrics.py | 714 ++++++++++++++++++ tools/auto_benchmark/metrics_megatron.py | 212 ------ tools/auto_benchmark/metrics_torchtitan.py | 224 ------ .../run_primus_autobenchmark.sh | 432 +++++++---- 4 files changed, 1003 insertions(+), 579 deletions(-) create mode 100644 tools/auto_benchmark/metrics.py delete mode 100644 tools/auto_benchmark/metrics_megatron.py delete mode 100644 tools/auto_benchmark/metrics_torchtitan.py diff --git a/tools/auto_benchmark/metrics.py b/tools/auto_benchmark/metrics.py new file mode 100644 index 000000000..f15ab0973 --- /dev/null +++ b/tools/auto_benchmark/metrics.py @@ -0,0 +1,714 @@ +#!/usr/bin/env python3 + +import argparse +import csv +import os +import re +import shutil +from collections import defaultdict +from datetime import datetime +from statistics import mean + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +PRIMUS_ROOT = os.path.normpath(os.path.join(SCRIPT_DIR, "..", "..")) +RESULTS_DIR = os.path.join(SCRIPT_DIR, "results") + +WARMUP_SKIP = 5 + +ERROR_STATUS = "Error found - check log" +OK_STATUS = "OK" + +ANSI_ESCAPE_REGEX = re.compile(r"\x1b\[[0-9;]*m") +LOG_EXIT_CODE_REGEX = re.compile(r"primus launcher exited with code (\d+)", re.IGNORECASE) + +LOG_ERROR_PATTERNS = ( + re.compile(r"Traceback \(most recent call last\)", re.IGNORECASE), + re.compile(r"\[ERROR\]", re.IGNORECASE), + re.compile(r"\bRuntimeError:", re.IGNORECASE), + re.compile(r"\bOutOfMemoryError:", re.IGNORECASE), + re.compile(r"\bCUDA out of memory\b", re.IGNORECASE), + re.compile(r"\bHIP out of memory\b", re.IGNORECASE), + re.compile(r"\bNCCL error\b", re.IGNORECASE), + re.compile(r"\bfatal error\b", re.IGNORECASE), +) + +LOG_ERROR_EXCLUSIONS = ( + re.compile(r"error_injection", re.IGNORECASE), + re.compile(r"\[SKIP\].*Import failed", re.IGNORECASE), + re.compile(r"avoid ImportError", re.IGNORECASE), + re.compile(r"TORCH_NCCL_ASYNC_ERROR_HANDLING", re.IGNORECASE), + re.compile(r"destroy_process_group\(\)", re.IGNORECASE), +) + +RUN_LABEL_REGEX = re.compile(r"_run(\d+)\.log$", re.IGNORECASE) +LEGACY_TIMESTAMP_REGEX = re.compile(r"_(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.log$") +LEGACY_SUFFIX_REGEX = re.compile(r"_MI\d+X?_(.+)\.log$", re.IGNORECASE) + +ENV_DEFAULT_REGEX = re.compile(r"\$\{[A-Z0-9_]+:(\d+)\}") +PLAIN_INT_REGEX = re.compile(r"^\d+$") + +MEGATRON_BS_KEYS = ("micro_batch_size",) +MEGATRON_SEQ_KEYS = ("seq_length",) +MEGATRON_GBS_KEYS = ("global_batch_size",) + +TORCHTITAN_BS_KEYS = ("local_batch_size",) +TORCHTITAN_SEQ_KEYS = ("seq_len", "seq_length") +TORCHTITAN_GBS_KEYS = ("global_batch_size",) + +MEGATRON_NOTE_TEXT = """ +NOTE: +- Results are saved to results/metrics_megatron.csv (latest) and a timestamped snapshot. +- "Run" is run1, run2, ... per model and device, ordered oldest to newest. +- "BS", "Seq", and "GBS" come from the benchmark yaml when available. +- Megatron logs often print each iteration twice (Primus log forwarding); metrics + deduplicate by iteration number before averaging. +- Timing/throughput fields use the current value before "/" (e.g. 5896.1/5913.1 -> 5896.1). +- Warm-up: the first five iterations are excluded before averaging. +- "Status" is "Error found - check log" when the log contains errors or metrics could not be computed. +- Numeric fields may contain commas in logs; commas are removed before averaging. +""" + +TORCHTITAN_NOTE_TEXT = """ +NOTE: +- Results are saved to results/metrics_torchtitan.csv (latest) and a timestamped snapshot. +- "Run" is run1, run2, ... per model and device, ordered oldest to newest. +- "BS", "Seq", and "GBS" come from the benchmark yaml when available. +- "Steps" is the number of training steps used after dropping the first five warm-up steps. +- TPS and TFLOPS values may contain commas in logs; commas are removed before averaging. +- "Status" is "Error found - check log" when the log contains errors or metrics could not be computed. +""" + +MEGATRON_NUM = r"[\d,]+(?:\.\d+)?" +MEGATRON_METRIC_VALUE = rf"({MEGATRON_NUM})(?:\s*/\s*{MEGATRON_NUM})?" + +MEGATRON_ITERATION_REGEX = re.compile( + rf"iteration\s+(\d+)/\s*\d+.*?" + rf"elapsed time per iteration \(ms\):\s*{MEGATRON_METRIC_VALUE}.*?" + rf"throughput per GPU \(TFLOP/s/GPU\):\s*{MEGATRON_METRIC_VALUE}.*?" + rf"(?:tokens per GPU \(tokens/s/GPU\):\s*{MEGATRON_METRIC_VALUE}.*?)?" + rf"global batch size:\s*(\d+)", + re.IGNORECASE, +) + +MEGATRON_FILENAME_REGEX = re.compile(r"(?P.+?)_megatron_(?PMI\d+X?)", re.IGNORECASE) + +TORCHTITAN_STEP_REGEX = re.compile( + r"step:\s*(\d+).*?" + r"memory:\s*([\d.]+)GiB.*?" + r"tps:\s*([\d,]+(?:\.\d+)?).*?" + r"tflops:\s*([\d,]+(?:\.\d+)?).*?" + r"mfu:\s*([\d.]+)%" +) + +TORCHTITAN_BS_REGEX = re.compile(r"training\.local_batch_size\s*\.{2,}\s*(\d+)") +TORCHTITAN_SEQ_REGEX = re.compile(r"training\.seq_len\s*\.{2,}\s*(\d+)") + +TORCHTITAN_FILENAME_REGEX = re.compile(r"(?P.+?)_torchtitan_(?PMI\d+X?)", re.IGNORECASE) + +PRECISION_REGEX = re.compile(r"(BF16|FP8)", re.IGNORECASE) + + +def _parse_scalar(value): + if value is None: + return None + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return int(value) if float(value).is_integer() else value + if isinstance(value, str): + stripped = value.strip() + if PLAIN_INT_REGEX.match(stripped): + return int(stripped) + m = ENV_DEFAULT_REGEX.search(stripped) + if m: + return int(m.group(1)) + return None + + +def _deep_find(node, keys): + if isinstance(node, dict): + for key in keys: + if key in node: + parsed = _parse_scalar(node[key]) + if parsed is not None: + return parsed + for child in node.values(): + found = _deep_find(child, keys) + if found is not None: + return found + elif isinstance(node, list): + for item in node: + found = _deep_find(item, keys) + if found is not None: + return found + return None + + +def _regex_find(text, keys): + for key in keys: + m = re.search( + rf"^\s*{re.escape(key)}:\s*(.+?)(?:\s+#.*)?$", + text, + re.MULTILINE, + ) + if not m: + continue + parsed = _parse_scalar(m.group(1).strip()) + if parsed is not None: + return parsed + return None + + +def load_yaml_dict(path): + try: + import yaml + except ImportError: + return None + + try: + with open(path, "r", errors="ignore") as f: + data = yaml.safe_load(f) + except Exception: + return None + + return data if isinstance(data, dict) else None + + +def resolve_config_yaml(log_fname, backend, device, model, log_dir): + base = log_fname[:-4] if log_fname.endswith(".log") else log_fname + + for suffix in ("_override", "_edited", ""): + candidate = os.path.join(log_dir, f"{base}{suffix}.yaml") + if os.path.isfile(candidate): + return candidate + + config_dir = os.path.join(PRIMUS_ROOT, "examples", backend, "configs", device) + return os.path.join(config_dir, f"{model}.yaml") + + +def load_training_params(backend, device, model, log_fname, log_dir): + path = resolve_config_yaml(log_fname, backend, device, model, log_dir) + if backend == "megatron": + bs_keys, seq_keys, gbs_keys = MEGATRON_BS_KEYS, MEGATRON_SEQ_KEYS, MEGATRON_GBS_KEYS + else: + bs_keys, seq_keys, gbs_keys = TORCHTITAN_BS_KEYS, TORCHTITAN_SEQ_KEYS, TORCHTITAN_GBS_KEYS + + bs = seq = gbs = None + + data = load_yaml_dict(path) + if data is not None: + bs = _deep_find(data, bs_keys) + seq = _deep_find(data, seq_keys) + gbs = _deep_find(data, gbs_keys) + + if os.path.isfile(path) and (bs is None or seq is None or gbs is None): + with open(path, "r", errors="ignore") as f: + text = f.read() + if bs is None: + bs = _regex_find(text, bs_keys) + if seq is None: + seq = _regex_find(text, seq_keys) + if gbs is None: + gbs = _regex_find(text, gbs_keys) + + return ( + bs if bs is not None else "-", + seq if seq is not None else "-", + gbs if gbs is not None else "-", + ) + + +def log_chronological_key(fname, path): + m = RUN_LABEL_REGEX.search(fname) + if m: + return (0, int(m.group(1)), fname) + + m = LEGACY_TIMESTAMP_REGEX.search(fname) + if m: + return (1, m.group(1), fname) + + m = LEGACY_SUFFIX_REGEX.search(fname) + if m and m.group(1): + return (1, m.group(1), fname) + + return (2, os.path.getmtime(path), fname) + + +def assign_run_labels(entries): + grouped = defaultdict(list) + for entry in entries: + grouped[(entry["model"], entry["device"])].append(entry) + + for group in grouped.values(): + group.sort(key=lambda entry: log_chronological_key(entry["fname"], entry["path"])) + for index, entry in enumerate(group, start=1): + entry["run"] = f"run{index}" + + return entries + + +def run_sort_key(run_label): + m = re.fullmatch(r"run(\d+)", run_label) + if m: + return int(m.group(1)) + return 0 + + +def terminal_width(default=120): + try: + return shutil.get_terminal_size().columns + except OSError: + return default + + +def print_table(headers, rows, max_col_width=36): + if not rows: + return + + widths = [] + for i, header in enumerate(headers): + col_width = max(len(str(header)), *(len(str(row[i])) for row in rows)) + widths.append(min(col_width, max_col_width)) + + def fmt(row): + cells = [] + for i, value in enumerate(row): + text = str(value) + if len(text) > widths[i]: + text = text[: widths[i] - 3] + "..." + cells.append(text.ljust(widths[i])) + return "| " + " | ".join(cells) + " |" + + sep = "+-" + "-+-".join("-" * w for w in widths) + "-+" + print(sep) + print(fmt(headers)) + print(sep) + for row in rows: + print(fmt(row)) + print(sep) + + +def save_csv(headers, rows, backend): + os.makedirs(RESULTS_DIR, exist_ok=True) + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + latest_path = os.path.join(RESULTS_DIR, f"metrics_{backend}.csv") + snapshot_path = os.path.join(RESULTS_DIR, f"metrics_{backend}_{timestamp}.csv") + + for path in (latest_path, snapshot_path): + with open(path, "w", newline="", encoding="utf-8") as handle: + writer = csv.writer(handle) + writer.writerow(headers) + writer.writerows(rows) + + return latest_path, snapshot_path + + +def strip_ansi(text): + return ANSI_ESCAPE_REGEX.sub("", text) + + +def log_has_error(path): + exit_code = None + saw_error_line = False + + with open(path, "r", errors="ignore") as f: + for line in f: + plain = strip_ansi(line) + + m = LOG_EXIT_CODE_REGEX.search(plain) + if m: + exit_code = int(m.group(1)) + + if any(pattern.search(plain) for pattern in LOG_ERROR_EXCLUSIONS): + continue + + if any(pattern.search(plain) for pattern in LOG_ERROR_PATTERNS): + saw_error_line = True + + if exit_code not in (None, 0): + return True + return saw_error_line + + +def render_results(backend, headers, rows, note_text): + if not rows: + print(f"No {backend} logs found.") + return None + + latest_path, snapshot_path = save_csv(headers, rows, backend) + print_table(headers, rows) + + error_rows = [row for row in rows if len(row) > 2 and row[2] == ERROR_STATUS] + if error_rows: + print( + f"\n{len(error_rows)} run(s) reported errors or incomplete metrics." + " Open the corresponding log file for details." + ) + + table_width = sum(len(str(header)) for header in headers) + 3 * len(headers) + if table_width > terminal_width(): + print("\n(Table may be wider than the terminal. Open the CSV for the full view.)") + + print(note_text) + print("\nMetrics saved to:") + print(f" Latest: {latest_path}") + print(f" Snapshot: {snapshot_path}") + + return latest_path + + +def is_megatron(filename): + name = filename.lower() + return "megatron" in name or "megatorn" in name + + +def megatron_parse_filename(filename): + m = MEGATRON_FILENAME_REGEX.search(filename) + if not m: + return None + + p = PRECISION_REGEX.search(filename) + precision = p.group(1).upper() if p else "-" + + return { + "model": m.group("model"), + "device": m.group("device"), + "precision": precision, + } + + +def megatron_to_float(num_str): + return float(num_str.replace(",", "")) + + +def megatron_parse_log_file(path): + records_by_iter = {} + + with open(path, "r", errors="ignore") as f: + for line in f: + if "iteration" not in line: + continue + + m = MEGATRON_ITERATION_REGEX.search(line) + if not m: + continue + + iter_num = int(m.group(1)) + if iter_num in records_by_iter: + continue + + records_by_iter[iter_num] = { + "iter": iter_num, + "elapsed_ms": megatron_to_float(m.group(2)), + "tflops_gpu": megatron_to_float(m.group(3)), + "tokens_gpu": megatron_to_float(m.group(4)) if m.group(4) else None, + "gbs": int(m.group(5)), + } + + return [records_by_iter[i] for i in sorted(records_by_iter)] + + +def megatron_compute_averages(records): + if len(records) <= WARMUP_SKIP: + return None + + records = records[WARMUP_SKIP:] + token_values = [r["tokens_gpu"] for r in records if r["tokens_gpu"] is not None] + + return { + "count": len(records), + "elapsed_ms": mean(r["elapsed_ms"] for r in records), + "tflops_gpu": mean(r["tflops_gpu"] for r in records), + "tokens_gpu": mean(token_values) if token_values else None, + "gbs": records[0]["gbs"], + } + + +def megatron_collect_rows(): + log_dir = os.path.join(RESULTS_DIR, "logs_megatron") + entries = [] + + if not os.path.isdir(log_dir): + return [] + + for fname in sorted(os.listdir(log_dir)): + if not fname.endswith(".log") or not is_megatron(fname): + continue + + meta = megatron_parse_filename(fname) + if not meta: + continue + + path = os.path.join(log_dir, fname) + has_error = log_has_error(path) + records = megatron_parse_log_file(path) + stats = megatron_compute_averages(records) + + bs, seq, gbs = load_training_params("megatron", meta["device"], meta["model"], fname, log_dir) + + if has_error or not stats: + if gbs == "-" and records: + gbs = records[0]["gbs"] + values = [ + ERROR_STATUS, + "megatron", + meta["device"], + bs, + seq, + gbs, + meta["precision"], + "-", + "-", + "-", + "-", + ] + else: + if gbs == "-": + gbs = stats["gbs"] + + tokens_gpu = f"{stats['tokens_gpu']:.2f}" if stats["tokens_gpu"] is not None else "-" + values = [ + OK_STATUS, + "megatron", + meta["device"], + bs, + seq, + gbs, + meta["precision"], + stats["count"], + f"{stats['elapsed_ms']:.2f}", + f"{stats['tflops_gpu']:.2f}", + tokens_gpu, + ] + + entries.append( + { + "model": meta["model"], + "device": meta["device"], + "fname": fname, + "path": path, + "values": values, + } + ) + + assign_run_labels(entries) + + rows = [[entry["model"], entry["run"], *entry["values"]] for entry in entries] + rows.sort(key=lambda row: (row[0], run_sort_key(row[1]), row[7])) + return rows + + +def megatron_main(): + headers = [ + "Model", + "Run", + "Status", + "Backend", + "Device", + "BS", + "Seq", + "GBS", + "Precision", + "Iterations", + "Iter Time (ms)", + "TFLOPS/GPU", + "Tokens/GPU", + ] + render_results("megatron", headers, megatron_collect_rows(), MEGATRON_NOTE_TEXT) + + +def is_torchtitan(filename): + return "torchtitan" in filename.lower() + + +def torchtitan_parse_filename(filename): + m = TORCHTITAN_FILENAME_REGEX.search(filename) + if not m: + return None + + p = PRECISION_REGEX.search(filename) + precision = p.group(1).upper() if p else "-" + + return { + "model": m.group("model"), + "device": m.group("device"), + "precision": precision, + } + + +def torchtitan_parse_log_file(path): + steps_by_num = {} + + with open(path, "r", errors="ignore") as f: + for line in f: + m = TORCHTITAN_STEP_REGEX.search(line) + if not m: + continue + + step_num = int(m.group(1)) + if step_num in steps_by_num: + continue + + steps_by_num[step_num] = { + "step": step_num, + "memory": float(m.group(2)), + "tps": float(m.group(3).replace(",", "")), + "tflops": float(m.group(4).replace(",", "")), + "mfu": float(m.group(5)), + } + + return [steps_by_num[i] for i in sorted(steps_by_num)] + + +def torchtitan_parse_log_training_fallback(path): + bs = None + seq = None + + with open(path, "r", errors="ignore") as f: + for line in f: + if bs is None: + m = TORCHTITAN_BS_REGEX.search(line) + if m: + bs = int(m.group(1)) + if seq is None: + m = TORCHTITAN_SEQ_REGEX.search(line) + if m: + seq = int(m.group(1)) + if bs is not None and seq is not None: + break + + return bs, seq + + +def torchtitan_compute_averages(steps): + if len(steps) <= WARMUP_SKIP: + return None + + steps = steps[WARMUP_SKIP:] + + return { + "count": len(steps), + "memory": mean(s["memory"] for s in steps), + "tps": mean(s["tps"] for s in steps), + "tflops": mean(s["tflops"] for s in steps), + "mfu": mean(s["mfu"] for s in steps), + } + + +def torchtitan_collect_rows(): + log_dir = os.path.join(RESULTS_DIR, "logs_torchtitan") + entries = [] + + if not os.path.isdir(log_dir): + return [] + + for fname in sorted(os.listdir(log_dir)): + if not fname.endswith(".log") or not is_torchtitan(fname): + continue + + meta = torchtitan_parse_filename(fname) + if not meta: + continue + + path = os.path.join(log_dir, fname) + has_error = log_has_error(path) + steps = torchtitan_parse_log_file(path) + stats = torchtitan_compute_averages(steps) + + bs, seq, gbs = load_training_params("torchtitan", meta["device"], meta["model"], fname, log_dir) + if bs == "-": + log_bs, _ = torchtitan_parse_log_training_fallback(path) + if log_bs is not None: + bs = log_bs + if seq == "-": + _, log_seq = torchtitan_parse_log_training_fallback(path) + if log_seq is not None: + seq = log_seq + + if has_error or not stats: + values = [ + ERROR_STATUS, + "torchtitan", + meta["device"], + bs, + seq, + gbs, + meta["precision"], + "-", + "-", + "-", + "-", + "-", + ] + else: + values = [ + OK_STATUS, + "torchtitan", + meta["device"], + bs, + seq, + gbs, + meta["precision"], + stats["count"], + f"{stats['memory']:.2f}", + f"{stats['tps']:.2f}", + f"{stats['tflops']:.2f}", + f"{stats['mfu']:.2f}", + ] + + entries.append( + { + "model": meta["model"], + "device": meta["device"], + "fname": fname, + "path": path, + "values": values, + } + ) + + assign_run_labels(entries) + + rows = [[entry["model"], entry["run"], *entry["values"]] for entry in entries] + rows.sort(key=lambda row: (row[0], run_sort_key(row[1]), row[7])) + return rows + + +def torchtitan_main(): + headers = [ + "Model", + "Run", + "Status", + "Backend", + "Device", + "BS", + "Seq", + "GBS", + "Precision", + "Steps", + "Mem(GiB)", + "TPS", + "TFLOPS", + "MFU(%)", + ] + render_results("torchtitan", headers, torchtitan_collect_rows(), TORCHTITAN_NOTE_TEXT) + + +BACKENDS = { + "megatron": megatron_main, + "torchtitan": torchtitan_main, +} + + +def main(): + parser = argparse.ArgumentParser(description="Generate benchmark metrics tables and CSVs.") + parser.add_argument( + "backend", + choices=sorted(BACKENDS), + help="Training backend to process logs for", + ) + args = parser.parse_args() + BACKENDS[args.backend]() + + +if __name__ == "__main__": + main() diff --git a/tools/auto_benchmark/metrics_megatron.py b/tools/auto_benchmark/metrics_megatron.py deleted file mode 100644 index ff70c9c79..000000000 --- a/tools/auto_benchmark/metrics_megatron.py +++ /dev/null @@ -1,212 +0,0 @@ -#!/usr/bin/env python3 - -import os -import re -from statistics import mean - -LOG_DIR = "/workspace/Primus/tools/auto_benchmark/results/logs_megatron" - -# ============================================================ -# NOTE ON ITERATIONS AND AVERAGING -# ============================================================ - -NOTE_TEXT = """ -NOTE: -- "Iterations" represents the number of training iterations USED to compute - the averages shown above (after warm-up removal). - -- All iteration records are extracted from the log file, sorted by iteration - number, and the FIRST TWO iterations are discarded to remove warm-up effects. - -- Iteration numbers do NOT need to be sequential. Valid examples: - * iteration 1, iteration 2, iteration 3 - * iteration 1, iteration 5, iteration 10 - * iteration 1, iteration 10, iteration 20 - -- Numeric fields (elapsed time, TFLOPS/GPU, tokens/GPU) may appear as: - 1234 - 1,234 - 1,234.56 - Commas are removed before averaging. -""" - -# ============================================================ -# Regex patterns (comma-safe) -# ============================================================ - -NUM = r"[\d,]+(?:\.\d+)?" - -ITERATION_REGEX = re.compile( - rf"iteration\s+(\d+)/\s*\d+.*?" - rf"elapsed time per iteration \(ms\):\s*({NUM}).*?" - rf"throughput per GPU \(TFLOP/s/GPU\):\s*({NUM}).*?" - rf"tokens per GPU \(tokens/s/GPU\):\s*({NUM}).*?" - rf"global batch size:\s*(\d+)", - re.IGNORECASE, -) - -# Filename metadata -FILENAME_REGEX = re.compile(r"(?P.+?)_megatron_(?PMI\d+X?)", re.IGNORECASE) - -PRECISION_REGEX = re.compile(r"(BF16|FP8)", re.IGNORECASE) - -# ============================================================ -# Helpers -# ============================================================ - - -def is_megatron(filename: str) -> bool: - name = filename.lower() - return "megatron" in name or "megatorn" in name - - -def parse_filename(filename: str): - m = FILENAME_REGEX.search(filename) - if not m: - return None - - p = PRECISION_REGEX.search(filename) - precision = p.group(1).upper() if p else "-" - - return { - "model": m.group("model"), - "device": m.group("device"), - "precision": precision, - } - - -def to_float(num_str: str) -> float: - """Convert numbers like '1,234.56' safely to float.""" - return float(num_str.replace(",", "")) - - -# ============================================================ -# Log parsing (FULL FILE SCAN) -# ============================================================ - - -def parse_log_file(path): - records = [] - - with open(path, "r", errors="ignore") as f: - for line in f: - m = ITERATION_REGEX.search(line) - if m: - records.append( - { - "iter": int(m.group(1)), - "elapsed_ms": to_float(m.group(2)), - "tflops_gpu": to_float(m.group(3)), - "tokens_gpu": to_float(m.group(4)), - "gbs": int(m.group(5)), - } - ) - - return records - - -def compute_averages(records): - if len(records) <= 2: - return None - - # Sort by iteration number (non-sequential safe) - records = sorted(records, key=lambda x: x["iter"]) - - # Drop first two warm-up iterations - records = records[2:] - - return { - "count": len(records), - "elapsed_ms": mean(r["elapsed_ms"] for r in records), - "tflops_gpu": mean(r["tflops_gpu"] for r in records), - "tokens_gpu": mean(r["tokens_gpu"] for r in records), - "gbs": records[0]["gbs"], - } - - -# ============================================================ -# Table printer -# ============================================================ - - -def print_table(headers, rows): - widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))] - - def fmt(row): - return "| " + " | ".join(str(row[i]).ljust(widths[i]) for i in range(len(row))) + " |" - - sep = "+-" + "-+-".join("-" * w for w in widths) + "-+" - - print(sep) - print(fmt(headers)) - print(sep) - for r in rows: - print(fmt(r)) - print(sep) - - -def print_note(): - print(NOTE_TEXT) - - -# ============================================================ -# Main -# ============================================================ - - -def main(): - rows = [] - - for fname in sorted(os.listdir(LOG_DIR)): - if not fname.endswith(".log"): - continue - - if not is_megatron(fname): - continue - - meta = parse_filename(fname) - if not meta: - continue - - path = os.path.join(LOG_DIR, fname) - - records = parse_log_file(path) - stats = compute_averages(records) - if not stats: - continue - - rows.append( - [ - meta["model"], - "megatron", - meta["device"], - meta["precision"], - stats["count"], - f"{stats['elapsed_ms']:.2f}", - f"{stats['tflops_gpu']:.2f}", - f"{stats['tokens_gpu']:.2f}", - stats["gbs"], - ] - ) - - headers = [ - "Model", - "Backend", - "Device", - "Precision", - "Iterations", - "Iter Time (ms)", - "TFLOPS/GPU", - "Tokens/GPU", - "GBS", - ] - - if rows: - print_table(headers, rows) - print_note() - else: - print("No Megatron logs found.") - - -if __name__ == "__main__": - main() diff --git a/tools/auto_benchmark/metrics_torchtitan.py b/tools/auto_benchmark/metrics_torchtitan.py deleted file mode 100644 index 88770f72f..000000000 --- a/tools/auto_benchmark/metrics_torchtitan.py +++ /dev/null @@ -1,224 +0,0 @@ -#!/usr/bin/env python3 - -import os -import re -from statistics import mean - -LOG_DIR = "/workspace/Primus/tools/auto_benchmark/results/logs_torchtitan" - -# ============================================================ -# Regex patterns -# ============================================================ - -STEP_REGEX = re.compile( - r"step:\s*(\d+).*?" - r"memory:\s*([\d.]+)GiB.*?" - r"tps:\s*([\d,]+(?:\.\d+)?).*?" - r"tflops:\s*([\d,]+(?:\.\d+)?).*?" - r"mfu:\s*([\d.]+)%" -) - -# Dot-aligned config values -BS_REGEX = re.compile(r"training\.local_batch_size\s*\.{2,}\s*(\d+)") - -SEQ_REGEX = re.compile(r"training\.seq_len\s*\.{2,}\s*(\d+)") - -# Filename metadata -FILENAME_REGEX = re.compile(r"(?P.+?)_torchtitan_(?PMI\d+X?)", re.IGNORECASE) - -PRECISION_REGEX = re.compile(r"(BF16|FP8)", re.IGNORECASE) - -# ============================================================ -# Helpers -# ============================================================ - - -def is_torchtitan(filename): - return "torchtitan" in filename.lower() - - -def parse_filename(filename): - m = FILENAME_REGEX.search(filename) - if not m: - return None - - p = PRECISION_REGEX.search(filename) - precision = p.group(1).upper() if p else "-" - - return { - "model": m.group("model"), - "device": m.group("device"), - "precision": precision, - } - - -# ============================================================ -# Log parsing (FULL FILE SCAN) -# ============================================================ - - -def parse_log_file(path): - """ - Fully scans the log file to extract: - - BS from training.local_batch_size - - SEQ from training.seq_len - - Per-step performance metrics - """ - bs = None - seq = None - steps = [] - - with open(path, "r", errors="ignore") as f: - for line in f: - # Batch size - if bs is None: - m = BS_REGEX.search(line) - if m: - bs = int(m.group(1)) - - # Sequence length - if seq is None: - m = SEQ_REGEX.search(line) - if m: - seq = int(m.group(1)) - - # Step metrics - m = STEP_REGEX.search(line) - if m: - steps.append( - { - "step": int(m.group(1)), - "memory": float(m.group(2)), - "tps": float(m.group(3).replace(",", "")), - "tflops": float(m.group(4).replace(",", "")), - "mfu": float(m.group(5)), - } - ) - - return bs if bs is not None else "-", seq if seq is not None else "-", steps - - -def compute_averages(steps): - if len(steps) <= 2: - return None - - steps = sorted(steps, key=lambda x: x["step"]) - steps = steps[2:] # drop first two warm-up steps - - return { - "count": len(steps), - "memory": mean(s["memory"] for s in steps), - "tps": mean(s["tps"] for s in steps), - "tflops": mean(s["tflops"] for s in steps), - "mfu": mean(s["mfu"] for s in steps), - } - - -# ============================================================ -# Table printer -# ============================================================ - - -def print_table(headers, rows): - widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))] - - def fmt(row): - return "| " + " | ".join(str(row[i]).ljust(widths[i]) for i in range(len(row))) + " |" - - sep = "+-" + "-+-".join("-" * w for w in widths) + "-+" - - print(sep) - print(fmt(headers)) - print(sep) - for r in rows: - print(fmt(r)) - print(sep) - - -# ============================================================ -# Note printer -# ============================================================ - - -def print_note(): - print( - """ -NOTE: -- "Steps" represents the number of training steps USED to compute averages. -- Steps are sorted by step number and the FIRST TWO steps are dropped - to remove warm-up effects. -- Step numbers do NOT need to be sequential. Valid examples: - * step 1, step 2, step 3 - * step 1, step 5, step 10 - * step 1, step 10, step 20 -- TPS and TFLOPS values may contain commas (e.g. 1,202.17); - commas are removed before averaging. -""" - ) - - -# ============================================================ -# Main -# ============================================================ - - -def main(): - rows = [] - - for fname in sorted(os.listdir(LOG_DIR)): - if not fname.endswith(".log"): - continue - - if not is_torchtitan(fname): - continue - - meta = parse_filename(fname) - if not meta: - continue - - path = os.path.join(LOG_DIR, fname) - - bs, seq, steps = parse_log_file(path) - stats = compute_averages(steps) - if not stats: - continue - - rows.append( - [ - meta["model"], - "torchtitan", - meta["device"], - bs, - seq, - meta["precision"], - stats["count"], - f"{stats['memory']:.2f}", - f"{stats['tps']:.2f}", - f"{stats['tflops']:.2f}", - f"{stats['mfu']:.2f}", - ] - ) - - headers = [ - "Model", - "Backend", - "Device", - "BS", - "Seq", - "Precision", - "Steps", - "Mem(GiB)", - "TPS", - "TFLOPS", - "MFU(%)", - ] - - if rows: - print_table(headers, rows) - print_note() - else: - print("No TorchTitan logs found.") - - -if __name__ == "__main__": - main() diff --git a/tools/auto_benchmark/run_primus_autobenchmark.sh b/tools/auto_benchmark/run_primus_autobenchmark.sh index 53201b858..4af17ecfe 100644 --- a/tools/auto_benchmark/run_primus_autobenchmark.sh +++ b/tools/auto_benchmark/run_primus_autobenchmark.sh @@ -1,6 +1,11 @@ #!/usr/bin/env bash # Removed: set -e (allow script to continue on errors in benchmarks) +if [[ -z "${BASH_VERSION:-}" ]]; then + echo "This script must be run with bash (not sh)." >&2 + exit 1 +fi + # Set up trap to debug unexpected exits trap 'echo "[DEBUG] Script exiting at line $LINENO with exit code $?"' EXIT @@ -29,6 +34,7 @@ PRIMUS_ROOT="/workspace/Primus" MEGATRON_BASE_DIR="${PRIMUS_ROOT}/examples/megatron/configs" TORCHTITAN_BASE_DIR="${PRIMUS_ROOT}/examples/torchtitan/configs" RUN_SCRIPT="${PRIMUS_ROOT}/examples/run_pretrain.sh" +VALID_DEVICES=(MI300X MI325X MI355X) # Check if run_pretrain.sh exists, otherwise try run_pretrain_1.sh if [[ ! -f "$RUN_SCRIPT" && -f "${PRIMUS_ROOT}/examples/run_pretrain_1.sh" ]]; then @@ -36,6 +42,213 @@ if [[ ! -f "$RUN_SCRIPT" && -f "${PRIMUS_ROOT}/examples/run_pretrain_1.sh" ]]; t echo "[DEBUG] Using run_pretrain_1.sh instead" fi +# ------------------------------------------ +# Helpers +# ------------------------------------------ +install_vim_editor() { + echo -e "${YELLOW}⚠ No editor found. Installing vim...${RESET}" + + if [[ $EUID -eq 0 ]]; then + apt-get update && apt-get install -y vim + elif command -v sudo &>/dev/null; then + sudo apt-get update && sudo apt-get install -y vim + else + echo -e "${RED}✗ Cannot install vim: not root and sudo is unavailable.${RESET}" + echo -e " ${DOT} Install vim manually, then re-run config editing:" + echo -e " ${CYAN}apt-get update && apt-get install -y vim${RESET}" + return 1 + fi +} + +open_config_editor() { + local config_file="$1" + local candidate editor_bin editor_args editor_label + + for candidate in \ + "${EDITOR:-}" \ + "${VISUAL:-}" \ + "vim" \ + "vi" \ + "nano" \ + "emacs -nw" \ + "code --wait" \ + "cursor --wait"; do + if [[ -z "$candidate" ]]; then + continue + fi + + editor_bin="${candidate%% *}" + if ! command -v "$editor_bin" &>/dev/null; then + continue + fi + + editor_args="${candidate#"$editor_bin"}" + editor_label="$candidate" + echo -e " ${DOT} Using editor: ${CYAN}$editor_label${RESET}" + # shellcheck disable=SC2086 + "$editor_bin" $editor_args "$config_file" + return 0 + done + + if install_vim_editor && command -v vim &>/dev/null; then + echo -e " ${DOT} Using editor: ${CYAN}vim${RESET}" + vim "$config_file" + return 0 + fi + + echo -e "${RED}✗ Failed to open an editor for:${RESET} ${CYAN}$config_file${RESET}" + return 1 +} + +next_run_number() { + local model_name="$1" + local prefix="${model_name}_${BACKEND}_${DEVICE}" + local max_run=0 + local f bn run_n legacy_count=0 + + shopt -s nullglob + for f in "$LOG_DIR"/"${prefix}"_run*.log; do + bn=$(basename "$f" .log) + if [[ "$bn" =~ _run([0-9]+)$ ]]; then + run_n="${BASH_REMATCH[1]}" + if (( run_n > max_run )); then + max_run=$run_n + fi + fi + done + + for f in "$LOG_DIR"/"${prefix}"_*.log; do + bn=$(basename "$f" .log) + if [[ "$bn" =~ _run[0-9]+$ ]]; then + continue + fi + legacy_count=$((legacy_count + 1)) + done + shopt -u nullglob + + echo $((max_run + legacy_count + 1)) +} + +prepare_benchmark_artifacts() { + local cfg_file="$1" + local model_name run_num artifact_prefix + + PREP_CFG_FILE="$cfg_file" + PREP_MODEL_NAME=$(basename "$cfg_file" .yaml) + model_name="$PREP_MODEL_NAME" + run_num=$(next_run_number "$model_name") + PREP_RUN_LABEL="run${run_num}" + artifact_prefix="${model_name}_${BACKEND}_${DEVICE}_${PREP_RUN_LABEL}" + + PREP_LOG_FILE="$LOG_DIR/${artifact_prefix}.log" + + if [[ -n "${EDITED_CONFIGS[$cfg_file]:-}" ]]; then + PREP_WORKING_CONFIG="${EDITED_CONFIGS[$cfg_file]}" + else + PREP_WORKING_CONFIG="$cfg_file" + fi + + if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then + PREP_WORKING_CONFIG="$LOG_DIR/${artifact_prefix}_override.yaml" + cp "${EDITED_CONFIGS[$cfg_file]:-$cfg_file}" "$PREP_WORKING_CONFIG" + for KEY in "${!PARAM_OVERRIDES[@]}"; do + sed -i "s|^\([[:space:]]*${KEY}:[[:space:]]*\).*|\1${PARAM_OVERRIDES[$KEY]}|g" "$PREP_WORKING_CONFIG" + done + elif [[ -n "${EDITED_CONFIGS[$cfg_file]:-}" ]]; then + PREP_WORKING_CONFIG="$LOG_DIR/${artifact_prefix}_edited.yaml" + cp "${EDITED_CONFIGS[$cfg_file]}" "$PREP_WORKING_CONFIG" + fi +} + +execute_benchmark_run() { + local cfg_file="$1" + local working_config="$2" + local log_file="$3" + local model_name="$4" + local current="$5" + local total="$6" + + local original_config_backup="" + local run_exit_code=0 + + echo -e "${STAR} ${BOLD}Starting Benchmark ${current}/${total}...${RESET}" + echo -e " ${DOT} Model: ${CYAN}$model_name${RESET}" + echo -e " ${DOT} Backend: ${CYAN}$BACKEND${RESET}" + echo -e " ${DOT} Device: ${CYAN}$DEVICE${RESET}" + echo -e " ${DOT} Config: ${YELLOW}$working_config${RESET}" + echo -e " ${DOT} Log: ${YELLOW}$log_file${RESET}\n" + + if [[ "$working_config" != "$cfg_file" ]]; then + original_config_backup="${cfg_file}.backup_$$" + cp "$cfg_file" "$original_config_backup" + cp "$working_config" "$cfg_file" + echo -e " ${CHECK} Copied edited/overridden config to: ${CYAN}$cfg_file${RESET}" + fi + + EXP="${BACKEND_BASE_DIR}/${DEVICE}/$(basename "$cfg_file")" + export EXP + echo -e " ${CHECK} EXP set to: ${CYAN}$EXP${RESET}\n" + + echo -e " ${DOT} Changing to Primus root directory: ${CYAN}$PRIMUS_ROOT${RESET}" + cd "$PRIMUS_ROOT" || return 1 + + set +e + bash "$RUN_SCRIPT" 2>&1 | tee "$log_file" || true + run_exit_code=$? + set +e + + cd "$SCRIPT_DIR" || return 1 + + if [[ -n "$original_config_backup" && -f "$original_config_backup" ]]; then + mv "$original_config_backup" "$cfg_file" + echo -e " ${CHECK} Restored original config file" + fi + + echo + echo -e "${GREEN}==========================================${RESET}" + if [[ $run_exit_code -eq 0 ]]; then + echo -e " ${BOLD}${GREEN}✓ Benchmark ${current}/${total} Completed Successfully!${RESET}" + else + echo -e " ${BOLD}${YELLOW}⚠ Benchmark ${current}/${total} Completed with Exit Code: $run_exit_code${RESET}" + fi + echo -e " Log saved at:" + echo -e " ${CYAN}$log_file${RESET}" + if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then + echo -e " Override config saved at:" + echo -e " ${CYAN}$working_config${RESET}" + fi + echo -e "${GREEN}==========================================${RESET}" + echo + + return "$run_exit_code" +} + +generate_metrics_table() { + local metrics_script="metrics.py" + + echo + echo -e "${STAR} ${BOLD}Generating Metrics Table...${RESET}\n" + + if [[ -f "$SCRIPT_DIR/$metrics_script" && ( "$BACKEND" == "megatron" || "$BACKEND" == "torchtitan" ) ]]; then + echo -e " ${CHECK} Running: ${CYAN}python $metrics_script $BACKEND${RESET}\n" + metrics_output=$(cd "$SCRIPT_DIR" && python "$metrics_script" "$BACKEND") + metrics_status=$? + printf '%s\n' "$metrics_output" + echo + if [[ $metrics_status -eq 0 ]]; then + csv_path=$(printf '%s\n' "$metrics_output" | awk '/^ Latest:/{print $2}') + echo -e " ${CHECK} ${GREEN}Metrics table generated successfully${RESET}" + if [[ -n "$csv_path" ]]; then + echo -e " ${DOT} CSV: ${CYAN}$csv_path${RESET}" + fi + else + echo -e " ${RED}✗ Metrics generation failed${RESET}" + fi + else + echo -e " ${RED}✗ Metrics script not found: ${metrics_script:-unknown}${RESET}" + fi +} + # ------------------------------------------ # Banner # ------------------------------------------ @@ -85,14 +298,49 @@ sleep 0.2 # ------------------------------------------ echo -e "${STAR} ${BOLD}Detecting Device...${RESET}" -DEVICE=$(/opt/rocm/bin/rocminfo | grep "AMD Instinct" | head -n1 | awk '{print $5}') -echo -e " ${DOT} Device found: ${CYAN}$DEVICE${RESET}" +ROCMINFO="" +for candidate in \ + "$(command -v rocminfo 2>/dev/null)" \ + "/opt/rocm/bin/rocminfo" \ + "${ROCM_PATH:+$ROCM_PATH/bin/rocminfo}" \ + /opt/rocm-*/bin/rocminfo; do + if [[ -n "$candidate" && -x "$candidate" ]]; then + ROCMINFO="$candidate" + break + fi +done -if [[ -z "$DEVICE" || "$DEVICE" != "MI300X" && "$DEVICE" != "MI355X" ]]; then - ARCH=$(/opt/rocm/bin/rocminfo | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') +is_valid_device() { + local candidate="$1" + for dev in "${VALID_DEVICES[@]}"; do + if [[ "$candidate" == "$dev" ]]; then + return 0 + fi + done + return 1 +} + +if [[ -z "$ROCMINFO" ]]; then + echo -e " ${YELLOW}⚠ rocminfo not found (checked PATH, /opt/rocm/bin, ROCM_PATH)${RESET}" + DEVICE="" +else + echo -e " ${DOT} Using rocminfo: ${CYAN}$ROCMINFO${RESET}" + DEVICE=$("$ROCMINFO" 2>/dev/null | grep -oE 'MI3[0-9]{2}X' | head -n1) + if [[ -z "$DEVICE" ]]; then + DEVICE=$("$ROCMINFO" 2>/dev/null | grep "AMD Instinct" | head -n1 | awk '{print $5}') + fi + echo -e " ${DOT} Device found: ${CYAN}$DEVICE${RESET}" +fi + +if ! is_valid_device "$DEVICE"; then + if [[ -n "$ROCMINFO" ]]; then + ARCH=$("$ROCMINFO" 2>/dev/null | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') + else + ARCH="" + fi case "$ARCH" in - "gfx942") DEVICE="MI300X" ;; "gfx950") DEVICE="MI355X" ;; + # gfx942 is shared by MI300X and MI325X; require manual selection if marketing name is missing *) DEVICE="" ;; esac fi @@ -101,7 +349,8 @@ if [[ -z "$DEVICE" ]]; then echo -e "${RED}✗ Could not detect device automatically${RESET}" echo -e "${STAR} ${BOLD}Please select Device manually:${RESET}" echo -e " ${DOT} 1) MI300X" - echo -e " ${DOT} 2) MI355X" + echo -e " ${DOT} 2) MI325X" + echo -e " ${DOT} 3) MI355X" echo -en " ${ARROW} Enter number or name: " read -r DEV_IN @@ -110,7 +359,10 @@ if [[ -z "$DEVICE" ]]; then 1|MI300X|mi300x|Mi300x) DEVICE="MI300X" ;; - 2|MI355X|mi355x|Mi355x) + 2|MI325X|mi325x|Mi325x) + DEVICE="MI325X" + ;; + 3|MI355X|mi355x|Mi355x) DEVICE="MI355X" ;; *) @@ -175,10 +427,10 @@ SELECTED_CONFIGS=() if [[ "$CFG_NUM" == "all" ]]; then # Select all configs SELECTED_CONFIGS=("${CONFIG_LIST[@]}") -elif [[ "$CFG_NUM" =~ ^([0-9]+)-([0-9]+)$ ]]; then +elif [[ "$CFG_NUM" =~ ^[0-9]+-[0-9]+$ ]]; then # Handle range input like 4-8 - START="${BASH_REMATCH[1]}" - END="${BASH_REMATCH[2]}" + START="${CFG_NUM%%-*}" + END="${CFG_NUM##*-}" if [[ $START -lt 1 || $END -gt ${#CONFIG_LIST[@]} || $START -gt $END ]]; then echo -e "${RED}✗ Invalid range: $START-$END${RESET}" @@ -190,7 +442,9 @@ elif [[ "$CFG_NUM" =~ ^([0-9]+)-([0-9]+)$ ]]; then done else # Handle comma-separated input + _saved_ifs=$IFS IFS=',' read -ra CFG_NUMS <<< "$CFG_NUM" + IFS=$_saved_ifs for num in "${CFG_NUMS[@]}"; do # Trim whitespace @@ -205,7 +459,8 @@ else done fi -echo -e " ${CHECK} Selected ${GREEN}${#SELECTED_CONFIGS[@]}${RESET} config(s):" +SELECTED_CONFIG_COUNT=${#SELECTED_CONFIGS[@]} +echo -e " ${CHECK} Selected ${GREEN}${SELECTED_CONFIG_COUNT}${RESET} configs:" for cfg in "${SELECTED_CONFIGS[@]}"; do echo -e " ${DOT} $(basename "$cfg")" done @@ -255,7 +510,9 @@ if [[ ${#SELECTED_CONFIGS[@]} -gt 1 ]]; then if [[ "$EDIT_SELECTION" == "all" ]]; then MODELS_TO_EDIT=("${!SELECTED_CONFIGS[@]}") else + _saved_ifs=$IFS IFS=',' read -ra EDIT_NUMS <<< "$EDIT_SELECTION" + IFS=$_saved_ifs MODELS_TO_EDIT=() for num in "${EDIT_NUMS[@]}"; do num=$(echo "$num" | xargs) @@ -277,18 +534,7 @@ if [[ ${#SELECTED_CONFIGS[@]} -gt 1 ]]; then TEMP_EDIT_CONFIG="/tmp/primus_edit_${model_name}_$$.yaml" cp "$cfg" "$TEMP_EDIT_CONFIG" - # Try to find an editor - if command -v nano &> /dev/null; then - nano "$TEMP_EDIT_CONFIG" - elif command -v vim &> /dev/null; then - vim "$TEMP_EDIT_CONFIG" - elif command -v vi &> /dev/null; then - vi "$TEMP_EDIT_CONFIG" - elif command -v code &> /dev/null; then - code --wait "$TEMP_EDIT_CONFIG" - else - ${EDITOR:-vi} "$TEMP_EDIT_CONFIG" - fi + open_config_editor "$TEMP_EDIT_CONFIG" # Store the edited config EDITED_CONFIGS["$cfg"]="$TEMP_EDIT_CONFIG" @@ -311,18 +557,7 @@ elif [[ ${#SELECTED_CONFIGS[@]} -eq 1 ]]; then TEMP_EDIT_CONFIG="/tmp/primus_edit_${model_name}_$$.yaml" cp "$cfg" "$TEMP_EDIT_CONFIG" - # Try to find an editor - if command -v nano &> /dev/null; then - nano "$TEMP_EDIT_CONFIG" - elif command -v vim &> /dev/null; then - vim "$TEMP_EDIT_CONFIG" - elif command -v vi &> /dev/null; then - vi "$TEMP_EDIT_CONFIG" - elif command -v code &> /dev/null; then - code --wait "$TEMP_EDIT_CONFIG" - else - ${EDITOR:-vi} "$TEMP_EDIT_CONFIG" - fi + open_config_editor "$TEMP_EDIT_CONFIG" # Store the edited config EDITED_CONFIGS["$cfg"]="$TEMP_EDIT_CONFIG" @@ -360,7 +595,8 @@ if [[ "$OVERRIDE_PARAMS" == "y" || "$OVERRIDE_PARAMS" == "Y" ]]; then done if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then - echo -e "\n ${CHECK} ${GREEN}${#PARAM_OVERRIDES[@]}${RESET} parameter(s) will be overridden\n" + PARAM_OVERRIDE_COUNT=${#PARAM_OVERRIDES[@]} + echo -e "\n ${CHECK} ${GREEN}${PARAM_OVERRIDE_COUNT}${RESET} parameters will be overridden\n" fi fi @@ -398,7 +634,8 @@ if [[ "$ADD_ENV_VARS" == "y" || "$ADD_ENV_VARS" == "Y" ]]; then done if [[ ${#DEVICE_ENV_VARS[@]} -gt 0 ]]; then - echo -e "\n ${CHECK} ${GREEN}${#DEVICE_ENV_VARS[@]}${RESET} environment variable(s) will be set\n" + DEVICE_ENV_VAR_COUNT=${#DEVICE_ENV_VARS[@]} + echo -e "\n ${CHECK} ${GREEN}${DEVICE_ENV_VAR_COUNT}${RESET} environment variables will be set\n" fi fi @@ -453,105 +690,30 @@ for CFG_FILE in "${SELECTED_CONFIGS[@]}"; do echo -e "${MAGENTA}${BOLD}║ CONFIG FILE: $(basename "$CFG_FILE")${RESET}" echo -e "${MAGENTA}${BOLD}╚════════════════════════════════════════════════════════════╝${RESET}\n" - # Extract full filename without extension to preserve all details - CONFIG_FILENAME=$(basename "$CFG_FILE" .yaml) - MODEL_NAME="$CONFIG_FILENAME" - TIMESTAMP=$(date +"%Y-%m-%d_%H-%M-%S") + prepare_benchmark_artifacts "$CFG_FILE" - # Use the full config filename in the log name to preserve all details - LOG_FILE="$LOG_DIR/${CONFIG_FILENAME}_${BACKEND}_${DEVICE}_${TIMESTAMP}.log" - - # Use edited config if available, otherwise use original - if [[ -n "${EDITED_CONFIGS[$CFG_FILE]}" ]]; then - WORKING_CONFIG="${EDITED_CONFIGS[$CFG_FILE]}" - echo -e "${INFO} ${BOLD}Using edited config for ${CYAN}$MODEL_NAME${RESET}" - else - WORKING_CONFIG="$CFG_FILE" + if [[ -n "${EDITED_CONFIGS[$CFG_FILE]:-}" && ${#PARAM_OVERRIDES[@]} -eq 0 ]]; then + echo -e "${INFO} ${BOLD}Using edited config for ${CYAN}$PREP_MODEL_NAME${RESET}" fi - - # Apply parameter overrides if any + echo -e " ${DOT} Run label: ${CYAN}$PREP_RUN_LABEL${RESET}" if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then - OVERRIDE_CONFIG="$LOG_DIR/${MODEL_NAME}_${BACKEND}_${DEVICE}_${TIMESTAMP}_override.yaml" - cp "$WORKING_CONFIG" "$OVERRIDE_CONFIG" - echo -e "${STAR} ${BOLD}Applying parameter overrides...${RESET}" for KEY in "${!PARAM_OVERRIDES[@]}"; do - VALUE="${PARAM_OVERRIDES[$KEY]}" - # Use sed to replace the parameter value in YAML (handles both 'key: value' and 'key:value') - sed -i "s|^\([[:space:]]*${KEY}:[[:space:]]*\).*|\1${VALUE}|g" "$OVERRIDE_CONFIG" - echo -e " ${DOT} ${CYAN}$KEY${RESET}: $VALUE" + echo -e " ${DOT} ${CYAN}$KEY${RESET}: ${PARAM_OVERRIDES[$KEY]}" done - WORKING_CONFIG="$OVERRIDE_CONFIG" - echo -e " ${CHECK} Override config saved: ${YELLOW}$OVERRIDE_CONFIG${RESET}\n" - elif [[ -n "${EDITED_CONFIGS[$CFG_FILE]}" ]]; then - # Save edited config to logs directory - SAVED_CONFIG="$LOG_DIR/${MODEL_NAME}_${BACKEND}_${DEVICE}_${TIMESTAMP}_edited.yaml" - cp "$WORKING_CONFIG" "$SAVED_CONFIG" - WORKING_CONFIG="$SAVED_CONFIG" + echo -e " ${CHECK} Override config saved: ${YELLOW}$PREP_WORKING_CONFIG${RESET}\n" fi - echo -e "${STAR} ${BOLD}Starting Benchmark ${CURRENT}/${TOTAL_CONFIGS}...${RESET}" - echo -e " ${DOT} Model: ${CYAN}$MODEL_NAME${RESET}" - echo -e " ${DOT} Backend: ${CYAN}$BACKEND${RESET}" - echo -e " ${DOT} Device: ${CYAN}$DEVICE${RESET}" - echo -e " ${DOT} Config: ${YELLOW}$WORKING_CONFIG${RESET}" - echo -e " ${DOT} Log: ${YELLOW}$LOG_FILE${RESET}\n" - - # Set EXP to the working config (edited or overridden version if available) - # For edited/overridden configs, we need to copy them to the expected location - if [[ "$WORKING_CONFIG" != "$CFG_FILE" ]]; then - # Config was edited or overridden, copy it to the original location temporarily - ORIGINAL_CONFIG_BACKUP="${CFG_FILE}.backup_$$" - cp "$CFG_FILE" "$ORIGINAL_CONFIG_BACKUP" - cp "$WORKING_CONFIG" "$CFG_FILE" - echo -e " ${CHECK} Copied edited/overridden config to: ${CYAN}$CFG_FILE${RESET}" - fi - - # Set EXP to the device-specific config path (now contains edited content if applicable) - EXP_CONFIG_PATH="${BACKEND_BASE_DIR}/${DEVICE}/$(basename "$CFG_FILE")" - export EXP="$EXP_CONFIG_PATH" - echo -e " ${CHECK} EXP set to: ${CYAN}$EXP${RESET}\n" - - # Change to Primus root directory before running the script - echo -e " ${DOT} Changing to Primus root directory: ${CYAN}$PRIMUS_ROOT${RESET}" - cd "$PRIMUS_ROOT" - - # Run the script and capture exit code, but don't stop on failure - # Use 'set +e' locally to ensure we continue even on errors - set +e - bash $RUN_SCRIPT 2>&1 | tee "$LOG_FILE" || true - RUN_EXIT_CODE=$? - set -e - - # Return to script directory - cd "$SCRIPT_DIR" - - # Restore original config if it was temporarily replaced - if [[ -n "$ORIGINAL_CONFIG_BACKUP" && -f "$ORIGINAL_CONFIG_BACKUP" ]]; then - mv "$ORIGINAL_CONFIG_BACKUP" "$CFG_FILE" - echo -e " ${CHECK} Restored original config file" - unset ORIGINAL_CONFIG_BACKUP - fi - - echo - echo -e "${GREEN}==========================================${RESET}" - if [[ $RUN_EXIT_CODE -eq 0 ]]; then - echo -e " ${BOLD}${GREEN}✓ Benchmark ${CURRENT}/${TOTAL_CONFIGS} Completed Successfully!${RESET}" - else - echo -e " ${BOLD}${YELLOW}⚠ Benchmark ${CURRENT}/${TOTAL_CONFIGS} Completed with Exit Code: $RUN_EXIT_CODE${RESET}" - fi - echo -e " Log saved at:" - echo -e " ${CYAN}$LOG_FILE${RESET}" - if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then - echo -e " Override config saved at:" - echo -e " ${CYAN}$OVERRIDE_CONFIG${RESET}" - fi - echo -e "${GREEN}==========================================${RESET}" - echo + execute_benchmark_run \ + "$PREP_CFG_FILE" \ + "$PREP_WORKING_CONFIG" \ + "$PREP_LOG_FILE" \ + "$PREP_MODEL_NAME" \ + "$CURRENT" \ + "$TOTAL_CONFIGS" || true CURRENT=$((CURRENT + 1)) - # Add a short delay between runs if [[ $CURRENT -le $TOTAL_CONFIGS ]]; then echo -e "${YELLOW}Preparing next benchmark...${RESET}\n" echo -e "${INFO} ${BOLD}Next: Config ${CURRENT}/${TOTAL_CONFIGS}${RESET}\n" @@ -561,26 +723,10 @@ done echo echo -e "${MAGENTA}${BOLD}=========================================${RESET}" -echo -e "${MAGENTA}${BOLD} All ${TOTAL_CONFIGS} Benchmark(s) Completed!${RESET}" +echo -e "${MAGENTA}${BOLD} All ${TOTAL_CONFIGS} benchmarks completed!${RESET}" echo -e "${MAGENTA}${BOLD}=========================================${RESET}" # ------------------------------------------ # 7. GENERATE METRICS TABLE # ------------------------------------------ -echo -echo -e "${STAR} ${BOLD}Generating Metrics Table...${RESET}\n" - -if [[ "$BACKEND" == "megatron" ]]; then - METRICS_SCRIPT="metrics_megatron.py" -elif [[ "$BACKEND" == "torchtitan" ]]; then - METRICS_SCRIPT="metrics_torchtitan.py" -fi - -if [[ -f "$METRICS_SCRIPT" ]]; then - echo -e " ${CHECK} Running: ${CYAN}python $METRICS_SCRIPT${RESET}\n" - python "$METRICS_SCRIPT" - echo - echo -e " ${CHECK} ${GREEN}Metrics table generated successfully${RESET}" -else - echo -e " ${RED}✗ Metrics script not found: $METRICS_SCRIPT${RESET}" -fi +generate_metrics_table From a6550770ef5560070e23807ff73850f012821090 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Tue, 14 Jul 2026 03:51:29 +0300 Subject: [PATCH 023/127] feat(flux): torch.compile + DDP-overlap compile patches (#815) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/flux` and also merges `feat/flux/turbo` — review after both. ## What this changes The `torch.compile` and DDP-overlap-compile patches for the Flux model, plus the compile / graph-break tests and the two backend-selection / TE-vs-local-spec tests (which top-level-import turbo modules). ## Why it has two parents The source only needs the turbo layer, but the co-located tests run the full compiled model under FSDP2 — importing the Flux model (`feat/flux/flux`) and the fp8 all-gather (via `feat/flux/turbo` → `feat/flux/opt`). ## Dependencies Builds on the CI-pins PR (`feat/flux/ci-env`) — it carries the required bumped Primus-Turbo pin for the CK launch path; that PR merges first. Also builds on `feat/flux/flux` + `feat/flux/turbo`. (No `MEGATRON_PATH` dependency: the compiled-vs-eager test runs in-process via a Dynamo reset before each compiled build, and the conftest puts the recursively-checked-out Megatron submodule on `sys.path`.) ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "compile or backend_selection or te_vs_local"`. Validated locally on an AMD GPU container: 41 passed (incl. the in-process `test_compiled_local_vs_eager`/`test_te_eager_vs_local_compiled`, confirmed passing with no Megatron entry on `PYTHONPATH`). ## Files 7 (compile + DDP-overlap patches, compile/graph-break + backend-selection tests). --------- Co-authored-by: Flux Split Trial Co-authored-by: luiza-amd --- .../patches/ddp_overlap_compile_patches.py | 173 ++ .../megatron/patches/torch_compile_patches.py | 82 + .../test_flux_compile_checkpoint_keys.py | 57 + .../diffusion/test_flux_compile_strategies.py | 128 + .../test_flux_layer_spec_backend_selection.py | 79 + .../test_fp8_compile_graph_breaks.py | 2316 +++++++++++++++++ .../test_te_vs_local_spec_attention.py | 381 +++ 7 files changed, 3216 insertions(+) create mode 100644 primus/backends/megatron/patches/ddp_overlap_compile_patches.py create mode 100644 primus/backends/megatron/patches/torch_compile_patches.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_compile_checkpoint_keys.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_compile_strategies.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_layer_spec_backend_selection.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_fp8_compile_graph_breaks.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_te_vs_local_spec_attention.py diff --git a/primus/backends/megatron/patches/ddp_overlap_compile_patches.py b/primus/backends/megatron/patches/ddp_overlap_compile_patches.py new file mode 100644 index 000000000..02ed61e12 --- /dev/null +++ b/primus/backends/megatron/patches/ddp_overlap_compile_patches.py @@ -0,0 +1,173 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +DDP overlap_param_gather + torch.compile compatibility patches. + +When using per_block torch.compile with Megatron DDP's overlap_param_gather, +the original per-sub-module hook registration causes Dynamo to trace through +hooks containing NCCL side effects, leading to assertion errors. + +This patch uses two-tier hook registration: + - Layer hooks (recurse=True): one hook per transformer layer (children of + ModuleList). These fire in the eager __call__ wrapper before the compiled + forward, causing zero graph breaks inside compiled layers. + - Non-layer hooks (recurse=False): hooks on all modules outside compiled + regions (embeddings, projections, norms). These preserve original Megatron + behavior and run entirely in eager mode. +""" + +import torch + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +@register_patch( + "megatron.ddp.overlap_param_gather_compile", + backend="megatron", + phase="before_train", + description=( + "Patch DDP forward pre-hooks for overlap_param_gather + torch.compile " + "compatibility using two-tier hook registration." + ), + # Distinct from the FSDP2 fp8-cache patch (priority=45) to avoid a + # registration-order-dependent tie-break. The two are mutually exclusive + # (DDP vs FSDP2), so the exact relative order is not load-bearing. + priority=44, + condition=lambda ctx: ( + getattr(get_args(ctx), "use_distributed_optimizer", False) + and getattr(get_args(ctx), "overlap_param_gather", False) + and getattr(getattr(get_args(ctx), "torch_compile", None), "enable", False) + and not getattr(get_args(ctx), "disable_ddp_compile_patches", False) + ), +) +def patch_ddp_overlap_param_gather_for_compile(ctx: PatchContext) -> None: + """Patch DistributedDataParallel to use two-tier hook registration. + + Monkey-patches class methods on DDP so that both the initial hook + registration in __init__ and subsequent enable/disable cycles in the + training loop use the patched logic. + """ + try: + from megatron.core.distributed.distributed_data_parallel import ( + DistributedDataParallel as DDP, + ) + from megatron.core.transformer.cuda_graphs import is_graph_capturing + + def _get_overlap_hook_modules(self): + """Partition modules into layer-tier and non-layer-tier for hooks. + + Layer modules: children of nn.ModuleList containers (transformer + layers). These get hooks with recurse=True. + + Non-layer modules: everything else (embeddings, projections, norms). + These get hooks with recurse=False (original Megatron behavior). + Modules that are descendants of a layer module are excluded (they + are covered by the layer hook's recurse=True). + """ + # Preserve module-traversal order (deterministic) while de-duplicating; + # iterating a plain set would make hook-registration order vary run to + # run, which hurts reproducibility/debuggability. + layer_modules = [] + seen = set() + for module in self.module.modules(): + if isinstance(module, torch.nn.ModuleList): + for child in module.children(): + if child not in seen: + seen.add(child) + layer_modules.append(child) + + inside_layer = set() + for layer in layer_modules: + for sub in layer.modules(): + inside_layer.add(sub) + + non_layer_modules = [m for m in self.module.modules() if m not in inside_layer] + + return layer_modules, non_layer_modules + + def _make_forward_pre_hook(self, recurse=False): + """Create a forward pre-hook parameterized by recurse depth.""" + + def hook(module, *unused): + if not self.use_forward_hook: + raise RuntimeError("Should use pre-hook only when overlap_param_gather is True") + + if is_graph_capturing(): + return + + for param in module.parameters(recurse=recurse): + if param not in self.param_to_bucket_group: + continue + if not param.requires_grad: + raise RuntimeError("Bucketed param in forward pre-hook must require grad") + + skip_next_bucket_dispatch = ( + self.ddp_config.align_param_gather or self.overlap_param_gather_with_optimizer_step + ) + self.param_to_bucket_group[param].finish_param_sync( + skip_next_bucket_dispatch=skip_next_bucket_dispatch + ) + + return hook + + def enable_forward_pre_hook(self): + """Register two-tier forward pre-hooks for param all-gather overlap.""" + if not self.use_forward_hook: + raise RuntimeError("enable_forward_pre_hook requires use_forward_hook=True") + if len(self.remove_forward_pre_hook_handles) != 0: + raise RuntimeError("Forward pre-hooks already registered") + + layer_modules, non_layer_modules = self._get_overlap_hook_modules() + + layer_hook = self._make_forward_pre_hook(recurse=True) + for module in layer_modules: + self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook(layer_hook) + + non_layer_hook = self._make_forward_pre_hook(recurse=False) + for module in non_layer_modules: + self.remove_forward_pre_hook_handles[module] = module.register_forward_pre_hook( + non_layer_hook + ) + + def disable_forward_pre_hook(self, param_sync: bool = True): + """Remove all forward pre-hooks (both tiers).""" + if not self.use_forward_hook: + raise RuntimeError("disable_forward_pre_hook requires use_forward_hook=True") + for module, handle in list(self.remove_forward_pre_hook_handles.items()): + handle.remove() + self.remove_forward_pre_hook_handles.clear() + + if param_sync: + self.start_param_sync(force_sync=True) + + DDP._get_overlap_hook_modules = _get_overlap_hook_modules + DDP._make_forward_pre_hook = _make_forward_pre_hook + DDP.enable_forward_pre_hook = enable_forward_pre_hook + DDP.disable_forward_pre_hook = disable_forward_pre_hook + + log_rank_0( + "[Patch:megatron.ddp.overlap_param_gather_compile] " + "Patched DDP with two-tier hook registration for " + "overlap_param_gather + torch.compile compatibility" + ) + + except Exception as e: + import traceback + + log_rank_0( + f"[Patch:megatron.ddp.overlap_param_gather_compile] " + f"ERROR: Failed to patch DDP: {type(e).__name__}: {e}" + ) + log_rank_0(f"Traceback: {traceback.format_exc()}") + # Re-raise: silently leaving DDP unpatched here would run training with + # incorrect forward pre-hooks (Dynamo graph breaks / NCCL side effects), + # so fail loudly instead of degrading correctness. + raise RuntimeError( + "Failed to apply DDP overlap_param_gather + torch.compile patch; " + "see log above for the underlying error." + ) from e diff --git a/primus/backends/megatron/patches/torch_compile_patches.py b/primus/backends/megatron/patches/torch_compile_patches.py new file mode 100644 index 000000000..0bc709d28 --- /dev/null +++ b/primus/backends/megatron/patches/torch_compile_patches.py @@ -0,0 +1,82 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Megatron torch.compile Patches + +This module contains patches that modify Megatron's setup_model_and_optimizer +to apply torch.compile after model setup. +""" + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +@register_patch( + "megatron.training.torch_compile", + backend="megatron", + phase="before_train", + description="Patch setup_model_and_optimizer to apply torch.compile after model setup", + priority=60, # Higher than optimizer patch to wrap last + condition=lambda ctx: ( + getattr(get_args(ctx), "torch_compile", None) is not None + and getattr(getattr(get_args(ctx), "torch_compile", None), "enable", False) + ) + or getattr(get_args(ctx), "enable_torch_compile", False), +) +def patch_setup_model_and_optimizer_for_torch_compile(ctx: PatchContext): + """ + Patch Megatron's setup_model_and_optimizer to apply torch.compile after model setup. + + Behavior: + - Wraps setup_model_and_optimizer() to call apply_torch_compile_if_enabled() + after model is created and wrapped + - If compilation fails, exception propagates (fails entire setup) + - Works for both FSDP2 and non-FSDP2 paths + """ + try: + from megatron.training import training + + from primus.backends.megatron.core.utils import ( + apply_torch_compile_if_enabled, + apply_torch_compile_to_optimizer_if_enabled, + ) + + # Save original function + original_setup_model_and_optimizer = training.setup_model_and_optimizer + + def patched_setup_model_and_optimizer(*args, **kwargs): + """Patched setup_model_and_optimizer that applies torch.compile after model setup.""" + from megatron.training import get_args + + result = original_setup_model_and_optimizer(*args, **kwargs) + + # Extract model from result tuple + model, optimizer, opt_param_scheduler = result + + # Apply torch.compile (raises exception on failure) + megatron_args = get_args() + apply_torch_compile_if_enabled(model, megatron_args) + apply_torch_compile_to_optimizer_if_enabled(optimizer, megatron_args) + + # Return original result + return result + + # Apply the patch + training.setup_model_and_optimizer = patched_setup_model_and_optimizer + log_rank_0( + "[Patch:megatron.training.torch_compile] " + "Patched setup_model_and_optimizer to apply torch.compile after model setup" + ) + + except Exception as e: + log_rank_0( + f"[Patch:megatron.training.torch_compile] " + f"WARNING: Failed to patch setup_model_and_optimizer: {type(e).__name__}: {e}" + ) + import traceback + + log_rank_0(f"Traceback: {traceback.format_exc()}") diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_checkpoint_keys.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_checkpoint_keys.py new file mode 100644 index 000000000..1db4b6457 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_checkpoint_keys.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests that torch.compile does not alter sharded_state_dict keys. + +The "whole_model" compile strategy compiles self.forward without wrapping +submodules, so state_dict keys must remain identical before and after +compilation. This test guards against regressions (e.g., _orig_mod prefixes +leaking into keys). +""" + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from tests.utils import PrimusUT + + +class TestFluxCompileCheckpointKeys(PrimusUT): + """Verify torch.compile does not change sharded_state_dict keys.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state for model tests.""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + def test_compile_does_not_change_sharded_state_dict_keys(self): + """Keys from sharded_state_dict() must be identical before and after compile_model(). + + Compile must actually run (enable_torch_compile=True), otherwise + compile_model() returns early and the test is vacuous. 'per_block' + rebinds each layer's .forward (no module re-wrap), so a regression that + wraps the module instead would leak '_orig_mod' prefixes into the keys + and fail this test. + """ + config = FluxConfig.flux_535m() + config.enable_torch_compile = True + config.torch_compile_strategy = "per_block" + model = Flux(config).cuda() + + keys_before = set(model.sharded_state_dict().keys()) + + model.compile_model() + + keys_after = set(model.sharded_state_dict().keys()) + + added = keys_after - keys_before + removed = keys_before - keys_after + + assert not added, f"Keys added after compile: {added}" + assert not removed, f"Keys removed after compile: {removed}" diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_strategies.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_strategies.py new file mode 100644 index 000000000..fbde0e0a4 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_compile_strategies.py @@ -0,0 +1,128 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for Flux compile_model() strategy validation. + +Tests that invalid configuration combinations raise the expected errors +and that fallback behavior works correctly. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from tests.utils import PrimusUT + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +class TestFluxCompileStrategies(PrimusUT): + """Tests for compile_model() error/warning behavior.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state for model tests.""" + + def test_whole_model_rejects_overlap_param_gather(self): + """whole_model strategy must reject overlap_param_gather.""" + config = FluxConfig.flux_535m() + config.enable_torch_compile = True + config.torch_compile_strategy = "whole_model" + model = Flux(config) + + with patch( + "megatron.training.get_args", + return_value=SimpleNamespace(overlap_param_gather=True), + ): + with pytest.raises(ValueError, match="overlap_param_gather"): + model.compile_model() + + def test_stack_rejects_recompute_granularity(self): + """stack strategy must reject recompute_granularity='full'.""" + config = FluxConfig.flux_535m() + config.enable_torch_compile = True + config.torch_compile_strategy = "stack" + config.recompute_granularity = "full" + model = Flux(config) + + with pytest.raises(ValueError, match="recompute_granularity"): + model.compile_model() + + def test_double_stack_rejects_recompute_granularity(self): + """double_stack strategy must reject recompute_granularity='full'.""" + config = FluxConfig.flux_535m() + config.enable_torch_compile = True + config.torch_compile_strategy = "double_stack" + config.recompute_granularity = "full" + model = Flux(config) + + with pytest.raises(ValueError, match="recompute_granularity"): + model.compile_model() + + def test_local_spec_stack_falls_back_to_per_block(self): + """Local spec (non-TE) with stack strategy should fallback to per_block.""" + config = FluxConfig.flux_535m() + config.enable_torch_compile = True + config.torch_compile_strategy = "stack" + config.transformer_impl = "local" + model = Flux(config) + + compile_calls = [] + + def mock_compile(fn, **kwargs): + compile_calls.append(fn) + return fn + + with patch("torch.compile", side_effect=mock_compile): + model.compile_model() + + assert len(compile_calls) == len(model.transformer.layers), ( + f"Expected {len(model.transformer.layers)} per-block compile calls, " f"got {len(compile_calls)}" + ) + + def test_per_block_with_cuda_graph_warns(self): + """per_block + enable_cuda_graph should emit a warning.""" + config = FluxConfig.flux_535m() + config.enable_torch_compile = True + config.torch_compile_strategy = "per_block" + # Build without cuda_graph (avoids CudaGraphManager compat issues), + # then set the flag before compile_model() which only reads config. + model = Flux(config) + model.config.enable_cuda_graph = True + + compile_calls = [] + + def mock_compile(fn, **kwargs): + compile_calls.append(fn) + return fn + + output_lines = [] + + def capture_log(msg): + output_lines.append(msg) + + with patch("torch.compile", side_effect=mock_compile), patch( + "primus.core.utils.module_utils.log_rank_0", side_effect=capture_log + ): + model.compile_model() + + output_text = "\n".join(output_lines) + assert "CUDA graph" in output_text, f"Expected CUDA graph warning in output, got: {output_text[:200]}" + + def test_invalid_strategy_raises_value_error(self): + """Invalid strategy name must raise ValueError.""" + config = FluxConfig.flux_535m() + config.enable_torch_compile = True + config.torch_compile_strategy = "nonexistent" + model = Flux(config) + + with pytest.raises(ValueError, match="Invalid torch_compile_strategy"): + model.compile_model() diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_layer_spec_backend_selection.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_layer_spec_backend_selection.py new file mode 100644 index 000000000..d1af8b474 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_layer_spec_backend_selection.py @@ -0,0 +1,79 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for Flux layer spec backend selection. + +Tests that get_flux_layer_spec() correctly selects backend based on transformer_impl. +This ensures alignment between backend selection and FSDP2 wrapping decisions. +""" + +import pytest + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.layer_spec import ( + get_flux_layer_spec, +) +from tests.utils import PrimusUT + + +class TestFluxLayerSpecBackendSelection(PrimusUT): + """Tests for backend selection in get_flux_layer_spec().""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state for layer spec tests.""" + + def test_backend_selection_fp8_local_spec(self): + """Test that local + fp8 selects PrimusTurboFloat8LocalSpecProvider.""" + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + Float8ColumnParallelLinear, + ) + + config = FluxConfig.flux_535m( + transformer_impl="local", + fp8="e4m3", + fp8_recipe="tensorwise", + ) + + block_submodules = get_flux_layer_spec(config, backend=None) + + # At least one linear spec should reference Float8ColumnParallelLinear + found_fp8 = False + for layer_spec in block_submodules.layer_specs: + attn_spec = layer_spec.submodules.self_attention + if hasattr(attn_spec, "submodules") and hasattr(attn_spec.submodules, "linear_qkv"): + if attn_spec.submodules.linear_qkv == Float8ColumnParallelLinear: + found_fp8 = True + break + assert found_fp8, "Expected Float8ColumnParallelLinear in layer specs for local+fp8" + + def test_backend_selection_local_no_fp8_uses_native_linear(self): + """Test that local without fp8 uses native ColumnParallelLinear.""" + + from megatron.core.tensor_parallel import ColumnParallelLinear + + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + Float8ColumnParallelLinear, + ) + + config = FluxConfig.flux_535m(transformer_impl="local", fp8=None) + + block_submodules = get_flux_layer_spec(config, backend=None) + + found_any = False + for layer_spec in block_submodules.layer_specs: + attn_spec = layer_spec.submodules.self_attention + if hasattr(attn_spec, "submodules") and hasattr(attn_spec.submodules, "linear_qkv"): + found_any = True + assert ( + attn_spec.submodules.linear_qkv != Float8ColumnParallelLinear + ), "Should NOT use Float8ColumnParallelLinear when fp8=None" + assert ( + attn_spec.submodules.linear_qkv == ColumnParallelLinear + ), f"Expected native ColumnParallelLinear when fp8=None, got {attn_spec.submodules.linear_qkv}" + assert found_any, "No attention linear_qkv specs found to validate backend selection" diff --git a/tests/unit_tests/backends/megatron/diffusion/test_fp8_compile_graph_breaks.py b/tests/unit_tests/backends/megatron/diffusion/test_fp8_compile_graph_breaks.py new file mode 100644 index 000000000..3f9a24d14 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_fp8_compile_graph_breaks.py @@ -0,0 +1,2316 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +torch.compile graph-break coverage for the FP8 linear / attention path. + +Motivation: the @allow_in_graph approach used by older FP8 ops created graph +break boundaries between AdaLN Triton kernels and FP8 quantize/GEMM ops, +costing ~144 ms (10.5%) GPU idle per iteration. Eliminating those breaks +(setup_context Functions + fused amax capture) is what this file guards. + +This file has two kinds of tests: + +1. Production-guarding tests -- these ``.apply`` the shipped Functions + (``OpaqueFP8LinearTensorwiseFunction``, ``DualFP8LinearTensorwiseFunction``, + and ``DelayedFP8LinearTensorwiseFunction``) and assert they trace under + torch.compile with zero graph breaks and stay numerically equivalent to + eager. A regression in the production op fails these directly. +2. Characterization / feasibility tests -- these use small in-file + autograd.Function replicas (``_Level1FP8Linear``, the ``_AllowInGraph*`` + baselines, ``_DelayedFP8*``) to pin the torch.compile *capabilities* the + production design relies on (buffer ``copy_`` inside a compiled forward, + tensor hooks firing, ``allow_in_graph`` vs setup_context parity, buffer + mutation visibility between compiled calls). They do not guard a production + symbol; they document why the production approach is sound and would surface + a torch/Inductor behavior change. + +Numerical forward/backward correctness of the production Functions lives in +``test_delayed_fp8_triton_op.py`` and the turbo float8 tests; this file is +specifically about compile behavior. +""" + +import math + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +import primus_turbo.pytorch as pt +import torch.nn as nn +from primus_turbo.pytorch.core.backend import BackendType, GlobalBackendManager +from primus_turbo.pytorch.core.low_precision import ( + Float8QuantConfig, + Format, + ScalingGranularity, + float8_e4m3, + float8_e5m2, +) +from primus_turbo.pytorch.kernels.gemm.gemm_fp8_impl import gemm_fp8_impl +from primus_turbo.pytorch.ops.quantization import quantize_fp8 + +from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + FP8UnshardedWeightTensor, +) +from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + DelayedFP8LinearTensorwiseFunction, + DualFP8LinearTensorwiseFunction, + OpaqueFP8LinearTensorwiseFunction, + _extract_fp8_weight, +) +from tests.utils import PrimusUT + +_has_cuda = torch.cuda.is_available() +requires_cuda = pytest.mark.skipif(not _has_cuda, reason="CUDA required") + +_DIM = 64 +_OUT_DIM = 128 +_BATCH = 4 +_GRAN_VALUE = ScalingGranularity.TENSORWISE.value +_BACKEND_VALUE = BackendType.HIPBLASLT.value + + +# --------------------------------------------------------------------------- +# Level 1: Bare autograd.Function -- C++ custom ops, primitive args only +# --------------------------------------------------------------------------- + + +class _Level1FP8Linear(torch.autograd.Function): + """FP8 linear without @allow_in_graph, using setup_context + primitives. + + Calls the C++ quantize_fp8_tensorwise and gemm_fp8_impl custom ops + directly. No Float8QuantConfig, no FP8UnshardedWeightTensor. + """ + + @staticmethod + def forward(input, weight, fp8_dtype, gran_value, backend_value): + out_dtype = input.dtype + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + a_fp8, a_scale_inv = torch.ops.primus_turbo_cpp_extension.quantize_fp8_tensorwise( + input_2d, + fp8_dtype, + None, + ) + b_fp8, b_scale_inv = torch.ops.primus_turbo_cpp_extension.quantize_fp8_tensorwise( + weight, + fp8_dtype, + None, + ) + + output = gemm_fp8_impl( + a_fp8, + a_scale_inv, + False, + b_fp8, + b_scale_inv, + True, + out_dtype, + False, + granularity=gran_value, + default_backend=backend_value, + ) + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + return output, a_fp8, a_scale_inv, b_fp8, b_scale_inv + + @staticmethod + def setup_context(ctx, inputs, output): + input, weight, fp8_dtype, gran_value, backend_value = inputs + output_val, a_fp8, a_scale_inv, b_fp8, b_scale_inv = output + + ctx.save_for_backward(a_fp8, a_scale_inv, b_fp8, b_scale_inv) + ctx.mark_non_differentiable(a_fp8, a_scale_inv, b_fp8, b_scale_inv) + ctx.out_dtype = input.dtype + ctx.orig_shape = input.shape + ctx.fp8_dtype = fp8_dtype + ctx.gran_value = gran_value + ctx.backend_value = backend_value + + @staticmethod + def backward(ctx, grad_output, *_): + a_fp8, a_scale_inv, b_fp8, b_scale_inv = ctx.saved_tensors + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + if not grad_2d.is_contiguous(): + grad_2d = grad_2d.contiguous() + + grad_fp8, grad_scale_inv = torch.ops.primus_turbo_cpp_extension.quantize_fp8_tensorwise( + grad_2d, + ctx.fp8_dtype, + None, + ) + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + b_fp8, + b_scale_inv, + False, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_weight = gemm_fp8_impl( + a_fp8, + a_scale_inv, + True, + grad_fp8, + grad_scale_inv, + False, + ctx.out_dtype, + True, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + return grad_input, grad_weight, None, None, None + + +class _Level1Module(nn.Module): + """Minimal module: LayerNorm + Level 1 FP8 linear.""" + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + + def forward(self, x): + x = self.norm(x) + out = _Level1FP8Linear.apply( + x, + self.weight, + float8_e4m3, + _GRAN_VALUE, + _BACKEND_VALUE, + ) + return out[0] + + +# --------------------------------------------------------------------------- +# Bonus: LayerNorm + FP8 linear single compiled graph +# --------------------------------------------------------------------------- + + +@requires_cuda +class TestFullGraphIntegration(PrimusUT): + """Verify LayerNorm + FP8 linear compiles into a single graph frame.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def test_single_compiled_frame(self): + """LayerNorm → FP8 linear should produce exactly 1 compiled frame.""" + model = _Level1Module(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda") + + x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda") + + explanation = torch._dynamo.explain(model)(x) + + assert explanation.graph_break_count == 0, ( + f"Expected single graph (0 breaks), got {explanation.graph_break_count}. " + f"Reasons: {explanation.break_reasons}" + ) + assert explanation.graph_count == 1, f"Expected 1 graph, got {explanation.graph_count}." + + +# --------------------------------------------------------------------------- +# Compiled vs eager numerical equivalence (regression for view_as aliasing fix) +# --------------------------------------------------------------------------- + + +class _OpaqueModule(nn.Module): + """Module using OpaqueFP8LinearTensorwiseFunction with pre-extracted weights.""" + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + + def forward(self, x): + x = self.norm(x) + w_fp8, w_scale = quantize_fp8(self.weight, float8_e4m3, ScalingGranularity.TENSORWISE) + result = OpaqueFP8LinearTensorwiseFunction.apply( + x, + self.weight, + w_fp8, + w_scale, + float8_e4m3, + float8_e5m2, + _GRAN_VALUE, + _BACKEND_VALUE, + ) + return result[0] + + +@requires_cuda +class TestCompiledVsEagerEquivalence(PrimusUT): + """Verify compiled and eager execution produce identical forward and backward results. + + Regression test for the view_as input-output aliasing bug where AOTAutograd's + view-replay mechanism could corrupt saved FP8 weight data. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def test_forward_equivalence(self): + """Compiled forward matches eager forward exactly.""" + model = _OpaqueModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda") + x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda") + + eager_out = model(x) + + torch._dynamo.reset() + compiled = torch.compile(model) + compiled_out = compiled(x) + + torch.testing.assert_close(compiled_out, eager_out, atol=0, rtol=0) + + def test_backward_equivalence(self): + """Compiled backward matches eager backward exactly.""" + model = _OpaqueModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda") + + x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda", requires_grad=True) + eager_out = model(x) + eager_out.sum().backward() + eager_x_grad = x.grad.clone() + eager_w_grad = model.weight.grad.clone() + + x.grad = None + model.weight.grad = None + torch._dynamo.reset() + compiled = torch.compile(model) + compiled_out = compiled(x) + compiled_out.sum().backward() + + torch.testing.assert_close(x.grad, eager_x_grad, atol=0, rtol=0) + torch.testing.assert_close(model.weight.grad, eager_w_grad, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# Multi-step convergence: compiled (setup_context) vs compiled (@allow_in_graph) +# vs eager, plus BF16 baseline +# --------------------------------------------------------------------------- + +import copy + + +@torch._dynamo.allow_in_graph +class _AllowInGraphFP8Linear(torch.autograd.Function): + """Old-style @allow_in_graph FP8 linear for baseline comparison. + + Quantizes both input and weight inside forward, saves FP8 tensors on ctx. + AOTAutograd never traces through this -- it's fully opaque. + """ + + @staticmethod + def forward(ctx, input, weight): + out_dtype = input.dtype + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + a_fp8, a_scale_inv = quantize_fp8(input_2d, float8_e4m3, ScalingGranularity.TENSORWISE) + b_fp8, b_scale_inv = quantize_fp8(weight, float8_e4m3, ScalingGranularity.TENSORWISE) + + output = gemm_fp8_impl( + a_fp8, + a_scale_inv, + False, + b_fp8, + b_scale_inv, + True, + out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + ctx.save_for_backward(a_fp8, a_scale_inv, b_fp8, b_scale_inv) + ctx.out_dtype = out_dtype + ctx.orig_shape = orig_shape + return output + + @staticmethod + def backward(ctx, grad_output): + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + a_fp8, a_scale_inv, b_fp8, b_scale_inv = ctx.saved_tensors + + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + grad_fp8, grad_scale_inv = quantize_fp8(grad_2d, float8_e5m2, ScalingGranularity.TENSORWISE) + + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + b_fp8, + b_scale_inv, + False, + ctx.out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_weight = gemm_fp8_impl( + a_fp8, + a_scale_inv, + True, + grad_fp8, + grad_scale_inv, + False, + ctx.out_dtype, + True, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + return grad_input, grad_weight + + +class _AllowInGraphModule(nn.Module): + """Module using old @allow_in_graph FP8 linear.""" + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + + def forward(self, x): + x = self.norm(x) + return _AllowInGraphFP8Linear.apply(x, self.weight) + + +class _BF16LinearModule(nn.Module): + """Pure BF16 linear for baseline comparison (no FP8 quantization).""" + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.linear = nn.Linear(dim, out_dim, bias=False) + + def forward(self, x): + x = self.norm(x) + return self.linear(x) + + +_CONV_DIM = 256 +_CONV_OUT = 512 +_CONV_BATCH = 16 +_N_STEPS = 200 + + +def _make_inputs(n, dim, batch, seed=42): + g = torch.Generator(device="cpu").manual_seed(seed) + return [torch.randn(batch, dim, dtype=torch.bfloat16, generator=g).cuda() for _ in range(n)] + + +def _run_training_loop(model, inputs, n_steps, lr=1e-3): + optimizer = torch.optim.AdamW(model.parameters(), lr=lr) + has_prepare = hasattr(model, "prepare_fp8_weight") + if not has_prepare and hasattr(model, "_orig_mod"): + has_prepare = hasattr(model._orig_mod, "prepare_fp8_weight") + losses = [] + grad_norms = [] + for step in range(n_steps): + if has_prepare: + m = model._orig_mod if hasattr(model, "_orig_mod") else model + m.prepare_fp8_weight() + optimizer.zero_grad() + out = model(inputs[step % len(inputs)]) + loss = out.sum() + loss.backward() + gn = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=float("inf")) + grad_norms.append(gn.item()) + optimizer.step() + losses.append(loss.item()) + return losses, grad_norms + + +def _clone_state(model): + return copy.deepcopy(model.state_dict()) + + +def _print_comparison(label_a, losses_a, label_b, losses_b, milestones=None): + """Print side-by-side loss comparison at milestones for diagnosis.""" + if milestones is None: + milestones = [0, 1, 5, 10, 20, 50, 100, 150, 199] + milestones = [m for m in milestones if m < len(losses_a) and m < len(losses_b)] + print(f"\n{'Step':>6} | {label_a:>20} | {label_b:>20} | {'Rel Diff':>10}") + print("-" * 65) + for m in milestones: + a, b = losses_a[m], losses_b[m] + rel = abs(a - b) / max(abs(a), 1e-12) + print(f"{m:>6} | {a:>20.6f} | {b:>20.6f} | {rel:>10.6f}") + + +@requires_cuda +class TestMultiStepConvergence(PrimusUT): + """Multi-step convergence comparison across compilation modes. + + Runs 200 optimizer steps on a small model and compares loss/grad_norm + trajectories to isolate whether AOTAutograd backward compilation + causes numerical divergence vs the old @allow_in_graph approach. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def _make_fp8_setup_ctx_model(self, state_dict): + model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + model.load_state_dict(state_dict) + return model + + def _make_fp8_allow_in_graph_model(self, state_dict): + model = _AllowInGraphModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + model.load_state_dict(state_dict) + return model + + def _make_bf16_model(self, state_dict): + model = _BF16LinearModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + # Map weight -> linear.weight for BF16 module + bf16_state = { + "norm.weight": state_dict["norm.weight"], + "norm.bias": state_dict["norm.bias"], + "linear.weight": state_dict["weight"], + } + model.load_state_dict(bf16_state) + return model + + def test_compiled_setup_ctx_vs_eager(self): + """setup_context compiled should track FP8 eager closely over 200 steps.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(ref_model) + + eager_losses, eager_gn = _run_training_loop(ref_model, inputs, _N_STEPS) + + torch._dynamo.reset() + compiled_model = self._make_fp8_setup_ctx_model(init_state) + compiled_model_c = torch.compile(compiled_model) + compiled_losses, compiled_gn = _run_training_loop(compiled_model_c, inputs, _N_STEPS) + + _print_comparison("FP8 Eager", eager_losses, "FP8 Compiled(setup_ctx)", compiled_losses) + + final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"FP8 compiled (setup_context) diverged from eager: " + f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + def test_allow_in_graph_vs_eager(self): + """@allow_in_graph compiled should track FP8 eager closely over 200 steps.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(ref_model) + + eager_losses, eager_gn = _run_training_loop(ref_model, inputs, _N_STEPS) + + torch._dynamo.reset() + aig_model = self._make_fp8_allow_in_graph_model(init_state) + aig_model_c = torch.compile(aig_model) + aig_losses, aig_gn = _run_training_loop(aig_model_c, inputs, _N_STEPS) + + _print_comparison("FP8 Eager", eager_losses, "FP8 @allow_in_graph", aig_losses) + + final_rel = abs(eager_losses[-1] - aig_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"FP8 @allow_in_graph diverged from eager: " + f"eager={eager_losses[-1]:.6f}, aig={aig_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + def test_bf16_compiled_vs_eager(self): + """BF16 compiled should match BF16 eager closely (no FP8 noise).""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(ref_model) + + eager_model = self._make_bf16_model(init_state) + eager_losses, _ = _run_training_loop(eager_model, inputs, _N_STEPS) + + torch._dynamo.reset() + compiled_model = self._make_bf16_model(init_state) + compiled_model_c = torch.compile(compiled_model) + compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS) + + _print_comparison("BF16 Eager", eager_losses, "BF16 Compiled", compiled_losses) + + final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.01, ( + f"BF16 compiled diverged from eager: " + f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + +# --------------------------------------------------------------------------- +# FSDP2 FP8UnshardedWeightTensor subclass variants +# --------------------------------------------------------------------------- + + +class _FP8SubclassModule(nn.Module): + """Module simulating FSDP2 FP8 all-gather: weight goes through + FP8UnshardedWeightTensor subclass before _extract_fp8_weight. + + In production, FSDP2 creates the FP8UnshardedWeightTensor outside the + compiled graph (in fsdp_post_all_gather). The compiled forward only sees + the subclass as an input and calls _extract_fp8_weight on it. We simulate + this by pre-creating the subclass in prepare_fp8_weight() which must be + called before each forward step. + """ + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + self._fp8_weight = None + + @torch.no_grad() + def prepare_fp8_weight(self): + """Simulate FSDP2 fsdp_post_all_gather: create FP8UnshardedWeightTensor + outside the compiled graph, just like FSDP2 does before forward.""" + w_fp8, w_scale = quantize_fp8(self.weight, float8_e4m3, ScalingGranularity.TENSORWISE) + fp8_config = Float8QuantConfig(format=Format.E4M3, granularity=ScalingGranularity.TENSORWISE) + self._fp8_weight = FP8UnshardedWeightTensor(w_fp8, w_scale, torch.bfloat16, fp8_config) + + def forward(self, x): + x = self.norm(x) + weight_fp8, weight_scale_inv = _extract_fp8_weight(self._fp8_weight, float8_e4m3) + result = OpaqueFP8LinearTensorwiseFunction.apply( + x, + self.weight, + weight_fp8, + weight_scale_inv, + float8_e4m3, + float8_e5m2, + _GRAN_VALUE, + _BACKEND_VALUE, + ) + return result[0] + + +@torch._dynamo.allow_in_graph +class _AllowInGraphFP8LinearWithSubclass(torch.autograd.Function): + """Old-style @allow_in_graph FP8 linear that receives an + FP8UnshardedWeightTensor and extracts data inside the opaque boundary.""" + + @staticmethod + def forward(ctx, input, weight): + out_dtype = input.dtype + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + a_fp8, a_scale_inv = quantize_fp8(input_2d, float8_e4m3, ScalingGranularity.TENSORWISE) + + if isinstance(weight, FP8UnshardedWeightTensor): + b_fp8, b_scale_inv = weight.get_fp8_data_and_scale_inv() + else: + b_fp8, b_scale_inv = quantize_fp8(weight, float8_e4m3, ScalingGranularity.TENSORWISE) + + output = gemm_fp8_impl( + a_fp8, + a_scale_inv, + False, + b_fp8, + b_scale_inv, + True, + out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + ctx.save_for_backward(a_fp8, a_scale_inv, b_fp8, b_scale_inv) + ctx.out_dtype = out_dtype + ctx.orig_shape = orig_shape + return output + + @staticmethod + def backward(ctx, grad_output): + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + a_fp8, a_scale_inv, b_fp8, b_scale_inv = ctx.saved_tensors + + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + grad_fp8, grad_scale_inv = quantize_fp8(grad_2d, float8_e5m2, ScalingGranularity.TENSORWISE) + + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + b_fp8, + b_scale_inv, + False, + ctx.out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_weight = gemm_fp8_impl( + a_fp8, + a_scale_inv, + True, + grad_fp8, + grad_scale_inv, + False, + ctx.out_dtype, + True, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + return grad_input, grad_weight + + +class _AllowInGraphSubclassModule(nn.Module): + """Module using old @allow_in_graph FP8 linear with FP8UnshardedWeightTensor + passed directly into the opaque boundary (matching old production code). + + Like _FP8SubclassModule, the subclass is created outside compile.""" + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + self._fp8_weight = None + + @torch.no_grad() + def prepare_fp8_weight(self): + w_fp8, w_scale = quantize_fp8(self.weight, float8_e4m3, ScalingGranularity.TENSORWISE) + fp8_config = Float8QuantConfig(format=Format.E4M3, granularity=ScalingGranularity.TENSORWISE) + self._fp8_weight = FP8UnshardedWeightTensor(w_fp8, w_scale, torch.bfloat16, fp8_config) + + def forward(self, x): + x = self.norm(x) + return _AllowInGraphFP8LinearWithSubclass.apply(x, self._fp8_weight) + + +@requires_cuda +class TestFP8SubclassConvergence(PrimusUT): + """Multi-step convergence with FP8UnshardedWeightTensor subclass path. + + Isolates whether the FSDP2 tensor subclass interaction under torch.compile + causes numerical divergence vs eager or the old @allow_in_graph path. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def test_fp8_subclass_compiled_vs_eager(self): + """setup_context + FP8UnshardedWeightTensor: compiled vs eager.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _FP8SubclassModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(ref_model) + + eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS) + + torch._dynamo.reset() + compiled_model = _FP8SubclassModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + compiled_model.load_state_dict(init_state) + compiled_model_c = torch.compile(compiled_model) + compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS) + + _print_comparison("FP8 Subclass Eager", eager_losses, "FP8 Subclass Compiled", compiled_losses) + + final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"FP8 subclass compiled diverged from eager: " + f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + def test_fp8_subclass_allow_in_graph_vs_eager(self): + """@allow_in_graph + FP8UnshardedWeightTensor: compiled vs eager.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _AllowInGraphSubclassModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(ref_model) + + eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS) + + torch._dynamo.reset() + compiled_model = _AllowInGraphSubclassModule(_CONV_DIM, _CONV_OUT).to( + dtype=torch.bfloat16, device="cuda" + ) + compiled_model.load_state_dict(init_state) + compiled_model_c = torch.compile(compiled_model) + compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS) + + _print_comparison( + "FP8 AIG+Subclass Eager", eager_losses, "FP8 AIG+Subclass Compiled", compiled_losses + ) + + final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"FP8 @allow_in_graph+subclass compiled diverged from eager: " + f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + def test_subclass_vs_plain_weight_eager(self): + """FP8UnshardedWeightTensor path vs plain quantize path (both eager). + Verifies the subclass wrapper doesn't change numerical results.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + plain_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(plain_model) + + plain_losses, _ = _run_training_loop(plain_model, inputs, _N_STEPS) + + subclass_model = _FP8SubclassModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + subclass_model.load_state_dict(init_state) + subclass_losses, _ = _run_training_loop(subclass_model, inputs, _N_STEPS) + + _print_comparison("FP8 Plain Eager", plain_losses, "FP8 Subclass Eager", subclass_losses) + + final_rel = abs(plain_losses[-1] - subclass_losses[-1]) / max(abs(plain_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.01, ( + f"FP8 subclass path diverged from plain path in eager mode: " + f"plain={plain_losses[-1]:.6f}, subclass={subclass_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + +# --------------------------------------------------------------------------- +# Dual FP8 linear convergence tests +# --------------------------------------------------------------------------- + + +class _DualFP8Module(nn.Module): + """Module using DualFP8LinearTensorwiseFunction (two GEMMs in one node). + + Simulates the JointSelfAttention dual output projection pattern: + input goes through LayerNorm, then two independent FP8 linear projections + are computed in a single autograd node. Outputs are summed for loss. + """ + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight_a = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + self.weight_b = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + + def forward(self, x): + x = self.norm(x) + w_fp8_a, w_scale_a = quantize_fp8(self.weight_a, float8_e4m3, ScalingGranularity.TENSORWISE) + w_fp8_b, w_scale_b = quantize_fp8(self.weight_b, float8_e4m3, ScalingGranularity.TENSORWISE) + result = DualFP8LinearTensorwiseFunction.apply( + x, + self.weight_a, + w_fp8_a, + w_scale_a, + x, + self.weight_b, + w_fp8_b, + w_scale_b, + float8_e4m3, + float8_e5m2, + _GRAN_VALUE, + _BACKEND_VALUE, + ) + return result[0] + result[1] + + +@torch._dynamo.allow_in_graph +class _AllowInGraphDualFP8Linear(torch.autograd.Function): + """Old-style @allow_in_graph dual FP8 linear for baseline comparison. + + Quantizes both inputs and weights inside forward, saves FP8 tensors on ctx. + Matches the old DualFP8LinearTensorwiseFunction behavior before setup_context. + """ + + @staticmethod + def forward(ctx, input_a, weight_a, input_b, weight_b): + out_dtype = input_a.dtype + + orig_shape_a = input_a.shape + orig_shape_b = input_b.shape + input_a_2d = input_a.reshape(-1, input_a.shape[-1]) + input_b_2d = input_b.reshape(-1, input_b.shape[-1]) + + a_fp8_a, a_scale_a = quantize_fp8(input_a_2d, float8_e4m3, ScalingGranularity.TENSORWISE) + b_fp8_a, b_scale_a = quantize_fp8(weight_a, float8_e4m3, ScalingGranularity.TENSORWISE) + a_fp8_b, a_scale_b = quantize_fp8(input_b_2d, float8_e4m3, ScalingGranularity.TENSORWISE) + b_fp8_b, b_scale_b = quantize_fp8(weight_b, float8_e4m3, ScalingGranularity.TENSORWISE) + + output_a = gemm_fp8_impl( + a_fp8_a, + a_scale_a, + False, + b_fp8_a, + b_scale_a, + True, + out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + output_b = gemm_fp8_impl( + a_fp8_b, + a_scale_b, + False, + b_fp8_b, + b_scale_b, + True, + out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + + output_a = output_a.reshape(*orig_shape_a[:-1], output_a.shape[-1]) + output_b = output_b.reshape(*orig_shape_b[:-1], output_b.shape[-1]) + + ctx.save_for_backward( + a_fp8_a, + a_scale_a, + b_fp8_a, + b_scale_a, + a_fp8_b, + a_scale_b, + b_fp8_b, + b_scale_b, + ) + ctx.out_dtype = out_dtype + ctx.orig_shape_a = orig_shape_a + ctx.orig_shape_b = orig_shape_b + return output_a, output_b + + @staticmethod + def backward(ctx, grad_output_a, grad_output_b): + if not grad_output_a.is_contiguous(): + grad_output_a = grad_output_a.contiguous() + if not grad_output_b.is_contiguous(): + grad_output_b = grad_output_b.contiguous() + + (a_fp8_a, a_scale_a, b_fp8_a, b_scale_a, a_fp8_b, a_scale_b, b_fp8_b, b_scale_b) = ctx.saved_tensors + + grad_a_2d = grad_output_a.reshape(-1, grad_output_a.shape[-1]) + grad_fp8_a, grad_scale_a = quantize_fp8(grad_a_2d, float8_e5m2, ScalingGranularity.TENSORWISE) + + grad_input_a = gemm_fp8_impl( + grad_fp8_a, + grad_scale_a, + False, + b_fp8_a, + b_scale_a, + False, + ctx.out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + grad_input_a = grad_input_a.reshape(ctx.orig_shape_a) + + grad_weight_a = gemm_fp8_impl( + a_fp8_a, + a_scale_a, + True, + grad_fp8_a, + grad_scale_a, + False, + ctx.out_dtype, + True, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + + grad_b_2d = grad_output_b.reshape(-1, grad_output_b.shape[-1]) + grad_fp8_b, grad_scale_b = quantize_fp8(grad_b_2d, float8_e5m2, ScalingGranularity.TENSORWISE) + + grad_input_b = gemm_fp8_impl( + grad_fp8_b, + grad_scale_b, + False, + b_fp8_b, + b_scale_b, + False, + ctx.out_dtype, + False, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + grad_input_b = grad_input_b.reshape(ctx.orig_shape_b) + + grad_weight_b = gemm_fp8_impl( + a_fp8_b, + a_scale_b, + True, + grad_fp8_b, + grad_scale_b, + False, + ctx.out_dtype, + True, + granularity=_GRAN_VALUE, + default_backend=_BACKEND_VALUE, + ) + + return grad_input_a, grad_weight_a, grad_input_b, grad_weight_b + + +class _AllowInGraphDualModule(nn.Module): + """Module using old @allow_in_graph dual FP8 linear.""" + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight_a = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + self.weight_b = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + + def forward(self, x): + x = self.norm(x) + out_a, out_b = _AllowInGraphDualFP8Linear.apply(x, self.weight_a, x, self.weight_b) + return out_a + out_b + + +class _TwoSinglesModule(nn.Module): + """Module using two separate OpaqueFP8LinearTensorwiseFunction calls. + + Same computation as _DualFP8Module but using two independent autograd + nodes instead of one combined node. Used to verify the dual node + doesn't introduce numerical differences. + """ + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight_a = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + self.weight_b = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + + def forward(self, x): + x = self.norm(x) + w_fp8_a, w_scale_a = quantize_fp8(self.weight_a, float8_e4m3, ScalingGranularity.TENSORWISE) + w_fp8_b, w_scale_b = quantize_fp8(self.weight_b, float8_e4m3, ScalingGranularity.TENSORWISE) + result_a = OpaqueFP8LinearTensorwiseFunction.apply( + x, + self.weight_a, + w_fp8_a, + w_scale_a, + float8_e4m3, + float8_e5m2, + _GRAN_VALUE, + _BACKEND_VALUE, + ) + result_b = OpaqueFP8LinearTensorwiseFunction.apply( + x, + self.weight_b, + w_fp8_b, + w_scale_b, + float8_e4m3, + float8_e5m2, + _GRAN_VALUE, + _BACKEND_VALUE, + ) + return result_a[0] + result_b[0] + + +@requires_cuda +class TestDualFP8Convergence(PrimusUT): + """Multi-step convergence for DualFP8LinearTensorwiseFunction. + + The dual function bundles two FP8 GEMMs into a single autograd node + (used by JointSelfAttention). This tests whether the combined node + under torch.compile causes divergence vs eager or the old @allow_in_graph. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def _make_dual_init_state(self): + model = _DualFP8Module(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + return model, _clone_state(model) + + def _load_dual_model(self, cls, state_dict): + model = cls(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + model.load_state_dict(state_dict) + return model + + def test_dual_compiled_setup_ctx_vs_eager(self): + """DualFP8 setup_context compiled should track eager over 200 steps.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model, init_state = self._make_dual_init_state() + eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS) + + torch._dynamo.reset() + compiled_model = self._load_dual_model(_DualFP8Module, init_state) + compiled_model_c = torch.compile(compiled_model) + compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS) + + _print_comparison("Dual FP8 Eager", eager_losses, "Dual FP8 Compiled", compiled_losses) + + final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"Dual FP8 compiled (setup_context) diverged from eager: " + f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + def test_dual_allow_in_graph_vs_eager(self): + """DualFP8 @allow_in_graph compiled should track eager over 200 steps.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model, init_state = self._make_dual_init_state() + eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS) + + torch._dynamo.reset() + aig_model = self._load_dual_model(_AllowInGraphDualModule, init_state) + aig_model_c = torch.compile(aig_model) + aig_losses, _ = _run_training_loop(aig_model_c, inputs, _N_STEPS) + + _print_comparison("Dual FP8 Eager", eager_losses, "Dual FP8 @allow_in_graph", aig_losses) + + final_rel = abs(eager_losses[-1] - aig_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"Dual FP8 @allow_in_graph diverged from eager: " + f"eager={eager_losses[-1]:.6f}, aig={aig_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + def test_dual_vs_two_singles_eager(self): + """Dual node should match two single OpaqueFP8 calls (both eager). + + Verifies that bundling two GEMMs into one autograd node doesn't + change numerical results compared to two separate nodes. + """ + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + dual_model, init_state = self._make_dual_init_state() + dual_losses, _ = _run_training_loop(dual_model, inputs, _N_STEPS) + + singles_model = self._load_dual_model(_TwoSinglesModule, init_state) + singles_losses, _ = _run_training_loop(singles_model, inputs, _N_STEPS) + + _print_comparison("Dual FP8 Eager", dual_losses, "Two Singles Eager", singles_losses) + + final_rel = abs(dual_losses[-1] - singles_losses[-1]) / max(abs(dual_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.01, ( + f"Dual FP8 node diverged from two single nodes: " + f"dual={dual_losses[-1]:.6f}, singles={singles_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + +# --------------------------------------------------------------------------- +# FP8 Attention Module: QKV projection + flash attention + output projection +# --------------------------------------------------------------------------- + +_ATTN_DIM = 256 +_ATTN_HEADS = 4 +_ATTN_SEQ = 32 +_ATTN_BATCH = 8 +_ATTN_N_STEPS = 200 + + +def _hermetic_compile_state(): + """Reset global compile / kernel-dispatch state for hermetic isolation. + + These convergence tests are sensitive to global state that leaks across + tests in the same process: + * primus-turbo's kernel-dispatch and origami autotune caches + (``GlobalBackendManager.reset()`` clears both), and + * torch dynamo's compiled-graph cache (``torch._dynamo.reset()``). + Without isolation, kernel/codegen choices made by an earlier test bleed + in and shift the compiled-vs-eager numerics, which then compounds over the + 200-step loop into order-dependent (flaky) failures. We also pin + auto-tune off so dispatch deterministically uses the default backend. + """ + GlobalBackendManager.reset() + GlobalBackendManager.set_auto_tune(False) + torch._dynamo.reset() + torch.manual_seed(0) + torch.cuda.manual_seed_all(0) + + +class _FP8AttentionModule(nn.Module): + """Minimal FP8 attention: QKV projection + flash attention + output projection. + + Mirrors the real Flux attention path: + - FP8 linear for QKV (OpaqueFP8LinearTensorwiseFunction) + - pt.ops.flash_attn_func for attention + - FP8 linear for output projection + """ + + def __init__(self, dim, num_heads, deterministic=False): + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.softmax_scale = 1.0 / math.sqrt(self.head_dim) + self.deterministic = deterministic + + self.norm = nn.LayerNorm(dim) + self.qkv_weight = nn.Parameter(torch.randn(3 * dim, dim, dtype=torch.bfloat16)) + self.proj_weight = nn.Parameter(torch.randn(dim, dim, dtype=torch.bfloat16)) + + def forward(self, x): + B, S, D = x.shape + x = self.norm(x) + + qkv_w_fp8, qkv_w_scale = quantize_fp8(self.qkv_weight, float8_e4m3, ScalingGranularity.TENSORWISE) + qkv = OpaqueFP8LinearTensorwiseFunction.apply( + x, + self.qkv_weight, + qkv_w_fp8, + qkv_w_scale, + float8_e4m3, + float8_e5m2, + _GRAN_VALUE, + _BACKEND_VALUE, + )[0] + + qkv = qkv.view(B, S, 3, self.num_heads, self.head_dim) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + + attn_out = pt.ops.flash_attn_func( + q, + k, + v, + dropout_p=0.0, + softmax_scale=self.softmax_scale, + causal=False, + window_size=(-1, -1), + deterministic=self.deterministic, + return_lse=False, + ) + + attn_out = attn_out.reshape(B, S, D) + + proj_w_fp8, proj_w_scale = quantize_fp8(self.proj_weight, float8_e4m3, ScalingGranularity.TENSORWISE) + out = OpaqueFP8LinearTensorwiseFunction.apply( + attn_out, + self.proj_weight, + proj_w_fp8, + proj_w_scale, + float8_e4m3, + float8_e5m2, + _GRAN_VALUE, + _BACKEND_VALUE, + )[0] + + return out + + +def _make_seq_inputs(n, seq, dim, batch, seed=42): + g = torch.Generator(device="cpu").manual_seed(seed) + return [torch.randn(batch, seq, dim, dtype=torch.bfloat16, generator=g).cuda() for _ in range(n)] + + +@requires_cuda +class TestFP8AttentionConvergence(PrimusUT): + """FP8 attention convergence: compiled vs eager with flash attention. + + Tests whether torch.compile interacts with FP8 QKV projection + + flash attention + FP8 output projection to cause divergence. + + Flash attention is non-deterministic by default, so the non-deterministic + test measures whether compile-induced variance stays within the natural + variance of flash attention itself. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + @pytest.fixture(autouse=True) + def hermetic_compile_state(self): + _hermetic_compile_state() + try: + yield + finally: + _hermetic_compile_state() + + def setup_method(self, method): + torch._dynamo.reset() + + def test_fp8_attention_nondeterministic_compiled_vs_eager(self): + """With non-deterministic flash attention, compile variance should be + within a small multiple of the natural flash attention variance.""" + torch._dynamo.reset() + inputs = _make_seq_inputs(20, _ATTN_SEQ, _ATTN_DIM, _ATTN_BATCH) + + ref_model = _FP8AttentionModule(_ATTN_DIM, _ATTN_HEADS, deterministic=False).to( + dtype=torch.bfloat16, device="cuda" + ) + init_state = _clone_state(ref_model) + + eager1 = _FP8AttentionModule(_ATTN_DIM, _ATTN_HEADS, deterministic=False).to( + dtype=torch.bfloat16, device="cuda" + ) + eager1.load_state_dict(init_state) + eager1_losses, _ = _run_training_loop(eager1, inputs, _ATTN_N_STEPS) + + eager2 = _FP8AttentionModule(_ATTN_DIM, _ATTN_HEADS, deterministic=False).to( + dtype=torch.bfloat16, device="cuda" + ) + eager2.load_state_dict(init_state) + eager2_losses, _ = _run_training_loop(eager2, inputs, _ATTN_N_STEPS) + + torch._dynamo.reset() + compiled_model = _FP8AttentionModule(_ATTN_DIM, _ATTN_HEADS, deterministic=False).to( + dtype=torch.bfloat16, device="cuda" + ) + compiled_model.load_state_dict(init_state) + compiled_model = torch.compile(compiled_model) + compiled_losses, _ = _run_training_loop(compiled_model, inputs, _ATTN_N_STEPS) + + _print_comparison( + "FP8 Attn Eager1", + eager1_losses, + "FP8 Attn Eager2", + eager2_losses, + ) + _print_comparison( + "FP8 Attn Eager1", + eager1_losses, + "FP8 Attn Compiled", + compiled_losses, + ) + + natural_var = abs(eager1_losses[-1] - eager2_losses[-1]) + compile_var = abs(eager1_losses[-1] - compiled_losses[-1]) + ref_loss = max(abs(eager1_losses[-1]), 1e-12) + + natural_rel = natural_var / ref_loss + compile_rel = compile_var / ref_loss + + print(f"\nNatural variance (eager vs eager): {natural_var:.6f} (rel: {natural_rel:.6f})") + print(f"Compile variance (eager vs compiled): {compile_var:.6f} (rel: {compile_rel:.6f})") + + if natural_var < 1e-6: + assert compile_rel < 0.05, ( + f"Flash attn showed no natural variance but compile diverged: " + f"compile_rel={compile_rel:.6f}" + ) + else: + ratio = compile_var / natural_var + print(f"Compile/natural variance ratio: {ratio:.2f}x") + assert ratio < 5.0, ( + f"Compile variance {ratio:.1f}x larger than natural flash attn variance: " + f"natural={natural_var:.6f}, compile={compile_var:.6f}" + ) + + +# --------------------------------------------------------------------------- +# Mock Megatron DDP: param/grad buffer remapping + backward hooks +# --------------------------------------------------------------------------- + + +class _MockMegatronDDP(nn.Module): + """Simulates Megatron DDP's param/grad buffer remapping and backward hooks. + + Reproduces the essential behavior from distributed_data_parallel.py + (lines 419-444) and param_and_grad_buffer.py (lines 760-785) without + any Megatron imports: + + 1. Remaps param.data into a contiguous param_buffer + (like use_distributed_optimizer=True) + 2. Assigns param.main_grad as views into a contiguous grad_buffer + 3. Registers gradient accumulator hooks that do: + param.main_grad.add_(param.grad.data); param.grad = None + """ + + def __init__(self, module): + super().__init__() + self.module = module + + params = [p for p in module.parameters() if p.requires_grad] + + total_numel = sum(p.numel() for p in params) + self.param_buffer = torch.zeros(total_numel, dtype=params[0].dtype, device=params[0].device) + self.grad_buffer = torch.zeros(total_numel, dtype=params[0].dtype, device=params[0].device) + + offset = 0 + self.grad_accs = [] + for param in params: + numel = param.numel() + new_data = self.param_buffer[offset : offset + numel].view(param.shape) + new_data.copy_(param.data) + param.data = new_data + param.main_grad = self.grad_buffer[offset : offset + numel].view(param.shape) + param.grad_added_to_main_grad = False + offset += numel + + param_tmp = param.expand_as(param) + grad_acc = param_tmp.grad_fn.next_functions[0][0] + grad_acc.register_hook(self._make_hook(param)) + self.grad_accs.append(grad_acc) + + @staticmethod + def _make_hook(param): + def hook(*unused): + if param.grad is not None and not param.grad_added_to_main_grad: + param.main_grad.add_(param.grad.data) + param.grad = None + + return hook + + def zero_grad_buffer(self): + self.grad_buffer.zero_() + for p in self.module.parameters(): + if p.requires_grad: + p.grad_added_to_main_grad = False + + def forward(self, *args, **kwargs): + return self.module(*args, **kwargs) + + +class _MainGradAdamW(torch.optim.AdamW): + """AdamW that reads from param.main_grad instead of param.grad. + + Simulates how Megatron's distributed optimizer operates on main_grad + buffers rather than the standard param.grad tensors. + """ + + @torch.no_grad() + def step(self, closure=None): + for group in self.param_groups: + for p in group["params"]: + if hasattr(p, "main_grad") and p.main_grad is not None: + p.grad = p.main_grad.clone() + result = super().step(closure=closure) + for group in self.param_groups: + for p in group["params"]: + p.grad = None + return result + + +def _run_mock_ddp_training(module, inputs, n_steps, compile_target=None, lr=1e-3): + """Run training with mock Megatron DDP. + + Args: + module: The model to wrap with mock DDP. + inputs: List of input tensors. + n_steps: Number of training steps. + compile_target: One of None, "whole_model", "inner_only". + None = eager, "whole_model" = compile ddp.forward, + "inner_only" = compile ddp.module.forward. + lr: Learning rate. + """ + ddp = _MockMegatronDDP(module) + optimizer = _MainGradAdamW(ddp.module.parameters(), lr=lr) + + if compile_target == "whole_model": + ddp.forward = torch.compile(ddp.forward) + elif compile_target == "inner_only": + ddp.module.forward = torch.compile(ddp.module.forward) + + losses = [] + grad_norms = [] + for step in range(n_steps): + ddp.zero_grad_buffer() + out = ddp(inputs[step % len(inputs)]) + loss = out.sum() + loss.backward() + gn = torch.nn.utils.clip_grad_norm_( + [p for p in ddp.module.parameters() if p.main_grad is not None], + max_norm=float("inf"), + ) + grad_norms.append(gn.item() if hasattr(gn, "item") else gn) + optimizer.step() + losses.append(loss.item()) + return losses, grad_norms + + +@requires_cuda +class TestMockDDPConvergence(PrimusUT): + """Mock Megatron DDP convergence: whole_model compile vs inner-only vs eager. + + Tests whether torch.compile scope interacts with Megatron DDP's backward + hooks (param.grad -> main_grad transfer + param.grad = None) to cause + numerical divergence. + + The local spec DDP config uses whole_model compile, while the converging + TE spec reference uses stack (inner-only) compile. Phase 1 proved that + FP8 compiled vs eager is bit-identical without DDP hooks. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def test_mock_ddp_whole_model_compiled_vs_eager(self): + """whole_model compile + mock DDP should track eager + mock DDP.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(ref_model) + + eager_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + eager_model.load_state_dict(init_state) + eager_losses, _ = _run_mock_ddp_training(eager_model, inputs, _N_STEPS) + + torch._dynamo.reset() + compiled_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + compiled_model.load_state_dict(init_state) + compiled_losses, _ = _run_mock_ddp_training( + compiled_model, inputs, _N_STEPS, compile_target="whole_model" + ) + + _print_comparison( + "MockDDP Eager", + eager_losses, + "MockDDP WholeModel", + compiled_losses, + ) + + final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"Mock DDP whole_model compile diverged from eager: " + f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + def test_mock_ddp_inner_only_compiled_vs_eager(self): + """inner-only compile + mock DDP should track eager + mock DDP.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(ref_model) + + eager_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + eager_model.load_state_dict(init_state) + eager_losses, _ = _run_mock_ddp_training(eager_model, inputs, _N_STEPS) + + torch._dynamo.reset() + compiled_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + compiled_model.load_state_dict(init_state) + compiled_losses, _ = _run_mock_ddp_training( + compiled_model, inputs, _N_STEPS, compile_target="inner_only" + ) + + _print_comparison( + "MockDDP Eager", + eager_losses, + "MockDDP InnerOnly", + compiled_losses, + ) + + final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"Mock DDP inner-only compile diverged from eager: " + f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + def test_mock_ddp_vs_vanilla_eager(self): + """Mock DDP eager should produce same results as vanilla eager.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(ref_model) + + vanilla_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + vanilla_model.load_state_dict(init_state) + vanilla_losses, _ = _run_training_loop(vanilla_model, inputs, _N_STEPS) + + ddp_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + ddp_model.load_state_dict(init_state) + ddp_losses, _ = _run_mock_ddp_training(ddp_model, inputs, _N_STEPS) + + _print_comparison( + "Vanilla Eager", + vanilla_losses, + "MockDDP Eager", + ddp_losses, + ) + + final_rel = abs(vanilla_losses[-1] - ddp_losses[-1]) / max(abs(vanilla_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.01, ( + f"Mock DDP eager diverged from vanilla eager: " + f"vanilla={vanilla_losses[-1]:.6f}, ddp={ddp_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + +# --------------------------------------------------------------------------- +# Phase 0: Delayed FP8 Scaling Feasibility Tests +# +# These tests resolve compile-interaction unknowns required before +# implementing delayed scaling. Each test answers a specific question +# about torch.compile behavior with buffer mutations, tensor hooks, +# and autograd Function side-outputs. +# --------------------------------------------------------------------------- + +_FP8_FWD_MAX = torch.finfo(float8_e4m3).max +_FP8_BWD_MAX = torch.finfo(float8_e5m2).max + + +# -- Test 0a helpers -------------------------------------------------------- + + +class _BufferCopyModule(nn.Module): + """Module that copies a computed scalar into a registered buffer. + + Mimics the real delayed scaling pattern where input_amax is computed + from the BF16 input (not the FP8 output) and stored via buffer copy_(). + """ + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + self.register_buffer("staged_amax", torch.tensor(0.0)) + + def forward(self, x): + x = self.norm(x) + input_amax = x.detach().abs().amax().float() + result = _Level1FP8Linear.apply( + x, + self.weight, + float8_e4m3, + _GRAN_VALUE, + _BACKEND_VALUE, + ) + self.staged_amax.copy_(input_amax) + return result[0] + + +# -- Test 0c helpers -------------------------------------------------------- + + +def _grad_amax_hook(grad, buf): + """Tensor hook that records gradient amax into a buffer. Returns None + to leave the gradient unchanged.""" + buf.copy_(grad.detach().abs().amax().float()) + return None + + +class _TensorHookModule(nn.Module): + """Module that registers a backward tensor hook to capture grad amax.""" + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + self.register_buffer("grad_amax", torch.tensor(0.0)) + + def forward(self, x): + x = self.norm(x) + result = _Level1FP8Linear.apply( + x, + self.weight, + float8_e4m3, + _GRAN_VALUE, + _BACKEND_VALUE, + ) + output = result[0] + buf = self.grad_amax + output.register_hook(lambda g, b=buf: _grad_amax_hook(g, b)) + return output + + +# -- Test 0d helpers -------------------------------------------------------- + + +class _DelayedFP8Linear(torch.autograd.Function): + """Minimal delayed-scaling FP8 linear for feasibility testing. + + Takes 3 pre-computed scales (input, weight, gradient). + Returns (output, a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, + input_amax, weight_amax) with the last 6 mark_non_differentiable. + """ + + @staticmethod + def forward( + input, + weight, + scale_input, + scale_weight, + scale_grad, + fp8_fwd_dtype, + fp8_bwd_dtype, + fp8_fwd_max, + fp8_bwd_max, + gran_value, + backend_value, + ): + out_dtype = input.dtype + orig_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]) + + scale_in_narrow = scale_input.clamp(max=torch.finfo(out_dtype).max).to(out_dtype) + a_fp8 = (input_2d * scale_in_narrow).clamp(-fp8_fwd_max, fp8_fwd_max).to(fp8_fwd_dtype) + a_scale_inv = (1.0 / scale_input).float() + + scale_w_narrow = scale_weight.clamp(max=torch.finfo(out_dtype).max).to(out_dtype) + w_fp8 = (weight * scale_w_narrow).clamp(-fp8_fwd_max, fp8_fwd_max).to(fp8_fwd_dtype) + w_scale_inv = (1.0 / scale_weight).float() + + output = gemm_fp8_impl( + a_fp8, + a_scale_inv, + False, + w_fp8, + w_scale_inv, + True, + out_dtype, + False, + granularity=gran_value, + default_backend=backend_value, + ) + output = output.reshape(*orig_shape[:-1], output.shape[-1]) + + input_amax = input_2d.detach().abs().amax().float() + weight_amax = weight.detach().abs().amax().float() + + a_t_fp8 = a_fp8.t().contiguous() + w_t_fp8 = w_fp8.t().contiguous() + + return (output, a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, input_amax, weight_amax) + + @staticmethod + def setup_context(ctx, inputs, output): + ( + input, + weight, + scale_input, + scale_weight, + scale_grad, + fp8_fwd_dtype, + fp8_bwd_dtype, + fp8_fwd_max, + fp8_bwd_max, + gran_value, + backend_value, + ) = inputs + (output_val, a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, input_amax, weight_amax) = output + + ctx.save_for_backward(a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, scale_grad) + ctx.mark_non_differentiable(a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, input_amax, weight_amax) + ctx.out_dtype = input.dtype + ctx.orig_shape = input.shape + ctx.fp8_bwd_dtype = fp8_bwd_dtype + ctx.fp8_bwd_max = fp8_bwd_max + ctx.gran_value = gran_value + ctx.backend_value = backend_value + + @staticmethod + def backward(ctx, grad_output, *_): + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + a_t_fp8, a_scale_inv, w_t_fp8, w_scale_inv, scale_grad = ctx.saved_tensors + + grad_2d = grad_output.reshape(-1, grad_output.shape[-1]) + + sg_narrow = scale_grad.clamp(max=torch.finfo(ctx.out_dtype).max).to(ctx.out_dtype) + grad_fp8 = (grad_2d * sg_narrow).clamp(-ctx.fp8_bwd_max, ctx.fp8_bwd_max).to(ctx.fp8_bwd_dtype) + grad_scale_inv = (1.0 / scale_grad).float() + + grad_input = gemm_fp8_impl( + grad_fp8, + grad_scale_inv, + False, + w_t_fp8, + w_scale_inv, + True, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + grad_input = grad_input.reshape(ctx.orig_shape) + + grad_t_fp8 = grad_fp8.t().contiguous() + grad_weight = gemm_fp8_impl( + grad_t_fp8, + grad_scale_inv, + False, + a_t_fp8, + a_scale_inv, + True, + ctx.out_dtype, + False, + granularity=ctx.gran_value, + default_backend=ctx.backend_value, + ) + + return (grad_input, grad_weight, None, None, None, None, None, None, None, None, None) + + +class _ProdDelayedModule(nn.Module): + """Drives the *production* ``DelayedFP8LinearTensorwiseFunction``. + + Unlike the earlier ``_DelayedFnModule`` (which wrapped the in-file + ``_DelayedFP8Linear`` replica), this calls the shipped Function so a + regression that introduces a graph break or breaks the fused amax capture + in production is actually caught. The function writes the current amaxes + in-place into the ``staged_*`` buffers during forward (fused capture), so + no side-output ``copy_`` is needed here. + + The scale/amax buffers stay fp32 even though the module runs in bf16: the + fused ``atomic_max`` amax kernel only supports fp32, mirroring how the + production layers keep these buffers fp32. + """ + + def __init__(self, dim, out_dim, force_nt): + super().__init__() + self._force_nt = force_nt + self.norm = nn.LayerNorm(dim) + self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + for name in ("scale_input", "scale_weight", "scale_grad"): + self.register_buffer(name, torch.tensor(1.0, dtype=torch.float32)) + for name in ("staged_input_amax", "staged_weight_amax", "staged_grad_amax"): + self.register_buffer(name, torch.tensor(0.0, dtype=torch.float32)) + + def forward(self, x): + x = self.norm(x) + result = DelayedFP8LinearTensorwiseFunction.apply( + x, + self.weight, + self.scale_input, + self.scale_weight, + self.scale_grad, + self.staged_input_amax, + self.staged_weight_amax, + self.staged_grad_amax, + float8_e4m3, + float8_e5m2, + _GRAN_VALUE, + _BACKEND_VALUE, + self._force_nt, + ) + return result[0] + + +def _build_prod_delayed_module(dim, out_dim, force_nt): + """Construct a ``_ProdDelayedModule`` on CUDA/bf16 with fp32 scale/amax + buffers (the blanket ``.to(bfloat16)`` would otherwise downcast them and + trip the fp32-only fused amax kernel).""" + model = _ProdDelayedModule(dim, out_dim, force_nt).to(dtype=torch.bfloat16, device="cuda") + for name in ( + "scale_input", + "scale_weight", + "scale_grad", + "staged_input_amax", + "staged_weight_amax", + "staged_grad_amax", + ): + model._buffers[name] = model._buffers[name].float() + return model + + +class _DelayedFP8E2EModule(nn.Module): + """Full delayed scaling module with scale update loop for e2e tests.""" + + def __init__(self, dim, out_dim, use_grad_hook=True, history_len=16): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + self._use_grad_hook = use_grad_hook + self._fp8_fwd_max = _FP8_FWD_MAX + self._fp8_bwd_max = _FP8_BWD_MAX + + self.register_buffer("scale_input", torch.tensor(1.0)) + self.register_buffer("scale_weight", torch.tensor(1.0)) + self.register_buffer("scale_grad", torch.tensor(1.0)) + + self.register_buffer("amax_history_input", torch.zeros(history_len)) + self.register_buffer("amax_history_weight", torch.zeros(history_len)) + self.register_buffer("amax_history_grad", torch.zeros(history_len)) + + self.register_buffer("staged_input_amax", torch.tensor(0.0)) + self.register_buffer("staged_weight_amax", torch.tensor(0.0)) + self.register_buffer("staged_grad_amax", torch.tensor(0.0)) + + self._history_idx = 0 + + def update_scales(self): + idx = self._history_idx + self.amax_history_input[idx] = self.staged_input_amax + self.amax_history_weight[idx] = self.staged_weight_amax + self.amax_history_grad[idx] = self.staged_grad_amax + self._history_idx = (idx + 1) % self.amax_history_input.shape[0] + + self._compute_scale(self.scale_input, self.amax_history_input, self._fp8_fwd_max) + self._compute_scale(self.scale_weight, self.amax_history_weight, self._fp8_fwd_max) + self._compute_scale(self.scale_grad, self.amax_history_grad, self._fp8_bwd_max) + + def _compute_scale(self, scale_buf, history, fp8_max): + amax = history.max() + sf = fp8_max / amax.clamp(min=1e-12) + sf = torch.where(amax > 0.0, sf, scale_buf) + sf = torch.where(torch.isfinite(amax), sf, scale_buf) + sf = sf.clamp(max=torch.finfo(torch.float32).max) + scale_buf.fill_(sf) + + def update_grad_amax_from_weight_grad(self): + if self.weight.grad is not None: + self.staged_grad_amax.fill_(self.weight.grad.detach().abs().amax().float()) + + def forward(self, x): + x = self.norm(x) + result = _DelayedFP8Linear.apply( + x, + self.weight, + self.scale_input, + self.scale_weight, + self.scale_grad, + float8_e4m3, + float8_e5m2, + self._fp8_fwd_max, + self._fp8_bwd_max, + _GRAN_VALUE, + _BACKEND_VALUE, + ) + output = result[0] + self.staged_input_amax.copy_(result[5]) + self.staged_weight_amax.copy_(result[6]) + + if self._use_grad_hook: + buf = self.staged_grad_amax + output.register_hook(lambda g, b=buf: _grad_amax_hook(g, b)) + return output + + +def _run_delayed_training_loop(model, inputs, n_steps, lr=1e-3): + optimizer = torch.optim.AdamW(model.parameters(), lr=lr) + mod = model._orig_mod if hasattr(model, "_orig_mod") else model + losses, scale_log = [], [] + for step in range(n_steps): + mod.update_scales() + optimizer.zero_grad() + out = model(inputs[step % len(inputs)]) + loss = out.sum() + loss.backward() + if not mod._use_grad_hook: + mod.update_grad_amax_from_weight_grad() + optimizer.step() + losses.append(loss.item()) + scale_log.append( + ( + mod.scale_input.item(), + mod.scale_weight.item(), + mod.scale_grad.item(), + ) + ) + return losses, scale_log + + +# -- Test 0e helpers -------------------------------------------------------- + + +class _BufferMutationModule(nn.Module): + """Module that uses a buffer in the forward path to test mutation visibility.""" + + def __init__(self, dim, out_dim): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.weight = nn.Parameter(torch.randn(out_dim, dim, dtype=torch.bfloat16)) + self.register_buffer("scale", torch.tensor(1.0)) + + def forward(self, x): + x = self.norm(x) + x_scaled = x * self.scale + result = _Level1FP8Linear.apply( + x_scaled, + self.weight, + float8_e4m3, + _GRAN_VALUE, + _BACKEND_VALUE, + ) + return result[0] + + +# --------------------------------------------------------------------------- +# Test 0a: Buffer copy_() of autograd Function side-output +# --------------------------------------------------------------------------- + + +@requires_cuda +class TestDelayedPhase0a_BufferCopy(PrimusUT): + """Can a compiled forward mutate a registered buffer via copy_()? + + This is the critical gate test for delayed scaling. If buffer copy_() + works inside compiled code, amax can be routed from the autograd + Function to the module without graph breaks. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def test_no_graph_break(self): + """torch._dynamo.explain reports zero graph breaks with buffer copy_().""" + model = _BufferCopyModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda") + x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda") + + explanation = torch._dynamo.explain(model)(x) + + assert explanation.graph_break_count == 0, ( + f"Test 0a: Expected 0 graph breaks, got {explanation.graph_break_count}. " + f"Reasons: {explanation.break_reasons}" + ) + + def test_buffer_updated(self): + """After compiled forward, staged_amax buffer should be non-zero.""" + model = _BufferCopyModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda") + compiled = torch.compile(model) + x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda") + + assert model.staged_amax.item() == 0.0 + _ = compiled(x) + assert model.staged_amax.item() > 0.0, "Buffer copy_() had no effect: staged_amax is still 0.0" + + def test_multistep_compiled_vs_eager(self): + """Multi-step: compiled with buffer copy_() should track eager.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _BufferCopyModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(ref_model) + eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS) + + torch._dynamo.reset() + compiled_model = _BufferCopyModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + compiled_model.load_state_dict(init_state) + compiled_model_c = torch.compile(compiled_model) + compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS) + + _print_comparison("BufferCopy Eager", eager_losses, "BufferCopy Compiled", compiled_losses) + final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"Buffer copy_() compiled diverged from eager: " + f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + +# Test 0c: Tensor hook with buffer copy_() in backward +# --------------------------------------------------------------------------- + + +@requires_cuda +class TestDelayedPhase0c_TensorHook(PrimusUT): + """Does register_hook on an intermediate tensor work under compile, + and can the hook callback mutate a registered buffer? + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def test_no_graph_break(self): + """torch._dynamo.explain reports zero graph breaks with tensor hook.""" + model = _TensorHookModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda") + x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda") + + explanation = torch._dynamo.explain(model)(x) + + print(f"\nTest 0c: graph_break_count = {explanation.graph_break_count}") + if explanation.graph_break_count > 0: + print(f" Break reasons: {explanation.break_reasons}") + assert explanation.graph_break_count == 0, ( + f"Test 0c: Expected 0 graph breaks, got {explanation.graph_break_count}. " + f"Reasons: {explanation.break_reasons}" + ) + + def test_hook_fires_and_updates_buffer(self): + """After forward + backward, grad_amax buffer should be non-zero.""" + model = _TensorHookModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda") + compiled = torch.compile(model) + x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda", requires_grad=True) + + assert model.grad_amax.item() == 0.0 + out = compiled(x) + out.sum().backward() + assert model.grad_amax.item() > 0.0, "Tensor hook did not fire: grad_amax is still 0.0" + + def test_multistep_compiled_vs_eager(self): + """Multi-step: compiled with tensor hook should track eager.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _TensorHookModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(ref_model) + eager_losses, _ = _run_training_loop(ref_model, inputs, _N_STEPS) + + torch._dynamo.reset() + compiled_model = _TensorHookModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + compiled_model.load_state_dict(init_state) + compiled_model_c = torch.compile(compiled_model) + compiled_losses, _ = _run_training_loop(compiled_model_c, inputs, _N_STEPS) + + _print_comparison("TensorHook Eager", eager_losses, "TensorHook Compiled", compiled_losses) + final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"Tensor hook compiled diverged from eager: " + f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + +# --------------------------------------------------------------------------- +# Test 0d-fn: _DelayedFP8Linear autograd Function (isolation) +# --------------------------------------------------------------------------- + + +@requires_cuda +class TestDelayedProductionFunctionCompile(PrimusUT): + """Graph-break + fused-amax-capture contract of the *production* + ``DelayedFP8LinearTensorwiseFunction`` under torch.compile. + + This guards the shipped Function directly (both the native and forced-NT + backward layouts), so a regression that re-introduces a graph break or + breaks the in-place staged-amax capture fails here. Numerical + forward/backward correctness of the same Function is covered separately in + ``test_delayed_fp8_triton_op.py``. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def test_function_no_graph_break(self): + """torch._dynamo.explain on the production module -> 0 breaks (both arms).""" + for force_nt in (True, False): + torch._dynamo.reset() + model = _build_prod_delayed_module(_DIM, _OUT_DIM, force_nt) + x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda") + + explanation = torch._dynamo.explain(model)(x) + + assert explanation.graph_break_count == 0, ( + f"force_nt={force_nt}: expected 0 graph breaks, got " + f"{explanation.graph_break_count}. Reasons: {explanation.break_reasons}" + ) + + def test_amax_side_outputs_populated(self): + """Compiled forward must capture the current amaxes into the staged + buffers in-place (both arms).""" + for force_nt in (True, False): + torch._dynamo.reset() + model = _build_prod_delayed_module(_DIM, _OUT_DIM, force_nt) + compiled = torch.compile(model) + x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda") + + assert model.staged_input_amax.item() == 0.0 + assert model.staged_weight_amax.item() == 0.0 + _ = compiled(x) + assert ( + model.staged_input_amax.item() > 0.0 + ), f"force_nt={force_nt}: staged_input_amax not populated" + assert ( + model.staged_weight_amax.item() > 0.0 + ), f"force_nt={force_nt}: staged_weight_amax not populated" + # The captured amax must match the actual input/weight amax (fused + # capture is the delayed-scaling contract, not an arbitrary write). + xn = model.norm(x) + expected_in = xn.detach().float().abs().amax().item() + expected_w = model.weight.detach().float().abs().amax().item() + assert abs(model.staged_input_amax.item() - expected_in) / max(expected_in, 1e-8) < 1e-2 + assert abs(model.staged_weight_amax.item() - expected_w) / max(expected_w, 1e-8) < 1e-2 + + +# --------------------------------------------------------------------------- +# Test 0d-e2e: Full delayed module with scale update loop +# --------------------------------------------------------------------------- + + +@requires_cuda +class TestDelayedPhase0d_E2E(PrimusUT): + """End-to-end delayed scaling: multi-step training with scale updates. + + Tests the complete interaction between eager scale updates and + compiled forward/backward, including scale convergence and + loss tracking against dynamic scaling. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def test_scale_convergence_eager(self): + """Scales must change from initial 1.0 within 3 steps (eager). + + Uses history_len=1 (most_recent scaling) so scales respond to + the very first recorded amax rather than being dominated by the + 16-slot history initialization. + """ + model = _DelayedFP8E2EModule(_CONV_DIM, _CONV_OUT, history_len=1).to( + dtype=torch.bfloat16, device="cuda" + ) + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + _, scale_log = _run_delayed_training_loop(model, inputs, 10) + + for i, (si, sw, sg) in enumerate(scale_log): + print(f" step {i}: scale_input={si:.4f}, " f"scale_weight={sw:.4f}, scale_grad={sg:.4f}") + + s_in_3, s_w_3, s_g_3 = scale_log[3] + assert s_in_3 != 1.0, f"scale_input stuck at 1.0 after 3 steps" + assert s_w_3 != 1.0, f"scale_weight stuck at 1.0 after 3 steps" + assert s_g_3 != 1.0, f"scale_grad stuck at 1.0 after 3 steps" + + for i, (si, sw, sg) in enumerate(scale_log): + assert math.isfinite(si) and si > 0, f"scale_input non-finite/negative at step {i}: {si}" + assert math.isfinite(sw) and sw > 0, f"scale_weight non-finite/negative at step {i}: {sw}" + assert math.isfinite(sg) and sg > 0, f"scale_grad non-finite/negative at step {i}: {sg}" + + def test_compiled_delayed_vs_eager_delayed(self): + """Compiled delayed should track eager delayed over 200 steps.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + ref_model = _DelayedFP8E2EModule(_CONV_DIM, _CONV_OUT, history_len=1).to( + dtype=torch.bfloat16, device="cuda" + ) + init_state = _clone_state(ref_model) + eager_losses, _ = _run_delayed_training_loop(ref_model, inputs, _N_STEPS) + + torch._dynamo.reset() + compiled_model = _DelayedFP8E2EModule(_CONV_DIM, _CONV_OUT, history_len=1).to( + dtype=torch.bfloat16, device="cuda" + ) + compiled_model.load_state_dict(init_state) + compiled_model_c = torch.compile(compiled_model) + compiled_losses, _ = _run_delayed_training_loop(compiled_model_c, inputs, _N_STEPS) + + _print_comparison("Delayed Eager", eager_losses, "Delayed Compiled", compiled_losses) + final_rel = abs(eager_losses[-1] - compiled_losses[-1]) / max(abs(eager_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.05, ( + f"Delayed compiled diverged from delayed eager: " + f"eager={eager_losses[-1]:.6f}, compiled={compiled_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + def test_compiled_delayed_vs_compiled_dynamic(self): + """Compiled delayed should be within 10% of compiled dynamic.""" + torch._dynamo.reset() + inputs = _make_inputs(20, _CONV_DIM, _CONV_BATCH) + + dynamic_model = _OpaqueModule(_CONV_DIM, _CONV_OUT).to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(dynamic_model) + dynamic_model_c = torch.compile(dynamic_model) + dynamic_losses, _ = _run_training_loop(dynamic_model_c, inputs, _N_STEPS) + + torch._dynamo.reset() + delayed_model = _DelayedFP8E2EModule(_CONV_DIM, _CONV_OUT, history_len=1).to( + dtype=torch.bfloat16, device="cuda" + ) + delayed_model.norm.load_state_dict( + {"weight": init_state["norm.weight"], "bias": init_state["norm.bias"]} + ) + delayed_model.weight.data.copy_(init_state["weight"]) + delayed_model_c = torch.compile(delayed_model) + delayed_losses, _ = _run_delayed_training_loop(delayed_model_c, inputs, _N_STEPS) + + _print_comparison("Dynamic Compiled", dynamic_losses, "Delayed Compiled", delayed_losses) + final_rel = abs(dynamic_losses[-1] - delayed_losses[-1]) / max(abs(dynamic_losses[-1]), 1e-12) + print(f"\nFinal loss relative diff: {final_rel:.6f}") + assert final_rel < 0.10, ( + f"Delayed diverged too far from dynamic: " + f"dynamic={dynamic_losses[-1]:.6f}, " + f"delayed={delayed_losses[-1]:.6f}, " + f"rel_diff={final_rel:.6f}" + ) + + +# --------------------------------------------------------------------------- +# Test 0e: Buffer mutation between compiled invocations +# --------------------------------------------------------------------------- + + +@requires_cuda +class TestDelayedPhase0e_BufferMutation(PrimusUT): + """Verifies that buffer mutations between compiled calls are visible. + + When an eager hook mutates a registered buffer between calls to a + compiled forward, the compiled forward must read the new value. + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + pass + + def setup_method(self, method): + torch._dynamo.reset() + + def test_buffer_mutation_between_compiled_calls(self): + """Buffer fill_() between compiled calls must change the output.""" + model = _BufferMutationModule(_DIM, _OUT_DIM).to(dtype=torch.bfloat16, device="cuda") + compiled = torch.compile(model) + x = torch.randn(_BATCH, _DIM, dtype=torch.bfloat16, device="cuda") + + out1 = compiled(x) + model.scale.fill_(2.0) + out2 = compiled(x) + + assert not torch.allclose( + out1, out2 + ), "Buffer mutation not reflected: compiled forward cached old value" diff --git a/tests/unit_tests/backends/megatron/diffusion/test_te_vs_local_spec_attention.py b/tests/unit_tests/backends/megatron/diffusion/test_te_vs_local_spec_attention.py new file mode 100644 index 000000000..5b37b6281 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_te_vs_local_spec_attention.py @@ -0,0 +1,381 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Numerical equivalence tests: TE spec vs Primus-Turbo local spec attention. + +Unlike the earlier version of this file, which compared in-file replicas of the +attention kernels, these tests build the *production* ``core_attention`` modules +returned by the two spec providers and feed them identical inputs: + + * TE spec -> ``TEDotProductAttention`` (FusedAttention/CK, SBHD native) + * local spec -> ``PrimusTurboLocalAttention`` (Primus-Turbo flash attention) + +Both modules are constructed via ``build_module(...)`` with the exact keyword +arguments Megatron's ``transformer/attention.py`` uses when it wires up +``submodules.core_attention``, so a regression in how either spec selects or +configures its attention kernel is caught here. + +The (formerly user-visible) contiguous-vs-non-contiguous split is now a +production-internal detail: ``PrimusTurboLocalAttention`` decides whether to +force a contiguous BSHD copy from the device capability (``force_contiguous_qkv`` +on gfx942). The four tests below therefore exercise the two real production +paths in eager and compiled form rather than permutations of hand-rolled +replicas. + +NOTE: this module is GPU-only (``skip_if_no_cuda()``) and additionally requires +the Megatron + Primus-Turbo stack that only imports inside the ROCm training +container, so it is validated on the GPU runner lane, not in CPU CI. +""" + +import copy +import os + +os.environ["NVTE_FUSED_ATTN"] = "1" +os.environ["NVTE_FUSED_ATTN_CK"] = "1" +os.environ["NVTE_CK_USES_FWD_V3"] = "1" +os.environ["NVTE_CK_USES_BWD_V3"] = "1" + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +import torch.nn as nn + +# NOTE: the diffusion ``conftest`` installs Primus' aiter RTLD_DEEPBIND import +# hook (``install_aiter_deepbind_hook``) before this module loads, so Turbo's +# attention backward binds the pinned ``aiter::mha_bwd`` instead of the stale one +# vendored by transformer_engine below (ROCm/aiter#1332). Without it, attention +# backward crashes on gfx942/gfx950. +import transformer_engine.pytorch # noqa: F401 (import side effects only) +from megatron.core.extensions.transformer_engine import TEDotProductAttention +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import build_module +from megatron.core.transformer.transformer_config import TransformerConfig + +from primus.backends.megatron.core.extensions.primus_turbo_local_spec import ( + PrimusTurboLocalAttention, +) +from tests.utils import PrimusUT + +_has_cuda = torch.cuda.is_available() +requires_cuda = pytest.mark.skipif(not _has_cuda, reason="CUDA required") + +# head_dim=128 is required for FAv3 eligibility in AITER. +# seq=256 > 128 so the Python API fmha_v3_fwd path is taken. +_DIM = 512 +_HEADS = 4 +_HEAD_DIM = _DIM // _HEADS # 128 +_SEQ = 256 +_BATCH = 2 +_N_STEPS = 100 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _compute_snr(x: torch.Tensor, y: torch.Tensor) -> float: + x, y = x.float(), y.float() + signal_power = torch.norm(x).pow(2) + noise_power = torch.norm(x - y).pow(2) + return 10 * torch.log10(signal_power / (noise_power + 1e-12)).detach().item() + + +def _make_sbhd_inputs(n, seq, dim, batch, seed=42): + g = torch.Generator(device="cpu").manual_seed(seed) + return [torch.randn(seq, batch, dim, dtype=torch.bfloat16, generator=g).cuda() for _ in range(n)] + + +def _clone_state(model): + return copy.deepcopy(model.state_dict()) + + +def _run_training_loop(model, inputs, n_steps, lr=1e-3): + optimizer = torch.optim.AdamW(model.parameters(), lr=lr) + losses = [] + for step in range(n_steps): + optimizer.zero_grad() + out = model(inputs[step % len(inputs)]) + loss = out.sum() + loss.backward() + optimizer.step() + losses.append(loss.item()) + return losses + + +def _print_comparison(label_a, losses_a, label_b, losses_b, milestones=None): + if milestones is None: + milestones = [0, 1, 5, 10, 20, 50, 99] + milestones = [m for m in milestones if m < len(losses_a) and m < len(losses_b)] + print(f"\n{'Step':>6} | {label_a:>20} | {label_b:>20} | {'Rel Diff':>10}") + print("-" * 65) + for m in milestones: + a, b = losses_a[m], losses_b[m] + rel = abs(a - b) / max(abs(a), 1e-12) + print(f"{m:>6} | {a:>20.6f} | {b:>20.6f} | {rel:>10.6f}") + + +def _make_attention_config() -> TransformerConfig: + """Minimal TransformerConfig matching the test attention dims. + + ``kv_channels`` is pinned to ``_HEAD_DIM`` and ``softmax_scale`` is left at + its default (None) so both production attention modules fall back to the same + ``1/sqrt(head_dim)`` scale, exactly as they do in a real Flux layer. + """ + return TransformerConfig( + num_layers=1, + hidden_size=_DIM, + num_attention_heads=_HEADS, + kv_channels=_HEAD_DIM, + attention_dropout=0.0, + hidden_dropout=0.0, + tensor_model_parallel_size=1, + context_parallel_size=1, + bf16=True, + params_dtype=torch.bfloat16, + ) + + +def _build_core_attention(core_attention_cls): + """Build a production ``core_attention`` module the way Megatron's + ``attention.py`` does (``build_module`` with the same kwargs). + """ + config = _make_attention_config() + return build_module( + core_attention_cls, + config=config, + layer_number=1, + attn_mask_type=AttnMaskType.no_mask, + attention_type="self", + softmax_scale=config.softmax_scale, + ) + + +# --------------------------------------------------------------------------- +# Wrapper that swaps only the production core_attention kernel +# --------------------------------------------------------------------------- + + +class _SpecAttentionModule(nn.Module): + """norm -> qkv -> production core_attention -> proj. + + The norm/qkv/proj linears are identical across instances; the only thing + that differs is ``core_attention_cls`` (the class a spec provider's + ``core_attention()`` returns), so a numerical divergence is attributable to + the attention kernel/wrapper alone. + """ + + def __init__(self, dim, num_heads, core_attention_cls): + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.norm = nn.LayerNorm(dim) + self.qkv = nn.Linear(dim, 3 * dim, bias=False) + self.proj = nn.Linear(dim, dim, bias=False) + self.core = _build_core_attention(core_attention_cls) + + def forward(self, x): + S, B, D = x.shape + x = self.norm(x) + qkv = self.qkv(x).view(S, B, 3, self.num_heads, self.head_dim) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + # Production core_attention forward: (q, k, v, attention_mask, attn_mask_type) + # q/k/v are SBHD [S, B, H, head_dim]; output is merged [S, B, H*head_dim]. + out = self.core(q, k, v, None, AttnMaskType.no_mask) + out = out.reshape(S, B, D) + return self.proj(out) + + +def _te_module(): + return _SpecAttentionModule(_DIM, _HEADS, TEDotProductAttention) + + +def _local_module(): + return _SpecAttentionModule(_DIM, _HEADS, PrimusTurboLocalAttention) + + +# --------------------------------------------------------------------------- +# Test class +# --------------------------------------------------------------------------- + + +@requires_cuda +class TestTEvsLocalSpecAttention(PrimusUT): + """Numerical equivalence: TE spec vs Primus-Turbo local spec attention. + + Part A: Single-pass SNR (fast, precise kernel-level comparison). + Part B: Eager training-loop convergence (detects accumulated drift). + Part C: Compiled vs eager for the local spec. + Part D: Cross-path end-to-end (TE eager vs local compiled). + """ + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + # init_parallel_state sets up TP=1 parallel state, the TP RNG tracker, + # and Megatron global args (enable_turbo_attention_float8=False), all of + # which PrimusTurboLocalAttention.__init__ relies on. + pass + + def setup_method(self, method): + torch._dynamo.reset() + + # ----------------------------------------------------------------------- + # Part A: Single-forward-pass SNR + # ----------------------------------------------------------------------- + + def _run_single_pass(self, build_fn): + """Run one forward + backward, return output, input grad, and state.""" + torch.manual_seed(123) + torch.cuda.manual_seed_all(123) + + model = build_fn().to(dtype=torch.bfloat16, device="cuda") + + x = torch.randn(_SEQ, _BATCH, _DIM, dtype=torch.bfloat16, device="cuda", requires_grad=True) + grad_out = torch.randn(_SEQ, _BATCH, _DIM, dtype=torch.bfloat16, device="cuda") + + out = model(x) + out.backward(grad_out) + torch.cuda.synchronize() + return out.detach(), x.grad.detach(), model.state_dict() + + def _run_single_pass_with_state(self, build_fn, state_dict): + """Run one forward + backward with pre-loaded norm/qkv/proj weights.""" + torch.manual_seed(123) + torch.cuda.manual_seed_all(123) + + model = build_fn().to(dtype=torch.bfloat16, device="cuda") + own_keys = set(model.state_dict().keys()) + filtered = {k: v for k, v in state_dict.items() if k in own_keys} + model.load_state_dict(filtered, strict=False) + + x = torch.randn(_SEQ, _BATCH, _DIM, dtype=torch.bfloat16, device="cuda", requires_grad=True) + grad_out = torch.randn(_SEQ, _BATCH, _DIM, dtype=torch.bfloat16, device="cuda") + + out = model(x) + out.backward(grad_out) + torch.cuda.synchronize() + return out.detach(), x.grad.detach() + + def _compare_snr(self, label, out_a, grad_a, out_b, grad_b, threshold=40.0): + out_snr = _compute_snr(out_a, out_b) + grad_snr = _compute_snr(grad_a, grad_b) + print(f"\n [{label}] output SNR={out_snr:.2f} dB, grad SNR={grad_snr:.2f} dB") + assert out_snr > threshold, f"[{label}] output SNR too low: {out_snr:.2f}" + assert grad_snr > threshold, f"[{label}] grad SNR too low: {grad_snr:.2f}" + + def test_single_pass_te_vs_local(self): + """TE-spec and local-spec attention produce near-identical single-pass + output and input gradients given identical weights/inputs.""" + out_te, grad_te, state_te = self._run_single_pass(_te_module) + out_local, grad_local = self._run_single_pass_with_state(_local_module, state_te) + self._compare_snr("TE vs Local", out_te, grad_te, out_local, grad_local) + + # ----------------------------------------------------------------------- + # Part B: Eager training-loop convergence + # ----------------------------------------------------------------------- + + def test_eager_te_vs_local(self): + """Eager: TE-spec and local-spec attention converge to the same loss.""" + torch._dynamo.reset() + inputs = _make_sbhd_inputs(20, _SEQ, _DIM, _BATCH) + + model_te = _te_module().to(dtype=torch.bfloat16, device="cuda") + init_state = _clone_state(model_te) + losses_te = _run_training_loop(model_te, inputs, _N_STEPS) + + model_local = _local_module().to(dtype=torch.bfloat16, device="cuda") + own_keys = set(model_local.state_dict().keys()) + filtered = {k: v for k, v in init_state.items() if k in own_keys} + model_local.load_state_dict(filtered, strict=False) + losses_local = _run_training_loop(model_local, inputs, _N_STEPS) + + _print_comparison("TE Eager", losses_te, "Local Eager", losses_local) + + final_rel = abs(losses_te[-1] - losses_local[-1]) / max(abs(losses_te[-1]), 1e-12) + print(f"\n Final loss relative diff: {final_rel:.6f}") + assert final_rel < 0.02, ( + f"TE vs Local eager: final loss diverged: " + f"te={losses_te[-1]:.6f}, local={losses_local[-1]:.6f}, rel={final_rel:.6f}" + ) + + # ----------------------------------------------------------------------- + # Part C/D: Compiled local spec (in-process) + # + # An earlier revision ran the compiled training in a subprocess because + # ``torch.compile`` + ``allow_in_graph(AiterFlashAttnFunc)`` could corrupt + # AOT-autograd view-replay metadata when sharing a process with the rest of + # the pytest/Megatron harness. That is avoided here by resetting Dynamo + # (``torch._dynamo.reset()``) immediately before each compiled build, so the + # local-attention graph is traced from a clean state regardless of what + # earlier tests in the session compiled. The compiled run builds the + # *production* ``PrimusTurboLocalAttention`` via the same ``_local_module`` + # helper as the eager paths, so the two paths differ only by compilation. + # ----------------------------------------------------------------------- + + @staticmethod + def _run_compiled_local(n_steps: int = 100, init_state=None): + """Compile + train the production local-spec attention in-process. + + Resets Dynamo and allows ``AiterFlashAttnFunc`` into the graph, then + compiles a fresh ``_local_module`` and runs the training loop. When + ``init_state`` is given the model is seeded with it so a comparison run + shares identical norm/qkv/proj weights; otherwise the freshly + initialized weights are captured. Returns ``(losses, init_state)``. + """ + from primus_turbo.pytorch.ops.attention.flash_attn_interface import ( + AiterFlashAttnFunc, + ) + + torch._dynamo.reset() + torch._dynamo.allow_in_graph(AiterFlashAttnFunc) + + model = _local_module().to(dtype=torch.bfloat16, device="cuda") + if init_state is not None: + own_keys = set(model.state_dict().keys()) + model.load_state_dict({k: v for k, v in init_state.items() if k in own_keys}, strict=False) + captured_state = _clone_state(model) + + inputs = _make_sbhd_inputs(20, _SEQ, _DIM, _BATCH) + compiled = torch.compile(model) + losses = _run_training_loop(compiled, inputs, n_steps) + return losses, captured_state + + def test_compiled_local_vs_eager(self): + """Compiled local-spec attention converges like eager local-spec.""" + losses_compiled, init_state = self._run_compiled_local(_N_STEPS) + + torch._dynamo.reset() + model_eager = _local_module().to(dtype=torch.bfloat16, device="cuda") + own_keys = set(model_eager.state_dict().keys()) + model_eager.load_state_dict({k: v for k, v in init_state.items() if k in own_keys}, strict=False) + inputs = _make_sbhd_inputs(20, _SEQ, _DIM, _BATCH) + losses_eager = _run_training_loop(model_eager, inputs, _N_STEPS) + + _print_comparison("Local Eager", losses_eager, "Local Compiled", losses_compiled) + + final_rel = abs(losses_eager[-1] - losses_compiled[-1]) / max(abs(losses_eager[-1]), 1e-12) + print(f"\n Final loss relative diff: {final_rel:.6f}") + assert final_rel < 0.02, f"Diverged: {final_rel:.6f}" + + def test_te_eager_vs_local_compiled(self): + """End-to-end cross-path: TE-spec eager and local-spec compiled, sharing + initial weights + inputs, converge to the same loss.""" + losses_local, init_state = self._run_compiled_local(_N_STEPS) + + torch._dynamo.reset() + model_te = _te_module().to(dtype=torch.bfloat16, device="cuda") + own_keys = set(model_te.state_dict().keys()) + model_te.load_state_dict({k: v for k, v in init_state.items() if k in own_keys}, strict=False) + inputs = _make_sbhd_inputs(20, _SEQ, _DIM, _BATCH) + losses_te = _run_training_loop(model_te, inputs, _N_STEPS) + + _print_comparison("TE Eager", losses_te, "Local Compiled", losses_local) + + final_rel = abs(losses_te[-1] - losses_local[-1]) / max(abs(losses_te[-1]), 1e-12) + print(f"\n Final loss relative diff: {final_rel:.6f}") + assert final_rel < 0.02, f"Diverged: {final_rel:.6f}" From 4648a4aa1b92977218dc4aec9856bedec62e0b86 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Tue, 14 Jul 2026 10:28:16 +0300 Subject: [PATCH 024/127] feat(flux): diffusion data pipeline (energon/synthetic providers, encoders) (#812) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/flux` — review after it. ## What this changes The diffusion data layer: dataloader + CUDA prefetch, energon and synthetic dataset providers, the image (VAE) and text (CLIP-L, T5-XXL) encoders, and the image task encoder. Also moves the energon/webdataset pins into `requirements.txt` here, since this layer's source and tests need them at runtime. ## Why it's stacked here One synthetic-dataset module top-level-imports `flux.utils`, so it bases on `feat/flux/flux` (this also transitively re-parents the prep/trainers/mlperf layers onto the Flux branch). ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/flux`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion/data` (needs `megatron-energon`/`webdataset` from the moved `requirements.txt`). Validated locally on an AMD GPU container: 50 passed. ## Files 32 (dataloader/prefetch, dataset providers, VAE/CLIP/T5 encoders, task encoder, `requirements.txt` + tests). --------- Co-authored-by: Flux Split Trial Co-authored-by: eshaw2 Co-authored-by: luiza-amd Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- primus/backends/megatron/data/__init__.py | 39 ++ .../backends/megatron/data/cuda_prefetch.py | 84 +++ primus/backends/megatron/data/dataloader.py | 163 +++++ .../megatron/data/dataset_provider.py | 63 ++ .../megatron/data/diffusion/__init__.py | 11 + .../data/diffusion/encoders/__init__.py | 162 +++++ .../megatron/data/diffusion/encoders/base.py | 285 ++++++++ .../data/diffusion/encoders/config.py | 160 +++++ .../data/diffusion/encoders/image/__init__.py | 15 + .../encoders/image/autoencoder_kl.py | 270 ++++++++ .../data/diffusion/encoders/text/__init__.py | 14 + .../data/diffusion/encoders/text/clip_l.py | 254 ++++++++ .../data/diffusion/encoders/text/t5_xxl.py | 229 +++++++ .../data/diffusion/task_encoders/__init__.py | 22 + .../data/diffusion/task_encoders/image.py | 396 +++++++++++ .../megatron/data/energon_dataset_provider.py | 282 ++++++++ .../megatron/data/synthetic/__init__.py | 32 + .../megatron/data/synthetic/mock_datasets.py | 614 ++++++++++++++++++ .../data/synthetic_dataset_provider.py | 217 +++++++ requirements.txt | 4 + .../megatron/diffusion/data/__init__.py | 2 + .../megatron/diffusion/data/conftest.py | 257 ++++++++ .../diffusion/data/encoders/__init__.py | 2 + .../diffusion/data/encoders/test_base.py | 130 ++++ .../diffusion/data/encoders/test_clip.py | 60 ++ .../diffusion/data/encoders/test_config.py | 30 + .../test_encoder_wrappers_consolidated.py | 405 ++++++++++++ .../diffusion/data/encoders/test_t5.py | 57 ++ .../diffusion/data/encoders/test_vae.py | 58 ++ .../diffusion/data/task_encoders/__init__.py | 2 + .../data/task_encoders/test_task_encoders.py | 192 ++++++ .../diffusion/data/test_synthetic_datasets.py | 129 ++++ .../test_megatron_bridge_adapter.py | 50 +- 33 files changed, 4674 insertions(+), 16 deletions(-) create mode 100644 primus/backends/megatron/data/__init__.py create mode 100644 primus/backends/megatron/data/cuda_prefetch.py create mode 100644 primus/backends/megatron/data/dataloader.py create mode 100644 primus/backends/megatron/data/dataset_provider.py create mode 100644 primus/backends/megatron/data/diffusion/__init__.py create mode 100644 primus/backends/megatron/data/diffusion/encoders/__init__.py create mode 100644 primus/backends/megatron/data/diffusion/encoders/base.py create mode 100644 primus/backends/megatron/data/diffusion/encoders/config.py create mode 100644 primus/backends/megatron/data/diffusion/encoders/image/__init__.py create mode 100644 primus/backends/megatron/data/diffusion/encoders/image/autoencoder_kl.py create mode 100644 primus/backends/megatron/data/diffusion/encoders/text/__init__.py create mode 100644 primus/backends/megatron/data/diffusion/encoders/text/clip_l.py create mode 100644 primus/backends/megatron/data/diffusion/encoders/text/t5_xxl.py create mode 100644 primus/backends/megatron/data/diffusion/task_encoders/__init__.py create mode 100644 primus/backends/megatron/data/diffusion/task_encoders/image.py create mode 100644 primus/backends/megatron/data/energon_dataset_provider.py create mode 100644 primus/backends/megatron/data/synthetic/__init__.py create mode 100644 primus/backends/megatron/data/synthetic/mock_datasets.py create mode 100644 primus/backends/megatron/data/synthetic_dataset_provider.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/__init__.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/conftest.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/encoders/__init__.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/encoders/test_base.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/encoders/test_clip.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/encoders/test_config.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/encoders/test_encoder_wrappers_consolidated.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/encoders/test_t5.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/encoders/test_vae.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/task_encoders/__init__.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/test_synthetic_datasets.py diff --git a/primus/backends/megatron/data/__init__.py b/primus/backends/megatron/data/__init__.py new file mode 100644 index 000000000..48c0e16f5 --- /dev/null +++ b/primus/backends/megatron/data/__init__.py @@ -0,0 +1,39 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Megatron data loading infrastructure. + +This module provides: + - DatasetProvider abstraction for pluggable data pipelines + - EnergonDatasetProvider for multimodal/diffusion (Energon) + - SyntheticDatasetProvider for mock/synthetic data + - MegatronDataloaderWrapper for generic dataloader compatibility + - Synthetic datasets for testing and development + - Diffusion-specific task encoders and preprocessing + +Architecture: + The strategy pattern allows different trainers to use different + data sources while sharing the same training infrastructure. +""" + +_LAZY_IMPORTS = { + "DatasetProvider": ".dataset_provider", + "EnergonDatasetProvider": ".energon_dataset_provider", + "SyntheticDatasetProvider": ".synthetic_dataset_provider", + "MegatronDataloaderWrapper": ".dataloader", +} + + +def __getattr__(name): + if name in _LAZY_IMPORTS: + import importlib + + module = importlib.import_module(_LAZY_IMPORTS[name], __name__) + value = getattr(module, name) + globals()[name] = value + return value + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = list(_LAZY_IMPORTS.keys()) diff --git a/primus/backends/megatron/data/cuda_prefetch.py b/primus/backends/megatron/data/cuda_prefetch.py new file mode 100644 index 000000000..3f4c35626 --- /dev/null +++ b/primus/backends/megatron/data/cuda_prefetch.py @@ -0,0 +1,84 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +CUDA stream-based data prefetcher for overlapping HtoD transfers with compute. + +Wraps an iterator (typically RerunDataIterator -> MegatronDataloaderWrapper -> +Energon SavableDataLoader) and transfers the next batch to GPU on a dedicated +secondary stream. + +Each ``__next__`` call waits for the in-flight HtoD to complete, returns the +GPU batch, and immediately kicks off the *next* CPU fetch + HtoD dispatch so +the transfer overlaps with forward/backward compute. + +Requirements: + - Upstream DataLoader must use pin_memory=True for truly async non_blocking + transfers. Energon's SavableDataLoader satisfies this. + - Batches must be dict[str, Tensor | Any]. +""" + +import torch + + +class CudaPrefetchIterator: + """Prefetch data batches to GPU on a secondary CUDA stream. + + On construction, eagerly fetches the first batch and dispatches its HtoD + transfer (cold start). Each subsequent ``__next__`` waits for the previous + HtoD, grabs the GPU batch, kicks off the *next* prefetch, and returns. + + ``wait_stream()`` (not CUDA events) is used because prefetch depth is 1. + ``record_stream()`` is not needed because ``self._next_batch`` holds GPU + tensor references until consumed, and the caller holds the returned batch + through forward/backward. + """ + + def __init__(self, iterator, compute_dtype=torch.bfloat16): + if isinstance(iterator, (list, tuple)): + # Virtual pipeline parallel hands Megatron a list of per-chunk + # iterators; this single-stream prefetcher only wraps one iterator. + raise TypeError( + "CudaPrefetchIterator expects a single data iterator, got " + f"{type(iterator).__name__}. Skip prefetch for virtual pipeline " + "parallel or wrap each per-chunk iterator separately." + ) + self._iterator = iterator + self._stream = torch.cuda.Stream() + self._dtype = compute_dtype + self._next_batch = None + self._prefetch() + + def _prefetch(self): + """Fetch next batch from upstream and dispatch HtoD on secondary stream.""" + try: + batch = next(self._iterator) + except StopIteration: + self._next_batch = None + return + with torch.cuda.stream(self._stream): + gpu_batch = {} + for k, v in batch.items(): + if isinstance(v, torch.Tensor): + if v.is_floating_point(): + gpu_batch[k] = v.to(dtype=self._dtype, device="cuda", non_blocking=True) + else: + gpu_batch[k] = v.cuda(non_blocking=True) + else: + gpu_batch[k] = v + self._next_batch = gpu_batch + + def __next__(self): + torch.cuda.current_stream().wait_stream(self._stream) + batch = self._next_batch + self._next_batch = None + if batch is None: + raise StopIteration + self._prefetch() + return batch + + def __iter__(self): + return self diff --git a/primus/backends/megatron/data/dataloader.py b/primus/backends/megatron/data/dataloader.py new file mode 100644 index 000000000..4a53ceb13 --- /dev/null +++ b/primus/backends/megatron/data/dataloader.py @@ -0,0 +1,163 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. +# +# Adapted from Megatron-LM multimodal examples +# Reference: examples/multimodal/dataloader_provider.py + +""" +Generic Megatron-compatible dataloader wrapper. + +This module provides utilities to adapt any iterable (PyTorch DataLoader, +Megatron Energon loader, synthetic data, etc.) to work with Megatron's +training loop. + +Key Features: + - Cyclic iteration (never raises StopIteration) + - Optional checkpoint support via duck typing + - Works with any Python iterable + - No dependencies on specific data sources + +Historical Note: + Originally named "EnergonDataloader" but has no Energon dependencies. + Renamed to MegatronDataloaderWrapper for clarity. +""" + +import logging +from typing import Any, Iterator + +logger = logging.getLogger(__name__) + + +def cyclic_iter(iterator: Iterator) -> Iterator: + """ + Create a cyclic iterator that restarts when exhausted. + + Megatron's training loop expects infinite iterators that never + raise StopIteration. This wrapper cycles any iterator infinitely. + + Args: + iterator: Any Python iterator + + Yields: + Items from iterator, cycling infinitely + + Example: + >>> loader = DataLoader(dataset, batch_size=4) + >>> infinite_loader = cyclic_iter(loader) + >>> # Never raises StopIteration + + Reference: + Megatron-LM examples/multimodal/dataloader_provider.py:cyclic_iter + """ + while True: + for item in iterator: + yield item + + +class MegatronDataloaderWrapper: + """ + Generic wrapper to make any iterable compatible with Megatron training loop. + + This wrapper is completely generic and has NO dependencies on specific + data sources. It works with: + - PyTorch DataLoader (for synthetic/mock data) + - Megatron Energon loaders (for WebDataset) + - Any Python iterable + + Features: + 1. Cyclic iteration - never raises StopIteration + 2. Optional state management - duck-typed for compatibility + 3. Megatron training loop compatibility + + The wrapper uses duck typing for checkpoint methods. If your dataloader + has save_state_rank() or restore_state_rank(), they'll be used. + Otherwise, state operations are no-ops (perfect for synthetic data). + + Example with PyTorch DataLoader (synthetic data): + >>> from torch.utils.data import DataLoader + >>> loader = DataLoader(mock_dataset, batch_size=4) + >>> megatron_loader = MegatronDataloaderWrapper(loader) + >>> for batch in megatron_loader: + ... # Never exhausts, cycles infinitely + ... pass + + Example with Energon (real data): + >>> from megatron.energon import get_loader + >>> energon_loader = get_loader(...) + >>> megatron_loader = MegatronDataloaderWrapper(energon_loader) + >>> # Checkpoint support via save_state_rank() is available + + Reference: + Megatron-LM examples/multimodal/dataloader_provider.py:EnergonDataloader + """ + + def __init__(self, dataloader): + """ + Initialize wrapper for any iterable. + + Args: + dataloader: Any iterable (PyTorch DataLoader, Energon loader, etc.) + """ + self._dataloader = dataloader + self._iter = iter(cyclic_iter(dataloader)) + logger.debug(f"Initialized MegatronDataloaderWrapper for {type(dataloader).__name__}") + + def __next__(self): + """Get next batch (never raises StopIteration).""" + return next(self._iter) + + def __iter__(self): + """Return iterator.""" + return self._iter + + def save_state(self) -> Any: + """ + Save dataloader state for checkpointing (if supported). + + Uses duck typing - checks for save_state_rank() method. + Returns None if not supported (e.g., for synthetic/mock data). + + Returns: + Dataloader state dictionary, or None if not supported + """ + if hasattr(self._dataloader, "save_state_rank"): + return self._dataloader.save_state_rank() + else: + logger.debug( + f"{type(self._dataloader).__name__} does not support save_state_rank() " + f"(OK for synthetic data)" + ) + return None + + def restore_state(self, state: Any): + """ + Restore dataloader state from checkpoint (if supported). + + Uses duck typing - checks for restore_state_rank() method. + No-op if not supported (e.g., for synthetic/mock data). + + Args: + state: Dataloader state dictionary + """ + if hasattr(self._dataloader, "restore_state_rank"): + self._dataloader.restore_state_rank(state) + # Recreate iterator after restore + self._iter = iter(cyclic_iter(self._dataloader)) + logger.info("Restored dataloader state from checkpoint") + else: + logger.debug( + f"{type(self._dataloader).__name__} does not support restore_state_rank() " + f"(OK for synthetic data)" + ) + + +# Backwards compatibility alias (DEPRECATED) +# NOTE: Deprecated alias kept for backward compatibility. +EnergonDataloader = MegatronDataloaderWrapper + + +__all__ = [ + "MegatronDataloaderWrapper", + "cyclic_iter", + "EnergonDataloader", # Deprecated, use MegatronDataloaderWrapper +] diff --git a/primus/backends/megatron/data/dataset_provider.py b/primus/backends/megatron/data/dataset_provider.py new file mode 100644 index 000000000..4955338b9 --- /dev/null +++ b/primus/backends/megatron/data/dataset_provider.py @@ -0,0 +1,63 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Abstract DatasetProvider for Megatron trainers. + +This module provides a strategy pattern for pluggable data pipelines, +allowing different trainers to use different data sources while sharing +the same training infrastructure. +""" + +from abc import ABC, abstractmethod +from typing import Any, List, Optional, Tuple + + +class DatasetProvider(ABC): + """ + Abstract strategy for dataset/dataloader creation. + + Implementations provide different data pipelines: + - EnergonDatasetProvider: Megatron Energon for multimodal/diffusion + - SyntheticDatasetProvider: synthetic/mock batches for smoke tests + + The provider pattern allows MegatronTrainer to remain agnostic about + the underlying data source while maintaining compatibility with + Megatron's build_train_valid_test_data_iterators(). + """ + + @abstractmethod + def create_dataloaders( + self, trainer_config: Any, train_val_test_num_samples: List[int], vp_stage: Optional[int] = None + ) -> Tuple[Any, Any, Any]: + """ + Create train/valid/test dataloaders. + + Args: + trainer_config: Megatron args namespace (from megatron.training.get_args()) + train_val_test_num_samples: [train_samples, valid_samples, test_samples] + vp_stage: Virtual pipeline stage (for VP parallelism) + + Returns: + Tuple of (train_dataloader, valid_dataloaders, test_dataloader) + - train_dataloader: Training data iterator + - valid_dataloaders: List of validation data iterators (or None) + - test_dataloader: Test data iterator (or None) + """ + + @property + @abstractmethod + def is_distributed(self) -> bool: + """ + Whether dataloaders are distributed across ranks. + + This flag tells Megatron whether to bypass indexed dataset logic: + False: GPTDataset (uses BlendedMegatronDatasetBuilder, rank-specific slicing) + True: Energon (handles distribution internally via WorkerConfig) + + Returns: + bool: True if dataloaders handle distribution internally + """ + + +__all__ = ["DatasetProvider"] diff --git a/primus/backends/megatron/data/diffusion/__init__.py b/primus/backends/megatron/data/diffusion/__init__.py new file mode 100644 index 000000000..57c8e928f --- /dev/null +++ b/primus/backends/megatron/data/diffusion/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Data pipeline components for diffusion models. + +This module contains diffusion-specific data handling: + - Encoders (VAE, T5, CLIP, etc.) + - TaskEncoders for Energon integration + - Preprocessing utilities +""" diff --git a/primus/backends/megatron/data/diffusion/encoders/__init__.py b/primus/backends/megatron/data/diffusion/encoders/__init__.py new file mode 100644 index 000000000..c379365e1 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/encoders/__init__.py @@ -0,0 +1,162 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Encoder implementations for diffusion models. + +This module contains encoders for various modalities: + - Image encoders (VAE, etc.) + - Text encoders (T5, CLIP, etc.) + - Encoder registry for config-driven selection + - Base encoder classes + +Hierarchical structure: + - encoders/image/autoencoder_kl.py - VAE (AutoencoderKL) + - encoders/text/t5_xxl.py - T5-XXL text encoder + - encoders/text/clip_l.py - CLIP-L text encoder + +Each encoder type can have multiple implementations registered +in the ENCODER_REGISTRY for flexible configuration. +""" + +import importlib +import logging +from typing import Dict, Optional, Type + +from .base import BaseEncoder, BaseTextEncoder, BaseVAE +from .config import ( + CLIPLConfig, + EncoderConfig, + FluxEncoderConfig, + T5XXLConfig, + TextEncoderConfig, + VAEConfig, +) + +logger = logging.getLogger(__name__) + +# Global encoder registry +ENCODER_REGISTRY: Dict[str, Type[BaseEncoder]] = {} + + +def register_encoder(name: str, encoder_class: Type[BaseEncoder]): + """ + Register an encoder class in the global registry. + + Args: + name: Unique identifier for the encoder (e.g., 'autoencoder_kl', 't5_xxl') + encoder_class: Encoder class to register + + Example: + >>> register_encoder('my_encoder', MyEncoder) + """ + if name in ENCODER_REGISTRY: + logger.warning(f"Encoder '{name}' already registered. Overwriting with {encoder_class}") + ENCODER_REGISTRY[name] = encoder_class + logger.debug(f"Registered encoder: {name} -> {encoder_class.__name__}") + + +def get_encoder(config: EncoderConfig) -> BaseEncoder: + """ + Factory function to get an encoder instance from config. + + Args: + config: Encoder configuration object with 'type' field + + Returns: + Instantiated encoder instance + + Raises: + ValueError: If encoder type not found in registry + + Example: + >>> config = VAEConfig(type='autoencoder_kl', model_path='...') + >>> encoder = get_encoder(config) + """ + _ensure_encoders_discovered() + + encoder_type = config.type + + if encoder_type not in ENCODER_REGISTRY: + raise ValueError( + f"Encoder type '{encoder_type}' not found in registry. " + f"Available encoders: {list(ENCODER_REGISTRY.keys())}" + ) + + encoder_class = ENCODER_REGISTRY[encoder_type] + logger.info(f"Loading encoder: {encoder_type} from {config.model_path}") + + # Use from_pretrained if available, otherwise use constructor + if hasattr(encoder_class, "from_pretrained") and callable(encoder_class.from_pretrained): + return encoder_class.from_pretrained(config.model_path, config=config) + else: + return encoder_class(config) + + +def list_encoders() -> Dict[str, Type[BaseEncoder]]: + """ + List all registered encoders. + + Returns: + Dictionary mapping encoder names to encoder classes + """ + _ensure_encoders_discovered() + return ENCODER_REGISTRY.copy() + + +_ENCODERS_DISCOVERED = False + + +def _auto_discover_encoders(): + """ + Auto-discover and register all encoder implementations. + + This function imports encoder modules which triggers their registration + through decorators or explicit register_encoder() calls. + """ + encoder_modules = [ + "primus.backends.megatron.data.diffusion.encoders.image.autoencoder_kl", + "primus.backends.megatron.data.diffusion.encoders.text.t5_xxl", + "primus.backends.megatron.data.diffusion.encoders.text.clip_l", + ] + + for module_path in encoder_modules: + try: + importlib.import_module(module_path) + logger.debug(f"Successfully imported encoder module: {module_path}") + except ImportError as e: + logger.debug(f"Could not import encoder module {module_path}: {e}") + except Exception as e: + logger.warning(f"Error importing encoder module {module_path}: {e}") + + +def _ensure_encoders_discovered(): + """Run encoder discovery once, lazily (on first get_encoder/list_encoders). + + Avoids importing all encoder backends (and their heavy deps) at package + import time. + """ + global _ENCODERS_DISCOVERED + if not _ENCODERS_DISCOVERED: + _auto_discover_encoders() + _ENCODERS_DISCOVERED = True + + +__all__ = [ + # Base classes + "BaseEncoder", + "BaseVAE", + "BaseTextEncoder", + # Config classes + "EncoderConfig", + "VAEConfig", + "TextEncoderConfig", + "T5XXLConfig", + "CLIPLConfig", + "FluxEncoderConfig", + # Registry functions + "register_encoder", + "get_encoder", + "list_encoders", + "ENCODER_REGISTRY", +] diff --git a/primus/backends/megatron/data/diffusion/encoders/base.py b/primus/backends/megatron/data/diffusion/encoders/base.py new file mode 100644 index 000000000..7f676f323 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/encoders/base.py @@ -0,0 +1,285 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Base encoder classes for diffusion models. + +This module defines abstract base classes for different encoder types. +All concrete encoder implementations should inherit from these base classes. +""" + +import logging +from abc import ABC, abstractmethod +from typing import Any, List, Optional, Tuple, Union + +import torch +import torch.nn as nn + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# Utility Functions +# ============================================================================ + + +def get_torch_dtype(precision: str) -> torch.dtype: + """ + Convert precision string to torch dtype. + + Args: + precision: One of 'bf16', 'fp16', 'fp32' + + Returns: + Corresponding torch dtype + + Example: + >>> dtype = get_torch_dtype('bf16') + >>> dtype + torch.bfloat16 + """ + DTYPE_MAP = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, + } + return DTYPE_MAP.get(precision, torch.bfloat16) + + +def load_pretrained_with_subfolder_fallback( + loader_class: Any, model_path: str, subfolder: Optional[str] = None, **kwargs +) -> Any: + """ + Load a pretrained model with automatic subfolder fallback. + + Tries to load from subfolder first, falls back to root path if that fails. + This handles both HuggingFace repos (with subfolders) and local paths. + + Args: + loader_class: Model class with from_pretrained method + model_path: Path to model (HuggingFace repo or local path) + subfolder: Optional subfolder within model_path + **kwargs: Additional arguments passed to from_pretrained + (e.g., torch_dtype, cache_dir, token) + + Returns: + Loaded model instance + + Example: + >>> from transformers import T5EncoderModel + >>> model = load_pretrained_with_subfolder_fallback( + ... T5EncoderModel, + ... "black-forest-labs/FLUX.1-dev", + ... subfolder="text_encoder_2", + ... torch_dtype=torch.bfloat16, + ... cache_dir="/custom/cache/path" + ... ) + """ + # Add token from environment if available and not already in kwargs + import os + + if "token" not in kwargs and "HF_TOKEN" in os.environ: + kwargs["token"] = os.environ["HF_TOKEN"] + + if subfolder: + try: + return loader_class.from_pretrained(model_path, subfolder=subfolder, **kwargs) + except Exception as e: + logger.debug(f"Failed loading from subfolder '{subfolder}': {e}") + logger.debug("Retrying without subfolder...") + + return loader_class.from_pretrained(model_path, **kwargs) + + +# ============================================================================ +# Base Encoder Classes +# ============================================================================ + + +class BaseEncoder(ABC, nn.Module): + """Abstract base class for all encoders.""" + + def __init__(self, config): + """ + Initialize encoder. + + Args: + config: Encoder configuration object + """ + super().__init__() + self.config = config + self._device = config.device if hasattr(config, "device") else "cuda" + self._dtype = get_torch_dtype(config.precision if hasattr(config, "precision") else "bf16") + + @property + def device(self) -> torch.device: + """Get encoder device.""" + if isinstance(self._device, str): + return torch.device(self._device) + return self._device + + @property + def dtype(self) -> torch.dtype: + """Get encoder dtype.""" + return self._dtype + + def _get_dtype(self, precision: str) -> torch.dtype: + """ + Convert precision string to torch dtype. + + Args: + precision: One of 'bf16', 'fp16', 'fp32' + + Returns: + Corresponding torch dtype + """ + return get_torch_dtype(precision) + + @abstractmethod + def encode(self, *args, **kwargs): + """ + Encode input data. + + Must be implemented by subclasses. + """ + raise NotImplementedError("Subclasses must implement encode()") + + @classmethod + @abstractmethod + def from_pretrained(cls, model_path: str, config=None): + """ + Load encoder from pretrained weights. + + Args: + model_path: Path to pretrained model + config: Optional encoder configuration + + Returns: + Loaded encoder instance + """ + raise NotImplementedError("Subclasses must implement from_pretrained()") + + def freeze(self): + """Freeze all encoder parameters.""" + for param in self.parameters(): + param.requires_grad = False + + def unfreeze(self): + """Unfreeze all encoder parameters.""" + for param in self.parameters(): + param.requires_grad = True + + +class BaseVAE(BaseEncoder): + """Abstract base class for VAE encoders.""" + + def __init__(self, config): + """ + Initialize VAE encoder. + + Args: + config: VAEConfig object + """ + super().__init__(config) + self.scale_factor = config.scale_factor if hasattr(config, "scale_factor") else 1.0 + self.shift_factor = config.shift_factor if hasattr(config, "shift_factor") else 0.0 + self.in_channels = config.in_channels if hasattr(config, "in_channels") else 3 + self.out_channels = config.out_channels if hasattr(config, "out_channels") else 16 + self.latent_downsample_factor = ( + config.latent_downsample_factor if hasattr(config, "latent_downsample_factor") else 8 + ) + + @abstractmethod + def encode(self, images: torch.Tensor) -> torch.Tensor: + """ + Encode images to latent representations. + + Args: + images: Input images tensor of shape (B, C, H, W) + + Returns: + Latent representations of shape (B, latent_channels, H/downsample, W/downsample) + """ + raise NotImplementedError("Subclasses must implement encode()") + + @abstractmethod + def decode(self, latents: torch.Tensor) -> torch.Tensor: + """ + Decode latent representations to images. + + Args: + latents: Latent representations of shape (B, latent_channels, H/downsample, W/downsample) + + Returns: + Reconstructed images of shape (B, C, H, W) + """ + raise NotImplementedError("Subclasses must implement decode()") + + def get_latent_shape(self, image_height: int, image_width: int) -> Tuple[int, int, int]: + """ + Calculate latent shape from image dimensions. + + Args: + image_height: Input image height + image_width: Input image width + + Returns: + Tuple of (channels, latent_height, latent_width) + """ + latent_h = image_height // self.latent_downsample_factor + latent_w = image_width // self.latent_downsample_factor + return (self.out_channels, latent_h, latent_w) + + +class BaseTextEncoder(BaseEncoder): + """Abstract base class for text encoders.""" + + def __init__(self, config): + """ + Initialize text encoder. + + Args: + config: TextEncoderConfig object + """ + super().__init__(config) + self.max_length = config.max_length if hasattr(config, "max_length") else 512 + self.embedding_dim = config.embedding_dim if hasattr(config, "embedding_dim") else 4096 + self.return_pooled = config.return_pooled if hasattr(config, "return_pooled") else False + self._tokenizer = None + + @abstractmethod + def encode(self, texts: Union[str, List[str]], **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + """ + Encode text(s) to embeddings. + + Args: + texts: Single text string or list of text strings + **kwargs: Additional encoding arguments + + Returns: + Text embeddings tensor(s) + - For standard encoders: (B, seq_len, embedding_dim) + - For pooled encoders: tuple of (sequence_embeddings, pooled_embeddings) + """ + raise NotImplementedError("Subclasses must implement encode()") + + @property + def tokenizer(self): + """Get tokenizer instance.""" + if self._tokenizer is None: + raise ValueError("Tokenizer not initialized. Call from_pretrained() first.") + return self._tokenizer + + def _prepare_texts(self, texts: Union[str, List[str]]) -> List[str]: + """ + Prepare text inputs for encoding. + + Args: + texts: Single text string or list of text strings + + Returns: + List of text strings + """ + if isinstance(texts, str): + return [texts] + return texts diff --git a/primus/backends/megatron/data/diffusion/encoders/config.py b/primus/backends/megatron/data/diffusion/encoders/config.py new file mode 100644 index 000000000..a3b3a8cdf --- /dev/null +++ b/primus/backends/megatron/data/diffusion/encoders/config.py @@ -0,0 +1,160 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Configuration dataclasses for diffusion encoders. + +This module defines configuration classes for various encoder types used in diffusion models. +""" + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Dict, Optional + + +class EncoderType(Enum): + """Categories of encoders for diffusion models.""" + + IMAGE = "image" + TEXT = "text" + + +@dataclass +class EncoderConfig: + """Base configuration for encoders.""" + + model_path: str # Path to pretrained model weights (REQUIRED) + type: str = None # Encoder type identifier (e.g., 'autoencoder_kl', 't5_xxl', 'clip_l') + encoder_type: Optional[EncoderType] = None # Category: IMAGE or TEXT + precision: str = "bf16" # Model precision: 'bf16', 'fp32', 'fp16' + device: str = "cuda" # Device to load model on + use_cached: bool = True # Whether to use cached/pre-encoded data + freeze_weights: bool = True # Whether to freeze encoder weights during training + cache_dir: Optional[str] = None # Directory to cache downloaded models (if None, uses HF default) + subfolder: Optional[str] = None # Subfolder containing model weights (e.g., 'vae', 'text_encoder_2') + # Opt-in only: allow executing custom modeling code from the HF repo. Defaults to False + # so a malicious/unexpected repo cannot run arbitrary code unless explicitly enabled. + trust_remote_code: bool = False + + # Optional parameters for specific encoders + extra_config: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + """Validate configuration.""" + if self.type is None: + raise ValueError("'type' must be specified in encoder config") + valid_precisions = ["bf16", "fp32", "fp16"] + if self.precision not in valid_precisions: + raise ValueError(f"precision must be one of {valid_precisions}, got {self.precision}") + + +@dataclass +class VAEConfig(EncoderConfig): + """Configuration for VAE encoders.""" + + type: str = "autoencoder_kl" # VAE type identifier + scale_factor: float = 0.3611 # Flux default scale factor + shift_factor: float = 0.1159 # Flux default shift factor + in_channels: int = 3 # Input image channels + out_channels: int = 16 # Latent channels + latent_downsample_factor: int = 8 # Spatial downsampling factor (H/8, W/8) + + def __post_init__(self): + """Set encoder type and validate.""" + if self.encoder_type is None: + self.encoder_type = EncoderType.IMAGE + super().__post_init__() + + +@dataclass +class TextEncoderConfig(EncoderConfig): + """Configuration for text encoders.""" + + max_length: int = 512 # Maximum sequence length + embedding_dim: int = 4096 # Output embedding dimension + return_pooled: bool = False # Whether to return pooled embeddings + tokenizer_path: Optional[str] = None # Optional separate tokenizer path + tokenizer_subfolder: Optional[str] = ( + None # Subfolder containing tokenizer files (e.g., 'tokenizer', 'tokenizer_2') + ) + + def __post_init__(self): + """Validate configuration.""" + if not self.tokenizer_path: + self.tokenizer_path = self.model_path # Default to same as model path + if self.encoder_type is None: + self.encoder_type = EncoderType.TEXT + super().__post_init__() + + +@dataclass +class T5XXLConfig(TextEncoderConfig): + """Configuration for T5-XXL encoder.""" + + type: str = "t5_xxl" # T5-XXL type identifier + max_length: int = 512 # Flux uses 512 for T5-XXL + embedding_dim: int = 4096 # T5-XXL hidden size + + def __post_init__(self): + """Validate configuration.""" + super().__post_init__() + + +@dataclass +class CLIPLConfig(TextEncoderConfig): + """Configuration for CLIP-L encoder.""" + + type: str = "clip_l" # CLIP-L type identifier + max_length: int = 77 # CLIP uses 77 max length + embedding_dim: int = 768 # CLIP-L hidden size + return_pooled: bool = True # CLIP returns both sequence and pooled embeddings + pooled_dim: int = 768 # CLIP-L pooled embedding dimension + + def __post_init__(self): + """Validate configuration.""" + super().__post_init__() + + +@dataclass +class FluxEncoderConfig: + """Configuration for all Flux encoders (VAE + T5-XXL + CLIP-L).""" + + vae: VAEConfig + t5: T5XXLConfig + clip: CLIPLConfig + use_preencoded: bool = True # Whether to use pre-encoded data + + @classmethod + def from_pretrained_flux( + cls, + model_path: str = "black-forest-labs/FLUX.1-dev", + precision: str = "bf16", + device: str = "cuda", + use_preencoded: bool = True, + cache_dir: Optional[str] = None, + ): + """Create FluxEncoderConfig from pretrained Flux model path.""" + return cls( + vae=VAEConfig( + type="autoencoder_kl", + model_path=f"{model_path}", + precision=precision, + device=device, + cache_dir=cache_dir, + ), + t5=T5XXLConfig( + type="t5_xxl", + model_path=f"{model_path}", + precision=precision, + device=device, + cache_dir=cache_dir, + ), + clip=CLIPLConfig( + type="clip_l", + model_path=f"{model_path}", + precision=precision, + device=device, + cache_dir=cache_dir, + ), + use_preencoded=use_preencoded, + ) diff --git a/primus/backends/megatron/data/diffusion/encoders/image/__init__.py b/primus/backends/megatron/data/diffusion/encoders/image/__init__.py new file mode 100644 index 000000000..c11e6c777 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/encoders/image/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Image encoder implementations for diffusion models. +""" + +from primus.backends.megatron.data.diffusion.encoders.base import BaseVAE + +from .autoencoder_kl import AutoencoderKL + +__all__ = [ + "BaseVAE", + "AutoencoderKL", +] diff --git a/primus/backends/megatron/data/diffusion/encoders/image/autoencoder_kl.py b/primus/backends/megatron/data/diffusion/encoders/image/autoencoder_kl.py new file mode 100644 index 000000000..fc8d3d403 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/encoders/image/autoencoder_kl.py @@ -0,0 +1,270 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. +# +# Adapted from NeMo's VAE implementation + +""" +AutoencoderKL implementation for Flux VAE. + +This module provides a wrapper around the diffusers AutoencoderKL model +for use with Flux diffusion models. It handles encoding images to latent +representations and decoding latents back to images. +""" + +import logging +from typing import Optional, Tuple + +import torch + +try: + from diffusers import AutoencoderKL as DiffusersAutoencoderKL +except ImportError: + DiffusersAutoencoderKL = None + +from primus.backends.megatron.data.diffusion.encoders.base import ( + BaseVAE, + get_torch_dtype, + load_pretrained_with_subfolder_fallback, +) +from primus.backends.megatron.data.diffusion.encoders.config import VAEConfig + +logger = logging.getLogger(__name__) + + +class AutoencoderKL(BaseVAE): + """ + AutoencoderKL VAE for Flux diffusion models. + + This encoder wraps the diffusers AutoencoderKL model and applies + Flux-specific scale and shift factors to the latent representations. + + For Flux, the default configuration is: + - scale_factor: 0.3611 + - shift_factor: 0.1159 + - in_channels: 3 (RGB images) + - out_channels: 16 (latent channels) + - latent_downsample_factor: 8 (spatial downsampling) + + This means a 1024x1024 image becomes a 16x128x128 latent. + """ + + def __init__(self, config: VAEConfig): + """ + Initialize AutoencoderKL encoder. + + Args: + config: VAEConfig with model_path, scale_factor, shift_factor, etc. + """ + super().__init__(config) + + if DiffusersAutoencoderKL is None: + raise ImportError( + "diffusers library is required for AutoencoderKL. " "Install with: pip install diffusers" + ) + + self.vae = None # Will be loaded in from_pretrained + + @classmethod + def from_pretrained( + cls, + model_path: str, + config: Optional[VAEConfig] = None, + subfolder: Optional[str] = None, + ) -> "AutoencoderKL": + """ + Load AutoencoderKL from pretrained weights. + + Args: + model_path: Path to pretrained model (local path or HuggingFace repo) + config: Optional VAEConfig. If None, uses defaults. + subfolder: Subfolder for VAE weights. Priority: param > config.subfolder + + Returns: + Loaded AutoencoderKL instance + + Examples: + >>> # Using config (recommended) + >>> config = VAEConfig( + ... model_path="black-forest-labs/FLUX.1-dev", + ... subfolder="vae" + ... ) + >>> vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", config=config) + + >>> # Using method parameter (backward compatible) + >>> vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae") + """ + if config is None: + config = VAEConfig( + type="autoencoder_kl", + model_path=model_path, + precision="bf16", + ) + + instance = cls(config) + + # Prepare kwargs for from_pretrained calls + pretrained_kwargs = {} + if config.cache_dir: + pretrained_kwargs["cache_dir"] = config.cache_dir + logger.info(f"Using cache directory: {config.cache_dir}") + + # Resolve model subfolder with priority: param > config.subfolder > error + model_subfolder = subfolder if subfolder is not None else getattr(config, "subfolder", None) + + # Require explicit configuration for model subfolder + if model_subfolder is None and not hasattr(config, "subfolder"): + raise ValueError( + f"subfolder must be specified for AutoencoderKL with model_path='{model_path}'. " + f"For FLUX models (e.g., black-forest-labs/FLUX.1-dev), use subfolder='vae'. " + f"Set it via config.subfolder or the subfolder parameter." + ) + + # Load VAE from diffusers + torch_dtype = get_torch_dtype(config.precision) + + logger.info(f"Loading AutoencoderKL from {model_path} (subfolder={model_subfolder})") + instance.vae = load_pretrained_with_subfolder_fallback( + DiffusersAutoencoderKL, + model_path, + subfolder=model_subfolder, + torch_dtype=torch_dtype, + **pretrained_kwargs, + ) + + instance.vae.to(instance.device) + + if config.freeze_weights: + instance.freeze() + instance.vae.eval() + + logger.info( + f"Loaded AutoencoderKL: in_channels={instance.in_channels}, " + f"out_channels={instance.out_channels}, scale={instance.scale_factor:.4f}, " + f"shift={instance.shift_factor:.4f}" + ) + + return instance + + @torch.no_grad() + def encode(self, images: torch.Tensor) -> torch.Tensor: + """ + Encode images to latent representations. + + Applies the Flux scale and shift factors: + latents = scale_factor * (vae_encode(images) - shift_factor) + + Args: + images: Input images tensor of shape (B, C, H, W) + Values should be in range [-1, 1] (normalized) + + Returns: + Latent representations of shape (B, 16, H/8, W/8) + + Example: + >>> images = torch.randn(2, 3, 1024, 1024) # B=2, 1024x1024 RGB + >>> latents = vae.encode(images) + >>> latents.shape + torch.Size([2, 16, 128, 128]) + """ + if self.vae is None: + raise RuntimeError("VAE not loaded. Call from_pretrained() first.") + + # Ensure images are on correct device and dtype + images = images.to(device=self.device, dtype=self.dtype) + + # Encode using diffusers VAE + latent_dist = self.vae.encode(images).latent_dist + latents = latent_dist.sample() + + # Apply Flux scale and shift factors + # Formula: z = scale_factor * (encoded - shift_factor) + latents = self.scale_factor * (latents - self.shift_factor) + + return latents + + @torch.no_grad() + def encode_for_resample(self, images: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Encode images returning posterior parameters for training-time resampling. + + Unlike ``encode()`` which draws a single stochastic sample, this method + returns the raw posterior parameters (mean, logvar) so the training loop + can re-draw latents via reparameterization at every step. + + The returned ``latents`` use the deterministic posterior mode (mean) + with scale/shift applied, for backward compatibility and debugging. + + Args: + images: Input images tensor (B, C, H, W) in range [-1, 1] + + Returns: + Tuple of (latents, mean, logvar): + - latents: scale * (mode - shift), shape (B, 16, H/8, W/8) + - mean: raw posterior mean, shape (B, 16, H/8, W/8) + - logvar: raw posterior log-variance, shape (B, 16, H/8, W/8) + """ + if self.vae is None: + raise RuntimeError("VAE not loaded. Call from_pretrained() first.") + + images = images.to(device=self.device, dtype=self.dtype) + + latent_dist = self.vae.encode(images).latent_dist + latents = self.scale_factor * (latent_dist.mode() - self.shift_factor) + + return latents, latent_dist.mean, latent_dist.logvar + + @torch.no_grad() + def decode(self, latents: torch.Tensor) -> torch.Tensor: + """ + Decode latent representations to images. + + Reverses the Flux scale and shift factors before decoding: + vae_decode(latents / scale_factor + shift_factor) + + Args: + latents: Latent representations of shape (B, 16, H/8, W/8) + + Returns: + Reconstructed images of shape (B, C, H, W) + Values in range [-1, 1] + + Example: + >>> latents = torch.randn(2, 16, 128, 128) + >>> images = vae.decode(latents) + >>> images.shape + torch.Size([2, 3, 1024, 1024]) + """ + if self.vae is None: + raise RuntimeError("VAE not loaded. Call from_pretrained() first.") + + # Ensure latents are on correct device and dtype + latents = latents.to(device=self.device, dtype=self.dtype) + + # Reverse Flux scale and shift factors + # Formula: decoded_input = z / scale_factor + shift_factor + latents = latents / self.scale_factor + self.shift_factor + + # Decode using diffusers VAE + images = self.vae.decode(latents, return_dict=False)[0] + + return images + + def forward(self, images: torch.Tensor) -> torch.Tensor: + """ + Forward pass: encode then decode (reconstruction). + + Args: + images: Input images tensor of shape (B, C, H, W) + + Returns: + Reconstructed images of shape (B, C, H, W) + """ + latents = self.encode(images) + reconstructed = self.decode(latents) + return reconstructed + + +# Register encoder in registry +from primus.backends.megatron.data.diffusion.encoders import register_encoder + +register_encoder("autoencoder_kl", AutoencoderKL) diff --git a/primus/backends/megatron/data/diffusion/encoders/text/__init__.py b/primus/backends/megatron/data/diffusion/encoders/text/__init__.py new file mode 100644 index 000000000..ce83875a3 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/encoders/text/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Text encoder implementations for diffusion models. +""" + +from .clip_l import CLIPLEncoder +from .t5_xxl import T5XXLEncoder + +__all__ = [ + "T5XXLEncoder", + "CLIPLEncoder", +] diff --git a/primus/backends/megatron/data/diffusion/encoders/text/clip_l.py b/primus/backends/megatron/data/diffusion/encoders/text/clip_l.py new file mode 100644 index 000000000..3079f272e --- /dev/null +++ b/primus/backends/megatron/data/diffusion/encoders/text/clip_l.py @@ -0,0 +1,254 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. +# +# Adapted from NeMo's text encoder implementation + +""" +CLIP-L text encoder implementation for Flux. + +This module provides a wrapper around the transformers CLIPTextModel +for encoding text prompts into embeddings for Flux diffusion models. +""" + +import logging +from typing import List, Optional, Tuple, Union + +import torch + +try: + from transformers import CLIPTextModel, CLIPTokenizer +except ImportError: + CLIPTextModel = None + CLIPTokenizer = None + +from primus.backends.megatron.data.diffusion.encoders.base import ( + BaseTextEncoder, + get_torch_dtype, + load_pretrained_with_subfolder_fallback, +) +from primus.backends.megatron.data.diffusion.encoders.config import CLIPLConfig + +logger = logging.getLogger(__name__) + + +class CLIPLEncoder(BaseTextEncoder): + """ + CLIP-L text encoder for Flux diffusion models. + + This encoder uses the CLIP ViT-L/14 text encoder from HuggingFace to encode + text prompts into embeddings. For Flux, the default configuration is: + - max_length: 77 tokens (CLIP default) + - embedding_dim: 768 (CLIP-L hidden size) + - pooled_dim: 768 (CLIP-L pooled embedding size) + - precision: bf16 + + The encoder returns both sequence embeddings (B, seq_len, 768) and + pooled embeddings (B, 768) which are used differently in Flux. + + Note: Sequence length is padded to a multiple of 8 for Flux compatibility. + """ + + def __init__(self, config: CLIPLConfig): + """ + Initialize CLIP-L encoder. + + Args: + config: CLIPLConfig with model_path, max_length, etc. + """ + super().__init__(config) + self.pooled_dim = config.pooled_dim if hasattr(config, "pooled_dim") else 768 + + if CLIPTextModel is None or CLIPTokenizer is None: + raise ImportError( + "transformers library is required for CLIPLEncoder. " "Install with: pip install transformers" + ) + + self.transformer = None # Will be loaded in from_pretrained + self._tokenizer = None + + @classmethod + def from_pretrained( + cls, + model_path: str, + config: Optional[CLIPLConfig] = None, + subfolder: Optional[str] = None, + ) -> "CLIPLEncoder": + """ + Load CLIP-L encoder from pretrained weights. + + Args: + model_path: Path to pretrained model (local path or HuggingFace repo) + config: Optional CLIPLConfig. If None, uses defaults. + subfolder: Subfolder for model weights. Priority: param > config.subfolder + + Returns: + Loaded CLIPLEncoder instance + + Examples: + >>> # Using config (recommended) + >>> config = CLIPLConfig( + ... model_path="black-forest-labs/FLUX.1-dev", + ... subfolder="text_encoder", + ... tokenizer_subfolder="tokenizer" + ... ) + >>> encoder = CLIPLEncoder.from_pretrained("black-forest-labs/FLUX.1-dev", config=config) + + >>> # Using method parameter (backward compatible) + >>> encoder = CLIPLEncoder.from_pretrained("openai/clip-vit-large-patch14", subfolder=None) + """ + if config is None: + config = CLIPLConfig( + type="clip_l", + model_path=model_path, + precision="bf16", + ) + + instance = cls(config) + + # Prepare kwargs for from_pretrained calls + pretrained_kwargs = {} + if config.cache_dir: + pretrained_kwargs["cache_dir"] = config.cache_dir + logger.info(f"Using cache directory: {config.cache_dir}") + + # Resolve model subfolder with priority: param > config.subfolder > error + model_subfolder = subfolder if subfolder is not None else getattr(config, "subfolder", None) + + # Require explicit configuration for model subfolder + # NOTE: config.subfolder can legitimately be None for standalone models like openai/clip-vit-large-patch14 + # but it must be explicitly set in the config. If neither param nor config specify it, that's an error. + if model_subfolder is None and not hasattr(config, "subfolder"): + raise ValueError( + f"subfolder must be specified for CLIPLEncoder with model_path='{model_path}'. " + f"For FLUX models (e.g., black-forest-labs/FLUX.1-dev), use subfolder='text_encoder'. " + f"For standalone CLIP models (e.g., openai/clip-vit-large-patch14), use subfolder=None. " + f"Set it via config.subfolder or the subfolder parameter." + ) + + # Resolve tokenizer subfolder with priority: config.tokenizer_subfolder > model_subfolder + tokenizer_subfolder = getattr(config, "tokenizer_subfolder", None) + if tokenizer_subfolder is None: + tokenizer_subfolder = model_subfolder # Use model subfolder as fallback + + # Load tokenizer + tokenizer_path = config.tokenizer_path or model_path + logger.info(f"Loading CLIP tokenizer from {tokenizer_path} (subfolder={tokenizer_subfolder})") + try: + instance._tokenizer = load_pretrained_with_subfolder_fallback( + CLIPTokenizer, + tokenizer_path, + subfolder=tokenizer_subfolder, + **pretrained_kwargs, + ) + except Exception as e: + # Try standard CLIP tokenizer if neither subfolder nor root work + logger.info(f"Failed loading from {tokenizer_path}, trying standard CLIP tokenizer: {e}") + instance._tokenizer = CLIPTokenizer.from_pretrained( + "openai/clip-vit-large-patch14", + **pretrained_kwargs, + ) + + # Load model + torch_dtype = get_torch_dtype(config.precision) + + logger.info(f"Loading CLIP-L model from {model_path} (subfolder={model_subfolder})") + instance.transformer = load_pretrained_with_subfolder_fallback( + CLIPTextModel, + model_path, + subfolder=model_subfolder, + torch_dtype=torch_dtype, + **pretrained_kwargs, + ) + + instance.transformer.to(instance.device) + + if config.freeze_weights: + instance.freeze() + instance.transformer.eval() + + logger.info( + f"Loaded CLIP-L: max_length={instance.max_length}, " + f"embedding_dim={instance.embedding_dim}, pooled_dim={instance.pooled_dim}, " + f"dtype={torch_dtype}" + ) + + return instance + + @torch.no_grad() + def encode( + self, + texts: Union[str, List[str]], + max_sequence_length: Optional[int] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Encode text(s) to embeddings. + + Args: + texts: Single text string or list of text strings + max_sequence_length: Optional max length override. If None, uses config max_length. + + Returns: + Tuple of (sequence_embeddings, pooled_embeddings): + - sequence_embeddings: (batch_size, padded_seq_len, 768) + Sequence length is padded to a multiple of 8 for Flux + - pooled_embeddings: (batch_size, 768) + Pooled representation from CLIP's [CLS] token + + Example: + >>> texts = ["A beautiful sunset over the ocean", "A cat sitting on a mat"] + >>> seq_embeds, pooled_embeds = encoder.encode(texts) + >>> seq_embeds.shape + torch.Size([2, 80, 768]) # Padded to multiple of 8 + >>> pooled_embeds.shape + torch.Size([2, 768]) + """ + if self.transformer is None or self._tokenizer is None: + raise RuntimeError("CLIP encoder not loaded. Call from_pretrained() first.") + + texts = self._prepare_texts(texts) + max_len = max_sequence_length if max_sequence_length is not None else self.max_length + + # Tokenize + batch_encoding = self._tokenizer( + texts, + truncation=True, + max_length=max_len, + return_length=True, + return_overflowing_tokens=False, + padding="max_length", + return_tensors="pt", + ) + + tokens = batch_encoding["input_ids"].to(self.device, non_blocking=True) + + # Encode + outputs = self.transformer(input_ids=tokens, output_hidden_states=False) + + # Get sequence embeddings (last hidden state) + sequence_embeddings = outputs.last_hidden_state + + # Pad sequence length to multiple of 8 (required for Flux) + seq_len = sequence_embeddings.shape[1] + padded_seq_len = ((seq_len + 7) // 8) * 8 # Round up to nearest multiple of 8 + + if padded_seq_len > seq_len: + sequence_embeddings = torch.nn.functional.pad( + sequence_embeddings, + (0, 0, 0, padded_seq_len - seq_len), # Pad on sequence dimension + value=0.0, + ) + + # Get pooled embeddings (from pooler output) + pooled_embeddings = outputs.pooler_output + + return sequence_embeddings, pooled_embeddings + + def forward(self, texts: Union[str, List[str]], **kwargs) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass (alias for encode).""" + return self.encode(texts, **kwargs) + + +# Register encoder in registry +from primus.backends.megatron.data.diffusion.encoders import register_encoder + +register_encoder("clip_l", CLIPLEncoder) diff --git a/primus/backends/megatron/data/diffusion/encoders/text/t5_xxl.py b/primus/backends/megatron/data/diffusion/encoders/text/t5_xxl.py new file mode 100644 index 000000000..437fe94b9 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/encoders/text/t5_xxl.py @@ -0,0 +1,229 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. +# +# Adapted from NeMo's text encoder implementation + +""" +T5-XXL text encoder implementation for Flux. + +This module provides a wrapper around the transformers T5EncoderModel +for encoding text prompts into embeddings for Flux diffusion models. +""" + +import logging +from typing import List, Optional, Union + +import torch + +try: + from transformers import T5EncoderModel, T5Tokenizer +except ImportError: + T5EncoderModel = None + T5Tokenizer = None + +from primus.backends.megatron.data.diffusion.encoders.base import ( + BaseTextEncoder, + get_torch_dtype, + load_pretrained_with_subfolder_fallback, +) +from primus.backends.megatron.data.diffusion.encoders.config import T5XXLConfig + +logger = logging.getLogger(__name__) + + +class T5XXLEncoder(BaseTextEncoder): + """ + T5-XXL text encoder for Flux diffusion models. + + This encoder uses the T5-v1.1-XXL model from HuggingFace to encode + text prompts into embeddings. For Flux, the default configuration is: + - max_length: 512 tokens + - embedding_dim: 4096 (T5-XXL hidden size) + - precision: bf16 + + The output embeddings have shape (batch_size, seq_len, 4096). + """ + + def __init__(self, config: T5XXLConfig): + """ + Initialize T5-XXL encoder. + + Args: + config: T5XXLConfig with model_path, max_length, etc. + """ + super().__init__(config) + + if T5EncoderModel is None or T5Tokenizer is None: + raise ImportError( + "transformers library is required for T5XXLEncoder. " "Install with: pip install transformers" + ) + + self.transformer = None # Will be loaded in from_pretrained + self._tokenizer = None + + @classmethod + def from_pretrained( + cls, + model_path: str, + config: Optional[T5XXLConfig] = None, + subfolder: Optional[str] = None, + ) -> "T5XXLEncoder": + """ + Load T5-XXL encoder from pretrained weights. + + Args: + model_path: Path to pretrained model (local path or HuggingFace repo) + config: Optional T5XXLConfig. If None, uses defaults. + subfolder: Subfolder for model weights. Priority: param > config.subfolder + + Returns: + Loaded T5XXLEncoder instance + + Examples: + >>> # Using config (recommended) + >>> config = T5XXLConfig( + ... model_path="black-forest-labs/FLUX.1-dev", + ... subfolder="text_encoder_2", + ... tokenizer_subfolder="tokenizer_2" + ... ) + >>> encoder = T5XXLEncoder.from_pretrained("black-forest-labs/FLUX.1-dev", config=config) + + >>> # Using method parameter (backward compatible) + >>> encoder = T5XXLEncoder.from_pretrained("google/t5-v1_1-xxl", subfolder=None) + """ + if config is None: + config = T5XXLConfig( + type="t5_xxl", + model_path=model_path, + precision="bf16", + ) + + instance = cls(config) + + # Prepare kwargs for from_pretrained calls + pretrained_kwargs = {} + if config.cache_dir: + pretrained_kwargs["cache_dir"] = config.cache_dir + logger.info(f"Using cache directory: {config.cache_dir}") + + # Resolve model subfolder with priority: param > config.subfolder > error + model_subfolder = subfolder if subfolder is not None else getattr(config, "subfolder", None) + + # Require explicit configuration for model subfolder + # NOTE: config.subfolder can legitimately be None for standalone models like google/t5-v1_1-xxl + # but it must be explicitly set in the config. If neither param nor config specify it, that's an error. + if model_subfolder is None and not hasattr(config, "subfolder"): + raise ValueError( + f"subfolder must be specified for T5XXLEncoder with model_path='{model_path}'. " + f"For FLUX models (e.g., black-forest-labs/FLUX.1-dev), use subfolder='text_encoder_2'. " + f"For standalone T5 models (e.g., google/t5-v1_1-xxl), use subfolder=None. " + f"Set it via config.subfolder or the subfolder parameter." + ) + + # Resolve tokenizer subfolder with priority: config.tokenizer_subfolder > model_subfolder + tokenizer_subfolder = getattr(config, "tokenizer_subfolder", None) + if tokenizer_subfolder is None: + tokenizer_subfolder = model_subfolder # Use model subfolder as fallback + + # Load tokenizer + tokenizer_path = config.tokenizer_path or model_path + logger.info(f"Loading T5 tokenizer from {tokenizer_path} (subfolder={tokenizer_subfolder})") + + # Load tokenizer from subfolder + import os + + tokenizer_kwargs = { + "token": os.environ.get("HF_TOKEN"), + "trust_remote_code": getattr(config, "trust_remote_code", False), + **pretrained_kwargs, # Merge cache_dir if provided + } + if tokenizer_subfolder: + tokenizer_kwargs["subfolder"] = tokenizer_subfolder + + instance._tokenizer = T5Tokenizer.from_pretrained( + tokenizer_path, + **tokenizer_kwargs, + ) + + # Load model + torch_dtype = get_torch_dtype(config.precision) + + logger.info(f"Loading T5-XXL model from {model_path} (subfolder={model_subfolder})") + instance.transformer = load_pretrained_with_subfolder_fallback( + T5EncoderModel, + model_path, + subfolder=model_subfolder, + torch_dtype=torch_dtype, + **pretrained_kwargs, + ) + + instance.transformer.to(instance.device) + + if config.freeze_weights: + instance.freeze() + instance.transformer.eval() + + logger.info( + f"Loaded T5-XXL: max_length={instance.max_length}, " + f"embedding_dim={instance.embedding_dim}, dtype={torch_dtype}" + ) + + return instance + + @torch.no_grad() + def encode( + self, + texts: Union[str, List[str]], + max_sequence_length: Optional[int] = None, + ) -> torch.Tensor: + """ + Encode text(s) to embeddings. + + Args: + texts: Single text string or list of text strings + max_sequence_length: Optional max length override. If None, uses config max_length. + + Returns: + Text embeddings tensor of shape (batch_size, seq_len, 4096) + Sequence length is padded to max_length. + + Example: + >>> texts = ["A beautiful sunset over the ocean", "A cat sitting on a mat"] + >>> embeddings = encoder.encode(texts) + >>> embeddings.shape + torch.Size([2, 512, 4096]) + """ + if self.transformer is None or self._tokenizer is None: + raise RuntimeError("T5 encoder not loaded. Call from_pretrained() first.") + + texts = self._prepare_texts(texts) + max_len = max_sequence_length if max_sequence_length is not None else self.max_length + + # Tokenize + batch_encoding = self._tokenizer( + texts, + truncation=True, + max_length=max_len, + return_length=False, + return_overflowing_tokens=False, + padding="max_length", + return_tensors="pt", + ) + + tokens = batch_encoding["input_ids"].to(self.device, non_blocking=True) + + # Encode + outputs = self.transformer(input_ids=tokens, output_hidden_states=None) + embeddings = outputs.last_hidden_state + + return embeddings + + def forward(self, texts: Union[str, List[str]], **kwargs) -> torch.Tensor: + """Forward pass (alias for encode).""" + return self.encode(texts, **kwargs) + + +# Register encoder in registry +from primus.backends.megatron.data.diffusion.encoders import register_encoder + +register_encoder("t5_xxl", T5XXLEncoder) diff --git a/primus/backends/megatron/data/diffusion/task_encoders/__init__.py b/primus/backends/megatron/data/diffusion/task_encoders/__init__.py new file mode 100644 index 000000000..a3ae2036f --- /dev/null +++ b/primus/backends/megatron/data/diffusion/task_encoders/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +TaskEncoders for diffusion models. +""" + +from .image import ( + DiffusionSample, + EncodedDiffusionTaskEncoder, + RawDiffusionTaskEncoder, + cook_preencoded_diffusion, + cook_raw_images, +) + +__all__ = [ + "DiffusionSample", + "EncodedDiffusionTaskEncoder", + "RawDiffusionTaskEncoder", + "cook_preencoded_diffusion", + "cook_raw_images", +] diff --git a/primus/backends/megatron/data/diffusion/task_encoders/image.py b/primus/backends/megatron/data/diffusion/task_encoders/image.py new file mode 100644 index 000000000..837bbcfdb --- /dev/null +++ b/primus/backends/megatron/data/diffusion/task_encoders/image.py @@ -0,0 +1,396 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Diffusion TaskEncoder for image-text pairs using Crude Data pattern. + +This module provides TaskEncoder implementations for diffusion models: +- EncodedDiffusionTaskEncoder: Loads pre-encoded data (latents, embeddings) +- RawDiffusionTaskEncoder: Loads raw images and text (no encoding) + +Encoding happens in the model, not in the TaskEncoder (following best practices). +Position IDs are generated in the model code, not here. +""" + +import io +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Optional + +import numpy as np +import torch +from megatron.energon import ( + Cooker, + DefaultTaskEncoder, + Sample, + SampleDecoder, + WorkerConfig, + basic_sample_keys, + stateless, +) + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# Sample Definition (with proper Sample inheritance) +# ============================================================================ + + +@dataclass +class DiffusionSample(Sample): + """ + Diffusion training sample with framework-standard field names. + + Inherits from megatron.energon.Sample to ensure __key__, __restore_key__, + and __subflavors__ are properly tracked for deterministic training resumption. + + Position IDs (img_ids, text_ids) are generated at runtime based on actual + tensor shapes, not stored in the dataset. This provides flexibility for + variable-resolution training. + + Attributes: + prompt_embeds: T5 text embeddings (seq_len, hidden_dim) + pooled_prompt_embeds: CLIP pooled embeddings (hidden_dim,) + latents: Image latents from VAE (C, H, W) — optional for resample-only datasets + mean: VAE posterior mean (C, H, W) — present in 'resample' mode datasets + logvar: VAE posterior log-variance (C, H, W) — present in 'resample' mode datasets + caption: Original text caption (optional, for debugging) + """ + + prompt_embeds: torch.Tensor + pooled_prompt_embeds: torch.Tensor + latents: Optional[torch.Tensor] = None + mean: Optional[torch.Tensor] = None + logvar: Optional[torch.Tensor] = None + caption: str = "" + timestep: Optional[torch.Tensor] = None + + +# ============================================================================ +# Cooker Functions +# ============================================================================ + + +@stateless +def cook_preencoded_diffusion(sample: dict) -> DiffusionSample: + """ + Cooker for preencoded diffusion features with framework-standard keys. + + Loads precalculated VAE latents and text embeddings from disk. + Position IDs are NOT loaded - they are generated at runtime based on + actual tensor shapes for flexibility. + + Required standard keys: + - 'latents.pth': VAE-encoded image latents + - 'prompt_embeds.pth': T5 text embeddings + - 'pooled_prompt_embeds.pth': CLIP pooled embeddings + + Args: + sample: Raw sample dict from WebDataset + + Returns: + DiffusionSample with all metadata properly forwarded + """ + + def load_tensor(data): + """Helper to load tensor from bytes.""" + if data is None: + return None + if isinstance(data, (str, Path)): + return torch.load(data, map_location="cpu") + elif isinstance(data, bytes): + return torch.load(io.BytesIO(data), map_location="cpu") + elif isinstance(data, torch.Tensor): + return data + else: + return data + + # Load required fields (position IDs not loaded) + latents = load_tensor(sample.get("latents.pth")) + prompt_embeds = load_tensor(sample.get("prompt_embeds.pth")) + pooled_prompt_embeds = load_tensor(sample.get("pooled_prompt_embeds.pth")) + + # Load optional resample-mode fields + mean = load_tensor(sample.get("mean.pth")) + logvar = load_tensor(sample.get("logvar.pth")) + + if prompt_embeds is None or pooled_prompt_embeds is None: + raise ValueError( + f"Sample missing required keys. Expected: 'prompt_embeds.pth', " + f"'pooled_prompt_embeds.pth'. Got: {list(sample.keys())}" + ) + if latents is None and mean is None: + raise ValueError( + f"Sample must have 'latents.pth' or 'mean.pth'/'logvar.pth'. " f"Got: {list(sample.keys())}" + ) + + # Load caption if available + caption = sample.get("caption.txt", b"") + if isinstance(caption, bytes): + caption = caption.decode("utf-8") + elif isinstance(caption, (str, Path)): + try: + if Path(caption).exists(): + with open(caption, "r") as f: + caption = f.read().strip() + else: + caption = str(caption) + except (OSError, ValueError): + caption = str(caption) + else: + caption = str(caption) if caption else "" + + # Extract timestep from JSON sidecar (MLPerf validation datasets) + timestep = None + if "json" in sample: + import json as json_mod + + raw = sample["json"] + if isinstance(raw, dict): + decoded = raw + elif isinstance(raw, bytes): + decoded = json_mod.loads(raw.decode("utf-8")) + else: + decoded = json_mod.loads(raw) + if "timestep" in decoded: + timestep = torch.tensor(decoded["timestep"]) + + return DiffusionSample( + **basic_sample_keys(sample), + latents=latents, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + mean=mean, + logvar=logvar, + caption=caption, + timestep=timestep, + ) + + +@stateless +def cook_raw_images(sample: dict) -> Dict[str, Any]: + """ + Cooker for raw images - just loads data, NO ENCODING. + + Encoding happens in the model's forward_step, not here. + This follows the pattern where TaskEncoders only load data. + + Standard data keys: + - 'images': Raw image data (or format-specific key like 'jpg', 'png', 'webp') + - 'txt': Text caption + + Args: + sample: Raw sample dict from WebDataset + + Returns: + Dict with raw data ready for model encoding + """ + return { + **basic_sample_keys(sample), + "images": sample.get("images"), + "txt": sample.get("txt"), + } + + +def load_numpy_tensor(data: Optional[bytes]) -> Optional[torch.Tensor]: + """Load a bfloat16 tensor from numpy uint16 bytes (MLPerf format). + + The MLPerf Flux dataset stores bfloat16 tensors by reinterpreting them + as uint16 numpy arrays serialized as .npy files. This function loads + the .npy buffer (preserving shape), then reinterprets uint16 as + bfloat16 via torch. + + Falls back to raw ``np.frombuffer`` for headerless byte buffers. + """ + if data is None: + return None + if data[:6] == b"\x93NUMPY": + arr = np.load(io.BytesIO(data)) + else: + arr = np.frombuffer(data, dtype=np.uint16) + return torch.from_numpy(arr.copy()).view(torch.bfloat16) + + +@stateless +def cook_preencoded_numpy_diffusion(sample: dict) -> DiffusionSample: + """Cooker for MLPerf pre-encoded numpy data (bfloat16 as uint16 bytes). + + Expected keys in the WebDataset tar shard: + - 't5.bytes': T5 text embeddings + - 'clip.bytes': CLIP pooled embeddings + - 'mean.bytes': VAE posterior mean + - 'logvar.bytes': VAE posterior log-variance + + No 'latents' are stored; the training loop resamples from mean/logvar. + """ + prompt_embeds = load_numpy_tensor(sample.get("t5.bytes")) + pooled_prompt_embeds = load_numpy_tensor(sample.get("clip.bytes")) + mean = load_numpy_tensor(sample.get("mean.bytes")) + logvar = load_numpy_tensor(sample.get("logvar.bytes")) + + if prompt_embeds is None or pooled_prompt_embeds is None: + raise ValueError( + f"Sample missing required keys. Expected: 't5.bytes', " + f"'clip.bytes'. Got: {list(sample.keys())}" + ) + if mean is None: + raise ValueError(f"Sample missing 'mean.bytes'. Got: {list(sample.keys())}") + + timestep = None + if "json" in sample: + import json as json_mod + + raw = sample["json"] + if isinstance(raw, dict): + decoded = raw + elif isinstance(raw, bytes): + decoded = json_mod.loads(raw.decode("utf-8")) + else: + decoded = json_mod.loads(raw) + if "timestep" in decoded: + timestep = torch.tensor(decoded["timestep"]) + + return DiffusionSample( + **basic_sample_keys(sample), + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + mean=mean, + logvar=logvar, + timestep=timestep, + ) + + +# ============================================================================ +# TaskEncoders +# ============================================================================ + + +class EncodedDiffusionTaskEncoder(DefaultTaskEncoder[DiffusionSample, DiffusionSample, dict, dict]): + """ + TaskEncoder for PRE-ENCODED diffusion data. + + Use this when your dataset contains pre-encoded features: + - latents.pth (VAE-encoded image latents) + - prompt_embeds.pth (T5 text embeddings) + - pooled_prompt_embeds.pth (CLIP pooled embeddings) + + Does NOT do any encoding - just loads from disk. + For raw data, use RawDiffusionTaskEncoder instead. + + Outputs batch with standard keys: + - 'latents' + - 'prompt_embeds' + - 'pooled_prompt_embeds' + + Use with dataset.yaml: + ```yaml + subflavors: + encoding: preencoded + ``` + """ + + decoder = SampleDecoder(image_decode="pil") + + cookers = [ + Cooker(cook_preencoded_diffusion, has_subflavors={"encoding": "preencoded"}), + Cooker(cook_preencoded_numpy_diffusion, has_subflavors={"encoding": "preencoded_numpy"}), + ] + + def __init__(self, worker_config: Optional[WorkerConfig] = None): + """Initialize pre-encoded TaskEncoder.""" + super().__init__() + self.worker_config = worker_config + logger.info("Initialized EncodedDiffusionTaskEncoder (preencoded / preencoded_numpy modes)") + + def batch(self, samples: List[DiffusionSample]) -> Dict[str, torch.Tensor]: + """ + Batch pre-encoded samples. + + Position IDs are NOT included - they are generated at runtime + in the forward step based on actual tensor shapes. + + Returns: + Dict with keys: prompt_embeds, pooled_prompt_embeds, + and conditionally latents and/or mean, logvar depending on + which fields are present in the samples. + """ + batch: Dict[str, torch.Tensor] = { + "prompt_embeds": torch.stack([s.prompt_embeds for s in samples]), + "pooled_prompt_embeds": torch.stack([s.pooled_prompt_embeds for s in samples]), + } + + if samples[0].latents is not None: + batch["latents"] = torch.stack([s.latents for s in samples]) + + if samples[0].mean is not None: + batch["mean"] = torch.stack([s.mean for s in samples]) + batch["logvar"] = torch.stack([s.logvar for s in samples]) + + if samples[0].timestep is not None: + batch["timestep"] = torch.stack([s.timestep for s in samples]) + + return batch + + +class RawDiffusionTaskEncoder(DefaultTaskEncoder): + """ + TaskEncoder for RAW diffusion data (images and text). + + Use this when your dataset contains raw files: + - images (raw image files) + - txt (text captions) + + This TaskEncoder: + - Loads raw images and captions from disk + - Does NOT do any encoding (no VAE, no T5, no CLIP) + - Encoding happens on-the-fly in model's forward_step + + Outputs batch with standard keys: + - 'images': List of PIL Images or image bytes + - 'txt': List of caption strings + + Use with dataset.yaml: + ```yaml + subflavors: + encoding: raw + ``` + """ + + decoder = SampleDecoder(image_decode="pil") + + cookers = [ + Cooker(cook_raw_images, has_subflavors={"encoding": "raw"}), + ] + + def __init__(self, worker_config: Optional[WorkerConfig] = None): + """Initialize raw diffusion TaskEncoder.""" + super().__init__() + self.worker_config = worker_config + logger.info("Initialized RawDiffusionTaskEncoder (no encoding, passes raw data)") + + def batch(self, samples: List[Dict]) -> Dict[str, Any]: + """ + Batch raw samples. + + Returns: + Dict with standard keys: + - 'images': List of PIL Images or image bytes + - 'txt': List of caption strings + """ + return { + "images": [s["images"] for s in samples], + "txt": [s["txt"] for s in samples], + } + + +__all__ = [ + "DiffusionSample", + "EncodedDiffusionTaskEncoder", + "RawDiffusionTaskEncoder", + "cook_preencoded_diffusion", + "cook_preencoded_numpy_diffusion", + "cook_raw_images", + "load_numpy_tensor", +] diff --git a/primus/backends/megatron/data/energon_dataset_provider.py b/primus/backends/megatron/data/energon_dataset_provider.py new file mode 100644 index 000000000..4ac1728dc --- /dev/null +++ b/primus/backends/megatron/data/energon_dataset_provider.py @@ -0,0 +1,282 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +EnergonDatasetProvider: Megatron Energon dataset provider. + +Provides Energon-based dataloaders for multimodal and diffusion models, +using WebDataset format with distributed loading via WorkerConfig. + +Reference: + Megatron-LM examples/multimodal/dataloader_provider.py +""" + +from typing import Any, Callable, List, Optional, Tuple + +from megatron.core import parallel_state +from megatron.core.num_microbatches_calculator import get_num_microbatches +from megatron.core.parallel_state import ( + get_pipeline_model_parallel_rank, + get_pipeline_model_parallel_world_size, + get_tensor_model_parallel_rank, +) +from megatron.energon import ( + LimitDataset, + RepeatDataset, + WorkerConfig, + get_loader, + get_savable_loader, + get_train_dataset, + get_val_datasets, +) + +from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper +from primus.backends.megatron.data.dataset_provider import DatasetProvider +from primus.core.utils.module_utils import log_rank_0 + + +class EnergonDatasetProvider(DatasetProvider): + """ + Dataset provider for multimodal/diffusion models using Megatron Energon. + + Uses: + - Megatron Energon (WebDataset format) + - Task encoders (model-specific, injected via factory) + - WorkerConfig for distributed data loading + + Architecture: + - Task encoder factory is called during setup() to create encoder + - WorkerConfig handles rank-based data sharding automatically + - Dataloaders are only created on specific ranks (first TP, first/last PP) + + Reference: + Megatron-LM examples/multimodal/dataloader_provider.py + """ + + def __init__(self, task_encoder_factory: Callable[[], Any]): + """ + Initialize Energon dataset provider. + + Args: + task_encoder_factory: Function that creates task encoder + Signature: () -> TaskEncoder + Example: lambda: EncodedDiffusionTaskEncoder(...) + + The factory pattern is used because task encoders may need + access to trainer state (e.g., self.module_config) which + isn't available during __init__. + """ + self.task_encoder_factory = task_encoder_factory + + def create_dataloaders( + self, trainer_config: Any, train_val_test_num_samples: List[int], vp_stage: Optional[int] = None + ) -> Tuple[Any, Any, Any]: + """ + Build train, validation, and test dataloaders using Energon. + + This closely follows the pattern from Megatron-LM multimodal examples, + with adaptations for Primus configuration style. + """ + from megatron.training import get_args + + args = get_args() + + # Check if we should create dataloaders on this rank + # (Only first TP rank and first/last PP stage) + if not self._is_dataloader_rank(): + log_rank_0( + "Skipping dataloader creation on this rank (not first TP rank or not first/last PP stage)" + ) + return None, None, None + + # Create task encoder using factory + task_encoder = self.task_encoder_factory() + log_rank_0(f"Created task encoder: {type(task_encoder).__name__}") + + # Create worker config for distributed loading + worker_config = self._create_worker_config(args) + + # Get data path + data_path = self._get_data_path(args) + + # Create training dataset using Energon + log_rank_0(f"Creating training dataset from: {data_path}") + train_dataset = get_train_dataset( + data_path, + batch_size=args.micro_batch_size, + task_encoder=task_encoder, + worker_config=worker_config, + virtual_epoch_length=getattr(args, "virtual_epoch_length", 1_000_000_000), + max_samples_per_sequence=getattr(args, "max_samples_per_sequence", 100), + shuffle_buffer_size=getattr(args, "shuffle_buffer_size", None), + handler=lambda *args: None, # Error handler (print errors but continue) + ) + + # Wrap in savable loader for checkpointing support + prefetch_factor = getattr(args, "prefetch_factor", 2) + log_rank_0(f"Dataloader prefetch_factor: {prefetch_factor}") + train_dataloader = get_savable_loader( + train_dataset, worker_config=worker_config, prefetch_factor=prefetch_factor + ) + train_dataloader = MegatronDataloaderWrapper(train_dataloader) + log_rank_0("Created training dataloader") + + # Create validation dataloaders if evaluation is enabled + valid_dataloaders = None + if args.eval_iters > 0: + try: + log_rank_0("Creating validation dataloaders...") + val_datasets = get_val_datasets( + data_path, + batch_size=args.micro_batch_size, + task_encoder=task_encoder, + worker_config=worker_config, + handler=lambda *args: None, + ) + + # Limit validation datasets to eval_iters * num_microbatches + val_datasets_limited = [ + LimitDataset( + RepeatDataset(val_ds, worker_config=worker_config), + length=args.eval_iters * get_num_microbatches(), + worker_config=worker_config, + reset_after_epoch=True, + ) + for val_ds, _src_ds in val_datasets + ] + + valid_dataloaders = [ + MegatronDataloaderWrapper( + get_loader(valid_ds, worker_config=worker_config, prefetch_factor=prefetch_factor) + ) + for valid_ds in val_datasets_limited + ] + log_rank_0(f"Created {len(valid_dataloaders)} validation dataloaders") + except Exception as e: + log_rank_0("=" * 80) + log_rank_0("WARNING: Could not create validation dataloaders") + log_rank_0(f"Reason: {e}") + log_rank_0("") + log_rank_0("This typically means:") + log_rank_0(" - The dataset does not have a validation split") + log_rank_0(" - The dataset path is incorrect") + log_rank_0("") + log_rank_0("AUTOMATIC FIX: Disabling evaluation (setting eval_iters = 0)") + log_rank_0("To enable evaluation, provide a dataset with validation split") + log_rank_0("=" * 80) + valid_dataloaders = None + + # Automatically disable evaluation when no validation data exists + args.eval_iters = 0 + + # Test dataloaders not implemented for Energon + test_dataloader = None + + return train_dataloader, valid_dataloaders, test_dataloader + + @property + def is_distributed(self) -> bool: + """ + Energon dataloaders are distributed (handle sharding internally). + + Returns True to tell Megatron to bypass indexed dataset logic. + """ + return True + + def _is_dataloader_rank(self) -> bool: + """ + Check if we should have the dataloader on this rank. + + Energon dataloaders should only run on: + - First tensor parallel rank (data will be broadcast to others) + - First or last pipeline parallel stage (where embeddings/outputs are) + + Reference: + Megatron-LM examples/multimodal/dataloader_provider.py:is_dataloader_rank() + """ + # Run dataloader only on first tensor parallel rank + is_first_tp_rank = get_tensor_model_parallel_rank() == 0 + + # Check pipeline parallel stage + pp_size = get_pipeline_model_parallel_world_size() + if pp_size == 1: + # No pipeline parallelism + is_valid_pp_stage = True + else: + # With pipeline parallelism, run on first and last stage + pp_rank = get_pipeline_model_parallel_rank() + is_valid_pp_stage = pp_rank in (0, pp_size - 1) + + return is_first_tp_rank and is_valid_pp_stage + + def _create_worker_config(self, args) -> WorkerConfig: + """ + Create Energon WorkerConfig for distributed loading. + + WorkerConfig tells Energon how to shard data across workers. + """ + rank = parallel_state.get_data_parallel_rank() + world_size = parallel_state.get_data_parallel_world_size() + data_parallel_group = parallel_state.get_data_parallel_group() + + return WorkerConfig( + rank=rank, + world_size=world_size, + num_workers=getattr(args, "num_workers", 4), + data_parallel_group=data_parallel_group, + ) + + def _get_data_path(self, args) -> str: + """ + Extract data path from args. + + Handles both string and list formats. + Energon's get_train_dataset() can accept either: + - Directory path (will look for .nv-meta/dataset.yaml or dataset.yaml) + - Direct path to dataset.yaml file + + This method returns the path as-is, letting Energon handle the resolution. + """ + from pathlib import Path + + data_path = args.data_path + if isinstance(data_path, list): + data_path = data_path[0] + + if not data_path: + raise ValueError("data_path not found in configuration") + + data_path = str(data_path) + path_obj = Path(data_path) + + # Check if path exists + if not path_obj.exists(): + raise ValueError( + f"data_path does not exist: {data_path}\n" + f"Please verify the path is correct and the dataset has been prepared." + ) + + # If it's a directory, check for dataset.yaml indicators + if path_obj.is_dir(): + dataset_yaml = path_obj / "dataset.yaml" + nv_meta_yaml = path_obj / ".nv-meta" / "dataset.yaml" + + if not dataset_yaml.exists() and not nv_meta_yaml.exists(): + log_rank_0("=" * 80) + log_rank_0("WARNING: No dataset.yaml found in dataset directory") + log_rank_0(f" Directory: {data_path}") + log_rank_0(f" Expected: {dataset_yaml} or {nv_meta_yaml}") + log_rank_0("") + log_rank_0("This dataset may not be properly indexed with Energon.") + log_rank_0("Run: energon prepare --num-workers 4") + log_rank_0("=" * 80) + elif nv_meta_yaml.exists(): + log_rank_0(f"Found Energon metadata at: {nv_meta_yaml}") + elif dataset_yaml.exists(): + log_rank_0(f"Found dataset.yaml at: {dataset_yaml}") + + log_rank_0(f"Using data path: {data_path}") + return data_path + + +__all__ = ["EnergonDatasetProvider"] diff --git a/primus/backends/megatron/data/synthetic/__init__.py b/primus/backends/megatron/data/synthetic/__init__.py new file mode 100644 index 000000000..ff7eb721b --- /dev/null +++ b/primus/backends/megatron/data/synthetic/__init__.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Synthetic data generation for diffusion models. + +This module provides mock datasets that generate synthetic data +with correct tensor shapes and formats for training and testing +without requiring real datasets or model weights. +""" + +from .mock_datasets import ( + LatentConfig, + MockDiffusionDataset, + MockFluxDataset, + MockFluxSchnellDataset, + ModelPreset, + PreGeneratedMockFluxDataset, + PreGeneratedMockFluxSchnellDataset, + TextEmbeddingConfig, +) + +__all__ = [ + "MockDiffusionDataset", + "MockFluxDataset", + "PreGeneratedMockFluxDataset", + "MockFluxSchnellDataset", + "PreGeneratedMockFluxSchnellDataset", + "LatentConfig", + "TextEmbeddingConfig", + "ModelPreset", +] diff --git a/primus/backends/megatron/data/synthetic/mock_datasets.py b/primus/backends/megatron/data/synthetic/mock_datasets.py new file mode 100644 index 000000000..b7aaf179a --- /dev/null +++ b/primus/backends/megatron/data/synthetic/mock_datasets.py @@ -0,0 +1,614 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Synthetic datasets for diffusion model training and testing. + +Provides mock datasets that generate random tensors with correct shapes +for development, testing, and benchmarking without requiring real data. + +The module provides a flexible preset-based system for creating mock datasets +for different diffusion model architectures. + +Quick Start +----------- + +Using a preset (easiest): + >>> from primus.backends.megatron.data.synthetic import MockFluxDataset + >>> dataset = MockFluxDataset(num_samples=100, image_size=1024) + >>> sample = dataset[0] + >>> print(sample.keys()) # latents, t5_text_embeddings, clip_pooled_embeddings, etc. + +Using the preset system directly: + >>> from primus.backends.megatron.data.synthetic import MockDiffusionDataset + >>> dataset = MockDiffusionDataset(num_samples=100, model_preset='flux') + +Custom configuration for other models: + >>> from primus.backends.megatron.data.synthetic import ( + ... MockDiffusionDataset, + ... LatentConfig, + ... TextEmbeddingConfig, + ... ) + >>> dataset = MockDiffusionDataset( + ... num_samples=100, + ... latent_config=LatentConfig(channels=4, downsample_factor=8), + ... text_embedding_configs=[ + ... TextEmbeddingConfig(key='text_embeddings', shape=(77, 768)), + ... TextEmbeddingConfig(key='pooled_embeddings', shape=(1280,)), + ... ], + ... ) + +Adding a New Model +------------------ + +To add support for a new diffusion model (e.g., SDXL): + +1. Define a position ID generator function (if needed): + >>> def sdxl_position_id_generator(image_size, latent_size, **kwargs): + ... # SDXL doesn't use position IDs + ... return {} + +2. Register the model preset: + >>> from primus.backends.megatron.data.synthetic import ( + ... MockDiffusionDataset, + ... ModelPreset, + ... LatentConfig, + ... TextEmbeddingConfig, + ... ) + >>> MockDiffusionDataset.register_model('sdxl', ModelPreset( + ... latent_config=LatentConfig(channels=4, downsample_factor=8), + ... text_embedding_configs=[ + ... TextEmbeddingConfig(key='text_embeddings', shape=(77, 2048)), + ... TextEmbeddingConfig(key='pooled_embeddings', shape=(1280,)), + ... ], + ... position_id_generator=sdxl_position_id_generator, + ... )) + +3. Create a convenience class (optional): + >>> class MockSDXLDataset(MockDiffusionDataset): + ... def __init__(self, num_samples=100, image_size=1024, seed=None): + ... super().__init__( + ... num_samples=num_samples, + ... image_size=image_size, + ... model_preset='sdxl', + ... seed=seed, + ... ) + +Architecture +------------ + +The module uses a hybrid registry + subclass approach: + +- MockDiffusionDataset: Generic base class with preset system +- MockFluxDataset: Convenience class for Flux models +- Model presets: Registered configurations for specific model architectures +- Configuration dataclasses: LatentConfig, TextEmbeddingConfig, ModelPreset + +This design provides flexibility for testing various diffusion models while +maintaining a simple API for common use cases. +""" + +import logging +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional, Tuple + +import torch +from torch.utils.data import Dataset + +from primus.backends.megatron.core.models.diffusion.flux.utils import ( + generate_image_position_ids, +) + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# Configuration Data Structures +# ============================================================================ + + +@dataclass +class LatentConfig: + """Configuration for latent generation.""" + + channels: int # Number of latent channels + downsample_factor: int = 8 # VAE spatial downsampling factor + + +@dataclass +class TextEmbeddingConfig: + """Configuration for a single text embedding.""" + + key: str # Output dictionary key (e.g., 't5_text_embeddings') + shape: Tuple[int, ...] # Shape without batch dimension + dtype: str = "float32" # Data type + + +@dataclass +class ModelPreset: + """Complete configuration preset for a diffusion model.""" + + latent_config: LatentConfig + text_embedding_configs: List[TextEmbeddingConfig] + position_id_generator: Optional[Callable] = None + + +# ============================================================================ +# Position ID Generator Functions +# ============================================================================ + + +def flux_position_id_generator( + image_size: int, + latent_size: int, + t5_seq_len: int, + seed: int, + idx: int, + device: torch.device = torch.device("cpu"), +) -> Dict[str, torch.Tensor]: + """ + Generate Flux-specific position IDs. + + Args: + image_size: Original image size + latent_size: Latent space size (after VAE downsampling) + t5_seq_len: T5 sequence length + seed: Random seed (unused but kept for consistency) + idx: Sample index (unused but kept for consistency) + device: Device to create tensors on (default: cpu) + + Returns: + Dictionary with 'img_ids' and 'txt_ids' tensors + """ + # Generate image position IDs using Flux's standard function + img_ids = generate_image_position_ids( + batch_size=1, + height=latent_size, + width=latent_size, + device=device, + dtype=torch.float32, + ).squeeze( + 0 + ) # Remove batch dimension: (1, seq_len, 3) -> (seq_len, 3) + + # Text position IDs (all zeros for text in Flux) + txt_ids = torch.zeros(t5_seq_len, 3, device=device) + + return {"img_ids": img_ids, "txt_ids": txt_ids} + + +class MockDiffusionDataset(Dataset): + """ + Generic mock dataset for diffusion models. + + Supports both preset-based initialization (via model_preset) and + custom configuration (via explicit configs). This allows the dataset + to generate synthetic data for various diffusion model architectures. + + Usage: + # Using a preset (e.g., 'flux') + dataset = MockDiffusionDataset(num_samples=100, model_preset='flux') + + # Using custom configuration + dataset = MockDiffusionDataset( + num_samples=100, + latent_config=LatentConfig(channels=4, downsample_factor=8), + text_embedding_configs=[ + TextEmbeddingConfig(key='text_embeddings', shape=(77, 768)), + ] + ) + + # Using legacy parameters (backward compatible) + dataset = MockDiffusionDataset( + num_samples=100, + latent_channels=16, + t5_seq_len=512, + t5_hidden_dim=4096, + clip_hidden_dim=768, + ) + """ + + # Registry of model presets + MODEL_PRESETS: Dict[str, ModelPreset] = {} + + @classmethod + def register_model(cls, name: str, preset: ModelPreset): + """ + Register a model preset for easy reuse. + + Args: + name: Name of the preset (e.g., 'flux', 'sdxl') + preset: ModelPreset configuration object + """ + cls.MODEL_PRESETS[name] = preset + logger.info(f"Registered model preset: {name}") + + def __init__( + self, + num_samples: int = 100, + image_size: int = 1024, + # Option 1: Use preset + model_preset: Optional[str] = None, + # Option 2: Custom configuration + latent_config: Optional[LatentConfig] = None, + text_embedding_configs: Optional[List[TextEmbeddingConfig]] = None, + position_id_generator: Optional[Callable] = None, + seed: Optional[int] = None, + dtype: torch.dtype = torch.bfloat16, + device: str = "cpu", + vae_latent_mode: str = "presampled", + is_validation: bool = False, + ): + """ + Initialize mock diffusion dataset. + + Note: Always returns preencoded format (latents + embeddings + position IDs). + This matches the output of EncodedDiffusionTaskEncoder. + + Args: + num_samples: Number of samples in dataset + image_size: Image size (assumes square images) + model_preset: Name of registered model preset to use + latent_config: Custom latent configuration (ignored if model_preset set) + text_embedding_configs: Custom text embedding configs (ignored if model_preset set) + position_id_generator: Custom position ID generator function (ignored if model_preset set) + seed: Random seed for reproducibility + dtype: Data type for tensors (default: torch.bfloat16) + device: Device to create tensors on ('cpu' or 'cuda', default: 'cpu') + """ + super().__init__() + + # Load from preset if specified + if model_preset: + if model_preset not in self.MODEL_PRESETS: + raise ValueError( + f"Unknown model preset: {model_preset}. " + f"Available presets: {list(self.MODEL_PRESETS.keys())}" + ) + preset = self.MODEL_PRESETS[model_preset] + self.latent_config = preset.latent_config + self.text_embedding_configs = preset.text_embedding_configs + self.position_id_generator = preset.position_id_generator + logger.info(f"Using model preset: {model_preset}") + # Otherwise use custom configuration + else: + if latent_config is None: + raise ValueError("Must provide either model_preset or latent_config") + self.latent_config = latent_config + self.text_embedding_configs = text_embedding_configs or [] + self.position_id_generator = position_id_generator + + # Common initialization + self.num_samples = num_samples + self.image_size = image_size + self.seed = seed if seed is not None else 0 + self.dtype = dtype + self.device = torch.device(device) + self.vae_latent_mode = vae_latent_mode + self.is_validation = is_validation + self.latent_size = image_size // self.latent_config.downsample_factor + + # Store legacy attributes for backward compatibility + self.latent_channels = self.latent_config.channels + + # Extract T5 and CLIP dims if present (for backward compatibility) + self.t5_seq_len = None + self.t5_hidden_dim = None + self.clip_hidden_dim = None + for emb_config in self.text_embedding_configs: + if emb_config.key in ("t5_text_embeddings", "prompt_embeds") and len(emb_config.shape) == 2: + self.t5_seq_len, self.t5_hidden_dim = emb_config.shape + elif ( + emb_config.key in ("clip_pooled_embeddings", "pooled_prompt_embeds") + and len(emb_config.shape) == 1 + ): + self.clip_hidden_dim = emb_config.shape[0] + + logger.info( + f"Initialized MockDiffusionDataset: {num_samples} samples, " + f"image_size={image_size}, latent_size={self.latent_size}" + ) + + def __len__(self) -> int: + """Return number of samples.""" + return self.num_samples + + def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: + """ + Get a mock sample. + + Args: + idx: Sample index + + Returns: + Dictionary with mock data matching DiffusionSample format + + Raises: + IndexError: If idx is out of bounds + """ + if idx >= self.num_samples or idx < 0: + raise IndexError(f"Index {idx} out of range for dataset with {self.num_samples} samples") + + # Set seed based on instance seed and idx for reproducibility + gen = torch.Generator(device=self.device) + gen.manual_seed(self.seed + idx) + + latent_shape = ( + self.latent_config.channels, + self.latent_size, + self.latent_size, + ) + + if self.vae_latent_mode == "resample": + mean = torch.randn(*latent_shape, generator=gen, dtype=self.dtype, device=self.device) + logvar = torch.randn(*latent_shape, generator=gen, dtype=self.dtype, device=self.device) + sample = {"mean": mean, "logvar": logvar} + else: + latents = torch.randn(*latent_shape, generator=gen, dtype=self.dtype, device=self.device) + sample = {"latents": latents} + + # Generate text embeddings based on configs + for emb_config in self.text_embedding_configs: + # Get dtype + dtype = getattr(torch, emb_config.dtype, torch.float32) + + # Generate tensor with the specified shape + embedding = torch.randn( + *emb_config.shape, + generator=gen, + dtype=dtype, + device=self.device, + ) + + sample[emb_config.key] = embedding + + # Generate position IDs if generator provided + if self.position_id_generator is not None: + position_ids = self.position_id_generator( + image_size=self.image_size, + latent_size=self.latent_size, + t5_seq_len=self.t5_seq_len if self.t5_seq_len else 512, + seed=self.seed, + idx=idx, + device=self.device, + ) + sample.update(position_ids) + + # Generate mock caption + sample["caption"] = f"Mock caption {idx}" + + if self.is_validation: + sample["timestep"] = torch.tensor(idx % 8) + + return sample + + +# ============================================================================ +# Register Model Presets +# ============================================================================ + +# Register Flux model preset (FLUX.1-dev: T5 max_sequence_length=512) +MockDiffusionDataset.register_model( + "flux", + ModelPreset( + latent_config=LatentConfig(channels=16, downsample_factor=8), + text_embedding_configs=[ + TextEmbeddingConfig( + key="prompt_embeds", shape=(512, 4096), dtype="bfloat16" + ), # T5 text embeddings + TextEmbeddingConfig( + key="pooled_prompt_embeds", shape=(768,), dtype="bfloat16" + ), # CLIP pooled embeddings + ], + position_id_generator=flux_position_id_generator, + ), +) + +# Register Flux Schnell preset (FLUX.1-schnell: T5 max_sequence_length=256) +# Matches the MLPerf Training v5.1 benchmark specification and NVIDIA reference. +MockDiffusionDataset.register_model( + "flux_schnell", + ModelPreset( + latent_config=LatentConfig(channels=16, downsample_factor=8), + text_embedding_configs=[ + TextEmbeddingConfig( + key="prompt_embeds", shape=(256, 4096), dtype="bfloat16" + ), # T5 text embeddings + TextEmbeddingConfig( + key="pooled_prompt_embeds", shape=(768,), dtype="bfloat16" + ), # CLIP pooled embeddings + ], + position_id_generator=flux_position_id_generator, + ), +) + + +# ============================================================================ +# Model-Specific Convenience Classes +# ============================================================================ + + +class MockFluxDataset(MockDiffusionDataset): + """ + Mock dataset specifically for Flux diffusion model. + + Convenience wrapper that uses the 'flux' preset with Flux-specific + default parameters: + - image_size=1024 + - latent_channels=16 + - t5_seq_len=512 + - t5_hidden_dim=4096 + - clip_hidden_dim=768 + + This class provides a simplified API for the common case of testing + Flux models without needing to understand the preset system. + """ + + def __init__( + self, + num_samples: int = 100, + image_size: int = 1024, + seed: Optional[int] = None, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", + **kwargs, + ): + """ + Initialize mock Flux dataset. + + Note: Always returns preencoded format (matches TaskEncoder output). + + Args: + num_samples: Number of samples in dataset + image_size: Image size (default 1024 for Flux) + seed: Random seed for reproducibility + dtype: Data type for tensors (default: torch.bfloat16) + device: Device to create tensors on ('cpu' or 'cuda', default: 'cuda') + """ + super().__init__( + num_samples=num_samples, + image_size=image_size, + model_preset="flux", + seed=seed, + dtype=dtype, + device=device, + **kwargs, + ) + + logger.info(f"Initialized MockFluxDataset with Flux preset on device={device}") + + +class PreGeneratedMockFluxDataset(MockFluxDataset): + """ + Pre-generated mock Flux dataset for maximum training throughput. + + This dataset generates all samples once during initialization and caches + them in memory. This eliminates on-the-fly random number generation overhead + and provides benchmark-quality performance matching latent caching approaches. + + Memory usage: For 1000 samples at 512x512 resolution with bf16: + - Latents: 1000 * 16 * 64 * 64 * 2 bytes = ~128 MB + - Text embeddings: 1000 * (512 * 4096 + 768) * 2 bytes = ~4 GB + - Total: ~4.2 GB per GPU (acceptable for 192GB MI300X) + + Use this instead of MockFluxDataset for: + - Performance benchmarking + - Maximum training throughput + - Consistent timing measurements + """ + + def __init__(self, *args, **kwargs): + """Initialize and pre-generate all samples.""" + from primus.core.utils.module_utils import log_rank_0 + + # Initialize parent class + super().__init__(*args, **kwargs) + + log_rank_0(f"Pre-generating {self.num_samples} mock Flux samples...") + log_rank_0(f" Image size: {self.image_size}x{self.image_size}") + log_rank_0(f" Latent size: {self.latent_size}x{self.latent_size}") + + # Pre-generate all samples using parent's __getitem__ + self._samples = [MockFluxDataset.__getitem__(self, i) for i in range(self.num_samples)] + + # Calculate memory usage (only count tensors) + sample_size = sum( + tensor.element_size() * tensor.numel() + for tensor in self._samples[0].values() + if hasattr(tensor, "element_size") # Only count torch tensors + ) + total_mb = (sample_size * self.num_samples) / (1024 * 1024) + + log_rank_0(f"Pre-generation complete!") + log_rank_0(f" Memory usage: {total_mb:.1f} MB") + log_rank_0(f" Data loading overhead: eliminated") + + def __getitem__(self, idx: int): + """ + Return pre-generated sample (no computation). + + This is 10-100x faster than on-the-fly generation as it's just + a memory lookup with no random number generation or tensor creation. + """ + if idx >= self.num_samples or idx < 0: + raise IndexError(f"Index {idx} out of range for dataset with {self.num_samples} samples") + return self._samples[idx] + + +class MockFluxSchnellDataset(MockDiffusionDataset): + """ + Mock dataset for FLUX.1-schnell (MLPerf Training v5.1 benchmark). + + Uses the 'flux_schnell' preset which matches the NVIDIA MLPerf reference: + - T5 max_sequence_length=256 (vs 512 for FLUX.1-dev) + - image_size=256 (benchmark default) + - latent_channels=16, context_dim=4096, vec_in_dim=768 + """ + + def __init__( + self, + num_samples: int = 100, + image_size: int = 256, + seed: Optional[int] = None, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", + **kwargs, + ): + super().__init__( + num_samples=num_samples, + image_size=image_size, + model_preset="flux_schnell", + seed=seed, + dtype=dtype, + device=device, + **kwargs, + ) + logger.info(f"Initialized MockFluxSchnellDataset on device={device}") + + +class PreGeneratedMockFluxSchnellDataset(MockFluxSchnellDataset): + """ + Pre-generated mock FLUX.1-schnell dataset for benchmark throughput. + + Generates all samples once at init and caches them in memory. + Matches the MLPerf Training v5.1 data shapes (T5 seq_len=256). + """ + + def __init__(self, *args, **kwargs): + from primus.core.utils.module_utils import log_rank_0 + + super().__init__(*args, **kwargs) + + log_rank_0(f"Pre-generating {self.num_samples} mock Flux-schnell samples...") + log_rank_0(f" Image size: {self.image_size}x{self.image_size}") + log_rank_0(f" Latent size: {self.latent_size}x{self.latent_size}") + + self._samples = [MockFluxSchnellDataset.__getitem__(self, i) for i in range(self.num_samples)] + + sample_size = sum( + tensor.element_size() * tensor.numel() + for tensor in self._samples[0].values() + if hasattr(tensor, "element_size") + ) + total_mb = (sample_size * self.num_samples) / (1024 * 1024) + + log_rank_0(f"Pre-generation complete!") + log_rank_0(f" Memory usage: {total_mb:.1f} MB") + log_rank_0(f" Data loading overhead: eliminated") + + def __getitem__(self, idx: int): + if idx >= self.num_samples or idx < 0: + raise IndexError(f"Index {idx} out of range for dataset with {self.num_samples} samples") + return self._samples[idx] + + +__all__ = [ + "LatentConfig", + "TextEmbeddingConfig", + "ModelPreset", + "MockDiffusionDataset", + "MockFluxDataset", + "PreGeneratedMockFluxDataset", + "MockFluxSchnellDataset", + "PreGeneratedMockFluxSchnellDataset", +] diff --git a/primus/backends/megatron/data/synthetic_dataset_provider.py b/primus/backends/megatron/data/synthetic_dataset_provider.py new file mode 100644 index 000000000..67588be92 --- /dev/null +++ b/primus/backends/megatron/data/synthetic_dataset_provider.py @@ -0,0 +1,217 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +SyntheticDatasetProvider: Dataset provider for synthetic/mock data. + +Provides dataloaders for synthetic data generation during development, +testing, and benchmarking without requiring real datasets. +""" + +import logging +from importlib import import_module +from typing import Any, Dict, List, Optional, Tuple + +from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper +from primus.backends.megatron.data.dataset_provider import DatasetProvider +from primus.core.utils.module_utils import log_rank_0 + +logger = logging.getLogger(__name__) + + +class SyntheticDatasetProvider(DatasetProvider): + """ + Dataset provider for synthetic/mock data. + + This provider creates synthetic datasets based on configuration. + The dataset class and parameters are specified in YAML config, + keeping dataset creation logic separate from trainer logic. + + Usage in YAML: + modules: + pre_trainer: + mock_data: true + mock_dataset: + class: "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset" + params: + num_samples: 1000 + image_size: 512 + """ + + DEFAULT_DATASETS = { + "flux": "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxDataset", + "flux_onthefly": "primus.backends.megatron.data.synthetic.MockFluxDataset", + "flux_schnell": "primus.backends.megatron.data.synthetic.PreGeneratedMockFluxSchnellDataset", + "flux_schnell_onthefly": "primus.backends.megatron.data.synthetic.MockFluxSchnellDataset", + } + + def __init__(self, dataset_config: Optional[Dict[str, Any]] = None, model_type: str = "flux"): + """ + Initialize synthetic dataset provider. + + Args: + dataset_config: Dictionary with 'class' and 'params' keys: + { + 'class': 'fully.qualified.ClassName', # Optional, uses default for model_type + 'params': { # Optional, dataset-specific parameters + 'num_samples': 1000, + 'image_size': 512, + # ... other dataset params + } + } + model_type: Model type for default dataset selection ('flux', etc.) + Used when dataset_config['class'] is not specified. + """ + self.dataset_config = dataset_config or {} + self.model_type = model_type + + # Determine dataset class to use + dataset_class_from_config = self.dataset_config.get("class") + + # If class is None or not specified, use default for model_type + if dataset_class_from_config is None or dataset_class_from_config == "null": + self.dataset_class_path = self.DEFAULT_DATASETS.get(model_type) + else: + self.dataset_class_path = dataset_class_from_config + + if not self.dataset_class_path: + raise ValueError( + f"No default dataset for model_type='{model_type}'. " + f"Available types: {list(self.DEFAULT_DATASETS.keys())}. " + f"Or specify dataset_config['class'] explicitly." + ) + + # Get dataset parameters (will be augmented with trainer config later) + self.dataset_params = self.dataset_config.get("params", {}) + + logger.debug(f"SyntheticDatasetProvider initialized: {self.dataset_class_path}") + + def _import_dataset_class(self): + """ + Dynamically import the dataset class. + + Returns: + Dataset class + + Raises: + ImportError: If class cannot be imported + AttributeError: If class doesn't exist in module + """ + try: + module_path, class_name = self.dataset_class_path.rsplit(".", 1) + module = import_module(module_path) + dataset_class = getattr(module, class_name) + return dataset_class + except (ValueError, ImportError, AttributeError) as e: + raise ImportError( + f"Failed to import dataset class '{self.dataset_class_path}': {e}\n" + f"Ensure the class path is correct and the module is installed." + ) from e + + def create_dataloaders( + self, trainer_config: Any, train_val_test_num_samples: List[int], vp_stage: Optional[int] = None + ) -> Tuple[Any, Any, Any]: + """ + Create synthetic dataloaders. + + Args: + trainer_config: Megatron args namespace (from megatron.training.get_args()) + train_val_test_num_samples: [train_samples, valid_samples, test_samples] + vp_stage: Virtual pipeline stage (for VP parallelism) + + Returns: + Tuple of (train_dataloader, None, None) + Note: Validation/test loaders not supported for synthetic data + """ + from megatron.training import get_args + from torch.utils.data import DataLoader + + args = get_args() + + log_rank_0("=" * 80) + log_rank_0("Creating SYNTHETIC/MOCK dataloaders") + log_rank_0(f"Dataset class: {self.dataset_class_path}") + log_rank_0("=" * 80) + + # Import dataset class + dataset_class = self._import_dataset_class() + + # Merge config: YAML params + runtime params from trainer + final_params = { + **self.dataset_params, # From YAML + "seed": getattr(trainer_config, "seed", 42), # Always use trainer seed + } + + # Override with trainer config if specified (backwards compatibility) + if hasattr(trainer_config, "image_size") and "image_size" not in self.dataset_params: + final_params["image_size"] = trainer_config.image_size + + # Forward vae_latent_mode from trainer config to mock dataset + if hasattr(trainer_config, "vae_latent_mode") and "vae_latent_mode" not in self.dataset_params: + final_params["vae_latent_mode"] = trainer_config.vae_latent_mode + + # Set defaults if not specified + final_params.setdefault("num_samples", 1000) + final_params.setdefault("image_size", 512) + + log_rank_0(f"Dataset parameters:") + for key, value in sorted(final_params.items()): + log_rank_0(f" {key}: {value}") + + # Create dataset + dataset = dataset_class(**final_params) + + log_rank_0(f"Created mock dataset: {dataset_class.__name__}") + log_rank_0(f" num_samples: {len(dataset)}") + log_rank_0(f" batch_size: {args.micro_batch_size}") + + # Create PyTorch DataLoader + # Note: num_workers=0 for synthetic data (generation is fast) + train_loader = DataLoader( + dataset, + batch_size=args.micro_batch_size, + shuffle=True, + num_workers=0, + drop_last=True, + ) + + # Wrap in MegatronDataloaderWrapper for: + # 1. Cyclic iteration (never exhausts) + # 2. Megatron training loop compatibility + # 3. Checkpoint interface (no-op for synthetic data) + train_loader = MegatronDataloaderWrapper(train_loader) + + log_rank_0("✓ Synthetic dataloader ready (infinite iteration)") + log_rank_0("=" * 80) + + # Create validation dataloader when eval_iters > 0 + val_loader = None + eval_iters = getattr(args, "eval_iters", 0) + if eval_iters > 0: + val_num_samples = max(256, eval_iters * args.micro_batch_size) + val_params = {**final_params, "is_validation": True, "num_samples": val_num_samples} + val_dataset = dataset_class(**val_params) + val_raw_loader = DataLoader( + val_dataset, + batch_size=args.micro_batch_size, + shuffle=False, + num_workers=0, + drop_last=True, + ) + val_loader = MegatronDataloaderWrapper(val_raw_loader) + log_rank_0(f"Created validation dataloader: {val_num_samples} samples, is_validation=True") + + return train_loader, val_loader, None + + @property + def is_distributed(self) -> bool: + """ + Synthetic data doesn't need distributed handling. + + Returns True to tell Megatron to bypass indexed dataset logic. + Each rank creates its own synthetic data independently. + """ + return True + + +__all__ = ["SyntheticDatasetProvider"] diff --git a/requirements.txt b/requirements.txt index 9d3c7ac6e..d8beb163c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,3 +16,7 @@ pyrsmi plotext hip-python git+https://github.com/mlcommons/logging.git@6.0.0-rc5 +megatron-energon==7.3.2 +webdataset==1.0.2 +tqdm==4.67.3 +pyarrow==21.0.0 diff --git a/tests/unit_tests/backends/megatron/diffusion/data/__init__.py b/tests/unit_tests/backends/megatron/diffusion/data/__init__.py new file mode 100644 index 000000000..89778402a --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. diff --git a/tests/unit_tests/backends/megatron/diffusion/data/conftest.py b/tests/unit_tests/backends/megatron/diffusion/data/conftest.py new file mode 100644 index 000000000..97001e263 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/conftest.py @@ -0,0 +1,257 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Shared fixtures for data module tests. + +Provides mocked encoders, sample data, and utilities for testing. +""" + +import io +import tempfile +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +import torch +from PIL import Image + +# ============================================================================ +# Mock Encoder Fixtures +# ============================================================================ + + +@pytest.fixture +def mock_vae_model(): + """Mock VAE model for testing.""" + mock = MagicMock() + + # Mock encode method + def mock_encode(images): + batch_size = images.shape[0] if isinstance(images, torch.Tensor) else 1 + # Return mock latent distribution + latent_dist = MagicMock() + latent_dist.sample = MagicMock(return_value=torch.randn(batch_size, 16, 64, 64)) + return MagicMock(latent_dist=latent_dist) + + mock.encode = mock_encode + mock.to = MagicMock(return_value=mock) + mock.eval = MagicMock(return_value=mock) + + return mock + + +@pytest.fixture +def mock_t5_model(): + """Mock T5 model for testing.""" + mock = MagicMock() + + # Mock encoder + def mock_forward(input_ids, **kwargs): + batch_size = input_ids.shape[0] + seq_len = input_ids.shape[1] + hidden_dim = 4096 # T5-XXL hidden dim + return MagicMock(last_hidden_state=torch.randn(batch_size, seq_len, hidden_dim)) + + mock.encoder = MagicMock() + mock.encoder.return_value = mock_forward(torch.zeros(1, 77, dtype=torch.long)) + mock.to = MagicMock(return_value=mock) + mock.eval = MagicMock(return_value=mock) + + return mock + + +@pytest.fixture +def mock_t5_tokenizer(): + """Mock T5 tokenizer for testing.""" + mock = MagicMock() + + def mock_tokenize(text, **kwargs): + if isinstance(text, str): + text = [text] + batch_size = len(text) + max_length = kwargs.get("max_length", 77) + return { + "input_ids": torch.randint(0, 32000, (batch_size, max_length)), + "attention_mask": torch.ones(batch_size, max_length), + } + + mock.return_value = mock_tokenize(["test"]) + mock.side_effect = None + mock.__call__ = mock_tokenize + + return mock + + +@pytest.fixture +def mock_clip_model(): + """Mock CLIP model for testing.""" + mock = MagicMock() + + # Mock text encoder + def mock_text_model(input_ids, **kwargs): + batch_size = input_ids.shape[0] + seq_len = input_ids.shape[1] + hidden_dim = 768 # CLIP-L hidden dim + pooled_dim = 768 + return MagicMock( + last_hidden_state=torch.randn(batch_size, seq_len, hidden_dim), + pooler_output=torch.randn(batch_size, pooled_dim), + ) + + mock.text_model = MagicMock() + mock.text_model.return_value = mock_text_model(torch.zeros(1, 77, dtype=torch.long)) + mock.to = MagicMock(return_value=mock) + mock.eval = MagicMock(return_value=mock) + + return mock + + +@pytest.fixture +def mock_clip_tokenizer(): + """Mock CLIP tokenizer for testing.""" + mock = MagicMock() + + def mock_tokenize(text, **kwargs): + if isinstance(text, str): + text = [text] + batch_size = len(text) + max_length = kwargs.get("max_length", 77) + return { + "input_ids": torch.randint(0, 49408, (batch_size, max_length)), + "attention_mask": torch.ones(batch_size, max_length), + } + + mock.return_value = mock_tokenize(["test"]) + mock.__call__ = mock_tokenize + + return mock + + +# ============================================================================ +# Sample Data Fixtures +# ============================================================================ + + +@pytest.fixture +def sample_image(): + """Create a sample PIL Image.""" + return Image.new("RGB", (512, 512), color="red") + + +@pytest.fixture +def sample_image_tensor(): + """Create a sample image tensor.""" + return torch.randn(3, 512, 512) + + +@pytest.fixture +def sample_image_bytes(): + """Create sample image as bytes.""" + img = Image.new("RGB", (512, 512), color="blue") + buf = io.BytesIO() + img.save(buf, format="PNG") + return buf.getvalue() + + +@pytest.fixture +def sample_text(): + """Sample text caption.""" + return "A beautiful landscape with mountains and trees" + + +@pytest.fixture +def sample_latents(): + """Sample VAE latents (without batch dimension for single sample).""" + return torch.randn(16, 64, 64) + + +@pytest.fixture +def sample_prompt_embeds(): + """Sample T5 text embeddings (without batch dimension for single sample).""" + return torch.randn(256, 4096) + + +@pytest.fixture +def sample_pooled_prompt_embeds(): + """Sample CLIP pooled embeddings (without batch dimension for single sample).""" + return torch.randn(768) + + +@pytest.fixture +def sample_text_ids(): + """Sample text position IDs (without batch dimension for single sample).""" + return torch.randn(256, 3) + + +# ============================================================================ +# WebDataset Sample Fixtures +# ============================================================================ + + +@pytest.fixture +def preencoded_sample(sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds): + """Sample with pre-encoded data (tensor format).""" + return { + "__key__": "sample_001", + "__restore_key__": lambda: "sample_001", + "__subflavors__": {"encoding": "preencoded"}, + "latents.pth": sample_latents, # Already correct shape (16, 64, 64) + "prompt_embeds.pth": sample_prompt_embeds, # Already correct shape (256, 4096) + "pooled_prompt_embeds.pth": sample_pooled_prompt_embeds, # Already correct shape (768,) + "text_ids.pth": None, # Optional + "caption.txt": b"A test caption", + } + + +@pytest.fixture +def preencoded_sample_bytes(sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds): + """Sample with pre-encoded data (bytes format).""" + + def tensor_to_bytes(tensor): + buf = io.BytesIO() + torch.save(tensor, buf) + buf.seek(0) + return buf.getvalue() + + return { + "__key__": "sample_002", + "__restore_key__": lambda: "sample_002", + "__subflavors__": {"encoding": "preencoded"}, + "latents.pth": tensor_to_bytes(sample_latents), # Convert to bytes + "prompt_embeds.pth": tensor_to_bytes(sample_prompt_embeds), # Convert to bytes + "pooled_prompt_embeds.pth": tensor_to_bytes(sample_pooled_prompt_embeds), # Convert to bytes + "caption.txt": b"Another test caption", + } + + +@pytest.fixture +def raw_sample(sample_image_bytes): + """Sample with raw image data.""" + return { + "__key__": "sample_003", + "__restore_key__": lambda: "sample_003", + "__subflavors__": {"encoding": "raw"}, + "images": sample_image_bytes, + "txt": b"Raw image caption", + } + + +# ============================================================================ +# Temporary Directory Fixtures +# ============================================================================ + + +@pytest.fixture +def temp_output_dir(): + """Create temporary directory for test outputs.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) + + +@pytest.fixture +def temp_model_path(temp_output_dir): + """Create a temporary model path.""" + model_dir = temp_output_dir / "model" + model_dir.mkdir() + return model_dir diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/__init__.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/__init__.py new file mode 100644 index 000000000..89778402a --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_base.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_base.py new file mode 100644 index 000000000..265ac19c0 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_base.py @@ -0,0 +1,130 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for base encoder classes. + +Tests abstract base classes for encoders, ensuring proper methods and behavior. +""" + +from unittest.mock import Mock + +import pytest +import torch + +from primus.backends.megatron.data.diffusion.encoders.base import ( + BaseEncoder, + BaseTextEncoder, + BaseVAE, +) +from primus.backends.megatron.data.diffusion.encoders.config import ( + EncoderConfig, + T5XXLConfig, + VAEConfig, +) +from tests.utils import PrimusUT + + +# Concrete implementations for testing abstract classes +class ConcreteEncoder(BaseEncoder): + """Concrete encoder for testing BaseEncoder.""" + + def encode(self, *args, **kwargs): + return torch.randn(1, 10) + + @classmethod + def from_pretrained(cls, model_path, config=None): + return cls(config or EncoderConfig(type="test", model_path=model_path)) + + +class ConcreteVAE(BaseVAE): + """Concrete VAE for testing BaseVAE.""" + + def encode(self, images): + batch_size = images.shape[0] + latent_h = images.shape[2] // self.latent_downsample_factor + latent_w = images.shape[3] // self.latent_downsample_factor + return torch.randn(batch_size, self.out_channels, latent_h, latent_w) + + def decode(self, latents): + batch_size = latents.shape[0] + img_h = latents.shape[2] * self.latent_downsample_factor + img_w = latents.shape[3] * self.latent_downsample_factor + return torch.randn(batch_size, self.in_channels, img_h, img_w) + + @classmethod + def from_pretrained(cls, model_path, config=None): + return cls(config or VAEConfig(model_path=model_path)) + + +class ConcreteTextEncoder(BaseTextEncoder): + """Concrete text encoder for testing BaseTextEncoder.""" + + def __init__(self, config): + super().__init__(config) + self._tokenizer = Mock() # Mock tokenizer + + def encode(self, texts, **kwargs): + texts_list = self._prepare_texts(texts) + batch_size = len(texts_list) + return torch.randn(batch_size, self.max_length, self.embedding_dim) + + @classmethod + def from_pretrained(cls, model_path, config=None): + return cls(config or T5XXLConfig(model_path=model_path)) + + +class TestBaseEncoder(PrimusUT): + """Tests for BaseEncoder abstract class.""" + + def test_get_dtype_invalid_defaults_to_bf16(self): + """Test that invalid precision defaults to bfloat16.""" + config = EncoderConfig(type="test", model_path="/path") + encoder = ConcreteEncoder(config) + + assert encoder._get_dtype("unknown") == torch.bfloat16 + + +class TestBaseVAE(PrimusUT): + """Tests for BaseVAE abstract class.""" + + def test_get_latent_shape_various_sizes(self): + """Test latent shape calculation with various image sizes.""" + config = VAEConfig(type="autoencoder_kl", model_path="/path/to/vae") + vae = ConcreteVAE(config) + + test_cases = [ + ((256, 256), (16, 32, 32)), + ((512, 512), (16, 64, 64)), + ((1024, 1024), (16, 128, 128)), + ((512, 1024), (16, 64, 128)), + ] + + for (h, w), expected_shape in test_cases: + assert vae.get_latent_shape(h, w) == expected_shape + + def test_get_latent_shape_custom_downsample(self): + """Test latent shape with custom downsample factor.""" + config = VAEConfig(type="autoencoder_kl", model_path="/path/to/vae", latent_downsample_factor=16) + vae = ConcreteVAE(config) + + latent_shape = vae.get_latent_shape(512, 512) + + assert latent_shape == (16, 32, 32) # 512/16 = 32 + + +class TestBaseTextEncoder(PrimusUT): + """Tests for BaseTextEncoder abstract class.""" + + def test_tokenizer_not_initialized_error(self): + """Test that accessing tokenizer before initialization raises error.""" + config = T5XXLConfig(type="t5_xxl", model_path="/path/to/t5") + encoder = ConcreteTextEncoder(config) + encoder._tokenizer = None + + with pytest.raises(ValueError, match="Tokenizer not initialized"): + _ = encoder.tokenizer + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_clip.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_clip.py new file mode 100644 index 000000000..96fdd9b31 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_clip.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for CLIP text encoder implementations. + +Tests CLIPLEncoder with mocked models for fast unit testing. + +NOTE: Common wrapper tests (initialization, from_pretrained, device handling, etc.) +have been moved to test_encoder_wrappers_consolidated.py. This file contains only +CLIP-specific tests that verify Primus's CLIP integration logic. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from primus.backends.megatron.data.diffusion.encoders.config import CLIPLConfig + + +class TestCLIPLEncoder: + """Tests for CLIPLEncoder Primus-specific logic.""" + + @patch("primus.backends.megatron.data.diffusion.encoders.text.clip_l.CLIPTextModel") + @patch("primus.backends.megatron.data.diffusion.encoders.text.clip_l.CLIPTokenizer") + def test_pooled_embeddings_extraction( + self, mock_tokenizer_cls, mock_model_cls, mock_clip_model, mock_clip_tokenizer + ): + """Test that pooled embeddings are correctly extracted (Primus-specific CLIP integration).""" + from primus.backends.megatron.data.diffusion.encoders.text import CLIPLEncoder + + mock_model_cls.from_pretrained.return_value = mock_clip_model + mock_tokenizer_cls.from_pretrained.return_value = mock_clip_tokenizer + + # Setup mock tokenizer + mock_clip_tokenizer.return_value = { + "input_ids": torch.randint(0, 49408, (2, 77)), + "attention_mask": torch.ones(2, 77), + } + + # Setup mock model output with specific pooled embeddings + expected_pooled = torch.randn(2, 768) + mock_clip_model.return_value = MagicMock( + last_hidden_state=torch.randn(2, 77, 768), pooler_output=expected_pooled + ) + + config = CLIPLConfig(type="clip_l", model_path="/path/to/clip") + encoder = CLIPLEncoder.from_pretrained("/path/to/clip", config=config) + + sequence_embeds, pooled_embeds = encoder.encode(["Text 1", "Text 2"]) + + # Must return the model's pooler_output verbatim (not e.g. a mean-pool of + # last_hidden_state, which would have the same shape but wrong values). + assert pooled_embeds.shape == expected_pooled.shape + assert torch.equal(pooled_embeds, expected_pooled), "pooled_embeds must equal pooler_output" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_config.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_config.py new file mode 100644 index 000000000..d21937f0d --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_config.py @@ -0,0 +1,30 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for encoder configuration validation logic. +""" + +import pytest + +from primus.backends.megatron.data.diffusion.encoders.config import ( + CLIPLConfig, + EncoderConfig, + T5XXLConfig, + VAEConfig, +) +from tests.utils import PrimusUT + + +class TestConfigValidation(PrimusUT): + """Tests for config validation across all types.""" + + def test_all_configs_validate_precision(self): + """Test that all config types validate precision.""" + for ConfigClass in [EncoderConfig, VAEConfig, T5XXLConfig, CLIPLConfig]: + with pytest.raises(ValueError, match="precision must be one of"): + ConfigClass(type="test", model_path="/path", precision="invalid_precision") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_encoder_wrappers_consolidated.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_encoder_wrappers_consolidated.py new file mode 100644 index 000000000..3baef0dbf --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_encoder_wrappers_consolidated.py @@ -0,0 +1,405 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Consolidated tests for encoder wrapper classes (CLIP, T5, VAE). + +This file contains parametrized tests for common wrapper functionality across +different encoder types to avoid duplication. Tests that are specific to individual +encoder implementations (e.g., VAE scale/shift, CLIP pooled embeddings) remain +in their respective test files. + +Encoder-specific tests remain in: +- test_clip.py: CLIP-specific logic (pooled embeddings extraction) +- test_t5.py: T5-specific logic (padding/truncation) +- test_vae.py: VAE-specific logic (scale/shift, latent shape calculation) +""" + +from unittest.mock import MagicMock, Mock, patch + +import pytest +import torch + +from primus.backends.megatron.data.diffusion.encoders.config import ( + CLIPLConfig, + T5XXLConfig, + VAEConfig, +) + + +@pytest.fixture +def encoder_specs(): + """Define specifications for all encoder types.""" + return { + "clip": { + "module_path": "primus.backends.megatron.data.diffusion.encoders.text.clip_l", + "encoder_cls_name": "CLIPLEncoder", + "config_cls": CLIPLConfig, + "config_kwargs": {"type": "clip_l", "model_path": "/path/to/clip"}, + "model_cls_name": "CLIPTextModel", + "tokenizer_cls_name": "CLIPTokenizer", + "max_length": 77, + "embedding_dim": 768, + "has_tokenizer": True, + }, + "t5": { + "module_path": "primus.backends.megatron.data.diffusion.encoders.text.t5_xxl", + "encoder_cls_name": "T5XXLEncoder", + "config_cls": T5XXLConfig, + "config_kwargs": {"type": "t5_xxl", "model_path": "/path/to/t5"}, + "model_cls_name": "T5EncoderModel", + "tokenizer_cls_name": "T5Tokenizer", + "max_length": 512, + "embedding_dim": 4096, + "has_tokenizer": True, + }, + "vae": { + "module_path": "primus.backends.megatron.data.diffusion.encoders.image.autoencoder_kl", + "encoder_cls_name": "AutoencoderKL", + "config_cls": VAEConfig, + "config_kwargs": {"type": "autoencoder_kl", "model_path": "/path/to/vae"}, + "model_cls_name": "DiffusersAutoencoderKL", + "tokenizer_cls_name": None, + "scale_factor": 0.3611, + "shift_factor": 0.1159, + "in_channels": 3, + "out_channels": 16, + "has_tokenizer": False, + }, + } + + +@pytest.fixture +def mock_clip_model(): + """Create a mock CLIP model.""" + model = Mock() + model.return_value = MagicMock( + last_hidden_state=torch.randn(1, 77, 768), pooler_output=torch.randn(1, 768) + ) + return model + + +@pytest.fixture +def mock_clip_tokenizer(): + """Create a mock CLIP tokenizer.""" + tokenizer = Mock() + tokenizer.return_value = { + "input_ids": torch.randint(0, 49408, (1, 77)), + "attention_mask": torch.ones(1, 77), + } + return tokenizer + + +@pytest.fixture +def mock_t5_model(): + """Create a mock T5 model.""" + model = Mock() + model.return_value = MagicMock(last_hidden_state=torch.randn(1, 512, 4096)) + return model + + +@pytest.fixture +def mock_t5_tokenizer(): + """Create a mock T5 tokenizer.""" + tokenizer = Mock() + tokenizer.return_value = { + "input_ids": torch.randint(0, 32000, (1, 512)), + "attention_mask": torch.ones(1, 512), + } + return tokenizer + + +@pytest.fixture +def mock_vae_model(): + """Create a mock VAE model.""" + model = Mock() + # Mock encode method + mock_dist = Mock() + mock_dist.sample.return_value = torch.randn(1, 16, 64, 64) + model.encode.return_value = mock_dist + # Mock decode method - return tuple like real diffusers AutoencoderKL + model.decode.return_value = (torch.randn(1, 3, 512, 512),) + return model + + +class TestEncoderFromPretrained: + """Test from_pretrained for all encoder types.""" + + @pytest.mark.parametrize("encoder_type", ["clip", "t5"]) + def test_text_encoder_from_pretrained_without_config( + self, + encoder_type, + encoder_specs, + mock_clip_model, + mock_clip_tokenizer, + mock_t5_model, + mock_t5_tokenizer, + ): + """Test loading text encoders without explicit config.""" + spec = encoder_specs[encoder_type] + + # Select appropriate mocks + if encoder_type == "clip": + mock_model = mock_clip_model + mock_tokenizer = mock_clip_tokenizer + else: # t5 + mock_model = mock_t5_model + mock_tokenizer = mock_t5_tokenizer + + with patch(f"{spec['module_path']}.{spec['model_cls_name']}") as mock_model_cls, patch( + f"{spec['module_path']}.{spec['tokenizer_cls_name']}" + ) as mock_tokenizer_cls: + + mock_model_cls.from_pretrained.return_value = mock_model + mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer + + # Import encoder class + module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]]) + encoder_cls = getattr(module, spec["encoder_cls_name"]) + + encoder = encoder_cls.from_pretrained("/path") + + assert encoder is not None + assert encoder.config.model_path == "/path" + assert encoder.config.precision == "bf16" + + +class TestEncoderPrecisionMapping: + """Test precision to dtype mapping for all encoders.""" + + @pytest.mark.parametrize( + "encoder_type,precision,expected_dtype", + [ + ("clip", "bf16", torch.bfloat16), + ("clip", "fp16", torch.float16), + ("clip", "fp32", torch.float32), + ("t5", "bf16", torch.bfloat16), + ("t5", "fp16", torch.float16), + ("t5", "fp32", torch.float32), + ("vae", "bf16", torch.bfloat16), + ("vae", "fp16", torch.float16), + ("vae", "fp32", torch.float32), + ], + ) + def test_precision_dtype_mapping( + self, + encoder_type, + precision, + expected_dtype, + encoder_specs, + mock_clip_model, + mock_clip_tokenizer, + mock_t5_model, + mock_t5_tokenizer, + mock_vae_model, + ): + """Test that precision config is correctly mapped to torch dtype.""" + spec = encoder_specs[encoder_type] + + # Select appropriate mocks + if encoder_type == "clip": + mock_model = mock_clip_model + mock_tokenizer = mock_clip_tokenizer + elif encoder_type == "t5": + mock_model = mock_t5_model + mock_tokenizer = mock_t5_tokenizer + else: # vae + mock_model = mock_vae_model + mock_tokenizer = None + + if spec["has_tokenizer"]: + with patch(f"{spec['module_path']}.{spec['model_cls_name']}") as mock_model_cls, patch( + f"{spec['module_path']}.{spec['tokenizer_cls_name']}" + ) as mock_tokenizer_cls: + + mock_model_cls.from_pretrained.return_value = mock_model + mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer + + # Import encoder class + module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]]) + encoder_cls = getattr(module, spec["encoder_cls_name"]) + + config = spec["config_cls"](**spec["config_kwargs"], precision=precision) + _ = encoder_cls.from_pretrained("/path", config=config) + + # Check that from_pretrained was called with correct dtype + call_kwargs = mock_model_cls.from_pretrained.call_args[1] + assert call_kwargs["torch_dtype"] == expected_dtype + else: + # VAE case + with patch(f"{spec['module_path']}.{spec['model_cls_name']}") as mock_model_cls: + mock_model_cls.from_pretrained.return_value = mock_model + + # Import encoder class + module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]]) + encoder_cls = getattr(module, spec["encoder_cls_name"]) + + config = spec["config_cls"](**spec["config_kwargs"], precision=precision) + encoder_cls.from_pretrained("/path", config=config) + + # Check that from_pretrained was called with correct dtype + call_kwargs = mock_model_cls.from_pretrained.call_args[1] + assert call_kwargs["torch_dtype"] == expected_dtype + + +class TestEncoderDeviceHandling: + """Test device handling for all encoder types.""" + + @pytest.mark.parametrize("encoder_type", ["clip", "t5", "vae"]) + def test_device_handling( + self, + encoder_type, + encoder_specs, + mock_clip_model, + mock_clip_tokenizer, + mock_t5_model, + mock_t5_tokenizer, + mock_vae_model, + ): + """Test encoder device configuration.""" + spec = encoder_specs[encoder_type] + + # Select appropriate mocks + if encoder_type == "clip": + mock_model = mock_clip_model + mock_tokenizer = mock_clip_tokenizer + elif encoder_type == "t5": + mock_model = mock_t5_model + mock_tokenizer = mock_t5_tokenizer + else: # vae + mock_model = mock_vae_model + mock_tokenizer = None + + if spec["has_tokenizer"]: + with patch(f"{spec['module_path']}.{spec['model_cls_name']}") as mock_model_cls, patch( + f"{spec['module_path']}.{spec['tokenizer_cls_name']}" + ) as mock_tokenizer_cls: + + mock_model_cls.from_pretrained.return_value = mock_model + mock_tokenizer_cls.from_pretrained.return_value = mock_tokenizer + + # Import encoder class + module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]]) + encoder_cls = getattr(module, spec["encoder_cls_name"]) + + config = spec["config_cls"](**spec["config_kwargs"], device="cpu") + encoder = encoder_cls.from_pretrained("/path", config=config) + + assert encoder.device.type == "cpu" + else: + # VAE case + with patch(f"{spec['module_path']}.{spec['model_cls_name']}") as mock_model_cls: + mock_model_cls.from_pretrained.return_value = mock_model + + # Import encoder class + module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]]) + encoder_cls = getattr(module, spec["encoder_cls_name"]) + + config = spec["config_cls"](**spec["config_kwargs"], device="cpu") + vae = encoder_cls.from_pretrained("/path", config=config) + + assert vae.device.type == "cpu" + + +class TestEncoderImportErrors: + """Test import error handling when external libraries are not available.""" + + @pytest.mark.parametrize("encoder_type", ["clip", "t5"]) + def test_text_encoder_import_error(self, encoder_type, encoder_specs): + """Test that ImportError is raised when transformers is not installed.""" + spec = encoder_specs[encoder_type] + + with patch(f"{spec['module_path']}.{spec['model_cls_name']}", None), patch( + f"{spec['module_path']}.{spec['tokenizer_cls_name']}", None + ): + + # Import encoder class + module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]]) + encoder_cls = getattr(module, spec["encoder_cls_name"]) + + config = spec["config_cls"](**spec["config_kwargs"]) + + with pytest.raises(ImportError, match="transformers library is required"): + encoder_cls(config) + + def test_vae_import_error(self, encoder_specs): + """Test that ImportError is raised when diffusers is not installed.""" + spec = encoder_specs["vae"] + + with patch(f"{spec['module_path']}.{spec['model_cls_name']}", None): + # Import encoder class + module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]]) + encoder_cls = getattr(module, spec["encoder_cls_name"]) + + config = spec["config_cls"](**spec["config_kwargs"]) + + with pytest.raises(ImportError, match="diffusers library is required"): + encoder_cls(config) + + +class TestEncoderCacheDir: + """Test cache_dir configuration support across all encoders.""" + + @pytest.mark.parametrize("encoder_type", ["clip", "t5", "vae"]) + def test_cache_dir_passed_to_from_pretrained(self, encoder_type, encoder_specs): + """Test that cache_dir is passed through to HuggingFace from_pretrained.""" + spec = encoder_specs[encoder_type] + config_cls = spec["config_cls"] + config_kwargs = spec["config_kwargs"].copy() + cache_dir = "/custom/cache/path" + + # Create config with cache_dir + config = config_cls(**config_kwargs, cache_dir=cache_dir) + + # Mock the model/tokenizer classes + mock_model = MagicMock() + mock_model_instance = Mock() + mock_model.from_pretrained.return_value = mock_model_instance + mock_model_instance.to.return_value = mock_model_instance + + # FIXED: Use just attribute names for patch.multiple (not full module paths) + patches = {spec["model_cls_name"]: mock_model} + + # Add tokenizer patch for text encoders + if spec["has_tokenizer"]: + mock_tokenizer = MagicMock() + mock_tokenizer.from_pretrained.return_value = Mock() + patches[spec["tokenizer_cls_name"]] = mock_tokenizer + + # Also mock the helper function + mock_helper = MagicMock(return_value=mock_model_instance) + patches["load_pretrained_with_subfolder_fallback"] = mock_helper + + with patch.dict("sys.modules", {k: MagicMock() for k in ["transformers", "diffusers"]}): + with patch.multiple(spec["module_path"], **patches): + # Import and call from_pretrained + module = __import__(spec["module_path"], fromlist=[spec["encoder_cls_name"]]) + encoder_cls = getattr(module, spec["encoder_cls_name"]) + + encoder_cls.from_pretrained("/path", config=config) + + # The helper must be invoked, and cache_dir must be forwarded to it. + assert mock_helper.called, "load_pretrained_with_subfolder_fallback was never called" + call_kwargs = mock_helper.call_args.kwargs + assert "cache_dir" in call_kwargs, "cache_dir not passed to helper function" + assert call_kwargs["cache_dir"] == cache_dir + + def test_flux_encoder_config_cache_dir(self): + """Test that FluxEncoderConfig.from_pretrained_flux accepts cache_dir.""" + from primus.backends.megatron.data.diffusion.encoders.config import ( + FluxEncoderConfig, + ) + + cache_dir = "/shared/hf_cache" + flux_config = FluxEncoderConfig.from_pretrained_flux( + model_path="black-forest-labs/FLUX.1-dev", cache_dir=cache_dir + ) + + # All sub-configs should have the same cache_dir + assert flux_config.vae.cache_dir == cache_dir + assert flux_config.t5.cache_dir == cache_dir + assert flux_config.clip.cache_dir == cache_dir + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_t5.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_t5.py new file mode 100644 index 000000000..37c583140 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_t5.py @@ -0,0 +1,57 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for T5 text encoder implementations. + +Tests T5XXLEncoder with mocked models for fast unit testing. + +NOTE: Common wrapper tests (initialization, from_pretrained, device handling, etc.) +have been moved to test_encoder_wrappers_consolidated.py. This file contains only +T5-specific tests that verify Primus's T5 integration logic. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from primus.backends.megatron.data.diffusion.encoders.config import T5XXLConfig + + +class TestT5XXLEncoder: + """Tests for T5XXLEncoder Primus-specific logic.""" + + @patch("primus.backends.megatron.data.diffusion.encoders.text.t5_xxl.T5EncoderModel") + @patch("primus.backends.megatron.data.diffusion.encoders.text.t5_xxl.T5Tokenizer") + def test_encode_with_padding_truncation( + self, mock_tokenizer_cls, mock_model_cls, mock_t5_model, mock_t5_tokenizer + ): + """Test that tokenizer is called with correct padding/truncation (Primus tokenization logic).""" + from primus.backends.megatron.data.diffusion.encoders.text import T5XXLEncoder + + mock_model_cls.from_pretrained.return_value = mock_t5_model + mock_tokenizer_cls.from_pretrained.return_value = mock_t5_tokenizer + + # Setup mock tokenizer + mock_t5_tokenizer.return_value = { + "input_ids": torch.randint(0, 32000, (1, 512)), + "attention_mask": torch.ones(1, 512), + } + + # Setup mock model output + mock_t5_model.return_value = MagicMock(last_hidden_state=torch.randn(1, 512, 4096)) + + config = T5XXLConfig(type="t5_xxl", model_path="/path/to/t5") + encoder = T5XXLEncoder.from_pretrained("/path/to/t5", config=config) + + encoder.encode("Test caption") + + # Verify tokenizer was called with correct Primus padding/truncation params + call_kwargs = mock_t5_tokenizer.call_args[1] + assert call_kwargs["padding"] == "max_length" + assert call_kwargs["truncation"] is True + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_vae.py b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_vae.py new file mode 100644 index 000000000..4c936c995 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/encoders/test_vae.py @@ -0,0 +1,58 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for VAE encoder implementations. + +Tests AutoencoderKL with mocked models for fast unit testing. + +NOTE: Common wrapper tests (initialization, from_pretrained, device handling, etc.) +have been moved to test_encoder_wrappers_consolidated.py. This file contains only +VAE-specific tests that verify Primus's VAE integration logic (scale/shift, latent shapes). +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from primus.backends.megatron.data.diffusion.encoders.config import VAEConfig + + +class TestAutoencoderKL: + """Tests for AutoencoderKL Primus-specific logic.""" + + @patch("primus.backends.megatron.data.diffusion.encoders.image.autoencoder_kl.DiffusersAutoencoderKL") + def test_encode_applies_scale_and_shift(self, mock_diffusers_vae, mock_vae_model): + """encode() must apply z = scale_factor * (sample - shift_factor) to the VAE sample.""" + from primus.backends.megatron.data.diffusion.encoders.image import AutoencoderKL + + mock_diffusers_vae.from_pretrained.return_value = mock_vae_model + + config = VAEConfig( + type="autoencoder_kl", model_path="/path/to/vae", scale_factor=0.5, shift_factor=0.1 + ) + vae = AutoencoderKL.from_pretrained("/path/to/vae", config=config) + + assert vae.config.scale_factor == 0.5 + assert vae.config.shift_factor == 0.1 + + # Pin the VAE's sampled latent so we can assert the exact scale/shift math + # (the shared fixture returns a fresh random sample each call). + known_sample = torch.randn(1, 16, 64, 64) + latent_dist = MagicMock() + latent_dist.sample = MagicMock(return_value=known_sample) + vae.vae.encode = MagicMock(return_value=MagicMock(latent_dist=latent_dist)) + + images = torch.randn(1, 3, 256, 256) + latents = vae.encode(images) + + expected = config.scale_factor * ( + known_sample.to(latents.device, latents.dtype) - config.shift_factor + ) + assert latents.shape == known_sample.shape + torch.testing.assert_close(latents, expected, rtol=1e-5, atol=1e-6) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/__init__.py b/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/__init__.py new file mode 100644 index 000000000..89778402a --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. diff --git a/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py b/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py new file mode 100644 index 000000000..2bac33b3f --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/task_encoders/test_task_encoders.py @@ -0,0 +1,192 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for diffusion TaskEncoders. + +Tests both EncodedDiffusionTaskEncoder (pre-encoded data) and +RawDiffusionTaskEncoder (raw images and text). +""" + + +import pytest + +from primus.backends.megatron.data.diffusion.task_encoders import ( + DiffusionSample, + EncodedDiffusionTaskEncoder, + RawDiffusionTaskEncoder, + cook_preencoded_diffusion, + cook_raw_images, +) + + +class TestCookPreencodedDiffusion: + """Tests for cook_preencoded_diffusion cooker function.""" + + def test_cook_with_tensor_data(self, preencoded_sample): + """Test cooking with direct tensor data.""" + result = cook_preencoded_diffusion(preencoded_sample) + + assert isinstance(result, DiffusionSample) + assert result.latents.shape == (16, 64, 64) + assert result.prompt_embeds.shape == (256, 4096) + assert result.pooled_prompt_embeds.shape == (768,) + assert result.caption == "A test caption" + + def test_cook_with_bytes_data(self, preencoded_sample_bytes): + """Test cooking with tensor data as bytes.""" + result = cook_preencoded_diffusion(preencoded_sample_bytes) + + assert isinstance(result, DiffusionSample) + assert result.latents.shape == (16, 64, 64) + assert result.prompt_embeds.shape == (256, 4096) + assert result.pooled_prompt_embeds.shape == (768,) + assert result.caption == "Another test caption" + + def test_cook_ignores_text_ids(self, preencoded_sample, sample_text_ids): + """Test that text_ids in sample are ignored (position IDs generated at runtime).""" + # Even if text_ids are provided in the sample, they should not be stored + # because position IDs are generated at runtime based on actual tensor shapes + preencoded_sample["text_ids.pth"] = sample_text_ids + result = cook_preencoded_diffusion(preencoded_sample) + + # Verify that text_ids attribute doesn't exist (not part of DiffusionSample) + assert not hasattr(result, "text_ids"), "text_ids should not be stored in DiffusionSample" + + def test_cook_missing_required_key_latents(self, preencoded_sample): + """Test error when latents are missing.""" + del preencoded_sample["latents.pth"] + + # Production now emits a distinct message for missing latents (vs + # missing prompt_embeds / pooled_prompt_embeds) since 'latents.pth' is + # interchangeable with 'mean.pth' / 'logvar.pth' for resample-mode. + with pytest.raises( + ValueError, + match=r"Sample must have 'latents\.pth' or 'mean\.pth'/'logvar\.pth'", + ): + cook_preencoded_diffusion(preencoded_sample) + + def test_cook_missing_required_key_prompt_embeds(self, preencoded_sample): + """Test error when prompt_embeds are missing.""" + del preencoded_sample["prompt_embeds.pth"] + + with pytest.raises(ValueError, match="Sample missing required keys"): + cook_preencoded_diffusion(preencoded_sample) + + def test_cook_missing_required_key_pooled(self, preencoded_sample): + """Test error when pooled_prompt_embeds are missing.""" + del preencoded_sample["pooled_prompt_embeds.pth"] + + with pytest.raises(ValueError, match="Sample missing required keys"): + cook_preencoded_diffusion(preencoded_sample) + + def test_cook_caption_as_str(self, preencoded_sample): + """Test caption handling when it's a string.""" + preencoded_sample["caption.txt"] = "String caption" + result = cook_preencoded_diffusion(preencoded_sample) + + assert result.caption == "String caption" + + def test_cook_caption_missing(self, preencoded_sample): + """Test caption handling when missing.""" + del preencoded_sample["caption.txt"] + result = cook_preencoded_diffusion(preencoded_sample) + + assert result.caption == "" + + def test_cook_preserves_sample_keys(self, preencoded_sample): + """Test that sample metadata keys are preserved.""" + result = cook_preencoded_diffusion(preencoded_sample) + + assert result.__key__ == "sample_001" + assert result.__subflavors__ == {"encoding": "preencoded"} + + +class TestCookRawImages: + """Tests for cook_raw_images cooker function.""" + + def test_cook_with_png(self, sample_image_bytes): + """Test cooking raw PNG image.""" + sample = { + "__key__": "test_png", + "__restore_key__": lambda: "test_png", + "__subflavors__": {"encoding": "raw"}, + "images": sample_image_bytes, + "txt": b"PNG caption", + } + result = cook_raw_images(sample) + + assert result["images"] == sample_image_bytes + assert result["txt"] == b"PNG caption" + + def test_cook_preserves_metadata(self, raw_sample): + """Test that metadata keys are preserved.""" + result = cook_raw_images(raw_sample) + + assert result["__key__"] == "sample_003" + assert result["__subflavors__"] == {"encoding": "raw"} + + +class TestEncodedDiffusionTaskEncoder: + """Tests for EncodedDiffusionTaskEncoder.""" + + def test_batch_with_all_fields(self, sample_latents, sample_prompt_embeds, sample_pooled_prompt_embeds): + """Test batching samples with all standard fields (position IDs generated at runtime).""" + encoder = EncodedDiffusionTaskEncoder(worker_config=None) + + samples = [ + DiffusionSample( + __key__="s1", + __restore_key__=lambda: "s1", + __subflavors__={"encoding": "preencoded"}, + latents=sample_latents.squeeze(0), + prompt_embeds=sample_prompt_embeds.squeeze(0), + pooled_prompt_embeds=sample_pooled_prompt_embeds.squeeze(0), + caption="Caption 1", + ), + DiffusionSample( + __key__="s2", + __restore_key__=lambda: "s2", + __subflavors__={"encoding": "preencoded"}, + latents=sample_latents.squeeze(0), + prompt_embeds=sample_prompt_embeds.squeeze(0), + pooled_prompt_embeds=sample_pooled_prompt_embeds.squeeze(0), + caption="Caption 2", + ), + ] + + batch = encoder.batch(samples) + + assert "latents" in batch + assert "prompt_embeds" in batch + assert "pooled_prompt_embeds" in batch + # text_ids not in batch - position IDs are generated at runtime + # Note: 'captions' may or may not be in batch depending on implementation + + assert batch["latents"].shape[0] == 2 + assert batch["prompt_embeds"].shape[0] == 2 + assert batch["pooled_prompt_embeds"].shape[0] == 2 + + +class TestRawDiffusionTaskEncoder: + """Tests for RawDiffusionTaskEncoder.""" + + def test_batch_with_raw_data(self, sample_image_bytes): + """Test batching raw image samples.""" + encoder = RawDiffusionTaskEncoder(worker_config=None) + + samples = [ + {"images": sample_image_bytes, "txt": b"Caption 1"}, + {"images": sample_image_bytes, "txt": b"Caption 2"}, + ] + + batch = encoder.batch(samples) + + assert "images" in batch + assert "txt" in batch + assert len(batch["images"]) == 2 + assert len(batch["txt"]) == 2 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/data/test_synthetic_datasets.py b/tests/unit_tests/backends/megatron/diffusion/data/test_synthetic_datasets.py new file mode 100644 index 000000000..62698df6b --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/test_synthetic_datasets.py @@ -0,0 +1,129 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Schema/shape/reproducibility contract tests for the synthetic Flux mock datasets. + +These guard the sample contract (keys, tensor shapes, seed reproducibility) that +the Flux forward step depends on for mock-data training runs. Relocated here from +tests/integration_tests/... because they are CPU-only unit tests of the synthetic +dataset providers and do not exercise an end-to-end TaskEncoder/dataloader path. +""" + +import pytest +import torch + +from primus.backends.megatron.data.synthetic import ( + MockFluxDataset, + MockFluxSchnellDataset, + PreGeneratedMockFluxSchnellDataset, +) +from primus.backends.megatron.data.synthetic_dataset_provider import ( + SyntheticDatasetProvider, +) +from tests.unit_tests.backends.megatron.diffusion.helpers import assert_tensor_shape +from tests.utils import PrimusUT + + +class TestMockFluxDataset(PrimusUT): + """Contract tests for MockFluxDataset (Flux dev).""" + + def test_mock_dataset_sample_structure(self): + dataset = MockFluxDataset(num_samples=10, seed=42) + sample = dataset[0] + + expected_keys = { + "latents", + "prompt_embeds", + "pooled_prompt_embeds", + "img_ids", + "txt_ids", + "caption", + } + assert set(sample.keys()) == expected_keys + + def test_mock_dataset_sample_shapes(self): + dataset = MockFluxDataset(num_samples=10, image_size=1024, seed=42) + sample = dataset[0] + + # Latents: (C, H, W) -- 1024/8 = 128 + assert_tensor_shape(sample["latents"], (16, 128, 128), "latents") + # T5 embeddings: (S, D) + assert_tensor_shape(sample["prompt_embeds"], (512, 4096), "t5_embeddings") + # CLIP pooled: (D,) + assert_tensor_shape(sample["pooled_prompt_embeds"], (768,), "clip_pooled") + # Image IDs: (N, 3) where N = (H/2) * (W/2) + expected_img_ids_len = (128 // 2) * (128 // 2) + assert_tensor_shape(sample["img_ids"], (expected_img_ids_len, 3), "img_ids") + # Text IDs: (S, 3) + assert_tensor_shape(sample["txt_ids"], (512, 3), "txt_ids") + + def test_mock_dataset_reproducibility(self): + dataset1 = MockFluxDataset(num_samples=10, seed=42) + dataset2 = MockFluxDataset(num_samples=10, seed=42) + + sample1 = dataset1[0] + sample2 = dataset2[0] + + assert torch.allclose(sample1["latents"], sample2["latents"]) + assert torch.allclose(sample1["prompt_embeds"], sample2["prompt_embeds"]) + assert torch.allclose(sample1["pooled_prompt_embeds"], sample2["pooled_prompt_embeds"]) + + def test_position_ids_format(self): + dataset = MockFluxDataset(num_samples=5, seed=42) + sample = dataset[0] + + img_ids = sample["img_ids"] + txt_ids = sample["txt_ids"] + + # 3D RoPE format for both image and text position IDs. + assert img_ids.shape[-1] == 3, "Image IDs should have 3 dimensions for RoPE" + assert txt_ids.shape[-1] == 3, "Text IDs should have 3 dimensions for RoPE" + # Text IDs are all zeros. + assert torch.allclose(txt_ids, torch.zeros_like(txt_ids)), "Text IDs should be all zeros" + + def test_latent_size_scales_with_image_size(self): + # Latent spatial size must be image_size/8 (VAE downsample contract). + for image_size, latent in ((512, 64), (2048, 256)): + dataset = MockFluxDataset(num_samples=5, image_size=image_size, seed=42) + sample = dataset[0] + assert_tensor_shape(sample["latents"], (16, latent, latent), "latents") + expected_img_ids_len = (latent // 2) * (latent // 2) + assert_tensor_shape(sample["img_ids"], (expected_img_ids_len, 3), "img_ids") + + +class TestMockFluxSchnellDataset(PrimusUT): + """Contract tests for MockFluxSchnellDataset (Flux Schnell / MLPerf v5.1).""" + + def test_schnell_sample_shapes(self): + # Schnell uses T5 seq_len=256 (vs 512 for dev). + dataset = MockFluxSchnellDataset(num_samples=10, image_size=256, seed=42) + sample = dataset[0] + + assert_tensor_shape(sample["latents"], (16, 32, 32), "latents") + assert_tensor_shape(sample["prompt_embeds"], (256, 4096), "prompt_embeds") + assert_tensor_shape(sample["pooled_prompt_embeds"], (768,), "pooled_prompt_embeds") + assert_tensor_shape(sample["txt_ids"], (256, 3), "txt_ids") + + def test_schnell_pregenerated_dataset(self): + dataset = PreGeneratedMockFluxSchnellDataset(num_samples=10, image_size=256, seed=42) + + sample0_a = dataset[0] + sample0_b = dataset[0] + + assert torch.allclose(sample0_a["latents"], sample0_b["latents"]) + assert torch.allclose(sample0_a["prompt_embeds"], sample0_b["prompt_embeds"]) + + +class TestSyntheticDatasetProviderLookup(PrimusUT): + """Tests for SyntheticDatasetProvider DEFAULT_DATASETS lookup.""" + + def test_schnell_class_resolves(self): + provider = SyntheticDatasetProvider(model_type="flux_schnell") + cls = provider._import_dataset_class() + assert cls is not None + assert cls.__name__ == "PreGeneratedMockFluxSchnellDataset" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron_bridge/test_megatron_bridge_adapter.py b/tests/unit_tests/backends/megatron_bridge/test_megatron_bridge_adapter.py index 43d8032a1..d5a928896 100644 --- a/tests/unit_tests/backends/megatron_bridge/test_megatron_bridge_adapter.py +++ b/tests/unit_tests/backends/megatron_bridge/test_megatron_bridge_adapter.py @@ -29,13 +29,26 @@ def adapter(): @pytest.fixture def sys_modules_guard(adapter): - """Drop any stub modules a test injected, keeping tests isolated.""" + """Isolate stub-eligible modules for the duration of a test. + + Drops any relevant modules BOTH before and after the test. The + before-drop matters because some entries (e.g. ``megatron.energon``) are + real, top-level packages that may already be installed and cached in + ``sys.modules`` from elsewhere in the session. If left in place, the stub + installer's ``if pkg_name in sys.modules: continue`` fast-path would skip + them, bypassing a test's forced-missing ``import_module`` mock and breaking + assertions that expect every package to be stubbed. + """ prefixes = ("modelopt",) + adapter._BRIDGE_OPTIONAL_PACKAGES - before = set(sys.modules) + + def _drop(): + for key in list(sys.modules): + if any(key == p or key.startswith(p + ".") for p in prefixes): + sys.modules.pop(key, None) + + _drop() yield - for key in set(sys.modules) - before: - if any(key == p or key.startswith(p + ".") for p in prefixes): - sys.modules.pop(key, None) + _drop() # --------------------------------------------------------------------------- @@ -118,29 +131,34 @@ def test_install_bridge_optional_stubs_is_idempotent(adapter, monkeypatch, sys_m # --------------------------------------------------------------------------- # _install_transformers_stub (transformers must be importable) # --------------------------------------------------------------------------- -def test_install_transformers_stub_adds_placeholders_for_missing(adapter): +def test_install_transformers_stub_adds_placeholders_for_missing(adapter, monkeypatch): transformers = pytest.importorskip("transformers") - name, _used_by = adapter._TRANSFORMERS_PLACEHOLDER_CLASSES[0] - had_attr = hasattr(transformers, name) - original = getattr(transformers, name, None) - if had_attr: - delattr(transformers, name) + # Append a synthetic always-missing class to the real list to drive the + # "missing class" branch deterministically. Deleting a real class won't + # work: transformers' lazy module re-caches it before the installer's + # hasattr check, and newer transformers ship every class anyway. + fake_name = "_PrimusFakeMissingTransformersClass_ForUnitTest" + monkeypatch.setattr( + adapter, + "_TRANSFORMERS_PLACEHOLDER_CLASSES", + adapter._TRANSFORMERS_PLACEHOLDER_CLASSES + ((fake_name, "models/fake/fake_bridge.py"),), + ) + assert not hasattr(transformers, fake_name) try: stubbed = adapter._install_transformers_stub() - assert name in stubbed - placeholder = getattr(transformers, name) + assert fake_name in stubbed + placeholder = getattr(transformers, fake_name) assert getattr(placeholder, "_primus_placeholder", False) is True with pytest.raises(RuntimeError, match="Primus placeholder"): placeholder() finally: - # restore: remove every placeholder we (or the call) injected + # Remove every placeholder the call injected: our synthetic entry plus + # any real classes that were genuinely missing in this environment. for cname, _ in adapter._TRANSFORMERS_PLACEHOLDER_CLASSES: obj = getattr(transformers, cname, None) if getattr(obj, "_primus_placeholder", False): delattr(transformers, cname) - if had_attr: - setattr(transformers, name, original) def test_install_transformers_stub_skips_existing_classes(adapter): From 01f4a4d882ff54e4a2304482c8d894f15a1aa03d Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Tue, 14 Jul 2026 16:52:58 +0300 Subject: [PATCH 025/127] feat(flux): diffusion training primitives (forward step, schedulers, loss) (#816) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/flux` — review after it. One of the parents of the trainers PR. ## What this changes The diffusion training primitives: the forward step, loss computation, noise utilities, timestep sampling, and the flow-matching schedulers. ## Why it's stacked here The forward step imports `flux.utils`, so it bases on `feat/flux/flux` rather than `feat/flux/model-common`. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/flux`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion/training -k "forward_step or loss or scheduler"`. Validated locally on an AMD GPU container: 20 passed. ## Files 13 (forward step, loss, noise/timestep sampling, flow-matching schedulers + tests). Co-authored-by: Flux Split Trial Co-authored-by: luiza-amd --- .../megatron/training/diffusion/__init__.py | 12 + .../training/diffusion/forward_step.py | 608 ++++++++++++++++++ .../training/diffusion/loss_computation.py | 94 +++ .../training/diffusion/noise_utils.py | 67 ++ .../training/diffusion/schedulers/__init__.py | 15 + .../training/diffusion/schedulers/base.py | 80 +++ .../diffusion/schedulers/flow_matching.py | 322 ++++++++++ .../training/diffusion/timestep_sampling.py | 329 ++++++++++ .../megatron/diffusion/training/__init__.py | 8 + .../megatron/diffusion/training/conftest.py | 36 ++ .../diffusion/training/test_forward_step.py | 219 +++++++ .../training/test_loss_computation.py | 56 ++ .../diffusion/training/test_scheduler.py | 257 ++++++++ 13 files changed, 2103 insertions(+) create mode 100644 primus/backends/megatron/training/diffusion/__init__.py create mode 100644 primus/backends/megatron/training/diffusion/forward_step.py create mode 100644 primus/backends/megatron/training/diffusion/loss_computation.py create mode 100644 primus/backends/megatron/training/diffusion/noise_utils.py create mode 100644 primus/backends/megatron/training/diffusion/schedulers/__init__.py create mode 100644 primus/backends/megatron/training/diffusion/schedulers/base.py create mode 100644 primus/backends/megatron/training/diffusion/schedulers/flow_matching.py create mode 100644 primus/backends/megatron/training/diffusion/timestep_sampling.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/__init__.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/conftest.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_forward_step.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_loss_computation.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_scheduler.py diff --git a/primus/backends/megatron/training/diffusion/__init__.py b/primus/backends/megatron/training/diffusion/__init__.py new file mode 100644 index 000000000..fb6c0a922 --- /dev/null +++ b/primus/backends/megatron/training/diffusion/__init__.py @@ -0,0 +1,12 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Diffusion training helpers. + +Import the concrete helpers from their submodules directly, e.g. +``from primus.backends.megatron.training.diffusion.forward_step import +flux_forward_step_func``. +""" diff --git a/primus/backends/megatron/training/diffusion/forward_step.py b/primus/backends/megatron/training/diffusion/forward_step.py new file mode 100644 index 000000000..1a112133d --- /dev/null +++ b/primus/backends/megatron/training/diffusion/forward_step.py @@ -0,0 +1,608 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. +# +# Adapted from NeMo's Flux training architecture. + +""" +Forward step functions for diffusion model training. + +This module provides forward step implementations for different +diffusion models, handling the training loop logic. + +Supported data formats (framework-standard keys): + - Pre-encoded: latents, prompt_embeds, pooled_prompt_embeds, text_ids (optional) + - Raw: images, txt - encodes on-the-fly + +Architecture follows functional composition for clarity and testability. +""" + +import logging +from typing import Optional, Tuple + +import torch + +from primus.backends.megatron.core.models.diffusion.flux.utils import ( + generate_image_position_ids, + generate_text_position_ids, + pack_latents, + unpack_latents, +) +from primus.backends.megatron.training.diffusion.noise_utils import ( + apply_flow_matching_noise, +) +from primus.backends.megatron.training.diffusion.timestep_sampling import ( + LogitNormalSampler, +) + +logger = logging.getLogger(__name__) + + +def prepare_flux_latents( + latents: torch.Tensor, + scheduler, + img_ids: Optional[torch.Tensor] = None, + guidance_scale: Optional[float] = None, + use_guidance_embed: bool = False, + timestep_sampler=None, # Optional: custom timestep sampler + pregenerated_noise: Optional[torch.Tensor] = None, + pregenerated_timesteps: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, ...]: + """ + Prepare latents for Flux training forward pass. + + This function: + 1. Generates img_ids if not provided (for robustness) + 2. Samples timesteps using configurable sampling strategy + 3. Adds noise to latents using flow matching + 4. Packs latents into sequence format + 5. Prepares guidance embeddings (if enabled) + + Args: + latents: Clean latent tensor (B, C, H, W) + scheduler: Flow matching scheduler + img_ids: Image position IDs (B, H*W/4, 3). If None, will be generated + guidance_scale: Guidance scale value (for CFG) + use_guidance_embed: Whether to use guidance embedding + timestep_sampler: Optional custom timestep sampler + (default: LogitNormalSampler) + pregenerated_noise: If provided, use this noise instead of sampling. + Used by deterministic comparison tests to ensure identical inputs. + pregenerated_timesteps: If provided, use these timesteps (in [0,1] range) + instead of sampling. Used by deterministic comparison tests. + + Returns: + Tuple containing: + - clean_latents: Original latents (for target computation) + - noise: Sampled noise + - packed_noisy_latents: Noisy latents in packed format + - img_ids: Image position IDs + - guidance_vec: Guidance vector (or None) + - timesteps: Sampled timesteps (in [0, num_train_timesteps] range) + - sigma_1d: Raw sigma in [0, 1] range, shape [B]. Pass directly to + the model as timesteps_norm to avoid bf16 round-trip precision loss. + + Reference: + NeMo's prepare_image_latent_like_reference() + """ + batch_size, num_channels, height, width = latents.shape + device = latents.device + dtype = latents.dtype + + # Use default sampler if not provided + if timestep_sampler is None: + timestep_sampler = LogitNormalSampler() + + # Generate img_ids if not provided (for robustness with variable sizes) + if img_ids is None: + img_ids = generate_image_position_ids(batch_size, height, width, device, dtype) + + if pregenerated_noise is not None: + noise = pregenerated_noise.to(device=device, dtype=dtype) + else: + noise = torch.randn_like(latents, device=device, dtype=dtype) + + if pregenerated_timesteps is not None: + sigma = pregenerated_timesteps.to(device=device, dtype=dtype) + timesteps = sigma * scheduler.num_train_timesteps + else: + timesteps, sigma = timestep_sampler.sample(batch_size, device, scheduler) + + # Convert sigma to correct dtype + sigma = sigma.to(dtype=dtype) + + # Save 1D sigma [B] before unsqueezing — used as timesteps_norm to avoid + # the bf16 round-trip (sigma * 1000 / 1000) that corrupts ~2.5% of values. + sigma_1d = sigma.clone() + + # Broadcast sigma to match latent dimensions + while len(sigma.shape) < latents.ndim: + sigma = sigma.unsqueeze(-1) + + # Flow matching forward process: x_t = (1 - sigma) * x_0 + sigma * noise + noisy_latents = apply_flow_matching_noise(latents, noise, sigma) + + # Pack latents into sequence format + packed_noisy_latents = pack_latents(noisy_latents) + + # Prepare guidance embedding (if enabled) + if use_guidance_embed and guidance_scale is not None: + guidance_vec = torch.full( + (batch_size,), + guidance_scale, + device=device, + dtype=dtype, + ) + else: + guidance_vec = None + + return ( + latents, + noise, + packed_noisy_latents, + img_ids, + guidance_vec, + timesteps, + sigma_1d, + ) + + +# NOTE: kept as an eager alias — torch.compile breaks CUDA RNG reproducibility +# (compiled torch.randn_like produces different values than eager mode with the +# same generator state). prepare_flux_latents only contains small ops (randn, +# rand, element-wise), so compile overhead exceeds any fusion benefit. Eager +# also matches NeMo's RNG sequence for cross-framework convergence comparison. +_eager_prepare_flux_latents = prepare_flux_latents + + +def flux_forward_step_func( + data_iterator, + model, + scheduler, + use_guidance_embed=False, + guidance_scale=None, + timestep_sampler=None, + cfg_dropout_prob=0.0, + empty_t5_encodings=None, + empty_clip_encodings=None, + vae_scale=None, + vae_shift=None, + vae_latent_mode="presampled", + per_step_rng_reseed=False, + step_count=0, +): + """ + Forward step function for Flux training with distributed data loading. + + Following Megatron's multimodal data loading pattern: + - When TP=1 (pure DP): each rank loads data directly, no broadcast needed + - When TP>1: only TP rank 0 has data_iterator, broadcast to other TP ranks + - Middle PP stages return early + + This function orchestrates the training step by: + 1. Handling distributed data loading (broadcast from rank 0) + 2. Loading or encoding images (via helper function) + 3. Loading or encoding text (via helper function) + 4. Optionally applying CFG dropout (replacing text embeddings with empty encodings) + 5. Preparing latents with noise and packing (via helper function) + 6. Running model forward pass + 7. Returning model output and loss computation inputs + + Supports two data formats (follows NeMo conventions): + 1. Pre-encoded: latents, prompt_embeds, pooled_prompt_embeds, text_ids (optional) + 2. Raw: images, txt - encodes on-the-fly + + Architecture follows NeMo conventions for Flux training. + + Args: + data_iterator: Iterator yielding training batches (None on non-dataloader ranks) + model: Flux model instance with encoders (config.params_dtype used for data broadcasting) + scheduler: Flow matching scheduler + use_guidance_embed: Whether model uses guidance embedding + guidance_scale: Guidance scale for CFG training + timestep_sampler: Optional custom timestep sampler (default: LogitNormalSampler) + cfg_dropout_prob: Probability of replacing text embeddings with empty encodings (default: 0.0) + empty_t5_encodings: Pre-generated fixed empty T5 encodings (seq_len, 1, context_dim) + empty_clip_encodings: Pre-generated fixed empty CLIP encodings (vec_in_dim,) + vae_scale: Optional VAE latent scale factor (default: None, MLPerf uses 0.3611) + vae_shift: Optional VAE latent shift factor (default: None, MLPerf uses 0.1159) + vae_latent_mode: How to obtain latents from the batch (default: "presampled"). + "presampled" — use stored latents directly. + "resample" — reconstruct latents from stored mean+logvar via + reparameterization at every step, then apply vae_scale/vae_shift. + per_step_rng_reseed: Reseed the default CUDA generator at each step + to isolate training random ops from model forward RNG consumption + (default: False). + step_count: Monotonically increasing counter identifying this forward + call. Used to derive a unique per-step RNG seed. Managed by the + caller (DiffusionPretrainTrainer) and reconstructed from checkpoint + state on resume as iteration * num_microbatches. + + Returns: + Tuple of (noise_pred, clean_latents, noise, loss_mask, metrics_dict, is_validation) + - noise_pred: Model output (predicted velocity) [B, C, H, W] + - clean_latents: Original clean latents [B, C, H, W] + - noise: Sampled noise [B, C, H, W] + - loss_mask: Optional mask for variable-length sequences [B] or None + - metrics_dict: Dictionary with training metrics + - is_validation: True when batch contains "timestep" key (MLPerf validation mode) + """ + # Reseed default CUDA generator per step to isolate training random ops + # (noise, timesteps, CFG dropout) from model forward RNG consumption. + # Required because TE fused attention advances the default generator even + # with dropout=0 when the DPA prologue patch is active. + if per_step_rng_reseed: + from megatron.core import parallel_state as _ps + from megatron.training import get_args as _get_args + + _seed = _get_args().seed + _per_rank_seed = _seed + 100 * _ps.get_data_parallel_rank() + _step_seed = (_per_rank_seed * 10000 + step_count) % (2**63) + torch.cuda.manual_seed(_step_seed) + + from megatron.core import tensor_parallel + from megatron.core.parallel_state import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + ) + + # Pipeline parallelism (pipeline_model_parallel_size > 1) is rejected at + # config construction for Flux (see BaseDiffusionConfig.__post_init__), so + # no middle-pipeline-stage handling is needed here. + # Derive compute dtype from bf16/fp16 flags rather than params_dtype. + # When use_fsdp2_fp32_param_optimizer is active, params_dtype is FP32 (for + # optimizer precision) but compute should still be BF16/FP16. + if model.config.bf16: + compute_dtype = torch.bfloat16 + elif model.config.fp16: + compute_dtype = torch.float16 + else: + compute_dtype = model.config.params_dtype + + tp_size = get_tensor_model_parallel_world_size() + + if tp_size == 1: + # Pure DP: every rank has its own data iterator (is_distributed=True). + # Skip broadcast_data overhead (~1.5ms of GPU idle from NCCL + # self-broadcasts, GPU->CPU transfers, and .item() sync stalls). + if data_iterator is None: + raise RuntimeError( + "data_iterator is None with TP=1; dataset provider must set is_distributed=True" + ) + batch = next(data_iterator) + if not isinstance(batch, dict): + raise TypeError( + f"[ForwardStep] Expected batch to be dict, got {type(batch)}. Batch value: {batch}" + ) + if vae_latent_mode == "resample": + required_keys = ["mean", "logvar", "prompt_embeds", "pooled_prompt_embeds"] + else: + required_keys = ["latents", "prompt_embeds", "pooled_prompt_embeds"] + missing_keys = [k for k in required_keys if k not in batch] + if missing_keys: + raise KeyError( + f"[ForwardStep] Batch missing required keys: {missing_keys}. " + f"Got keys: {list(batch.keys())}. " + f"vae_latent_mode={vae_latent_mode}" + ) + + # Cast to compute_dtype and move to CUDA in one pass + for key in batch: + if isinstance(batch[key], torch.Tensor): + if batch[key].is_floating_point(): + batch[key] = batch[key].to(dtype=compute_dtype, device="cuda", non_blocking=True) + elif not batch[key].is_cuda: + batch[key] = batch[key].cuda(non_blocking=True) + + prompt_embeds = batch["prompt_embeds"] + pooled_prompt_embeds = batch["pooled_prompt_embeds"] + + if vae_latent_mode == "resample": + mean = batch["mean"] + logvar = batch["logvar"] + else: + latents = batch["latents"] + + loss_mask = batch.get("loss_mask") + else: + # TP > 1: only rank 0 loads data, broadcast to other TP ranks + if data_iterator is not None and get_tensor_model_parallel_rank() == 0: + try: + batch = next(data_iterator) + if not isinstance(batch, dict): + raise TypeError( + f"[ForwardStep] Expected batch to be dict, got {type(batch)}. " + f"Batch value: {batch}" + ) + if vae_latent_mode == "resample": + required_keys = ["mean", "logvar", "prompt_embeds", "pooled_prompt_embeds"] + else: + required_keys = ["latents", "prompt_embeds", "pooled_prompt_embeds"] + missing_keys = [k for k in required_keys if k not in batch] + if missing_keys: + raise KeyError( + f"[ForwardStep] Batch missing required keys: {missing_keys}. " + f"Got keys: {list(batch.keys())}. " + f"vae_latent_mode={vae_latent_mode}" + ) + except StopIteration: + raise RuntimeError( + "[ForwardStep] Data iterator exhausted (should be infinite with " + "MegatronDataloaderWrapper). This indicates a bug in the dataloader wrapper." + ) + except Exception as e: + logger.error(f"[ForwardStep] Error getting batch: {type(e).__name__}: {e}") + import traceback + + logger.error(f"[ForwardStep] Traceback: {traceback.format_exc()}") + raise + else: + batch = None + + if batch is not None: + for key in batch: + if isinstance(batch[key], torch.Tensor) and batch[key].is_floating_point(): + batch[key] = batch[key].to(dtype=compute_dtype) + + try: + prompt_embeds = tensor_parallel.broadcast_data(["prompt_embeds"], batch, compute_dtype).get( + "prompt_embeds" + ) + pooled_prompt_embeds = tensor_parallel.broadcast_data( + ["pooled_prompt_embeds"], batch, compute_dtype + ).get("pooled_prompt_embeds") + + if vae_latent_mode == "resample": + mean = tensor_parallel.broadcast_data(["mean"], batch, compute_dtype).get("mean") + logvar = tensor_parallel.broadcast_data(["logvar"], batch, compute_dtype).get("logvar") + else: + latents = tensor_parallel.broadcast_data(["latents"], batch, compute_dtype).get("latents") + except Exception as e: + logger.error(f"[ForwardStep] Error broadcasting data: {type(e).__name__}: {e}") + logger.error( + f"[ForwardStep] batch type: {type(batch)}, " + f"batch keys: {list(batch.keys()) if isinstance(batch, dict) else 'N/A'}" + ) + if isinstance(batch, dict): + for key, value in batch.items(): + logger.error( + f"[ForwardStep] {key}: type={type(value)}, " + f"shape={value.shape if hasattr(value, 'shape') else 'N/A'}" + ) + import traceback + + logger.error(f"[ForwardStep] Traceback: {traceback.format_exc()}") + raise + + loss_mask = None + if batch is not None and "loss_mask" in batch: + loss_mask = tensor_parallel.broadcast_data(["loss_mask"], batch, compute_dtype).get("loss_mask") + if not loss_mask.is_cuda: + loss_mask = loss_mask.cuda(non_blocking=True) + + if not prompt_embeds.is_cuda: + prompt_embeds = prompt_embeds.cuda(non_blocking=True) + if not pooled_prompt_embeds.is_cuda: + pooled_prompt_embeds = pooled_prompt_embeds.cuda(non_blocking=True) + + # Obtain latents based on vae_latent_mode + if vae_latent_mode == "resample": + # Resample mode: reconstruct latents from posterior parameters each step + if not mean.is_cuda: + mean = mean.cuda(non_blocking=True) + if not logvar.is_cuda: + logvar = logvar.cuda(non_blocking=True) + std = torch.exp(0.5 * logvar) + vae_eps = torch.randn_like(mean) + + latents = mean + std * vae_eps + # Scale/shift is always applied after resampling (raw posterior -> normalized latents) + latents = vae_scale * (latents - vae_shift) + else: + # Presampled mode: use stored latents directly + if not latents.is_cuda: + latents = latents.cuda(non_blocking=True) + + # Validation detection. + # + # MLPerf v5.1 Flux1 validation spec (flux1/nemo/README.md §6 "Evaluation"): + # - Per-sample fixed timestep t ∈ {0/8, 1/8, ..., 7/8} + # - Equal sample count per timestep (29 696 / 8 = 3 712) + # - val_loss = mean over per-timestep means (equivalent to flat mean given + # equal counts). + # + # NeMo's official to_webdataset preserves a `timestep` integer per sample + # from the MLCommons Arrow source. Our `primus-cli data diffusion-ingest` + # path (pipelines/ingest.py:33 `ARROW_COLUMNS`) ingests only the 4 tensor + # columns and writes `{"key": ...}` to the json sidecar — so our val shards + # are MISSING the timestep field, which used to make this branch fall + # through to the training path with uniform-random timesteps via the + # `timestep_sampler`. That produced a *different* val_loss estimator than + # the spec's: E_t~U[0,1][MSE] (Monte Carlo over [0,1]) vs the spec's + # left-Riemann sum over t∈{0/8..7/8}. The two estimators are not + # comparable, so a uniform-random val path can make val_loss converge + # spuriously fast relative to the reference convergence point. + # + # Fix: when batch is in eval mode (model.training=False, set by + # the evaluation harness via `model_module.eval()`) and lacks a `timestep` + # field, inject equidistant timesteps deterministically by within-batch + # index. With MBS=64, each micro-batch covers each t∈{0..7} exactly 8 + # times. Across 58 micro-batches × 8 DP ranks = 464 micro-batches → exactly + # 3 712 samples per timestep, matching the MLPerf v5.1 spec count. + # + # CFG dropout during val: SUPPRESSED. + # + # Reference-implementation tally for "apply CFG dropout during validation": + # NeMo MLPerf reference (custom_flux.py): ON + # AMD's MLPerf submission: OFF + # TorchTitan flux training script: OFF + # + # CFG-off during validation is MLPerf-compliant under the v6.0 rules even + # though NeMo (which generated the reference convergence point) has it on. + # Empirically, applying CFG-during-val structurally inflates val_loss by + # ~0.015-0.030 (the 10% unconditional samples pay a ~0.15-0.30 MSE + # penalty), which is enough to materially shift the convergence-crossing + # step, so we keep it off to match the submission configuration. + is_validation = False + if batch is not None and "timestep" in batch: + is_validation = True + val_timesteps = batch["timestep"].float() / 8.0 + batch["timesteps"] = val_timesteps + elif batch is not None and not model.training: + is_validation = True + batch_size_val = pooled_prompt_embeds.shape[0] + val_idx = torch.arange(batch_size_val, device="cuda") % 8 + batch["timestep"] = val_idx + val_timesteps = val_idx.to(dtype=compute_dtype) / 8.0 + batch["timesteps"] = val_timesteps + + # Matches NeMo's forward_step which wraps prepare_image_latent_like_reference + # in torch.no_grad() — no gradients needed for position IDs, noise sampling, + # timestep sampling, or latent packing. + with torch.no_grad(): + # Generate img_ids based on latent spatial dimensions + # NOTE: When RoPE fusion is enabled, we use batch_size=1 to satisfy Transformer Engine's + # fused kernel constraints (freqs must have shape [S, 1, 1, D]). PyTorch broadcasting + # applies the same position grid across all batch samples. This requires all images in + # the batch to have the same resolution (same height/width). + rope_fusion_batch_size = 1 if model.config.apply_rope_fusion else latents.shape[0] + img_ids = generate_image_position_ids( + batch_size=rope_fusion_batch_size, + height=latents.shape[2], + width=latents.shape[3], + device=latents.device, + dtype=latents.dtype, + ) + + # Generate text_ids (Flux convention: zeros for text position IDs) + # NOTE: When RoPE fusion is enabled, use batch_size=1 for consistency with img_ids + # (broadcasting will handle the actual batch dimension). This matches NVIDIA's MLPerf + # implementation strategy: both txt_ids and img_ids have shape [1, seq_len, 3] with + # RoPE fusion, allowing proper concatenation before the fused RoPE kernel. + text_ids = generate_text_position_ids( + batch_size=rope_fusion_batch_size, + seq_len=prompt_embeds.shape[1], + device=latents.device, + dtype=latents.dtype, + ) + + # Extract pre-generated noise/timesteps from batch (deterministic tests) + batch_noise = None + batch_timesteps = None + if batch is not None: + if tp_size == 1: + batch_noise = batch.get("noise") + batch_timesteps = batch.get("timesteps") + else: + if "noise" in batch: + batch_noise = tensor_parallel.broadcast_data(["noise"], batch, compute_dtype).get("noise") + if not batch_noise.is_cuda: + batch_noise = batch_noise.cuda(non_blocking=True) + if "timesteps" in batch: + batch_timesteps = tensor_parallel.broadcast_data(["timesteps"], batch, compute_dtype).get( + "timesteps" + ) + if not batch_timesteps.is_cuda: + batch_timesteps = batch_timesteps.cuda(non_blocking=True) + + # Prepare latents (noise, packing, scheduling). + # Eager wrapper — see _eager_prepare_flux_latents NOTE for why compile + # is intentionally disabled (CUDA RNG reproducibility). + ( + clean_latents, + noise, + packed_noisy_latents, + img_ids, + guidance_vec, + timesteps, + sigma_1d, + ) = _eager_prepare_flux_latents( + latents=latents, + scheduler=scheduler, + img_ids=img_ids, + guidance_scale=guidance_scale, + use_guidance_embed=use_guidance_embed, + timestep_sampler=timestep_sampler, + pregenerated_noise=batch_noise, + pregenerated_timesteps=batch_timesteps, + ) + + # CFG dropout: randomly replace text embeddings with fixed empty encodings. + # Placed after prepare_flux_latents so the RNG consumption order matches NeMo: + # VAE resample → noise → timesteps → CFG dropout. + # Applied during training only — validation uses fixed per-sample timesteps. + if ( + not is_validation + and cfg_dropout_prob > 0.0 + and empty_t5_encodings is not None + and empty_clip_encodings is not None + ): + batch_size_cfg = pooled_prompt_embeds.shape[0] + dropout_mask = torch.rand(batch_size_cfg, device="cuda") < cfg_dropout_prob + + empty_t5 = empty_t5_encodings.to(device="cuda", dtype=prompt_embeds.dtype, non_blocking=True) + empty_t5 = empty_t5.squeeze(1).unsqueeze(0) + + if empty_t5.shape[1] != prompt_embeds.shape[1]: + raise ValueError( + f"Empty T5 encoding seq_len ({empty_t5.shape[1]}) does not match " + f"data T5 seq_len ({prompt_embeds.shape[1]}). " + f"Regenerate empty encodings with matching t5_max_length." + ) + + t5_mask = dropout_mask.view(-1, 1, 1).expand_as(prompt_embeds) + prompt_embeds = torch.where(t5_mask, empty_t5.expand_as(prompt_embeds), prompt_embeds) + + empty_clip = empty_clip_encodings.to( + device="cuda", dtype=pooled_prompt_embeds.dtype, non_blocking=True + ) + clip_mask = dropout_mask.view(-1, 1).expand_as(pooled_prompt_embeds) + pooled_prompt_embeds = torch.where( + clip_mask, empty_clip.expand_as(pooled_prompt_embeds), pooled_prompt_embeds + ) + + # Transpose for Megatron format (sequence-first) + packed_noisy_latents = packed_noisy_latents.transpose(0, 1) + prompt_embeds = prompt_embeds.transpose(0, 1) + + # Use raw sigma directly instead of timesteps/1000 to avoid bf16 round-trip + timesteps_norm = sigma_1d.to(dtype=packed_noisy_latents.dtype) + + with torch.amp.autocast("cuda", enabled=True, dtype=compute_dtype): + noise_pred = model( + img=packed_noisy_latents, + txt=prompt_embeds, + y=pooled_prompt_embeds, + timesteps=timesteps_norm, + img_ids=img_ids, + txt_ids=text_ids, + guidance=guidance_vec, + ) + + # Unpack latents from sequence format + noise_pred = noise_pred.transpose(0, 1) # (S, B, C*4) -> (B, S, C*4) + noise_pred = unpack_latents( + noise_pred, + height=clean_latents.shape[2], + width=clean_latents.shape[3], + ) # -> (B, C, H, W) + + # Create metrics dict for logging + metrics = { + "batch_size": latents.shape[0], + "image_height": clean_latents.shape[2] * 8, # VAE 8x downsampling + "image_width": clean_latents.shape[3] * 8, + "latent_channels": clean_latents.shape[1], + "avg_timestep": timesteps.float().mean(), + "text_seq_len": prompt_embeds.shape[0], # After transpose to (S, B, C) + "img_seq_len": packed_noisy_latents.shape[0], # After transpose to (S, B, C) + } + + # Return model output and loss computation inputs (matching Megatron's pattern) + return noise_pred, clean_latents, noise, loss_mask, metrics, is_validation + + +__all__ = [ + "prepare_flux_latents", + "flux_forward_step_func", +] diff --git a/primus/backends/megatron/training/diffusion/loss_computation.py b/primus/backends/megatron/training/diffusion/loss_computation.py new file mode 100644 index 000000000..5702e6f5f --- /dev/null +++ b/primus/backends/megatron/training/diffusion/loss_computation.py @@ -0,0 +1,94 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Loss computation for diffusion models. + +This module provides reusable loss computation logic for different +diffusion training objectives. All functions are pure (no state), +making them easy to test and reusable across different model architectures. + +Supported loss types: + - Flow Matching: Used by Flux, SD3, and video models + - Epsilon Prediction: Used by DDPM and older models + - V-Prediction: Alternative parameterization + +All loss functions accept tensors of any shape and reduce to scalar. +""" + +from typing import Optional + +import torch.nn.functional as F +from torch import Tensor + + +def compute_flow_matching_loss( + prediction: Tensor, + clean_latents: Tensor, + noise: Tensor, + loss_mask: Optional[Tensor] = None, +) -> Tensor: + """ + Compute flow matching loss with optional masking. + + Flow matching models predict the velocity field that transforms + noise to clean latents. The training objective is: + + Formula: target = noise - clean_latents + loss = MSE(prediction, target) + + This loss is model-agnostic and works with any tensor shape, + making it reusable across 2D (images) and 3D (video) models. + + Args: + prediction: Model output (predicted velocity) [any shape] + clean_latents: Original clean latents [same shape as prediction] + noise: Sampled noise [same shape as prediction] + loss_mask: Optional mask for variable-length sequences [batch_size] + or broadcast-compatible shape. Default None (no masking). + + Returns: + Scalar loss value (mean squared error, optionally masked). + + Reference: + Flow Matching for Generative Modeling + https://arxiv.org/abs/2210.02747 + + Example (no masking): + >>> pred = torch.randn(2, 16, 64, 64) # Batch of 2D latents + >>> clean = torch.randn(2, 16, 64, 64) + >>> noise = torch.randn(2, 16, 64, 64) + >>> loss = compute_flow_matching_loss(pred, clean, noise) + >>> loss.backward() + + Example (with packing/masking): + >>> pred = torch.randn(2, 16, 64, 64) + >>> clean = torch.randn(2, 16, 64, 64) + >>> noise = torch.randn(2, 16, 64, 64) + >>> mask = torch.tensor([1.0, 0.0]) # Second sample is padding + >>> loss = compute_flow_matching_loss(pred, clean, noise, mask) + """ + target = noise - clean_latents + + if loss_mask is None: + # Simple case: direct mean reduction + loss = F.mse_loss(prediction.float(), target.float(), reduction="mean") + else: + # With masking for variable-length sequences (packing support) + loss_per_element = F.mse_loss(prediction.float(), target.float(), reduction="none") + + # Broadcast mask to match loss_per_element shape if needed + if loss_mask.dim() == 1 and loss_per_element.dim() > 1: + mask_shape = [loss_mask.shape[0]] + [1] * (loss_per_element.dim() - 1) + loss_mask = loss_mask.view(*mask_shape) + + # Apply mask and compute mean over valid elements + loss = (loss_per_element * loss_mask).sum() / loss_mask.sum() + + return loss + + +__all__ = [ + "compute_flow_matching_loss", +] diff --git a/primus/backends/megatron/training/diffusion/noise_utils.py b/primus/backends/megatron/training/diffusion/noise_utils.py new file mode 100644 index 000000000..70bda1e2a --- /dev/null +++ b/primus/backends/megatron/training/diffusion/noise_utils.py @@ -0,0 +1,67 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Noise application utilities for diffusion training. + +This module provides pure functions for applying noise according to +different diffusion forward processes. All functions are stateless +and work with tensors of any shape. + +Supported processes: + - Flow Matching: Linear interpolation between clean and noise + - DDPM: Variance-preserving noise schedule + +These utilities encapsulate the mathematical formulas for the forward +diffusion process, making the code self-documenting and reusable +across different model architectures (2D images, 3D video, etc.). +""" + +from torch import Tensor + + +def apply_flow_matching_noise( + clean_latents: Tensor, + noise: Tensor, + sigma: Tensor, +) -> Tensor: + """ + Apply noise using flow matching forward process. + + Flow matching uses a simple linear interpolation between clean + latents and noise, controlled by the sigma parameter: + + Formula: noisy = (1 - sigma) * clean + sigma * noise + + This formulation ensures: + - When sigma=0: noisy = clean (no noise) + - When sigma=1: noisy = noise (pure noise) + - Linear interpolation in between + + Args: + clean_latents: Clean latents [any shape] + noise: Sampled noise [same shape as clean_latents] + sigma: Noise schedule values [broadcast compatible] + Should be in range [0, 1] + + Returns: + Noisy latents [same shape as clean_latents] + + Reference: + Flow Matching for Generative Modeling + https://arxiv.org/abs/2210.02747 + + Example: + >>> clean = torch.randn(2, 16, 64, 64) + >>> noise = torch.randn(2, 16, 64, 64) + >>> sigma = torch.tensor([0.3, 0.7]).reshape(2, 1, 1, 1) + >>> noisy = apply_flow_matching_noise(clean, noise, sigma) + >>> # sigma[0]=0.3 means 30% noise, 70% clean for first sample + """ + return (1.0 - sigma) * clean_latents + sigma * noise + + +__all__ = [ + "apply_flow_matching_noise", +] diff --git a/primus/backends/megatron/training/diffusion/schedulers/__init__.py b/primus/backends/megatron/training/diffusion/schedulers/__init__.py new file mode 100644 index 000000000..2b917c927 --- /dev/null +++ b/primus/backends/megatron/training/diffusion/schedulers/__init__.py @@ -0,0 +1,15 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +from primus.backends.megatron.training.diffusion.schedulers.base import BaseScheduler +from primus.backends.megatron.training.diffusion.schedulers.flow_matching import ( + FlowMatchEulerDiscreteScheduler, +) + +__all__ = [ + "BaseScheduler", + "FlowMatchEulerDiscreteScheduler", +] diff --git a/primus/backends/megatron/training/diffusion/schedulers/base.py b/primus/backends/megatron/training/diffusion/schedulers/base.py new file mode 100644 index 000000000..85366d255 --- /dev/null +++ b/primus/backends/megatron/training/diffusion/schedulers/base.py @@ -0,0 +1,80 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Base scheduler interface for diffusion models. +""" + +from abc import ABC, abstractmethod +from typing import Optional, Tuple, Union + +import torch + + +class BaseScheduler(ABC): + """ + Abstract base class for diffusion schedulers. + + Schedulers handle: + 1. Noise schedule generation + 2. Forward process (adding noise to samples) + 3. Reverse process (denoising step) + """ + + @abstractmethod + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + noise: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + """ + Add noise to sample (forward process). + + Args: + sample: Clean sample + timestep: Current timestep + noise: Optional noise tensor (generated if None) + + Returns: + Noisy sample + """ + raise NotImplementedError("Subclasses must implement scale_noise()") + + @abstractmethod + def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None, **kwargs): + """ + Set discrete timesteps for inference. + + Args: + num_inference_steps: Number of diffusion steps + device: Device to place timesteps on + **kwargs: Additional scheduler-specific arguments + """ + raise NotImplementedError("Subclasses must implement set_timesteps()") + + @abstractmethod + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + **kwargs, + ) -> Tuple[torch.FloatTensor, ...]: + """ + Perform one denoising step (reverse process). + + Args: + model_output: Model prediction (noise or velocity) + timestep: Current timestep + sample: Current noisy sample + **kwargs: Additional scheduler-specific arguments + + Returns: + Tuple containing denoised sample (and optionally other values) + """ + raise NotImplementedError("Subclasses must implement step()") + + +__all__ = ["BaseScheduler"] diff --git a/primus/backends/megatron/training/diffusion/schedulers/flow_matching.py b/primus/backends/megatron/training/diffusion/schedulers/flow_matching.py new file mode 100644 index 000000000..4f0fb85cf --- /dev/null +++ b/primus/backends/megatron/training/diffusion/schedulers/flow_matching.py @@ -0,0 +1,322 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# Licensed under the Apache License, Version 2.0. +# +# Adapted from NeMo's flow matching scheduler + +""" +Flow Matching Euler Discrete Scheduler for Flux. + +This scheduler implements the Euler discrete sampling method for flow matching +models like Flux. It supports both training and inference modes with optional +dynamic timestep shifting for variable resolution. +""" + +import math +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch + +from .base import BaseScheduler + + +class FlowMatchEulerDiscreteScheduler(BaseScheduler): + """ + Euler scheduler for flow matching diffusion models. + + This scheduler is used for Flux and implements: + - Flow matching training objective + - Euler discrete sampling for inference + - Dynamic timestep shifting for variable resolution + + Args: + num_train_timesteps: Number of diffusion steps for training (default: 1000) + shift: Base shift value for timestep schedule (default: 1.0) + use_dynamic_shifting: Whether to use dynamic shifting based on image resolution + base_shift: Base shift for dynamic shifting (default: 0.5) + max_shift: Maximum shift for dynamic shifting (default: 1.15) + base_image_seq_len: Base image sequence length for dynamic shifting (default: 256) + max_image_seq_len: Maximum image sequence length for dynamic shifting (default: 4096) + + Reference: + - Flux paper: https://blackforestlabs.ai/flux-1-tools/ + - Adapted from NeMo's flow matching scheduler + """ + + _compatibles = [] + order = 1 + + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + use_dynamic_shifting: bool = False, + base_shift: Optional[float] = 0.5, + max_shift: Optional[float] = 1.15, + base_image_seq_len: Optional[int] = 256, + max_image_seq_len: Optional[int] = 4096, + ): + """Initialize Flow Matching Euler Discrete Scheduler.""" + # Generate initial timesteps + timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() + timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) + + # Convert to sigmas (normalized timesteps in [0, 1]) + sigmas = timesteps / num_train_timesteps + + # Apply static shifting if not using dynamic shifting + if not use_dynamic_shifting: + # Shift formula: shift * sigma / (1 + (shift - 1) * sigma) + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.timesteps = sigmas * num_train_timesteps + + self._step_index = None + self._begin_index = None + + # Move sigmas to CPU to avoid too much CPU/GPU communication + self.sigmas = sigmas.to("cpu") + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + # Store parameters + self.base_shift = base_shift + self.max_shift = max_shift + self.base_image_seq_len = base_image_seq_len + self.max_image_seq_len = max_image_seq_len + self.use_dynamic_shifting = use_dynamic_shifting + self.num_train_timesteps = num_train_timesteps + self.shift = shift + + @property + def step_index(self): + """The index counter for current timestep.""" + return self._step_index + + @property + def begin_index(self): + """The index for the first timestep.""" + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + """ + Set the begin index for the scheduler. + + Args: + begin_index: The begin index for the scheduler + """ + self._begin_index = begin_index + + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + noise: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + """ + Add noise to sample using flow matching forward process. + + Flow matching interpolation: x_t = sigma * noise + (1 - sigma) * sample + + Args: + sample: Clean sample tensor + timestep: Current timestep(s) + noise: Optional noise tensor (generated if None) + + Returns: + Noisy sample + """ + # Generate noise if not provided + if noise is None: + noise = torch.randn_like(sample) + + # Ensure sigmas and timesteps have same device and dtype as sample + sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) + + # Handle MPS device (doesn't support float64) + if sample.device.type == "mps" and torch.is_floating_point(timestep): + schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) + timestep = timestep.to(sample.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(sample.device) + timestep = timestep.to(sample.device) + + # Get step indices for given timesteps + # begin_index is None during training + if self.begin_index is None: + step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep] + elif self.step_index is not None: + # add_noise called after first denoising step (for inpainting) + step_indices = [self.step_index] * timestep.shape[0] + else: + # add noise called before first denoising step (img2img) + step_indices = [self.begin_index] * timestep.shape[0] + + # Get sigma values for these indices + sigma = sigmas[step_indices].flatten() + + # Broadcast sigma to match sample shape + while len(sigma.shape) < len(sample.shape): + sigma = sigma.unsqueeze(-1) + + # Flow matching interpolation + noisy_sample = sigma * noise + (1.0 - sigma) * sample + + return noisy_sample + + def _sigma_to_t(self, sigma: float) -> float: + """Convert sigma to timestep.""" + return sigma * self.num_train_timesteps + + def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor: + """ + Apply dynamic time shifting. + + Formula: exp(mu) / (exp(mu) + (1/t - 1)^sigma) + + Args: + mu: Shift parameter (logarithmic) + sigma: Exponent parameter (fixed at 1.0 for Flux) + t: Timesteps to shift + + Returns: + Shifted timesteps + """ + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + def set_timesteps( + self, + num_inference_steps: Optional[int] = None, + device: Union[str, torch.device] = None, + sigmas: Optional[List[float]] = None, + mu: Optional[float] = None, + ): + """ + Set discrete timesteps for inference. + + Args: + num_inference_steps: Number of diffusion steps for inference + device: Device to place timesteps on + sigmas: Optional precomputed sigma values + mu: Shift parameter for dynamic shifting + """ + # Validate dynamic shifting requirements + if self.use_dynamic_shifting and mu is None: + raise ValueError("Must pass a value for `mu` when `use_dynamic_shifting` is True") + + # Generate sigmas if not provided + if sigmas is None: + self.num_inference_steps = num_inference_steps + timesteps = np.linspace( + self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps + ) + sigmas = timesteps / self.num_train_timesteps + + # Apply shifting (dynamic or static) + if self.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) + + # Convert to tensors + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) + timesteps = sigmas * self.num_train_timesteps + + self.timesteps = timesteps.to(device=device) + # Append zero sigma at the end + self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + + self._step_index = None + self._begin_index = None + + def index_for_timestep( + self, timestep: Union[float, torch.FloatTensor], schedule_timesteps: Optional[torch.Tensor] = None + ) -> int: + """ + Get the index for a given timestep. + + Args: + timestep: Timestep value + schedule_timesteps: Optional schedule to search in (uses self.timesteps if None) + + Returns: + Index of the timestep + """ + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index taken for the very first step is always the second index + # (or the last index if there is only 1). This ensures we don't accidentally + # skip a sigma when starting in the middle of the schedule. + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep: Union[float, torch.FloatTensor]): + """Initialize step index from timestep.""" + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + generator: Optional[torch.Generator] = None, + **kwargs, + ) -> Tuple[torch.FloatTensor]: + """ + Perform one Euler denoising step. + + Args: + model_output: Model prediction (velocity field for flow matching) + timestep: Current timestep + sample: Current noisy sample + generator: Optional random number generator (unused) + **kwargs: Additional arguments (unused) + + Returns: + Tuple containing the denoised sample + """ + # Validate timestep type + if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): + raise ValueError( + "Passing integer indices as timesteps is not supported. " + "Pass one of scheduler.timesteps as a timestep." + ) + + # Initialize step index if needed + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues + sample = sample.to(torch.float32) + + # Get current and next sigma + sigma = self.sigmas[self.step_index] + sigma_next = self.sigmas[self.step_index + 1] + + # Euler step: x_{t-1} = x_t + (sigma_next - sigma) * model_output + prev_sample = sample + (sigma_next - sigma) * model_output + + # Cast back to model dtype + prev_sample = prev_sample.to(model_output.dtype) + + # Increment step index + self._step_index += 1 + + return (prev_sample,) + + def __len__(self): + """Return number of training timesteps.""" + return self.num_train_timesteps + + +__all__ = ["FlowMatchEulerDiscreteScheduler"] diff --git a/primus/backends/megatron/training/diffusion/timestep_sampling.py b/primus/backends/megatron/training/diffusion/timestep_sampling.py new file mode 100644 index 000000000..4a95a212e --- /dev/null +++ b/primus/backends/megatron/training/diffusion/timestep_sampling.py @@ -0,0 +1,329 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Timestep sampling strategies for diffusion training. + +Training timestep sampling is a hyperparameter optimization technique +that is separate from inference denoising algorithms. Different sampling +strategies can improve training convergence and model quality. + +Key Insight: + - Training: Use logit-normal sampling (hyperparameter optimization) + - Inference: Use linear schedule (denoising algorithm requirement) + +These are fundamentally different concerns and should be separated. + +Supported strategies: + - LogitNormal: Better convergence (from SD3 paper) + - Uniform: Baseline approach for comparison + - Mode: Alternative from SD3 paper + +Reference: + Stable Diffusion 3 paper: https://arxiv.org/abs/2403.03206v1 + Section 3.1: Timestep sampling +""" + +import math +from abc import ABC, abstractmethod +from typing import Tuple + +import torch +from torch import Tensor + + +class TimestepSampler(ABC): + """ + Abstract base class for training timestep sampling strategies. + + Different sampling strategies can improve training convergence by + emphasizing certain timesteps where the model needs to learn more. + + All samplers must implement the `sample()` method that returns + both timesteps and their corresponding sigma values. + """ + + @abstractmethod + def sample( + self, + batch_size: int, + device: torch.device, + scheduler, + ) -> Tuple[Tensor, Tensor]: + """ + Sample timesteps and compute sigmas for training. + + Args: + batch_size: Number of samples in the batch + device: Target device for the tensors + scheduler: Scheduler instance (for accessing timesteps/sigmas) + + Returns: + Tuple of (timesteps, sigmas): + - timesteps: Sampled timestep values [batch_size] + - sigmas: Corresponding noise schedule values [batch_size] + + Example: + >>> sampler = LogitNormalSampler() + >>> timesteps, sigmas = sampler.sample( + ... batch_size=32, + ... device='cuda', + ... scheduler=flow_scheduler + ... ) + """ + + +class LogitNormalSampler(TimestepSampler): + """ + Logit-normal timestep sampling (current approach, from SD3 paper). + + This sampling strategy uses a logit-normal distribution to sample + timesteps. With default parameters (mean=0, std=1), the distribution + emphasizes mid-range timesteps (around t≈500). The mean and std + parameters control the distribution shape. + + The distribution is created by: + 1. Sample u from Normal(mean, std) + 2. Apply sigmoid: u = sigmoid(u) + 3. Map to timestep indices: indices = u * num_train_timesteps + + This approach has been shown to improve training convergence + compared to uniform sampling. + + Args: + mean: Mean of the normal distribution (default: 0.0) + std: Standard deviation of the normal distribution (default: 1.0) + + Reference: + Stable Diffusion 3 paper: https://arxiv.org/abs/2403.03206v1 + Section 3.1: "We use rf/lognorm(0.00,1.00)" + + Example: + >>> sampler = LogitNormalSampler(mean=0.0, std=1.0) + >>> timesteps, sigmas = sampler.sample(32, 'cuda', scheduler) + >>> # More samples near the middle (t~500) than the extremes + """ + + def __init__(self, mean: float = 0.0, std: float = 1.0): + self.mean = mean + self.std = std + + def sample( + self, + batch_size: int, + device: torch.device, + scheduler, + ) -> Tuple[Tensor, Tensor]: + """Sample using logit-normal distribution.""" + # Sample from Normal(mean, std) + u = torch.normal( + mean=self.mean, + std=self.std, + size=(batch_size,), + device="cpu", # CPU to match scheduler.timesteps device + ) + + # Apply sigmoid to get logit-normal distribution + u = torch.sigmoid(u) + + # Map to timestep indices + indices = (u * scheduler.num_train_timesteps).long() + indices = torch.clamp(indices, 0, scheduler.num_train_timesteps - 1) + + # Get timesteps + timesteps = scheduler.timesteps[indices].to(device=device) + + # Get corresponding sigmas + sigmas = scheduler.sigmas.to(device=device) + schedule_timesteps = scheduler.timesteps.to(device=device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + sigma = sigmas[step_indices].flatten() + + return timesteps, sigma + + +class UniformSampler(TimestepSampler): + """ + Uniform timestep sampling (baseline approach). + + Samples timesteps uniformly from [0, num_train_timesteps), giving + equal probability to all timesteps. This is simpler but may converge + slower than logit-normal sampling. + + Useful as a baseline for comparison with more sophisticated + sampling strategies. + + Example: + >>> sampler = UniformSampler() + >>> timesteps, sigmas = sampler.sample(32, 'cuda', scheduler) + >>> # Equal probability for all timesteps + """ + + def sample( + self, + batch_size: int, + device: torch.device, + scheduler, + ) -> Tuple[Tensor, Tensor]: + """Sample uniformly from all timesteps.""" + indices = torch.randint(0, scheduler.num_train_timesteps, (batch_size,), device="cpu") + + timesteps = scheduler.timesteps[indices].to(device=device) + + sigmas = scheduler.sigmas.to(device=device) + schedule_timesteps = scheduler.timesteps.to(device=device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + sigma = sigmas[step_indices].flatten() + + return timesteps, sigma + + +class DirectUniformSampler(TimestepSampler): + """ + Direct uniform timestep sampling without scheduler indirection (NVIDIA MLPerf-compatible). + + Samples timesteps continuously from [0, num_train_timesteps) using torch.rand, + bypassing the scheduler's discrete timestep/sigma lookup tables. This matches + NVIDIA's MLPerf Flux implementation exactly: + + timesteps = torch.rand((batch_size,)) * num_train_timesteps + sigma = timesteps / num_train_timesteps (i.e., sigma == normalized timestep) + + Advantages over UniformSampler: + - Continuous distribution (no 1000-bin discretization) + - No Python loop for sigma lookup (avoids per-element .nonzero() calls) + - Produces tensors directly on target device + + Note: Assumes shift=1.0 in the scheduler (no timestep shifting). With shift=1.0, + sigma == normalized_timestep, which matches the scheduler's linear mapping. + + Example: + >>> sampler = DirectUniformSampler() + >>> timesteps, sigmas = sampler.sample(32, 'cuda', scheduler) + """ + + def sample( + self, + batch_size: int, + device: torch.device, + scheduler, + ) -> Tuple[Tensor, Tensor]: + """Sample uniformly from [0, num_train_timesteps) without scheduler lookup.""" + sigma = torch.rand((batch_size,), device=device) + timesteps = sigma * scheduler.num_train_timesteps + return timesteps, sigma + + +class ModeSampler(TimestepSampler): + """ + Mode-based sampling from SD3 paper (alternative to logit-normal). + + This sampling strategy uses a mode-based distribution: + + Formula: u = 1 - u - mode_scale * (cos(π*u/2)² - 1 + u) + + Where u is initially sampled uniformly from [0, 1). + + This creates a distribution that emphasizes certain timesteps + differently than logit-normal, potentially offering better + performance for some models. + + Args: + mode_scale: Scaling factor for mode distribution (default: 1.29) + Value from SD3 paper + + Reference: + Stable Diffusion 3 paper: https://arxiv.org/abs/2403.03206v1 + Section 3.1: Alternative sampling strategy + + Example: + >>> sampler = ModeSampler(mode_scale=1.29) + >>> timesteps, sigmas = sampler.sample(32, 'cuda', scheduler) + """ + + def __init__(self, mode_scale: float = 1.29): + self.mode_scale = mode_scale + + def sample( + self, + batch_size: int, + device: torch.device, + scheduler, + ) -> Tuple[Tensor, Tensor]: + """Sample using mode distribution.""" + u = torch.rand(size=(batch_size,), device="cpu") + u = 1 - u - self.mode_scale * (torch.cos(math.pi * u / 2) ** 2 - 1 + u) + + indices = (u * scheduler.num_train_timesteps).long() + indices = torch.clamp(indices, 0, scheduler.num_train_timesteps - 1) + + timesteps = scheduler.timesteps[indices].to(device=device) + + sigmas = scheduler.sigmas.to(device=device) + schedule_timesteps = scheduler.timesteps.to(device=device) + step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps] + sigma = sigmas[step_indices].flatten() + + return timesteps, sigma + + +def create_timestep_sampler(strategy: str = "logit_normal", **kwargs) -> TimestepSampler: + """ + Factory function for creating timestep samplers. + + This function provides a convenient way to create samplers + by name, allowing easy experimentation with different strategies. + + Args: + strategy: Sampling strategy name + Options: "logit_normal", "uniform", "direct_uniform", "mode" + **kwargs: Additional arguments for the sampler + (e.g., mean/std for LogitNormal, mode_scale for Mode) + + Returns: + TimestepSampler instance + + Raises: + ValueError: If strategy is unknown + + Example: + >>> # Default logit-normal + >>> sampler = create_timestep_sampler("logit_normal") + >>> + >>> # Custom parameters + >>> sampler = create_timestep_sampler( + ... "logit_normal", + ... mean=0.5, + ... std=0.5 + ... ) + >>> + >>> # Uniform baseline + >>> sampler = create_timestep_sampler("uniform") + >>> + >>> # Mode sampling + >>> sampler = create_timestep_sampler("mode", mode_scale=1.5) + """ + if strategy == "logit_normal": + return LogitNormalSampler(**kwargs) + elif strategy == "uniform": + return UniformSampler(**kwargs) + elif strategy == "direct_uniform": + return DirectUniformSampler(**kwargs) + elif strategy == "mode": + return ModeSampler(**kwargs) + else: + raise ValueError( + f"Unknown sampling strategy: {strategy}. " + f"Choose from: logit_normal, uniform, direct_uniform, mode" + ) + + +__all__ = [ + "TimestepSampler", + "LogitNormalSampler", + "UniformSampler", + "DirectUniformSampler", + "ModeSampler", + "create_timestep_sampler", +] diff --git a/tests/unit_tests/backends/megatron/diffusion/training/__init__.py b/tests/unit_tests/backends/megatron/diffusion/training/__init__.py new file mode 100644 index 000000000..948522468 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for Flux training components. + +Tests for trainer classes, forward step, schedulers, and loss computation. +""" diff --git a/tests/unit_tests/backends/megatron/diffusion/training/conftest.py b/tests/unit_tests/backends/megatron/diffusion/training/conftest.py new file mode 100644 index 000000000..b773e6e1b --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/conftest.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Pytest fixtures for diffusion training tests. + +Sets up the Primus logger (required by log_rank_0) and re-exports +the parallel state fixture from the parent conftest. +""" + +import os + +import pytest + +from primus.core.utils import logger +from tests.unit_tests.backends.megatron.conftest import ( # noqa: F401 + init_parallel_state, +) + + +@pytest.fixture(autouse=True, scope="session") +def setup_logger(): + """Initialize Primus logger for tests that use log_rank_0.""" + logger_cfg = logger.LoggerConfig( + exp_root_path=os.environ.get("UT_LOG_PATH", "ut_out"), + work_group="develop", + user_name="root", + exp_name="unittest", + module_name="UT-training", + file_sink_level="DEBUG", + stderr_sink_level="INFO", + node_ip="localhost", + rank=os.environ.get("RANK", 0), + world_size=os.environ.get("WORLD_SIZE", 1), + ) + logger.setup_logger(logger_cfg, is_head=False) diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step.py b/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step.py new file mode 100644 index 000000000..288acf498 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step.py @@ -0,0 +1,219 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for prepare_flux_latents function. + +Tests the higher-level prepare_flux_latents function that orchestrates +latent preparation for training. Lower-level pack/unpack utilities are +tested in test_flux_utils.py. +""" + +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.training.diffusion.forward_step import ( + prepare_flux_latents, +) +from primus.backends.megatron.training.diffusion.schedulers.flow_matching import ( + FlowMatchEulerDiscreteScheduler, +) +from tests.unit_tests.backends.megatron.diffusion.constants import ( + DEFAULT_NUM_TRAIN_TIMESTEPS, + VAE_LATENT_CHANNELS, +) +from tests.unit_tests.backends.megatron.diffusion.helpers import ( + assert_tensor_shape, + create_mock_latents, +) +from tests.utils import PrimusUT + + +class TestPrepareFluxLatents(PrimusUT): + """Tests for prepare_flux_latents function.""" + + def test_prepare_flux_latents_shapes(self): + """Test that prepare_flux_latents produces correct output shapes.""" + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS) + + batch_size = 2 + channels = VAE_LATENT_CHANNELS # VAE output channels + height, width = 128, 128 + + latents = create_mock_latents(batch_size, height, width, channels=channels, seed=42) + h_ids, w_ids = height // 2, width // 2 + img_ids = torch.zeros(batch_size, h_ids * w_ids, 3) + + ( + clean_latents, + noise, + packed_noisy_latents, + img_ids_out, + guidance_vec, + timesteps, + sigma_1d, + ) = prepare_flux_latents( + latents=latents, + scheduler=scheduler, + img_ids=img_ids, + ) + + # Check shapes + assert_tensor_shape(clean_latents, (batch_size, channels, height, width), "clean_latents") + assert_tensor_shape(noise, (batch_size, channels, height, width), "noise") + + # Packed latents: (B, H*W/4, C*4) + expected_packed_shape = ( + batch_size, + (height // 2) * (width // 2), + channels * 4, + ) + assert_tensor_shape(packed_noisy_latents, expected_packed_shape, "packed_noisy_latents") + + # Timesteps: (B,) + assert len(timesteps) == batch_size + + # Sigma: (B,) in [0, 1] + assert_tensor_shape(sigma_1d, (batch_size,), "sigma_1d") + + # Guidance should be None when not used + assert guidance_vec is None + + def test_prepare_flux_latents_with_guidance(self): + """Test prepare_flux_latents with guidance embedding enabled.""" + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS) + + batch_size = 2 + latents = create_mock_latents(batch_size, 128, 128, seed=42) + img_ids = torch.zeros(batch_size, 64 * 64, 3) + guidance_scale = 3.5 + + ( + clean_latents, + noise, + packed_noisy_latents, + img_ids_out, + guidance_vec, + timesteps, + sigma_1d, + ) = prepare_flux_latents( + latents=latents, + scheduler=scheduler, + img_ids=img_ids, + guidance_scale=guidance_scale, + use_guidance_embed=True, + ) + + # Guidance vector should be created + assert guidance_vec is not None + assert_tensor_shape(guidance_vec, (batch_size,), "guidance_vec") + + # All guidance values should be the scale + assert torch.allclose(guidance_vec, torch.tensor(guidance_scale)) + + def test_prepare_flux_latents_timestep_sampling(self): + """Test that timesteps are sampled correctly.""" + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS) + + batch_size = 8 + latents = create_mock_latents(batch_size, 128, 128, seed=42) + img_ids = torch.zeros(batch_size, 64 * 64, 3) + + ( + clean_latents, + noise, + packed_noisy_latents, + img_ids_out, + guidance_vec, + timesteps, + sigma_1d, + ) = prepare_flux_latents( + latents=latents, + scheduler=scheduler, + img_ids=img_ids, + ) + + # Check timesteps are valid + assert len(timesteps) == batch_size + + # All timesteps should be in scheduler's timestep range + for t in timesteps: + assert t >= 0 and t <= scheduler.num_train_timesteps + + def test_prepare_flux_latents_deterministic_with_seed(self): + """Test that results are deterministic with fixed random seed.""" + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS) + + latents = create_mock_latents(2, 64, 64, seed=42) + img_ids = torch.zeros(2, 32 * 32, 3) + + # Set seed and run + torch.manual_seed(12345) + result1 = prepare_flux_latents(latents, scheduler, img_ids) + + # Set same seed and run again + torch.manual_seed(12345) + result2 = prepare_flux_latents(latents, scheduler, img_ids) + + # Results should be identical + assert torch.allclose(result1[2], result2[2]), "Results should be deterministic with same seed" + + def test_timestep_sampling_range(self): + """Test that timestep sampling covers full [0, 1] range with logit-normal.""" + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS) + + # Sample many timesteps to check distribution + num_samples = 1000 + batch_size = 100 + latents = create_mock_latents(batch_size, 64, 64, seed=42) + + all_timesteps = [] + for i in range(num_samples // batch_size): + torch.manual_seed(42 + i) + result = prepare_flux_latents( + latents=latents, + scheduler=scheduler, + img_ids=None, # Test with auto-generation + ) + # Use sigma_1d (index 6) which is already in [0, 1] range + all_timesteps.append(result[6].float()) + + all_timesteps = torch.cat(all_timesteps) + + # Check that we sample across the full range + min_t = all_timesteps.min().item() + max_t = all_timesteps.max().item() + + # With logit-normal, we should get values close to 0 and 1 + assert min_t < 0.1, f"Min timestep {min_t} should be < 0.1 (covers early diffusion)" + assert max_t > 0.9, f"Max timestep {max_t} should be > 0.9 (covers late diffusion)" + + # Check distribution is reasonable (not all in one region) + mean_t = all_timesteps.mean().item() + assert 0.3 < mean_t < 0.7, f"Mean timestep {mean_t} should be roughly centered" + + def test_img_ids_optional_parameter(self): + """Test that img_ids can be None and will be generated automatically.""" + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=DEFAULT_NUM_TRAIN_TIMESTEPS) + + batch_size = 2 + height, width = 128, 128 + latents = create_mock_latents(batch_size, height, width, seed=42) + + # Call without img_ids (should auto-generate) + result = prepare_flux_latents( + latents=latents, + scheduler=scheduler, + img_ids=None, # Test auto-generation + ) + + # Should succeed and return valid img_ids + img_ids_out = result[3] + assert img_ids_out is not None + + # Check shape is correct + expected_shape = (batch_size, (height // 2) * (width // 2), 3) + assert_tensor_shape(img_ids_out, expected_shape, "auto_generated_img_ids") diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_loss_computation.py b/tests/unit_tests/backends/megatron/diffusion/training/test_loss_computation.py new file mode 100644 index 000000000..e5e5e9c39 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_loss_computation.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for loss computation utilities. + +These tests verify the correctness of loss computation functions, +ensuring they produce expected results and maintain backward compatibility. +""" + +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.training.diffusion.loss_computation import ( + compute_flow_matching_loss, +) +from tests.utils import PrimusUT + + +class TestFlowMatchingLoss(PrimusUT): + """Tests for flow matching loss computation.""" + + def test_basic_computation(self): + """Test flow matching loss matches manual calculation.""" + prediction = torch.randn(2, 16, 64, 64) + clean = torch.randn(2, 16, 64, 64) + noise = torch.randn(2, 16, 64, 64) + + loss = compute_flow_matching_loss(prediction, clean, noise) + + # Manual calculation + target = noise - clean + expected = torch.nn.functional.mse_loss(prediction.float(), target.float()) + + assert torch.allclose(loss, expected) + assert loss.dim() == 0 # Scalar + + def test_loss_with_partial_mask(self): + """Test loss computation with partial masking.""" + prediction = torch.randn(2, 16, 64, 64) + clean = torch.randn(2, 16, 64, 64) + noise = torch.randn(2, 16, 64, 64) + + # Create partial mask (half valid) + loss_mask = torch.zeros(2, 16, 64, 64) + loss_mask[:, :, :32, :] = 1.0 # First half valid + + loss_with_mask = compute_flow_matching_loss(prediction, clean, noise, loss_mask) + + # Should be different from unmasked loss + loss_without_mask = compute_flow_matching_loss(prediction, clean, noise) + + assert not torch.allclose(loss_with_mask, loss_without_mask) diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_scheduler.py b/tests/unit_tests/backends/megatron/diffusion/training/test_scheduler.py new file mode 100644 index 000000000..953879aa8 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_scheduler.py @@ -0,0 +1,257 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for diffusion schedulers. + +These tests validate the FlowMatchEulerDiscreteScheduler and other scheduler +implementations used by Flux and future diffusion models. + +The tests follow NeMo's scheduler API: + - scale_noise(): Add noise during training (forward process) + - set_timesteps(): Setup inference schedule + - step(): Perform denoising step (reverse process) +""" + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.training.diffusion.schedulers import ( + FlowMatchEulerDiscreteScheduler, +) +from tests.unit_tests.backends.megatron.diffusion.constants import ( + BATCH_SIZE_PAIR, + BATCH_SIZE_QUAD, + BATCH_SIZE_SINGLE, + DEFAULT_SHIFT, + IMG_SIZE_MICRO, + IMG_SIZE_MINI, + IMG_SIZE_SMALL, + TENSOR_CHANNELS_RGB, + TRAINING_STEPS_MODERATE, +) +from tests.utils import PrimusUT + + +class TestFlowMatchEulerDiscreteScheduler(PrimusUT): + """Tests for FlowMatchEulerDiscreteScheduler.""" + + # ======================================================================== + # Initialization Tests + # ======================================================================== + + def test_scheduler_initialization_with_shift(self): + """Test scheduler with custom shift parameter.""" + scheduler = FlowMatchEulerDiscreteScheduler(shift=2.0) + + assert scheduler.shift == 2.0 + + # With higher shift, sigmas should be different + scheduler_no_shift = FlowMatchEulerDiscreteScheduler(shift=DEFAULT_SHIFT) + assert not torch.allclose(scheduler.sigmas, scheduler_no_shift.sigmas) + + # ======================================================================== + # Sigma Schedule Tests + # ======================================================================== + + def test_sigma_range(self): + """Test sigma values are in valid range [0, 1].""" + scheduler = FlowMatchEulerDiscreteScheduler() + + assert (scheduler.sigmas >= 0.0).all() + assert (scheduler.sigmas <= 1.0).all() + + # Sigmas should be monotonically decreasing + assert (scheduler.sigmas[:-1] >= scheduler.sigmas[1:]).all() + + def test_sigma_schedule_with_shift(self): + """Test shift parameter affects sigma distribution.""" + scheduler_shift_1 = FlowMatchEulerDiscreteScheduler(shift=DEFAULT_SHIFT) + scheduler_shift_2 = FlowMatchEulerDiscreteScheduler(shift=2.0) + + # Higher shift moves distribution towards higher values + mean_shift_1 = scheduler_shift_1.sigmas.mean() + mean_shift_2 = scheduler_shift_2.sigmas.mean() + + assert mean_shift_2 > mean_shift_1 + + # ======================================================================== + # Training Tests (Forward Process - scale_noise) + # ======================================================================== + + def test_scale_noise_shape(self): + """Test scale_noise produces correct shape.""" + scheduler = FlowMatchEulerDiscreteScheduler() + + batch_size = BATCH_SIZE_QUAD # Quad sample tests + channels = 64 + height = IMG_SIZE_SMALL # Small size for standard tests + width = IMG_SIZE_SMALL + + sample = torch.randn(batch_size, channels, height, width) + noise = torch.randn(batch_size, channels, height, width) + timesteps = scheduler.timesteps[[100, 200, 300, 400]] + + noisy = scheduler.scale_noise(sample, timesteps, noise) + + assert noisy.shape == sample.shape + assert noisy.dtype == sample.dtype + + def test_scale_noise_flow_matching_interpolation(self): + """Test flow matching interpolation: x_t = (1-σ)*x_0 + σ*noise.""" + scheduler = FlowMatchEulerDiscreteScheduler() + + sample = torch.ones( + BATCH_SIZE_SINGLE, + TENSOR_CHANNELS_RGB, + IMG_SIZE_MICRO, + IMG_SIZE_MICRO, + ) + noise = torch.zeros( + BATCH_SIZE_SINGLE, + TENSOR_CHANNELS_RGB, + IMG_SIZE_MICRO, + IMG_SIZE_MICRO, + ) + + # At σ≈0 (t=0), should be mostly sample + timesteps = scheduler.timesteps[-1:] # Last timestep (smallest sigma) + noisy = scheduler.scale_noise(sample, timesteps, noise) + assert torch.allclose(noisy, sample, atol=0.1) + + # At σ≈1 (t=max), should be mostly noise + timesteps = scheduler.timesteps[0:1] # First timestep (largest sigma) + noisy = scheduler.scale_noise(sample, timesteps, noise) + assert torch.allclose(noisy, noise, atol=0.1) + + def test_scale_noise_without_provided_noise(self): + """Test scale_noise generates noise if not provided.""" + scheduler = FlowMatchEulerDiscreteScheduler() + + sample = torch.randn(BATCH_SIZE_PAIR, TENSOR_CHANNELS_RGB, IMG_SIZE_MINI, IMG_SIZE_MINI) + timesteps = scheduler.timesteps[[100, 200]] + + # Should generate noise internally + noisy = scheduler.scale_noise(sample, timesteps, noise=None) + + assert noisy.shape == sample.shape + # Should be different from original (has noise added) + assert not torch.allclose(noisy, sample) + + def test_scale_noise_dtype_consistency(self): + """Test scale_noise preserves dtype.""" + scheduler = FlowMatchEulerDiscreteScheduler() + + for dtype in [torch.float32, torch.float16]: + sample = torch.randn( + BATCH_SIZE_PAIR, + TENSOR_CHANNELS_RGB, + IMG_SIZE_MINI, + IMG_SIZE_MINI, + dtype=dtype, + ) + noise = torch.randn( + BATCH_SIZE_PAIR, + TENSOR_CHANNELS_RGB, + IMG_SIZE_MINI, + IMG_SIZE_MINI, + dtype=dtype, + ) + timesteps = scheduler.timesteps[[100, 200]] + + noisy = scheduler.scale_noise(sample, timesteps, noise) + + assert noisy.dtype == dtype + + # ======================================================================== + # Inference Tests (Reverse Process - set_timesteps, step) + # ======================================================================== + + def test_set_timesteps_basic(self): + """Test set_timesteps creates inference schedule.""" + scheduler = FlowMatchEulerDiscreteScheduler() + + num_steps = TRAINING_STEPS_MODERATE # Moderate training steps + scheduler.set_timesteps(num_inference_steps=num_steps, device="cpu") + + assert len(scheduler.timesteps) == num_steps + assert scheduler.timesteps[0] > scheduler.timesteps[-1] # Descending + assert len(scheduler.sigmas) == num_steps + 1 # +1 for final zero + + # Verify sigma at end is zero + assert scheduler.sigmas[-1] == 0.0 + + def test_set_timesteps_with_dynamic_shift(self): + """Test set_timesteps with dynamic shifting.""" + scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True) + + mu = 0.5 # Shift parameter based on image resolution + scheduler.set_timesteps(num_inference_steps=10, device="cpu", mu=mu) + + assert len(scheduler.timesteps) == 10 + + # Should raise error if mu not provided + scheduler2 = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True) + with pytest.raises(ValueError, match="Must pass a value for `mu`"): + scheduler2.set_timesteps(num_inference_steps=10) + + def test_step_basic(self): + """Test step performs Euler denoising step.""" + scheduler = FlowMatchEulerDiscreteScheduler() + scheduler.set_timesteps(num_inference_steps=5, device="cpu") + + sample = torch.randn(BATCH_SIZE_SINGLE, TENSOR_CHANNELS_RGB, IMG_SIZE_MINI, IMG_SIZE_MINI) + model_output = torch.randn( + BATCH_SIZE_SINGLE, TENSOR_CHANNELS_RGB, IMG_SIZE_MINI, IMG_SIZE_MINI + ) # Velocity prediction + timestep = scheduler.timesteps[0] + + result = scheduler.step(model_output, timestep, sample) + + # step() returns a tuple (prev_sample,) + prev_sample = result[0] + + assert prev_sample.shape == sample.shape + assert prev_sample.dtype == sample.dtype + # Should be different from input + assert not torch.equal(prev_sample, sample) + + def test_step_consistency(self): + """Test multiple steps are consistent.""" + scheduler = FlowMatchEulerDiscreteScheduler() + scheduler.set_timesteps(num_inference_steps=5, device="cpu") + + sample = torch.randn(BATCH_SIZE_SINGLE, TENSOR_CHANNELS_RGB, IMG_SIZE_MINI, IMG_SIZE_MINI) + model_output = torch.randn(BATCH_SIZE_SINGLE, TENSOR_CHANNELS_RGB, IMG_SIZE_MINI, IMG_SIZE_MINI) + + # Perform multiple steps + current_sample = sample + for timestep in scheduler.timesteps: + result = scheduler.step(model_output, timestep, current_sample) + # step() returns a tuple (prev_sample,) + current_sample = result[0] + + # Final sample should be different from initial + assert not torch.equal(current_sample, sample) + assert current_sample.shape == sample.shape + + def test_dynamic_shifting_behavior(self): + """Test dynamic shifting changes timesteps based on resolution.""" + scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=True) + + # Test with different mu values (resolution-dependent) + mu_small = 0.3 # Smaller resolution + mu_large = 0.8 # Larger resolution + + scheduler.set_timesteps(num_inference_steps=10, device="cpu", mu=mu_small) + timesteps_small = scheduler.timesteps.clone() + + scheduler.set_timesteps(num_inference_steps=10, device="cpu", mu=mu_large) + timesteps_large = scheduler.timesteps.clone() + + # Timesteps should be different for different mu + assert not torch.allclose(timesteps_small, timesteps_large) From c8e5c09dd1a6409310d7dd984f3eb550bc08b35f Mon Sep 17 00:00:00 2001 From: Fuyuan Jing <167437074+amd-fuyuajin@users.noreply.github.com> Date: Tue, 14 Jul 2026 21:22:36 -0400 Subject: [PATCH 026/127] Dev/production doc (#824) working on production documentation --------- Co-authored-by: Cursor Co-authored-by: Peter Park --- .gitignore | 1 + README.md | 18 +- docs/.gitignore | 2 + docs/.readthedocs.yaml | 16 + docs/01-getting-started/README.md | 12 + docs/01-getting-started/glossary.md | 250 +++++ docs/01-getting-started/installation.md | 265 ++++++ docs/01-getting-started/overview.md | 137 +++ docs/01-getting-started/quickstart.md | 171 ++++ docs/02-user-guide/README.md | 18 + docs/02-user-guide/benchmarking.md | 204 +++++ docs/02-user-guide/cli-reference.md | 233 +++++ docs/02-user-guide/configuration-system.md | 170 ++++ docs/02-user-guide/posttraining.md | 216 +++++ docs/02-user-guide/preflight.md | 136 +++ docs/02-user-guide/pretraining.md | 291 ++++++ docs/02-user-guide/primus-tools.md | 29 + docs/02-user-guide/projection.md | 177 ++++ docs/02-user-guide/training-recipes.md | 260 ++++++ .../tuning-agent.md} | 140 +-- docs/03-configuration-reference/README.md | 13 + .../environment-variables.md | 239 +++++ .../maxtext-parameters.md | 146 +++ .../megatron-bridge-parameters.md | 167 ++++ .../megatron-parameters.md | 854 ++++++++++++++++++ .../torchtitan-parameters.md | 362 ++++++++ docs/04-technical-guides/README.md | 22 + .../checkpoint-management.md | 203 +++++ .../collective-operations.md | 289 ++++++ docs/04-technical-guides/data-preparation.md | 159 ++++ .../determinism-and-reproducibility.md | 119 +++ .../diffusion-models}/README.md | 62 +- .../diffusion-models}/STRUCTURE.md | 45 +- .../diffusion-models}/adding_new_models.md | 54 +- .../diffusion-models}/api_reference.md | 80 +- .../architecture_overview.md | 128 +-- .../diffusion-models}/data_preprocessing.md | 86 +- .../diffusion-models}/energon_integration.md | 78 +- .../diffusion-models}/flux_architecture.md | 118 +-- .../diffusion-models}/fp8_training.md | 64 +- .../diffusion-models}/mxfp4_training.md | 48 +- .../fault-tolerance-and-elastic-training.md | 123 +++ .../logging-and-experiment-tracking.md | 163 ++++ docs/04-technical-guides/moe-training.md | 193 ++++ .../multi-node-networking.md | 203 +++++ .../native-sft-lora.md} | 33 +- .../parallelism-configuration.md | 255 ++++++ .../parallelism-strategies.md | 361 ++++++++ .../04-technical-guides/performance-tuning.md | 242 +++++ .../profiling-and-observability.md | 170 ++++ docs/05-operations/README.md | 12 + docs/05-operations/deployment.md | 237 +++++ docs/05-operations/monitoring-logging.md | 222 +++++ docs/05-operations/security.md | 154 ++++ docs/05-operations/troubleshooting.md | 193 ++++ docs/06-developer-guide/README.md | 17 + docs/06-developer-guide/adding-models.md | 379 ++++++++ docs/06-developer-guide/architecture.md | 119 +++ .../06-developer-guide/backend-patch-notes.md | 142 +++ .../cli-architecture.md} | 55 +- docs/06-developer-guide/contributing.md | 142 +++ docs/06-developer-guide/extending-backends.md | 419 +++++++++ .../model-support-matrix.md | 210 +++++ docs/06-developer-guide/testing.md | 122 +++ docs/06-developer-guide/tooling.md | 24 + docs/README.md | 172 +++- docs/conf.py | 37 + docs/license.md | 4 + docs/sphinx/_toc.yml.in | 139 +++ docs/sphinx/requirements.in | 1 + docs/sphinx/requirements.txt | 279 ++++++ docs_deprecated/README.md | 82 ++ .../megatron-upstream-main-2026-04-30.json | 0 .../torchtitan-upstream-main-2026-04-21.json | 0 .../reports/megatron/upstream-main/report.md | 0 .../reports/megatron/upstream-main/summary.md | 0 .../torchtitan/upstream-main/report.md | 0 .../torchtitan/upstream-main/summary.md | 0 .../backends/adding-megatron-models.md | 0 .../backends/adding-torchtitan-models.md | 0 .../backends/extending-backends.md | 0 .../backends/maxtext/patch-notes.md | 0 .../backends/megatron/patch-notes.md | 0 .../backends/overview.md | 0 .../backends/torchtitan/patch-notes.md | 0 {docs => docs_deprecated}/benchmark.md | 0 .../cli/PRIMUS-CLI-GUIDE.md | 2 +- {docs => docs_deprecated}/cli/README.md | 4 +- {docs => docs_deprecated}/install-on-host.md | 0 {docs => docs_deprecated}/posttraining.md | 0 {docs => docs_deprecated}/preflight.md | 0 {docs => docs_deprecated}/projection.md | 2 +- {docs => docs_deprecated}/quickstart.md | 2 +- .../primus_cli_unified_entry_rocm.md | 0 .../primus_pipeline/imgs/actual-perf.png | Bin .../primus_pipeline/imgs/llama2-7B-perf.png | Bin .../primus_pipeline/imgs/qwen-235B-perf.png | Bin .../primus_pipeline/imgs/simulation.png | Bin .../primus_pipeline/imgs/simulator_shell.png | Bin .../primus_pipeline/primus_pipeline.md | 0 .../tech_blogs/projection/projection.md | 0 .../weekly_reports/2026-W17-primus-weekly.md | 0 .../weekly_reports/2026-W18-primus-weekly.md | 0 .../weekly_reports/2026-W19-primus-weekly.md | 0 .../dashboard-data/reports/2026-W17.json | 0 .../dashboard-data/reports/2026-W18.json | 0 .../dashboard-data/reports/2026-W19.json | 0 examples/README.md | 4 +- .../diffusion/flux_535m_pretrain_fp8.yaml | 4 +- .../diffusion/flux_535m_pretrain_fp8.yaml | 4 +- examples/megatron/diffusion/README.md | 12 +- primus/agents/tuning_agent/README.md | 2 +- .../configs/data/megatron/diffusion/README.md | 16 +- .../diffusion/templates/metadataset.yaml | 2 +- .../models/megatron/diffusion/encoders.yaml | 2 +- .../megatron/primus_megatron_module.yaml | 2 +- tests/README.md | 49 + tools/README.md | 24 + 118 files changed, 10535 insertions(+), 582 deletions(-) create mode 100644 docs/.gitignore create mode 100644 docs/.readthedocs.yaml create mode 100644 docs/01-getting-started/README.md create mode 100644 docs/01-getting-started/glossary.md create mode 100644 docs/01-getting-started/installation.md create mode 100644 docs/01-getting-started/overview.md create mode 100644 docs/01-getting-started/quickstart.md create mode 100644 docs/02-user-guide/README.md create mode 100644 docs/02-user-guide/benchmarking.md create mode 100644 docs/02-user-guide/cli-reference.md create mode 100644 docs/02-user-guide/configuration-system.md create mode 100644 docs/02-user-guide/posttraining.md create mode 100644 docs/02-user-guide/preflight.md create mode 100644 docs/02-user-guide/pretraining.md create mode 100644 docs/02-user-guide/primus-tools.md create mode 100644 docs/02-user-guide/projection.md create mode 100644 docs/02-user-guide/training-recipes.md rename docs/{tuning_agent.md => 02-user-guide/tuning-agent.md} (89%) create mode 100644 docs/03-configuration-reference/README.md create mode 100644 docs/03-configuration-reference/environment-variables.md create mode 100644 docs/03-configuration-reference/maxtext-parameters.md create mode 100644 docs/03-configuration-reference/megatron-bridge-parameters.md create mode 100644 docs/03-configuration-reference/megatron-parameters.md create mode 100644 docs/03-configuration-reference/torchtitan-parameters.md create mode 100644 docs/04-technical-guides/README.md create mode 100644 docs/04-technical-guides/checkpoint-management.md create mode 100644 docs/04-technical-guides/collective-operations.md create mode 100644 docs/04-technical-guides/data-preparation.md create mode 100644 docs/04-technical-guides/determinism-and-reproducibility.md rename docs/{backends/megatron/diffusion => 04-technical-guides/diffusion-models}/README.md (91%) rename docs/{backends/megatron/diffusion => 04-technical-guides/diffusion-models}/STRUCTURE.md (89%) rename docs/{backends/megatron/diffusion => 04-technical-guides/diffusion-models}/adding_new_models.md (95%) rename docs/{backends/megatron/diffusion => 04-technical-guides/diffusion-models}/api_reference.md (97%) rename docs/{backends/megatron/diffusion => 04-technical-guides/diffusion-models}/architecture_overview.md (81%) rename docs/{backends/megatron/diffusion => 04-technical-guides/diffusion-models}/data_preprocessing.md (93%) rename docs/{backends/megatron/diffusion => 04-technical-guides/diffusion-models}/energon_integration.md (92%) rename docs/{backends/megatron/diffusion => 04-technical-guides/diffusion-models}/flux_architecture.md (93%) rename docs/{backends/megatron/diffusion => 04-technical-guides/diffusion-models}/fp8_training.md (93%) rename docs/{backends/megatron/diffusion => 04-technical-guides/diffusion-models}/mxfp4_training.md (82%) create mode 100644 docs/04-technical-guides/fault-tolerance-and-elastic-training.md create mode 100644 docs/04-technical-guides/logging-and-experiment-tracking.md create mode 100644 docs/04-technical-guides/moe-training.md create mode 100644 docs/04-technical-guides/multi-node-networking.md rename docs/{README_NATIVE_SFT_LORA_EN.md => 04-technical-guides/native-sft-lora.md} (95%) create mode 100644 docs/04-technical-guides/parallelism-configuration.md create mode 100644 docs/04-technical-guides/parallelism-strategies.md create mode 100644 docs/04-technical-guides/performance-tuning.md create mode 100644 docs/04-technical-guides/profiling-and-observability.md create mode 100644 docs/05-operations/README.md create mode 100644 docs/05-operations/deployment.md create mode 100644 docs/05-operations/monitoring-logging.md create mode 100644 docs/05-operations/security.md create mode 100644 docs/05-operations/troubleshooting.md create mode 100644 docs/06-developer-guide/README.md create mode 100644 docs/06-developer-guide/adding-models.md create mode 100644 docs/06-developer-guide/architecture.md create mode 100644 docs/06-developer-guide/backend-patch-notes.md rename docs/{cli/CLI-ARCHITECTURE.md => 06-developer-guide/cli-architecture.md} (93%) create mode 100644 docs/06-developer-guide/contributing.md create mode 100644 docs/06-developer-guide/extending-backends.md create mode 100644 docs/06-developer-guide/model-support-matrix.md create mode 100644 docs/06-developer-guide/testing.md create mode 100644 docs/06-developer-guide/tooling.md create mode 100644 docs/conf.py create mode 100644 docs/license.md create mode 100644 docs/sphinx/_toc.yml.in create mode 100644 docs/sphinx/requirements.in create mode 100644 docs/sphinx/requirements.txt create mode 100644 docs_deprecated/README.md rename {docs => docs_deprecated}/backend-gap/dashboard-data/reports/megatron-upstream-main-2026-04-30.json (100%) rename {docs => docs_deprecated}/backend-gap/dashboard-data/reports/torchtitan-upstream-main-2026-04-21.json (100%) rename {docs => docs_deprecated}/backend-gap/reports/megatron/upstream-main/report.md (100%) rename {docs => docs_deprecated}/backend-gap/reports/megatron/upstream-main/summary.md (100%) rename {docs => docs_deprecated}/backend-gap/reports/torchtitan/upstream-main/report.md (100%) rename {docs => docs_deprecated}/backend-gap/reports/torchtitan/upstream-main/summary.md (100%) rename {docs => docs_deprecated}/backends/adding-megatron-models.md (100%) rename {docs => docs_deprecated}/backends/adding-torchtitan-models.md (100%) rename {docs => docs_deprecated}/backends/extending-backends.md (100%) rename {docs => docs_deprecated}/backends/maxtext/patch-notes.md (100%) rename {docs => docs_deprecated}/backends/megatron/patch-notes.md (100%) rename {docs => docs_deprecated}/backends/overview.md (100%) rename {docs => docs_deprecated}/backends/torchtitan/patch-notes.md (100%) rename {docs => docs_deprecated}/benchmark.md (100%) rename {docs => docs_deprecated}/cli/PRIMUS-CLI-GUIDE.md (99%) rename {docs => docs_deprecated}/cli/README.md (95%) rename {docs => docs_deprecated}/install-on-host.md (100%) rename {docs => docs_deprecated}/posttraining.md (100%) rename {docs => docs_deprecated}/preflight.md (100%) rename {docs => docs_deprecated}/projection.md (99%) rename {docs => docs_deprecated}/quickstart.md (96%) rename {docs => docs_deprecated}/tech_blogs/primus_cli_unified_entry_rocm.md (100%) rename {docs => docs_deprecated}/tech_blogs/primus_pipeline/imgs/actual-perf.png (100%) rename {docs => docs_deprecated}/tech_blogs/primus_pipeline/imgs/llama2-7B-perf.png (100%) rename {docs => docs_deprecated}/tech_blogs/primus_pipeline/imgs/qwen-235B-perf.png (100%) rename {docs => docs_deprecated}/tech_blogs/primus_pipeline/imgs/simulation.png (100%) rename {docs => docs_deprecated}/tech_blogs/primus_pipeline/imgs/simulator_shell.png (100%) rename {docs => docs_deprecated}/tech_blogs/primus_pipeline/primus_pipeline.md (100%) rename {docs => docs_deprecated}/tech_blogs/projection/projection.md (100%) rename {docs => docs_deprecated}/weekly_reports/2026-W17-primus-weekly.md (100%) rename {docs => docs_deprecated}/weekly_reports/2026-W18-primus-weekly.md (100%) rename {docs => docs_deprecated}/weekly_reports/2026-W19-primus-weekly.md (100%) rename {docs => docs_deprecated}/weekly_reports/dashboard-data/reports/2026-W17.json (100%) rename {docs => docs_deprecated}/weekly_reports/dashboard-data/reports/2026-W18.json (100%) rename {docs => docs_deprecated}/weekly_reports/dashboard-data/reports/2026-W19.json (100%) diff --git a/.gitignore b/.gitignore index 3353a8bdb..1d04b58e5 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,4 @@ output experiment /data/* pp_simulation_result +.cursor/ diff --git a/README.md b/README.md index 481149035..12b8810f9 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ ## ✨ Key Features - **🔄 Multi-Backend Support**: Seamlessly switch between Megatron-LM, TorchTitan, and other training frameworks -- **🚀 Unified CLI**: One command interface for local development, containers, and Slurm clusters ([Docs](./docs/README.md)) +- **🚀 Unified CLI**: One command interface for local development, containers, and Slurm clusters ([Docs](./docs/02-user-guide/cli-reference.md)) - **⚡ ROCm Optimized**: Deep integration with AMD ROCm stack and optimized kernels from Primus-Turbo - **📦 Production Ready**: Battle-tested on large-scale training with hundreds of GPUs - **🔌 Extensible Architecture**: Plugin-based design for easy integration of custom models and workflows @@ -27,7 +27,7 @@ - **TorchTitan**: LLaMA3 / LLaMA4, DeepSeek-V3, and related decoder-only architectures - **MaxText (JAX)**: LLaMA3.x and other MaxText-supported transformer models (subset; see MaxText docs for details) -For the full and up-to-date model matrix, see [Supported Models](./docs/backends/overview.md#supported-models). +For the full and up-to-date model matrix, see [Supported Models](./docs/06-developer-guide/model-support-matrix.md). --- @@ -93,7 +93,7 @@ primus-cli deps sync --dir ~/.cache/Primus/third_party # For Megatron-LM and TorchTitan backends docker pull rocm/primus:v26.3 # For MaxText backend - docker pull rocm/jax-training:v26.3 + docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0 ``` 2. **Clone the repository** @@ -118,7 +118,7 @@ primus-cli deps sync --dir ~/.cache/Primus/third_party -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml ``` -For more detailed usage instructions, see the [CLI User Guide](./docs/cli/PRIMUS-CLI-GUIDE.md). +For more detailed usage instructions, see the [CLI User Guide](./docs/02-user-guide/cli-reference.md). #### Option 2: wheel installation of Primus and run training in container @@ -137,7 +137,7 @@ For more detailed usage instructions, see the [CLI User Guide](./docs/cli/PRIMUS >**Note**: this will only install the Primus CLI in your virtual environment under the `site-packages` directory, without other dependencies. The third party submodules will be downloaded on the first run. The complete dependencies and training software stack is provided in the AMD published training Docker images. You can use `primus-cli` to launch the training in container from any directory. - >**Note**: If you don't want to use docker container to run training, and want to install the complete dependencies and training software stack on your host machine, please refer to the instruction: [Install training environment on your host machine](docs/install-on-host.md). The automated installation script is under development and will be released soon. + >**Note**: If you don't want to use docker container to run training, and want to install the complete dependencies and training software stack on your host machine, please refer to the instruction: [Install training environment on your host machine](docs/01-getting-started/installation.md#bare-metal-host-setup). The automated installation script is under development and will be released soon. 2. **Run training in container using pip-installed Primus** @@ -162,10 +162,10 @@ For more detailed usage instructions, see the [CLI User Guide](./docs/cli/PRIMUS Comprehensive documentation is available in the [`docs/`](./docs/) directory: -- **[Quick Start Guide](./docs/quickstart.md)** - Get started in 5 minutes -- **[Primus CLI User Guide](./docs/cli/PRIMUS-CLI-GUIDE.md)** - Complete CLI reference and usage -- **[CLI Architecture](./docs/cli/CLI-ARCHITECTURE.md)** - Technical design and architecture -- **[Backend Patch Notes](./docs/backends/overview.md)** - Primus-specific backend arguments +- **[Quick Start Guide](./docs/01-getting-started/quickstart.md)** - Get started in 5 minutes +- **[Primus CLI User Guide](./docs/02-user-guide/cli-reference.md)** - Complete CLI reference and usage +- **[CLI Architecture](./docs/06-developer-guide/cli-architecture.md)** - Technical design and architecture +- **[Backend Patch Notes](./docs/06-developer-guide/backend-patch-notes.md)** - Primus-specific backend arguments - **[Full Documentation Index](./docs/README.md)** - Browse all available documentation --- diff --git a/docs/.gitignore b/docs/.gitignore new file mode 100644 index 000000000..c6cf22b0f --- /dev/null +++ b/docs/.gitignore @@ -0,0 +1,2 @@ +_build +sphinx/_toc.yml diff --git a/docs/.readthedocs.yaml b/docs/.readthedocs.yaml new file mode 100644 index 000000000..75a0bb1c4 --- /dev/null +++ b/docs/.readthedocs.yaml @@ -0,0 +1,16 @@ +# Read the Docs configuration file +# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details + +version: 2 + +build: + os: ubuntu-24.04 + tools: + python: "3.12" + +sphinx: + configuration: docs/conf.py + +python: + install: + - requirements: docs/sphinx/requirements.txt diff --git a/docs/01-getting-started/README.md b/docs/01-getting-started/README.md new file mode 100644 index 000000000..559b9b52f --- /dev/null +++ b/docs/01-getting-started/README.md @@ -0,0 +1,12 @@ +# Getting started + +Start here if you are new to Primus. + +- [Project overview](overview.md): what Primus does, who it is for, key capabilities +- [Installation guide](installation.md): prerequisites, Docker/bare-metal/Slurm setup +- [Quickstart](quickstart.md): first training run in 5 minutes +- [Glossary](glossary.md): terms, acronyms, and domain concepts + +--- + +[← Documentation home](../README.md) diff --git a/docs/01-getting-started/glossary.md b/docs/01-getting-started/glossary.md new file mode 100644 index 000000000..4ede3de7b --- /dev/null +++ b/docs/01-getting-started/glossary.md @@ -0,0 +1,250 @@ +# Glossary + +Alphabetical reference for terms used in Primus documentation and configuration. Cross-links point to other production docs where applicable. + +--- + +### AINIC + +**AMD AI NIC**—AMD’s AI-optimized network interface for multi-node GPU communication (for example, the **AMD Pensando™ Pollara 400 AI NIC**). + +--- + +### Backend + +A training framework integrated into Primus (for example **Megatron-LM**, **TorchTitan**, **MaxText**, **Megatron Bridge**, **HummingbirdXT**). + +--- + +### BackendAdapter + +Abstract class in Primus that connects a backend: discovery of setup paths, config conversion, and trainer loading. + +--- + +### BackendRegistry + +Registry mapping backend names to adapter classes, often with **lazy import** to avoid loading unused frameworks. + +--- + +### BaseTrainer + +Abstract trainer defining the lifecycle: **setup** → **init** → **train** → **cleanup**. + +--- + +### BF16 / FP16 / FP8 / FP4 + +Floating-point precisions: **Brain Float 16**, **IEEE half**, **8-bit float**, and **4-bit float** training or inference formats (exact support depends on backend and hardware). + +--- + +### CP (context parallelism) + +Parallelism that **splits the sequence dimension** across devices for long-context training. + +--- + +### DP (data parallelism) + +Replicates the model across GPUs; each rank processes **different data batches**. + +--- + +### DeepEP + +**Deep Expert Parallelism**—Primus-Turbo’s acceleration path for **MoE token dispatch** and related expert-parallel work. + +--- + +### EP (expert parallelism) + +Distributes **Mixture-of-Experts** expert networks across devices. + +--- + +### Experiment config + +Top-level **YAML** describing `work_group`, **modules**, and **overrides** for a training run. + +--- + +### FSDP + +**Fully Sharded Data Parallel**—shards parameters, gradients, and optimizer states across devices (PyTorch FSDP and similar concepts per backend). + +--- + +### GBS (global batch size) + +Total **batch size across all data-parallel ranks** for one optimizer step (may combine micro-batching and gradient accumulation). + +--- + +### Gradient accumulation + +Accumulates gradients over **multiple micro-batches** before an optimizer update. + +--- + +### HipBLASLt + +AMD’s high-performance **BLAS** library with **autotuning** for GEMM and related kernels. + +--- + +### Hook + +Shell or Python scripts under `runner/helpers/hooks/` executed at defined **lifecycle** points. + +--- + +### LoRA + +**Low-Rank Adaptation**—parameter-efficient fine-tuning that trains small adapter matrices. + +--- + +### MBS (micro batch size) + +Batch size **per GPU** (per rank) for **one forward/backward pass** within a gradient-accumulation window. + +--- + +### MLA (multi-latent attention) + +Compressed **KV-cache** attention architecture used in models such as DeepSeek. + +--- + +### MoE (mixture of experts) + +Architecture with **multiple expert** sub-networks and a **router** that assigns tokens to experts. + +--- + +### Model config + +YAML **preset** describing architecture (hidden size, layers, attention heads, and so on). + +--- + +### Module config + +YAML **preset** for training behavior: learning rate, batch sizes, optimizer, schedules. + +--- + +### NCCL / RCCL + +**NVIDIA Collective Communications Library** / **ROCm** equivalent—libraries for **GPU collective** operations in distributed training. + +--- + +### PP (pipeline parallelism) + +Splits **model layers** into **stages** on different devices. + +--- + +### Patch + +Runtime **monkey-patch** registered in **PatchRegistry** and applied at a named training phase. + +--- + +### PatchRegistry + +Registry of **phase-aware** patches (for example `build_args`, `setup`, `before_train`, `after_train`). + +--- + +### Platform config + +YAML describing **cluster environment** mappings (for example `platform_azure.yaml`): env vars, paths, and scheduler hints. + +--- + +### Preflight + +Cluster **diagnostic** tooling that checks host, GPU, network, and baseline performance before long jobs. See `primus/tools/preflight/` in the repository. + +--- + +### Preset + +Reusable YAML fragment under `primus/configs/` (**module**, **model**, or **platform**). + +--- + +### PrimusRuntime + +Core **orchestrator**: loads configuration, resolves the backend, applies patches, and drives the **trainer lifecycle**. + +--- + +### Primus-SaFE + +**Stability and Fault-tolerance Engine**—external ecosystem component for **cluster management** and resilience. This repository references it in auxiliary tooling but does not include a production integration guide. + +--- + +### Primus-Turbo + +High-performance **operator** library (for example FlashAttention-style kernels, GEMM, collectives, grouped GEMM). + +--- + +### Projection + +Tools that **estimate memory** and **training performance** without requiring a full production cluster. + +--- + +### ROCm + +**Radeon Open Compute**—AMD’s GPU computing platform (drivers, compilers, libraries). + +--- + +### SFT (supervised fine-tuning) + +Supervised fine-tuning that typically **updates all** (or a defined subset of) model parameters, as opposed to adapter-only methods. + +--- + +### SP (sequence parallelism) + +Parallelism that extends tensor-parallel regions to **non-TP** parts of the model to **reduce activation memory**. + +--- + +### TP (tensor parallelism) + +Splits **layer weights** across GPUs within a node (or defined process group). + +--- + +### Transformer engine (TE) + +Library stack for **FP8** and related training optimizations (availability depends on backend and build). + +--- + +### VPP (virtual pipeline parallelism) + +**Interleaved** pipeline parallelism with **multiple virtual stages** per device to improve utilization. + +--- + +### Zero-bubble + +Pipeline scheduling that **reduces or eliminates pipeline bubbles** (idle time between micro-batches). + +--- + +## Related documentation + +- [Overview](./overview.md) +- [Configuration system](../02-user-guide/configuration-system.md) diff --git a/docs/01-getting-started/installation.md b/docs/01-getting-started/installation.md new file mode 100644 index 000000000..f660fbba7 --- /dev/null +++ b/docs/01-getting-started/installation.md @@ -0,0 +1,265 @@ +# Installation and setup + +This guide covers supported platforms, prerequisites, and how to set up the training environment: **container (recommended)** and **bare metal**, plus **multi-node distributed training** (Slurm recommended). + +--- + +## Supported platforms + + +| Requirement | Notes | +| ----------- | ----------------------------------------------------------------------------------------------------------- | +| **OS** | Linux (ROCm-supported distributions per AMD documentation). | +| **ROCm** | **≥ 7.0** recommended. | +| **GPUs** | AMD Instinct™ **MI300X**, **MI325X**, **MI355X** (or other ROCm-supported Instinct SKUs your site supports). | + + +--- + +## Prerequisites + + +| Prerequisite | Purpose | +| ------------------------------------------------------------- | --------------------------------------------- | +| **AMD Instinct GPUs** | Training and benchmarks execute on GPU. | +| **ROCm drivers and user-space stack** | Required for HIP, RCCL, and ML frameworks. | +| **Docker ≥ 24.0** (or Podman with compatible GPU passthrough) | Container mode and reproducible environments. | +| **git** | Clone the repository and submodules. | + + +### Quick environment checks + +```bash +rocm-smi +docker --version +``` + +`rocm-smi` should list your GPUs; `docker --version` should report **24.0** or newer. + +--- + +## Container setup (recommended) + +AMD publishes training Docker images monthly, providing a consistent, ready-to-run environment optimized for AMD GPUs. It is recommended to use the AMD-published training Docker images together with this Primus-LM repository to run your training jobs. The images support pre-training and post-training workflows with multiple backends including Megatron-LM, TorchTitan, and JAX MaxText, alongside ROCm-optimized components. + +Check the AMD-published training Docker images here: + +- For Megatron-LM and TorchTitan backends: [https://hub.docker.com/r/rocm/primus/tags](https://hub.docker.com/r/rocm/primus/tags) +- For MaxText backend: [https://hub.docker.com/r/rocm/jax-training/tags](https://hub.docker.com/r/rocm/jax-training/tags) + +### 1. Pull the image + +```bash +# For Megatron-LM and TorchTitan backends +docker pull rocm/primus:v26.3 +# For MaxText backend +docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0 +``` + +### 2. Clone the repository + +Submodules are required for third-party backends and tools: + +```bash +git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git +cd Primus +# checkout the branch for the specific release +git checkout release/v26.3 +git submodule update --init --recursive +``` + +### 3. Run a verification benchmark + +From the repository root: + +```bash +./primus-cli container --image rocm/primus:v26.3 -- \ + benchmark gemm --M 4096 --N 4096 --K 4096 +``` + +A successful run validates the GPU stack and Primus CLI wiring without launching a full training job. + +--- + +## Bare-metal (host) setup + +> **The [container setup](#container-setup-recommended) above is strongly recommended.** The AMD-published training Docker image is the tested, reproducible, and best-supported path. Build the full stack on a bare-metal host only when containers are not an option (for example, due to policy or operational constraints). + +### What to expect + +Reproducing the training environment on the host means building the **same stack the Docker image ships**, mostly from source. This is a **long, build-heavy process** that requires: + +- Several **source-built kernel libraries** (Flash Attention, TransformerEngine, AITER, Primus-Turbo, Grouped GEMM, causal-conv1d, Mamba) compiled against ROCm. +- A **machine with many CPU cores, ample RAM, and tens of GB of free disk**. +- A long time (expect a **multi-hour first build**). + +### What a complete host environment needs + +| Layer | What it provides | How it is installed | +| ----------------------- | -------------------------------------------------------------------------------- | -------------------------------- | +| Kernel/hardware | AMD GPU driver (amdgpu KMD) and device access (`/dev/kfd`, `/dev/dri`) | OS/administrator (root, one-time) | +| OS libraries | Build toolchain and runtime libraries (`g++`, `git`, RDMA, hwloc, etc.) | `apt` (root, one-time) | +| ROCm user-space | `rocm-sdk-devel` + device wheels—**no system-wide ROCm install required** | `pip` (TheRock wheels, in `venv`) | +| Deep learning framework | ROCm-enabled PyTorch (`torch`, `torchvision`, `torchaudio`, `apex`) | `pip` (TheRock wheels, in `venv`) | +| Accelerated kernels | Flash Attention, TransformerEngine, AITER, Primus-Turbo, Grouped GEMM, Mamba | build from source (in `venv`) | +| Multi-node communications | UCX, OpenMPI, rocSHMEM, AMD AINIC—only for distributed (multi-node) training | build from source or `apt` | +| Primus + Python dependencies | Primus, submodules, and training libraries (datasets, transformers, wandb, etc.) | `git` + `pip` (in `venv`) | + +### General approach + +1. **System packages (root, one-time):** install the build toolchain and, for multi-node, the RDMA/networking libraries via `apt`. The GPU kernel driver must already be loaded. +2. **Python virtual environment (no root):** create a `venv`, then install ROCm and PyTorch from AMD's TheRock multi-arch wheels—this replaces a system ROCm install and keeps everything unprivileged. +3. **Build the accelerated kernels from source** against the ROCm in `venv` and your GPU architecture (`gfx942` for MI300X/MI325X, `gfx950` for MI350X/MI355X). +4. **Install Primus and its Python dependencies**, then persist the required environment variables (ROCm paths, `NVTE_*` flags) in your `venv` activation script. +5. **(Optional) Build the multi-node communication stack** (UCX, OpenMPI, rocSHMEM) only if you require RDMA-based distributed training. + +### Detailed instructions + +Follow the full, step-by-step guide here, which includes the exact pinned versions, environment variables, and automated install scripts: + +- **[Installing the Primus training environment on a host (no Docker)](../../docs_deprecated/install-on-host.md)** + +### Verify + +After the build, validate the environment and run a benchmark directly on the host (no container): + +```bash +./primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096 +``` + +--- + +## Multi-node distributed training (Slurm recommended) + +For training jobs that span **multiple nodes**, we recommend using **[Slurm](https://slurm.schedmd.com/)**. Slurm is a cluster workload manager and job scheduler: it allocates nodes and GPUs, places your job on them, launches one task per node, and injects the topology information (node list, node count, per-node rank) that distributed PyTorch needs. `primus-cli` has a built-in **`slurm` mode** that wraps `srun`/`sbatch` and wires this topology into the training launcher for you. + +### Cluster baseline + +Before launching distributed jobs, ensure every participating node has: + +- The **same software stack**—use the **same container image** on all nodes (recommended), or an identical bare-metal install (see sections above). +- A **shared filesystem** for code, datasets, checkpoints, and logs (e.g. NFS/Lustre), mounted at the same path on every node. +- **Working inter-node networking**: for best performance, use a high-speed RDMA fabric (InfiniBand, RoCE, or AMD AINIC) and ensure RCCL can select the right interface. + +### Setting up Slurm + +Setting up Slurm itself is a cluster-administration task and is outside Primus's scope. Follow the official documentation: + +- [Slurm Quick Start (users)](https://slurm.schedmd.com/quickstart.html) +- [Slurm Quick Start Administrator Guide (install & configure)](https://slurm.schedmd.com/quickstart_admin.html) + +Once `sinfo` and `srun` work on your login node, Primus can submit jobs to it. If you don't administer the cluster, your site administrator typically provides the partition, account, and reservation names you need. + +### Launching with `primus-cli slurm` + +The Slurm wrapper uses a single `--` separator: + +- **Before the `--`**: put the launcher (`srun` or `sbatch`, default `srun`) and Slurm flags (`-N`, `-p`, `--nodelist`, `--account`, `--qos`, `--reservation`, …). +- **After the `--`**: put the Primus command to run (`train` / `benchmark` / …). It runs inside the container image (see [Selecting the container image](#selecting-the-container-image) below). + +```bash +cd /path/to/Primus + +# Pretrain on 2 nodes via srun +./primus-cli slurm srun -N 2 \ + -- train pretrain \ + --config examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml +``` + +Add the global `--dry-run` flag before the launcher to print the exact command without executing it: + +```bash +./primus-cli --dry-run slurm srun -N 2 \ + -- train pretrain --config .yaml +``` + +> See `runner/README.md` and the [CLI reference](../02-user-guide/cli-reference.md) for the full set of launcher flags and scenarios. + +#### Selecting the container image + +The image used for the job is resolved according to the following priority order (highest first): + +1. **`DOCKER_IMAGE` environment variable**—overrides everything else. This is the simplest way to switch images and it propagates to all nodes: + +```bash +export DOCKER_IMAGE=rocm/primus:v26.3 +./primus-cli slurm srun -N 2 \ + -- train pretrain --config .yaml +``` + +2. **`--image` CLI flag**—place it immediately after the `--`, before the Primus command (ignored if `DOCKER_IMAGE` is set): + +```bash +./primus-cli slurm srun -N 2 \ + -- --image rocm/primus:v26.3 train pretrain --config .yaml +``` + +3. **Config file default**—`container.options.image` in `runner/.primus.yaml` (or your `~/.primus.yaml`), which is set to `rocm/primus:v26.3` by default. + +### Distributed environment variables + +Primus launches training with `torchrun`, which needs to know the cluster topology. These are the key variables: + + +| Variable | Role | Default | +| --------------- | -------------------------------------------------------- | ----------- | +| `MASTER_ADDR` | Hostname/IP of rank 0; all ranks rendezvous here. | `localhost` | +| `MASTER_PORT` | Port on the master used for rendezvous. | `1234` | +| `NNODES` | Number of nodes in the job. | `1` | +| `NODE_RANK` | Index of this node (0-based, unique per node). | `0` | +| `GPUS_PER_NODE` | GPUs (processes) to launch per node. | `8` | + +The total number of training processes (world size) is `NNODES × GPUS_PER_NODE`. + +**Under Slurm, the values of these variables are derived automatically.** `primus-cli slurm` reads Slurm's own variables and sets the Primus ones for you: `NNODES` from `SLURM_NNODES`, `NODE_RANK` from `SLURM_NODEID`, and `MASTER_ADDR` from the first host in `SLURM_NODELIST` (with `MASTER_PORT` defaulting to `1234`). You normally only need to set `GPUS_PER_NODE` if your nodes don't have 8 GPUs. You can still override `MASTER_PORT` (e.g. to avoid a port clash) via `--env`. + +### Without Slurm: Kubernetes or `parallel-ssh` + +Slurm is recommended but not required. The mechanism underneath is simple: **set the same distributed environment variables on every node, point them all at the same `MASTER_ADDR`, give each node a unique `NODE_RANK`, and run the same `primus-cli direct` command on each node.** Any tool that can run a command across nodes works—for example Kubernetes (e.g. a `PyTorchJob` / indexed Job) or `parallel-ssh`/`pdsh`. + +For a 2-node job you would run, on the master (rank 0): + +```bash +export NNODES=2 GPUS_PER_NODE=8 NODE_RANK=0 MASTER_ADDR= MASTER_PORT=1234 +./primus-cli direct -- train pretrain --config .yaml +``` + +and on the worker (rank 1), the same command with `NODE_RANK=1` and the same `MASTER_ADDR`: + +```bash +export NNODES=2 GPUS_PER_NODE=8 NODE_RANK=1 MASTER_ADDR= MASTER_PORT=1234 +./primus-cli direct -- train pretrain --config .yaml +``` + +With Kubernetes, inject these as container environment variables (deriving `NODE_RANK` from the pod's index); with `parallel-ssh`, pass them per host. The training command itself is identical on every node. + +### Other important considerations + +- **Cluster validation.** Run the built-in preflight check across your nodes before a long job: `./primus-cli slurm srun -N -- preflight`. +- **Networking/RCCL.** On RDMA fabrics, make sure the correct interface is selected (Primus should auto-detect this; if it doesn't, set `NCCL_SOCKET_IFNAME` / `NCCL_IB_HCA`). Use `NCCL_DEBUG=INFO` (passed via `--env`) to diagnose hangs at startup. +- **RDMA limits.** High-performance networking usually needs locked-memory limits raised (`ulimit -l unlimited`) and sometimes hugepages—configured by your admin. +- **`MASTER_PORT` must be free** on the master node and reachable from all workers; firewalls between nodes will cause rendezvous timeouts. +- **Hugging Face access.** If your configuration downloads gated models or tokenizers, export `HF_TOKEN` (and ensure it is propagated to all nodes and into the containers). + +--- + +## Post-installation verification checklist (for all setup approaches) + + +| Step | Check | +| -------------------- | ------------------------------------------------------------------------------------------------------- | +| **ROCm** | `rocm-smi` shows expected GPUs and no driver errors. | +| **Container engine** | `docker run --rm ... rocm/primus:v26.3` (or your site’s GPU test) succeeds. | +| **GEMM benchmark** | `./primus-cli` **container** or **direct** benchmark completes (see sections above). | +| **Preflight** | Run preflight diagnostics: `./primus-cli direct -- preflight` (single node) or `./primus-cli slurm srun -N -- preflight` (cluster). | + + +If your training pulls models or tokenizers from Hugging Face Hub, configure tokens (for example `HF_TOKEN`) in the environment or container flags as required by your configuration. + +--- + +## Related documentation + +- [Overview](./overview.md) +- [Quickstart](./quickstart.md) +- [CLI reference](../02-user-guide/cli-reference.md) diff --git a/docs/01-getting-started/overview.md b/docs/01-getting-started/overview.md new file mode 100644 index 000000000..7c35fc096 --- /dev/null +++ b/docs/01-getting-started/overview.md @@ -0,0 +1,137 @@ +# Primus overview + +## Executive summary + +**Primus** is a YAML-driven training framework for large-scale foundation model work on AMD GPUs. It targets **machine learning engineers**, **researchers**, and **platform/operations teams** who need reproducible, multi-backend training pipelines on AMD Instinct™ hardware. Within the broader Primus ecosystem, this repository is the training component, sometimes referred to as **Primus-LM** (see [Primus ecosystem](#primus-ecosystem) below). + +**Repository:** [https://github.com/AMD-AGI/Primus](https://github.com/AMD-AGI/Primus) + +--- + +## What Primus provides + +| Area | Description | +|------|-------------| +| **Multi-backend training** | One workflow surface over Megatron-LM, TorchTitan, JAX MaxText, Megatron Bridge, and HummingbirdXT. | +| **Unified CLI** | `primus-cli` with **direct** (bare metal), **container** (Docker/Podman), and **slurm** (cluster) execution modes. | +| **YAML-driven configuration** | Experiment, model, module, and platform presets composed from reusable fragments under `primus/configs/`. | +| **Benchmark suite** | Built-in benchmarks (for example, GEMM) for quick hardware and stack validation. | +| **Preflight diagnostics** | Cluster-oriented checks for host, GPU, and network health before long jobs. | +| **Performance projection** | Tools to estimate memory use and throughput without occupying a full cluster. | + +Workflows span **pretraining** and **post-training** (including SFT and LoRA). Some Megatron configuration files expose RL-related parameters, but reinforcement-learning workflows are outside the scope of this documentation set: they are not part of the tested, supported paths described here, and this documentation does not cover how to run them. + +--- + +## Primus ecosystem + +The training component (Primus-LM) sits between the stability/platform services above it and the low-level operator libraries below it: + +``` + +------------------+ + | Primus-SaFE | + | (stability / | + | cluster mgmt) | + +--------+---------+ + | + +--------v---------+ + | Primus-LM | + | (this repo: | + | training) | + +--------+---------+ + | + +--------v---------+ + | Primus-Turbo | + | (operators / | + | kernels) | + +------------------+ +``` + +- **Primus-SaFE**: Stability and fault-tolerance oriented cluster management referenced by auxiliary tooling (not documented here). +- **Primus-LM**: Training orchestration, backends, CLI, and configurations (maintained in the [AMD-AGI/Primus](https://github.com/AMD-AGI/Primus/) repository). +- **Primus-Turbo**: High-performance operators (for example, FlashAttention-style kernels, GEMM, and collectives). + +--- + +## Supported backends + +| Backend | Typical use | +|---------|-------------| +| **Megatron-LM** | Broadest model coverage; default for many GPT-style and MoE recipes. | +| **TorchTitan** | PyTorch-native large-model training paths. | +| **JAX MaxText** | JAX/Flax training stacks aligned with MaxText. | +| **Megatron Bridge** | Post-training and bridge workflows on top of Megatron-related stacks. | +| **HummingbirdXT** | Additional integrated training path when enabled by your deployment. | + +Backend choice is expressed in the configuration YAML and resolved through Primus’s adapter layer (see [glossary](./glossary.md)). + +--- + +## Supported hardware and stack + +| Item | Requirement | +|------|----------------| +| **GPUs** | AMD Instinct™ **MI300X**, **MI325X**, **MI355X** | +| **Platform** | **ROCm** (version **≥ 7.0** recommended) | +| **Container image (reference)** | `docker.io/rocm/primus:v26.3` | + +Exact kernel and driver packages should match AMD’s documentation for your GPU SKU and ROCm release. + +--- + +## Key dependencies + +Primus depends on the following categories of software (the list is non-exhaustive): + +| Category | Examples | +|----------|----------| +| **Framework** | PyTorch (Megatron-LM, TorchTitan paths); JAX/Flax (MaxText path). | +| **AMD stack** | ROCm, RCCL, HipBLASLt (GEMM), GPU drivers. | +| **Execution** | Docker or Podman for container mode; Slurm for cluster mode. | +| **Observability & tooling** | **loguru** (logging), **Weights & Biases** (`wandb`) and other optional trackers (see `requirements.txt`). | + +Install specifics are covered in [Installation and setup](./installation.md). + +--- + +## Runtime model + +At a high level, a run follows this pipeline: + +1. **YAML configuration** defines work group, modules, model preset, and overrides. +2. **`primus-cli`** selects execution mode (direct, container, or slurm) and forwards to the runner. +3. **Backend adapter** maps the resolved configuration to the target framework (Megatron-LM, TorchTitan, and so on). +4. **Distributed launch** typically uses **`torchrun`** (or the backend’s equivalent) to start workers across GPUs and nodes. + +For CLI shape and options, see [Quickstart](./quickstart.md) and [CLI reference](../02-user-guide/cli-reference.md). + +--- + +## Repository layout (top level) + +The following table shows the top-level layout of the Primus project repository, [AMD-AGI/Primus](https://github.com/AMD-AGI/Primus/): + +| Path | Role | +|------|------| +| `primus/` | Core library: configurations, runtime, trainers, backend adapters, tools (including preflight). | +| `runner/` | CLI implementation, helpers, hooks, and launch glue. | +| `examples/` | End-to-end example YAML and recipes per backend and GPU SKU. | +| `docs/` | Project documentation (this documentation set). | +| `tests/` | Automated tests. | +| `tools/` | Auxiliary scripts and utilities. | +| `benchmark/` | Benchmark drivers and related assets. | +| `third_party/` | Vendored or submodule dependencies. | + +--- + +## Next steps + +- [Installation and setup](./installation.md): ROCm, Docker, pip, and Slurm setup. +- [Quickstart](./quickstart.md): run a minimal training job in minutes. +- [Glossary](./glossary.md): terms used across Primus documentation. + +--- + +## Licensing + +Primus is distributed under the terms described in the project's `LICENSE` file and `README`. If you encounter differing license references between the `README` and the repository root `LICENSE` file, treat licensing as project-specific: confirm the intended terms with the maintainers and your own compliance process before redistributing. diff --git a/docs/01-getting-started/quickstart.md b/docs/01-getting-started/quickstart.md new file mode 100644 index 000000000..2157e8206 --- /dev/null +++ b/docs/01-getting-started/quickstart.md @@ -0,0 +1,171 @@ +# Quickstart (about five minutes) + +This guide runs a **small Megatron-LM pretraining example** with **mock data** so you can validate the stack without preparing a full dataset. The same sample YAML works across **direct**, **container**, and **Slurm** modes. + +> **Recommended: run this example inside AMD-published training Docker images.** AMD publishes ready-to-run ROCm training images (`rocm/primus` for Megatron-LM and TorchTitan, `rocm/jax-training` for MaxText) with all dependencies and the complete training software stack already installed and validated. Using them means you don't have to build or tune the environment yourself, and—most important for multi-node jobs—**every node runs an identical, tested environment**. That consistency helps avoid version-skew and configuration issues that often occur with per-host installations. Host-based installation is supported, but is recommended only for advanced users (see [Installation and setup](./installation.md)). + +See [Installation and setup](./installation.md) for prerequisites and environment setup. + +--- + +## Prerequisites + +- AMD ROCm drivers (version ≥ 7.0 recommended) +- Docker (version ≥ 24.0) with ROCm support +- ROCm-compatible AMD GPUs (e.g., Instinct MI300 series) +- Proper permissions for Docker and GPU device access + +--- + +## Option 1: Clone the repository and run training in a container (recommended) + +### Step 1: Pull the container image + +Check the AMD published training Docker images: + +- Megatron-LM and TorchTitan backends: +- MaxText backend: + +```bash +# For Megatron-LM and TorchTitan backends +docker pull rocm/primus:v26.3 +# For MaxText backend +docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0 +``` + +### Step 2: Clone the repository + +```bash +git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git +cd Primus +# checkout the branch for the specific release +git checkout release/v26.3 +git submodule update --init --recursive +``` + +### Step 3: Run training inside container + +Run the training from the repository root. If your configuration downloads weights or tokenizers from Hugging Face Hub, pass `HF_TOKEN` into the container: + +```bash +./primus-cli container --image rocm/primus:v26.3 \ + --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \ + -- train pretrain \ + --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +``` + +--- + +## Option 2: Install Primus from a wheel and run training in a container + +### Step 1: Install Primus as a Python package + +Install Primus in a virtual environment: + +```bash +python -m venv primus-env +source primus-env/bin/activate +pip install "primus==26.3.1" --no-deps --extra-index-url https://amd-agi.github.io/Primus/simple/ +``` + +> **Note:** This installs only the Primus CLI into your virtual environment (under `site-packages`), without other dependencies. Third-party submodules are downloaded on the first run of the container, and the complete training software stack is provided in the AMD-published Docker images. You can launch `primus-cli` from any directory. + +### Step 2: Run training in a container using the pip-installed Primus + +```bash +primus-cli container --image rocm/primus:v26.3 \ + --env HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" \ + --volume /path/to/your/data:/data -- --log_file /data/run.log \ + -- train pretrain --config /data/your/config.yaml +``` + +> **Note:** `--volume` mounts a local data directory into the container. `--log_file` writes the training log there; if omitted, logs go to the Primus install directory (`site-packages/primus/logs` by default). + +--- + +## Expected output + +You should see the backend initialize distributed processes, load the training configuration, and emit **iteration-level logs** (with loss, throughput, step index, etc.). Exact fields depend on the backend and logging configuration; a typical pattern resembles: + +``` +... [INFO] starting training ... +... iteration 1 | loss: 10.xxx | ... +... iteration 2 | loss: 9.xxx | ... +``` + +Let the job run briefly to confirm stability; stop with `Ctrl+C` when satisfied. + +--- + +## Same configuration, three execution modes + +Use one configuration file: `examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml`. + +| Mode | Example command | +|------|------------------| +| **Container** | `./primus-cli container --image rocm/primus:v26.3 -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml` | +| **Direct** | `./primus-cli direct -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml` | +| **Slurm** | `./primus-cli slurm srun -N ... -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml` | + +Replace `` and Slurm resource flags with values appropriate for your cluster. **Container** and **Slurm** runs execute inside a Docker image (Slurm dispatches through the container launcher on each node); **Direct** runs execute inside a Docker container (you start the container yourself and run the command) or directly on the host if the training environment is already set up. + +> **Multi-node networking:** Primus auto-detects RDMA settings (`NCCL_IB_HCA`, `NCCL_SOCKET_IFNAME`, …) on each node. If auto-detection selects the wrong NIC, or your fabric needs specific values (RoCE `NCCL_IB_GID_INDEX`, AMD AINIC, etc.), override them via `--env` or your config. See [Multi-node networking](../04-technical-guides/multi-node-networking.md). + +### Selecting the container image + +For container and Slurm runs, Primus resolves which Docker image to use in the following order (**highest priority first**): + +1. **`DOCKER_IMAGE` environment variable**—if set, it overrides every other source (including `--image`): + +```bash +export DOCKER_IMAGE=rocm/primus:v26.3 +./primus-cli container -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +``` + +2. **`--image` command-line argument**—the usual per-run override, passed as a container mode argument (before `--`): + +```bash +./primus-cli container --image rocm/primus:v26.3 -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +``` + +3. **`image` field in a config file** (`container.options.image`)—used when neither of the above is set: + +```yaml +container: + options: + image: "rocm/primus:v26.3" +``` + +Primus loads a **single** config file—the first that exists among `--config `, then `~/.primus.yaml`, then the shipped `runner/.primus.yaml` (these files are **not** merged). Because `runner/.primus.yaml` ships with a default image, a bare `./primus-cli container -- ...` works out of the box. + +>**Check the logs to make sure actual image being used is the one you wanted.** + + +--- + +## Command structure + +`primus-cli` parses **global options**, a **mode** (`direct`, `container`, `slurm`, …), optional **mode-specific arguments**, then a **`--` separator** followed by the **subcommand and its arguments** (for example `train` or `benchmark`). + +``` +primus-cli [global-options] [mode-args] -- [command-args...] +``` + +Example: + +```text +primus-cli container --image rocm/primus:v26.3 -- train pretrain --config path/to/experiment.yaml + ^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^^^^^^ ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + mode mode args command + args +``` + +--- + +## Next steps + +| Topic | Document | +|-------|----------| +| Full CLI flags and subcommands | [CLI reference](../02-user-guide/cli-reference.md) | +| YAML presets, overrides, and composition | [Configuration system](../02-user-guide/configuration-system.md) | +| Pretraining workflows and backend notes | [Pretraining workflows](../02-user-guide/pretraining.md) | +| Terminology | [Glossary](./glossary.md) | diff --git a/docs/02-user-guide/README.md b/docs/02-user-guide/README.md new file mode 100644 index 000000000..4a89c2c42 --- /dev/null +++ b/docs/02-user-guide/README.md @@ -0,0 +1,18 @@ +# User guide + +Core workflows and day-to-day usage. + +- [Primus tools](primus-tools.md): start here—an at-a-glance catalog of all Primus tools and ecosystem projects with how-to starting points +- [CLI reference](cli-reference.md): `primus-cli` modes, flags, and subcommands +- [Configuration system](configuration-system.md): YAML configuration model, presets, overrides, inheritance +- [Pretraining](pretraining.md): pretraining **concepts**: backends, YAML structure, parallelism, configuration inventory +- [Backend training recipes](training-recipes.md): pretraining **commands**: copy-paste, GPU-arch-specific run commands +- [Post-training](posttraining.md): SFT and LoRA fine-tuning via Megatron Bridge +- [Benchmarking](benchmarking.md): GEMM, RCCL, and dense-GEMM benchmark suites +- [Preflight](preflight.md): cluster diagnostics and environment validation +- [Projection](projection.md): memory and performance projection tools +- [Tuning agent](tuning-agent.md): LLM-driven search for an optimal training configuration (uses projection as an oracle) + +--- + +[← Documentation home](../README.md) diff --git a/docs/02-user-guide/benchmarking.md b/docs/02-user-guide/benchmarking.md new file mode 100644 index 000000000..77fd309ea --- /dev/null +++ b/docs/02-user-guide/benchmarking.md @@ -0,0 +1,204 @@ +# Benchmark suite + +Primus ships microbenchmarks for GPU compute and distributed communication. They are exposed as the `benchmark` subcommand of the Primus CLI. Use them to sanity-check a node or cluster before long training jobs. + +**Implementation:** `primus/cli/subcommands/benchmark.py` (initializes distributed execution, runs the selected suite, then finalizes). + +Related documentation: [Preflight diagnostics](./preflight.md) (broader cluster checks), [Memory and performance projection](./projection.md) (training-scale estimates), [Installation](../01-getting-started/installation.md) (environment setup). + +--- + +## Overview and command syntax + +```bash +primus-cli [global-options] [mode-args] -- benchmark [suite-specific-args] +``` + +- **``** is typically `direct`, `container`, or `slurm` so that `WORLD_SIZE`, `RANK`, `MASTER_ADDR`, and related variables are set consistently. +- **`benchmark`** runs inside the Primus Python CLI; the runner wires up the process environment the same way as training. + +The CLI also registers an `attention` suite; the subsections below cover **`gemm`**, **`gemm-dense`**, **`gemm-deepseek`**, **`strided-allgather`**, and **`rccl`**. + +--- + +## Quick start + +Single-node GEMM: + +```bash +primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096 --dtype bf16 --duration 10 +``` + +Multi-node RCCL on Slurm: + +```bash +primus-cli slurm srun -N 4 -- benchmark rccl --op all_reduce --min-bytes 1M --max-bytes 128M +``` + +--- + +## Suite reference + +### `gemm` + +Single-shape general matrix multiply (GEMM) microbenchmark. + +| Argument | Description | +|----------|-------------| +| `--M`, `--N`, `--K` | Matrix dimensions (defaults: 4096 / 4096 / 4096). | +| `--trans_a` | Transpose the A matrix. | +| `--trans_b` | Transpose the B matrix. | +| `--dtype` | `bf16`, `fp16`, `fp32`, or `fp8` (`fp8` requires torchao). Default: `bf16`. | +| `--duration` | Run duration in seconds (default: 10). | +| `--output-file` | Destination for results (`.md`, `.csv`, `.tsv`, `.jsonl`, `.jsonl.gz`). Default: `./gemm_report.md`. Use `-` or omit for Markdown on stdout. | + +**Example** + +```bash +primus-cli direct -- benchmark gemm --M 8192 --N 8192 --K 8192 --dtype bf16 --duration 10 --output-file ./gemm_report.md +``` + +--- + +### `gemm-dense` + +Dense GEMM workload using Llama-like shape parameters (model-derived GEMMs). + +| Argument | Description | +|----------|-------------| +| `--model` | Optional label (for example `Llama3.1_8B`). | +| `--seqlen` | Sequence length (default: 2048). | +| `--hidden-size` | Hidden size (default: 4096). | +| `--intermediate-size` | FFN intermediate size (default: 11008). | +| `--num-attention-heads` | Attention heads (default: 32). | +| `--num-key-value-heads` | KV heads (default: 32). | +| `--head-dim` | Per-head dimension (default: 128). | +| `--vocab-size` | Vocabulary size (default: 32000). | +| `--dtype` | `bf16`, `fp16`, `fp32`, or `fp8` (`fp8` requires torchao). Default: `bf16`. | +| `--mbs` | Microbatch size (default: 1). | +| `--duration` | Seconds per shape (default: 3). | +| `--output-file` | Report path (default: `./gemm-dense_report.md`). | + +**Example** + +```bash +primus-cli direct -- benchmark gemm-dense --model Llama3.1_8B --seqlen 4096 --dtype bf16 +``` + +--- + +### `gemm-deepseek` + +Dense GEMM workload using DeepSeek-style shapes (MoE / MLA-related dimensions). + +| Argument | Description | +|----------|-------------| +| `--model` | Label (for example `Deepseek_V2`, `Deepseek_V3`). | +| `--seqlen` | Sequence length (default: 4096). | +| `--hidden-size` | Hidden size (default: 4096). | +| `--intermediate-size` | Dense FFN intermediate (default: 12288). | +| `--kv-lora-rank` | KV LoRA rank (default: 512). | +| `--moe-intermediate-size` | MoE expert intermediate (default: 1536). | +| `--num-attention-heads` | Attention heads (default: 64). | +| `--num-experts-per-tok` | Experts per token (default: 6). | +| `--n-routed-experts` | Number of routed experts (default: 128). | +| `--n-shared-experts` | Shared experts (default: 2). | +| `--q-lora-rank` | Optional Q LoRA rank. | +| `--qk-nope-head-dim`, `--qk-rope-head-dim`, `--v-head-dim` | Head dimensions for MLA-style attention (defaults: 128 / 64 / 128). | +| `--vocab-size` | Vocabulary size (default: 128256). | +| `--dtype` | `bf16` or `fp16` (default: `bf16`). | +| `--mbs` | Microbatch size (default: 1). | +| `--duration` | Seconds per shape (default: 3). | +| `--output-file` | Report path (default: `./gemm-deepseek_report.md`). | +| `--append` | Append to an existing report instead of overwriting. | + +**Example** + +```bash +primus-cli direct -- benchmark gemm-deepseek --model Deepseek_V3 --dtype bf16 --append +``` + +--- + +### `strided-allgather` + +Strided all-gather microbenchmark (useful for multi-rank communication patterns). + +| Argument | Description | +|----------|-------------| +| `--sizes-mb` | Comma-separated message sizes in MB per rank (default: `64,128,256`). | +| `--stride` | Rank stride for group formation (default: 8). | +| `--parallel` | Run multiple groups’ all-gathers in parallel. | +| `--iters` | Timed iterations per size (default: 50). | +| `--warmup` | Warmup iterations per size (default: 10). | +| `--dtype` | `fp16`, `bf16`, or `fp32` (default: `bf16`). | +| `--backend` | `nccl`, `gloo`, or `mpi` (default: `nccl`). | + +**Example** + +```bash +primus-cli slurm srun -N 2 -- benchmark strided-allgather --sizes-mb 64,128 --stride 8 --iters 50 +``` + +--- + +### `rccl` + +RCCL collective benchmark: sweeps message sizes and reports bandwidth and latency statistics. + +| Argument | Description | +|----------|-------------| +| `--op` | One or more of: `all_reduce`, `broadcast`, `reduce_scatter`, `all_gather`, `alltoall` (default: `all_reduce`). | +| `--sizes` | Explicit size list (for example `1K,2K,4K,8K,1M`). Overrides generated sweep. | +| `--min-bytes` | Minimum message size for generated sweep (default: `1K`). | +| `--max-bytes` | Maximum message size (default: `128M`). | +| `--num-sizes` | Number of points in generated sweep (default: 12). | +| `--scale` | `log2` or `linear` for generated sweeps (default: `log2`). | +| `--dtype` | `bf16`, `fp16`, or `fp32` (default: `bf16`). | +| `--warmup` | Warmup iterations (default: 20). | +| `--iters` | Timed iterations (default: 100). | +| `--repeat` | Repeat each `(op, size)` for stability (default: 1). | +| `--aggregate-repeat` | Emit an extra summary row aggregating repeat runs. | +| `--check` | Enable lightweight correctness checks. | +| `--output-file` | Report path (`.md`, `.csv`, `.tsv`, `.jsonl`, `.jsonl.gz`; default: `./rccl_report.md`). | +| `--append` | Append instead of overwrite. | +| `--per-rank` | Per-rank summary lines. | +| `--per-rank-file` | Path for per-rank stats (if empty, derived from `--output-file` with `_rank` suffix). | +| `--per-iter-trace` | Emit per-iteration trace (can be large). | +| `--trace-file` | Trace output path (if empty, derived from `--output-file`). | +| `--trace-limit` | Max iterations to record per `(op, size)`; `0` means all. | +| `--trace-ops` | Comma-separated ops to include in trace (empty = all). | +| `--trace-sizes` | Comma-separated sizes to include in trace (empty = all). | +| `--cluster` | Label for the report preamble. Defaults to `$PRIMUS_CLUSTER`, falling back to a built-in placeholder (`amd-aig-poolside`) when it is unset—set `PRIMUS_CLUSTER` or pass `--cluster` to record your own cluster name. | + +**Example** + +```bash +primus-cli slurm srun -N 4 -- benchmark rccl --op all_reduce --min-bytes 1M --max-bytes 128M --dtype bf16 +``` + +--- + +## Understanding results + +- **GEMM suites** emit throughput-oriented metrics suitable for comparing dtypes, shapes, and durations across runs. Keep `duration` long enough to smooth variance on shared clusters. +- **`rccl`** reports collective latency and bandwidth across a size sweep; use it to verify inter-node behavior and to compare against expected NIC bandwidth. +- **Markdown / CSV / TSV / JSONL** output formats support post-processing in notebooks or CI; gzip JSONL is supported for large traces. + +--- + +## Tips + +1. **Distributed initialization:** If jobs hang or report uninitialized distributed state, launch through `primus-cli` (`direct` / `container` / `slurm`) rather than calling Python entrypoints manually without the right environment. +2. **Paths:** Prefer **absolute** paths for `--output-file` when using containers or Slurm so the working directory matches your expectations. +3. **Multi-node:** Use your scheduler integration (`primus-cli slurm …`) so rank and address assignment matches your cluster. +4. **Full cluster validation:** Combine targeted `benchmark` runs with [Preflight](./preflight.md) for host, GPU, network, and integrated perf checks. + +--- + +## Related documentation + +- [Preflight diagnostics](./preflight.md) +- [Memory and performance projection](./projection.md) +- [Post-training workflows](./posttraining.md) +- [Installation](../01-getting-started/installation.md) diff --git a/docs/02-user-guide/cli-reference.md b/docs/02-user-guide/cli-reference.md new file mode 100644 index 000000000..b8f708da3 --- /dev/null +++ b/docs/02-user-guide/cli-reference.md @@ -0,0 +1,233 @@ +# CLI reference + +This section describes the unified Primus launcher (`runner/primus-cli`) and how it invokes the Python CLI (`primus/cli/main.py`). For deeper background, see [CLI architecture](../06-developer-guide/cli-architecture.md). + +--- + +## Command structure + +```text +primus-cli [global-options] [mode-args] -- [command] +``` + +- **Global options** go before the mode name and they affect configuration loading and logging for the whole run. +- **Mode** is one of `direct`, `container`, or `slurm`. +- **`--` (required)** separates launcher options from the Primus Python CLI. Everything after the first `--` is passed to `primus/cli/main.py` (or another script if you override it in direct mode). + +From the repository root, invoke the launcher as `./runner/primus-cli` (or install/link it as `primus-cli` on your `PATH`). + +--- + +## Global options + +These flags are parsed in `runner/primus-cli` before the mode name is read and passed on to `runner/primus-cli-.sh`. + +| Option | Description | +| --- | --- | +| `--config FILE` | Load a YAML file for launcher defaults (see [Configuration precedence](#configuration-precedence-launcher-yaml)). | +| `--debug` | Verbose logging; sets `PRIMUS_LOG_LEVEL=DEBUG`. | +| `--dry-run` | Print the command that would run and exit without executing the mode script. | +| `--version` | Print the CLI version and exit. | +| `-h`, `--help` | Show top-level usage and exit. | + +Mode-specific help: + +```bash +./runner/primus-cli direct --help +./runner/primus-cli container --help +./runner/primus-cli slurm --help +``` + +Primus Python CLI help (after `--`): + +```bash +./runner/primus-cli direct -- --help +./runner/primus-cli direct -- train --help +./runner/primus-cli direct -- benchmark --help +``` + +--- + +## Direct mode + +Run training, benchmarks, or diagnostics on the current host (or inside an environment you already prepared). GPU-specific tuning is applied via `runner/helpers/envs/.sh` when present. + +### Syntax + +```bash +primus-cli direct [options] -- +``` + +### Options + +| Option | Description | +| --- | --- | +| `--config FILE` | Launcher YAML (same resolution as [global `--config`](#configuration-precedence-launcher-yaml)). | +| `--debug` | Debug logging for the direct launcher. | +| `--dry-run` | Show the resolved command that would be launched without running training. | +| `--single` | Run with `python3` instead of `torchrun` (single process). | +| `--script PATH` | Python entry script (default: `primus/cli/main.py`). | +| `--env KEY=VALUE` | Set an environment variable before launch (repeatable). A path without `=` is treated as an env file (`--env_file`), loaded later in the launch sequence. | +| `--patch script.sh` | Run a shell snippet before the main script (repeatable). | +| `--log_file PATH` | Redirect logs to a file. | +| `--numa` | Force NUMA binding on. | +| `--no-numa` | Force NUMA binding off. | + +### Distributed environment variables + +For multi-node or multi-process runs, set these via `export` or `--env`: + +| Variable | Role | Typical default | +| --- | --- | --- | +| `NNODES` | Number of nodes | `1` | +| `NODE_RANK` | Rank of this node | `0` | +| `GPUS_PER_NODE` | GPUs per node | `8` (see `runner/.primus.yaml` `direct.gpus_per_node`) | +| `MASTER_ADDR` | Hostname or IP of rank 0 | `localhost` | +| `MASTER_PORT` | TCP port for the process group | `1234` | + +--- + +## Container mode + +Run the same Python CLI inside Docker or Podman with ROCm-oriented defaults from `runner/.primus.yaml`. + +### Syntax + +```bash +primus-cli container [options] -- +``` + +### Common options + +| Option | Description | +| --- | --- | +| `--image NAME` | Image tag (default from config: `rocm/primus:v26.3`). | +| `--volume HOST[:CONTAINER]` | Bind mount (repeatable). | +| `--env KEY=VALUE` | Pass into the **inner** `primus-cli direct` as `--env` (repeatable). | +| `--device PATH` | Extra device nodes (repeatable; defaults include GPU/RDMA devices). | +| `--name`, `--user`, `--network`, `--ipc` | Standard container runtime options. | +| `--clean` | Remove all containers before launch. | +| `--cpus N` | CPU limit. | +| `--memory SIZE` | Memory limit (e.g. `128G`). | +| `--shm-size SIZE` | Shared memory size. | +| `--gpus N` | GPU limit (when using a runtime that supports this flag). | + +### Auto-mounted devices + +When using `runner/.primus.yaml`, the default container section includes: + +- `/dev/kfd`—ROCm kernel fusion driver +- `/dev/dri`—GPU render nodes +- `/dev/infiniband`—InfiniBand character devices (when present) + +### Environment forwarding + +`container.options.env` in `runner/.primus.yaml` lists **names** that are forwarded into the container as inner `--env` arguments when the variable is set in the host environment (for example `MASTER_ADDR`, `HF_TOKEN`, `NCCL_SOCKET_IFNAME`). The container script also auto-forwards host variables whose names start with `PRIMUS_`, `NCCL_`, `RCCL_`, `GLOO_`, `IONIC_`, or `HIPBLASLT_` when not already listed. + +--- + +## Slurm mode + +Launch distributed jobs with `srun` or `sbatch`. The Slurm launcher builds `srun` or `sbatch` flags, merges them with `slurm.*` entries from the loaded YAML, then runs `runner/primus-cli-slurm-entry.sh` on allocated nodes. + +### Syntax + +```text +primus-cli slurm [--config FILE] [--debug] [--dry-run] [srun|sbatch] [SLURM_FLAGS...] -- +``` + +| Part | Meaning | +| --- | --- | +| First `--` | Separates Slurm launcher flags from the Primus Python CLI command (for example `train pretrain ...`). | +| Default launcher | If you omit `srun` and `sbatch`, **`srun` is used** (`LAUNCH_CMD` in `runner/primus-cli-slurm.sh`). | + +### Examples + +```bash +# Interactive multi-node training +./runner/primus-cli slurm srun -N 4 -p gpu -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml + +# Batch job +./runner/primus-cli slurm sbatch -N 8 -t 8:00:00 -o train.log -- train pretrain --config exp.yaml +``` + +On each node, `primus-cli-slurm-entry.sh` sets `NNODES`, `NODE_RANK`, `GPUS_PER_NODE`, `MASTER_ADDR`, and `MASTER_PORT` from Slurm and invokes `primus-cli-container.sh` with matching `--env` injections (see `runner/primus-cli-slurm-entry.sh`). Container options such as `--image` should come from `runner/.primus.yaml` or the launcher config file rather than appearing as an inner `container` command after the Slurm separator. + +--- + +## Python subcommands (after `--`) + +These run under `primus/cli/main.py` unless you change `--script` in direct mode. + +| Subcommand | Purpose | +| --- | --- | +| `train pretrain --config ` | Pretraining (Megatron-LM, TorchTitan, MaxText, Megatron Bridge, etc., per configuration YAML). | +| `train posttrain --config ` | Post-training (SFT or LoRA-style workflows; same top-level flags as pretrain in the parser). | +| `benchmark [args]` | Performance microbenchmarks (see table below). | +| `preflight [--host] [--gpu] [--network] [--perf-test]` | Cluster and node diagnostics. | +| `projection memory --config ` | Memory estimation from a merged config. | +| `projection performance --config ` | Performance projection from a merged config. | +| `projection both --config ` | Single benchmark → both performance and memory projections (cluster sizing). | + +### Benchmark suites + +Implemented in `primus/cli/subcommands/benchmark.py`: + +| Suite | Notes | +| --- | --- | +| `gemm` | General GEMM microbenchmark. | +| `gemm-dense` | Dense GEMM variant. | +| `gemm-deepseek` | DeepSeek-style dense GEMM. | +| `strided-allgather` | Communication microbenchmark. | +| `rccl` | RCCL collective microbenchmark. | + +The same file also registers an `attention` suite for attention microbenchmarks. + +--- + +## Configuration precedence (launcher YAML) + +Resolution is implemented in `runner/lib/config.sh` (functions `resolve_config_file` and `load_config_auto`): + +1. **`--config FILE`** on the command line (if given). +2. **`~/.primus.yaml`** if it exists. +3. **`runner/.primus.yaml`** (system default). + +Within a chosen file, nested keys follow normal YAML structure. Slurm and container scripts merge CLI flags with their sections so that **explicit CLI arguments override file values** where applicable. + +**Note:** This precedence applies to the **shell launcher** YAML. Training YAML merge order for configurations is documented in [Configuration system](configuration-system.md). + +--- + +## Common examples + +| Goal | Example | +| --- | --- | +| Direct pretrain | `./runner/primus-cli direct -- train pretrain --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml` | +| Direct GEMM | `./runner/primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096` | +| Container pretrain | `./runner/primus-cli container --volume /data:/data -- train pretrain --config /data/exp.yaml` | +| Slurm training | `./runner/primus-cli slurm srun -N 4 -- train pretrain --config exp.yaml` | +| Preflight (fast) | `./runner/primus-cli slurm srun -N 4 -- preflight --host --gpu --network` | +| Inspect launch command | `./runner/primus-cli --dry-run direct -- train pretrain --config exp.yaml` | +| Dry-run Slurm | `./runner/primus-cli --dry-run slurm srun -N 2 -- train pretrain --config exp.yaml` | + +--- + +## Exit codes + +From `runner/primus-cli`: + +| Code | Meaning | +| --- | --- | +| 0 | Success | +| 1 | Library or dependency failure | +| 2 | Invalid arguments or configuration | +| 3 | Runtime execution failure | + +--- + +## Related documentation + +- [Getting started: Quickstart](../01-getting-started/quickstart.md): installation and first steps +- [Configuration system](configuration-system.md): YAML configuration model, presets, overrides, inheritance +- [Pretraining](pretraining.md): pretraining workflows and backend notes diff --git a/docs/02-user-guide/configuration-system.md b/docs/02-user-guide/configuration-system.md new file mode 100644 index 000000000..90b8cb6fe --- /dev/null +++ b/docs/02-user-guide/configuration-system.md @@ -0,0 +1,170 @@ +# Configuration system + +Primus experiments are described in YAML. The loader resolves **environment variables**, **`extends:` inheritance**, and **module/model/platform presets** before training starts. This document focuses on the Python configuration pipeline (`primus/core/config/` and `primus/core/launcher/parser.py`). + +**Related documentation** + +| Topic | Location | +| --- | --- | +| CLI launcher and `--config` | [CLI Reference](cli-reference.md) | +| Backend parameter references | [Megatron parameters](../03-configuration-reference/megatron-parameters.md), [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md), [MaxText parameters](../03-configuration-reference/maxtext-parameters.md) | + +--- + +## Overview: Three-layer YAML + +A typical experiment ties together: + +1. **Experiment YAML**—your run: identity, workspace, and a `modules` section naming framework presets plus overrides. +2. **Module preset**—training defaults for a backend (optimizer, schedule, parallelism hooks) under `primus/configs/modules//`. +3. **Model preset**—architecture and tokenizer metadata under `primus/configs/models//`. + +All of these are **deep-merged** (see `primus/core/config/yaml_loader.py` and `primus/core/config/merge_utils.py`). A **platform preset** (`primus/configs/platforms/`) maps distributed environment variable names and logging defaults; if omitted, the parser injects `platform_azure.yaml` (see `PrimusParser.parse_platform` in `primus/core/launcher/parser.py`). + +--- + +## Experiment config structure + +```yaml +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:my_experiment} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron # backend name (megatron, torchtitan, maxtext, megatron_bridge, …) + config: pre_trainer.yaml # module preset file under primus/configs/modules// + model: llama3_8B.yaml # model preset file under primus/configs/models// + overrides: # training overrides (deep-merged last) + train_iters: 50 + micro_batch_size: 4 +``` + +Required top-level keys are validated in `PrimusParser.parse_meta_info`: `work_group`, `user_name`, `exp_name`, `workspace`. + +--- + +## Module presets + +- **Location:** `primus/configs/modules//` (for example `primus/configs/modules/megatron/pre_trainer.yaml`). +- **Purpose:** Default training behavior (iterations, batching, optimizer, logging, parallelism-related flags for that backend). +- **Inheritance:** Use `extends:` to compose files in the same directory (or relative paths). Multiple entries merge in order; the **current file wins** on conflicts (`_apply_extends` in `yaml_loader.py`). + +Example chain (excerpt): `pre_trainer.yaml` extends `trainer_base.yaml`, which extends `../module_base.yaml` and other shared fragments (`primus/configs/modules/megatron/trainer_base.yaml`). + +--- + +## Model presets + +- **Location:** `primus/configs/models//` (for example `primus/configs/models/megatron/llama3_8B.yaml`). +- **Purpose:** Architecture dimensions, tokenizer identifiers, and other model metadata. +- **Inheritance:** Same `extends:` mechanism as modules (for example `llama3_8B.yaml` → `llama3_base.yaml` → …). + +--- + +## Platform presets + +- **Location:** `primus/configs/platforms/` (for example `primus/configs/platforms/platform_azure.yaml`). +- **Purpose:** Names of environment variables used for distributed launch (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, …) and defaults such as `master_sink_level` and `workspace`. +- **Default:** If the experiment omits `platform`, the parser sets `config: platform_azure.yaml` (`primus/core/launcher/parser.py`). + +--- + +## Environment variable substitution + +`primus/core/config/yaml_loader.py` expands: + +| Pattern | Behavior | +| --- | --- | +| `${VAR}` | **Required.** Raises if `VAR` is unset. | +| `${VAR:default}` | Uses `default` when `VAR` is unset. | + +After substitution, purely numeric strings may be converted to `int` or `float`. + +--- + +## `extends:` inheritance + +For each YAML file: + +1. Each path in `extends:` is resolved **relative to the directory of the current file**. +2. Presets are loaded recursively (each may have its own `extends:`). +3. Merge order: earlier presets in the list are merged first; **later presets override earlier ones**; the **current file overrides all** (`_apply_extends`). + +`PresetLoader.load` (`primus/core/config/preset_loader.py`) resolves `primus/configs///.yaml` and runs the same `parse_yaml` pipeline. + +--- + +## CLI overrides (training) + +After the main arguments are parsed, unknown tokens are interpreted as **key=value overrides** and deep-merged into the active training module namespace (`module_cfg.params`). The core runtime applies them in `PrimusRuntime._apply_overrides` (`primus/core/runtime/train_runtime.py`) using `parse_cli_overrides` (`primus/core/utils/arg_utils.py`) followed by `deep_merge`. Both `key=value` and `--key value` forms are accepted. Unknown keys are merged in (not rejected) and forwarded to the backend; the stricter key-existence check in `parse_args` / `_check_keys_exist` (`primus/core/launcher/parser.py`) belongs to a legacy path that the `train` subcommand does not exercise. + +Example (conceptual): + +```bash +./runner/primus-cli direct -- train pretrain --config exp.yaml \ + --train_iters 100 --micro_batch_size 2 +``` + +--- + +## Merge priority (training config) + +The effective ordering of training parameters is: + +1. **CLI overrides** (key=value after the main `train` arguments)—highest. Applied last by the runtime (`PrimusRuntime._apply_overrides`), after preset and experiment merging. +2. **`modules.pre_trainer.overrides`** in the experiment YAML (applied in `PrimusParser.parse_trainer_module`). +3. **Module preset with model preset additions**: the module preset is loaded first, then the model preset is merged in with `allow_override=False`, so duplicate top-level module keys are preserved while non-duplicate model keys are added (`merge_namespace` in `parse_trainer_module`). +4. **Preset chains** via `extends:` inside those files—base layers first, specialized layers later, file body last. + +A concise mental model: + +**CLI overrides > experiment `overrides` > module preset with model preset additions > each preset's own `extends:` chain > shared bases such as `module_base.yaml`.** + +--- + +## Config resolution walkthrough: `llama2_7B-BF16-pretrain.yaml` + +Example experiment: `examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml`. + +1. **Load experiment**—`parse_yaml` reads the file; `${PRIMUS_*:…}` placeholders resolve. +2. **Meta**—`work_group`, `user_name`, `exp_name`, `workspace` are checked. +3. **Platform**—If not present, default `platform_azure.yaml` loads from `primus/configs/platforms/`. +4. **Module preset**—`PresetLoader.load("pre_trainer.yaml", "megatron", "modules")` loads `primus/configs/modules/megatron/pre_trainer.yaml` and applies its `extends:` chain (for example `trainer_base.yaml` → …). +5. **Model preset**—`PresetLoader.load("llama2_7B.yaml", "megatron", "models")` loads `primus/configs/models/megatron/llama2_7B.yaml` (which extends `llama2_base.yaml` → `llama_base.yaml` → …). +6. **Merge**—Module and model namespaces are merged for `pre_trainer`. +7. **Experiment overrides**—Keys under `modules.pre_trainer.overrides` in the example file (for example `mock_data: true`, parallelism, LR) are applied on top. +8. **CLI overrides**—Any key=value pairs from the command line are merged last. + +--- + +## How to write a new config for a new model + +1. **Add or reuse a model preset** under `primus/configs/models//`, using `extends:` from the closest existing architecture (for example copy `llama3_8B.yaml` and adjust hidden size, layers, tokenizer). +2. **Point an experiment YAML at it**—set `modules.pre_trainer.framework`, `config: .yaml`, and `model: .yaml`. +3. **Set overrides** in the experiment file for run-specific values (batch sizes, paths, `mock_data`, parallelism). Prefer small experiment files that reference presets instead of duplicating hundreds of keys. +4. **Validate** with `--dry-run` and by tracing the referenced experiment, module, and model presets (see below). +5. **Optional:** add an example under `examples//configs//` for others to copy. + +--- + +## Debugging config issues + +| Technique | What it does | +| --- | --- | +| `--export_config PATH` | Parsed by the training config parser, but the default core `PrimusRuntime` path does not currently write the resolved YAML. Treat this as legacy/future functionality unless your deployment implements it. | +| `./runner/primus-cli --dry-run …` | Shows the launcher command without executing (shell layer). | +| `--debug` | Enables verbose logging for launcher and Python (`PRIMUS_LOG_LEVEL=DEBUG`). | +| Inspect presets | Open the resolved `extends:` chain under `primus/configs/modules/` and `primus/configs/models/` for the framework you use. | + +If `${VAR}` substitution fails, set the variable or switch to `${VAR:default}` in the YAML. + +--- + +## Cross-references + +- Default launcher YAML: `runner/.primus.yaml` +- YAML loader (env + extends): `primus/core/config/yaml_loader.py` +- Parser and merge: `primus/core/launcher/parser.py` +- Preset paths: `primus/core/config/preset_loader.py` diff --git a/docs/02-user-guide/posttraining.md b/docs/02-user-guide/posttraining.md new file mode 100644 index 000000000..03367e412 --- /dev/null +++ b/docs/02-user-guide/posttraining.md @@ -0,0 +1,216 @@ +# Post-training workflows + +Post-training (supervised fine-tuning) adapts a pre-trained foundation model to new tasks or domains. In Primus, post-training runs through the **Megatron Bridge** backend using the `train posttrain` subcommand. Example YAML configurations live under `examples/megatron_bridge/configs/` in the [Primus repository](https://github.com/AMD-AGI/Primus). + +For YAML field details, see [Megatron Bridge parameters](../03-configuration-reference/megatron-bridge-parameters.md). For related tooling, see [Benchmark suite](./benchmarking.md), [Preflight diagnostics](./preflight.md), [Memory and performance projection](./projection.md). + +--- + +## Overview: SFT vs LoRA + +The following table shows how SFT (Supervised Fine-Tuning) and LoRA (Low-Rank Adaptation) differ in different aspects. + +| Aspect | SFT (full fine-tuning) | LoRA (parameter-efficient) | +|--------|------------------------|----------------------------| +| **PEFT setting** | `peft: "none"` | `peft: lora` | +| **What is trained** | All model parameters | Low-rank adapters only | +| **Memory** | Higher | Lower | +| **Throughput** | Typically slower per step | Often faster iteration | +| **Learning rate** | Lower: roughly `5e-6` to `1e-5` | Higher: roughly `1e-4` to `5e-4` | +| **Typical use** | Maximum adaptation when memory allows | Limited GPU memory, many task-specific adapters, rapid experimentation | + +--- + +## Quick start commands + +General form: + +```bash +./primus-cli -- train posttrain --config +``` + +From a clone of the Primus repository, the same entrypoint is often invoked as `./runner/primus-cli`. + +**Prerequisites:** AMD ROCm (recommended ≥ 7.0) and Docker with ROCm support (optional but typical) installed on systems with AMD Instinct™ GPUs (for example MI300X, MI355X). Run this for a quick check on these prerequisites: `rocm-smi && docker --version`. See [Installation and setup](../01-getting-started/installation.md) for the full prerequisites and container setup. + +### Direct mode (bare metal or inside a Docker container) + +```bash +# SFT — example: Qwen3 32B on MI355X +./runner/primus-cli direct -- train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml + +# LoRA — same model family +./runner/primus-cli direct -- train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml +``` + +### Container mode + +```bash +./runner/primus-cli container --image rocm/primus:v26.3 -- \ + train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml +``` + +--- + +## Configuration reference + +These keys are commonly set under `modules.post_trainer.overrides` in your configuration YAML (see the examples in the repository under `examples/megatron_bridge/configs/`). + +| Area | Parameters | Notes | +|------|------------|--------| +| **Method** | `peft` | `"none"` for SFT; `lora` for LoRA. | +| **Learning rate** | `finetune_lr`, `min_lr`, `lr_warmup_iters`, `lr_decay_iters` | LoRA usually needs a higher `finetune_lr` than SFT. | +| **Precision** | `precision_config` | Typical: `bf16_mixed`. Alternatives include `fp16_mixed` and `fp32` depending on backend support. | +| **Parallelism** | `tensor_model_parallel_size`, `pipeline_model_parallel_size`, `context_parallel_size`, `sequence_parallel` | Increase TP/PP when the model does not fit on fewer GPUs. | +| **Recompute (memory)** | `recompute_granularity`, `recompute_method`, `recompute_num_layers` | Use to trade compute for activation memory (for example `recompute_granularity: full` with uniform recompute). | +| **Batching / length** | `train_iters`, `global_batch_size`, `micro_batch_size`, `seq_length` | `micro_batch_size` is per-GPU; tune with sequence length and memory. | + +Snippet of an SFT configuration (illustrative only): + +```yaml +modules: + post_trainer: + framework: megatron_bridge + config: sft_trainer.yaml + model: qwen3_32b.yaml + overrides: + peft: "none" + finetune_lr: 5.0e-6 + precision_config: bf16_mixed + tensor_model_parallel_size: 1 + global_batch_size: 8 + micro_batch_size: 1 + seq_length: 8192 +``` + +Snippet of a LoRA configuration (illustrative only): + +```yaml +modules: + post_trainer: + framework: megatron_bridge + config: sft_trainer.yaml + model: qwen3_32b.yaml + overrides: + peft: lora + finetune_lr: 1.0e-4 + precision_config: bf16_mixed + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 +``` + +--- + +> The reference configurations below are organized by GPU architecture. **MI325X** uses the same configurations as **MI300X** (both are `gfx942`), and **MI350X** uses the same configurations as **MI355X** (both are `gfx950`). + +## MI300X configurations + +Paths are relative to `examples/megatron_bridge/configs/` in the [Primus repository](https://github.com/AMD-AGI/Primus). + +| Model | Method | Config path | TP | GBS | MBS | Seq len | +|-------|--------|-------------|----|-----|-----|---------| +| Qwen3 32B | SFT | `MI300X/qwen3_32b_sft_posttrain.yaml` | 2 | 8 | 2 | 8192 | +| Qwen3 32B | LoRA | `MI300X/qwen3_32b_lora_posttrain.yaml` | 1 | 32 | 2 | 8192 | + +**Legend:** TP = tensor parallel size; GBS = global batch size; MBS = micro batch size per GPU; Seq len = `seq_length`. + +Sample command for running the post-training: + +```bash +./runner/primus-cli direct -- train posttrain \ + --config ./examples/megatron_bridge/configs/MI300X/qwen3_32b_sft_posttrain.yaml +``` + +--- + +## MI355X configurations + +Paths are relative to `examples/megatron_bridge/configs/` in the [Primus repository](https://github.com/AMD-AGI/Primus). + +| Model | Method | Config path | TP | GBS | MBS | Seq len | +|-------|--------|-------------|----|-----|-----|---------| +| Qwen3 32B | SFT | `MI355X/qwen3_32b_sft_posttrain.yaml` | 1 | 8 | 1 | 8192 | +| Qwen3 32B | LoRA | `MI355X/qwen3_32b_lora_posttrain.yaml` | 1 | 32 | 4 | 8192 | + +Sample command for running the post-training: + +```bash +./runner/primus-cli direct -- train posttrain \ + --config ./examples/megatron_bridge/configs/MI355X/qwen3_32b_lora_posttrain.yaml +``` + +--- + +## Best practices + +### Use SFT or LoRA? + +- **SFT is preferred** when you need the strongest possible task fit, have enough GPU memory, and can afford longer runs. +- **LoRA is preferred** when memory is tight, you want fast iteration, or you plan to maintain multiple adapters for different tasks. + +### Learning rates + +- **SFT:** start in the `5e-6`–`1e-5` range; adjust with validation loss. +- **LoRA:** often `1e-4`–`5e-4`; still use warmup (`lr_warmup_iters`) for stability. + +### Batch sizes + +- **SFT**: starting with `global_batch_size: 8` is a reasonable default for development; scale up when stable (for example to 64, 128, or higher) if memory and throughput allow. +- **LoRA**: larger global batches are often feasible (for example 32 in the reference configs); align `micro_batch_size` with sequence length and available HBM. +- Very long sequences (for example 8192) may require smaller micro-batches or more parallelism. + +### Parallelism + +- **SFT:** large models may need higher `tensor_model_parallel_size` (for example TP 8 for very large models). The bundled 32B examples use TP 2 on MI300X and TP 1 on MI355X for SFT. +- **LoRA:** adapters reduce memory pressure; lower TP is often sufficient for a given model size. + +--- + +## Troubleshooting + +### Out of memory (OOM) + +**SFT** + +1. Increase `tensor_model_parallel_size` (and/or pipeline parallelism for very large models). +2. Reduce `micro_batch_size` or `seq_length`. +3. Enable activation recomputation (`recompute_granularity`, `recompute_method`, `recompute_num_layers`). + +**LoRA** + +1. Confirm `peft: lora` is set. +2. Reduce `micro_batch_size` if OOM persists. +3. Apply the same recompute settings as for SFT. + +### Training instability (loss spikes, NaNs) + +1. Decrease `finetune_lr`. +2. Increase `lr_warmup_iters`. +3. Keep mixed precision stable (`precision_config: bf16_mixed` where supported). +4. Monitor gradients and clipping settings if exposed by your trainer config. + +### Slow training + +1. Increase effective batch size where memory allows (`global_batch_size` / `micro_batch_size` tuning). +2. Revisit TP/PP/CP for your cluster topology. +3. Run [benchmarks](./benchmarking.md) or [preflight](./preflight.md) to isolate network or GPU issues. + +### Configuration errors + +1. Verify YAML paths and indentation. +2. Set `PRIMUS_WORKSPACE` and other environment variables expected by your team’s templates. +3. Confirm checkpoint and data paths by reviewing the experiment YAML and any presets it references. The current core training runtime parses `--export_config`, but resolved-config export is not implemented on the default `PrimusRuntime` path. + +--- + +## Related documentation + +- [Megatron Bridge parameters](../03-configuration-reference/megatron-bridge-parameters.md) +- [Native SFT / LoRA quick start](../04-technical-guides/native-sft-lora.md) +- [Benchmark suite](./benchmarking.md) +- [Preflight diagnostics](./preflight.md) +- [Memory and performance projection](./projection.md) diff --git a/docs/02-user-guide/preflight.md b/docs/02-user-guide/preflight.md new file mode 100644 index 000000000..d88e0d8ec --- /dev/null +++ b/docs/02-user-guide/preflight.md @@ -0,0 +1,136 @@ +# Preflight diagnostics + +`preflight` is Primus’s cluster diagnostic command. It can produce a **fast environment report** (host, GPU, and network facts) and optionally run **performance tests** (GEMM plus intra- and inter-node communication) to catch misconfiguration or outliers before large distributed training jobs. + +`preflight` is implemented by `primus/cli/subcommands/preflight.py`, which in turn delegates to `primus.tools.preflight`. + +--- + +## Overview: What preflight checks + +| Category | What are checked | +|----------|----------------| +| **Host** | CPU, memory, PCIe, and related system context | +| **GPU** | ROCm-visible GPU inventory and key attributes | +| **Network** | Network configuration relevant to distributed training | +| **Performance tests** | Heavier GEMM and communication tests (slower than information-only) | + +--- + +## Quick start + +### Information only (fast) + +```bash +primus-cli direct -- preflight --host --gpu --network +``` + +### Full preflight (information and performance tests) + +```bash +primus-cli direct -- preflight +``` + +### Performance tests only + +Skips the host, GPU, and network information report and runs GEMM + communication tests. + +```bash +primus-cli direct -- preflight --perf-test +``` + +--- + +## CLI flags reference + +| Flag | Purpose | +|------|---------| +| `--host` | Include host information (CPU, memory, PCIe). Alias: `--check-host`. | +| `--gpu` | Include GPU information. Alias: `--check-gpu`. | +| `--network` | Include network information. Alias: `--check-network`. | +| `--perf-test` | Run **only** performance tests (GEMM plus intra- and inter-node communication); skip the information report. | +| `--plot` | Generate plots when used with `--perf-test`. | +| `--dist-timeout-sec` | Timeout in seconds for `torch.distributed` process-group initialization (default: 120). On failure, `preflight` still attempts to write the information report and exits with a non-zero status. | +| `--dump-path` | Output directory for reports (default: `output/preflight`). | +| `--report-file-name` | Base filename for reports (default: `preflight_report`). | +| `--disable-pdf` | Disable PDF generation (PDF is enabled by default when the toolchain allows). | + +**Behavior notes** + +- With **no** `--host`, `--gpu`, or `--network` flags and **no** `--perf-test`, `preflight` runs in the **full** workflow (information plus performance tests). +- Combine `--host`, `--gpu`, and `--network` to limit the information report to include only those sections. + +--- + +## Usage modes + +### Single-node + +```bash +primus-cli direct -- preflight --host --gpu --network +``` + +### Multi-node (Slurm mode) + +Info report: + +```bash +primus-cli slurm srun -N 4 -- preflight --host --gpu --network +``` + +Full `preflight`: + +```bash +primus-cli slurm srun -N 4 -- preflight +``` + +Performance tests only: + +```bash +primus-cli slurm srun -N 4 -- preflight --perf-test +``` + +Use the same launcher pattern you rely on for training to ensure that distributed environment variables (`WORLD_SIZE`, `RANK`, `MASTER_ADDR`, etc.) are consistent. + +--- + +## Output files and contents + +Default output directory: `output/preflight` (override with `--dump-path`). + +| File(s) | Contents | +|---------|----------| +| `.md` / `.pdf` | **Information** report: host, GPU, and network sections when those checks are enabled. | +| `_perf.md` / `_perf.pdf` | **Performance** report: GEMM and communication results from the performance test path. | + +The base `` comes from `--report-file-name` (default: `preflight_report`). + +--- + +## Interpreting results + +1. **Information report:** Confirm GPU count, model match expectations, and PCIe topology is sensible for your workload. Network sections should reflect the interfaces you intend for distributed training. +2. **Performance report:** Compare GEMM and collective results across nodes. Large outliers on one node often indicate driver, fabric, or process placement issues. +3. **Timeouts:** If `--dist-timeout-sec` is exceeded, inspect firewall rules, interface bindings, `MASTER_ADDR`, and `MASTER_PORT` before scaling up training. + +--- + +## Common issues preflight helps detect + +| Symptom | What to verify in reports | +|---------|---------------------------| +| Missing or wrong GPU count | GPU section: ROCm health on the node | +| Wrong network device or address | Network section: NCCL/RCCL environment | +| Slow or asymmetric inter-node comm | Performance report: compare ranks or nodes | +| Hangs at distributed process group initialization | Use `--dist-timeout-sec` to avoid, then check rendezvous and Slurm network setup | + +For deeper, single-purpose measurements, see the [Benchmark suite](./benchmarking.md). + +--- + +## Related documentation + +- [Benchmark suite](./benchmarking.md) +- [Memory and performance projection](./projection.md) +- [Post-training workflows](./posttraining.md) +- [Installation and setup](../01-getting-started/installation.md) diff --git a/docs/02-user-guide/pretraining.md b/docs/02-user-guide/pretraining.md new file mode 100644 index 000000000..110c478ec --- /dev/null +++ b/docs/02-user-guide/pretraining.md @@ -0,0 +1,291 @@ +# Pretraining workflows + +Primus is a YAML-driven training stack for AMD GPUs. You select a **backend** (Megatron-LM, TorchTitan, JAX MaxText, Megatron Bridge), point `train pretrain` at a **configuration YAML**, and launch Primus with the unified CLI (`runner/primus-cli`) in **direct**, **container**, or **Slurm** mode. See [CLI reference](cli-reference.md) and [Configuration system](configuration-system.md). + +This section helps you understand concepts related to the Primus workflow: how backends work, YAML structure and inheritance, parallelism vocabulary, the full per-backend configuration inventory, and so on. If you already understand the concepts and just need the specific commands to run your training with Primus, see [Backend training recipes](training-recipes.md). + +--- + +## Overview + +The following table describes the four backend types supported by Primus and their typical uses. + +| Backend | Framework | Typical use | +| --- | --- | --- | +| Megatron-LM | `framework: megatron` | Large-scale transformer pretraining with Megatron-style parallelism (TP/PP/EP). | +| TorchTitan | `framework: torchtitan` | PyTorch-native scaled training (FSDP / tensor / pipeline / expert parallelism per config). | +| MaxText (JAX) | `framework: maxtext` | JAX/MaxText single- and multi-node runs; parallelism via MaxText `ici_*` / `dcn_*` settings. | +| Megatron Bridge | `framework: megatron_bridge` | Bridge-oriented workflows (configure like other backends; see parameter reference). | + +> Several setup steps apply to **all** backends (mock vs. real data, Hugging Face tokens, scaling to multiple nodes, and HipBLASLt autotuning). After you read the backend section that applies to you, see [Common patterns](#common-patterns) below. + +--- + +## Megatron-LM pretraining + +### Quick start (container mode) + +From the root of the clone of the [Primus repository](https://github.com/AMD-AGI/Primus), with Docker or Podman available, the following command starts the training in container mode: + +```bash +./runner/primus-cli container -- train pretrain \ + --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +``` + +This uses the default image from `runner/.primus.yaml` (`rocm/primus:v26.3` unless overridden). The project tree is mounted into the container automatically by `runner/primus-cli-container.sh`. + +### Example configurations under `examples/megatron/configs/MI300X/` + +The following files ship in the repository (sorted by name). Parallelism columns are taken from `tensor_model_parallel_size` / `pipeline_model_parallel_size` / `expert_model_parallel_size` in each file (literals or `${PRIMUS_TP:…}` defaults). + +| Config | TP | PP | EP | +| --- | --- | --- | --- | +| `deepseek_v2-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:4}` | `${PRIMUS_EP:8}` | +| `deepseek_v2-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:4}` | `${PRIMUS_EP:8}` | +| `deepseek_v2_lite-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `deepseek_v2_lite-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `deepseek_v3-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `deepseek_v3-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `gpt_oss_20B-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `gpt_oss_20B-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `grok1-BF16-pretrain.yaml` | `1` | `4` | `8` | +| `grok1-FP8-pretrain.yaml` | `1` | `4` | `8` | +| `grok2-BF16-pretrain.yaml` | `1` | `4` | `8` | +| `grok2-FP8-pretrain.yaml` | `1` | `4` | `8` | +| `llama2_13B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `llama2_13B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `llama2_70B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `llama2_70B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `llama2_7B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `llama2_7B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `llama3.1_405B-BF16-pretrain.yaml` | `8` | `8` | `1` | +| `llama3.1_405B-FP8-pretrain.yaml` | `8` | `8` | `1` | +| `llama3.1_70B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `llama3.1_70B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `llama3.1_8B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `llama3.1_8B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `llama3.2_1B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `llama3.2_1B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `llama3.2_3B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `llama3.2_3B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `llama3.3_70B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `llama3.3_70B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `llama3_70B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `llama3_70B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `llama3_8B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `llama3_8B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `llama4_17B128E-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `llama4_17B128E-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `llama4_17B16E-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `llama4_17B16E-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `mamba_370M-pretrain.yaml` | `1` | `1` | `1` | +| `mixtral_8x22B_v0.1-BF16-pretrain.yaml` | `1` | `4` | `8` | +| `mixtral_8x22B_v0.1-FP8-pretrain.yaml` | `1` | `4` | `8` | +| `mixtral_8x7B_v0.1-BF16-pretrain.yaml` | `1` | `1` | `8` | +| `mixtral_8x7B_v0.1-FP8-pretrain.yaml` | `1` | `1` | `8` | +| `qwen2.5_14B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `qwen2.5_14B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `qwen2.5_32B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `qwen2.5_32B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `qwen2.5_3B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `qwen2.5_3B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `qwen2.5_72B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `qwen2.5_72B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `qwen2.5_7B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `qwen2.5_7B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `qwen3_14B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `qwen3_14B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `qwen3_235B_A22B-BF16-pretrain.yaml` | `1` | `1` | `8` | +| `qwen3_235B_A22B-FP8-pretrain.yaml` | `1` | `1` | `8` | +| `qwen3_30B_A3B-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `qwen3_30B_A3B-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `qwen3_32B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `qwen3_32B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `qwen3_4B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `qwen3_4B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `qwen3_5_35B_A3B-BF16-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `qwen3_5_35B_A3B-FP8-pretrain.yaml` | `${PRIMUS_TP:1}` | `${PRIMUS_PP:1}` | `${PRIMUS_EP:8}` | +| `qwen3_8B-BF16-pretrain.yaml` | `1` | `1` | `1` | +| `qwen3_8B-FP8-pretrain.yaml` | `1` | `1` | `1` | +| `zebra_llama_1B-pretrain.yaml` | `1` | `1` | `1` | +| `zebra_llama_3B-pretrain.yaml` | `1` | `1` | `1` | +| `zebra_llama_8B-pretrain.yaml` | `1` | `1` | `1` | + +### Sample YAML file (`llama2_7B-BF16-pretrain.yaml`) explained + +Path: `examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml` + +| Section | Role | +| --- | --- | +| `work_group`, `user_name`, `exp_name`, `workspace` | Run identity and output root (supports `${VAR:default}` substitution). | +| `modules.pre_trainer.framework` | `megatron` selects Megatron-LM integration. | +| `config: pre_trainer.yaml` | Module preset under `primus/configs/modules/megatron/`. | +| `model: llama2_7B.yaml` | Model preset under `primus/configs/models/megatron/` (extends `llama2_base.yaml` → …). | +| `overrides` | Run-specific training knobs: iterations, batching, LR, **parallelism** (`tensor_model_parallel_size`, `pipeline_model_parallel_size`, `expert_model_parallel_size`), data paths, checkpoints, Primus Turbo flags, etc. | + +The sample sets `mock_data: true` and `train_data_path: null` so you can validate the stack without real corpora. + +### Mock data versus real data + +- **Mock data:** Set `mock_data: true` and leave `train_data_path` / `valid_data_path` empty (as in `llama2_7B-BF16-pretrain.yaml`). +- **Real data:** Set `mock_data: false` and populate Megatron-compatible data paths (and tokenizer assets) in `overrides`. Use paths visible inside your container mounts. + +### Multi-node training with Slurm + +```bash +./runner/primus-cli slurm srun -N 4 -p -- train pretrain \ + --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +``` + +`runner/primus-cli-slurm-entry.sh` derives `MASTER_ADDR`, `NNODES`, and `NODE_RANK` from Slurm and forwards them into the container. Align `tensor_model_parallel_size`, `pipeline_model_parallel_size`, and `expert_model_parallel_size` with your cluster width and job size. + +--- + +## TorchTitan pretraining + +### Quick start + +```bash +./runner/primus-cli container -- train pretrain \ + --config examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml +``` + +### Example configurations under `examples/torchtitan/configs/MI300X/` + +| File | +| --- | +| `deepseek_v3_16b-BF16-pretrain.yaml` | +| `deepseek_v3_16b-FP8-pretrain.yaml` | +| `deepseek_v3_236b-BF16-pretrain.yaml` | +| `deepseek_v3_236b-FP8-pretrain.yaml` | +| `deepseek_v3_671b-pretrain.yaml` | +| `llama3.1_405B-BF16-pretrain.yaml` | +| `llama3.1_405B-FP8-pretrain.yaml` | +| `llama3.1_70B-BF16-pretrain.yaml` | +| `llama3.1_70B-FP8-pretrain.yaml` | +| `llama3.1_8B-BF16-pretrain.yaml` | +| `llama3.1_8B-FP8-pretrain.yaml` | +| `llama4_17Bx128E-BF16-pretrain.yaml` | +| `llama4_17Bx128E-FP8-pretrain.yaml` | +| `llama4_17Bx16E-BF16-pretrain.yaml` | +| `llama4_17Bx16E-FP8-pretrain.yaml` | +| `qwen3_0.6B-pretrain.yaml` | +| `qwen3_1.7B-pretrain.yaml` | +| `qwen3_14B-pretrain.yaml` | +| `qwen3_32B-pretrain.yaml` | +| `qwen3_4B-pretrain.yaml` | +| `qwen3_8B-pretrain.yaml` | + +### Sample YAML file (`llama3.1_8B-BF16-pretrain.yaml`) explained + +Path: `examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml` + +| Section | Role | +| --- | --- | +| `framework: torchtitan` | Selects the TorchTitan integration. | +| `config: pre_trainer.yaml` | Module preset under `primus/configs/modules/torchtitan/`. | +| `model: llama3.1_8B.yaml` | Model preset under `primus/configs/models/torchtitan/`. | +| `overrides.training`, `lr_scheduler`, `activation_checkpoint`, `primus_turbo` | Run-specific batching, steps, checkpointing, and Turbo options. | + +Some configurations omit an explicit `parallelism:` block; in that case the default values come from the **module and model presets** (`primus/configs/modules/torchtitan/pre_trainer.yaml` and the chosen model YAML). Other examples (for example DeepSeek and Qwen) set `parallelism:` inline with `tensor_parallel_degree`, `pipeline_parallel_degree`, `expert_parallel_degree`, etc. + +--- + +## MaxText (JAX) pretraining + +### Quick start + +```bash +./runner/primus-cli container -- train pretrain \ + --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml +``` + +### JAX-specific requirements + +Install JAX/MaxText dependencies from the repository root: + +```bash +pip install -r requirements-jax.txt +``` + +### Example configurations under `examples/maxtext/configs/MI300X/` + +| File | Key parallelism (`ici_*` intra-node, `dcn_*` inter-node) | +| --- | --- | +| `deepseek_v2_16B-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` | +| `grok1-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` | +| `llama2_70B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` | +| `llama2_7B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` | +| `llama3.3_70B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` | +| `llama3_70B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` | +| `llama3_8B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` | +| `mixtral_8x7B-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` | +| `qwen3_14B-pretrain.yaml` | `ici_fsdp_parallelism: 8`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` | +| `qwen3_30B_A3B-pretrain.yaml` | `ici_fsdp_parallelism: 1`, `ici_data_parallelism: 1`, `dcn_fsdp_parallelism: 1`, `dcn_data_parallelism: -1` | + +The `llama2_7B-pretrain.yaml` example also sets `dataset_type: "synthetic"` and `hf_access_token: ${HF_TOKEN:""}` for gated Hugging Face assets when you switch to real data. + +--- + +## Common patterns + +### Testing with mock data + +Set `mock_data: true` (Megatron/TorchTitan) or synthetic dataset settings (MaxText) to validate the configurations and infrastructure without I/O-heavy datasets. + +### Real training data + +- Megatron: Configure `train_data_path` / `valid_data_path` and tokenizer assets in `overrides` once `mock_data` is false. +- For **all backends**, ensure host paths are mounted in **container** mode (`--volume` or `container.options.volume` in YAML). +- TorchTitan/MaxText: Follow backend-specific dataset fields in the `overrides` and presets. + +### Scaling from single-node to multi-node + +- Use **Slurm** mode for allocation; keep the **container** entry if you want the same image on every node. +- Set environment variables consistently (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, `MASTER_PORT`, `GPUS_PER_NODE`); the Slurm entry script injects them when using `primus-cli slurm`. +- Increase values in the parallelism fields (Megatron TP/PP/EP; TorchTitan `parallelism`; MaxText `ici_*` / `dcn_*`) to match topology. + +### Hugging Face token for gated models + +Export `HF_TOKEN` on the host before launching **container** mode; `runner/.primus.yaml` lists `HF_TOKEN` under `container.options.env` so it can be forwarded into the container. MaxText configurations may reference `${HF_TOKEN:""}` directly. + +### HipBLASLt autotuning (three stages) + +Controlled with `PRIMUS_HIPBLASLT_TUNING_STAGE` (see `examples/README.md`): + +| Stage | Purpose | +| --- | --- | +| 1 | Dump GEMM shapes seen during training (reduce `train_iters` for faster collection). | +| 2 | Tune kernels from dumped shapes (offline tooling under `examples/offline_tune`). | +| 3 | Train using tuned kernel artifacts from `./output/tune_hipblaslt/...`. | + +Example (from in-repo docs): + +```bash +export PRIMUS_HIPBLASLT_TUNING=1 # master switch (required; tuning is skipped without it) +export PRIMUS_HIPBLASLT_TUNING_STAGE=1 +./runner/primus-cli direct -- train pretrain \ + --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +``` + +--- + +## Supported models + +The tables above in the Megatron, TorchTitan, and MaxText sections are curated MI300X examples from the [Primus repository](https://github.com/AMD-AGI/Primus). Use `examples//configs/` in the repository as the authoritative inventory, as new presets and hardware-specific examples may be added there before this document is updated to reflect their additions. + +| Backend | Example region | Parallelism vocabulary | +| --- | --- | --- | +| Megatron-LM | `examples/megatron/configs/MI300X/` | `tensor_model_parallel_size`, `pipeline_model_parallel_size`, `expert_model_parallel_size` (and env-driven `${PRIMUS_TP:…}` variants). | +| TorchTitan | `examples/torchtitan/configs/MI300X/` | `parallelism.*` (e.g. `tensor_parallel_degree`, `pipeline_parallel_degree`, `expert_parallel_degree`, FSDP shard settings). | +| MaxText | `examples/maxtext/configs/MI300X/` | `ici_fsdp_parallelism`, `ici_data_parallelism`, `dcn_fsdp_parallelism`, `dcn_data_parallelism`. | + +For scripting patterns that predate `primus-cli`, the repository still documents `examples/run_local_pretrain.sh` and `examples/run_slurm_pretrain.sh` in `examples/README.md`; equivalent launches are shown above using `./runner/primus-cli`. + +--- + +## Related documentation + +- [CLI reference](cli-reference.md): launcher usage +- [Configuration system](configuration-system.md): YAML merge rules +- Backend parameter references: [Megatron parameters](../03-configuration-reference/megatron-parameters.md), [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md), [MaxText parameters](../03-configuration-reference/maxtext-parameters.md) diff --git a/docs/02-user-guide/primus-tools.md b/docs/02-user-guide/primus-tools.md new file mode 100644 index 000000000..212a3fd73 --- /dev/null +++ b/docs/02-user-guide/primus-tools.md @@ -0,0 +1,29 @@ +# Primus tools + +A quick catalog of the tools that ship with Primus and the sibling projects +around it—command-line tools, the tuning agent, ecosystem projects, and +auxiliary utilities. Each row gives a short description and a how-to starting +point; follow a tool's link for the full reference. + +| Tool | Type | What it does | How to use | +|------|------|--------------|------------| +| [`train`](./pretraining.md) | CLI | Launch pretraining or post-training on any backend from a YAML configuration. | `primus-cli -- train pretrain --config ` | +| [`benchmark`](./benchmarking.md) | CLI | GEMM, RCCL, and attention microbenchmarks for hardware and stack validation. | `primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096` | +| [`preflight`](./preflight.md) | CLI | Host, GPU, and network health checks before long jobs. | `primus-cli slurm srun -N 4 -- preflight --host --gpu --network` | +| [`projection`](./projection.md) | CLI | Estimate per-GPU memory and throughput without occupying a full cluster. | `primus-cli direct -- projection both --config ` | +| [Tuning agent](./tuning-agent.md) | Agent | LLM-driven search for a near-optimal training configuration, scored by projection. | `python -m primus.agents.tuning_agent --workload --target-cluster ` | +| [Primus-LM](../01-getting-started/quickstart.md) | Ecosystem | The training framework in this repository (multi-backend, unified CLI). | See the [Quickstart](../01-getting-started/quickstart.md) | +| [Primus-Turbo](https://github.com/AMD-AGI/Primus-Turbo) | Ecosystem | High-performance ROCm operators (attention, GEMM, grouped GEMM, DeepEP, FP8/FP4). | Bundled in the `rocm/primus` image; enabled via configuration flags | +| [Primus-SaFE](https://github.com/AMD-AGI/Primus-SaFE) | Ecosystem | Kubernetes-native stability, scheduling, and fault-tolerance platform. | Deployed separately on Kubernetes (Helm) | +| [IRLens](../../tools/IRLens/README.md) | Auxiliary | Parse XLA HLO dumps into a communication-vs-compute execution skeleton. | See the [README](../../tools/IRLens/README.md) | +| [model_stats](../../tools/model_stats/README.md) | Auxiliary | Chart model dimensions from the config registry. | See the [README](../../tools/model_stats/README.md) | +| [Pipeline visualization](../../tools/visualization/pp_vis/README.md) | Auxiliary | Render pipeline-parallel schedules in a local web UI. | See the [README](../../tools/visualization/pp_vis/README.md) | +| [Auto benchmark](../../tools/auto_benchmark/Primus_Auto_Benchmark_README.md) | Auxiliary | Interactive Megatron/TorchTitan benchmark menu with metrics collection. | See the [README](../../tools/auto_benchmark/Primus_Auto_Benchmark_README.md) | + +--- + +## Related documentation + +- [CLI reference](./cli-reference.md)—full launcher grammar and subcommand options. +- [Tooling](../06-developer-guide/tooling.md)—developer-guide index of the `tools/` utilities. +- [Project overview](../01-getting-started/overview.md#primus-ecosystem)—how the Primus ecosystem layers fit together. diff --git a/docs/02-user-guide/projection.md b/docs/02-user-guide/projection.md new file mode 100644 index 000000000..7b087df65 --- /dev/null +++ b/docs/02-user-guide/projection.md @@ -0,0 +1,177 @@ +# Memory and performance projection + +Primus projection tools estimate **per-GPU memory** and **training throughput** for large-scale distributed jobs without requiring the full target cluster. Two modes are available: analytical **memory** projection and **performance** projection that combines profiling with simulation. + +**Implementation:** `primus/cli/subcommands/projection.py` + +| Mode | Command | Role | +|------|---------|------| +| **Memory** | `projection memory` | Estimates per-GPU memory (parameters, optimizer state, activations) using analytical formulas. | +| **Performance** | `projection performance` | Benchmarks on a single node (or sub-node), then projects training time to multi-node configurations. | +| **Both** | `projection both` | Runs a single benchmark and produces **both** the performance and (benchmark-anchored) memory projections from it. Recommended for cluster-sizing workflows. | + +**Core logic** + +- Memory: `primus/core/projection/memory_projection/` +- Performance: `primus/core/projection/performance_projection/` + +Related: [Benchmark suite](./benchmarking.md), [Preflight diagnostics](./preflight.md), [Megatron parameters](../03-configuration-reference/megatron-parameters.md). + +--- + +## Memory projection + +### Quick start + +```bash +export NNODES=1 +export HSA_NO_SCRATCH_RECLAIM=1 + +./runner/primus-cli direct --script primus/cli/main.py -- \ + projection memory \ + --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml +``` + +Adjust `--config` to your experiment YAML. Memory estimation is analytical; the CLI still expects a normal Primus launch path (including distributed initialization where applicable). + +### What it estimates + +| Component | Meaning | +|-----------|---------| +| **Parameter memory** | Model weights assigned to this GPU (respecting parallelism). | +| **Optimizer memory** | Optimizer state (for example Adam moments), accounting for sharding across data-parallel groups. | +| **Activation memory** | Activations retained for the backward pass for a given microbatch and sequence length. | + +The tool walks a hierarchical profiler structure aligned with the model (embeddings, dense and MoE layers, output head, loss) and aggregates per-component contributions. + +### How to interpret results + +Console output includes per-component breakdowns and a summary such as parameter count, param+optimizer memory, activation memory for the configured batch size and sequence length, and a projected total. Use these to answer whether a configuration fits in HBM before you allocate large clusters. + +--- + +## Performance projection + +### Quick start + +Minimum required nodes (derived from parallelism): + +```bash +export NNODES=1 +export HSA_NO_SCRATCH_RECLAIM=1 + +./runner/primus-cli direct --script primus/cli/main.py -- \ + projection performance \ + --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml +``` + +### How it works + +1. **Profile** layer-level behavior on **one node** (or a subset of GPUs with automatic scaling rules). +2. **Simulate** pipeline scheduling, data parallelism, and communication using analytical models. +3. **Project** iteration time and tokens/s to a **target** node count when you specify one. + +### Projecting to a specific node count + +```bash +./runner/primus-cli direct --script primus/cli/main.py -- \ + projection performance \ + --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml \ + --target-nodes 4 +``` + +If `--target-nodes` is omitted, the tool defaults to the **minimum** number of nodes implied by your parallelism configuration (TP, PP, EP, CP, GPUs per node). + +### Parallelism overrides (environment) + +You can override parallelism for what-if analysis: + +```bash +export PRIMUS_TP=1 +export PRIMUS_PP=3 +export PRIMUS_EP=8 + +./runner/primus-cli direct --script primus/cli/main.py -- \ + projection performance \ + --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml \ + --target-nodes 6 +``` + +--- + +## Command reference + +### Syntax + +```bash +primus-cli [global-options] [mode-args] -- projection {memory,performance,both} [options] +``` + +### Shared options (both modes) + +| Option | Description | +|--------|-------------| +| `--config` / `--exp` | Path to the Primus YAML configuration (**required**). | +| `--data_path` | Data directory (default `./data` when included on the parser). | +| `--backend_path` | Optional Megatron/TorchTitan import path appended to `PYTHONPATH`. | +| `--export_config` | Accepted by the shared pretrain parser, but the default core runtime does not currently write a resolved YAML file. | + +### Performance-only options + +| Option | Description | +|--------|-------------| +| `--target-nodes` | Target number of nodes for scaling projection. Defaults to the minimum nodes required by TP/PP/EP/CP and GPUs per node. | +| `--target-num-nodes` | Alias-style projection override for target node count. | +| `--target-ep-size` | Override `expert_model_parallel_size` for the projection target. | +| `--benchmark-gpus` | Use fewer than `GPUS_PER_NODE` GPUs for benchmarking; results are scaled analytically back to a full node. | +| `--hardware-config` | YAML file with hardware parameters for communication modeling. | +| `--profiling-mode` | `benchmark` (default, uses GPU), `simulate` (analytical / Origami GEMM + SDPA models, no GPU), or `both` (side-by-side). | +| `--gemm-backend` | GEMM simulation backend when profiling is simulated (`origami`). | +| `--gpu-arch` | Target architecture for simulation (for example `mi300x`, `gfx942`, `mi355x`, `gfx950`); can use `PRIMUS_GPU_ARCH`. | +| `--gpu-clock-mhz` | Override GPU clock in MHz for simulation; can use `PRIMUS_GPU_CLOCK_MHZ`. | +| `--pipeline-schedule-algorithm` | Pipeline simulation scheduler (`auto`, zero-bubble variants, or `all` for comparison). | +| `--enable-zero-bubble` | Enable zero-bubble pipeline scheduling for projection. | +| `--enable-deepep` | Enable DeepEP overlap modeling. | +| `--sync-free-stage` | Override Sync-Free MoE stage (`0` off; stages `1`-`3` enable additional modeling assumptions). | +| `--num-virtual-stages-per-pipeline-rank` | Override virtual pipeline stage count for projection. | +| `--micro-batch-size`, `--global-batch-size` | Override batch sizes for projection without editing the YAML. | + +--- + +## Assumptions and limitations + +### Assumptions (performance projection) + +1. **Data-parallel scaling**—Compute time scales with ideal weak-scaling assumptions versus data-parallel width. +2. **Communication Model**—Uses simplified bandwidth and latency models (defaults such as efficiency factors may apply). +3. **Pipeline scheduling**—Bubble and overlap behavior is modeled with fixed splits; real frameworks may differ. +4. **Gradients and MoE**—Gradient all-reduce overlap and MoE all-to-all behavior follow the implemented model (for example overlap flags, EP scaling). + +### Limitations + +1. **Single-node benchmark accuracy**—Reduced PP/EP on the benchmark GPU count may not capture every production behavior. +2. **Contention**—Network contention between jobs is not modeled. +3. **Memory vs speed**—Activation recomputation reduces memory but adds compute; performance projection may not fully reflect that trade-off unless modeled. +4. **Heterogeneity**—Assumes homogeneous nodes; GPU frequency drift across nodes is not modeled. + +--- + +## Tips + +1. Run **`projection memory`** first to confirm a configuration is feasible in HBM before spending time on performance projection. +2. Always establish a **single-node** baseline before interpreting multi-node projections. +3. **Data-parallel scaling** is bounded by batching: if you run out of microbatches (`global_batch_size` / `micro_batch_size`), adding nodes may not increase throughput. +4. If the YAML **requires** multiple nodes (for example large PP), the performance path may automatically reduce parallelism for benchmarking and restore it analytically—read the console summary carefully. +5. **No GPU available:** use `--profiling-mode simulate` for CPU-side analytical timing. +6. **Validate models:** use `--profiling-mode both` to compare GPU benchmark timing with simulation on the same config. +7. For **MoE** models, activation memory from MoE layers often dominates; memory projection highlights when recomputation is worth considering. + +--- + +## Related documentation + +- [Benchmark suite](./benchmarking.md) +- [Preflight diagnostics](./preflight.md) +- [Post-training workflows](./posttraining.md) +- [Tuning agent](./tuning-agent.md) +- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) diff --git a/docs/02-user-guide/training-recipes.md b/docs/02-user-guide/training-recipes.md new file mode 100644 index 000000000..1d36ba20c --- /dev/null +++ b/docs/02-user-guide/training-recipes.md @@ -0,0 +1,260 @@ +# Backend training recipes + +Task-oriented, copy-paste commands for launching pretraining runs with each Primus backend on AMD Instinct™ GPUs. + +This section is for users who already know what they want to run and need the specific command for a given model, precision, and GPU. To understand concepts related to the Primus workflow (how backends work, YAML structure and inheritance, parallelism vocabulary, the full per-backend configuration inventory, etc.), see **[Pretraining](pretraining.md)**. + +> **Authoritative full matrices.** AMD publishes per-model reproduction pages with verified images, commits, and tuned batch sizes for every supported model at the following locations—treat them as the source of truth for achieving the expected performance; this page only gives the canonical *pattern* plus a representative example per backend and links to other reference materials. +> +> - [Training with Primus + Megatron-LM](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/benchmark-docker/primus-megatron.html) +> - [Training with Primus + PyTorch (TorchTitan)](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/benchmark-docker/primus-pytorch.html) +> - [Training with Primus + JAX MaxText](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/benchmark-docker/jax-maxtext.html) + +--- + +## How recipes are structured + +Every recipe follows the same four-step pattern: + +1. **Pull and launch the AMD Docker image** (for a reproducible environment). +2. **Set the GPU-architecture environment** (performance environment variable settings differ by GPU). +3. **Pick the configuration YAML** for your GPU architecture under `examples//configs//` in the [Primus repository](https://github.com/AMD-AGI/Primus). +4. **Launch** with `runner/primus-cli` in `direct`, `container`, or `slurm` mode. + +### GPU-architecture config folders + +Configuration YAMLs are organized by GPU architecture. Always pick the folder that matches your hardware: + + +| Backend | `MI300X` | `MI325X` | `MI355X` / `MI350X` | +| ----------------------------------- | -------- | -------- | ------------------- | +| `examples/megatron/configs/` | yes | yes | yes | +| `examples/torchtitan/configs/` | yes | yes | yes | +| `examples/maxtext/configs/` | yes | — | yes | +| `examples/megatron_bridge/configs/` | yes | — | yes | + + +> MI350X uses the same configurations as MI355X because both are based on the gfx950 architecture. If a configuration for your model is not available in the architecture-specific folder, use the closest match from the same generation as a starting point. + +### GPU-architecture environment variables + +MI300X and MI325X benefit from the following performance settings; MI355X/MI350X do **not** need them: + +```bash +# MI300X / MI325X only -- improves performance +export HSA_NO_SCRATCH_RECLAIM=1 +export PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1 +export NVTE_CK_IS_V3_ATOMIC_FP32=1 +``` + +### Choosing the Docker image + +For **container** and **Slurm** modes (direct mode runs in whatever environment you launched it from), the default image is `rocm/primus:v26.3` (`runner/.primus.yaml`). For reproducing published benchmarks, use the AMD-published tag for your release (the AMD pages under **Authoritative full matrices** above list the most current tag). JAX MaxText has its own separate image family of `rocm/jax-training:maxtext-...`. + +Image is picked in the priority order of `DOCKER_IMAGE` environment variable > `--image` CLI argument > config file. See [Selecting the container image](../01-getting-started/quickstart.md#selecting-the-container-image) for a full explanation, and [Configuration system](configuration-system.md) for configuration loading. + +--- + +## Shared setup for all backends + +These apply across all backends. Set them up before running the recipes below. + +### Hugging Face token (for gated models or real data) + +```bash +export HF_TOKEN= +``` + +`runner/.primus.yaml` forwards `HF_TOKEN` into the container automatically. MaxText configurations may also read `${HF_TOKEN:""}` directly. + +### Mock vs. real data + +- **Mock/synthetic data** (default for most examples): validates the stack without datasets. Megatron and TorchTitan set `mock_data: true`; MaxText sets `dataset_type: "synthetic"`. +- **Real data:** set `mock_data: false` and point `train_data_path` (for Megatron) or the backend's dataset fields at paths visible *inside* your container mounts. + +### Multi-node networking checklist + +The `primus-cli` launcher sets sensible `NCCL_`* defaults, but auto-detection can pick the wrong device on multi-NIC nodes. Before multi-node, confirm and export if needed: + +```bash +export NCCL_IB_HCA= # from `ibv_devices` +export NCCL_SOCKET_IFNAME= # from `ip a` +export GLOO_SOCKET_IFNAME= +export NCCL_IB_GID_INDEX=3 # 3 for RoCE (1 for AMD AINIC) +``` + +For AMD AINIC clusters also set `USING_AINIC=1`, `NCCL_PXN_DISABLE=0`, `NCCL_IB_GID_INDEX=1`. See [Multi-Node Networking](../04-technical-guides/multi-node-networking.md) for the full reference. + +--- + +## Megatron-LM + +**Image:** `rocm/primus`  |  **Configurations:** `examples/megatron/configs//`  |  **Precisions:** BF16, FP8 + +### 1. Launch the container + +```bash +docker pull rocm/primus:v26.3 +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.3 +``` + +Access the container later with `docker start primus_training_env && docker exec -it primus_training_env bash`. + +### 2. Run pretraining (direct mode, inside the container) + +Pretrain Llama 3.1 8B BF16 on **MI355X / MI350X**: + +```bash +./runner/primus-cli direct \ + --log_file /tmp/primus_llama3.1_8B.log \ + -- train pretrain \ + --config examples/megatron/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml +``` + +Pretrain the same model on **MI300X / MI325X** (add the performance environment variables): + +```bash +export HSA_NO_SCRATCH_RECLAIM=1 +export PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1 +export NVTE_CK_IS_V3_ATOMIC_FP32=1 + +./runner/primus-cli direct \ + --log_file /tmp/primus_llama3.1_8B.log \ + -- train pretrain \ + --config examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml +``` + +Switch model or precision by changing the config filename (e.g. `llama3.1_70B-FP8-pretrain.yaml`, `mixtral_8x7B_v0.1-BF16-pretrain.yaml`). The full configuration inventory is the repository's `examples/megatron/configs//` directory. See the parallelism table in [Pretraining](pretraining.md#example-configurations-under-examplesmegatronconfigsmi300x). + +**Model-specific notes:** + +- **Zebra-Llama** (hybrid Mamba+MLA) pretrain presets ship at `examples/megatron/configs//zebra_llama_{1B,3B,8B}-pretrain.yaml` and run via the standard core runtime; Megatron Bridge SFT variants live under `examples/megatron_bridge/configs//`. +- **MoE models** (DeepSeek-V2-Lite, Mixtral) may need extra grouped-GEMM or router flags; [Training with Primus + Megatron-LM](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/training/benchmark-docker/primus-megatron.html) lists the exact flags per model. + +### 3. Multi-node (Slurm mode) + +```bash +./runner/primus-cli slurm srun -N 8 -p -- train pretrain \ + --config examples/megatron/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml \ + --micro_batch_size 4 --global_batch_size 1024 +``` + +Scale batch size with node count and align `tensor_model_parallel_size`, `pipeline_model_parallel_size`, and `expert_model_parallel_size` to your topology. See the [multi-node networking checklist](#multi-node-networking-checklist) above. + +--- + +## TorchTitan (PyTorch) + +**Image:** `rocm/primus`  |  **Configurations:** `examples/torchtitan/configs//`  |  **Precisions:** BF16, FP8 + +Use the same `rocm/primus` container as Megatron (step 1 above). TorchTitan parameters use a dotted namespace (e.g. `--training.local_batch_size`). + +### Run pretraining (direct mode) + +Pretrain Llama 3.1 8B BF16 on **MI355X / MI350X**: + +```bash +./runner/primus-cli direct \ + --log_file /tmp/primus_llama3.1_8B.log \ + -- train pretrain \ + --config examples/torchtitan/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml +``` + +On **MI300X / MI325X**, export the performance environment variables first (see above) and use the `MI300X` config path. + +### Multi-node (Slurm mode) + +```bash +./runner/primus-cli slurm srun -N 4 -- train pretrain \ + --config examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml \ + --training.local_batch_size 6 \ + --training.global_batch_size 192 \ + --training.mock_data True +``` + +Available models include Llama 3.1 (8B/70B/405B), Llama 4, DeepSeek V3, and Qwen 3. See the `examples/torchtitan/configs//` directory in the repository. + +--- + +## JAX MaxText + +**Image:** `rocm/jax-training:maxtext-...` (separate family from the other backends)  |  **Configurations:** `examples/maxtext/configs//` + +MaxText uses a different Docker image than Megatron and TorchTitan, and it is **not** the default image pointed to in `runner/.primus.yaml`. In container or Slurm mode, you must point Primus at your MaxText image explicitly. + +### 1. Launch the container + +```bash +docker pull rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0 +docker run -it \ + --device /dev/dri --device /dev/kfd \ + --network host --ipc host \ + --group-add video --cap-add SYS_PTRACE \ + --security-opt seccomp=unconfined --privileged \ + -v $HOME:$HOME -v $HOME/.ssh:/root/.ssh \ + --shm-size 64G \ + --name training_env \ + rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0 +``` + +If you run Primus directly on the host instead of inside the prebuilt Docker image, install the JAX dependencies first by `pip install -r requirements-jax.txt`. + +### 2. Run pretraining + +Direct mode (inside the container)—pretraining Llama 3 8B on **MI355X**: + +```bash +./runner/primus-cli direct \ + -- train pretrain \ + --config examples/maxtext/configs/MI355X/llama3_8B-pretrain.yaml +``` + +Container mode—passing the MaxText image with `--image`: + +```bash +./runner/primus-cli container --image rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0 \ + -- train pretrain \ + --config examples/maxtext/configs/MI355X/llama3_8B-pretrain.yaml +``` + +Slurm mode—supplying the image (and any environment variables) via a config file: + +```bash +./runner/primus-cli --config my_maxtext_config.yaml slurm srun -N 8 \ + -- train pretrain \ + --config examples/maxtext/configs/MI300X/llama3_8B-pretrain.yaml +``` + +MaxText parallelism is set with `ici_*` (intra-node) and `dcn_*` (inter-node) fields—see the [MaxText config table](pretraining.md#maxtext-jax-pretraining) and [MaxText parameters](../03-configuration-reference/maxtext-parameters.md). + +--- + +## Megatron Bridge (post-training) + +Megatron Bridge configurations are under `examples/megatron_bridge/configs//` in the repository and are primarily **SFT and LoRA post-training** recipes (e.g. `qwen3_32b_sft_posttrain.yaml`, `llama31_70b_lora_posttrain.yaml`). Launch with `train posttrain`: + +```bash +./runner/primus-cli direct \ + --log_file /tmp/primus_qwen3_32b_sft.log \ + -- train posttrain \ + --config examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml +``` + +See [Post-training](posttraining.md) for the full SFT/LoRA workflow. + +--- + +## Related documentation + +- [Pretraining](pretraining.md): backend concepts, configuration walkthroughs, parallelism vocabulary, full configuration inventories. +- [Post-training](posttraining.md): SFT and LoRA via Megatron Bridge. +- [CLI reference](cli-reference.md): `direct` / `container` / `slurm` modes and flags. +- [Configuration system](configuration-system.md): YAML inheritance, overrides, image/env precedence. +- [Performance tuning](../04-technical-guides/performance-tuning.md): HipBLASLt autotuning, Primus-Turbo, FP8, MoE. diff --git a/docs/tuning_agent.md b/docs/02-user-guide/tuning-agent.md similarity index 89% rename from docs/tuning_agent.md rename to docs/02-user-guide/tuning-agent.md index e346de1d2..08baa04d2 100644 --- a/docs/tuning_agent.md +++ b/docs/02-user-guide/tuning-agent.md @@ -1,18 +1,18 @@ -# Tuning Agent +# Tuning agent The **Tuning Agent** is an LLM-driven search for a near-optimal Primus -**training configuration** — the full parallelism strategy *plus* the coupled +**training configuration**—the full parallelism strategy *plus* the coupled batching, pipeline-schedule, memory, MoE-communication, and precision knobs — on a target GPU cluster, **without running the workload at scale**. It drives the [Primus Projection](./projection.md) tool as an evaluation oracle. Projection -provides two estimates — **memory** and **performance** — each of which runs +provides two estimates—**memory** and **performance**—each of which runs **benchmark-anchored by default** (measuring what fits on a sub-node run and scaling the rest analytically) with a fully analytical **no-GPU `simulate`** fallback. The agent returns the configuration that maximizes `tokens/s/GPU` subject to a per-GPU memory safety margin. See [Knobs Searched](#knobs-searched) for the full set of levers it tunes. -- **Package**: [`primus/agents/tuning_agent/`](../primus/agents/tuning_agent/) +- **Package**: [`primus/agents/tuning_agent/`](../../primus/agents/tuning_agent/) - **Entry point**: `python -m primus.agents.tuning_agent` This document is both the **user/operator guide** (installation, configuration, @@ -22,36 +22,36 @@ features) for the agent. --- -## Table of Contents +## Table of contents -1. [Why a Tuning Agent](#why-a-tuning-agent) -2. [How It Works](#how-it-works) -3. [Knobs Searched](#knobs-searched) +1. [Why a tuning agent](#why-a-tuning-agent) +2. [How it works](#how-it-works) +3. [Knobs searched](#knobs-searched) 4. [Installation](#installation) -5. [LLM Setup](#llm-setup) +5. [LLM setup](#llm-setup) 6. [Quickstart](#quickstart) -7. [Execution Modes](#execution-modes) -8. [CLI Reference](#cli-reference) -9. [Target-Cluster YAML](#target-cluster-yaml) -10. [The Search Loop](#the-search-loop) -11. [Evaluator and Projection Modes](#evaluator-and-projection-modes) -12. [Output Artefacts](#output-artefacts) -13. [Worked Example](#worked-example) +7. [Execution modes](#execution-modes) +8. [CLI reference](#cli-reference) +9. [Target-cluster YAML](#target-cluster-yaml) +10. [The search loop](#the-search-loop) +11. [Evaluator and projection modes](#evaluator-and-projection-modes) +12. [Output artefacts](#output-artefacts) +13. [Worked example](#worked-example) 14. [Troubleshooting](#troubleshooting) 15. [Limitations](#limitations) -16. [Design Notes & Future Features](#design-notes--future-features) +16. [Design notes and future features](#design-notes-and-future-features) --- -## Why a Tuning Agent +## Why a tuning agent Choosing a training configuration for a large training (or inference) workload is a combinatorial problem. The configuration is the joint choice of the parallelism dimensions: -- **Data Parallel (DP)** — derived from world size and the other axes, +- **Data Parallel (DP)**—derived from world size and the other axes, - **Tensor Parallel (TP)**, -- **Expert Parallel (EP)** — for MoE models, +- **Expert Parallel (EP)**—for MoE models, - **Context Parallel (CP)**, - **Pipeline Parallel (PP)**, with virtual pipeline (**VPP**) and the **pipeline schedule** (1F1B / interleaved / zero-bubble / ZBV-\* / @@ -62,7 +62,7 @@ dimensions translate into in-flight work and memory pressure: global batch size (**GBS**), micro batch size (**MBS**), activation recomputation (`recompute_granularity`, `recompute_num_layers`), the overlap flags (`overlap_grad_reduce`, `overlap_param_gather`), and a set of higher-impact -levers — FP8 precision, MoE DeepEP / sync-free communication, fused +levers—FP8 precision, MoE DeepEP / sync-free communication, fused cross-entropy, and optimizer-state sharding (distributed optimizer / FSDP2). The full set is enumerated in [Knobs Searched](#knobs-searched). @@ -77,7 +77,7 @@ legal configuration within a user-specified budget. --- -## How It Works +## How it works ``` ┌────────────────────────────────────────────────────────────────────────┐ @@ -108,7 +108,7 @@ recompute for MBS, whether CP helps a given MoE shape). --- -## Knobs Searched +## Knobs searched The agent sweeps far more than the five parallelism dimensions. Its trial configuration (`TrialConfig` in `legality.py`) carries the full set of knobs @@ -117,7 +117,7 @@ agent leaves it unset) or **overridden for a trial** and translated into the corresponding `projection` flags by the evaluator. Each is legality-checked in code before it ever reaches the projection tool. -### Parallelism & batching +### Parallelism and batching | Knob | Legal values | What it controls | |------|--------------|------------------| @@ -143,11 +143,11 @@ code before it ever reaches the projection tool. |------|--------------|------------------| | `recompute_granularity` | `none` / `selective` / `full` | Activation recomputation strategy | | `recompute_num_layers` | int (≤ layers per VPP stage) | Layers recomputed per stage under `full` | -| `cross_entropy_loss_fusion` | `true` / `false` / inherit | Fused cross-entropy — large-vocab memory + compute win | +| `cross_entropy_loss_fusion` | `true` / `false` / inherit | Fused cross-entropy—large-vocab memory + compute win | | `use_distributed_optimizer` | `true` / `false` / inherit | ZeRO-1 optimizer-state sharding across DP | | `use_torch_fsdp2` | `true` / `false` / inherit | FSDP2 sharding (mutually exclusive with `use_distributed_optimizer`) | -### MoE communication — *MoE only, high impact* +### MoE communication—*MoE only, high impact* | Knob | Legal values | What it controls | |------|--------------|------------------| @@ -155,11 +155,11 @@ code before it ever reaches the projection tool. | `sync_free_stage` | `0` / `1` / `2` / `3` | Sync-free MoE pipelining; stage ≥ 2 auto-enables DeepEP | | `target_ep_size` | positive int / inherit | EP override used for All-to-All modeling | -### Precision — *high impact* +### Precision—*high impact* | Knob | Legal values | What it controls | |------|--------------|------------------| -| `fp8` | `none` / `hybrid` (also `e4m3`, `delayed`) | FP8 on linear layers — roughly 2× compute on GEMMs | +| `fp8` | `none` / `hybrid` (also `e4m3`, `delayed`) | FP8 on linear layers—roughly 2× compute on GEMMs | ### Coupling rules enforced in code @@ -206,10 +206,10 @@ Origami) are only needed by the evaluator paths you actually use: --- -## LLM Setup +## LLM setup The agent uses [DSPy](https://dspy.ai), which routes LLM calls through -[LiteLLM](https://docs.litellm.ai/docs/providers) internally — **no separate +[LiteLLM](https://docs.litellm.ai/docs/providers) internally—**no separate proxy process is required**. Set credentials for whichever provider you use: ```bash @@ -266,7 +266,7 @@ python -m primus.agents.tuning_agent \ --- -## Execution Modes +## Execution modes The agent has two orthogonal mode switches: **what the evaluator does** (`--mode`) and **whether the LLM stage runs** (`--seed-only` / `--agent-only`). @@ -293,7 +293,7 @@ The agent has two orthogonal mode switches: **what the evaluator does** --- -## CLI Reference +## CLI reference ```bash python -m primus.agents.tuning_agent \ @@ -322,15 +322,15 @@ python -m primus.agents.tuning_agent \ --- -## Target-Cluster YAML +## Target-cluster YAML A thin wrapper around the existing Primus `hardware_config` convention, so no new networking format has to be invented; topology, bandwidths, and latencies are consumed by the analytical communication model (see -[`projection.md` → Communication Modeling](./projection.md#communication-modeling)). +[`projection.md` → Assumptions (performance projection)](./projection.md#assumptions-performance-projection)). A complete example ships at -[`examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml`](../examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml): +[`examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml`](../../examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml): ```yaml target_cluster: @@ -396,7 +396,7 @@ agent: --- -## The Search Loop +## The search loop 1. **Resolve the workload.** Load the workload YAML, follow `modules.pre_trainer.model` into `primus/configs/models/megatron/.yaml`, @@ -437,7 +437,7 @@ agent: --- -## Evaluator and Projection Modes +## Evaluator and projection modes The evaluator wraps the Primus Projection CLI behind a uniform interface, so the agent does not need to know which mode produced a number: @@ -463,21 +463,21 @@ pre-filter to reject infeasible configs before paying for a performance call, then `simulate` (or `benchmark` for promising candidates when a GPU is available). The tool belt exposed to the LLM mirrors this: -- `evaluate_memory_only(config_json)` — cheap pre-filter -- `evaluate_simulate(config_json)` — primary scoring path -- `evaluate_with_benchmark(config_json)` — only if `has_gpu: true` +- `evaluate_memory_only(config_json)`—cheap pre-filter +- `evaluate_simulate(config_json)`—primary scoring path +- `evaluate_with_benchmark(config_json)`—only if `has_gpu: true` - `get_history`, `get_best`, `get_legal_axes`, `get_architecture`, `get_cluster`, `get_budget_status` - `note_to_scratchpad`, `read_scratchpad` -- `query_llm(prompt, system?)` — one-shot "LLM-inside-LLM" consultation +- `query_llm(prompt, system?)`—one-shot "LLM-inside-LLM" consultation -For the underlying projection math — memory components, the simulate vs. +For the underlying projection math—memory components, the simulate vs. benchmark trade-off, and the benchmark-based memory projection the agent uses -for OOM-accurate feasibility — see [`projection.md`](./projection.md). +for OOM-accurate feasibility—see [`projection.md`](./projection.md). --- -## Output Artefacts +## Output artefacts Everything lands in `--out-dir`: @@ -509,7 +509,7 @@ The run also prints the best configuration and ready-to-paste exports: --- -## Worked Example +## Worked example Search for the best Mixtral 8×22B configuration on a 4-node MI355X pod, with a single idle 8-GPU node available for benchmarking: @@ -566,17 +566,17 @@ These are honest caveats, not future features: `memory_safety_margin` compensates conservatively; the benchmark-based memory path (see `projection.md`) closes most of this gap. 3. **Search-space explosion** with all axes on. The agent mitigates with an - impact-ordered deterministic seed plan — high-leverage levers first + impact-ordered deterministic seed plan—high-leverage levers first (recompute, MoE DeepEP / sync-free, FP8, then schedule and VPP/MBS neighbors), with the broad TP×PP×EP×CP grid evaluated last and capped by - `--seed-budget` — so the LLM starts from an informed incumbent and spends + `--seed-budget`—so the LLM starts from an informed incumbent and spends its budget on polish. 4. **Cluster-description lossiness**: averaged bandwidth/latency cannot capture contention or per-rail asymmetry. --- -## Design Notes & Future Features +## Design notes and future features > This section captures the design rationale and the paper-ready problem > statement, plus the list of features deliberately deferred from v1. @@ -584,8 +584,8 @@ These are honest caveats, not future features: ### Problem statement (paper-ready) We address the problem of **automatically selecting a near-optimal -training configuration — the parallelism strategy plus the coupled batching, -schedule, memory, MoE-communication, and precision knobs — for a large-scale +training configuration—the parallelism strategy plus the coupled batching, +schedule, memory, MoE-communication, and precision knobs—for a large-scale distributed training workload on a target GPU cluster, without executing the workload at scale**. The configuration space is combinatorial: for each axis only a small set of values @@ -608,7 +608,7 @@ We formulate the search as **LLM-as-policy over a hybrid analytical / benchmark-driven evaluator**. The evaluator is the Primus Projection tool, which provides three calls of increasing fidelity and cost: (i) a **memory projection** that runs either analytically (no GPU) or, by default, -benchmark-anchored — measuring the real per-rank peak on a sub-node run and +benchmark-anchored—measuring the real per-rank peak on a sub-node run and extrapolating it for an OOM-accurate estimate; (ii) a fully analytical **performance projection** built on the Origami GEMM model and an SDPA simulator; and (iii) a **hybrid benchmark** that measures per-layer compute on @@ -622,7 +622,7 @@ to a configurable per-GPU memory safety margin. The contribution is a configuration-search methodology that exploits an LLM's ability to reason over architectural priors (topk dominance of MoE activations, the impact of MQA on attention activation, the inter-node/intra-node boundary -for All-to-All) to direct an analytical oracle — replacing exhaustive sweeps on +for All-to-All) to direct an analytical oracle—replacing exhaustive sweeps on real hardware with a small number of informed analytical evaluations. The system runs either entirely on a CPU-only host (simulate backend only) or in a mixed mode where a small number of real-hardware benchmark runs calibrate the @@ -640,7 +640,7 @@ analytical predictions. 4. Agent search over **all** parallelism and coupled axes (TP, PP, EP, CP, MBS, GBS, VPP, pipeline schedule, recompute, overlap flags) plus the higher-impact levers (FP8, MoE DeepEP / sync-free, fused cross-entropy, - distributed-optimizer / FSDP2) — restricted by per-architecture legality. + distributed-optimizer / FSDP2)—restricted by per-architecture legality. See [Knobs Searched](#knobs-searched) for the complete list. 5. Two evaluator paths: a **no-GPU path** (`projection memory --memory-mode simulate` + `projection performance --profiling-mode simulate`), always @@ -657,60 +657,60 @@ analytical predictions. These are recorded so they are not lost; they are deliberately left out to keep the agent small. -- **F1. Multi-objective Pareto** — replace the scalar `tokens/s/GPU` objective +- **F1. Multi-objective Pareto**—replace the scalar `tokens/s/GPU` objective with a Pareto frontier over (throughput, MFU, memory headroom, projected $/token, projected energy/token). -- **F2. Online cluster-spec retrieval** — pull the cluster description from an +- **F2. Online cluster-spec retrieval**—pull the cluster description from an internal registry or known-archs catalog (MI300X / MI325X / MI355X reference pods) and fall back to user overrides; optionally infer topology from a small DCGM / ROCm-SMI dump. -- **F3. Persistent memory / configuration cache** — cache +- **F3. Persistent memory / configuration cache**—cache `(model_signature, cluster_signature) → best_known_configs` across runs; invalidate when ROCm / hipBLASLt / framework versions change. -- **F4. Agent-proposed scale-downs and microbenchmarks** — let the agent +- **F4. Agent-proposed scale-downs and microbenchmarks**—let the agent *propose* reduced-model proxies and targeted microbenchmarks to reduce the uncertainty of its current top-k (reusing the `moe_proxy_single_node.yaml` pattern). -- **F5. Telemetry plug-ins (rocprofiler / TraceLens / Magpie)** — after a +- **F5. Telemetry plug-ins (rocprofiler / TraceLens / Magpie)**—after a benchmark run, optionally extract per-kernel time, GEMM efficiency, A2A bytes, NIC utilization to calibrate the analytical models and explain underperformance back to the agent. Exposed as a `SKILL.md`-described plug-in. References: [TraceLens](https://github.com/AMD-AGI/TraceLens-internal), [Magpie](https://github.com/AMD-AGI/Magpie). -- **F6. Calibration learning** — under `--profiling-mode both`, record +- **F6. Calibration learning**—under `--profiling-mode both`, record per-(model, arch, dim) residuals between simulate and benchmark, fit a small correction model, and report a confidence band on subsequent simulate runs. -- **F7. Robustness / sensitivity report** — for the winning config, sweep ±1 +- **F7. Robustness / sensitivity report**—for the winning config, sweep ±1 step on each axis and report whether the optimum is a sharp peak or a broad basin (cheap; only `simulate` calls). -- **F8. Cross-axis priors as DSPy modules** — a library of tunable "rules of +- **F8. Cross-axis priors as DSPy modules**—a library of tunable "rules of thumb" that DSPy's optimizer can refine over time using the trial logs. -- **F9. Sub-agent "test-proposer"** — delegate targeted experiments (e.g. a +- **F9. Sub-agent "test-proposer"**—delegate targeted experiments (e.g. a 2-layer scale-down forward+backward microbenchmark, or a stand-alone A2A probe at the proposed EP × hidden_size × topk), profiling the *test* with explicit synchronisation rather than the sandbox. A `run_proposed_experiment(plan, code)` tool can be added to `tools.py` without restructuring the loop. -- **F10. Sub-LLM expert router** — extend the existing `query_llm` tool into a +- **F10. Sub-LLM expert router**—extend the existing `query_llm` tool into a *named expert* router (`query_llm(expert='moe', …)`). ### Known design holes -1. **Simulator-vs-reality gap** — see Limitations above; in no-GPU mode the +1. **Simulator-vs-reality gap**—see Limitations above; in no-GPU mode the agent reports a confidence caveat. -2. **Memory-projection blind spots** — A2A buffers, allocator fragmentation, +2. **Memory-projection blind spots**—A2A buffers, allocator fragmentation, and comm scratch are not modeled analytically; the benchmark-based memory projection closes most of this gap by anchoring on a measured peak. -3. **Search-space explosion** — mitigated by an impact-ordered deterministic +3. **Search-space explosion**—mitigated by an impact-ordered deterministic seed plan (high-leverage levers first, broad parallelism grid last) plus LLM-guided search from the seed incumbent. -4. **Cluster-description lossiness** — averaged bandwidth/latency cannot capture +4. **Cluster-description lossiness**—averaged bandwidth/latency cannot capture contention or per-rail asymmetry. --- -## Related Documentation +## Related documentation -- [Projection](./projection.md) — memory + performance projection internals, +- [Projection](./projection.md)—memory + performance projection internals, including the benchmark-based memory projection the agent relies on. -- [Tuning Agent package README](../primus/agents/tuning_agent/README.md) — +- [Tuning Agent package README](../../primus/agents/tuning_agent/README.md) — quickstart reference inside the source tree. diff --git a/docs/03-configuration-reference/README.md b/docs/03-configuration-reference/README.md new file mode 100644 index 000000000..487d6e526 --- /dev/null +++ b/docs/03-configuration-reference/README.md @@ -0,0 +1,13 @@ +# Configuration reference + +Parameter references for Primus presets, backend-facing keys, and commonly used environment variables. + +- [Megatron parameters](megatron-parameters.md): Megatron-LM backend YAML parameters and Primus overrides +- [TorchTitan parameters](torchtitan-parameters.md): Primus TorchTitan preset keys and common JobConfig fields +- [MaxText parameters](maxtext-parameters.md): Primus MaxText overlay defaults and common fields +- [Megatron Bridge parameters](megatron-bridge-parameters.md): Megatron Bridge recipe, SFT, and pretraining fields surfaced through Primus +- [Environment variables](environment-variables.md): practical reference for commonly encountered environment variables + +--- + +[← Documentation home](../README.md) diff --git a/docs/03-configuration-reference/environment-variables.md b/docs/03-configuration-reference/environment-variables.md new file mode 100644 index 000000000..d18ee8ce2 --- /dev/null +++ b/docs/03-configuration-reference/environment-variables.md @@ -0,0 +1,239 @@ +# Environment variables reference + +This document catalogs the main environment variables you may encounter when running Primus on AMD GPUs: distributed launchers, Primus runners and CLI, YAML substitution, libraries (NCCL/RCCL, ROCm, PyTorch, JAX), and optional integrations (Hugging Face, WandB, MLflow). It is a practical reference, not a complete list of every variable accepted by upstream libraries. + +**Legend** + +- **Required**: Must be set for the stated workflow; otherwise the job fails or mis-ranks. +- **Optional**: Has a safe default or is only needed for specific features. +- **Set by**: Typical source (launcher, `runner/helpers/envs/*.sh`, user shell, container host). +- **Used in**: Representative Primus paths; many variables are also read by NVIDIA NCCL, AMD RCCL, PyTorch, or JAX without Primus wrapping them. + +--- + +## 1. PyTorch distributed + +Set by `torchrun`, Slurm launchers, or `runner/primus-cli-direct.sh` / `runner/primus-cli-slurm-entry.sh`. Consumed by PyTorch distributed, RCCL, and Primus helpers. + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `MASTER_ADDR` | `localhost` (direct / `base_env.sh`) | User, Slurm entry (`primus-cli-slurm-entry.sh`), or validation fallback (`runner/lib/validation.sh`) | `primus/pretrain.py`, `primus/core/base_module.py`, `primus/core/utils/env.py`, `primus/tools/preflight/network/network_probe.py`, PyTorch rendezvous | Rendezvous hostname or IP for process group initialization. **Required** for multi-node if not using Slurm auto-detection. | +| `MASTER_PORT` | `1234` (direct), `29500` in some Python defaults | Config / CLI / user | Same as `MASTER_ADDR`; `validation.sh` enforces 1024–65535 | TCP port for the store backing `torch.distributed`. | +| `RANK` | `0` if unset in helpers | `torchrun` | `primus/tools/utils.py`, `primus/tools/preflight/global_vars.py`, projection and profiler code | Global rank index. | +| `WORLD_SIZE` | `1` | `torchrun` | Preflight, projection, `primus/core/base_module.py` | Total number of processes. | +| `LOCAL_RANK` | `0` | `torchrun` | `primus/core/base_module.py`, GPU selection in benchmarks and trainers | GPU index on this node. | +| `LOCAL_WORLD_SIZE` | `1` (Python) / `8` in benchmarks default | `torchrun` | `primus/tools/preflight/*.py`, `strided_allgather_bench.py` | Processes (GPUs) per node. | +| `NODE_RANK` | `0` | `primus-cli-direct` / `primus-cli-slurm-entry.sh` | `primus/pretrain.py`, logging in `runner/lib/common.sh` | Zero-based node index in multi-node jobs. | +| `NNODES` | `1` | Direct config (`runner/.primus.yaml`), `primus-cli-slurm-entry.sh` | `primus/pretrain.py`, `primus/core/projection/training_config.py` | Number of nodes in the job. | +| `GPUS_PER_NODE` | `8` | `runner/.primus.yaml` direct section, `primus-cli-slurm-entry.sh`, `validation.sh` | `primus/core/projection/module_profilers/*.py`, training config helpers | GPUs per node for world-size math and binding. | + +--- + +## 2. Primus core + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `PRIMUS_PATCHES` | `""` / `"all"` | User | `primus/core/patches/patch_runner.py` | `"all"` or empty enables all patches; `"none"` disables; comma list enables subset. | +| `PRIMUS_LOG_LEVEL` | `INFO` | User; debug paths in `runner/primus-cli-*.sh` set `DEBUG` | `runner/lib/common.sh` | Log verbosity: `DEBUG`, `INFO`, `WARN`, `ERROR`. | +| `PRIMUS_LOG_TIMESTAMP` | `1` | User | `runner/lib/common.sh` | `1` prefixes logs with timestamps; `0` disables. | +| `PRIMUS_LOG_COLOR` | `1` (auto-off if not a TTY) | User; tests may set `0` | `runner/lib/common.sh` | ANSI colors in runner logs. | +| `PRIMUS_DEBUG` | `0` | User | `runner/helpers/envs/primus-env.sh` | `1` enables `set -x` in the env loader for shell tracing. | +| `PRIMUS_SKIP_VALIDATION` | `0` | User / tests | `runner/helpers/envs/primus-env.sh` | `1` skips `validate_distributed_params` (not recommended). | +| `PRIMUS_EXPECT_IB` | (unset) | User | `primus/tools/preflight/network/network_standard.py` | When `1`, preflight treats InfiniBand as expected for validation. | +| `PRIMUS_CLUSTER` | `amd-aig-poolside` (CLI default) | User | `primus/tools/benchmark/rccl_bench_args.py` | Cluster label for RCCL benchmark tooling. | +| `PRIMUS_GPU_ARCH` | (auto / `"mi300x"` in simulators) | User / CLI | `primus/core/projection/simulation_backends/origami_backend.py`, `sdpa_simulator.py`, `projection.py` CLI | GPU architecture string for performance projection. | +| `PRIMUS_GPU_CLOCK_MHZ` | (unset) | User | Same as `PRIMUS_GPU_ARCH` | Optional clock override for projection. | +| `PRIMUS_GPU_DEVICE` | `0` | User | `origami_backend.py` | GPU index for hardware detection in projection. | +| `PRIMUS_GEMM_BACKEND` | (unset) | User | `primus/core/projection/simulation_backends/factory.py` | Selects GEMM simulation backend by name. | +| `PRIMUS_PREFLIGHT_MIN_FREE_MEM_GB` | `1` | User | `primus/tools/preflight/gpu/utils.py` | Minimum free GPU memory (GB) for preflight checks. | +| `PRIMUS_PREFLIGHT_MIN_TFLOPS` | `10.0` | User | `primus/tools/preflight/gpu/utils.py` | Minimum TFLOPS threshold for preflight GEMM checks. | +| `PRIMUS_TURBO_AUTO_TUNE` | (unset) | User / tests | `tests/trainer/test_megatron_trainer.py` (integration) | Enables Turbo auto-tuning in supported Turbo/Megatron test flows; not referenced in core `primus/` Python outside tests. **Optional**. | +| `PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND` | `TURBO` | User; hooks may set `DEEP_EP` | `primus/backends/megatron/patches/args/rocm_arg_validation.py`, `examples/run_pretrain.sh`, `runner/helpers/hooks/05_using_uep.sh` | MoE dispatch/combine backend selector. | + +--- + +## 3. Primus YAML substitution + +Parsed by `primus/core/config/yaml_loader.py` for patterns `${VAR}` (required) and `${VAR:default}` (optional). Typical experiment YAMLs under `examples/` use these for sweep-friendly overrides. + +| Variable | Typical default in YAML | Where set | Where used | Description | +|----------|-------------------------|-----------|------------|-------------| +| `PRIMUS_TEAM` | `"amd"` | User | Resolved before module merge in experiment YAML | Work group / team segment in paths. | +| `PRIMUS_USER` | `"root"` | User | Experiment YAML | User name segment. | +| `PRIMUS_EXP_NAME` | per-example | User | Experiment YAML | Experiment folder name. | +| `PRIMUS_WORKSPACE` | `"./output"` | User | Experiment YAML | Root workspace for artifacts. | +| `PRIMUS_TP` | `1` | User | Megatron example YAMLs | `tensor_model_parallel_size` override. | +| `PRIMUS_PP` | `1` | User | Megatron example YAMLs | `pipeline_model_parallel_size` override. | +| `PRIMUS_EP` | `1` | User | Megatron example YAMLs | `expert_model_parallel_size` override. | +| `PRIMUS_SEQ_LENGTH` | per-model | User | Megatron example YAMLs | Sequence length override. | +| `PRIMUS_MAX_POSITION_EMBEDDINGS` | `4096` or `131072` | User | `examples/megatron/**/*.yaml`, tests | Position embedding cap override. | +| `PRIMUS_GLOBAL_BATCH_SIZE` | per-model | User | Megatron example YAMLs | Global batch override. | +| `PRIMUS_NUM_LAYERS` | per-model | User | Tests and MoE examples | Transformer layer count override. | +| `PRIMUS_MOE_LAYER_FREQ` | MoE patterns | User | MoE examples / tests | MoE layer frequency pattern. | +| `PRIMUS_TOKENIZED_DATA_PATH` | `null` | User | Megatron examples | Path to tokenized training data. | +| `PRIMUS_MODEL` | per-stack | User | Megatron examples | Model preset stem (e.g. `llama3_8B`). | +| `PRIMUS_VPP` | `null` | User | `tests/trainer/test_megatron_trainer.yaml` | Virtual pipeline stages override. | + +--- + +## 4. NCCL / RCCL + +Primus seeds many of these in `runner/helpers/envs/base_env.sh`. RCCL honors NCCL-compatible variables on AMD GPUs. See [NCCL environment](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html) and [RCCL environment](https://rocm.docs.amd.com/projects/rccl/en/develop/api-reference/env-variables.html). + +| Variable | Default (Primus base) | Where set | Where used | Description | +|----------|------------------------|-----------|------------|-------------| +| `NCCL_DEBUG` | unset | User / `base_env.sh` empty default | Preflight reports, RCCL runtime | Log verbosity: `NONE`, `WARN`, `INFO`, `TRACE`, etc. **Optional** unless debugging comms. | +| `NCCL_SOCKET_IFNAME` | derived from `IP_INTERFACE` | `base_env.sh` | `primus/tools/preflight/network/*.py`, GPU topology helpers | Socket NIC for host networking. | +| `GLOO_SOCKET_IFNAME` | same as NCCL if unset | `base_env.sh` | Preflight | Gloo TCP backend interface. | +| `NCCL_IB_HCA` | auto via `get_nccl_ib_hca.sh` if empty | `base_env.sh`, container passthrough | Preflight, multi-node tuning | InfiniBand HCAs to use. | +| `NCCL_IB_GID_INDEX` | `3` | `base_env.sh` | RCCL | GID index for IB/RoCE; many sites use `1` for RoCE v2 (override as needed). | +| `NCCL_IB_TC` | (unset) | User | RCCL | InfiniBand traffic class. | +| `NCCL_IB_FIFO_TC` | (unset) | User | RCCL | InfiniBand FIFO traffic class. | +| `NCCL_IB_ROCE_VERSION_NUM` | (unset) | User | RCCL | RoCE version selection. | +| `NCCL_PXN_DISABLE` | `1` | `base_env.sh` | RCCL | Disable PXN (PCIe cross-NIC); set `0` to enable. | +| `NCCL_P2P_NET_CHUNKSIZE` | `524288` | `base_env.sh` | RCCL | P2P network chunk size tuning. | +| `NCCL_PROTO` | (unset) | User | RCCL | Protocol selection (e.g. `Simple`, `LL`, `LL128`). | +| `NCCL_CROSS_NIC` | `0` | `base_env.sh` | RCCL | Cross-NIC communication policy. | +| `NCCL_IB_RETRY_CNT` | (unset) | User | RCCL | IB retry count. | +| `NCCL_IB_TIMEOUT` | (unset) | User | RCCL | IB timeout. | +| `NCCL_NET_GDR_LEVEL` | (unset) | User | Preflight summaries | GPUDirect RDMA level. | +| `NCCL_IB_DISABLE` | `0` | User / env | Preflight | Disable IB; use sockets only. | +| `NCCL_DMABUF_ENABLE` | (unset) | User | RCCL | DMA-BUF registration path. | +| `NCCL_IGNORE_CPU_AFFINITY` | (unset) | User | RCCL | Ignore CPU affinity hints. | +| `NCCL_IB_QPS_PER_CONNECTION` | (unset) | User | RCCL | IB QPs per connection. | +| `NCCL_MAX_P2P_CHANNELS` | (unset) | User | RCCL | Cap P2P channels. | +| `NCCL_GDR_FLUSH_DISABLE` | (unset) | User | RCCL | Disable GDR flush. | +| `NCCL_IB_USE_INLINE` | (unset) | User | RCCL | Inline IB sends. | +| `NCCL_NET_PLUGIN` | (unset) | User | RCCL | Alternate network plugin (e.g. `librccl-anp.so`). | +| `RCCL_MSCCL_ENABLE` | `0` | `base_env.sh` | RCCL | Enable MSCCL algorithms. | +| `RCCL_MSCCLPP_THRESHOLD` | `1GiB` default | `base_env.sh` | RCCL | MSCCL++ message-size threshold. | +| `RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING` | `0` in hooks | `runner/helpers/hooks/03_enable_ainic.sh` | RCCL | Stricter GDR flush memory ordering; relevant for some NIC/GPU combos. | +| `TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK` | `0` | `base_env.sh` | PyTorch + RCCL | Tensor allocator hook for NCCL registration. | +| `TORCH_NCCL_HIGH_PRIORITY` | `1` | `base_env.sh` | PyTorch | High-priority NCCL streams. | + +--- + +## 5. ROCm / HSA / HIP + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `HSA_ENABLE_SDMA` | `1` | `base_env.sh` | ROCm runtime | Enable SDMA engines for copies. | +| `HSA_NO_SCRATCH_RECLAIM` | `1` | `base_env.sh`, container passthrough | ROCm runtime; documented for MoE stability | `1` keeps scratch allocated (often used for MoE stability). See [ROCR environment](https://rocm.docs.amd.com/projects/ROCR-Runtime/en/docs-7.1.1/environment_variables.html). | +| `HIP_VISIBLE_DEVICES` | `0..GPUS_PER_NODE-1` | `base_env.sh` | ROCm device visibility | Restricts which GPU indices ROCm exposes. | +| `ROCBLAS_DEFAULT_ATOMICS_MODE` | (unset) | User | `primus/backends/megatron/patches/args/rocm_arg_validation.py` | Read for deterministic / accuracy-sensitive GEMM behavior. | + +--- + +## 6. CUDA / PyTorch + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `CUDA_DEVICE_MAX_CONNECTIONS` | `1` | `base_env.sh`; Megatron patches may adjust | `primus/backends/megatron/patches/env_patches.py`, Megatron patches | Limits concurrent CUDA connections; often `1` for TP/PP overlap. | +| `TORCH_COMPILE_DISABLE` | `0` | User | `primus/backends/megatron/patches/args/rocm_arg_validation.py` | Disable `torch.compile` when `1`. | + +--- + +## 7. Transformer engine + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `NVTE_ROCM_ENABLE_MXFP8` | `1` | `base_env.sh` | Transformer Engine on ROCm | Enable MXFP8 paths. | +| `NVTE_CK_USES_BWD_V3` | `1` | `base_env.sh`, container passthrough | TE / CK | Use CK backward v3 kernels. | +| `NVTE_CK_IS_V3_ATOMIC_FP32` | (unset; examples print `0`) | User / `examples/run_pretrain.sh`, container passthrough | TE / CK | Atomic FP32 mode for CK v3 backward. | +| `PATCH_TE_FLASH_ATTN` | `0` | `base_env.sh`, container passthrough | `runner/helpers/hooks/01_patch_te_flash_attn_max_version.sh` | Trigger TE flash-attn patch hook when `1`. | + +--- + +## 8. Caches and authentication + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `HF_HOME` | `${DATA_PATH}/huggingface` | `base_env.sh`, `primus/core/utils/env_setup.py`, `primus/pretrain.py` | Hugging Face libraries | Cache for models and datasets. | +| `HF_TOKEN` | (unset) | User, container passthrough | Hugging Face Hub | Auth for gated models. **Required** for private/gated assets. | +| `TORCH_HOME` | under workspace | `primus/core/utils/env_setup.py` | PyTorch Hub | Torch Hub cache root. | +| `TRANSFORMERS_CACHE` | aligned with HF layout | `primus/core/utils/env_setup.py` | `transformers` | Model cache for Transformers. | +| `WANDB_API_KEY` | (unset) | User, container passthrough | WandB client, Megatron trainer checks | API key for logging. **Required** for WandB when enabled. | +| `WANDB_PROJECT` | (unset) | User / TorchTitan patch | `primus/backends/torchtitan/patches/wandb_patches.py` | Project name. | +| `WANDB_RUN_NAME` | (unset) | User / patches | Same | Run display name. | +| `WANDB_TEAM` | (unset) | User | TorchTitan metrics (entity) | WandB team/entity. | +| `DATABRICKS_HOST` | (unset) | User | `mlflow` client (via `primus/backends/megatron/training/global_vars.py` MLflow setup) | Required for Databricks-hosted MLflow when MLflow logging is enabled. | +| `DATABRICKS_TOKEN` | (unset) | User | Databricks APIs | Auth token paired with host. | +| `MLFLOW_TRACKING_URI` | (unset) | User | `mlflow` (via Megatron integrations) | MLflow tracking server URI. **Optional** unless using MLflow. | +| `MLFLOW_REGISTRY_URI` | (unset) | User | MLflow | Model registry endpoint. | +| `NLTK_DATA` | (unset) | User | `runner/helpers/hooks/train/pretrain/megatron/preprocess_data.py`, Megatron-LM tools | Punkt and other tokenizer data for preprocessing. | +| `TOKENIZED_DATA_PATH` | per-hook default | User | `runner/helpers/hooks/train/pretrain/megatron/prepare.py` | Pre-tokenized dataset location for Megatron data prep hooks. | + +--- + +## 9. HipBLASLt tuning + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `PRIMUS_HIPBLASLT_TUNING` | `0` | User | `examples/run_pretrain.sh` | **Master switch** for the HipBLASLt tuning flow (`1` enables). Must be set before `PRIMUS_HIPBLASLT_TUNING_STAGE` takes effect, and is mutually exclusive with deterministic mode (`PRIMUS_DETERMINISTIC=1`). | +| `PRIMUS_HIPBLASLT_TUNING_STAGE` | `0` | User | `examples/run_pretrain.sh` | Stages `0` off, `1` dump shapes, `2` offline tune, `3` apply tuned kernels. | +| `HIPBLASLT_TUNING_OVERRIDE_FILE` | (unset) | User / tuning scripts | `examples/run_pretrain.sh` | Path to tuned-kernel override file for stage `3`. | +| `TE_HIPBLASLT_TUNING_RUN_COUNT` | varies | User | `examples/run_pretrain.sh` | Number of benchmark runs per shape during TE HipBLASLt tuning. | +| `TE_HIPBLASLT_TUNING_ALGO_COUNT` | varies | User | `examples/run_pretrain.sh` | Transformer Engine HipBLASLt search breadth. | +| `TE_HIPBLASLT_TUNING_ALGO_FILE` | (unset) | User | TE + HipBLASLt | Algorithm file for TE tuning flows. | +| `TE_HIPBLASLT_TUNING` | (unset) | User | `examples/run_pretrain.sh` | When set, interacts with deterministic mode and tuning stages (disable conflicting modes per script comments). | +| `HIPBLASLT_LOG_LEVEL` | (unset) | User | HipBLASLt | Library log level. | +| `HIPBLASLT_LOG_MASK` | (unset) | User | HipBLASLt | Bitmask for log categories. | + +--- + +## 10. Build and rebuild + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `REBUILD_PRIMUS_TURBO` | `0` | User, container passthrough | `runner/helpers/hooks/00_rebuild_primus_turbo.sh` | `1` rebuilds Primus-Turbo on startup. | +| `REBUILD_BNXT` | `0` | User, container passthrough | `runner/helpers/hooks/02_rebuild_bnxt.sh` | `1` rebuilds BNXT driver artifacts when packaged. | +| `USING_AINIC` | (unset) | User | `runner/helpers/hooks/03_enable_ainic.sh` | `1` enables AINIC-oriented networking hooks. | +| `MAX_JOBS` | (unset) | User / tooling | `tools/daily/safe_wrapper.py` | Parallel compile jobs for pip builds. | +| `BACKEND_PATH` | (unset) | User | `primus/pretrain.py`, `primus/core/backend/backend_adapter.py` | Override checkout path for third-party backends (Megatron, TorchTitan, MaxText). | + +--- + +## 11. Container passthrough + +`runner/.primus.yaml` lists names forwarded from the host into training containers (`container.options.env`). Primus does not assign values here; it only whitelists keys for `--env` forwarding. + +Forwarded keys: + +`MASTER_ADDR`, `MASTER_PORT`, `NNODES`, `NODE_RANK`, `GPUS_PER_NODE`, `DOCKER_IMAGE`, `HF_TOKEN`, `WANDB_API_KEY`, `ENABLE_NUMA_BINDING`, `REBUILD_PRIMUS_TURBO`, `USING_AINIC`, `PATCH_TE_FLASH_ATTN`, `REBUILD_BNXT`, `HSA_NO_SCRATCH_RECLAIM`, `NVTE_CK_USES_BWD_V3`, `GPU_MAX_HW_QUEUES`, `HSA_KERNARG_POOL_SIZE`, `PRIMUS_TURBO_DEEPEP_TIMEOUT`, `NCCL_IB_HCA`, `NCCL_SOCKET_IFNAME`, `GLOO_SOCKET_IFNAME`, `NCCL_IB_GID_INDEX`, `PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32`, `NVTE_CK_IS_V3_ATOMIC_FP32`, `PATH_TO_BNXT_TAR_PACKAGE`, `ANP_HOME_DIR`, `RCCL_HOME_DIR`, `MPI_HOME_DIR`, `DUMP_HLO`, `DUMP_HLO_DIR`, `PRIMUS_DETERMINISTIC`, `PRIMUS_HIPBLASLT_TUNING`, `PRIMUS_HIPBLASLT_TUNING_STAGE`, `TE_HIPBLASLT_TUNING_RUN_COUNT`, `TE_HIPBLASLT_TUNING_ALGO_COUNT`, `HIPBLASLT_LOG_MASK`, `HIPBLASLT_LOG_FILE`, `HIPBLASLT_LOG_LEVEL`, `HIPBLASLT_TUNING_OVERRIDE_FILE` + +--- + +## 12. Slurm + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `SLURM_NNODES` / `SLURM_JOB_NUM_NODES` | job-dependent | Slurm | `primus-cli-slurm-entry.sh` (`NNODES` export), preflight probes | Node count for the allocation. | +| `SLURM_NODEID` | job-dependent | Slurm | Mapped to `NODE_RANK` in `primus-cli-slurm-entry.sh` | Node index. | +| `SLURM_PROCID` | job-dependent | Slurm | Fallback for `NODE_RANK` when `SLURM_NODEID` is unset | Process id within the Slurm step (entry script). | +| `SLURM_JOB_ID` | job-dependent | Slurm | `primus/tools/preflight/host/host_probe.py` | Job identifier string. | + +--- + +## 13. Debug and pipeline + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `DUMP_PP_DIR` | `output/pp_data` | User | `primus/backends/megatron/megatron_pretrain_trainer.py`, `primus/backends/megatron/patches/pp_dump_data_patches.py` | Directory for pipeline-parallel debug dumps. | +| `DEBUG_SIMULATOR` | `0` | User | `primus/core/projection/performance_projection/simulator.py` | `1` enables verbose projection simulator logging. | +| `RECORD_OFFLOAD_MEMORY_INFO` | `0` | User | `primus/core/pipeline_parallel/handler/offload_handler.py` | Record offload memory stats when `1`. | +| `RECORD_OFFLOAD_MEMORY_INFO_DIR` | `output` | User | `primus/core/pipeline_parallel/scheduler/scheduler.py` | Output directory for offload memory logs. | +| `USE_PINNED_OFFLOAD` | `0` | User | `offload_handler.py` | Use pinned host memory for offload buffers when `1`. | + +--- + +## 14. JAX / XLA (MaxText) + +Primus MaxText hooks print recommended values in `runner/helpers/hooks/train/pretrain/maxtext/prepare.py`; MaxText and JAX read them directly. + +| Variable | Default | Where set | Where used | Description | +|----------|---------|-----------|------------|-------------| +| `XLA_PYTHON_CLIENT_MEM_FRACTION` | e.g. `.97` in prepare hook | User / hook output | JAX / XLA allocator | Fraction of GPU memory pre-allocated for JAX. | +| `DUMP_HLO_DIR` | `${PRIMUS_PATH}/output/xla_dump_hlo` (example) | User | XLA via `XLA_FLAGS` composition | Directory for HLO dumps when enabled. | +| `DUMP_HLO` | `0` | User | Prepare hook → XLA flags | Gate HLO dumping (`1` enables in hook samples). | + +**Note:** MaxText also propagates many knobs through `XLA_FLAGS` and `LIBTPU_INIT_ARGS` upstream; see MaxText sources for the full matrix. diff --git a/docs/03-configuration-reference/maxtext-parameters.md b/docs/03-configuration-reference/maxtext-parameters.md new file mode 100644 index 000000000..b64bb332e --- /dev/null +++ b/docs/03-configuration-reference/maxtext-parameters.md @@ -0,0 +1,146 @@ +# MaxText backend configuration reference + +Primus routes experiment YAML into the [MaxText](https://maxtext.readthedocs.io/) stack (JAX / XLA). Configuration is a **flat map of keys** (no nested `training.`* trees like TorchTitan): Primus merges module and model presets, writes a temporary YAML, and MaxText’s `pyconfig.initialize` loads it on top of upstream defaults. + +The Primus overlay keeps `base_config: "base.yml"` so MaxText loads its own [`configs/base.yml`](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/configs/base.yml) at runtime. This page lists **Primus-defined defaults and commonly overridden Primus fields**. For the full upstream parameter set (hundreds of keys), see the [MaxText documentation](https://maxtext.readthedocs.io/) and upstream `base.yml`. + +--- + +## How parameters flow + +1. YAML presets under `primus/configs/modules/maxtext/` and `primus/configs/models/maxtext/` are merged with CLI overrides. +2. `MaxTextAdapter.convert_config` passes the merged namespace through `MaxTextConfigBuilder` (currently a thin pass-through). +3. `export_params_to_yaml` writes a flat YAML file; MaxText ignores unknown Primus-private keys via pydantic filtering. +4. Unknown keys from upstream still resolve through environment overrides inside MaxText (`pyconfig`), not shown here. + +--- + +## 1. Base module parameters + +Shared with all Primus modules via `module_base.yaml` and trainer extensions. + + +| Parameter | Default (Primus) | Description | +| ------------------- | ----------------------------------------------------------------- | -------------------------------------------------------------------------------------------- | +| `trainable` | `true` in `trainer_base.yaml` (overrides `module_base`’s `false`) | When `true`, the module participates in training orchestration. | +| `sink_level` | `null` | Structured logging sink level for the module (if the logging stack is configured to use it). | +| `file_sink_level` | `DEBUG` | File sink verbosity. | +| `stderr_sink_level` | `INFO` | Stderr sink verbosity. | + + +--- + +## 2. Training + +From `pre_trainer.yaml` (extends `trainer_base.yaml`). + + +| Parameter | Default | Description | +| ------------- | ------------ | ------------------------------------------------------------- | +| `base_config` | `"base.yml"` | Upstream MaxText base file loaded by `pyconfig._load_config`. | +| `hardware` | `"gpu"` | Hardware target string consumed by MaxText. | +| `steps` | `1000` | Global optimizer steps for the run. | +| `log_period` | `100` | Steps between log emissions. | + + +--- + +## 3. Data + + +| Parameter | Default | Description | +| ---------------- | -------------- | -------------------------------------------------------------------- | +| `dataset_type` | `"hf"` | Dataset backend selector (Hugging Face in the default path). | +| `hf_path` | `"allenai/c4"` | Hugging Face dataset repo or identifier. | +| `hf_data_dir` | `"en"` | Subdirectory / config slice within the HF dataset. | +| `hf_train_files` | `""` | Optional explicit train file list (format per MaxText HF loader). | +| `packing` | `true` | Sequence packing for efficiency when supported by the data pipeline. | + + +--- + +## 4. Checkpointing + +These are Primus overlay defaults. MaxText also loads upstream `base.yml` at runtime through `base_config: "base.yml"`, where upstream checkpoint defaults may differ. When debugging effective behavior, distinguish the Primus YAML written by the adapter from the upstream MaxText defaults loaded afterward. + + +| Parameter | Default | Description | +| ---------------------- | ------- | ------------------------------------------------------------------ | +| `enable_checkpointing` | `false` | See Training section. | +| `async_checkpointing` | `false` | When `enable_checkpointing` is true, use async checkpoint workers. | + + +--- + +## 5. Profiling + + +| Parameter | Default | Description | +| --------------------------------- | ---------- | --------------------------------------- | +| `profiler` | `"xplane"` | Profiler backend (e.g. XPlane for JAX). | +| `skip_first_n_steps_for_profiler` | `3` | Warmup steps excluded from capture. | +| `profiler_steps` | `1` | Number of steps to profile once active. | + + +--- + +## 6. Memory and recomputation + + +| Parameter | Default | Description | +| ------------------------------- | -------- | ---------------------------------------------------------------------------------- | +| `remat_policy` | `'full'` | Activation rematerialization policy (`none`, `minimal`, `full`, etc.—see MaxText). | +| `optimizer_memory_host_offload` | `false` | Offload optimizer state to host memory when supported. | +| `scan_layers` | `true` | Use scanned layer implementation where applicable. | +| `param_scan_axis` | `1` | Axis for parameter scanning / partitioning layout. | + + +--- + +## 7. Precision and quantization + + +| Parameter | Default | Description | +| ------------------------- | ----------------- | --------------------------------------------------------------------- | +| `dtype` | `"bfloat16"` | Default compute dtype for many ops. | +| `quantization` | `""` | Quantization mode string (empty = none; set per MaxText AQT recipes). | +| `quantize_kvcache` | `false` | Quantize KV cache tensors. | +| `kv_quant_axis` | `"heads_and_dkv"` | KV quantization axis naming for kernels. | +| `kv_quant_dtype` | `"int8"` | Storage dtype for KV cache when quantization is on. | +| `weight_dtype` | `bfloat16` | Weight storage / compute dtype for non-quantized paths. | +| `checkpoint_is_quantized` | `false` | Set `true` when loading an AQT-quantized checkpoint. | +| `logits_dot_in_fp32` | `false` | Compute logits matmul in `float32` for numerical stability. | + + +--- + +## 8. Model + +From `model_base.yaml` and per-model files such as `llama3_8B.yaml`. + + +| Parameter | Default | Description | +| ----------------------- | ------------------------------------------------------------------- | ----------------------------------------------------------------------- | +| `model_name` | `"default"` in `model_base`; e.g. `"llama3-8b"` in `llama3_8B.yaml` | Selects MaxText’s bundled model YAML when present. | +| `override_model_config` | `true` | When `true`, CLI / kwargs override values from the loaded model config. | +| `attention` | `"cudnn_flash_te"` | Attention implementation (Primus default favors TE flash on AMD GPUs). | +| `use_iota_embed` | `true` | Use iota-based embedding for performance on accelerator backends. | +| `tokenizer_path` | e.g. `"meta-llama/Meta-Llama-3-8B"` | Hugging Face tokenizer id or local path. | + + +--- + +## 9. Advanced + + +| Parameter | Default | Description | +| --------- | ------- | --------------------------------------------------------------------- | +| `shardy` | `false` | Enable Shardy-related integration in MaxText when building shardings. | + + +--- + +## Related reading + +- [MaxText documentation](https://maxtext.readthedocs.io/)—full parameter reference and recipes. +- Primus implementation: `primus/backends/maxtext/argument_builder.py`, `maxtext_pretrain_trainer.py`, `maxtext_adapter.py`. diff --git a/docs/03-configuration-reference/megatron-bridge-parameters.md b/docs/03-configuration-reference/megatron-bridge-parameters.md new file mode 100644 index 000000000..620f0a2be --- /dev/null +++ b/docs/03-configuration-reference/megatron-bridge-parameters.md @@ -0,0 +1,167 @@ +# Megatron Bridge backend configuration reference + +Megatron Bridge integrates [Megatron-Core](https://github.com/NVIDIA/Megatron-LM) training with Hugging Face–centric workflows. In Primus, the **`megatron_bridge`** framework is used for post-training with module preset `sft_trainer.yaml`, and the repository also ships a pretraining preset at `primus/configs/modules/megatron_bridge/pretrain_trainer.yaml`. + +## Recipe system + +Megatron Bridge resolves training defaults through a **recipe** and **flavor**: + +- `recipe` is a Python module path under `megatron.bridge.recipes` (e.g. `qwen.qwen3`). +- `flavor` is the function name inside that module (e.g. `qwen3_8b_finetune_config`) that returns a `ConfigContainer`. + +At runtime, `load_recipe_config` in `primus/backends/megatron_bridge/config_utils.py`: + +1. Imports `megatron.bridge.recipes.` and calls `(**filtered_backend_args)` to build the baseline `ConfigContainer`. +2. **Deep-merges** Primus `backend_args` (from YAML + CLI) into that dataclass via `_merge_dict_to_dataclass`, so user overrides sit on top of recipe defaults. + +You normally specify `recipe`, `flavor`, `hf_path`, and `dataset` in the model YAML; training hyperparameters and parallelism go in module overrides or experiment module overrides (`modules.post_trainer.overrides` for SFT/post-training, `modules.pre_trainer.overrides` for pretraining examples). + +--- + +## 1. Base module parameters + +From `primus/configs/modules/megatron_bridge/sft_trainer.yaml` (extends `module_base.yaml`). Pretraining examples use `pretrain_trainer.yaml` instead. + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `trainable` | `true` | Module participates in the training graph. | +| `sink_level` | `null` | Inherited from `module_base.yaml`; structured logging sink level. | +| `file_sink_level` | `DEBUG` | File sink verbosity. | +| `stderr_sink_level` | `INFO` | Stderr sink verbosity. | + +--- + +## 2. Training + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `stage` | `"sft"` | Backend stage selector. Primus dispatches post-training via `primus train posttrain` and loads the Megatron Bridge posttrain trainer when this module is used under `post_trainer`. | +| `trainable` | `true` | See Base module parameters. | + +**CLI note:** The user-facing suite is **`posttrain`** (`primus train posttrain --config ...`). The YAML `stage` field selects the Megatron Bridge trainer implementation (`sft`), not the CLI suite name. + +For Bridge pretraining, use the normal pretraining suite (`primus train pretrain --config ...`) with experiments that reference `modules.pre_trainer.config: pretrain_trainer.yaml`. + +--- + +## 3. Fine-tuning method (PEFT) + +Primus examples set these under `modules.post_trainer.overrides` (see `examples/megatron_bridge/configs/`). + +| Parameter | Example | Description | +|-----------|---------|-------------| +| `peft` | `"none"`, `"lora"` | Parameter-efficient fine-tuning mode. | +| `peft_dim` | `16` | LoRA rank (example: `llama31_70b_lora_posttrain.yaml`). | +| `peft_alpha` | `32` | LoRA scaling alpha (same example). | +| `packed_sequence` | `false` | Pack multiple short sequences per microbatch when supported. | + +Additional keys such as `pretrained_checkpoint`, `use_distributed_optimizer`, or `cross_entropy_loss_fusion` appear in larger examples and are merged into the recipe `ConfigContainer` when the dataclass exposes matching fields. + +--- + +## 4. Parallelism + +Typical overrides from Megatron Bridge examples: + +| Parameter | Example | Description | +|-----------|---------|-------------| +| `tensor_model_parallel_size` | `1`, `2`, `8` | Tensor parallelism degree. | +| `pipeline_model_parallel_size` | `1` | Pipeline parallelism degree. | +| `virtual_pipeline_model_parallel_size` | `null` | Virtual pipeline stages per rank when PP > 1. | +| `context_parallel_size` | `1` | Context parallelism degree. | +| `sequence_parallel` | `false` | Sequence parallelism within TP groups. | +| `use_megatron_fsdp` | `false` | Optional Megatron FSDP path. | + +--- + +## 5. Training hyperparameters + +| Parameter | Example | Description | +|-----------|---------|-------------| +| `train_iters` | `200`, `1000` | Total training iterations. | +| `global_batch_size` | `8`, `128` | Global batch across data-parallel groups. | +| `micro_batch_size` | `1`, `2` | Per-GPU microbatch before gradient accumulation. | +| `seq_length` | `2048`, `8192` | Training sequence length. | +| `eval_interval` | `30` | Steps between evaluations. | +| `save_interval` | `50` | Steps between checkpoint saves. | + +--- + +## 6. Learning rate + +| Parameter | Example | Description | +|-----------|---------|-------------| +| `finetune_lr` | `1.0e-4`, `5.0e-6` | Peak learning rate for fine-tuning. | +| `min_lr` | `0.0` | Floor learning rate after decay. | +| `lr_warmup_iters` | `50` | Linear warmup length in iterations. | +| `lr_decay_iters` | `null` | Optional decay span; `null` defers to recipe defaults. | + +--- + +## 7. Precision + +| Parameter | Example | Description | +|-----------|---------|-------------| +| `precision_config` | `bf16_mixed`, `fp16_mixed`, `fp32` | Mixed-precision recipe for Megatron Bridge. | +| `comm_overlap_config` | `null` | Optional communication/compute overlap policy object. | +| `pipeline_dtype` | `null` | Dtype for pipeline stages when PP is enabled. | + +--- + +## 8. Memory optimization + +| Parameter | Example | Description | +|-----------|---------|-------------| +| `recompute_granularity` | `full` | Activation recomputation granularity. | +| `recompute_method` | `uniform` | How recomputation is scheduled across layers. | +| `recompute_num_layers` | `1` | Number of layers per recompute group (workload-dependent). | + +--- + +## 9. Primus-Turbo + +From `sft_trainer.yaml` (defaults shown). + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `enable_primus_turbo` | `true` | Master flag for Primus-Turbo optimized kernels and paths. | +| `use_turbo_attention` | `false` | Turbo attention implementation. | +| `use_turbo_parallel_linear` | `false` | Turbo parallel linear layers. | +| `use_turbo_grouped_gemm` | `false` | Turbo grouped GEMM flag for MoE paths (the preset ships this key). The former `use_turbo_grouped_mlp` alias has been removed. | +| `moe_use_fused_router_with_aux_score` | `false` | Fused MoE router with auxiliary loss handling. | +| `enable_turbo_attention_float8` | `false` | FP8 path inside Turbo attention. | +| `use_turbo_deepep` | `false` | DeepEP-style expert-parallel integration. | +| `turbo_deepep_num_cu` | `32` | Compute-unit count hint for DeepEP. | +| `turbo_deepep_use_comm_stream` | `false` | Use dedicated communication streams. | +| `turbo_sync_free_moe_stage` | `0` | Sync-free MoE scheduling stage. | +| `use_turbo_fused_act_with_probs` | `false` | Fuse activation with probability tensors where applicable. | +| `use_turbo_rms_norm` | `false` | Turbo RMSNorm path. | + +**Environment:** `PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND` (default `TURBO`) is read in `primus/backends/megatron/patches/args/rocm_arg_validation.py` to select MoE dispatch/combine behavior when Turbo MoE is active. + +--- + +## 10. Model and dataset + +Model YAML files (`qwen3_8b.yaml`, `qwen3_32b.yaml`, `llama31_70b.yaml`) supply: + +| Parameter | Example | Description | +|-----------|---------|-------------| +| `recipe` | `qwen.qwen3`, `llama.llama3` | Recipe module under `megatron.bridge.recipes`. | +| `flavor` | `qwen3_8b_finetune_config`, `llama31_70b_finetune_config` | Flavor function producing the baseline `ConfigContainer`. | +| `hf_path` | `Qwen/Qwen3-8B`, `meta-llama/Meta-Llama-3.1-70B` | Hugging Face model id for weights/tokenizer flows. | +| `dataset` | nested | Example: `dataset_name: "rajpurkar/squad"` for SQuAD-style fine-tuning. | + +**Logging (optional overrides in examples):** `wandb_project`, `wandb_entity`, `wandb_exp_name` may be set under `overrides` for experiment tracking when Weights & Biases is configured. + +--- + +## Argument merge mechanics + +`MegatronBridgeArgBuilder` (`primus/backends/megatron_bridge/argument_builder.py`) performs a **deep merge** of CLI and YAML into a single dict/namespace before `load_recipe_config` runs. Nested dicts (for example dataset or optimizer sections) combine recursively; explicit `None` in the merged structure can clear fields depending on merge rules in `_merge_dict_to_dataclass`. + +--- + +## Example layouts + +Under `examples/megatron_bridge/configs/`, per-GPU directories (for example `MI300X/`, `MI355X/`) contain full experiment YAMLs that set `work_group`, `user_name`, `exp_name`, `workspace`, and Megatron Bridge modules. Post-training examples use `modules.post_trainer` with `config: sft_trainer.yaml`; MI300X pretraining examples use `modules.pre_trainer` with `config: pretrain_trainer.yaml`. Both patterns set `framework: megatron_bridge`, `model: .yaml`, and an `overrides` block for parallelism, LR, precision, and related options. diff --git a/docs/03-configuration-reference/megatron-parameters.md b/docs/03-configuration-reference/megatron-parameters.md new file mode 100644 index 000000000..25c880542 --- /dev/null +++ b/docs/03-configuration-reference/megatron-parameters.md @@ -0,0 +1,854 @@ +# Megatron backend configuration reference + +This page lists the flat configuration keys exposed by Primus when `framework: megatron`. Unless a section says otherwise, values are the defaults from `primus/configs/modules/megatron/trainer_base.yaml` and related model presets. The effective pretraining preset is `pre_trainer.yaml`, which extends `trainer_base.yaml` and overrides several high-impact training defaults. + +**Where parameters live.** Set overrides under `modules.pre_trainer.overrides:` in your experiment YAML. Model architecture keys usually come from `models..overrides:` (or your chosen model preset), but the same names map to Megatron’s argparse namespace either way. + +**Presets.** + +- Module presets: `primus/configs/modules/megatron/` (the main pretraining bundle is `pre_trainer.yaml`, which extends `trainer_base.yaml` and Primus Megatron add-ons). +- Model presets: `primus/configs/models/megatron/` (for example `language_model.yaml`). + +**Mapping to Megatron-LM.** Keys are passed through **1:1** to Megatron’s training arguments (same names as `argparse` / `Namespace`). Primus builds that namespace with `MegatronArgBuilder`. + +**Upstream reference.** Full flag semantics and newer options are defined in Megatron-LM: [`megatron/training/arguments.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/arguments.py). + +### Example (experiment YAML) + +```yaml +framework: megatron + +modules: + pre_trainer: + overrides: + global_batch_size: 256 + train_iters: 50000 + tensor_model_parallel_size: 2 + +models: + pre_train: + overrides: + hidden_size: 2048 + num_layers: 32 +``` + +--- + +## 1. Base module parameters + +*Source: `primus/configs/modules/module_base.yaml` (merged into Megatron presets; `trainer_base.yaml` sets `trainable: true`).* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `trainable` | `true` | When `true`, this module participates in training workflows. (`module_base.yaml` alone defaults to `false`; Megatron `trainer_base.yaml` overrides to `true`.) | +| `sink_level` | `null` | Log level for the structured sink (Primus module plumbing); `null` uses framework default. | +| `file_sink_level` | `DEBUG` | Minimum level for file-backed logging. | +| `stderr_sink_level` | `INFO` | Minimum level for stderr logging. | + +--- + +## 2. Training and batching + +*Source: `primus/configs/modules/megatron/trainer_base.yaml`; effective `pre_trainer.yaml` overrides are noted where they differ.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `yaml_cfg` | `null` | Reserved; not supported as a Megatron override in this preset. | +| `spec` | `null` | Optional trainer spec hook (unused in defaults). | +| `micro_batch_size` | `2` | Samples per microbatch per data-parallel rank (per forward/backward step before gradient accumulation). | +| `batch_size` | `null` | Deprecated; use `micro_batch_size` / `global_batch_size`. | +| `global_batch_size` | `128` (`16` in `pre_trainer.yaml`) | Total batch size across the data-parallel world (before or after splitting, per Megatron semantics). | +| `rampup_batch_size` | `null` | Optional batch-size ramp schedule string / config. | +| `decrease_batch_size_if_needed` | `false` | Allow shrinking batch if memory is insufficient. | +| `check_for_nan_in_loss_and_grad` | `true` | Abort on NaNs in loss or gradients. | +| `check_for_spiky_loss` | `false` | Detect abnormal loss spikes. | +| `check_for_large_grads` | `false` | Detect abnormally large gradients. | +| `make_vocab_size_divisible_by` | `128` | Pads vocabulary size for efficient kernels / partitioning. | +| `exit_signal_handler` | `false` | Install handlers for graceful shutdown signals. | +| `exit_duration_in_mins` | `null` | Stop training after this many minutes. | +| `exit_interval` | `null` | Exit after this many iterations (if set). | +| `onnx_safe` | `null` | ONNX export compatibility tweaks. | +| `bert_binary_head` | `true` | Use BERT binary classification head when applicable. | +| `use_flash_attn` | `false` (`true` in `pre_trainer.yaml`) | Prefer FlashAttention kernels when available. | +| `seed` | `1234` | RNG seed for reproducibility. | +| `data_parallel_random_init` | `false` | Random init that varies across data-parallel ranks. | +| `init_method_xavier_uniform` | `false` | Use Xavier uniform for some weights. | +| `test_mode` | `false` | Lightweight test path (fewer steps / checks). | +| `train_iters` | `null` (`1000` in `pre_trainer.yaml`) | Total training iterations (mutually exclusive with sample-based stopping in typical setups). | +| `train_samples` | `null` | Total training samples (when using sample-based training). | +| `eval_iters` | `32` (`0` in `pre_trainer.yaml`) | Validation iterations per eval. | +| `eval_interval` | `2000` (`1000` in `pre_trainer.yaml`) | Run validation every this many iterations. | +| `full_validation` | `false` | Run a full pass over validation data. | +| `multiple_validation_sets` | `false` | Multiple validation datasets / passes. | +| `skip_train` | `false` | Only run eval / test, no training updates. | +| `train_sync_interval` | `null` | Periodic distributed sync barrier for debugging. | +| `adlr_autoresume` | `false` | ADLR autoresume integration. | +| `adlr_autoresume_interval` | `1000` | Autoresume checkpoint interval. | +| `manual_gc` | `false` | Force Python GC on a schedule. | +| `manual_gc_interval` | `1` | GC every N steps when `manual_gc` is enabled. | +| `manual_gc_eval` | `false` | Run manual GC during evaluation. | +| `mask_type` | `random` | Masking strategy for MLM / similar objectives. | +| `mask_factor` | `1.0` | Masking strength multiplier. | +| `iter_per_epoch` | `1250` | Iterations interpreted as one “epoch” for logging. | + +--- + +## 3. Mixed precision + +*Source: `trainer_base.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `fp16` | `false` | Enable FP16 mixed precision training. | +| `bf16` | `true` | Enable BF16 mixed precision training. | +| `grad_reduce_in_bf16` | `false` | All-reduce gradients in BF16 (saves bandwidth). | +| `calculate_per_token_loss` | `false` | Normalize loss per token instead of per sample. | +| `loss_scale` | `null` | Static loss scale for FP16; `null` uses dynamic scaling. | +| `initial_loss_scale` | `4294967296` | Initial dynamic loss scale. | +| `min_loss_scale` | `1.0` | Floor for dynamic loss scale. | +| `loss_scale_window` | `1000` | Window for dynamic loss scaling updates. | +| `hysteresis` | `2` | Hysteresis steps for loss-scale decreases. | +| `accumulate_allreduce_grads_in_fp32` | `false` | Accumulate and reduce gradients in FP32. | +| `fp16_lm_cross_entropy` | `false` | Compute LM cross-entropy in FP16. | +| `fp8` | `null` | FP8 recipe selection (`e4m3`, `hybrid`, etc.); `null` disables. | +| `fp8_margin` | `0` | FP8 scaling margin. | +| `fp8_recipe` | `delayed` | FP8 recipe variant (e.g. delayed scaling). | +| `fp8_interval` | `1` | Deprecated FP8 interval (kept for compatibility). | +| `fp8_amax_history_len` | `1024` | History length for FP8 amax statistics. | +| `fp8_amax_compute_algo` | `"max"` | How to combine amax history (`max`, etc.). | +| `fp8_wgrad` | `true` | Run weight gradients in FP8 where supported. | +| `fp8_param_gather` | `false` | FP8 parameter gather for distributed optimizer paths. | +| `te_rng_tracker` | `false` | Transformer Engine RNG tracker for FP8. | +| `inference_rng_tracker` | `false` | Separate RNG tracker for inference FP8. | +| `fp4` | `null` | FP4 mode; `null` disables. | +| `fp4_recipe` | `nvfp4` | FP4 recipe name. | +| `fp4_param` | `false` | Store parameters in FP4. | +| `first_last_layers_bf16` | `false` | Keep first/last layers in BF16 for stability. | +| `num_layers_at_start_in_bf16` | `1` | Count of early layers forced to BF16 when enabled. | +| `num_layers_at_end_in_bf16` | `1` | Count of final layers forced to BF16 when enabled. | +| `no_fp8_weight_transpose_cache` | `false` | *Primus:* disable FP8 weight transpose cache (see `primus_megatron_module.yaml`). | + +--- + +## 4. Optimizer and learning rate + +*Source: `trainer_base.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `optimizer` | `adam` | Optimizer family (`adam`, `sgd`, etc.). | +| `lr` | `2.5e-4` (`2.0e-05` in `pre_trainer.yaml`) | Peak learning rate. | +| `lr_decay_style` | `cosine` | LR decay schedule (`cosine`, `linear`, `constant`, WSD, etc.). | +| `lr_decay_iters` | `null` | Decay duration in iterations. | +| `lr_decay_samples` | `null` | Decay duration in samples. | +| `lr_warmup_fraction` | `null` | Warmup as a fraction of total train steps. | +| `lr_warmup_iters` | `0` (`40` in `pre_trainer.yaml`) | Linear warmup steps. | +| `lr_warmup_samples` | `0` | Warmup in samples. | +| `lr_warmup_init` | `0.0` | LR at the start of warmup. | +| `min_lr` | `2.5e-5` (`0.0` in `pre_trainer.yaml`) | Minimum LR after decay. | +| `lr_wsd_decay_style` | `exponential` | Weight-decay schedule style for WSD when used. | +| `lr_wsd_decay_samples` | `null` | WSD decay window in samples. | +| `lr_wsd_decay_iters` | `null` | WSD decay window in iterations. | +| `head_lr_mult` | `1.0` | LR multiplier for attention/head modules when supported. | +| `weight_decay` | `0.01` (`0.0` in `pre_trainer.yaml`) | AdamW / L2-style weight decay. | +| `start_weight_decay` | `null` | Starting weight decay for schedules. | +| `end_weight_decay` | `null` | Ending weight decay for schedules. | +| `weight_decay_incr_style` | `constant` | How weight decay changes between start/end. | +| `clip_grad` | `1.0` | Global gradient norm clip. | +| `adam_beta1` | `0.9` | Adam first moment decay. | +| `adam_beta2` | `0.95` (`0.999` in `pre_trainer.yaml`) | Adam second moment decay. | +| `adam_eps` | `1.0e-08` | Adam epsilon. | +| `sgd_momentum` | `0.9` | SGD momentum when `optimizer` is SGD. | +| `override_opt_param_scheduler` | `false` (`true` in `pre_trainer.yaml`) | Override optimizer parameter groups’ schedulers. | +| `use_checkpoint_opt_param_scheduler` | `false` | Load optimizer scheduler state strictly from checkpoint. | +| `warmup` | `null` | Alternate warmup specification (legacy / schedule hooks). | +| `decoupled_lr` | `null` | Decoupled LR for certain param groups. | +| `decoupled_min_lr` | `null` | Minimum for decoupled LR. | +| `muon_extra_scale_factor` | `1.0` | Muon optimizer scaling. | +| `muon_scale_mode` | `"spectral"` | Muon scaling mode. | +| `muon_fp32_matmul_prec` | `"medium"` | Muon matmul precision hint. | +| `muon_num_ns_steps` | `5` | Muon Newton–Schulz iterations. | +| `muon_tp_mode` | `"blockwise"` | Muon tensor-parallel mode. | +| `muon_use_nesterov` | `false` | Muon Nesterov momentum. | +| `muon_split_qkv` | `true` | Split QKV for Muon. | +| `muon_momentum` | `0.95` | Muon momentum. | +| `muon_weight_decay` | `0.01` | Muon-specific decay. | +| `muon_weight_decay_method` | `"decoupled"` | How Muon applies decay. | +| `optimizer_cpu_offload` | `false` | Offload optimizer state to CPU. | +| `optimizer_offload_fraction` | `1.0` | Fraction of optimizer state offloaded. | +| `use_torch_optimizer_for_cpu_offload` | `false` | Use PyTorch optimizer for offload path. | +| `overlap_cpu_optimizer_d2h_h2d` | `false` | Overlap CPU optimizer device transfers. | +| `pin_cpu_grads` | `true` | Pin memory for CPU gradients. | +| `pin_cpu_params` | `true` | Pin memory for CPU params in offload. | +| `use_precision_aware_optimizer` | `false` | Use precision-aware optimizer (main grads/params in lower precision). | +| `main_grads_dtype` | `fp32` | Dtype for main gradients (`fp32`, `bf16`). | +| `main_params_dtype` | `fp32` | Dtype for master params. | +| `exp_avg_dtype` | `fp32` | Optimizer first moment dtype (`fp32`, `fp16`, `fp8`). | +| `exp_avg_sq_dtype` | `fp32` | Optimizer second moment dtype. | + +--- + +## 5. Parallelism and distribution + +*Sources: `trainer_base.yaml` (distributed runtime) and `primus/configs/models/megatron/language_model.yaml` (model-parallel sizes and TP communication).* + +### 5.1 Data / distributed runtime (trainer) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `overlap_p2p_comm` | `true` | Overlap pipeline P2P with compute. | +| `distributed_backend` | `nccl` | Process-group backend (`nccl`, `gloo`, …). | +| `distributed_timeout_minutes` | `10` (`60` in `pre_trainer.yaml`) | Collective timeout. | +| `defer_embedding_wgrad_compute` | `false` | Defer embedding weight gradients. | +| `wgrad_deferral_limit` | `0` | Max deferred embedding wgrad steps. | +| `align_grad_reduce` | `true` | Align gradient reductions for efficiency. | +| `ddp_num_buckets` | `null` | Number of DDP buckets. | +| `ddp_bucket_size` | `null` | DDP bucket size in elements. | +| `ddp_pad_buckets_for_high_nccl_busbw` | `false` | Pad buckets for NCCL bus bandwidth. | +| `ddp_average_in_collective` | `false` | Average inside collective vs outside. | +| `overlap_grad_reduce` | `false` | Overlap gradient all-reduce with backward. | +| `overlap_param_gather` | `false` | Overlap param all-gather (distributed optimizer). | +| `overlap_param_gather_with_optimizer_step` | `false` | Overlap param gather with optimizer step. | +| `align_param_gather` | `true` | Align param gather for distributed optimizer. | +| `scatter_gather_tensors_in_pipeline` | `true` | Scatter/gather tensors across PP ranks. | +| `use_ring_exchange_p2p` | `false` | Ring-exchange P2P for PP. | +| `local_rank` | `null` | Local rank override (normally from launcher). | +| `lazy_mpu_init` | `null` | Defer Megatron parallel state init. | +| `account_for_embedding_in_pipeline_split` | `false` | Account for embedding in PP partition. | +| `account_for_loss_in_pipeline_split` | `false` | Account for loss partition in PP. | +| `empty_unused_memory_level` | `0` | Aggressiveness of `torch.cuda.empty_cache`. | +| `standalone_embedding_stage` | `false` | Dedicated PP stage for embeddings. | +| `use_distributed_optimizer` | `false` (`true` in `pre_trainer.yaml`) | Shard optimizer state across data parallel. | +| `use_sharp` | `false` | Use SHARP for collectives when available. | +| `sharp_enabled_group` | `null` | Which group SHARP applies to (`dp`, `dp_replica`). | +| `use_custom_fsdp` | `false` | Custom FSDP integration path. | +| `use_megatron_fsdp` | `false` | Megatron FSDP path. | +| `init_model_with_meta_device` | `false` | Build model on `meta` device first. | +| `data_parallel_sharding_strategy` | `no_shard` | FSDP / ZeRO style sharding (`no_shard`, `optim`, …). | +| `gradient_reduce_div_fusion` | `true` | Fuse division into reduce-scatter. | +| `suggested_communication_unit_size` | `400000000` | Suggested communication chunk size. | +| `keep_fp8_transpose_cache_when_using_custom_fsdp` | `false` | Keep FP8 transpose cache with custom FSDP. | +| `num_distributed_optimizer_instances` | `1` | Sharded optimizer instances per rank group. | +| `use_torch_fsdp2` | `false` | Use PyTorch FSDP2 integration. | +| `nccl_communicator_config_path` | `null` | JSON config for NCCL communicators. | +| `use_tp_pp_dp_mapping` | `false` | Custom TP/PP/DP process mapping. | +| `replication` | `false` | Data replication mode for certain schedules. | +| `replication_jump` | `null` | Stride between replicated ranks. | +| `replication_factor` | `null` | Replication factor. | +| `deterministic_mode` | `false` | Prefer deterministic algorithms (slower). | +| `check_weight_hash_across_dp_replicas_interval` | `null` | Periodically hash weights across DP replicas for debugging. | +| `overlap_moe_expert_parallel_comm` | `false` | Overlap MoE expert-parallel communication. | +| `decoder_pipeline_manual_split_list` | `null` | *Primus:* manual PP split points for decoder (list of ints). | +| `patch_moe_overlap` | `false` | *Primus:* patch MoE compute/comm overlap. | + +### 5.2 Model parallelism (model preset) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `model_parallel_size` | `null` | Legacy combined MP size override. | +| `tensor_model_parallel_size` | `1` | Tensor parallelism degree (intra-layer split). | +| `encoder_tensor_model_parallel_size` | `0` | Encoder TP size when encoder/decoder differ. | +| `pipeline_model_parallel_size` | `1` | Pipeline parallelism stages. | +| `pipeline_model_parallel_layout` | `null` | Optional explicit PP layout string. | +| `pipeline_model_parallel_comm_backend` | `null` | `nccl` or `ucc` for PP collectives. | +| `encoder_pipeline_model_parallel_size` | `0` | Encoder PP stages (encoder–decoder models). | +| `pipeline_model_parallel_split_rank` | `null` | Rank where encoder/decoder split. | +| `decoder_first_pipeline_num_layers` | `null` | Layers on first decoder PP stage. | +| `decoder_last_pipeline_num_layers` | `null` | Layers on last decoder PP stage. | +| `virtual_pipeline_model_parallel_size` | `null` | Virtual PP (interleaved) depth. | +| `num_layers_per_virtual_pipeline_stage` | `null` | Layers per virtual stage. | +| `num_virtual_stages_per_pipeline_rank` | `null` | Virtual stages per physical PP rank. | +| `microbatch_group_size_per_vp_stage` | `null` | Microbatch grouping for interleaved PP. | +| `sequence_parallel` | `true` | Sequence parallelism when TP > 1. | +| `context_parallel_size` | `1` | Context (sequence) parallelism degree. | +| `cp_comm_type` | `p2p` | Context-parallel comm pattern (`p2p`, `a2a`, `allgather`, `a2a+p2p`). | +| `hierarchical_context_parallel_sizes` | `null` | Hierarchical CP group sizes. | +| `expert_model_parallel_size` | `1` | Expert parallelism for MoE. | +| `expert_tensor_parallel_size` | `null` | Expert tensor-parallel degree. | +| `high_priority_stream_groups` | `[]` | Named groups that get high-priority CUDA streams. | + +### 5.3 Tensor-parallel communication overlap (model) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `async_tensor_model_parallel_allreduce` | `true` | Async TP all-reduces for column-parallel layers. | +| `tp_comm_overlap` | `false` | Enable TP communication overlap planner. | +| `tp_comm_overlap_cfg` | `null` | Extra JSON / path for overlap configuration. | +| `tp_comm_overlap_ag` | `true` | Overlap all-gather in TP backward. | +| `tp_comm_overlap_rs` | `true` | Overlap reduce-scatter in TP backward. | +| `tp_comm_overlap_rs_dgrad` | `false` | Overlap RS for data-grad path. | +| `tp_comm_split_ag` | `true` | Split all-gather for overlap. | +| `tp_comm_split_rs` | `true` | Split reduce-scatter for overlap. | +| `tp_comm_bulk_wgrad` | `true` | Bulk weight-gradient path for TP comm. | +| `tp_comm_bulk_dgrad` | `true` | Bulk data-gradient path for TP comm. | +| `barrier_with_L1_time` | `true` | Barrier using L1 timing hooks for TP comm profiling. | +| `tp_comm_bootstrap_backend` | `nccl` | Backend used to bootstrap TP communicators. | + +--- + +## 6. Checkpointing + +*Source: `trainer_base.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `save` | `null` | Path prefix / pattern for checkpoints to write. | +| `save_interval` | `20000` (`1000` in `pre_trainer.yaml`) | Save every N iterations. | +| `save_retain_interval` | `null` | Retain checkpoints at this interval. | +| `no_save_optim` | `null` | Skip optimizer state in checkpoints when truthy. | +| `no_save_rng` | `null` | Skip RNG state in checkpoints when truthy. | +| `load` | `null` | Checkpoint path to load. | +| `load_main_params_from_ckpt` | `false` | Load only main parameters. | +| `no_load_optim` | `null` | Skip loading optimizer state. | +| `no_load_rng` | `null` | Skip loading RNG state. | +| `finetune` | `false` (`true` in `pre_trainer.yaml`) | Finetune mode (do not require full optimizer match). | +| `use_checkpoint_args` | `false` | When `true`, restore training args from checkpoint metadata. | +| `use_mp_args_from_checkpoint_args` | `false` | Restore model-parallel args from checkpoint. | +| `use_tokenizer_model_from_checkpoint_args` | `true` | Restore tokenizer path from checkpoint args. | +| `exit_on_missing_checkpoint` | `true` | Fail if `load` is set but checkpoint is missing. | +| `non_persistent_save_interval` | `null` | Ephemeral checkpoint interval. | +| `non_persistent_ckpt_type` | `null` | `global`, `local`, `in_memory`, or `null`. | +| `non_persistent_global_ckpt_dir` | `null` | Directory for non-persistent global checkpoints. | +| `non_persistent_local_ckpt_dir` | `null` | Directory for non-persistent local checkpoints. | +| `non_persistent_local_ckpt_algo` | `"fully_parallel"` | `fully_parallel` or `atomic`. | +| `pretrained_checkpoint` | `null` | Load weights from a pretrained checkpoint path. | +| `ckpt_step` | `null` | Specific step to load within a distributed checkpoint. | +| `use_dist_ckpt_deprecated` | `false` | Use deprecated distributed checkpoint format. | +| `use_persistent_ckpt_worker` | `false` | Background worker for checkpoint IO. | +| `auto_detect_ckpt_format` | `false` | Infer checkpoint format automatically. | +| `dist_ckpt_format_deprecated` | `null` | Legacy format hint. | +| `ckpt_format` | `torch_dist` | `torch`, `torch_dist`, or `zarr`. | +| `ckpt_convert_format` | `null` | Target format for one-shot conversion. | +| `ckpt_convert_save` | `null` | Output path for conversion. | +| `ckpt_convert_update_legacy_dist_opt_format` | `false` | Update legacy distributed-optimizer layout when converting. | +| `ckpt_fully_parallel_save_deprecated` | `false` | Deprecated fully-parallel save toggle. | +| `ckpt_fully_parallel_save` | `true` | Save shards in parallel across ranks. | +| `async_save` | `null` | Async checkpoint save (`null` = framework default). | +| `ckpt_fully_parallel_load` | `false` | Load shards in parallel. | +| `ckpt_assume_constant_structure` | `false` | Assume identical layer structure across ranks. | +| `dist_ckpt_strictness` | `assume_ok_unexpected` | How to handle unexpected keys in distributed ckpt. | +| `dist_ckpt_save_pre_mcore_014` | `null` | Compatibility flag for older Megatron-Core checkpoints. | +| `dist_ckpt_optim_fully_reshardable` | `null` | Optimizer state fully reshardable layout. | +| `auto_continue_train` | `false` | *Primus:* resume from latest checkpoint in the save directory when enabled. | +| `disable_last_saving` | `false` | *Primus:* skip writing the final checkpoint at shutdown. | + +--- + +## 7. Data + +*Source: `trainer_base.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `data_path` | `null` | Single blended dataset path / list. | +| `data_sharding` | `true` | Shard data across ranks. | +| `split` | `"99,1,0"` (`null` in `pre_trainer.yaml`) | Train/valid/test split ratios as comma string. | +| `train_data_path` | `null` | Training data blend. | +| `valid_data_path` | `null` | Validation data blend. | +| `test_data_path` | `null` | Test data blend. | +| `data_args_path` | `null` | External JSON/YAML of dataset arguments. | +| `per_split_data_args_path` | `null` | Per-split dataset args file. | +| `data_cache_path` | `null` | On-disk cache for indexed datasets. | +| `mock_data` | `false` | Use synthetic data (no real files). | +| `merge_file` | `null` | Merge file for blended datasets. | +| `seq_length` | `4096` (`1024` in `pre_trainer.yaml`) | Training sequence length. | +| `encoder_seq_length` | `null` | Encoder sequence length (encoder–decoder). | +| `decoder_seq_length` | `null` | Decoder sequence length. | +| `retriever_seq_length` | `256` | Sequence length for retriever models. | +| `sample_rate` | `1.0` | Sampling rate for dataset blending. | +| `mask_prob` | `0.15` | MLM mask probability. | +| `short_seq_prob` | `0.1` | Probability of shorter sequences in BERT-style data. | +| `num_workers` | `8` | DataLoader worker processes per rank. | +| `reset_position_ids` | `false` | Reset position IDs at document boundaries. | +| `reset_attention_mask` | `false` | Reset attention mask at boundaries. | +| `eod_mask_loss` | `false` | Mask loss at end-of-document tokens. | +| `dataloader_type` | `null` (`cyclic` in `pre_trainer.yaml`) | Dataloader implementation (`single`, `cyclic`, `external`, …). | +| `mmap_bin_files` | `true` | Memory-map `.bin` index files when supported. | +| `create_attention_mask_in_dataloader` | `true` | Build attention masks in the dataloader. | +| `num_dataset_builder_threads` | `1` | Threads to build dataset indices. | + +--- + +## 8. Recomputation (activation checkpointing) + +*Sources: `trainer_base.yaml` and `primus_megatron_module.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `recompute_activations` | `false` | Enable activation recomputation globally. | +| `recompute_granularity` | `null` | `full` or `selective` checkpointing. | +| `recompute_method` | `null` | `uniform` or `block` selective recomputation. | +| `recompute_num_layers` | `null` | Layers to recompute per block / schedule. | +| `recompute_layer_ids` | `null` | *Primus:* explicit **global** layer indices to recompute (`0 … num_layers-1`). | +| `distribute_saved_activations` | `false` | Distribute saved activations across TP/PP for memory balance. | +| `checkpoint_activations` | `false` | Deprecated alias for activation checkpointing. | +| `moe_layer_recompute` | `false` | Recompute MoE layer activations (model preset). | + +--- + +## 9. Logging and profiling + +*Sources: `trainer_base.yaml` and `primus_megatron_module.yaml`.* + +### 9.1 Logging + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `log_avg_skip_iterations` | `2` | Skip first N iterations for throughput averaging. | +| `log_avg_reset_interval` | `10` | Reset moving averages periodically. | +| `log_params_norm` | `false` | Log L2 norms of parameters. | +| `log_num_zeros_in_grad` | `false` | Log fraction of zero gradients. | +| `log_throughput` | `false` (`true` in `pre_trainer.yaml`) | Log tokens/sec and timing. | +| `log_progress` | `false` | Verbose progress logging. | +| `timing_log_level` | `0` | Verbosity for timing logs. | +| `timing_log_option` | `minmax` | Aggregate style for timing (`minmax`, `all`, …). | +| `tensorboard_log_interval` | `1` | Steps between TensorBoard scalars. | +| `tensorboard_queue_size` | `1000` | TensorBoard event queue size. | +| `log_timers_to_tensorboard` | `false` (`true` in `pre_trainer.yaml`) | Write timer stats to TensorBoard. | +| `log_batch_size_to_tensorboard` | `false` (`true` in `pre_trainer.yaml`) | Log batch size. | +| `log_learning_rate_to_tensorboard` | `true` | Log LR. | +| `log_validation_ppl_to_tensorboard` | `false` | Log validation perplexity. | +| `log_memory_to_tensorboard` | `false` | Log memory usage. | +| `log_world_size_to_tensorboard` | `false` | Log distributed world size. | +| `log_loss_scale_to_tensorboard` | `true` | Log FP16/FP8 loss scale. | +| `wandb_project` | `null` | Weights & Biases project name. | +| `wandb_exp_name` | `null` | W&B run name. | +| `wandb_save_dir` | `null` | W&B local directory. | +| `wandb_entity` | `null` | W&B entity / team. | +| `enable_one_logger` | `true` | Enable NVIDIA OneLogger integration. | +| `one_logger_project` | `megatron-lm` | OneLogger project string. | +| `one_logger_run_name` | `null` | OneLogger run name. | +| `log_interval` | `100` (`1` in `pre_trainer.yaml`) | Console log interval in iterations. | +| `tensorboard_dir` | `null` | TensorBoard output directory. | +| `logging_level` | `null` | Python logging level override. | +| `config_logger_dir` | `""` | Directory for dumped config logs. | +| `one_logger_async` | `false` | Async OneLogger flushing. | +| `app_tag_run_name` | `null` | Application tag for telemetry. | +| `app_tag_run_version` | `0.0.0` | Application tag version. | +| `disable_tensorboard` | `true` | *Primus:* disable TensorBoard integration in Primus-wrapped runs. | +| `disable_wandb` | `true` | *Primus:* disable W&B. | +| `disable_mlflow` | `true` | *Primus:* disable MLflow. | +| `mlflow_run_name` | `null` | *Primus:* MLflow run name. | +| `mlflow_experiment_name` | `null` | *Primus:* MLflow experiment name. | +| `use_rocm_mem_info` | `false` | *Primus:* collect ROCm memory info via `rocm-smi` every step when `true`. | +| `use_rocm_mem_info_iters` | `[1, 2]` | *Primus:* iterations at which to log memory if `use_rocm_mem_info` is `false`. | + +### 9.2 Profiling + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `profile` | `false` | Enable lightweight Nsight / CUDA profiling hooks. | +| `use_pytorch_profiler` | `false` | Enable `torch.profiler` regions. | +| `profile_ranks` | `[0]` | Ranks to profile. | +| `profile_step_start` | `10` | First step to profile. | +| `profile_step_end` | `12` | Last step to profile. | +| `iterations_to_skip` | `null` | Skip listed iterations in profiling. | +| `result_rejected_tracker_filename` | `null` | Log rejected samples to this file. | +| `enable_gloo_process_groups` | `true` | Create auxiliary Gloo groups for CPU-side ops. | +| `record_memory_history` | `false` | Record CUDA memory history (debug). | +| `memory_snapshot_path` | `snapshot.pickle` | Path for memory snapshot dumps. | +| `disable_profiler_activity_cpu` | `false` | *Primus:* omit CPU activities from profiler traces. | +| `torch_profiler_record_shapes` | `true` | *Primus:* record tensor shapes in PyTorch profiler. | +| `torch_profiler_with_stack` | `true` | *Primus:* capture Python stacks in profiler. | +| `torch_profiler_use_gzip` | `false` | *Primus:* gzip profiler outputs. | + +--- + +## 10. Model architecture + +*Sources: `primus/configs/models/megatron/language_model.yaml` and `primus/configs/models/megatron/primus_megatron_model.yaml`.* + +### 10.1 Core architecture + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `use_legacy_models` | `false` | Use legacy Megatron model code paths. | +| `deprecated_use_mcore_models` | `false` | Deprecated flag for Megatron-Core models; prefer current `transformer_impl` + stack. | +| `model_type` | `gpt` | `gpt` or `mamba` family. | +| `num_layers` | `24` | Transformer layers (decoder or unified stack). | +| `encoder_num_layers` | `null` | Encoder depth (encoder–decoder). | +| `decoder_num_layers` | `null` | Decoder depth. | +| `hidden_size` | `1024` | Hidden / model width. | +| `num_attention_heads` | `16` | Attention heads. | +| `attention_backend` | `auto` | Attention kernel backend selection. | +| `group_query_attention` | `false` | Enable grouped-query attention (GQA). | +| `qk_layernorm` | `false` | LayerNorm on Q/K projections. | +| `qk_l2_norm` | `false` | L2-normalize Q/K vectors. | +| `num_query_groups` | `null` | Number of query groups for GQA; `null` means MHA. | +| `add_position_embedding` | `false` | Add absolute position embeddings (non-RoPE stacks). | +| `position_embedding_type` | `learned_absolute` | Position embedding style. | +| `max_position_embeddings` | `null` | Maximum sequence positions (context length cap). | +| `original_max_position_embeddings` | `null` | Original pretrained length for interpolation / scaling. | +| `untie_embeddings_and_output_weights` | `true` | Separate input embedding and LM head weights. | +| `ffn_hidden_size` | `null` | FFN hidden size; `null` often defaults via `hidden_size` heuristics. | +| `kv_channels` | `null` | Per-head KV channels override. | +| `hidden_dropout` | `0.1` | Dropout on residual / hidden states. | +| `attention_dropout` | `0.1` | Attention dropout. | +| `fp32_residual_connection` | `false` | Accumulate residuals in FP32. | +| `apply_residual_connection_post_layernorm` | `false` | Apply residual after (vs before) norm where supported. | +| `add_bias_linear` | `false` | Biases in linear / column-parallel layers. | +| `add_qkv_bias` | `false` | Biases in QKV projections. | +| `swiglu` | `true` | SwiGLU activation in FFN. | +| `quick_geglu` | `false` | Faster GeGLU path. | +| `openai_gelu` | `false` | OpenAI GELU variant. | +| `squared_relu` | `false` | Squared ReLU activation. | +| `rotary_base` | `10000` | RoPE base frequency. | +| `rotary_percent` | `1.0` | Fraction of head dim spanned by RoPE. | +| `rotary_interleaved` | `false` | Interleaved RoPE layout. | +| `rotary_seq_len_interpolation_factor` | `null` | Positional interpolation factor for long contexts. | +| `use_rotary_position_embeddings` | `null` | Force RoPE on/off; `null` follows model type. | +| `use_rope_scaling` | `false` | Enable LLaMA-style rope scaling. | +| `rope_scaling_factor` | `8.0` | Scaling factor for extended contexts (LLaMA-3 style). | +| `transformer_impl` | `transformer_engine` | Backend library (`transformer_engine`, `local`, …). | +| `rope_type` | `null` | `rope` or `yarn` style extensions. | +| `norm_epsilon` | `1.0e-05` | LayerNorm / RMSNorm epsilon. | +| `normalization` | `"LayerNorm"` | Norm type (`LayerNorm`, `RMSNorm` with TE, …). | +| `apply_layernorm_1p` | `false` | LayerNorm with +1 offset trick. | +| `clone_scatter_output_in_embedding` | `true` | Clone embedding scatter for autograd safety. | +| `perform_initialization` | `true` | Run weight initialization. | +| `use_cpu_initialization` | `null` | Initialize on CPU then move to GPU. | +| `use_te_activation_func` | `false` | Use Transformer Engine activation kernels. | +| `gradient_accumulation_fusion` | `true` | Fuse gradient accumulation kernels. | +| `delay_wgrad_compute` | `false` | Delay weight-gradient computation for scheduling. | + +### 10.2 Tokenizer and vocabulary + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `tokenizer_type` | `null` | Tokenizer class name (`GPT2BPETokenizer`, `HuggingFaceTokenizer`, …). | +| `tokenizer_model` | `null` | Path to tokenizer model / vocabulary file. | +| `vocab_size` | `null` | Vocabulary size (often inferred from tokenizer). | +| `vocab_file` | `null` | Vocabulary file path for BPE/WP tokenizers. | +| `vocab_extra_ids` | `0` | Extra reserved token slots. | +| `tiktoken_pattern` | `null` | Regex pattern for tiktoken. | +| `tiktoken_num_special_tokens` | `1000` | Special token count for tiktoken setup. | +| `tiktoken_special_tokens` | `null` | Serialized special tokens for tiktoken. | +| `legacy_tokenizer` | `false` | Legacy tokenizer behavior. | +| `trust_remote_code` | `false` | `trust_remote_code` for Hugging Face tokenizers. | + +### 10.3 Initialization and attention numerics + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `init_method_std` | `0.02` | Standard deviation for weight init. | +| `apply_query_key_layer_scaling` | `false` | Scale Q/K by layer index (deprecated GPT-3 trick). | +| `attention_softmax_in_fp32` | `false` | Force softmax in FP32. | + +### 10.4 Kernel fusion flags + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `bias_gelu_fusion` | `true` | Fuse bias + GELU. | +| `cross_entropy_loss_fusion` | `false` | Fused cross-entropy + softmax. | +| `cross_entropy_fusion_impl` | `"native"` | `native` or `te` fused CE. | +| `bias_swiglu_fusion` | `true` | Fuse bias + SwiGLU. | +| `masked_softmax_fusion` | `true` | Fused masked softmax. | +| `no_persist_layer_norm` | `false` | Non-persistent LayerNorm mode in TE. | +| `bias_dropout_fusion` | `true` | Fuse bias + dropout. | +| `apply_rope_fusion` | `true` | Fused RoPE kernels. | + +### 10.5 Multi-latent attention (MLA) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `multi_latent_attention` | `false` | Enable MLA blocks instead of standard MHA. | +| `q_lora_rank` | `null` | Low-rank query projection rank. | +| `kv_lora_rank` | `32` | Low-rank KV compression rank. | +| `qk_head_dim` | `128` | Q/K head dimension for MLA. | +| `qk_pos_emb_head_dim` | `64` | Positional head dimension for MLA. | +| `v_head_dim` | `128` | Value head dimension for MLA. | +| `rotary_scaling_factor` | `1.0` | RoPE scaling inside MLA (distinct from `rope_scaling_factor` above). | +| `mscale` | `1.0` | Yarn / scaling m-factor. | +| `mscale_all_dim` | `1.0` | Yarn scaling on all dims. | + +### 10.6 Mixture-of-experts (MoE) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `num_experts` | `null` | Experts per MoE layer; `null` means dense model. | +| `moe_layer_freq` | `1` | Every Nth layer is MoE (1 = every layer). | +| `moe_ffn_hidden_size` | `null` | Expert FFN hidden size. | +| `moe_shared_expert_overlap` | `false` | Shared expert overlaps routing. | +| `moe_shared_expert_intermediate_size` | `null` | Shared expert FFN size. | +| `moe_grouped_gemm` | `false` | Grouped GEMM for experts. | +| `moe_router_load_balancing_type` | `"aux_loss"` | Router balancing (`aux_loss`, `seq_aux_loss`, `sinkhorn`, `none`). | +| `moe_router_dtype` | `null` | Router activation dtype (`fp32`, `fp64`). | +| `moe_router_score_function` | `softmax` | `softmax` or `sigmoid` routing scores. | +| `moe_router_topk` | `2` | Experts to select per token. | +| `moe_router_pre_softmax` | `false` | Apply softmax before top-k. | +| `moe_router_num_groups` | `null` | Group-limited routing: number of expert groups. | +| `moe_router_group_topk` | `null` | Groups to pick before top-k inside groups. | +| `moe_router_topk_scaling_factor` | `null` | Scaling for routing logits. | +| `moe_router_enable_expert_bias` | `false` | Learnable per-expert bias. | +| `moe_router_bias_update_rate` | `1.0e-03` | Update rate for expert bias. | +| `moe_use_legacy_grouped_gemm` | `false` | Legacy grouped GEMM path. | +| `moe_aux_loss_coeff` | `0.0` | Auxiliary load-balancing loss weight. | +| `moe_z_loss_coeff` | `null` | Router z-loss coefficient. | +| `moe_input_jitter_eps` | `null` | Input jitter for router stability. | +| `moe_token_dispatcher_type` | `allgather` | Token dispatch algorithm (`allgather`, `alltoall`, `flex`, `alltoall_seq`). | +| `moe_enable_deepep` | `false` | DeepEP-style expert parallelism. | +| `moe_per_layer_logging` | `false` | Per-layer MoE statistics logging. | +| `moe_expert_capacity_factor` | `null` | Capacity factor for token dropping / padding. | +| `moe_pad_expert_input_to_capacity` | `false` | Pad expert batches to capacity. | +| `moe_token_drop_policy` | `probs` | Token dropping policy when over capacity. | +| `moe_extended_tp` | `false` | Extended tensor-parallel for experts. | +| `moe_use_upcycling` | `false` | Expert upcycling initialization. | +| `moe_permute_fusion` | `false` | Fuse token permutation for MoE. | +| `disable_primus_topk_router` | `false` | *Primus:* disable Primus top-k router patch. | +| `moe_router_force_load_balancing` | `false` | *Primus:* force load-balanced routing. | +| `use_deprecated_20241209_moe_layer` | `false` | *Primus:* legacy MoE layer implementation. | + +### 10.7 Logit softcapping (Primus / Grok-style) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `final_logit_softcapping` | `null` | Softcap value for final logits; `null` disables. | +| `attn_logit_softcapping` | `null` | Softcap for attention logits. | +| `router_logit_softcapping` | `null` | Softcap for MoE router logits. | + +--- + +## 11. Primus extensions + +### 11.1 Build and compile + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `disable_compile_dependencies` | `true` | *Primus:* avoid compiling dependency stacks in the trainer wrapper. | + +### 11.2 Primus-Turbo (`primus_turbo.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `enable_primus_turbo` | `false` | Master switch for Primus-Turbo integrations. Many sub-features require this plus specific kernels. | +| `use_turbo_attention` | `false` | Turbo attention implementation. | +| `use_sink_attention` | `false` | GPT-OSS-style learned sink attention. | +| `sink_sliding_window` | `0` | Sliding-window size for sink attention (GPT-OSS uses `128`). | +| `sink_window_even_layers_only` | `true` | Apply the sliding window only to even layers (GPT-OSS pattern). | +| `use_turbo_parallel_linear` | `false` | Turbo parallel linear layers. | +| `use_turbo_grouped_gemm` | `false` | Active Turbo grouped GEMM flag for MoE paths. | +| `use_turbo_grouped_mlp` | *(removed)* | Removed—use `use_turbo_grouped_gemm`. Passing this key now raises an assertion error (`use_turbo_grouped_mlp has been removed; please use use_turbo_grouped_gemm instead`). | +| `moe_use_fused_router_with_aux_score` | `false` | Fused MoE router with auxiliary scores. | +| `enable_turbo_attention_float8` | `false` | FP8 path inside Turbo attention (spacing in YAML is normalized to this key). | +| `use_turbo_deepep` | `false` | Turbo DeepEP expert communication. | +| `turbo_deepep_num_cu` | `32` | DeepEP compute units / channels. | +| `turbo_deepep_use_comm_stream` | `false` | Use a dedicated communication stream for DeepEP. | +| `turbo_sync_free_moe_stage` | `0` | Stage selector for sync-free MoE. | +| `use_turbo_fused_act_with_probs` | `false` | Fuse activation + probability tensors to remove redundant work. | +| `use_turbo_rms_norm` | `false` | Turbo RMSNorm kernels. | + +### 11.3 Zero-bubble pipeline (`zero_bubble.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `patch_zero_bubble` | `false` | Install Primus zero-bubble PP patches when `true`. | +| `debug_scheduler_table` | `false` | Print PP scheduler tables (also in `primus_pipeline.yaml`; last merge wins—defaults match). | +| `enable_zb_runtime` | `true` | Unified runtime for zero-bubble and related schedules. | +| `pre_communication_optimization` | `false` | Issue a tiny comm before real comm to tune overlap. | +| `zero_bubble_pipeline_timers_start_iter` | `100` | Start iter for auto-scheduler timers. | +| `zero_bubble_pipeline_timers_end_iter` | `110` | End iter for auto-scheduler timers. | +| `zero_bubble_max_pending_backward` | `auto` | Max pending backward ops (ZB1p vs ZB2p style); `auto` adapts. | +| `zero_bubble_adaptive_memory_limit_percentile` | `85` | GPU memory percentile cap for adaptive ZB. | +| `enable_optimizer_post_validation` | `false` | Post-optimizer validation step (needs FSDP path). | +| `enable_exactly_numeric_match` | `true` | Require bitwise match in post validation when enabled. | +| `enable_zero_bubble` | `true` | Enable zero-bubble schedule features in the ZB runtime. | +| `zero_bubble_v_schedule` | `false` | Zero-bubble “V” schedule without extra memory vs some baselines. | +| `zero_bubble_v_schedule_mem_setup` | `half` | Memory setup variant: `half`, `min`, or `zb`. | +| `enable_1f1b_v` | `false` | 1F1B-V schedule variant. | +| `allow_padding_num_layers` | `true` | Allow PP layer padding for divisibility. | +| `profile_memory_iter` | `-1` | Iteration to profile memory (`-1` disables). | +| `interleave_group_size` | `0` | Interleaved PP group size. | +| `offload_chunk_num` | `0` | Activation offload chunk count. | +| `offload_time` | `1.0` | Time budget for offload (scheduler hint). | +| `auto_offload_time` | `true` | Auto-tune offload timing. | +| `offload_overlap_sr` | `true` | Overlap save/resume in offload path. | +| `num_seq_splits` | `1` | Splits along sequence dimension for ZB. | +| `cpu_offload` | `false` | CPU offload of activations in ZB path. | + +### 11.4 Primus pipeline (`primus_pipeline.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `patch_primus_pipeline` | `false` | Enable Primus pipeline scheduling patches. | +| `pp_algorithm` | `"1f1b-interleaved"` | Schedule name (`1f1b`, `1f1b-interleaved`, `zero-bubble`, `zero-bubble-heuristic`, `zbv-formatted`, `v-half`, `v-min`). | +| `communication_method` | `"async_p2p"` | `async_p2p` or `batch_p2p` PP transfers. | +| `offload` | `false` | Generic PP activation offload toggle in Primus pipeline. | +| `offload_ops` | `""` | Comma-separated offload targets (`attn` today; other ops listed in-file are not supported yet). | +| `pp_max_mem` | `null` | `zero-bubble-heuristic` only: max activation memory per stage (`null` = unlimited). | +| `pp_cost_f` | `null` | `zero-bubble-heuristic` only: forward cost per stage (scalar or list; `null` = default 1000). | +| `pp_cost_b` | `null` | `zero-bubble-heuristic` only: backward cost per stage (scalar or list; `null` = default 1000). | +| `pp_cost_w` | `null` | `zero-bubble-heuristic` only: weight-grad cost per stage (scalar or list; `null` = default 1000). | + +`pp_warmup` and `dump_pp_data` are *Primus* helpers defined in `primus_megatron_module.yaml` (not `primus_pipeline.yaml`): + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `pp_warmup` | `false` | *Primus:* warm-up PP stages to reduce first-iteration latency. | +| `dump_pp_data` | `false` | *Primus:* dump PP tensors for debugging. | + +--- + +## 12. Reinforcement learning and GRPO-related settings + +*Source: `trainer_base.yaml`. Names follow Megatron’s `grpo_*` / `rl_*` prefixes (there is no `rl_grpo` single flag in these presets).* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `perform_rl_step` | `false` | Run RL / preference optimization steps (GRPO / LangRL integration). | +| `rl_prompts_per_eval` | `32` | Prompts per RL evaluation pass. | +| `grpo_prompts_per_step` | `32` | GRPO prompts sampled per training step. | +| `grpo_group_size` | `2` | Samples per prompt group for GRPO. | +| `grpo_iterations` | `2` | Inner GRPO iterations. | +| `grpo_clamp_eps_lower` | `0.01` | PPO-style lower clip epsilon. | +| `grpo_clamp_eps_upper` | `0.01` | Upper clip epsilon. | +| `grpo_kl_beta` | `0.001` | KL penalty weight toward reference policy. | +| `grpo_entropy_term_weight` | `0.0` | Entropy bonus weight. | +| `grpo_filter_groups_with_same_reward` | `false` | Drop groups with identical rewards. | +| `grpo_default_temperature` | `1.0` | Default softmax temperature for rollouts. | +| `grpo_default_top_p` | `0` | Top-p sampling (`0` often means disabled / greedy—see Megatron RL docs). | +| `langrl_inference_server_type` | `inplace_megatron` | LangRL inference backend. | +| `langrl_inference_server_conversation_template` | `null` | Conversation template path / name. | +| `langrl_env_config` | `null` | Environment / task YAML for LangRL. | +| `rl_offload_optimizer_during_inference` | `false` | Offload optimizer to CPU during rollout inference. | +| `rl_offload_kv_cache_during_training` | `false` | Offload KV cache while training forward runs. | +| `rl_remove_kv_cache_during_training` | `false` | Drop KV cache between RL phases to save memory. | +| `rl_reset_cuda_graphs` | `false` | Reset CUDA graphs when switching RL phases. | +| `rl_partial_rollouts` | `false` | Partial sequence rollouts. | +| `rl_inference_logprobs_is_correction` | `false` | Interpret inference logprobs as IS correction term. | +| `rl_importance_sampling_truncation_coef` | `null` | Truncate importance ratios at this value. | +| `rl_calculate_intra_group_similarity` | `false` | Log similarity within GRPO groups. | + +--- + +## 13. Additional specialized parameters + +*Source: `trainer_base.yaml` (remaining domains).* + +### 13.1 Vision pretraining + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `vision_pretraining` | `false` | Enable vision backbone pretraining. | +| `vision_pretraining_type` | `classify` | Objective (`classify`, etc.). | +| `vision_backbone_type` | `vit` | Vision backbone family. | +| `swin_backbone_type` | `tiny` | Swin variant size. | +| `num_classes` | `1000` | Classification classes. | +| `img_h` | `224` | Image height. | +| `img_w` | `224` | Image width. | +| `num_channels` | `3` | Input channels. | +| `patch_dim` | `16` | ViT patch size. | +| `classes_fraction` | `1.0` | Fraction of classes used. | +| `data_per_class_fraction` | `1.0` | Fraction of data per class. | + +### 13.2 RETRO + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `retro_project_dir` | `null` | RETRO project directory with indices. | +| `retro_add_retriever` | `false` | Add frozen retriever tower. | +| `retro_cyclic_train_iters` | `null` | Cyclic iterator length. | +| `retro_encoder_layers` | `2` | Retriever encoder layers. | +| `retro_encoder_hidden_dropout` | `0.1` | Retriever dropout. | +| `retro_encoder_attention_dropout` | `0.1` | Retriever attention dropout. | +| `retro_num_neighbors` | `2` | Neighbors per query chunk. | +| `retro_num_retrieved_chunks` | `2` | Chunks concatenated per neighbor set. | +| `retro_attention_gate` | `1` | Gating between retrieval and LM. | +| `retro_verify_neighbor_count` | `true` | Assert neighbor counts for debugging. | + +### 13.3 DINO self-supervised + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `dino_local_img_size` | `96` | Local crop size. | +| `dino_local_crops_number` | `10` | Number of local crops. | +| `dino_head_hidden_size` | `2048` | Projection head width. | +| `dino_bottleneck_size` | `256` | Bottleneck dimension. | +| `dino_freeze_last_layer` | `1` | Freeze last layer epochs. | +| `dino_norm_last_layer` | `false` | Normalize last layer weights. | +| `dino_warmup_teacher_temp` | `0.04` | Teacher temperature warmup start. | +| `dino_teacher_temp` | `0.07` | Teacher temperature. | +| `dino_warmup_teacher_temp_epochs` | `30` | Epochs to warm teacher temperature. | + +### 13.4 Biencoder / ICT / retriever utilities + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `ict_head_size` | `null` | ICT projection head width. | +| `biencoder_projection_dim` | `0` | Biencoder shared projection dimension. | +| `biencoder_shared_query_context_model` | `false` | Share query/context encoders. | +| `ict_load` | `null` | ICT checkpoint path. | +| `bert_load` | `null` | BERT encoder checkpoint for biencoder. | +| `titles_data_path` | `null` | Titles file for ICT datasets. | +| `query_in_block_prob` | `0.1` | Probability of in-block queries. | +| `use_one_sent_docs` | `false` | Single-sentence pseudo documents. | +| `evidence_data_path` | `null` | Evidence passages for open-domain QA. | +| `retriever_report_topk_accuracies` | `[]` | k values for top-k accuracy logging. | +| `retriever_score_scaling` | `false` | Scale retriever scores. | +| `block_data_path` | `null` | Block JSON data for retrieval. | +| `embedding_path` | `null` | Precomputed embeddings path. | +| `indexer_batch_size` | `128` | Batch size when building ANN index. | +| `indexer_log_interval` | `1000` | Indexer progress log interval. | + +### 13.5 Straggler detection + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `log_straggler` | `false` | Log straggler diagnostics. | +| `disable_straggler_on_startup` | `false` | Skip straggler detection at startup. | +| `straggler_ctrlr_port` | `65535` | Controller port for straggler service. | +| `straggler_minmax_count` | `1` | Min/max samples for straggler stats. | + +### 13.6 Inference-oriented options + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `inference_batch_times_seqlen_threshold` | `-1` | Heuristic threshold tying batch and sequence length. | +| `inference_dynamic_batching` | `false` | Dynamic batching for inference server. | +| `inference_dynamic_batching_buffer_size_gb` | `40.0` | GPU buffer budget (GB). | +| `inference_dynamic_batching_buffer_guaranteed_fraction` | `0.2` | Minimum reserved fraction of buffer. | +| `inference_dynamic_batching_buffer_overflow_factor` | `null` | Overflow growth factor. | +| `inference_dynamic_batching_max_requests_override` | `null` | Hard cap on concurrent requests. | +| `inference_dynamic_batching_max_tokens_override` | `null` | Hard cap on tokens in flight. | +| `max_tokens_to_oom` | `12000` | Token limit guard before OOM abort. | +| `output_bert_embeddings` | `false` | Return BERT pooled embeddings. | +| `bert_embedder_type` | `megatron` | `megatron` or `huggingface` embedder. | +| `flash_decode` | `false` | Flash decode kernels for incremental generation. | +| `enable_cuda_graph` | `false` | Capture CUDA graphs for inference. | +| `cuda_graph_warmup_steps` | `3` | Warm-up steps before capturing graphs. | +| `external_cuda_graph` | `false` | External graph provider hooks. | +| `cuda_graph_scope` | `full` | Graph scope (`full` or `attn`). | +| `inference_max_requests` | `8` | Max concurrent requests. | +| `inference_max_seq_length` | `2560` | Max prefill + decode tokens per request. | + +### 13.7 Fault tolerance package and tooling + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `enable_ft_package` | `false` | NVIDIA fault-tolerance package hooks. | +| `calc_ft_timeouts` | `false` | Auto-calculate FT timeouts. | +| `run_workload_inspector_server` | `false` | Run workload inspector sidecar. | + +### 13.8 Heterogeneous layers and process resilience + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `heterogeneous_layers_config_path` | `null` | JSON describing variable layer widths/types per layer. | +| `heterogeneous_layers_config_encoded_json` | `null` | Inline base64/JSON blob for heterogeneous layers. | +| `inprocess_restart` | `false` | In-process restart for fault recovery experiments. | + +### 13.9 Experimental and rerun controls + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `enable_experimental` | `false` | Gate experimental Megatron features. | +| `error_injection_rate` | `0` | Fraction of iterations with injected errors (testing). | +| `error_injection_type` | `transient_error` | `correct_result`, `transient_error`, or `persistent_error`. | +| `rerun_mode` | `disabled` | `disabled`, `validate_results`, or `report_stats` for rerun harness. | + +--- + +### Related documentation + +- Megatron-LM argument definitions: [`megatron/training/arguments.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/training/arguments.py) +- Primus Megatron presets: `primus/configs/modules/megatron/` +- Primus Megatron model presets: `primus/configs/models/megatron/` diff --git a/docs/03-configuration-reference/torchtitan-parameters.md b/docs/03-configuration-reference/torchtitan-parameters.md new file mode 100644 index 000000000..d5e8775d4 --- /dev/null +++ b/docs/03-configuration-reference/torchtitan-parameters.md @@ -0,0 +1,362 @@ +# TorchTitan backend configuration reference + +This page lists Primus preset keys and common TorchTitan `JobConfig` fields used when `framework: torchtitan`. Defaults are taken from the TorchTitan module preset (`pre_trainer.yaml`), its `extends` chain (`module_base.yaml`, `quantize.yaml`), and the example model preset `llama3_8B.yaml`. It is not a complete upstream TorchTitan `JobConfig` reference. + +**Where parameters live.** Provide overrides under `modules.pre_trainer.overrides:` in your experiment YAML. TorchTitan’s `JobConfig` is hierarchical: use **dot notation** for flat overrides, or nest YAML objects under `overrides`—both are equivalent when merged. + +**Example (flat dot paths):** + +```yaml +framework: torchtitan + +modules: + pre_trainer: + overrides: + training.steps: 20000 + training.global_batch_size: 512 + optimizer.lr: 0.00015 + parallelism.tensor_parallel_degree: 2 +``` + +**Example (nested YAML):** + +```yaml +modules: + pre_trainer: + overrides: + training: + steps: 20000 + global_batch_size: 512 + optimizer: + lr: 0.00015 + parallelism: + tensor_parallel_degree: 2 +``` + +**Presets.** + +- Module presets: `primus/configs/modules/torchtitan/` (main entry: `pre_trainer.yaml`). +- Model presets: `primus/configs/models/torchtitan/` (example: `llama3_8B.yaml`). + +**Mapping to TorchTitan.** Keys are translated into TorchTitan’s `JobConfig` via `TorchTitanJobConfigBuilder` (same nested structure as upstream TorchTitan). + +**Upstream reference.** TorchTitan repository and documentation: [https://github.com/pytorch/torchtitan](https://github.com/pytorch/torchtitan) (vendored as the `third_party/torchtitan` submodule). + +--- + +## 1. Base module parameters + +*Source: `primus/configs/modules/module_base.yaml` (merged before TorchTitan-specific keys; `pre_trainer.yaml` does not override `trainable`).* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `trainable` | `false` | Whether this module is active in training orchestration. (TorchTitan preset inherits `false` from `module_base.yaml`.) | +| `sink_level` | `null` | Structured logging sink level; `null` uses defaults. | +| `file_sink_level` | `DEBUG` | Minimum level for file logging. | +| `stderr_sink_level` | `INFO` | Minimum level for stderr logging. | + +--- + +## 2. Training (`training.*`) + +*Source: `primus/configs/modules/torchtitan/pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `training.mock_data` | `true` | Primus preset extension: use synthetic data instead of reading `dataset_path`. | +| `training.debug_moe_force_load_balance` | `false` | Primus preset extension/debug helper to force MoE load balancing behavior. | +| `training.dataset` | `c4` | Dataset name key for TorchTitan dataset loaders. | +| `training.dataset_path` | `null` | Filesystem or remote path to dataset assets. | +| `training.deterministic` | `false` | Prefer deterministic algorithms (often slower). | +| `training.enable_cpu_offload` | `false` | Offload optimizer or activations to CPU when supported. | +| `training.gc_debug` | `false` | Extra garbage-collection diagnostics. | +| `training.gc_freq` | `50` | Run Python GC every N steps when enabled. | +| `training.global_batch_size` | `-1` | Global batch size across all ranks (`-1` often means “auto” / unset in TorchTitan). | +| `training.local_batch_size` | `8` | Per-rank microbatch size before gradient accumulation. | +| `training.max_norm` | `1.0` | Gradient clipping max norm (global). | +| `training.mixed_precision_param` | `bfloat16` | Parameter dtype for mixed precision (`bfloat16`, `float16`, etc.). | +| `training.mixed_precision_reduce` | `float32` | Dtype for reduction / gradient accumulation. | +| `training.seed` | `null` | RNG seed; `null` lets the framework choose. | +| `training.seq_len` | `2048` | Sequence length per sample. | +| `training.steps` | `10000` | Total optimizer steps. | + +--- + +## 3. Optimizer (`optimizer.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `optimizer.name` | `AdamW` | Optimizer class (`AdamW`, `Adam`, …). | +| `optimizer.lr` | `0.0008` | Base learning rate. | +| `optimizer.beta1` | `0.9` | First moment decay. | +| `optimizer.beta2` | `0.95` | Second moment decay. | +| `optimizer.eps` | `1.0e-08` | Numerical stability term. | +| `optimizer.weight_decay` | `0.1` | Weight decay coefficient. | +| `optimizer.implementation` | `fused` | Kernel implementation (`fused`, `foreach`, …). | +| `optimizer.early_step_in_backward` | `false` | Experimental: step optimizer during backward when supported. | + +--- + +## 4. Learning rate scheduler (`lr_scheduler.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `lr_scheduler.decay_ratio` | `null` | Fraction of training at the end used for decay; `null` uses framework default. | +| `lr_scheduler.decay_type` | `linear` | LR decay curve (`linear`, `cosine`, etc.). | +| `lr_scheduler.min_lr_factor` | `0.0` | LR floor as a fraction of base LR after decay. | +| `lr_scheduler.warmup_steps` | `200` | Linear warmup steps before decay. | + +--- + +## 5. Parallelism (`parallelism.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `parallelism.tensor_parallel_degree` | `1` | Tensor parallelism (intra-layer) degree. | +| `parallelism.pipeline_parallel_degree` | `1` | Pipeline parallelism stages. | +| `parallelism.pipeline_parallel_microbatch_size` | `1` | Microbatches per pipeline round. | +| `parallelism.pipeline_parallel_schedule` | `1F1B` | Pipeline schedule name (`1F1B`, `GPipe`, …). | +| `parallelism.pipeline_parallel_schedule_csv` | `''` | Optional CSV schedule definition. | +| `parallelism.pipeline_parallel_split_points` | `[]` | Layer indices for manual PP splits. | +| `parallelism.pipeline_parallel_layers_per_stage` | `null` | Layers per stage when auto-balanced. | +| `parallelism.pipeline_parallel_first_stage_less_layers` | `1` | Fewer layers on first PP stage (for imbalance). | +| `parallelism.pipeline_parallel_last_stage_less_layers` | `1` | Fewer layers on last PP stage. | +| `parallelism.data_parallel_shard_degree` | `-1` | FSDP / shard degree (`-1` = auto). | +| `parallelism.data_parallel_replicate_degree` | `1` | Replicated data-parallel groups. | +| `parallelism.expert_parallel_degree` | `1` | Expert parallelism for MoE models. | +| `parallelism.expert_tensor_parallel_degree` | `1` | Tensor parallelism inside experts. | +| `parallelism.context_parallel_degree` | `1` | Context (sequence) parallelism degree. | +| `parallelism.context_parallel_rotate_method` | `allgather` | Communication pattern for context parallel. | +| `parallelism.disable_loss_parallel` | `false` | Disable loss parallel layout when TP is used. | +| `parallelism.enable_async_tensor_parallel` | `false` | Overlap TP collectives with compute. | +| `parallelism.enable_compiled_autograd` | `false` | Use `torch.compile` on autograd regions. | +| `parallelism.fsdp_reshard_after_forward` | `default` | FSDP reshard policy (`default`, `always`, `never`). | +| `parallelism.module_fqns_per_model_part` | `null` | Map of pipeline stage → module FQNs for multi-part models. | + +--- + +## 6. Checkpoint (`checkpoint.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `checkpoint.enable` | `false` | Master switch for checkpointing. | +| `checkpoint.folder` | `checkpoint` | Output directory for checkpoints. | +| `checkpoint.interval` | `500` | Save every N steps. | +| `checkpoint.initial_load_path` | `null` | Path to load from at startup. | +| `checkpoint.initial_load_model_only` | `true` | Load weights only (skip optimizer/scheduler). | +| `checkpoint.initial_load_in_hf` | `false` | Load initial weights from Hugging Face format. | +| `checkpoint.last_save_model_only` | `true` | Final save stores weights only. | +| `checkpoint.last_save_in_hf` | `false` | Export final weights in Hugging Face format. | +| `checkpoint.export_dtype` | `float32` | Dtype for exported checkpoints. | +| `checkpoint.async_mode` | `disabled` | Async checkpoint (`disabled`, `async`, …). | +| `checkpoint.keep_latest_k` | `10` | Retain only the newest k checkpoints. | +| `checkpoint.load_step` | `-1` | Step index to load (`-1` = latest). | +| `checkpoint.exclude_from_loading` | `[]` | FQNs or keys to skip when loading. | +| `checkpoint.enable_first_step_checkpoint` | `false` | Save checkpoint at step 0 for debugging. | +| `checkpoint.create_seed_checkpoint` | `false` | Save a seed checkpoint before training starts. | + +--- + +## 7. Activation checkpoint (`activation_checkpoint.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `activation_checkpoint.mode` | `none` | Activation checkpointing mode (`none`, `selective`, `full`). | +| `activation_checkpoint.selective_ac_option` | `"2"` | Selective AC policy string (TorchTitan-specific). | +| `activation_checkpoint.per_op_sac_force_recompute_mm_shapes_by_fqns` | `["moe.router.gate"]` | FQNs that always recompute matmuls in selective AC. | +| `activation_checkpoint.early_stop` | `false` | Stop AC early in certain subgraphs. | + +--- + +## 8. Metrics and profiling + +### 8.1 Metrics (`metrics.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `metrics.disable_color_printing` | `false` | Disable ANSI colors in logs. | +| `metrics.enable_tensorboard` | `false` | Write TensorBoard scalars. | +| `metrics.enable_wandb` | `false` | Log to Weights & Biases. | +| `metrics.log_freq` | `10` | Steps between metric logs. | +| `metrics.save_for_all_ranks` | `false` | Save metric files per rank (not just rank 0). | +| `metrics.save_tb_folder` | `tb` | TensorBoard subdirectory / name. | + +### 8.2 Profiling (`profiling.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `profiling.enable_profiling` | `false` | Enable PyTorch profiler traces. | +| `profiling.enable_memory_snapshot` | `false` | Capture CUDA memory snapshots. | +| `profiling.profile_freq` | `10` | Steps between profiler activations. | +| `profiling.save_traces_folder` | `profile_traces` | Directory for profiler traces. | +| `profiling.save_memory_snapshot_folder` | `memory_snapshot` | Directory for memory snapshots. | + +--- + +## 9. Quantization (`quantize.*`) + +*Source: `primus/configs/modules/torchtitan/quantize.yaml` (merged into the module preset).* + +### 9.1 Linear FP8 (`quantize.linear.float8.*`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `quantize.linear.float8.enable_fsdp_float8_all_gather` | `false` | FP8 all-gather for FSDP sharded params (recommended for tensorwise scaling). | +| `quantize.linear.float8.precompute_float8_dynamic_scale_for_fsdp` | `false` | Precompute dynamic scales for FSDP FP8. | +| `quantize.linear.float8.recipe_name` | `null` | Recipe (`tensorwise`, `rowwise`, `rowwise_with_gw_hp`); `null` disables. | +| `quantize.linear.float8.filter_fqns` | `[]` | Module FQNs to skip for FP8 training. | +| `quantize.linear.float8.emulate` | `false` | Emulate FP8 in FP32 (no FP8 HW); not compatible with `torch.compile`. | + +### 9.2 Linear MX (`quantize.linear.mx.*`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `quantize.linear.mx.mxfp8_dim1_cast_kernel_choice` | `"triton"` | Kernel backend for MXFP8 dim-1 cast (`triton`, `cuda`, `torch`). | +| `quantize.linear.mx.recipe_name` | `"mxfp8_cublas"` | MX recipe name (see torchao `mx_formats`). | +| `quantize.linear.mx.filter_fqns` | `["output"]` | FQNs to skip; output layer skipped by default. | + +### 9.3 Grouped GEMM FP8 (`quantize.grouped_mm.float8.*`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `quantize.grouped_mm.float8.fqns` | `[]` | MoE layer FQNs for FP8 grouped GEMM (prototype; may require torchao nightly). | + +### 9.4 Grouped GEMM MX (`quantize.grouped_mm.mx.*`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `quantize.grouped_mm.mx.recipe_name` | `"mxfp8"` | MX recipe for grouped GEMMs. | +| `quantize.grouped_mm.mx.fqns` | `[]` | MoE module FQNs for MXFP8 grouped GEMM (prototype). | + +--- + +## 10. Compile (`compile.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `compile.enable` | `true` | Enable `torch.compile` on selected subsystems. | +| `compile.components` | `["model", "loss"]` | Which components to compile. | + +--- + +## 11. Communication and fault tolerance + +### 11.1 Communicator (`comm.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `comm.init_timeout_seconds` | `300` | Timeout for initial process-group setup. | +| `comm.train_timeout_seconds` | `100` | Timeout for training collectives. | +| `comm.trace_buf_size` | `20000` | Flight recorder buffer size for NCCL traces. | +| `comm.save_traces_folder` | `comm_traces` | Where to dump communication traces. | + +### 11.2 Fault tolerance (`fault_tolerance.*`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `fault_tolerance.enable` | `false` | Enable fault-tolerant training hooks. | +| `fault_tolerance.process_group` | `gloo` | Backend for control-plane process group. | +| `fault_tolerance.process_group_timeout_ms` | `10000` | Control-plane timeout. | +| `fault_tolerance.replica_id` | `0` | Replica index in elastic setups. | +| `fault_tolerance.group_size` | `0` | Group size (0 = unset / default). | +| `fault_tolerance.min_replica_size` | `1` | Minimum replicas to continue. | +| `fault_tolerance.semi_sync_method` | `null` | Optional semi-synchronous strategy name. | + +### 11.3 Memory estimation (`memory_estimation.*`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `memory_estimation.enable` | `false` | Run memory-estimation fake-mode passes. | +| `memory_estimation.disable_fake_mode` | `false` | Disable fake tensor mode inside estimation. | + +### 11.4 Experimental (`experimental.*`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `experimental.custom_import` | `""` | Optional Python module import path for custom extensions. | +| `experimental.custom_args_module` | `"primus.backends.torchtitan.primus_turbo_extensions.config_extension"` | Module providing extra `JobConfig` fields for Primus-Turbo. | + +--- + +## 12. Primus-Turbo (`primus_turbo.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `primus_turbo.enable_primus_turbo` | `true` | Master switch for Primus-Turbo integrations in TorchTitan. | +| `primus_turbo.enable_attention_float8` | `false` | FP8 attention path inside Turbo attention. | +| `primus_turbo.use_turbo_attention` | `true` | Use Turbo attention kernels. | +| `primus_turbo.use_classic_attention` | `false` | Fall back to classic attention implementation. | +| `primus_turbo.use_turbo_async_tp` | `true` | Async tensor-parallel communication in Turbo. | +| `primus_turbo.use_turbo_mx_linear` | `true` | MX linear layers via Turbo. | +| `primus_turbo.use_turbo_float8_linear` | `true` | FP8 linear layers via Turbo. | +| `primus_turbo.use_turbo_grouped_mm` | `false` | Turbo grouped GEMM for MoE (off by default). | +| `primus_turbo.use_moe_fp8` | `true` | FP8 paths for MoE experts when applicable. | +| `primus_turbo.enable_embedding_autocast` | `true` | Autocast policy around embeddings for Turbo. | + +--- + +## 13. Model (`model.*` and `job.*`) + +### 13.1 Model preset (`models.*` / `model.*`) + +*Example defaults from `primus/configs/models/torchtitan/llama3_8B.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `model.name` | `"llama3"` | Model family key for TorchTitan recipes. | +| `model.flavor` | `"8B"` | Size / variant within the family. | +| `model.hf_assets_path` | `"meta-llama/Meta-Llama-3-8B"` | Hugging Face Hub repo or local path for weights/tokenizer. | +| `model.converters` | `["primus_turbo"]` | Weight converter pipeline stages applied at load. | + +### 13.2 Job metadata (`job.*`) + +*Source: `llama3_8B.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `job.dump_folder` | `"./outputs"` | Root directory for logs, checkpoints, and exports. | +| `job.description` | `"Llama 3 8B training"` | Human-readable label for run metadata. | + +--- + +## 14. Validation (`validation.*`) + +*Source: `pre_trainer.yaml`.* + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `validation.enable` | `false` | Run periodic validation loops. | +| `validation.dataset` | `c4_validation` | Validation dataset key. | +| `validation.dataset_path` | `null` | Filesystem path to validation data. | +| `validation.local_batch_size` | `8` | Per-rank validation batch size. | +| `validation.seq_len` | `2048` | Validation sequence length. | +| `validation.freq` | `10` | Run validation every N training steps. | +| `validation.steps` | `-1` | Max validation steps (`-1` = full pass / framework default). | + +--- + +### Related documentation + +- TorchTitan repository and documentation: [https://github.com/pytorch/torchtitan](https://github.com/pytorch/torchtitan) +- Primus TorchTitan presets: `primus/configs/modules/torchtitan/` +- Primus TorchTitan model presets: `primus/configs/models/torchtitan/` diff --git a/docs/04-technical-guides/README.md b/docs/04-technical-guides/README.md new file mode 100644 index 000000000..4ca56d079 --- /dev/null +++ b/docs/04-technical-guides/README.md @@ -0,0 +1,22 @@ +# Technical guides + +Deep technical topics for advanced users. + +- [Parallelism strategies](parallelism-strategies.md): DP, TP, PP, SP, CP, EP, FSDP explained +- [Parallelism configuration](parallelism-configuration.md): per-backend parallelism setup and batch size relationships +- [Collective operations](collective-operations.md): NCCL/RCCL operations and their role in each parallelism strategy +- [Performance tuning](performance-tuning.md): HipBLASLt, Primus-Turbo, FP8, MoE optimization +- [MoE training deep-dive](moe-training.md): bottlenecks and Primus-Turbo optimizations for Mixture-of-Experts models +- [Data preparation](data-preparation.md): tokenization, data formats, mock data +- [Checkpoint management](checkpoint-management.md): formats, save/load, distributed checkpointing +- [Multi-node networking](multi-node-networking.md): InfiniBand, RoCE, AINIC configuration +- [Profiling and observability](profiling-and-observability.md): Torch profiler, TraceLens, memory snapshots, projection, pp_vis +- [Logging and experiment tracking](logging-and-experiment-tracking.md): TensorBoard, WandB, MLflow setup per backend +- [Fault tolerance and elastic training](fault-tolerance-and-elastic-training.md): graceful exit, auto-resume, in-process restart, torchft +- [Determinism and reproducibility](determinism-and-reproducibility.md): deterministic mode, seeds, trade-offs +- [Diffusion models](diffusion-models/README.md): Flux diffusion architecture, data pipeline, and FP8 / MXFP4 training +- [Native SFT and LoRA](native-sft-lora.md): Megatron-native SFT/LoRA runbook (BF16 / FP8 / FP4), no Megatron-Bridge dependency + +--- + +[← Documentation home](../README.md) diff --git a/docs/04-technical-guides/checkpoint-management.md b/docs/04-technical-guides/checkpoint-management.md new file mode 100644 index 000000000..a3fd38a8b --- /dev/null +++ b/docs/04-technical-guides/checkpoint-management.md @@ -0,0 +1,203 @@ +# Checkpoint management + +Checkpoints capture **model state**, **optimizer state**, and **training progress** (iteration or step counters, schedulers, and related metadata). They are essential for **fault tolerance** (resume after failure), **experiment management** (reproducibility and comparison), and **hand-offs** between pretraining, fine-tuning, and conversion workflows. + +Primus is YAML-driven: checkpoint behavior is configured per backend. **Megatron-LM**, **TorchTitan**, and **MaxText** each expose their own checkpoint surfaces; this guide maps the knobs you set in Primus configs. + +**Primary sources in this repository** + +| Area | File | +|------|------| +| Megatron trainer defaults | `primus/configs/modules/megatron/trainer_base.yaml` | +| Primus Megatron extensions | `primus/configs/modules/megatron/primus_megatron_module.yaml` | +| TorchTitan defaults | `primus/configs/modules/torchtitan/pre_trainer.yaml` | +| Megatron checkpoint benchmark | `benchmark/megatron/checkpoint/README.md` | + +--- + +## 1. Overview + +- **What is saved:** Typically model parameters, optimizer state, RNG state, and iteration/step tracking—exact contents depend on flags such as `no_save_optim` / `no_save_rng` (Megatron) or `initial_load_model_only` (TorchTitan). +- **Why it matters:** Long runs on AMD GPU clusters benefit from periodic saves to durable storage; resuming or branching experiments requires consistent paths and formats. +- **Backend-specific systems:** Each training backend integrates its own checkpoint pipeline; Primus wires YAML into those backends without forcing a single universal format across Megatron, TorchTitan, and MaxText. + +--- + +## 2. Megatron checkpoint configuration + +Megatron-related options live on the trainer configuration merged from `trainer_base.yaml` and `primus_megatron_module.yaml`. Defaults below are taken from `trainer_base.yaml` unless noted. + +### Core paths and cadence + +| Parameter | Default (`trainer_base.yaml`) | Description | +|-----------|-------------------------------|-------------| +| `save` | `null` | Directory where new checkpoints are written. | +| `load` | `null` | Directory to load from when **resuming** training. | +| `save_interval` | `20000` | Save every *N* iterations. | +| `finetune` | `false` | When `true`, loads weights but **resets** the iteration counter (typical fine-tune entry). | + +### Format and detection + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `ckpt_format` | `torch_dist` | Checkpoint format: `torch` (legacy single-file style), `torch_dist` (distributed), or `zarr`. | +| `auto_detect_ckpt_format` | `false` | When loading, infer format automatically. | +| `pretrained_checkpoint` | `null` | Path to a **pretrained** checkpoint. | +| `ckpt_step` | `null` | Load a specific step from the pretrained checkpoint when applicable. | + +### Optimizer and RNG inclusion + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `no_save_optim` | `null` | When set truthy, **omit** optimizer state from saves. | +| `no_save_rng` | `null` | When set truthy, **omit** RNG state from saves. | +| `no_load_optim` | `null` | When set truthy, **do not** restore optimizer from checkpoint. | +| `no_load_rng` | `null` | When set truthy, **do not** restore RNG from checkpoint. | + +### Performance and distributed I/O + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `async_save` | `null` | Asynchronous checkpoint saving to reduce time blocking the training loop. | +| `ckpt_fully_parallel_save` | `true` | Parallel save path for distributed checkpoints. | +| `ckpt_fully_parallel_load` | `false` | Parallel load for distributed checkpoints. | +| `ckpt_assume_constant_structure` | `false` | Optimization when model structure is fixed across saves/loads. | +| `non_persistent_save_interval` | `null` | Save to **fast local** storage on a different cadence than persistent saves. | + +Related keys in `trainer_base.yaml` for non-persistent checkpoints include `non_persistent_ckpt_type`, `non_persistent_global_ckpt_dir`, `non_persistent_local_ckpt_dir`, and `non_persistent_local_ckpt_algo` (default `"fully_parallel"`). + +### Format conversion + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `ckpt_convert_format` | `null` | Target format for conversion (`torch`, `torch_dist`, or `zarr`). | +| `ckpt_convert_save` | `null` | Output directory for converted checkpoints. | +| `ckpt_convert_update_legacy_dist_opt_format` | `false` | Update legacy distributed optimizer layout when converting. | + +### Primus extensions + +Defined in `primus/configs/modules/megatron/primus_megatron_module.yaml` and implemented in `primus/backends/megatron/patches/checkpoint_patches.py` and `primus/backends/megatron/patches/args/checkpoint_path_patches.py`. + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `auto_continue_train` | `false` | When `true`, **automatically resume** from the latest checkpoint under `save` (adjusts load/finetune and related flags). | +| `disable_last_saving` | `false` | When `true`, **skip** the final checkpoint at shutdown (useful for benchmarking or when only periodic saves matter). | + +--- + +## 3. TorchTitan checkpoint configuration + +TorchTitan checkpoint options are grouped under `checkpoint` in `primus/configs/modules/torchtitan/pre_trainer.yaml`. + +| Parameter | Default (`pre_trainer.yaml`) | Description | +|-----------|------------------------------|-------------| +| `checkpoint.enable` | `false` | Master switch for checkpointing. | +| `checkpoint.folder` | `checkpoint` | Output directory (relative to run layout unless given as absolute). | +| `checkpoint.interval` | `500` | Save every *N* **steps**. | +| `checkpoint.initial_load_path` | `null` | Path for **initial** load (cold start or migration). | +| `checkpoint.initial_load_model_only` | `true` | Load **weights only**, not optimizer state. | +| `checkpoint.initial_load_in_hf` | `false` | Load initial weights from **Hugging Face** layout. | +| `checkpoint.last_save_model_only` | `true` | On last save, write **model only**. | +| `checkpoint.last_save_in_hf` | `false` | Write final checkpoint in **Hugging Face** format. | +| `checkpoint.export_dtype` | `float32` | Dtype for exported checkpoints. | +| `checkpoint.async_mode` | `disabled` | Asynchronous checkpoint mode. | +| `checkpoint.keep_latest_k` | `10` | Retain only the **K** most recent checkpoints. | +| `checkpoint.load_step` | `-1` | Load a specific step (`-1` typically means latest or default behavior per backend). | +| `checkpoint.exclude_from_loading` | `[]` | Glob or pattern list to **exclude** from restore. | +| `checkpoint.enable_first_step_checkpoint` | `false` | Optional checkpoint at step 0. | +| `checkpoint.create_seed_checkpoint` | `false` | Create a seed checkpoint when enabled. | + +TorchTitan also defines `activation_checkpoint` (activation recomputation) separately from persistent training checkpoints—do not confuse the two sections in `pre_trainer.yaml`. + +--- + +## 4. MaxText checkpoint configuration + +MaxText (JAX) uses configuration keys surfaced in Primus documentation and MaxText configs under `third_party/maxtext`. Primus overlay presets set the defaults shown below, while upstream MaxText `base.yml` is still loaded at runtime via `base_config: "base.yml"` and may define different upstream defaults. Typical training flags: + +| Parameter | Typical default | Description | +|-----------|-----------------|-------------| +| `enable_checkpointing` | `false` | Enable Orbax (or configured) checkpoint saves. | +| `async_checkpointing` | `false` | When checkpointing is enabled, use **async** checkpoint workers. | + +See `docs/03-configuration-reference/maxtext-parameters.md` for the full MaxText parameter table and interaction with training runs. + +--- + +## 5. Checkpoint formats (Megatron) + +| Format | Behavior | Notes | +|--------|----------|--------| +| `torch` | Classic PyTorch save/load; often **one file per rank** in distributed settings. | Simple but less flexible for topology changes. | +| `torch_dist` | **Distributed** checkpoint format with **resharding** support (e.g., changing tensor/pipeline parallel degree between save and load). | **Recommended** for many production flows that may change parallelism. | +| `zarr` | Zarr-backed checkpoint storage. | Useful when the stack and storage backend support it. | + +**Recommendation:** Prefer `torch_dist` for production when you need **flexibility across parallel layouts** and scalable I/O (see `ckpt_fully_parallel_save` / `ckpt_fully_parallel_load` in Megatron config). + +--- + +## 6. Common workflows + +**Resume training** + +- Set `load` to the checkpoint directory produced by a previous run. +- Keep `save` pointed at the directory for **new** checkpoints (often the same tree with a new run id, depending on your layout). +- Ensure `finetune` is `false` when you want to **continue** iteration counts. + +**Fine-tune from a pretrained checkpoint** + +- Set `load` (and optionally `pretrained_checkpoint` / `ckpt_step` as appropriate). +- Set `finetune: true` so iteration counters reset while weights load. + +**Auto-resume (Primus Megatron extension)** + +- Set `auto_continue_train: true` in the Megatron module config. +- Primus searches for the latest checkpoint under `save` and aligns load/optimizer flags; see `primus/backends/megatron/patches/checkpoint_patches.py` for behavior details. + +**Convert checkpoint format** + +- Set `ckpt_convert_format` (for example `torch_dist`) and `ckpt_convert_save` to the output directory. + +**Import Hugging Face weights (TorchTitan)** + +- Set `checkpoint.initial_load_in_hf: true` and `checkpoint.initial_load_path` to the HF model directory. + +--- + +## 7. Benchmarking checkpoints + +The Megatron checkpoint benchmark lives in `benchmark/megatron/checkpoint/`. + +**Entry points** + +- `benchmark/megatron/checkpoint/ckpt_launch.py`—main launcher (requires a Primus YAML config). +- `benchmark/megatron/checkpoint/ckpt_report.py`—reporting utility (can be run separately). + +**Example** (from `benchmark/megatron/checkpoint/README.md`): + +```bash +export DATA_PATH=/PATH/TO/DATA +python3 benchmark/megatron/checkpoint/ckpt_launch.py \ + --yaml-config-path examples/megatron/configs/MI300X/mixtral_8x7B_v0.1-pretrain.yaml \ + --nnodes 1 +``` + +The tool reports save/load times, bandwidth, and configuration echoes (world size, `ckpt_format`, `async_save`, paths, and more). Truncate or clean leftover output directories between runs if permissions or stale outputs cause issues. + +--- + +## 8. Best practices + +- Enable **`async_save`** (Megatron) for large models when supported, to limit training stalls during checkpoint windows. +- Set **`save_interval`** from **economic** criteria: frequent enough to limit lost work, infrequent enough to avoid storage and throughput bottlenecks (Megatron default in `trainer_base.yaml` is `20000`—override per job). +- Use **`non_persistent_save_interval`** with fast **local SSD** for frequent snapshots and a slower interval to **NFS** or object storage for durability. +- **Validate** resume and fine-tune paths on short runs before multi-week jobs; confirm `finetune` and `auto_continue_train` behave as intended. +- For TorchTitan, enable **`checkpoint.enable`** explicitly and set **`checkpoint.keep_latest_k`** to bound disk usage. + +--- + +## Related documentation + +- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) +- [MaxText parameters](../03-configuration-reference/maxtext-parameters.md) +- [Benchmark suite](../02-user-guide/benchmarking.md) diff --git a/docs/04-technical-guides/collective-operations.md b/docs/04-technical-guides/collective-operations.md new file mode 100644 index 000000000..43637dda3 --- /dev/null +++ b/docs/04-technical-guides/collective-operations.md @@ -0,0 +1,289 @@ +# NCCL/RCCL collective operations guide + +Distributed training spends a large fraction of wall time in **collective communication**: many GPUs must exchange gradients, parameters, or activations in coordinated patterns. On AMD GPUs, **RCCL** (ROCm Collective Communications Library) provides these operations with an API aligned to **NCCL** (NVIDIA Collective Communications Library), so most concepts and environment variables carry over between vendors. + +This guide explains core collectives, how they map to parallelism strategies in Primus (Megatron-LM, TorchTitan), and how to benchmark and troubleshoot communication. + +For Megatron knobs like `overlap_grad_reduce` and TorchTitan parallelism flags, see [Megatron parameters](../03-configuration-reference/megatron-parameters.md). For `NCCL_*` / `RCCL_*` environment variables, see [Environment variables](../03-configuration-reference/environment-variables.md). + +--- + +## 1. Introduction + +### What are collective operations? + +A **collective** is a multi-party communication pattern where **every participant** (or a defined **process group**) follows the same operation: combine tensors, broadcast, scatter pieces, or exchange shards. Unlike a single **Send/Recv** pair, collectives are **synchronized** by construction and are implemented with optimized algorithms (ring, tree, etc.). + +### NCCL vs RCCL + +| | NCCL | RCCL | +|---|------|------| +| Vendor | NVIDIA CUDA | AMD ROCm | +| Role | GPU collective communication | GPU collective communication | +| Typical API surface | C/C++ and bindings used by PyTorch distributed | ROCm stack; PyTorch uses similar backends | + +Application code written for **PyTorch distributed** (e.g. `torch.distributed`) generally selects the backend provided by the stack (**nccl** on NVIDIA, **rccl** on AMD). **Operation names and semantics** (AllReduce, AllGather, …) align so that **framework-level** code and tuning guides are largely **portable**. + +### Process groups + +Not every rank talks to every other rank in every step. **Process groups** define **subsets of ranks** that participate in a collective (e.g. only **tensor-parallel** ranks, only **data-parallel** ranks). Correctness and performance depend on **matching ranks** to the same group for each layer or phase of training. + +--- + +## 2. Core collective operations + +Below, \(n\) is the number of ranks in the process group, and \(S\) is the size of the logical tensor being reduced or moved (per-rank message size in ring formulations). **Complexity** expressions are **standard ring-style** approximations for **amount of data moved per rank** relative to \(S\); real implementations pick algorithms based on message size, topology, and environment. + +--- + +### AllReduce + +**What it does:** Each rank contributes a tensor; the **element-wise reduction** (typically **sum**) is applied across ranks, and the **full result** is **replicated** on every rank. + +``` +Rank 0: [a0] Rank 1: [a1] Rank 2: [a2] + \ | / + \ | / + --> REDUCE(sum) <-- + | + All ranks: [a0+a1+a2] +``` + +**Complexity (ring, per-rank data moved):** about \(\frac{2(n-1)}{n} S\). + +**Where used:** **Data-parallel** gradient synchronization; **tensor-parallel** partial sums; any step that needs **identical** tensors on all ranks after a reduction. + +--- + +### AllGather + +**What it does:** Each rank holds **one shard**; every rank receives the **concatenation** (or stacked layout) of **all shards**. + +``` +Rank 0: [x0] Rank 1: [x1] Rank 2: [x2] + \ | / + \ | / + --> ALL GATHER --> +Each rank: [x0 | x1 | x2] +``` + +**Complexity (ring):** about \(\frac{n-1}{n} S\) **if** each rank contributes \(S/n\); more generally scales with gathering \(n-1\) other shards of comparable size. + +**Where used:** **FSDP / ZeRO-3** parameter gather before forward; **TP** weight or activation assembly depending on layout. + +--- + +### ReduceScatter + +**What it does:** Conceptually **AllReduce** then **split**: each rank ends with **one shard** of the reduced result (each rank’s shard is the reduction over corresponding positions from all ranks’ inputs). + +``` +Inputs per rank: full-sized chunks (partial sums local) + | + Reduce + partition + | +Rank i gets shard i of the fully reduced tensor +``` + +**Complexity (ring):** about \(\frac{n-1}{n} S\) for the common balanced case. + +**Where used:** **FSDP** gradient **sharding** after backward; **sequence parallelism** with TP (activation distribution); distributed optimizer flows that **scatter** reduced pieces. + +--- + +### AllToAll + +**What it does:** Each rank sends a **distinct slice** to every other rank; every rank receives from every rank (matrix transpose of data ownership). + +``` + From rank 0..n-1 + | + +---------+---------+ + | scatter per dest | + v v v +Each rank receives its column/row of the logical matrix +``` + +**Where used:** **Expert parallelism** token **dispatch** and **combine** in MoE; some **sparsity** and **parallel embedding** layouts. + +--- + +### Broadcast + +**What it does:** One **root** rank’s tensor is copied to all other ranks. + +``` +Root: [w] ----copy----> all other ranks: [w] +``` + +**Where used:** **Weight initialization**, **loading checkpoints** to a group, distributing hyperparameters or small metadata. + +--- + +### Reduce + +**What it does:** Like AllReduce, but the **full result** appears on **one root** rank only. + +**Where used:** **Logging** or **metrics** where only rank 0 needs the scalar (e.g. reduced loss on one process). + +--- + +### Send / Recv (point-to-point) + +**What it does:** **One** rank sends a buffer to **one** other rank (possibly bidirectional with two ops). + +``` +Stage i ----Send/Recv----> Stage i+1 +``` + +**Where used:** **Pipeline parallelism** activations in forward, gradients in backward; **ring attention** steps in **context parallelism** (often implemented as a ring of Send/Recv with careful ordering). + +--- + +## 3. Which collectives are used in each parallelism strategy + +| Parallelism | Forward Pass | Backward Pass | Optimizer Step | +|---------------|--------------|----------------|----------------| +| Data Parallel | — | AllReduce (gradients) | — | +| FSDP/ZeRO-3 | AllGather (params) | ReduceScatter (grads) + AllGather (params) | — | +| Tensor Parallel | AllReduce or AllGather+ReduceScatter | AllReduce or AllGather+ReduceScatter | — | +| Sequence Parallel | AllGather (activations) | ReduceScatter (activations) | — | +| Pipeline Parallel | Send/Recv (activations) | Send/Recv (gradients) | — | +| Expert Parallel | AllToAll (token dispatch) | AllToAll (gradient dispatch) | — | +| Context Parallel | Ring Send/Recv (KV chunks) | Ring Send/Recv | — | + +Exact fusion and overlap depend on the backend (Megatron vs TorchTitan) and flags such as async TP or overlapped gradient reduction. + +--- + +## 4. Communication patterns in Megatron-LM + +Primus trains with **Megatron-LM** patches and configurations. Typical patterns: + +| Mode | Pattern | +|------|---------| +| **TP** | Column-parallel and row-parallel **linear** layers use **AllReduce** or **ReduceScatter/AllGather** sequences; with **sequence_parallel**, activations follow the Megatron **scatter/gather** pattern around TP regions. | +| **PP** | **Point-to-point** Send/Recv (or backend equivalents) between **pipeline stages** for activations and backward tensors. | +| **DP** | **AllReduce** for gradients when not using distributed optimizer; with **distributed optimizer**, **ReduceScatter**-style paths for shard-sized gradients. | +| **EP** | **AllToAll** for MoE **routing** (dispatch/combine) when experts are parallelized. | + +### Overlap knobs + +Megatron integrates **communication/compute overlap** options such as: + +- `overlap_grad_reduce`—overlap gradient reduction with computation where supported. +- `overlap_param_gather`—overlap parameter gathering (e.g. with distributed optimizer / FSDP-style paths) with computation. + +See [Megatron parameters](../03-configuration-reference/megatron-parameters.md) for defaults and compatibility with `use_distributed_optimizer`, `use_torch_fsdp2`, and checkpoint formats. + +--- + +## 5. Communication patterns in TorchTitan + +TorchTitan (used as a backend in Primus) relies on **PyTorch** distributed primitives and **DTensor**-style layouts: + +| Area | Pattern | +|------|---------| +| **FSDP / sharding** | **AllGather** / **ReduceScatter** orchestrated by **FSDP2** (`fully_shard` and related APIs) when `data_parallel_shard_degree` and sharding are enabled. | +| **TP** | Tensor parallelism is integrated with **DTensor** and model parallel helpers; schedules may use **collectives** inside module forward/backward. | +| **PP** | **Pipeline schedules** (`parallelism.pipeline_parallel_schedule`, `parallelism.pipeline_parallel_degree`) determine stage boundaries and buffering; communication is managed by the pipeline implementation. | + +### Async tensor parallelism + +Set `parallelism.enable_async_tensor_parallel: true` (where supported) to **overlap** TP communication with computation in eligible layers. + +--- + +## 6. RCCL-specific features and tuning + +The following appear in ROCm / AMD deployments and partner integrations; availability depends on your **driver**, **RCCL build**, and **network** stack. + +| Feature | Notes | +|---------|--------| +| **MSCCL** | Microsoft Collective Communication Library: **custom algorithms** and patterns; may be used when the stack is built and configured for them. | +| **MSCCL++** | User-space collective paths aimed at **lower latency** for specific patterns and hardware. | +| **ANP (AMD Network Plugin)** | Network backend integration (e.g. **AINIC**-oriented paths). Example: `NCCL_NET_PLUGIN` may point to `librccl-anp.so` or similar when installed (see Primus `examples/run_pretrain.sh` patterns). | + +### Environment variables + +Many deployments tune behavior with **NCCL-prefixed** variables (honored by RCCL for compatibility), for example: + +- `NCCL_PROTO`—protocol selection hints. +- `NCCL_P2P_NET_CHUNKSIZE`—chunking for P2P/network paths. +- `NCCL_IB_*`—InfiniBand / RDMA-related settings when applicable. +- `NCCL_SOCKET_IFNAME`—**socket** interface selection for TCP fallback or hybrid setups. + +Document your cluster’s recommended values in [Environment variables](../03-configuration-reference/environment-variables.md). + +--- + +## 7. Benchmarking collectives with Primus + +Primus includes an **RCCL microbenchmark** suite to measure **latency and bandwidth** for common collectives across message sizes. + +### Command + +Invoke through **`primus-cli`** (after `runner/` / container setup per your installation): + +```bash +./primus-cli direct -- benchmark rccl --op all_reduce --min-bytes 1M --max-bytes 128M +``` + +Useful flags (see `primus/tools/benchmark/rccl_bench_args.py`): + +| Flag | Purpose | +|------|---------| +| `--op` | One or more of: `all_reduce`, `broadcast`, `reduce_scatter`, `all_gather`, `alltoall` | +| `--min-bytes`, `--max-bytes` | Sweep range (e.g. `1K`, `1M`, `128M`) | +| `--num-sizes`, `--scale` | Generated sweep (`log2` or `linear`) | +| `--dtype` | `bf16`, `fp16`, `fp32` | +| `--output-file` | Write Markdown/CSV/JSONL report (default `./rccl_report.md`) | +| `--check` | Lightweight correctness checks | + +Example with multiple ops: + +```bash +./primus-cli direct -- benchmark rccl --op all_reduce all_gather reduce_scatter --min-bytes 1M --max-bytes 128M +``` + +### Reading results + +- **Bandwidth** (GB/s or similar): higher is better for large messages; compare against **peak NIC** or **GPU-GPU** limits for your topology. +- **Latency** (µs): dominates for **small** messages; important for **frequent small** collectives (e.g. some TP patterns). + +Use results to spot **unexpected drops** (wrong NIC, congestion, fallback to TCP) before scaling full training. + +--- + +## 8. Troubleshooting communication issues + +| Symptom | Checks | +|---------|--------| +| Hangs / timeouts | Enable **`NCCL_DEBUG=INFO`** (or `TRACE` for deep dives) and inspect which collective stalls. | +| Wrong interface | Set **`NCCL_SOCKET_IFNAME`** to the intended **cluster** interface; verify with `ip link` / admin docs. | +| IB / RDMA not used | Confirm **`NCCL_IB_*`**, HCA names, and permissions; run **preflight** (below). | +| Slow AllReduce | Compare **`benchmark rccl`** to baseline; check **topology** (NVLink vs network), **contention**, **message sizes**. | + +### Preflight: Network validation + +Primus **preflight** can aggregate host/GPU/network info: + +```bash +./primus-cli direct -- preflight --network +``` + +Combine with GPU checks as needed: + +```bash +./primus-cli direct -- preflight --gpu --network +``` + +Use this to confirm **RCCL/NCCL-related environment** snapshots and **connectivity expectations** before long jobs. + +--- + +## Related documentation + +- [Parallelism strategies](./parallelism-strategies.md)—how TP, PP, DP, FSDP, EP, and CP fit together. +- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) +- [Environment variables](../03-configuration-reference/environment-variables.md) diff --git a/docs/04-technical-guides/data-preparation.md b/docs/04-technical-guides/data-preparation.md new file mode 100644 index 000000000..75dca8f2a --- /dev/null +++ b/docs/04-technical-guides/data-preparation.md @@ -0,0 +1,159 @@ +# Data preparation guide + +Primus routes training through **Megatron-LM**, **TorchTitan**, and **MaxText**. Each backend expects its own data format and preprocessing pipeline. This guide summarizes how to prepare data, how to use **mock** data for smoke tests, and which environment variables commonly apply. + +Scripts referenced below live under the Primus repository root, for example: + +- `examples/megatron/preprocess_data.py` +- `examples/megatron/prepare.py` +- `examples/megatron/prepare_bookcorpus_megatron_dataset.py` +- `examples/torchtitan/prepare.py` + +--- + +## 1. Overview + +| Backend | Format | Typical entry | +|---------|--------|----------------| +| Megatron | Indexed `.bin` + `.idx` datasets | `data_path`, `train_data_path`, tokenizer args | +| TorchTitan | Hugging Face datasets + local assets | `training.dataset`, `training.dataset_path`, `model.hf_assets_path` | +| MaxText | TFDS / Hugging Face / Grain / synthetic | `dataset_type`, paths per pipeline | + +All backends support **synthetic or mock** data for configuration and scaling tests without large downloads. + +--- + +## 2. Mock data (testing) + +### Megatron + +Set in the trainer module: + +```yaml +mock_data: true +``` + +Default in `primus/configs/modules/megatron/trainer_base.yaml` is `false`. When `true`, training uses generated data matching configured dimensions so you can validate YAML, parallelism, and throughput without real corpora. + +### TorchTitan + +```yaml +training: + mock_data: true +``` + +Default in `primus/configs/modules/torchtitan/pre_trainer.yaml` is `true` (useful for quick runs; set `false` and supply real datasets for production). + +### MaxText + +Use `dataset_type: synthetic` (or other synthetic paths in MaxText configs). See `third_party/maxtext/src/MaxText/configs/base.yml` and model YAMLs under `third_party/maxtext/src/MaxText/configs/`. + +--- + +## 3. Megatron data pipeline + +### Inputs + +- Raw **JSON** or **JSONL** text (one JSON object per line for JSONL). +- Optional **sentence splitting** via NLTK when `--split-sentences` is used (requires NLTK data; see [Environment variables](#6-environment-variables-for-data)). + +### Preprocessing: `examples/megatron/preprocess_data.py` + +The script tokenizes input and writes **Megatron indexed datasets** (`.bin` + `.idx`). It uses `build_tokenizer` from Primus’s Megatron tokenizer integration and accepts tokenizer flags from `_add_tokenizer_args`. + +**Important arguments** (from the script’s argparse): + +| Argument | Description | +|----------|-------------| +| `--input` | Path to input JSON (required). | +| `--json-keys` | Keys to read (default `text`). | +| `--output-prefix` | Output path **without** suffix; produces `{prefix}_{key}_{document|sentence}.bin` and `.idx`. | +| `--workers` | Number of worker processes (required). | +| `--partitions` | Split input for parallel preprocessing (default `1`). | +| `--split-sentences` | Run NLTK sentence splitting before encode. | +| `--append-eod` | Append end-of-document token. | + +**Example** (mirrors `examples/megatron/prepare.py` for BookCorpus-style flows): + +```bash +python3 examples/megatron/preprocess_data.py \ + --input /path/to/train.json \ + --tokenizer-type HuggingFaceTokenizer \ + --tokenizer-model /path/to/tokenizer \ + --output-prefix /path/to/out/bookcorpus_train \ + --workers "$(nproc)" \ + --split-sentences \ + --partitions 2 +``` + +### Configuring training runs + +| Parameter | Notes | +|-----------|--------| +| `data_path` | Single path or **weighted blend**: `0.5 /path/a 0.5 /path/b` | +| `train_data_path`, `valid_data_path`, `test_data_path` | Separate splits when used | +| `split` | Train/valid/test ratio string, e.g. `"99,1,0"` (default in `trainer_base.yaml`) or `"98,2,0"` for train/valid/test | +| `dataloader_type` | Megatron dataloader type; default in `trainer_base.yaml` is `null` (set explicitly in experiments as needed) | + +### BookCorpus example scripts + +- **`examples/megatron/prepare_bookcorpus_megatron_dataset.py`**—downloads BookCorpus to JSON via Hugging Face `datasets`, optional `--out-dir`. +- **`examples/megatron/prepare.py`**—orchestrates download, train/valid split, and calls `preprocess_data.py` with tokenizer settings from Primus config; respects `TOKENIZED_TRAIN_DATA_PATH` / `TOKENIZED_EVAL_DATA_PATH` for output locations. + +### Tokenizers + +Tokenizer type and model path are set on the model preset (for example `tokenizer_type`, `tokenizer_model` in `primus/configs/models/megatron/language_model.yaml` comments list `Llama2Tokenizer`, `HuggingFaceTokenizer`, etc.). + +--- + +## 4. TorchTitan data pipeline + +TorchTitan uses **Hugging Face datasets** style identifiers and local paths. + +| Key | Default (`pre_trainer.yaml`) | Description | +|-----|------------------------------|-------------| +| `training.dataset` | `c4` | Dataset identifier for TorchTitan loaders. | +| `training.dataset_path` | `null` | Local directory for dataset assets when needed. | + +Tokenizer and model assets are resolved from **`model.hf_assets_path`** (or equivalent in your model preset). The preparation script **`examples/torchtitan/prepare.py`**: + +- Resolves the TorchTitan checkout path. +- Runs `scripts/download_hf_assets.py` inside TorchTitan to fetch tokenizer assets for a given `repo_id`. +- Uses `HF_TOKEN` when the model or dataset is gated. + +--- + +## 5. MaxText data pipeline + +MaxText configuration is defined in upstream YAML (for example `third_party/maxtext/src/MaxText/configs/base.yml`). + +| Parameter | Meaning | +|-----------|---------| +| `dataset_type` | One of `synthetic`, `hf`, `grain`, `tfds` (per `base.yml` comments). | +| `hf_path`, `hf_data_dir`, `hf_train_files` | Hugging Face pipeline inputs when `dataset_type: hf`. | +| `per_device_batch_size` | Batch sizing on each device. | +| `packing` | Sequence packing for efficiency (default `True` in `base.yml`). | + +See MaxText’s data input documentation for Grain and TFDS specifics. + +--- + +## 6. Environment variables for data + +| Variable | Usage | +|----------|--------| +| `TOKENIZED_DATA_PATH` / `PRIMUS_TOKENIZED_DATA_PATH` | Tokenized dataset locations for Megatron hooks and examples (see `docs/03-configuration-reference/environment-variables.md`). | +| `TOKENIZED_TRAIN_DATA_PATH`, `TOKENIZED_EVAL_DATA_PATH` | Override output paths in `examples/megatron/prepare.py`. | +| `DATA_PATH` | General data root used in scripts and CI-style launches. | +| `HF_TOKEN` | **Required** for gated Hugging Face models and some datasets (TorchTitan `prepare.py`, Kubernetes examples in `examples/README.md`). | +| `HF_HOME` | Hugging Face cache directory (used in `examples/megatron/prepare.py`). | +| `NLTK_DATA` | NLTK tokenizer data directory for sentence splitting in `preprocess_data.py` when `NLTK_DATA` is set. | + +--- + +## Summary + +1. Use **mock** / **synthetic** data to validate configs and performance before investing in large preprocessing jobs. +2. For **Megatron**, convert JSON/JSONL to `.bin`/`.idx` with `preprocess_data.py` and point `data_path` or split paths at the outputs. +3. For **TorchTitan**, set `training.dataset` / `dataset_path` and run **`examples/torchtitan/prepare.py`** to fetch tokenizer assets. +4. For **MaxText**, configure `dataset_type` and `per_device_batch_size` per upstream `base.yml` and model YAMLs. diff --git a/docs/04-technical-guides/determinism-and-reproducibility.md b/docs/04-technical-guides/determinism-and-reproducibility.md new file mode 100644 index 000000000..fc903bde2 --- /dev/null +++ b/docs/04-technical-guides/determinism-and-reproducibility.md @@ -0,0 +1,119 @@ +# Determinism and reproducibility + +Reproducibility—getting bit-identical (or run-to-run stable) results—matters for debugging divergence, validating optimizations, and regression testing. This guide covers Primus's deterministic mode, the environment variables it sets, the per-backend seed/determinism knobs, and the performance trade-offs. Parameters and behavior are grounded in `examples/run_pretrain.sh`, `primus/configs/modules/megatron/trainer_base.yaml`, and `primus/configs/modules/torchtitan/pre_trainer.yaml`. + +--- + +## 1. What "deterministic" means here + +There are two distinct goals: + +- **Reproducible (seeded)**—same seed + same config + same hardware/software gives the *same trajectory*. Achieved with fixed seeds; cheap. +- **Bitwise-deterministic**—kernels avoid non-deterministic reductions/atomics and tuning so results don't vary between runs. Requires deterministic algorithms and disabling autotuning; **slower**. + +Full determinism also generally requires the **same world size, parallelism layout, and library versions**. Changing TP/PP/DP, GPU count, or ROCm/Megatron versions can change numerics even with everything else fixed. + +--- + +## 2. Primus deterministic mode (`PRIMUS_DETERMINISTIC`) + +Setting `PRIMUS_DETERMINISTIC=1` configures the GPU/communication stack for deterministic behavior. The CLI/runner path applies this through the hook `runner/helpers/hooks/05_deterministic.sh`; the `examples/run_pretrain.sh` script applies an equivalent inline block. The exported variables are: + +```bash +# when PRIMUS_DETERMINISTIC=1 (runner/helpers/hooks/05_deterministic.sh) +export NCCL_ALGO="Ring" # deterministic collective algorithm +export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 # Transformer Engine: forbid non-deterministic kernels +export ROCBLAS_DEFAULT_ATOMICS_MODE=0 # rocBLAS: disable atomic (non-deterministic) reductions +export TORCH_COMPILE_DISABLE=1 # avoid torch.compile/Triton race conditions +export PRIMUS_TURBO_AUTO_TUNE=0 # disable Primus-Turbo autotuning (stable kernel choice) +``` + +> `PRIMUS_TURBO_AUTO_TUNE` also defaults to `0` in `runner/helpers/envs/base_env.sh`. The inline block in `examples/run_pretrain.sh` sets the first four variables and relies on that default for the fifth. + +Additionally, **HipBLASLt autotuning is disabled** in deterministic mode: tuning only runs when `PRIMUS_DETERMINISTIC != 1` *and* `PRIMUS_HIPBLASLT_TUNING=1` (`examples/run_pretrain.sh`). This prevents run-to-run kernel-selection differences. See [Performance tuning](./performance-tuning.md). + +`PRIMUS_DETERMINISTIC` is on the container passthrough whitelist (`runner/.primus.yaml`), so it reaches the training container. See [Environment variables](../03-configuration-reference/environment-variables.md). + +```bash +export PRIMUS_DETERMINISTIC=1 +./runner/primus-cli direct -- train pretrain \ + --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +``` + +> The MoE example scripts explicitly set `PRIMUS_DETERMINISTIC=0` because deterministic mode disables the performance kernels/tuning they rely on. + +--- + +## 3. Seeds and deterministic algorithms (Megatron) + +In `primus/configs/modules/megatron/trainer_base.yaml`: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `seed` | `1234` | Master RNG seed (Python/NumPy/Torch, data order, init). | +| `deterministic_mode` | `false` | Force deterministic kernels/algorithms inside Megatron (slower; pairs with `PRIMUS_DETERMINISTIC`). | +| `data_parallel_random_init` | `false` | When `false`, parameters are initialized identically and broadcast across DP ranks; keep `false` for reproducible init. | + +For a fully reproducible Megatron run: set a fixed `seed`, `deterministic_mode: true`, and launch with `PRIMUS_DETERMINISTIC=1`. + +> **Startup assertion.** When `deterministic_mode: true`, Primus validates (`primus/backends/megatron/patches/args/rocm_arg_validation.py`, `validate_args_on_rocm`) that these environment variables are set, and **fails fast** otherwise: `TORCH_COMPILE_DISABLE=1`, `ROCBLAS_DEFAULT_ATOMICS_MODE=0`, `PRIMUS_TURBO_AUTO_TUNE=0`, and `PRIMUS_DETERMINISTIC=1`. Launching with `PRIMUS_DETERMINISTIC=1` (above) sets all of them, so always pair `deterministic_mode: true` with `PRIMUS_DETERMINISTIC=1`. + +--- + +## 4. Seeds and determinism (TorchTitan) + +Under `training:` in `primus/configs/modules/torchtitan/pre_trainer.yaml`: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `seed` | `null` | RNG seed; set an integer for reproducible runs. | +| `deterministic` | `false` | Enable deterministic algorithms (disables some optimized kernels; slower). | + +Related: `checkpoint.create_seed_checkpoint` (`false`) creates a deterministic seed checkpoint that all ranks load, ensuring identical initialization across a distributed run. + +Note `compile.enable: true` is the TorchTitan default; for strict determinism prefer launching with `PRIMUS_DETERMINISTIC=1` (which sets `TORCH_COMPILE_DISABLE=1`) or disable compilation. + +--- + +## 5. MaxText + +MaxText determinism is governed by the upstream MaxText seed/data options surfaced through the MaxText config (see [MaxText parameters](../03-configuration-reference/maxtext-parameters.md)). The GPU-stack environment effects of `PRIMUS_DETERMINISTIC` (rocBLAS atomics, deterministic collectives) still apply at the launcher level. + +--- + +## 6. Performance trade-offs + +Determinism is not free: + +| Setting | Cost | +|---------|------| +| `NCCL_ALGO=Ring` | Forgoes faster topology-aware collective algorithms. | +| `ROCBLAS_DEFAULT_ATOMICS_MODE=0` | Disables atomic reductions—slower GEMMs. | +| `NVTE_ALLOW_NONDETERMINISTIC_ALGO=0` | Restricts TE to deterministic (often slower) kernels. | +| `TORCH_COMPILE_DISABLE=1` | No `torch.compile` fusion/codegen speedups. | +| `PRIMUS_TURBO_AUTO_TUNE=0` | No Primus-Turbo kernel autotuning. | +| HipBLASLt tuning disabled | No autotuned GEMM kernels. | +| `deterministic_mode` / `deterministic` | Deterministic algorithm variants are generally slower. | + +**Use deterministic mode for debugging and validation, not production throughput runs.** Once a result is reproduced/diagnosed, disable it to recover performance. + +--- + +## 7. Reproducibility checklist + +1. **Pin the environment**—same container image, ROCm version, and backend (Megatron/TorchTitan) commit. +2. **Fix seeds**—Megatron `seed`; TorchTitan `training.seed`. +3. **Hold the layout constant**—same world size and TP/PP/DP/EP/CP degrees. +4. **Enable determinism**—`PRIMUS_DETERMINISTIC=1` plus backend `deterministic_mode`/`deterministic`. +5. **Disable autotuning**—automatic in deterministic mode (HipBLASLt tuning off). +6. **Use mock or fixed data ordering**—ensure the data pipeline is seeded; see [Data preparation](./data-preparation.md). +7. **Record everything**—log the full resolved config and env (see [Logging & experiment tracking](./logging-and-experiment-tracking.md)). + +--- + +## Related documentation + +- [Performance tuning](./performance-tuning.md)—HipBLASLt tuning and its interaction with deterministic mode. +- [Environment variables](../03-configuration-reference/environment-variables.md)—`PRIMUS_DETERMINISTIC` and related flags. +- [Data preparation](./data-preparation.md)—deterministic data ordering. +- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) and [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md). diff --git a/docs/backends/megatron/diffusion/README.md b/docs/04-technical-guides/diffusion-models/README.md similarity index 91% rename from docs/backends/megatron/diffusion/README.md rename to docs/04-technical-guides/diffusion-models/README.md index 606a25b3a..26daa6cb3 100644 --- a/docs/backends/megatron/diffusion/README.md +++ b/docs/04-technical-guides/diffusion-models/README.md @@ -1,14 +1,14 @@ -# Diffusion Models in Primus - Developer & Architecture Guide +# Diffusion models in Primus - developer and architecture guide **Purpose:** Developer-focused documentation for understanding Primus diffusion architecture, design decisions, and implementation details. -**For training/usage instructions, see:** [examples/megatron/diffusion/README.md](../../../../examples/megatron/diffusion/README.md) +**For training/usage instructions, see:** [examples/megatron/diffusion/README.md](../../../examples/megatron/diffusion/README.md) -**For test documentation, see:** [tests/unit_tests/backends/megatron/diffusion/](../../../../tests/unit_tests/backends/megatron/diffusion/) +**For test documentation, see:** [tests/unit_tests/backends/megatron/diffusion/](../../../tests/unit_tests/backends/megatron/diffusion/) --- -## Architecture Philosophy +## Architecture philosophy Primus diffusion models are built as **Megatron-Core native implementations**, designed for: - Production-scale distributed training @@ -16,7 +16,7 @@ Primus diffusion models are built as **Megatron-Core native implementations**, d - Advanced checkpoint management with heterogeneous layers - Clean separation of concerns (no framework dependencies like PyTorch Lightning) -### Key Design Decisions +### Key design decisions **1. Megatron-Core Integration** - Models in `core/models/diffusion/` follow Megatron-Core patterns @@ -43,9 +43,9 @@ Primus diffusion models are built as **Megatron-Core native implementations**, d --- -## Supported Models +## Supported models -### Flux ✅ Production Ready +### Flux ✅ production ready Flow-based diffusion model with MMDiT (Multimodal Diffusion Transformer) architecture. - **Architecture**: Dual-stream with joint and single transformer blocks @@ -53,14 +53,14 @@ Flow-based diffusion model with MMDiT (Multimodal Diffusion Transformer) archite - **Reference**: [Black Forest Labs FLUX.1](https://huggingface.co/black-forest-labs/FLUX.1-dev) - **Status**: Fully implemented and tested (390 tests) -### Future Models ⏳ Planned +### Future models ⏳ planned - **DiT**: Diffusion Transformer for image generation - **MovieGen**: Video diffusion models - **Custom Models**: Extensible framework for new architectures --- -## Project Structure +## Project structure ``` primus/backends/megatron/ @@ -115,7 +115,7 @@ tests/unit_tests/backends/megatron/diffusion/ # Comprehensive test suite (390 t ├── functional/ # End-to-end functional tests └── checkpointing/ # Checkpoint tests -docs/backends/megatron/diffusion/ # This directory +docs/04-technical-guides/diffusion-models/ # This directory ├── README.md # This file (developer guide) ├── architecture_overview.md # Detailed architecture ├── data_preprocessing.md # Data pipeline guide (includes Flux-specific section) @@ -129,9 +129,9 @@ docs/backends/megatron/diffusion/ # This directory --- -## Key Technical Features +## Key technical features -### 1. DiffusionModule Base Class +### 1. DiffusionModule base class All diffusion models inherit from `DiffusionModule`, which provides: - Megatron-Core integration (process groups, parallelism) @@ -151,7 +151,7 @@ Configuration class extending `TransformerConfig`: **Location:** `primus/backends/megatron/core/models/diffusion/common/config.py` -### 3. Hierarchical Encoder Registry +### 3. Hierarchical encoder registry Organized by modality → type → variant: ``` @@ -173,7 +173,7 @@ Benefits: - Lazy loading (encoders loaded only when needed) - Shared base classes for common functionality -### 4. Training Utilities Structure +### 4. Training utilities structure **Noise Application** (`noise_utils.py`): - `apply_flow_matching_noise()`: For flow matching models (Flux) @@ -192,7 +192,7 @@ Benefits: - `ModeSampler`: Mode-focused sampling - Base class for custom samplers -### 5. Shared Energon Infrastructure +### 5. Shared Energon infrastructure Located in `data/energon/` for reusability across models: - Shared data loading utilities @@ -200,7 +200,7 @@ Located in `data/energon/` for reusability across models: - WebDataset integration - Model-specific TaskEncoders in `data/diffusion/task_encoders/` -### 6. Precalculated Data Support +### 6. Precalculated data support **Performance**: 5-10x faster training than on-the-fly encoding @@ -219,7 +219,7 @@ Located in `data/energon/` for reusability across models: - Lower GPU memory (no encoders loaded during training) - Better reproducibility -### 7. MLPerf Streaming Ingest Pipeline +### 7. MLPerf streaming ingest pipeline **Location:** `data/diffusion/preprocessing/pipelines/ingest.py` @@ -245,9 +245,9 @@ The `StreamingIngestPipeline` downloads Apache Arrow IPC files from MLCommons R2 --- -## Implementation Status +## Implementation status -### Core Infrastructure ✅ +### Core infrastructure ✅ - ✅ Directory structure with 25+ directories - ✅ Base classes (DiffusionModule, BaseDiffusionConfig, BaseScheduler) - ✅ DiffusionModule with Megatron-Core integration @@ -257,7 +257,7 @@ The `StreamingIngestPipeline` downloads Apache Arrow IPC files from MLCommons R2 - ✅ Testing framework (390 tests) - ✅ Comprehensive documentation -### Flux Model Implementation ✅ +### Flux model implementation ✅ - ✅ Flux model architecture (dual-stream MMDiT) - ✅ MMDiT layers and attention (joint + single blocks) - ✅ Embeddings (3D RoPE, timestep, vector) @@ -268,9 +268,9 @@ The `StreamingIngestPipeline` downloads Apache Arrow IPC files from MLCommons R2 --- -## Documentation Map +## Documentation map -### Core Guides +### Core guides 📖 **[Architecture Overview](architecture_overview.md)** High-level design, directory structure, and architectural decisions. @@ -287,7 +287,7 @@ Megatron-Energon patterns and TaskEncoder implementation. 📖 **[Adding New Models](adding_new_models.md)** Step-by-step guide for implementing new diffusion models. -### Advanced Documentation +### Advanced documentation 📖 **[Flux Architecture Deep Dive](flux_architecture.md)** Mathematical formulation, detailed component descriptions, and performance optimizations. @@ -298,17 +298,17 @@ Complete API documentation with function signatures and usage examples. 📖 **[FP8 Training Guide](fp8_training.md)** FP8 precision training on AMD MI300X: configuration, benchmarks, tuning recipes, and troubleshooting. -### Related Documentation +### Related documentation -📖 **[Training Guide](../../../../examples/megatron/diffusion/README.md)** +📖 **[Training Guide](../../../examples/megatron/diffusion/README.md)** User-facing guide for training Flux models (quick start, configurations, troubleshooting). -📖 **[Test Directory](../../../../tests/unit_tests/backends/megatron/diffusion/)** +📖 **[Test Directory](../../../tests/unit_tests/backends/megatron/diffusion/)** Test suite for diffusion models. --- -## Testing Architecture +## Testing architecture **Test Organization** (following Megatron-LM patterns): - One comprehensive file per model (`test_flux_model.py`) @@ -318,18 +318,18 @@ Test suite for diffusion models. **Test Status**: ✅ 390 tests passing -See [tests/unit_tests/backends/megatron/diffusion/](../../../../tests/unit_tests/backends/megatron/diffusion/) for details. +See [tests/unit_tests/backends/megatron/diffusion/](../../../tests/unit_tests/backends/megatron/diffusion/) for details. --- -## Hardware Requirements +## Hardware requirements -### Flux 535M (Testing) +### Flux 535M (testing) - **Training**: 1x MI300X 192GB (compatible with H100/A100) - **Inference**: 1x MI300X 192GB - **Batch Size**: 1-8 per GPU -### Flux 12B (Production) +### Flux 12B (production) - **Training**: 8x MI300X 192GB (recommended) or 4x MI300X 192GB with TP=2 - **Inference**: 1x MI300X 192GB - **Batch Size**: 1-2 per GPU for training, 1-4 for inference diff --git a/docs/backends/megatron/diffusion/STRUCTURE.md b/docs/04-technical-guides/diffusion-models/STRUCTURE.md similarity index 89% rename from docs/backends/megatron/diffusion/STRUCTURE.md rename to docs/04-technical-guides/diffusion-models/STRUCTURE.md index 382ab7eb4..9bed81156 100644 --- a/docs/backends/megatron/diffusion/STRUCTURE.md +++ b/docs/04-technical-guides/diffusion-models/STRUCTURE.md @@ -1,4 +1,4 @@ -# Flux Diffusion Infrastructure - Directory Structure +# Flux diffusion infrastructure - directory structure **Created**: December 5, 2025 **Status**: ✓ Implementation Complete @@ -9,7 +9,7 @@ This document describes the directory structure created for Flux diffusion model --- -## Directory Tree +## Directory tree ``` Primus/ @@ -18,7 +18,7 @@ Primus/ │ │ ├── common/diffusion_module/ # DiffusionModule base class │ │ │ └── diffusion_module.py │ │ └── diffusion/ # Diffusion models (Megatron-Core convention) -│ │ ├── common/ # Shared components (MMDiT layers, attention) +│ │ ├── common/ # Shared building blocks (config, embeddings, normalization) │ │ │ ├── __init__.py │ │ │ ├── config.py # ✓ BaseDiffusionConfig │ │ │ ├── embeddings.py # ✓ TimeStepEmbedder, MLPEmbedder @@ -31,7 +31,6 @@ Primus/ │ │ │ ├── layer_spec.py # ✓ get_flux_layer_spec, get_flux_*_spec_for_backend, MMDiTLayer │ │ │ ├── attention.py # ✓ JointSelfAttention, FluxSingleAttention │ │ │ ├── utils.py # ✓ generate_image_position_ids -│ │ │ ├── checkpoint_utils.py # ✓ Checkpoint utilities │ │ │ └── checkpoint_converter.py # ✓ HF <-> Megatron conversion │ │ └── __init__.py │ │ @@ -74,9 +73,8 @@ Primus/ │ │ │ └── __init__.py │ -├── primus/modules/trainer/megatron/ -│ └── diffusion/ # Diffusion trainer -│ └── __init__.py # ✓ DiffusionTrainer +├── primus/backends/megatron/ +│ └── megatron_pretrain_trainer.py # ✓ Shared Megatron pretrain trainer (drives diffusion pretraining) │ ├── primus/configs/models/megatron/ │ └── diffusion/ # YAML configs @@ -135,9 +133,9 @@ Primus/ --- -## Completed Components +## Completed components -### ✓ Base Classes +### ✓ Base classes 1. **DiffusionModule** (`core/models/common/diffusion_module/diffusion_module.py`) - Base class for all diffusion models (extends MegatronModule) @@ -167,7 +165,7 @@ Primus/ - Linear interpolation: `x_t = (1-t)*noise + t*data` - Velocity target: `v = data - noise` -### ✓ Directory Structure +### ✓ Directory structure - **25 `__init__.py` files** with comprehensive docstrings - **Multiple implementation files** (models, configs, schedulers, data pipeline) @@ -175,39 +173,38 @@ Primus/ --- -## Architectural Decisions +## Architectural decisions ### 1. Models under `core/models/` - Follows Megatron-Core convention (`megatron/core/models/gpt/`, etc.) - Easier upstream tracking when Megatron-Core adds diffusion support -### 2. Shared Components in `common/` +### 2. Shared components in `common/` - Standard approach stores shared code in model-specific directories -- Primus: `common/` for MMDiT layers, attention, shared utilities -- Flux-specific: Only `EmbedND` and Flux model class +- Primus: `common/` for shared config, embeddings, and normalization +- Flux-specific: model class, MMDiT/single-block layer specs, joint attention, and `EmbedND` -### 3. Hierarchical Encoder Structure +### 3. Hierarchical encoder structure - `encoders/image/vae/`, `encoders/text/t5/`, `encoders/text/clip/` - Registry pattern for config-driven selection - Easy to add new encoder variants (5+ planned per modality) -### 4. Shared Energon Infrastructure +### 4. Shared Energon infrastructure - `data/energon/` for cross-model utilities (VLM, diffusion, future) - `data/diffusion/task_encoders/` for diffusion-specific TaskEncoders - Traditional approach nests Energon under model-specific directories -### 5. Mock Data in Tests -- `tests/fixtures/diffusion/` (not production code) -- Traditional approach mixes test utilities with production code -- Follows pytest best practices +### 5. Synthetic (mock) data +- Synthetic datasets live in `primus/backends/megatron/data/synthetic/mock_datasets.py`, wired through `primus/backends/megatron/data/synthetic_dataset_provider.py`, so training can run without real data +- Unit tests exercise them under `tests/unit_tests/backends/megatron/diffusion/data/` -### 6. No PyTorch Lightning +### 6. No PyTorch lightning - Pure Megatron patterns (no PTL DataModules) - Better integration with Megatron training loop --- -## Import Examples +## Import examples ```python # Base classes @@ -235,7 +232,7 @@ timesteps = scheduler.sample_timesteps(batch_size=8, device='cuda') --- -## Validation Status +## Validation status ✓ All Python files syntactically correct ✓ No linter errors detected @@ -245,7 +242,7 @@ timesteps = scheduler.sample_timesteps(batch_size=8, device='cuda') --- -## Files Summary +## Files summary All infrastructure files, model implementations, data pipeline components, tests, and documentation are complete and ready for production use. diff --git a/docs/backends/megatron/diffusion/adding_new_models.md b/docs/04-technical-guides/diffusion-models/adding_new_models.md similarity index 95% rename from docs/backends/megatron/diffusion/adding_new_models.md rename to docs/04-technical-guides/diffusion-models/adding_new_models.md index a29ed8fd0..e975bae64 100644 --- a/docs/backends/megatron/diffusion/adding_new_models.md +++ b/docs/04-technical-guides/diffusion-models/adding_new_models.md @@ -1,17 +1,17 @@ -# Adding New Diffusion Models +# Adding new diffusion models This guide explains how to add new diffusion models to Primus, following the established patterns and architecture. --- -## Table of Contents +## Table of contents 1. [Overview](#overview) 2. [Prerequisites](#prerequisites) -3. [Step-by-Step Guide](#step-by-step-guide) +3. [Step-by-step guide](#step-by-step-guide) 4. [Example: Adding DiT](#example-adding-dit) -5. [Testing Your Model](#testing-your-model) -6. [Best Practices](#best-practices) +5. [Testing your model](#testing-your-model) +6. [Best practices](#best-practices) --- @@ -45,9 +45,9 @@ Before adding a new model, ensure you have: --- -## Step-by-Step Guide +## Step-by-step guide -### Step 1: Create Model Directory +### Step 1: Create model directory Create a directory for your model under `core/models/diffusion/`: @@ -64,7 +64,7 @@ touch model.py touch layers.py # If model-specific layers needed ``` -### Step 2: Implement Configuration +### Step 2: Implement configuration **File**: `config.py` @@ -161,7 +161,7 @@ class DiTConfig(BaseDiffusionConfig): return cls(**defaults) ``` -### Step 3: Implement Model Class +### Step 3: Implement model class **File**: `model.py` @@ -340,11 +340,11 @@ class DiTBlock(nn.Module): pass ``` -### Step 4: Add Model-Specific Layers (If Needed) +### Step 4: Add model-specific layers (if needed) If your model has unique layers not shared with other models, add them to `layers.py`. -### Step 5: Export Model +### Step 5: Export model **File**: `__init__.py` @@ -363,7 +363,7 @@ __all__ = [ ] ``` -### Step 6: Create Configuration Files +### Step 6: Create configuration files **File**: `primus/configs/models/megatron/diffusion/dit_xl_2.yaml` @@ -407,7 +407,7 @@ global_batch_size: 256 learning_rate: 1.0e-4 ``` -### Step 7: Write Tests +### Step 7: Write tests **File**: `tests/unit_tests/backends/megatron/diffusion/test_dit_model.py` @@ -492,13 +492,13 @@ if __name__ == "__main__": pytest.main([__file__, "-v"]) ``` -### Step 8: Update Documentation +### Step 8: Update documentation 1. Add model to `README.md` supported models list 2. Update `architecture_overview.md` with model-specific details 3. Create model-specific training guide (e.g., `dit_training.md`) -### Step 9: Add Example Scripts +### Step 9: Add example scripts **File**: Use `examples/run_pretrain.sh` with appropriate config @@ -544,9 +544,9 @@ See the complete example in the step-by-step guide above. --- -## Testing Your Model +## Testing your model -### Unit Tests +### Unit tests Run tests to verify implementation: @@ -558,7 +558,7 @@ pytest tests/unit_tests/backends/megatron/diffusion/test_dit_model.py -v pytest tests/unit_tests/backends/megatron/diffusion/test_dit_model.py::TestDiTModel::test_dit_forward_shapes -v ``` -### Integration Tests +### Integration tests Test with actual data: @@ -576,9 +576,9 @@ Compare with reference implementation: --- -## Best Practices +## Best practices -### 1. Code Organization +### 1. Code organization - ✅ Separate configuration from model code - ✅ Reuse shared components from `common/` - ✅ Keep model-specific code minimal @@ -590,7 +590,7 @@ Compare with reference implementation: - ✅ Implement validation - ✅ Document all parameters -### 3. Model Implementation +### 3. Model implementation - ✅ Extend `DiffusionModule` from `primus.backends.megatron.core.models.common.diffusion_module.diffusion_module` - ✅ Implement required method: `forward()` - ✅ Use standalone loss functions from `loss_computation.py` @@ -619,9 +619,9 @@ Compare with reference implementation: --- -## Common Pitfalls +## Common pitfalls -### 1. Import Errors +### 1. Import errors ❌ **Wrong**: Absolute imports ```python from primus.backends.megatron.core.models.common.diffusion_module.diffusion_module import DiffusionModule @@ -632,7 +632,7 @@ from primus.backends.megatron.core.models.common.diffusion_module.diffusion_modu from ...common.diffusion_module.diffusion_module import DiffusionModule ``` -### 2. Configuration Validation +### 2. Configuration validation ❌ **Wrong**: No validation ```python class DiTConfig(BaseDiffusionConfig): @@ -647,7 +647,7 @@ def validate(self): raise ValueError(f"num_layers must be positive") ``` -### 3. Shape Mismatches +### 3. Shape mismatches ❌ **Wrong**: Assuming fixed shapes ```python def forward(self, x): @@ -698,11 +698,11 @@ Before submitting your new model: --- -## Getting Help +## Getting help If you encounter issues: 1. Review existing models (Flux) for patterns -2. Check documentation in `docs/backends/megatron/diffusion/` +2. Check documentation in `docs/04-technical-guides/diffusion-models/` 3. Run tests in debug mode: `pytest --pdb` 4. Consult architecture overview for design principles diff --git a/docs/backends/megatron/diffusion/api_reference.md b/docs/04-technical-guides/diffusion-models/api_reference.md similarity index 97% rename from docs/backends/megatron/diffusion/api_reference.md rename to docs/04-technical-guides/diffusion-models/api_reference.md index f2f84651a..c08525eb2 100644 --- a/docs/backends/megatron/diffusion/api_reference.md +++ b/docs/04-technical-guides/diffusion-models/api_reference.md @@ -1,4 +1,4 @@ -# Flux Model API Reference +# Flux model API reference ## Overview @@ -6,7 +6,7 @@ This document provides comprehensive API reference for the Flux diffusion model --- -## Base Classes +## Base classes ### DiffusionModule @@ -37,9 +37,9 @@ from primus.backends.megatron.core.models.common.diffusion_module import Diffusi --- -## Model Architecture +## Model architecture -### Flux Class +### Flux class **Location**: `primus/backends/megatron/core/models/diffusion/flux/model.py` @@ -52,7 +52,7 @@ config = FluxConfig.flux_535m() model = Flux(config) ``` -#### Architecture Diagram +#### Architecture diagram ``` ┌─────────────────────────────────────────────────────────────────┐ @@ -134,7 +134,7 @@ model = Flux(config) └─────────────────────────────────────────────────────────────────┘ ``` -#### Parameter Counts +#### Parameter counts | Variant | Joint Layers | Single Layers | Total Parameters | Use Case | |---------|-------------|---------------|------------------|----------| @@ -151,7 +151,7 @@ model = Flux(config) Complete configuration class for Flux models, inheriting from `BaseDiffusionConfig`. -#### Key Parameters +#### Key parameters | Parameter | Type | Default | Description | |-----------|------|---------|-------------| @@ -174,7 +174,7 @@ Complete configuration class for Flux models, inheriting from `BaseDiffusionConf | `hidden_dropout` | float | 0.0 | Hidden layer dropout rate | | `attention_dropout` | float | 0.0 | Attention dropout rate | -#### Configuration Examples +#### Configuration examples **Flux 535M (Testing)**: ```python @@ -254,7 +254,7 @@ embedded = embedder(clip_pooled) # [B, 3072] ### Normalization -#### AdaLN (Adaptive Layer Normalization) +#### AdaLN (adaptive layer normalization) **Location**: `primus/backends/megatron/core/models/diffusion/common/normalization.py` @@ -304,7 +304,7 @@ x_norm = norm(x) --- -### Position Embeddings +### Position embeddings #### EmbedND (3D RoPE) @@ -340,7 +340,7 @@ rope_freqs = embed_nd(img_ids) # [3, B, H*W, 3072] --- -### Attention Mechanisms +### Attention mechanisms #### JointSelfAttention @@ -388,7 +388,7 @@ output = single_attn(img_tokens, attention_mask=None) --- -### Layer Specifications +### Layer specifications #### MMDiTLayer @@ -442,9 +442,9 @@ output, _ = single_block(img_tokens, emb=emb) --- -## Training Utilities +## Training utilities -### Noise Application +### Noise application **Location**: `primus/backends/megatron/training/diffusion/noise_utils.py` @@ -496,7 +496,7 @@ noisy = apply_ddpm_noise(clean, noise, alpha_bar) --- -### Loss Computation +### Loss computation **Location**: `primus/backends/megatron/training/diffusion/loss_computation.py` @@ -572,7 +572,7 @@ loss = compute_v_prediction_loss(prediction, clean, noise, sigma) --- -### Timestep Sampling +### Timestep sampling **Location**: `primus/backends/megatron/training/diffusion/timestep_sampling.py` @@ -651,7 +651,7 @@ sampler = create_timestep_sampler("mode", mode_scale=1.5) --- -### Training Workflow Example +### Training workflow example Complete training step using the utilities: @@ -724,9 +724,9 @@ for batch in dataloader: --- -## Usage Examples +## Usage examples -### Basic Inference +### Basic inference ```python import torch @@ -764,7 +764,7 @@ with torch.no_grad(): print(f"Output shape: {predicted_velocity.shape}") # [1, 64, 128, 128] ``` -### Training Step +### Training step ```python import torch @@ -816,7 +816,7 @@ optimizer.step() print(f"Loss: {loss.item():.4f}") ``` -### With Guidance (Classifier-Free Guidance) +### With guidance (classifier-free guidance) ```python # Enable guidance in config @@ -839,7 +839,7 @@ with torch.no_grad(): ) ``` -### Different Resolutions +### Different resolutions ```python # Flux can handle different resolutions @@ -861,7 +861,7 @@ for height, width in resolutions: --- -## Methods Reference +## Methods reference ### Flux.forward() @@ -880,7 +880,7 @@ def forward( ) -> Tensor: # Returns: [B, C, H, W] Predicted velocity ``` -## Loss Computation +## Loss computation Loss is computed using standalone functions from `loss_computation.py`: @@ -909,7 +909,7 @@ def get_num_params( ## Testing -### Running Tests +### Running tests ```bash # Run all Flux tests @@ -927,7 +927,7 @@ pytest tests/unit_tests/backends/megatron/diffusion/test_flux_layer_spec_backend pytest tests/unit_tests/backends/megatron/diffusion/ --cov=primus.backends.megatron.core.models.diffusion --cov-report=html ``` -### Test Coverage +### Test coverage - **Component tests**: 80+ tests covering all components - **Model tests**: 17+ tests for full model @@ -935,22 +935,22 @@ pytest tests/unit_tests/backends/megatron/diffusion/ --cov=primus.backends.megat --- -## Performance Considerations +## Performance considerations -### Memory Usage +### Memory usage | Configuration | Model Weights | Training (bf16) | Training (fp32) | |--------------|---------------|-----------------|-----------------| | Flux 535M | ~2 GB | ~6-8 GB | ~10-12 GB | | Flux 12B | ~24 GB | ~60-80 GB | ~100-120 GB | -### Throughput (Estimated) +### Throughput (estimated) On MI300X (192GB): - **Flux 535M**: ~5-10 samples/sec (depends on resolution) - **Flux 12B**: ~0.5-1 samples/sec (requires multi-GPU) -### Optimization Tips +### Optimization tips 1. **Use mixed precision**: `torch.autocast(device_type='cuda', dtype=torch.bfloat16)` 2. **Enable CUDA graphs**: Set `enable_cuda_graph=True` in config @@ -959,9 +959,9 @@ On MI300X (192GB): --- -## Common Issues +## Common issues -### Issue: Import Errors +### Issue: Import errors ```python # Old import style @@ -971,7 +971,7 @@ from some_other_library import Flux from primus.backends.megatron.core.models.diffusion.flux.model import Flux ``` -### Issue: Position ID Shape Mismatch +### Issue: Position ID shape mismatch ```python # Position IDs must be [B, seq, 3] for 3D RoPE @@ -979,7 +979,7 @@ img_ids = generate_image_position_ids(batch_size, height, width) # Not: img_ids = torch.randn(batch_size, height * width, 2) # Wrong! ``` -### Issue: Timestep Range +### Issue: Timestep range ```python # Flux expects timesteps in [0, 1] @@ -989,9 +989,9 @@ timesteps = torch.rand(batch_size) # Correct: [0, 1] --- -## API Compatibility +## API compatibility -### Primus Architecture Features +### Primus architecture features | Aspect | Primus Implementation | |--------|----------------------| @@ -1002,7 +1002,7 @@ timesteps = torch.rand(batch_size) # Correct: [0, 1] | Checkpoint format | `transformer.layers.{0-56}` unified namespace | | Process groups | Via `pg_collection` parameter | -### Key Design Choices +### Key design choices **TransformerBlock Architecture**: - Primus uses Megatron-Core's `TransformerBlock` with heterogeneous layer specifications @@ -1035,14 +1035,14 @@ For more advanced examples, see `examples/run_pretrain.sh`. - **RoPE**: "RoFormer: Enhanced Transformer with Rotary Position Embedding" - **DiT**: "Scalable Diffusion Models with Transformers" (Peebles & Xie, 2023) -### Source Code +### Source code - **Primus Implementation**: `primus/backends/megatron/core/models/diffusion/flux/` - **Megatron-Core**: `megatron/core/transformer/` - **Official Flux**: Black Forest Labs (HuggingFace) --- -## Version Information +## Version information - **Primus Version**: Current - **Megatron-Core Version**: Latest @@ -1054,6 +1054,6 @@ For more advanced examples, see `examples/run_pretrain.sh`. For issues or questions: 1. Check this API reference -2. Read architecture guide: `docs/backends/megatron/diffusion/flux_architecture.md` +2. Read architecture guide: `docs/04-technical-guides/diffusion-models/flux_architecture.md` 3. See examples in docstrings 4. Check test files for usage patterns diff --git a/docs/backends/megatron/diffusion/architecture_overview.md b/docs/04-technical-guides/diffusion-models/architecture_overview.md similarity index 81% rename from docs/backends/megatron/diffusion/architecture_overview.md rename to docs/04-technical-guides/diffusion-models/architecture_overview.md index c58ffe47b..8ef25764f 100644 --- a/docs/backends/megatron/diffusion/architecture_overview.md +++ b/docs/04-technical-guides/diffusion-models/architecture_overview.md @@ -1,23 +1,23 @@ -# Architecture Overview +# Architecture overview This document provides a detailed overview of the diffusion model architecture in Primus, including design decisions, directory structure, and implementation patterns. --- -## Table of Contents +## Table of contents -1. [Design Philosophy](#design-philosophy) -2. [Directory Structure](#directory-structure) -3. [Architectural Decisions](#architectural-decisions) -4. [Component Hierarchy](#component-hierarchy) -5. [Data Flow](#data-flow) -6. [Comparison with Alternative Implementations](#comparison-with-alternative-implementations) +1. [Design philosophy](#design-philosophy) +2. [Directory structure](#directory-structure) +3. [Architectural decisions](#architectural-decisions) +4. [Component hierarchy](#component-hierarchy) +5. [Data flow](#data-flow) +6. [Comparison with alternative implementations](#comparison-with-alternative-implementations) --- -## Design Philosophy +## Design philosophy -### Core Principles +### Core principles 1. **Megatron-Core Native**: Built on Megatron-Core patterns and conventions 2. **Extensibility**: Easy to add new models (DiT, MovieGen, custom) @@ -25,7 +25,7 @@ This document provides a detailed overview of the diffusion model architecture i 4. **Reusability**: Shared components across multiple models 5. **Performance**: Support for precalculated data and multi-GPU training -### Architectural Advantages +### Architectural advantages | Aspect | Primus Design Choice | |--------|---------------------| @@ -34,14 +34,14 @@ This document provides a detailed overview of the diffusion model architecture i | Shared Code | Dedicated `common/` directory | | Encoders | Hierarchical `encoders/{type}/{variant}/` | | Energon | Shared `data/energon/` for all models | -| Mock Data | Separated in `tests/fixtures/` | +| Mock Data | Synthetic providers under `data/synthetic/` | | Framework | Pure Megatron (no PyTorch Lightning dependency) | --- -## Directory Structure +## Directory structure -### High-Level Organization +### High-level organization ``` primus/backends/megatron/ @@ -52,31 +52,35 @@ primus/backends/megatron/ └── diffusion/ # Diffusion-specific data ``` -### Detailed Breakdown +### Detailed breakdown -#### 1. Core Models (`core/models/diffusion/`) +#### 1. Core models (`core/models/diffusion/`) Following Megatron-Core convention (`megatron/core/models/gpt/`, `megatron/core/models/multimodal/`): ``` core/models/diffusion/ -├── common/ # Shared across all diffusion models +├── common/ # Shared building blocks │ ├── __init__.py │ ├── config.py # BaseDiffusionConfig -│ ├── attention.py # JointSelfAttention, FluxSingleAttention -│ └── layers.py # MMDiTLayer, FluxSingleTransformerBlock +│ ├── embeddings.py # TimeStepEmbedder, Timesteps, MLPEmbedder +│ └── normalization.py # RMSNorm, AdaLN, AdaLNContinuous │ └── flux/ # Flux-specific components ├── __init__.py ├── config.py # FluxConfig ├── model.py # Flux (main model class) - └── layers.py # EmbedND (3D RoPE), embedders + ├── layer_spec.py # MMDiTLayer, FluxSingleTransformerBlock + ├── attention.py # JointSelfAttention, FluxSingleAttention + ├── layers.py # EmbedND (RoPE), embedders + ├── checkpoint_converter.py # HF <-> Megatron conversion + └── utils.py # image position ids, helpers ``` **Rationale**: -- `common/`: Shared MMDiT patterns (used by both DiT and Flux) -- `flux/`: Only truly Flux-specific code (3D RoPE, Flux model class) -- Clear distinction enables easy DiT implementation (reuse `common/`) +- `common/`: shared building blocks (base config, timestep/positional embeddings, normalization) reusable across diffusion models +- `flux/`: Flux-specific code (model class, MMDiT and single-block layer specs, joint attention, RoPE embedders) +- Clear separation keeps model-specific code isolated and shared code reusable for future models #### 2. Training (`training/diffusion/`) @@ -97,7 +101,7 @@ training/diffusion/ - Separate from model code for clarity - Easy to add new schedulers (DDPM, EDM, etc.) -#### 3. Data Pipeline (`data/`) +#### 3. Data pipeline (`data/`) ``` data/ @@ -145,9 +149,9 @@ data/ --- -## Architectural Decisions +## Architectural decisions -### Decision 1: Models Under `core/models/` +### Decision 1: Models under `core/models/` **Choice**: `primus/backends/megatron/core/models/diffusion/` **Not**: `primus/backends/megatron/models/diffusion/` @@ -158,7 +162,7 @@ data/ - Clear that these are Megatron-Core compatible - Consistent with existing Primus structure (`core/models/gpt/`) -### Decision 2: Separate `common/` Directory +### Decision 2: Separate `common/` directory **Choice**: Shared components in `common/` **Not**: Everything in `flux/` or flat structure @@ -179,7 +183,7 @@ data/ - `EmbedND`: 3D RoPE position embedding (Flux-specific) - `Flux` model class -### Decision 3: Hierarchical Encoder Structure +### Decision 3: Hierarchical encoder structure **Choice**: `encoders/image/vae/`, `encoders/text/t5/`, `encoders/text/clip/` **Not**: Flat `encoders/conditioner.py` @@ -200,7 +204,7 @@ t5 = get_encoder('t5_xxl', config=t5_config) clip = get_encoder('clip_l', config=clip_config) ``` -### Decision 4: Shared Energon Infrastructure +### Decision 4: Shared Energon infrastructure **Choice**: `data/energon/` for shared utilities **Not**: Nested under `data/diffusion/` @@ -223,18 +227,16 @@ class EncodedDiffusionTaskEncoder: pass ``` -### Decision 5: Model Provider at Adapter Level +### Decision 5: Model provider at adapter level -**Choice**: `diffusion_model_provider.py` at `primus/backends/megatron/` -**Not**: Under `core/` +**Choice**: construct the model through provider functions in the Megatron trainer/adapter layer under `primus/backends/megatron/`, above `core/models/`. **Reasoning**: -- Follows existing pattern (`primus/backends/megatron/model_provider.py`) -- Model providers are adapter/wrapper functions -- Sit above core models to add Primus-specific functionality -- Example: Wrap model with logit softcapping, custom loss, etc. +- Model providers are adapter/wrapper functions rather than a standalone module; they live alongside the trainers (for example `primus/backends/megatron/megatron_pretrain_trainer.py`) +- They sit above the core models to add Primus-specific functionality +- Example: wrap the model with a custom loss, precision handling, or checkpoint logic -### Decision 6: No PyTorch Lightning +### Decision 6: No PyTorch lightning **Choice**: Pure Megatron patterns **Not**: PyTorch Lightning DataModules @@ -247,9 +249,9 @@ class EncodedDiffusionTaskEncoder: --- -## Component Hierarchy +## Component hierarchy -### 1. Model Hierarchy +### 1. Model hierarchy ``` nn.Module (PyTorch) @@ -266,7 +268,7 @@ nn.Module (PyTorch) └── EmbedND (Flux-specific, 3D RoPE) ``` -### 2. Configuration Hierarchy +### 2. Configuration hierarchy ``` TransformerConfig (Megatron-Core) @@ -276,7 +278,7 @@ TransformerConfig (Megatron-Core) └── flux_12b() factory ``` -### 3. Scheduler Hierarchy +### 3. Scheduler hierarchy ``` BaseScheduler (abstract) @@ -286,7 +288,7 @@ BaseScheduler (abstract) └── EulerDiscreteScheduler (future) ``` -### 4. Encoder Hierarchy +### 4. Encoder hierarchy ``` BaseEncoder (abstract) @@ -305,9 +307,9 @@ BaseEncoder (abstract) --- -## Data Flow +## Data flow -### Training Pipeline +### Training pipeline ``` Raw Data (images + captions) @@ -348,7 +350,7 @@ Loss Computation: MSE(v_pred, v_target) Backward Pass & Optimizer Step ``` -### Inference Pipeline +### Inference pipeline ``` Text Prompt @@ -373,31 +375,31 @@ Generated Image [1, 3, H*8, W*8] --- -## Comparison with Alternative Implementations +## Comparison with alternative implementations -### Primus Architectural Advantages +### Primus architectural advantages -#### 1. TransformerBlock Architecture (Primus Innovation) +#### 1. TransformerBlock architecture (Primus innovation) - **Primus**: Unified `TransformerBlock` with heterogeneous layer specs - **Others**: Separate `nn.ModuleList` containers for double/single blocks - **Benefit**: Better PP slicing, unified checkpointing, future-proof -#### 2. Megatron-Core Native +#### 2. Megatron-core native - **Primus**: Pure Megatron-Core, no framework dependencies - **Others**: Often integrated with PyTorch Lightning or other frameworks - **Benefit**: Tighter integration, simpler training loops -#### 3. Checkpoint Format +#### 3. Checkpoint format - **Primus**: Unified `transformer.layers.{0-56}` structure - **Others**: Separate `double_blocks.{i}` and `single_blocks.{j}` - **Benefit**: Simpler distributed checkpointing -#### 4. Encoder Architecture +#### 4. Encoder architecture - **Primus**: Registry-based, hierarchical organization - **Others**: Direct imports from monolithic files - **Benefit**: Easy extensibility for new encoder variants -### File Organization Comparison +### File organization comparison | Component | Standard Location | Primus Location | Improvement | |-----------|------------------|-----------------|-------------| @@ -408,9 +410,9 @@ Generated Image [1, 3, H*8, W*8] --- -## Implementation Status +## Implementation status -### Core Infrastructure ✅ +### Core infrastructure ✅ - ✅ Directory structure - ✅ Base classes (DiffusionModule, BaseDiffusionConfig, BaseScheduler) - ✅ DiffusionModule with Megatron-Core integration @@ -420,7 +422,7 @@ Generated Image [1, 3, H*8, W*8] - ✅ Testing framework (290+ tests) - ✅ Documentation structure -### Flux Model Implementation ✅ +### Flux model implementation ✅ - ✅ Flux model architecture - ✅ MMDiT layers and attention - ✅ Embeddings (RoPE, timestep, vector) @@ -430,9 +432,9 @@ Generated Image [1, 3, H*8, W*8] --- -## Extension Points +## Extension points -### Adding a New Model (e.g., DiT) +### Adding a new model (e.g., DiT) 1. **Create model directory**: `core/models/diffusion/dit/` 2. **Add config**: Extend `BaseDiffusionConfig` @@ -444,7 +446,7 @@ Generated Image [1, 3, H*8, W*8] 5. **Add tests**: `tests/unit_tests/backends/megatron/diffusion/test_dit_model.py` 6. **Update configs**: Add `dit_config.yaml` -### Adding a New Encoder Variant +### Adding a new encoder variant 1. **Create encoder file**: e.g., `data/diffusion/encoders/text/t5/t5_large.py` 2. **Implement encoder class**: Extend `BaseEncoder` @@ -452,7 +454,7 @@ Generated Image [1, 3, H*8, W*8] 4. **Add config**: Update `encoders.yaml` 5. **Add tests**: Test in `tests/unit_tests/backends/megatron/diffusion/data/encoders/` -### Adding a New Scheduler +### Adding a new scheduler 1. **Create scheduler file**: `training/diffusion/schedulers/ddpm.py` 2. **Implement**: Extend `BaseScheduler` @@ -462,21 +464,21 @@ Generated Image [1, 3, H*8, W*8] --- -## Performance Considerations +## Performance considerations -### Memory Optimization +### Memory optimization - **Precalculated data**: 5-10x faster, lower memory - **Frozen encoders**: Only train diffusion model - **Gradient checkpointing**: Trade compute for memory - **Mixed precision**: bf16 on MI300X (compatible with H100/A100) -### Multi-GPU Scaling +### Multi-GPU scaling - **Tensor Parallelism**: Split model across GPUs - **Pipeline Parallelism**: Split layers across GPUs - **Data Parallelism**: Replicate model, split data - **Sequence Parallelism**: For very long sequences -### Best Practices +### Best practices 1. Use precalculated mode for training 2. Freeze encoders (standard practice) 3. Use bf16 on modern hardware diff --git a/docs/backends/megatron/diffusion/data_preprocessing.md b/docs/04-technical-guides/diffusion-models/data_preprocessing.md similarity index 93% rename from docs/backends/megatron/diffusion/data_preprocessing.md rename to docs/04-technical-guides/diffusion-models/data_preprocessing.md index b5a592729..a99a9754f 100644 --- a/docs/backends/megatron/diffusion/data_preprocessing.md +++ b/docs/04-technical-guides/diffusion-models/data_preprocessing.md @@ -1,19 +1,19 @@ -# Data Preprocessing Guide +# Data preprocessing guide This guide explains how to prepare datasets for Flux and other diffusion models in Primus, including pre-encoding of VAE latents and text embeddings into Energon WebDataset format. --- -## Table of Contents +## Table of contents 1. [Overview](#overview) -2. [Quick Start](#quick-start) -3. [Two Pipelines](#two-pipelines) -4. [Running Preprocessing](#running-preprocessing) +2. [Quick start](#quick-start) +3. [Two pipelines](#two-pipelines) +4. [Running preprocessing](#running-preprocessing) 5. [Configuration](#configuration) 6. [Authentication](#authentication) 7. [Finalization](#finalization) -8. [Output Format](#output-format) +8. [Output format](#output-format) 9. [Validation](#validation) 10. [Troubleshooting](#troubleshooting) @@ -26,7 +26,7 @@ Diffusion models require three types of encodings: 2. **Text embeddings**: Captions encoded with T5-XXL (sequence) 3. **Pooled embeddings**: Captions encoded with CLIP-L (pooled) -### Why Pre-encode? +### Why pre-encode? **Benefits**: - 5-10x faster training (no online encoding) @@ -46,7 +46,7 @@ Diffusion models require three types of encodings: --- -## Quick Start +## Quick start Preprocess the Pokemon dataset with a single command: @@ -66,11 +66,11 @@ The default encoder model (`black-forest-labs/FLUX.1-dev`) is gated and requires --- -## Two Pipelines +## Two pipelines Primus provides two preprocessing pipelines via the `primus data` CLI: -### `diffusion-encoded` (Recommended for Training) +### `diffusion-encoded` (recommended for training) Pre-encodes images with VAE and text with T5/CLIP. Produces larger datasets but enables faster training since encoders are not needed at training time. @@ -104,9 +104,9 @@ primus-cli direct -- data diffusion-raw \ --- -## Running Preprocessing +## Running preprocessing -### Using a Config File (Recommended) +### Using a config file (recommended) The `--config` flag is supported by `diffusion-encoded` only. The simplest approach uses a YAML config file: @@ -127,7 +127,7 @@ Available example configs in `primus/configs/data/megatron/diffusion/preprocessi | `example_base.yaml` | N/A | Comprehensive reference with all fields | | `text_to_image_2m_10k.yaml` | HuggingFace | 10K subset of text-to-image-2M (1024px) | -### Using CLI Arguments Directly +### Using CLI arguments directly All config values can be provided as CLI arguments: @@ -142,7 +142,7 @@ primus-cli direct -- data diffusion-encoded \ --hf-token-file /path/to/.hf_token ``` -### CLI Overrides Config Values +### CLI overrides config values When using both `--config` and CLI arguments, CLI arguments take priority: @@ -160,7 +160,7 @@ Priority order (highest to lowest): 2. YAML config file values 3. CLI default values -### Multi-GPU Processing +### Multi-GPU processing Use `--nproc-per-node` for data-parallel preprocessing across multiple GPUs: @@ -176,7 +176,7 @@ Each GPU processes a subset of the data. Shards are named to avoid conflicts acr ## Configuration -### YAML Config Structure +### YAML config structure Preprocessing configs have four sections: @@ -206,7 +206,7 @@ image: center_crop: false ``` -### Source Types +### Source types **HuggingFace** (`type: huggingface`): ```yaml @@ -244,7 +244,7 @@ source: input_path: /data/existing_shards/*.tar ``` -### Model Configuration +### Model configuration The `model_path` defaults to `black-forest-labs/FLUX.1-dev`, which downloads VAE, T5-XXL, and CLIP-L encoders from HuggingFace. This model is gated and requires authentication (see [Authentication](#authentication)). @@ -264,7 +264,7 @@ model: The default encoder model (`FLUX.1-dev`) is gated on HuggingFace and requires authentication. Primus supports three authentication methods, checked in priority order: -### 1. Token File (Recommended) +### 1. Token file (recommended) ```bash primus-cli direct -- data diffusion-encoded \ @@ -279,14 +279,14 @@ echo "hf_your_token_here" > /path/to/.hf_token chmod 600 /path/to/.hf_token ``` -### 2. Environment Variable +### 2. Environment variable ```bash export HF_TOKEN=hf_your_token_here primus-cli direct -- data diffusion-encoded --config your_config.yaml ``` -### 3. HuggingFace CLI Login +### 3. HuggingFace CLI login ```bash huggingface-cli login @@ -305,7 +305,7 @@ Finalization is **automatic by default**. After preprocessing completes, Primus 2. **Runs `energon prepare`** to index the tar shards and create split assignments 3. **Validates the dataset** using Primus's custom validation (metadata checks, sample count verification, energon API spot-check) -### Skipping Finalization +### Skipping finalization To skip automatic finalization (e.g., for manual post-processing): @@ -316,7 +316,7 @@ primus-cli direct -- data diffusion-encoded \ --no-finalize ``` -### Custom Train/Val/Test Splits +### Custom train/val/test splits By default, 100% of data goes to the training split. To create validation and test splits: @@ -331,9 +331,9 @@ This creates an 80% train / 10% val / 10% test split. --- -## Output Format +## Output format -### Directory Structure +### Directory structure After preprocessing and finalization, the output directory contains: @@ -349,7 +349,7 @@ encoded_pokemon/ └── .info.json # Shard counts and sample counts ``` -### Pre-encoded Shard Contents (`diffusion-encoded`) +### Pre-encoded shard contents (`diffusion-encoded`) Each tar shard contains samples with these keys: @@ -369,7 +369,7 @@ Tensor shapes: - `prompt_embeds.pth`: `[seq_len, 4096]` (T5-XXL hidden dim) - `pooled_prompt_embeds.pth`: `[768]` (CLIP-L pooled dim) -### Raw Shard Contents (`diffusion-raw`) +### Raw shard contents (`diffusion-raw`) ``` 000000.tar: @@ -395,7 +395,7 @@ subflavors: ## Validation -### Automatic Validation +### Automatic validation Validation runs automatically as part of finalization. It performs four checks: @@ -404,7 +404,7 @@ Validation runs automatically as part of finalization. It performs four checks: 3. **Sample load check**: Loads one sample through Energon's Python API (same code path as training) 4. **Summary report**: Prints dataset statistics (encoding, total samples, splits, data shapes, size) -### Standalone Validation +### Standalone validation To validate a dataset independently: @@ -415,7 +415,7 @@ python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/t python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset --encoding raw ``` -### Programmatic Validation +### Programmatic validation ```python from primus.backends.megatron.data.diffusion.preprocessing.validate import validate_energon_dataset @@ -423,7 +423,7 @@ from primus.backends.megatron.data.diffusion.preprocessing.validate import valid ok = validate_energon_dataset('/path/to/dataset', encoding='preencoded') ``` -### Known Energon CLI Limitations +### Known Energon CLI limitations The standard Energon CLI tools (`energon info`, `energon preview`, `energon lint`) do **not** work correctly with `CrudeWebdataset` format. Primus uses custom validation instead: @@ -437,7 +437,7 @@ Use Primus's built-in validation or the standalone script above. ## Troubleshooting -### HuggingFace Authentication Failure +### HuggingFace authentication failure **Symptoms**: Error mentioning "token", "gated", "401", or "403" when downloading encoders. @@ -447,7 +447,7 @@ Use Primus's built-in validation or the standalone script above. 3. Run `huggingface-cli login` 4. Accept the model's license on https://huggingface.co/black-forest-labs/FLUX.1-dev -### Out of Memory During Preprocessing +### Out of memory during preprocessing **Symptoms**: CUDA out of memory error during encoding. @@ -457,7 +457,7 @@ Use Primus's built-in validation or the standalone script above. 3. Use fp16 precision: `--precision fp16` 4. Use multi-GPU to distribute work: `--nproc-per-node=8` -### Missing Dependencies +### Missing dependencies **Symptoms**: `ModuleNotFoundError` for `webdataset`, `megatron-energon`, `tqdm`, etc. @@ -468,7 +468,7 @@ pip install -r requirements.txt Or set `PRIMUS_AUTO_INSTALL=1` in the container to auto-install missing packages. -### Finalization Fails +### Finalization fails **Symptoms**: Error during `energon prepare` or validation after preprocessing. @@ -477,7 +477,7 @@ Or set `PRIMUS_AUTO_INSTALL=1` in the container to auto-install missing packages 2. Ensure `megatron-energon` is installed (`pip install megatron-energon`) 3. Re-run with `--no-finalize`, then manually inspect the output before finalizing -### Slow Preprocessing +### Slow preprocessing **Symptoms**: Low throughput (< 10 samples/sec). @@ -489,7 +489,7 @@ Or set `PRIMUS_AUTO_INSTALL=1` in the container to auto-install missing packages --- -## Best Practices +## Best practices 1. **Always pre-encode for production training**: 5-10x speedup is worth the storage 2. **Test on small dataset first**: Use `--max-samples 100` to verify the pipeline works @@ -500,7 +500,7 @@ Or set `PRIMUS_AUTO_INSTALL=1` in the container to auto-install missing packages --- -## Flux-Specific Data Preparation (torchrun) +## Flux-specific data preparation (torchrun) This section covers preparing datasets for Flux training using `torchrun` directly inside a Docker container, as an alternative to the `primus-cli` workflow above. @@ -521,7 +521,7 @@ Override the image with `DOCKER_IMAGE`: DOCKER_IMAGE=docker.io/rocm/primus:v26.1 bash tools/docker/start_container.sh ``` -### Input Directory Structure +### Input directory structure When using `--source-type directory`, organize your data as follows: @@ -545,7 +545,7 @@ Other supported source types: - **huggingface** -- Load directly from HuggingFace Hub. Requires `--hf-dataset`. - **webdataset** -- Read from existing WebDataset tar archives. Requires `--input-path`. -### Image Sizing +### Image sizing **Fixed size (default):** Every image is resized to `--image-size` pixels square (default 1024). Use `--center-crop` to control center-cropping. @@ -555,7 +555,7 @@ longest side to `--max-size` (default 1024) and rounding dimensions to multiples of 16. Only use this when all images share the same dimensions -- mixed tensor sizes cause load imbalance across GPUs. -### Key Parameters +### Key parameters | Parameter | Default | Description | |-----------|---------|-------------| @@ -613,7 +613,7 @@ your images have various sizes. --- -## Next Steps +## Next steps After preprocessing: 1. **Training**: Use the preprocessed dataset path in your training config @@ -621,7 +621,7 @@ After preprocessing: See: - [Energon Integration](energon_integration.md) for TaskEncoder and dataloader details -- [Config Directory Guide](../../../../primus/configs/data/megatron/diffusion/README.md) for config file reference +- [Config Directory Guide](../../../primus/configs/data/megatron/diffusion/README.md) for config file reference - Example configs in `primus/configs/data/megatron/diffusion/preprocessing/` --- diff --git a/docs/backends/megatron/diffusion/energon_integration.md b/docs/04-technical-guides/diffusion-models/energon_integration.md similarity index 92% rename from docs/backends/megatron/diffusion/energon_integration.md rename to docs/04-technical-guides/diffusion-models/energon_integration.md index 032e7222f..97941a67c 100644 --- a/docs/backends/megatron/diffusion/energon_integration.md +++ b/docs/04-technical-guides/diffusion-models/energon_integration.md @@ -1,19 +1,19 @@ -# Energon Integration Guide +# Energon integration guide This guide explains how Megatron-Energon is integrated with Primus diffusion models, including the Cooker/TaskEncoder pattern, dataloader configuration, and dataset format. --- -## Table of Contents +## Table of contents 1. [Overview](#overview) -2. [Energon Architecture](#energon-architecture) -3. [Dataset Format](#dataset-format) -4. [TaskEncoder and Cooker Pattern](#taskencoder-and-cooker-pattern) -5. [Dataloader Setup](#dataloader-setup) -6. [Dataset Configuration](#dataset-configuration) -7. [Implementation Examples](#implementation-examples) -8. [Best Practices](#best-practices) +2. [Energon architecture](#energon-architecture) +3. [Dataset format](#dataset-format) +4. [TaskEncoder and cooker pattern](#taskencoder-and-cooker-pattern) +5. [Dataloader setup](#dataloader-setup) +6. [Dataset configuration](#dataset-configuration) +7. [Implementation examples](#implementation-examples) +8. [Best practices](#best-practices) --- @@ -28,14 +28,14 @@ Megatron-Energon is NVIDIA's data loading framework for large-scale multimodal t - **Deterministic iteration**: Reproducible training - **Checkpoint resumption**: Resume from any step -### Why Energon for Diffusion? +### Why Energon for diffusion? - **Proven at scale**: Used by NVIDIA for LLM and multimodal training - **Flexible**: Supports various data formats and transformations - **Efficient**: Optimized for multi-GPU training - **Compatible**: Works with Megatron parallelism (TP, PP, DP) -### Primus Integration Strategy +### Primus integration strategy ``` Shared Infrastructure (data/) @@ -53,9 +53,9 @@ Model-Specific TaskEncoders (data/diffusion/task_encoders/) --- -## Energon Architecture +## Energon architecture -### Component Stack +### Component stack ``` Training Loop @@ -71,7 +71,7 @@ Megatron-Energon Core .tar Shards (CrudeWebdataset format) ``` -### Data Flow +### Data flow ``` 1. Load from .tar shard @@ -101,7 +101,7 @@ Megatron-Energon Core --- -## Dataset Format +## Dataset format ### CrudeWebdataset @@ -120,7 +120,7 @@ The `subflavors.encoding` field tells the TaskEncoder which Cooker to use: - `preencoded`: Uses `cook_preencoded_diffusion` -- loads pre-encoded tensors - `raw`: Uses `cook_raw_images` -- loads raw images and text -### Pre-encoded Shard Contents +### Pre-encoded shard contents Each tar shard contains samples with these keys: @@ -131,14 +131,14 @@ Each tar shard contains samples with these keys: | `pooled_prompt_embeds.pth` | Tensor | `[768]` | CLIP-L pooled embeddings | | `caption.txt` | str | N/A | Original caption text | -### Raw Shard Contents +### Raw shard contents | Key | Type | Description | |-----|------|-------------| | `jpg` / `png` / `webp` | bytes | Preprocessed image | | `txt` | str | Caption text | -### Known Energon CLI Limitations +### Known Energon CLI limitations The standard Energon CLI tools have issues with `CrudeWebdataset`: - `energon info` raises `KeyError: 'sample_type'` @@ -149,11 +149,11 @@ Primus includes custom validation (`primus.backends.megatron.data.diffusion.prep --- -## TaskEncoder and Cooker Pattern +## TaskEncoder and cooker pattern Primus uses Energon's Cooker pattern rather than the older `encode_sample()` approach. Cookers are `@stateless` functions that transform raw sample dicts into typed dataclass instances, dispatched based on `subflavors`. -### DiffusionSample Dataclass +### DiffusionSample dataclass Location: `primus/backends/megatron/data/diffusion/task_encoders/image.py` @@ -176,7 +176,7 @@ class DiffusionSample(Sample): caption: str = "" ``` -### Cooker Functions +### Cooker functions Cookers are `@stateless` functions registered with a `Cooker` wrapper that specifies which `subflavors` they handle: @@ -275,7 +275,7 @@ class RawDiffusionTaskEncoder(DefaultTaskEncoder): } ``` -### How Cooker Dispatch Works +### How cooker dispatch works The Cooker framework matches samples to cooker functions based on `subflavors`: @@ -289,7 +289,7 @@ This decouples data format (what's in the tar) from data loading logic (how to i --- -## Dataloader Setup +## Dataloader setup ### MegatronDataloaderWrapper @@ -308,7 +308,7 @@ wrapper = MegatronDataloaderWrapper(pytorch_or_energon_loader) Note: Originally named `EnergonDataloader`, renamed to `MegatronDataloaderWrapper` to reflect its generic nature (it has no Energon dependencies). The old name is available as a deprecated alias. -### Usage Example +### Usage example ```python from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper @@ -330,9 +330,9 @@ for batch in dataloader: --- -## Dataset Configuration +## Dataset configuration -### Per-Dataset dataset.yaml +### Per-dataset dataset.yaml Each dataset directory has a `.nv-meta/dataset.yaml` that specifies the Energon dataset type: @@ -345,7 +345,7 @@ subflavors: This file is auto-generated by Primus during finalization. See [Data Preprocessing Guide](data_preprocessing.md#finalization). -### Metadataset (Multi-Dataset Mixing) +### Metadataset (multi-dataset mixing) For combining multiple datasets with different weights: @@ -366,9 +366,9 @@ Energon samples proportionally to weights (70% from LAION, 30% from COCO). --- -## Implementation Examples +## Implementation examples -### Example 1: Pre-encoded Training +### Example 1: Pre-encoded training ```python from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper @@ -387,7 +387,7 @@ for batch in dataloader: ) ``` -### Example 2: Raw Data Training (On-the-Fly Encoding) +### Example 2: Raw data training (on-the-fly encoding) ```python from primus.backends.megatron.data.diffusion.task_encoders import RawDiffusionTaskEncoder @@ -403,7 +403,7 @@ for batch in dataloader: # Model handles VAE/T5/CLIP encoding in forward_step ``` -### Example 3: Multi-GPU Training +### Example 3: Multi-GPU training ```python import torch.distributed as dist @@ -419,25 +419,25 @@ for batch in dataloader: --- -## Best Practices +## Best practices -### 1. Pre-encoded Mode +### 1. Pre-encoded mode - **Always use for production training**: 5-10x faster - **Validate first**: Test with small dataset using `quickstart_pokemon.yaml` - **Version datasets**: Track which preprocessing config produced each dataset -### 2. Cooker Design +### 2. Cooker design - **Use `@stateless`**: Cookers must be stateless and side-effect-free - **Use `basic_sample_keys()`**: Always forward `__key__`, `__restore_key__`, `__subflavors__` - **Keep it simple**: Cookers should only deserialize and restructure data, not transform it - **Use subflavors for dispatch**: Let the framework choose the right cooker -### 3. TaskEncoder Design +### 3. TaskEncoder design - **Single responsibility**: One task encoder per data format family - **Minimal `batch()`**: Only stack tensors, avoid computation in the batch method - **Separate concerns**: Data loading in Cooker, encoding in model forward_step -### 4. Dataloader Configuration +### 4. Dataloader configuration - **Num workers**: Match CPU cores (typically 4-8) - **Batch size**: Max out GPU memory - **Shuffle**: Always true for training @@ -450,9 +450,9 @@ for batch in dataloader: --- -## Customization Patterns +## Customization patterns -### Custom Cooker Function +### Custom cooker function Add a new cooker for a different data format: @@ -487,7 +487,7 @@ class CustomDiffusionTaskEncoder(DefaultTaskEncoder[DiffusionSample, DiffusionSa } ``` -### Multiple Cookers in One TaskEncoder +### Multiple cookers in one TaskEncoder A single TaskEncoder can register multiple cookers for different subflavors: diff --git a/docs/backends/megatron/diffusion/flux_architecture.md b/docs/04-technical-guides/diffusion-models/flux_architecture.md similarity index 93% rename from docs/backends/megatron/diffusion/flux_architecture.md rename to docs/04-technical-guides/diffusion-models/flux_architecture.md index 02189af37..61dcb6bb6 100644 --- a/docs/backends/megatron/diffusion/flux_architecture.md +++ b/docs/04-technical-guides/diffusion-models/flux_architecture.md @@ -1,15 +1,15 @@ -# Flux Architecture Deep Dive +# Flux architecture deep dive -## Table of Contents +## Table of contents 1. [Overview](#overview) -2. [Architecture Principles](#architecture-principles) -3. [Model Components](#model-components) -4. [Data Flow](#data-flow) -5. [Mathematical Formulation](#mathematical-formulation) -6. [Implementation Details](#implementation-details) -7. [Megatron-Core Integration](#megatron-core-integration) -8. [Performance Optimizations](#performance-optimizations) +2. [Architecture principles](#architecture-principles) +3. [Model components](#model-components) +4. [Data flow](#data-flow) +5. [Mathematical formulation](#mathematical-formulation) +6. [Implementation details](#implementation-details) +7. [Megatron-core integration](#megatron-core-integration) +8. [Performance optimizations](#performance-optimizations) --- @@ -17,14 +17,14 @@ Flux is a **flow-based diffusion model** for high-quality text-to-image generation. It uses an innovative **MMDiT (Multimodal Diffusion Transformer)** architecture that jointly processes image and text tokens through shared transformer blocks. -### Key Innovations +### Key innovations 1. **Flow Matching**: Uses rectified flow instead of traditional diffusion 2. **MMDiT Architecture**: Joint image-text attention in early layers 3. **3D RoPE**: Multi-dimensional rotary position embeddings for spatial awareness 4. **Two-Stage Processing**: Joint layers followed by image-only layers -### Model Variants +### Model variants | Variant | Joint Layers | Single Layers | Parameters | Use Case | |---------|--------------|---------------|------------|----------| @@ -33,9 +33,9 @@ Flux is a **flow-based diffusion model** for high-quality text-to-image generati --- -## Architecture Principles +## Architecture principles -### 1. Flow Matching Framework +### 1. Flow matching framework Unlike traditional diffusion (which adds Gaussian noise), Flux uses **rectified flow**: @@ -58,7 +58,7 @@ v_θ(z_t, t, c) ≈ z_1 - z_0 - Faster sampling (fewer steps needed) - Better training stability -### 2. MMDiT (Multimodal Diffusion Transformer) +### 2. MMDiT (multimodal diffusion transformer) Traditional DiT processes image tokens independently. MMDiT jointly processes image and text: @@ -87,7 +87,7 @@ Traditional DiT processes image tokens independently. MMDiT jointly processes im - Richer cross-modal interactions - Improved compositional understanding -### 3. Two-Stage Processing +### 3. Two-stage processing Flux uses a unique two-stage architecture: @@ -103,11 +103,11 @@ Flux uses a unique two-stage architecture: --- -## Model Components +## Model components -### 1. Input Embeddings +### 1. Input embeddings -#### Image Path +#### Image path ``` Image (RGB) → VAE Encoder → Latents [B, 64, H/8, W/8] @@ -123,7 +123,7 @@ Image (RGB) → VAE Encoder → Latents [B, 64, H/8, W/8] **Linear Projection**: Maps 64 channels → 3072 hidden dim -#### Text Path +#### Text path ``` Caption → T5-XXL Encoder → Embeddings [B, S, 4096] @@ -159,7 +159,7 @@ Timestep t ∈ [0, 1] → Sinusoidal Encoding → [B, 256] vec = timestep_emb + clip_pooled_emb + [guidance_emb] ``` -### 2. Position Embeddings (3D RoPE) +### 2. Position embeddings (3D RoPE) Flux uses **3D Rotary Position Embeddings** for spatial awareness: @@ -191,7 +191,7 @@ sin_freq = sin(freqs) - Encodes channel relationships - Works for any resolution (generalization) -### 3. MMDiT Layer (Double Block) +### 3. MMDiT layer (double block) Each MMDiT layer performs: @@ -234,7 +234,7 @@ Output: img [B, H*W, D], txt [B, S, D] - **Adaptive gating**: Timestep-conditioned residual connections - **Separate MLPs**: Modality-specific processing -### 4. Flux Single Block +### 4. Flux single block After joint processing, image tokens go through single blocks: @@ -270,7 +270,7 @@ Output: img [B, H*W, D] (text tokens discarded) - Output focuses on image generation - More efficient than full joint processing -### 5. Output Processing +### 5. Output processing ``` Image Tokens [B, H*W, D] @@ -288,9 +288,9 @@ Image Tokens [B, H*W, D] --- -## Data Flow +## Data flow -### Complete Forward Pass +### Complete forward pass ``` Input @@ -355,7 +355,7 @@ Image Tokens [B, H*W, D] Predicted Velocity ``` -### Training Data Flow +### Training data flow ``` Original Image @@ -386,9 +386,9 @@ Original Image --- -## Mathematical Formulation +## Mathematical formulation -### Flow Matching Objective +### Flow matching objective **Forward Process**: ``` @@ -411,7 +411,7 @@ where: v_θ = Flux model # Predicted velocity ``` -### Sampling (Inference) +### Sampling (inference) **Euler Integration** (first-order ODE solver): ``` @@ -425,7 +425,7 @@ for t in [1.0, 0.9, ..., 0.1, 0.0]: - Heun's method (2nd order) - DPM-Solver (adaptive) -### Classifier-Free Guidance +### Classifier-free guidance During inference, use guidance scale `w`: @@ -443,7 +443,7 @@ where: config = FluxConfig(guidance_embed=True, guidance_scale=3.5) ``` -### 3D RoPE Mathematics +### 3D RoPE mathematics For position `(h, w)` in image: @@ -470,9 +470,9 @@ k_rot = [k[:d/2] * cos(freq) - k[d/2:] * sin(freq), --- -## Implementation Details +## Implementation details -### Memory Layout +### Memory layout Megatron-Core uses **sequence-first format**: `[seq, batch, hidden]` @@ -491,7 +491,7 @@ img_tokens = linear(img_seq) # [B, H*W, 3072] img_megatron = rearrange(img_tokens, 'b s d -> s b d') ``` -### Adaptive Layer Normalization +### Adaptive layer normalization **Standard AdaLN**: ```python @@ -514,7 +514,7 @@ x_attn = attention(x_norm) x = x + gate * x_attn # Gated residual ``` -### Attention Implementation +### Attention implementation **Using Megatron SelfAttention**: ```python @@ -546,7 +546,7 @@ txt_out = joint_output[seq_img:] --- -## Megatron-Core Integration +## Megatron-core integration ### TransformerConfig @@ -568,7 +568,7 @@ config = TransformerConfig( ) ``` -### Layer Specs +### Layer specs Flux provides factory functions for layer specs: @@ -589,7 +589,7 @@ single_spec = get_flux_single_transformer_spec_for_backend(backend) layer_specs = get_flux_layer_spec(config, backend=backend) ``` -### Distributed Training Support +### Distributed training support Flux inherits Megatron's parallelism: @@ -610,9 +610,9 @@ at config construction). --- -## Performance Optimizations +## Performance optimizations -### 1. Transformer Engine +### 1. Transformer engine Flux uses NVIDIA Transformer Engine for FP8 training: @@ -632,7 +632,7 @@ linear_proj = TERowParallelLinear(...) - Fused operations (LayerNorm + Linear) - Automatic scaling for numerical stability -### 2. Flash Attention +### 2. Flash attention Enabled via Megatron: @@ -644,7 +644,7 @@ config = TransformerConfig( **Speedup**: 2-3x faster attention, 4x less memory -### 3. Gradient Checkpointing +### 3. Gradient checkpointing For large models: @@ -658,7 +658,7 @@ config = TransformerConfig( **Memory Savings**: ~40% reduction, ~20% slower -### 4. Fused Operations +### 4. Fused operations ```python config = TransformerConfig( @@ -668,7 +668,7 @@ config = TransformerConfig( ) ``` -### 5. Mixed Precision +### 5. Mixed precision ```python from primus.backends.megatron.training.diffusion.loss_computation import compute_flow_matching_loss @@ -688,7 +688,7 @@ scaler.update() --- -## Comparison with Other Models +## Comparison with other models ### Flux vs DiT @@ -699,7 +699,7 @@ scaler.update() | Position Encoding | Learned 2D | 3D RoPE | | Diffusion Type | DDPM | Flow matching | -### Flux vs Stable Diffusion +### Flux vs stable diffusion | Aspect | Stable Diffusion (UNet) | Flux (Transformer) | |--------|-------------------------|---------------------| @@ -710,9 +710,9 @@ scaler.update() --- -## Design Decisions +## Design decisions -### Why Two-Stage (Joint + Single)? +### Why two-stage (joint + single)? **Joint Blocks**: - Deep semantic understanding @@ -726,7 +726,7 @@ scaler.update() **Alternative**: All joint blocks → slower, marginal quality gain -### Why Flow Matching? +### Why flow matching? **Advantages over DDPM**: 1. **Simpler training**: Straight-line interpolation (no schedule design) @@ -746,9 +746,9 @@ scaler.update() --- -## Primus Implementation Highlights +## Primus implementation highlights -### TransformerBlock Architecture +### TransformerBlock architecture Primus's key architectural enhancement is the use of Megatron-Core's **TransformerBlock with heterogeneous layer specifications**: @@ -774,7 +774,7 @@ graph TD 2. **Future-Proof**: Native support for new Megatron-Core features 3. **Cleaner Code**: No manual iteration over separate block lists -### Layer Specification Pattern +### Layer specification pattern ```python # Primus approach @@ -794,7 +794,7 @@ transformer = TransformerBlock( # No manual offset calculation needed ``` -### Checkpoint Format Comparison +### Checkpoint format comparison **Traditional Format**: ``` @@ -816,9 +816,9 @@ Benefits: Simpler distributed checkpointing, easier layer inspection, consistent --- -## Future Enhancements +## Future enhancements -### Planned Features +### Planned features 1. **ControlNet Support**: - Spatial conditioning (pose, depth, edges) @@ -836,7 +836,7 @@ Benefits: Simpler distributed checkpointing, easier layer inspection, consistent - Low-rank adaptation for custom styles - Efficient personalization -### Research Directions +### Research directions - **Sparse Attention**: Reduce quadratic complexity for high-res - **Mixture of Experts**: Conditional computation for efficiency @@ -854,7 +854,7 @@ Benefits: Simpler distributed checkpointing, easier layer inspection, consistent 4. **RoPE**: Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding", 2021 5. **Transformer Engine**: NVIDIA, "Transformer Engine: Accelerating Transformer Training", 2022 -### Code References +### Code references - **Primus Flux**: `primus/backends/megatron/core/models/diffusion/flux/` - **Megatron-Core**: `megatron/core/transformer/` @@ -865,7 +865,7 @@ Benefits: Simpler distributed checkpointing, easier layer inspection, consistent ## Appendix: Hyperparameters -### Flux 535M Training +### Flux 535M training ```yaml model: @@ -890,7 +890,7 @@ optimization: gradient_clip: 1.0 ``` -### Flux 12B Training +### Flux 12B training ```yaml model: diff --git a/docs/backends/megatron/diffusion/fp8_training.md b/docs/04-technical-guides/diffusion-models/fp8_training.md similarity index 93% rename from docs/backends/megatron/diffusion/fp8_training.md rename to docs/04-technical-guides/diffusion-models/fp8_training.md index 8f846dfc5..9067bf201 100644 --- a/docs/backends/megatron/diffusion/fp8_training.md +++ b/docs/04-technical-guides/diffusion-models/fp8_training.md @@ -1,4 +1,4 @@ -# FP8 Training for Flux Models +# FP8 training for Flux models Complete guide for training Flux diffusion models with FP8 (8-bit floating point) precision on AMD MI300X GPUs using Transformer Engine's delayed scaling recipe. @@ -11,28 +11,28 @@ FP8 training provides significant memory and speed improvements while maintainin - **Maintains numerical stability** via delayed scaling - **Enables larger batch sizes** or higher resolutions -## Table of Contents +## Table of contents - [Prerequisites](#prerequisites) -- [Quick Start](#quick-start) +- [Quick start](#quick-start) - [Configuration](#configuration) -- [Performance Benchmarks](#performance-benchmarks) +- [Performance benchmarks](#performance-benchmarks) - [Troubleshooting](#troubleshooting) -- [Best Practices](#best-practices) -- [AMD MI300X Specific](#amd-mi300x-specific) +- [Best practices](#best-practices) +- [AMD MI300X specific](#amd-mi300x-specific) --- ## Prerequisites -### Hardware Requirements +### Hardware requirements - **AMD MI300X GPUs** with ROCm 6.0+ support - **Minimum GPUs:** - Flux 535M: 1x MI300X (testing) - Flux 12B: 2x MI300X with TP=2 (can train with FP8) -### Software Requirements +### Software requirements 1. **ROCm 6.0+** with FP8 tensor core support 2. **Transformer Engine 2.1.0+** with ROCm backend @@ -47,9 +47,9 @@ Verify your environment has: --- -## Quick Start +## Quick start -### Test FP8 with Flux 535M (Recommended First Step) +### Test FP8 with Flux 535M (recommended first step) ```bash # 1. Prepare test dataset (or use existing) @@ -61,7 +61,7 @@ GPUS_PER_NODE=1 \ bash examples/run_pretrain.sh ``` -### Production Training with Flux 12B +### Production training with Flux 12B ```bash # After validating with 535M, scale to 12B (TransformerEngine FP8) @@ -81,7 +81,7 @@ bash examples/run_slurm_pretrain.sh ## Configuration -### FP8 Model Configuration +### FP8 model configuration FP8 is configured at the model level. Two pre-configured files are available: @@ -113,7 +113,7 @@ fp8_dot_product_attention: false # Keep attention in higher precision fp8_multi_head_attention: false # Keep MHA in higher precision ``` -### Training Configuration Adjustments +### Training configuration adjustments **Batch Sizes with FP8:** @@ -148,9 +148,9 @@ context_parallel_size: 1 --- -## Performance Benchmarks +## Performance benchmarks -### Memory Usage +### Memory usage | Model | Precision | Memory/GPU | Batch Size | Notes | |-------|-----------|------------|------------|-------| @@ -159,7 +159,7 @@ context_parallel_size: 1 | Flux 12B | BF16 | ~40-50GB | 1 | TP=2 required | | Flux 12B | FP8 | ~20-25GB | 2 | TP=2, ~50% reduction | -### Training Speed +### Training speed | Model | Precision | Steps/sec | Speedup | Hardware | |-------|-----------|-----------|---------|----------| @@ -168,7 +168,7 @@ context_parallel_size: 1 | Flux 12B | BF16 | ~0.5-1.0 | 1.0x | 32x MI300X | | Flux 12B | FP8 | ~0.8-1.5 | 1.5-2x | 32x MI300X | -### Expected Results +### Expected results - **Memory:** ~50% reduction vs BF16 - **Speed:** 1.5-2x faster training @@ -179,7 +179,7 @@ context_parallel_size: 1 ## Troubleshooting -### NaN or Inf in Losses +### NaN or inf in losses **Problem:** Training becomes unstable with NaN/Inf values @@ -212,7 +212,7 @@ context_parallel_size: 1 num_layers_at_end_in_bf16: 2 ``` -### Out of Memory Even with FP8 +### Out of memory even with FP8 **Problem:** Still hitting OOM errors with FP8 enabled @@ -239,7 +239,7 @@ context_parallel_size: 1 seq_length: 2048 # if applicable ``` -### FP8 Not Available +### FP8 not available **Problem:** Setup script shows "FP8 not available" @@ -269,7 +269,7 @@ context_parallel_size: 1 print(te.fp8.is_fp8_available()) # Should be True ``` -### Slower Than Expected +### Slower than expected **Problem:** FP8 training is not faster than BF16 @@ -295,9 +295,9 @@ context_parallel_size: 1 --- -## Best Practices +## Best practices -### Recommended Workflow +### Recommended workflow 1. **Start with 535M:** - Validate FP8 works correctly @@ -314,7 +314,7 @@ context_parallel_size: 1 - Watch for numerical issues - Compare checkpoints with BF16 -### Training Configuration +### Training configuration **Conservative (stable):** ```yaml @@ -370,7 +370,7 @@ fp8_wgrad: true --- -## Autotune (Local Spec FP8) +## Autotune (local spec FP8) > **Scope:** This section covers the **local-spec** FP8 path (`PrimusTurboFloat8LocalSpecProvider`, no TransformerEngine), e.g. `flux_12b_ddp_energon_schnell_resample_local_spec_fp8.yaml`. The TransformerEngine prerequisites and checks elsewhere in this guide (`te.fp8.is_fp8_available()`, "FP8 Not Available") do **not** apply here -- this path quantizes via Primus Turbo directly. @@ -396,9 +396,9 @@ For MXFP4/FP4 + AITER with a tuned CSV, do the **opposite**: leave `PRIMUS_TURBO --- -## AMD MI300X Specific +## AMD MI300X specific -### Environment Variables +### Environment variables ```bash # Optional: set for better performance @@ -407,7 +407,7 @@ export NCCL_DEBUG=INFO # For debugging export HSA_ENABLE_SDMA=0 # Disable SDMA for stability ``` -### ROCm Optimization +### ROCm optimization 1. **HipBLASLt tuning:** ```bash @@ -426,7 +426,7 @@ export HSA_ENABLE_SDMA=0 # Disable SDMA for stability export HSA_OVERRIDE_GFX_VERSION=9.4.2 # For MI300X ``` -### Known Issues +### Known issues 1. **Transformer Engine ROCm support:** - Verify TE version supports ROCm FP8 @@ -444,7 +444,7 @@ export HSA_ENABLE_SDMA=0 # Disable SDMA for stability ## Testing -### Integration Test +### Integration test ```bash # Quick 100-step validation run @@ -453,7 +453,7 @@ GPUS_PER_NODE=1 \ bash examples/run_pretrain.sh ``` -### Convergence Test +### Convergence test 1. Train both BF16 and FP8 for 5000 steps 2. Compare loss curves (should be within 5%) @@ -468,7 +468,7 @@ To set expectations for external users, the precision/convergence claims in this codebase fall into two tiers: - **Backed by in-repo tests** (CI-runnable on supported hardware): structural and - convention checks — attention TE-vs-local-spec equivalence, RNG/seed + convention checks—attention TE-vs-local-spec equivalence, RNG/seed determinism, chimera init, VAE resample reproducibility, fused delayed-scale update, and MLPerf warmup FP8 state. - **Asserted, not yet backed by an in-repo test:** end-to-end *tensor parity* diff --git a/docs/backends/megatron/diffusion/mxfp4_training.md b/docs/04-technical-guides/diffusion-models/mxfp4_training.md similarity index 82% rename from docs/backends/megatron/diffusion/mxfp4_training.md rename to docs/04-technical-guides/diffusion-models/mxfp4_training.md index b0dfbf86d..d45f64f52 100644 --- a/docs/backends/megatron/diffusion/mxfp4_training.md +++ b/docs/04-technical-guides/diffusion-models/mxfp4_training.md @@ -1,4 +1,4 @@ -# MXFP4 Training for Flux Models +# MXFP4 training for Flux models Guide for training Flux diffusion models in **MXFP4** (E2M1 mantissa + E8M0 block-of-32 scales) on AMD MI355X GPUs using Primus's local-spec MXFP4 implementation backed by Primus-Turbo and AITER. @@ -6,20 +6,20 @@ Guide for training Flux diffusion models in **MXFP4** (E2M1 mantissa + E8M0 bloc MXFP4 stores activations and weights in 4-bit microscale floating-point with one E8M0 exponent shared per block of 32 elements. The Primus integration: -- Uses a **local spec** (`PrimusTurboMXFP4LocalSpecProvider`) with **no Transformer Engine dependency** — MXFP4 linear layers are self-contained autograd `Function`s that call Primus-Turbo's `gemm_fp4_impl` directly, so the path is `torch.compile`-friendly with minimal graph breaks. +- Uses a **local spec** (`PrimusTurboMXFP4LocalSpecProvider`) with **no Transformer Engine dependency**—MXFP4 linear layers are self-contained autograd `Function`s that call Primus-Turbo's `gemm_fp4_impl` directly, so the path is `torch.compile`-friendly with minimal graph breaks. - Keeps **attention, optimizer state / main params, and inter-rank communication in BF16**. Only the MMA inputs of the column- and row-parallel linears are quantized. - Supports two backward modes via `mxfp4_backward_precision`: pure **MXFP4** (default) or **FP8** hybrid (E5M2 backward with tensorwise scaling on HipBLASLt). - Dispatches the FP4 GEMM through Primus-Turbo's pluggable backend layer, which can route to either AITER (recommended for MI355X) or HipBLASLt. -## Table of Contents +## Table of contents - [Prerequisites](#prerequisites) -- [Quick Start](#quick-start) +- [Quick start](#quick-start) - [Configuration](#configuration) -- [Primus-Turbo Backend Selection](#primus-turbo-backend-selection) +- [Primus-Turbo backend selection](#primus-turbo-backend-selection) - [Tuned GEMMs](#tuned-gemms) - [Troubleshooting](#troubleshooting) -- [Verification Status](#verification-status) +- [Verification status](#verification-status) --- @@ -27,7 +27,7 @@ MXFP4 stores activations and weights in 4-bit microscale floating-point with one ### Hardware -- **AMD Instinct MI355X** (gfx950) with FP4 tensor-core support. The MXFP4 linear-layer modules assert `check_mxfp4_support()` at construction and will refuse to initialize on unsupported devices ([`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`](../../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). +- **AMD Instinct MI355X** (gfx950) with FP4 tensor-core support. The MXFP4 linear-layer modules assert `check_mxfp4_support()` at construction and will refuse to initialize on unsupported devices ([`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). - Single node (the local-spec layers require `tensor_model_parallel_size: 1`). ### Software @@ -38,7 +38,7 @@ MXFP4 stores activations and weights in 4-bit microscale floating-point with one --- -## Quick Start +## Quick start The verified MXFP4 config is `examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml`. Launch with the AITER backend and the pre-tuned GEMM CSV: @@ -61,7 +61,7 @@ The pre-tuned CSV is distributed via an internal tuned-config source (`tuned_gem ## Configuration -The relevant overrides in [`examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml`](../../../../examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml): +The relevant overrides in [`examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml`](../../../examples/megatron/configs/MI355X/diffusion/flux_12b_ddp_energon_schnell_resample_local_spec_mxfp4.yaml): ```yaml # MXFP4 precision @@ -85,13 +85,13 @@ gradient_accumulation_fusion: false | Knob | Values | Notes | |------|--------|-------| | `fp4` | `"mxfp4"` | Top-level switch to enable FP4. | -| `fp4_recipe` | `"mxfp4"` for this guide | Default in [`primus/configs/modules/megatron/trainer_base.yaml`](../../../../primus/configs/modules/megatron/trainer_base.yaml) is `nvfp4`; the MXFP4 config overrides it. | -| `mxfp4_backward_precision` | `"mxfp4"` or `"fp8"` | Exhaustive set (checked by branch in [`primus_turbo_mxfp4_local.py`](../../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). `"fp8"` uses E5M2 with tensorwise HipBLASLt for backward. | +| `fp4_recipe` | `"mxfp4"` for this guide | Default in [`primus/configs/modules/megatron/trainer_base.yaml`](../../../primus/configs/modules/megatron/trainer_base.yaml) is `nvfp4`; the MXFP4 config overrides it. | +| `mxfp4_backward_precision` | `"mxfp4"` or `"fp8"` | Exhaustive set (checked by branch in [`primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). `"fp8"` uses E5M2 with tensorwise HipBLASLt for backward. | | `mxfp4_gradient_stochastic_rounding` | `true` / `false` | Optional. Enables SR on FP4 gradient quantization. | --- -## Primus-Turbo Backend Selection +## Primus-Turbo backend selection The FP4 GEMM call is routed by `GEMMFP4KernelDispatcher` in `Primus-Turbo/primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py`. Backends are selected with the precision-scoped env var `PRIMUS_TURBO_GEMM_BACKEND` (declared in `Primus-Turbo/primus_turbo/common/constants.py`): @@ -141,7 +141,7 @@ shape is M:16384, N:9216, K:3072, found padded_M: 16384, N:9216, K:3072 is tuned shape is M:..., N:..., K:..., not found tuned config in /path/to/flux_12b.csv, will use default config! ``` -Any miss line means the CSV needs re-tuning for that shape — follow the runbook in `tuned_gemm_configs/README.md`. +Any miss line means the CSV needs re-tuning for that shape—follow the runbook in `tuned_gemm_configs/README.md`. ### First-run JIT compile @@ -157,7 +157,7 @@ The (M, N, K) shape is missing from your CSV. AITER will fall back to its compil ### Slow first iteration (~minutes), normal afterwards -Expected — the first call to a CK-based `a4w4_blockscale_*` kernel triggers JIT compilation. Cached `.so` files are reused on subsequent starts. +Expected—the first call to a CK-based `a4w4_blockscale_*` kernel triggers JIT compilation. Cached `.so` files are reused on subsequent starts. ### `User specified backend AITER cannot handle the given inputs` @@ -171,7 +171,7 @@ Workaround: switch to `PRIMUS_TURBO_GEMM_BACKEND=FP4:HIPBLASLT` for unsupported ### `MXFP4ColumnParallelLinear requires tensor_model_parallel_size=1` -The MXFP4 linear-layer modules assert on `tensor_model_parallel_size == 1`, `gradient_accumulation_fusion == False`, and `sequence_parallel == False` ([`primus_turbo_mxfp4_local.py`](../../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). Adjust the config accordingly. +The MXFP4 linear-layer modules assert on `tensor_model_parallel_size == 1`, `gradient_accumulation_fusion == False`, and `sequence_parallel == False` ([`primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py)). Adjust the config accordingly. ### NaN losses @@ -185,7 +185,7 @@ If NaNs persist, also try `mxfp4_gradient_stochastic_rounding: true`. --- -## Verification Status +## Verification status The public config has been smoke-tested end-to-end: 1000 iters on 8x MI355X (single node, micro-batch 64 / global 512, sequence length 512) completes in ~16-20 minutes with `PRIMUS_TURBO_GEMM_BACKEND=FP4:AITER` and the tuned CSV. No errors across ranks; `pretrain() completed successfully`. @@ -193,18 +193,18 @@ Formal A/B benchmarks vs BF16 and FP8 (delayed and tensorwise) are pending and w --- -## Source Code Pointers +## Source code pointers -- MXFP4 spec provider: [`primus/backends/megatron/core/extensions/primus_turbo_local_spec.py`](../../../../primus/backends/megatron/core/extensions/primus_turbo_local_spec.py) (`PrimusTurboMXFP4LocalSpecProvider`). -- MXFP4 linear-layer autograd / fwd-bwd: [`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`](../../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py). -- Config schema defaults: [`primus/configs/modules/megatron/trainer_base.yaml`](../../../../primus/configs/modules/megatron/trainer_base.yaml). -- Dataclass field `mxfp4_backward_precision`: [`primus/backends/megatron/core/models/diffusion/common/config.py`](../../../../primus/backends/megatron/core/models/diffusion/common/config.py). +- MXFP4 spec provider: [`primus/backends/megatron/core/extensions/primus_turbo_local_spec.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_local_spec.py) (`PrimusTurboMXFP4LocalSpecProvider`). +- MXFP4 linear-layer autograd / fwd-bwd: [`primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py`](../../../primus/backends/megatron/core/extensions/primus_turbo_mxfp4_local.py). +- Config schema defaults: [`primus/configs/modules/megatron/trainer_base.yaml`](../../../primus/configs/modules/megatron/trainer_base.yaml). +- Dataclass field `mxfp4_backward_precision`: [`primus/backends/megatron/core/models/diffusion/common/config.py`](../../../primus/backends/megatron/core/models/diffusion/common/config.py). - FP4 backend selection (Primus-Turbo): `primus_turbo/common/constants.py`, `primus_turbo/pytorch/core/backend.py`, `primus_turbo/pytorch/kernels/gemm/gemm_fp4_impl.py`. - AITER tuned-config loader: `aiter/jit/core.py` (`AITER_CONFIG_GEMM_A4W4`). - AITER A4W4 dispatch + hit/miss logging: `aiter/ops/gemm_op_a4w4.py`. -## Related Documentation +## Related documentation -- [FP8 Training Guide](fp8_training.md) — companion guide for FP8. +- [FP8 Training Guide](fp8_training.md)—companion guide for FP8. - [Diffusion Architecture / Developer Guide](README.md). -- [Diffusion Examples README](../../../../examples/megatron/diffusion/README.md). +- [Diffusion Examples README](../../../examples/megatron/diffusion/README.md). diff --git a/docs/04-technical-guides/fault-tolerance-and-elastic-training.md b/docs/04-technical-guides/fault-tolerance-and-elastic-training.md new file mode 100644 index 000000000..1c107f072 --- /dev/null +++ b/docs/04-technical-guides/fault-tolerance-and-elastic-training.md @@ -0,0 +1,123 @@ +# Fault tolerance and elastic training + +Large-scale jobs run for days across thousands of GPUs, where hardware faults, NIC flaps, and node loss are routine. This guide covers the mechanisms Primus exposes to survive and recover from failures: graceful exit + checkpoint-based resume, Megatron's fault-tolerance package and in-process restart, and TorchTitan's [torchft](https://github.com/pytorch/torchft)-based elastic training. Parameters are grounded in `primus/configs/modules/megatron/trainer_base.yaml`, `primus_megatron_module.yaml`, and `primus/configs/modules/torchtitan/pre_trainer.yaml`. + +The foundation of all recovery is checkpointing—read [Checkpoint management](./checkpoint-management.md) first. + +--- + +## 1. The recovery model + +There are three layers, from simplest to most advanced: + +1. **Checkpoint + restart**—periodically save state; on failure, relaunch the job and resume from the last checkpoint. Works on every backend; relies on the scheduler (Slurm `--requeue`, Kubernetes restart policy) to relaunch. +2. **Graceful exit**—detect a signal or time/iteration budget, save a final checkpoint, and exit cleanly so the restart resumes with no lost work. +3. **In-job fault tolerance / elastic**—detect a failed rank and restart in-process (Megatron) or continue with a reduced/replaced replica group (TorchTitan + torchft) without tearing down the whole job. + +--- + +## 2. Graceful exit and auto-resume (Megatron) + +Controls in `trainer_base.yaml` let a run stop cleanly at a boundary so the next launch resumes seamlessly: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `exit_signal_handler` | `false` | Install a signal handler that saves a checkpoint and exits gracefully on SIGTERM (e.g. Slurm preemption). | +| `exit_duration_in_mins` | `null` | Exit (after saving) once the job has run this many minutes—useful to fit scheduler time limits. | +| `exit_interval` | `null` | Exit after this many iterations. | +| `adlr_autoresume` | `false` | Enable ADLR auto-resume integration. | +| `adlr_autoresume_interval` | `1000` | Iterations between auto-resume checks. | + +Primus-level continuation (`primus_megatron_module.yaml`): + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `auto_continue_train` | `false` | Automatically continue from the latest checkpoint in the experiment's save directory on relaunch. | +| `disable_last_saving` | `false` | Disable the final end-of-run checkpoint save (leave `false` so resume points exist). | + +**Pattern:** enable `exit_signal_handler` + `exit_duration_in_mins` (or rely on preemption signals), set a reasonable checkpoint `save_interval`, and turn on `auto_continue_train` so requeued jobs pick up where they left off. + +--- + +## 3. Megatron fault-tolerance package and in-process restart + +Megatron integrates an optional fault-tolerance package and in-process restart (`trainer_base.yaml`): + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `enable_ft_package` | `false` | Enable the Megatron fault-tolerance package (rank monitoring / heartbeat). | +| `calc_ft_timeouts` | `false` | Auto-calculate fault-tolerance timeouts from observed step times. | +| `run_workload_inspector_server` | `false` | Run the workload inspector server for health/diagnostics. | +| `inprocess_restart` | `false` | Restart failed ranks **in process** to avoid a full job teardown. | + +In-process restart reduces recovery time by re-initializing the process group and reloading state without re-scheduling the whole allocation. Combine with frequent checkpoints so the restarted ranks have a recent resume point. + +--- + +## 4. Numerical safety nets (Megatron) + +Detecting corruption early prevents wasted compute and divergence (`trainer_base.yaml`): + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `check_for_nan_in_loss_and_grad` | `true` | Abort/handle on NaN/Inf in loss or gradients. | +| `check_for_spiky_loss` | `false` | Detect anomalous loss spikes. | +| `check_for_large_grads` | `false` | Detect abnormally large gradients. | +| `decrease_batch_size_if_needed` | `false` | Reduce batch size when needed instead of failing. | + +These don't recover from hardware faults but stop a corrupted run before it pollutes downstream checkpoints. + +--- + +## 5. Elastic training with torchft (TorchTitan) + +TorchTitan supports semi-synchronous, replica-based fault tolerance via [torchft](https://github.com/pytorch/torchft). Configured under `fault_tolerance:` in `primus/configs/modules/torchtitan/pre_trainer.yaml`: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `enable` | `false` | Enable torchft fault tolerance. | +| `process_group` | `gloo` | Process group backend for fault-tolerance coordination. | +| `process_group_timeout_ms` | `10000` | Coordination timeout (ms). | +| `replica_id` | `0` | This replica's ID. | +| `group_size` | `0` | Number of replica groups (`0` = auto). | +| `min_replica_size` | `1` | Minimum replicas required to keep training. | +| `semi_sync_method` | `null` | Semi-synchronous algorithm (e.g. DiLoCo-style), `null` = standard. | + +With replica groups, the loss of one replica can be tolerated as long as `min_replica_size` is still satisfied—training continues while the failed replica recovers/rejoins, rather than crashing the whole job. + +**Install** the optional dependencies before enabling (`requirements-torchft.txt`): + +```bash +pip install -r requirements-torchft.txt # torchft-nightly + OpenTelemetry exporters +``` + +TorchTitan also exposes communication timeouts under `comm:` (`init_timeout_seconds: 300`, `train_timeout_seconds: 100`) that govern how long collectives wait before declaring a fault. + +--- + +## 6. Scheduler integration + +In-job mechanisms still need the scheduler to relaunch on full-job failure: + +- **Slurm**—submit with `--requeue` so preempted/failed jobs are re-queued; pair with `exit_signal_handler` to checkpoint on SIGTERM. See [Deployment](../05-operations/deployment.md). +- **Kubernetes**—use a restart policy / operator that recreates pods; mount checkpoint storage on a shared/persistent volume. +- **Shared checkpoint storage**—all ranks must read the same checkpoint directory after relaunch (NFS, Lustre, or object storage). See [Checkpoint management](./checkpoint-management.md). + +--- + +## 7. Recommended setup + +1. **Always checkpoint**—set a `save_interval` matched to your mean-time-between-failures; use async/distributed checkpointing to keep overhead low. +2. **Exit cleanly**—`exit_signal_handler: true` (+ `exit_duration_in_mins` for time-boxed allocations). +3. **Resume automatically**—`auto_continue_train: true` (Megatron) and `--requeue` (Slurm). +4. **Reduce recovery time at scale**—`enable_ft_package` + `inprocess_restart` (Megatron) or torchft replica groups (TorchTitan). +5. **Guard numerics**—keep `check_for_nan_in_loss_and_grad` on; consider spiky/large-grad checks for unstable configs. + +--- + +## Related documentation + +- [Checkpoint management](./checkpoint-management.md)—save/load formats, async and distributed checkpointing. +- [Deployment](../05-operations/deployment.md)—Slurm/Kubernetes restart and requeue. +- [Multi-node networking](./multi-node-networking.md)—NIC faults and collective timeouts. +- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) and [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md). diff --git a/docs/04-technical-guides/logging-and-experiment-tracking.md b/docs/04-technical-guides/logging-and-experiment-tracking.md new file mode 100644 index 000000000..4a9b2e2be --- /dev/null +++ b/docs/04-technical-guides/logging-and-experiment-tracking.md @@ -0,0 +1,163 @@ +# Logging and experiment tracking + +This guide covers how Primus emits training metrics and logs, and how to wire up the supported experiment trackers—**TensorBoard**, **Weights & Biases (WandB)**, and **MLflow** (including Databricks-hosted MLflow)—across the Megatron and TorchTitan backends. Parameters are grounded in `primus/configs/modules/megatron/trainer_base.yaml`, `primus_megatron_module.yaml`, and `primus/configs/modules/torchtitan/pre_trainer.yaml`. + +For an operations-oriented overview, see [Monitoring and logging](../05-operations/monitoring-logging.md). For required credentials/keys, see [Environment variables](../03-configuration-reference/environment-variables.md). + +--- + +## 1. Tracker toggles at a glance (Megatron) + +All three trackers are **opt-in** and disabled by default (`primus/configs/modules/megatron/primus_megatron_module.yaml`): + +```yaml +disable_tensorboard: true +disable_wandb: true +disable_mlflow: true +``` + +Set the relevant `disable_*` to `false` to enable a tracker. Primus performs sanity checks at startup—e.g. it warns if WandB is enabled but `WANDB_API_KEY` is unset (`primus/backends/megatron/patches/args/wandb_config_patches.py`). MLflow logging is initialized in `primus/backends/megatron/training/global_vars.py`; Databricks-hosted MLflow additionally requires `DATABRICKS_HOST` (read by the `mlflow` client). + +--- + +## 2. Console / step logging (Megatron) + +Core logging cadence and content (`trainer_base.yaml`, overridden by `pre_trainer.yaml`): + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `log_interval` | `100` (`pre_trainer.yaml` sets `1`) | Steps between log lines. | +| `log_throughput` | `false` (`pre_trainer.yaml` sets `true`) | Log tokens/s and TFLOP/s throughput. | +| `log_progress` | `false` | Log progress/ETA. | +| `log_params_norm` | `false` | Log parameter L2 norm. | +| `log_num_zeros_in_grad` | `false` | Log gradient sparsity. | +| `log_avg_skip_iterations` | `2` | Warmup iterations excluded from averages. | +| `log_avg_reset_interval` | `10` | Reset window for running averages. | +| `timing_log_level` | `0` | Verbosity of timer breakdowns. | +| `timing_log_option` | `minmax` | Timer aggregation across ranks. | +| `logging_level` | `null` | Python logging level override. | + +--- + +## 3. TensorBoard (Megatron) + +Enable with `disable_tensorboard: false` and set an output directory: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `tensorboard_dir` | `null` | Output directory for event files (**required** when enabled). | +| `tensorboard_log_interval` | `1` | Steps between TensorBoard writes. | +| `tensorboard_queue_size` | `1000` | Event queue size before flush. | +| `log_learning_rate_to_tensorboard` | `true` | Log LR. | +| `log_loss_scale_to_tensorboard` | `true` | Log loss scale (mixed precision). | +| `log_timers_to_tensorboard` | `false` | Log per-stage timers. | +| `log_batch_size_to_tensorboard` | `false` | Log batch size. | +| `log_memory_to_tensorboard` | `false` | Log GPU memory. | +| `log_world_size_to_tensorboard` | `false` | Log world size. | +| `log_validation_ppl_to_tensorboard` | `false` | Log validation perplexity. | + +--- + +## 4. Weights and biases (Megatron) + +Enable with `disable_wandb: false`. Configuration: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `wandb_project` | `null` | WandB project name. | +| `wandb_exp_name` | `null` | Run/experiment name. | +| `wandb_save_dir` | `null` | Local directory for WandB files. | +| `wandb_entity` | `null` | Team/entity. | + +**Credentials** (see [Environment variables](../03-configuration-reference/environment-variables.md)): + +```bash +export WANDB_API_KEY=... # required when WandB is enabled +export WANDB_PROJECT=... # optional +export WANDB_RUN_NAME=... # optional +export WANDB_TEAM=... # optional (entity) +``` + +`WANDB_API_KEY` is on the container passthrough whitelist (`runner/.primus.yaml`), so it propagates into the training container. + +--- + +## 5. MLflow (Megatron) + +Enable with `disable_mlflow: false`. Run identification and upload behavior: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `mlflow_run_name` | `null` | MLflow run name. | +| `mlflow_experiment_name` | `null` | MLflow experiment name. | +| `mlflow_upload_traces` | `false` | Upload profiler trace files. | +| `mlflow_upload_logs` | `false` | Upload training log files. | +| `mlflow_upload_performance_metrics` | `false` | Upload the comprehensive perf/memory/utilization metric set (implicitly enables throughput calc). | +| `mlflow_upload_tracelens_report` | `false` | Generate + upload TraceLens reports (see [Profiling & observability](./profiling-and-observability.md)). | + +**Credentials / endpoints:** + +```bash +export MLFLOW_TRACKING_URI=... # tracking server URI +export MLFLOW_REGISTRY_URI=... # optional model registry +# Databricks-hosted MLflow: +export DATABRICKS_HOST=... # checked at startup when MLflow is enabled +export DATABRICKS_TOKEN=... +``` + +> The `mlflow_upload_*` flags are designed so MLflow stays opt-in: they only take effect when `disable_mlflow: false`. + +--- + +## 6. One-logger (Megatron) + +NVIDIA One-Logger telemetry is enabled by default in `trainer_base.yaml`: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `enable_one_logger` | `true` | Enable One-Logger collection. | +| `one_logger_project` | `megatron-lm` | Project tag. | +| `one_logger_run_name` | `null` | Run name. | +| `one_logger_async` | `false` | Async upload. | +| `app_tag_run_name` / `app_tag_run_version` | `null` / `0.0.0` | Application tags. | + +--- + +## 7. Metrics and logging (TorchTitan) + +Configured under `metrics:` in `primus/configs/modules/torchtitan/pre_trainer.yaml`: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `enable_tensorboard` | `false` | Enable TensorBoard logging. | +| `enable_wandb` | `false` | Enable WandB logging. | +| `log_freq` | `10` | Steps between metric logs. | +| `save_tb_folder` | `tb` | TensorBoard output subfolder. | +| `save_for_all_ranks` | `false` | Write metrics from every rank (default: rank 0 only). | +| `disable_color_printing` | `false` | Disable ANSI colors in console output. | + +TorchTitan reads WandB settings from the environment (`WANDB_PROJECT`, `WANDB_RUN_NAME`, `WANDB_TEAM`) via `primus/backends/torchtitan/patches/wandb_patches.py` and `third_party/torchtitan/torchtitan/components/metrics.py`. + +--- + +## 8. Logging (MaxText) + +MaxText logging cadence is controlled by `log_period` (`primus/configs/modules/maxtext/pre_trainer.yaml`, default `100`). See [MaxText parameters](../03-configuration-reference/maxtext-parameters.md). + +--- + +## 9. Recommended setup + +1. **Local-only:** enable TensorBoard (`disable_tensorboard: false`, set `tensorboard_dir`)—no credentials required. +2. **Team tracking:** enable WandB (`disable_wandb: false`) + export `WANDB_API_KEY` and `wandb_project`/`wandb_entity`. +3. **Enterprise / scaling studies:** enable MLflow (`disable_mlflow: false`) + `MLFLOW_TRACKING_URI` (or Databricks host/token), and turn on `mlflow_upload_performance_metrics` for throughput/memory/utilization dashboards. +4. **Keep `WANDB_API_KEY` and tokens out of YAML**—pass them as environment variables (whitelisted for container passthrough). See [Security](../05-operations/security.md). + +--- + +## Related documentation + +- [Monitoring and logging](../05-operations/monitoring-logging.md)—operational view of trackers. +- [Profiling & observability](./profiling-and-observability.md)—traces, TraceLens, perf metrics. +- [Environment variables](../03-configuration-reference/environment-variables.md)—credentials and passthrough. +- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) and [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md). diff --git a/docs/04-technical-guides/moe-training.md b/docs/04-technical-guides/moe-training.md new file mode 100644 index 000000000..b31af0b2c --- /dev/null +++ b/docs/04-technical-guides/moe-training.md @@ -0,0 +1,193 @@ +# MoE training deep-dive + +This guide covers Mixture-of-Experts (MoE) training in Primus on AMD Instinct GPUs: the bottlenecks unique to sparse models, the Primus/Primus-Turbo optimizations that address them, and a model-by-model tuning walkthrough. It is adapted from the AMD blog [MoE Training Best Practices on AMD GPU](https://rocm.blogs.amd.com/software-tools-optimization/primus-moe-package/README.html) (`examples/moe_package/README.md`) and grounded in the actual Primus configs and run scripts. + +All flags shown here are the **real CLI/YAML keys** used by `examples/moe_package/run_*_pretrain_mi355x.sh` and the Megatron module configs (`primus/configs/modules/megatron/`). The Primus-Turbo MoE optimizations in this guide (DeepEP, sync-free MoE, Turbo grouped GEMM) are **Megatron-backend** features. TorchTitan also supports MoE via expert parallelism (`expert_parallel_degree`, `expert_tensor_parallel_degree`), but its tuning is out of scope here. + +--- + +## 1. Why MoE training is different + +MoE scales model capacity by routing each token through a small subset of "expert" sub-networks instead of activating the whole network. A gating/router picks the top-k experts per token, so a model can hold many billions of parameters while only a fraction are active per token. + +This sparsity creates performance challenges that dense models do not have: + +- **Grouped GEMM overhead**—each expert is a separate GEMM; naive multi-stream execution leaves scheduling gaps. +- **All-to-all (A2A) communication**—token dispatch/combine across expert-parallel ranks can dominate runtime, especially with `EP >= 8` and multi-node. +- **CPU sync & launch delays**—dynamic shapes (token counts per expert) force device-to-host syncs that stall the kernel launch queue. +- **Too many small kernels**—fine-grained MoE ops stress the CPU launch path. +- **Pipeline load imbalance**—uneven layer distribution across pipeline stages quietly degrades throughput. +- **Memory pressure**—activations dominate memory at large scale, forcing recomputation. + +--- + +## 2. Representative model configs + +Primus ships Megatron model presets for DeepSeek-style MoE models plus two ultra-large research configs: + +| Model | Total / Active params | Model config (`primus/configs/models/megatron/`) | +|-------|-----------------------|--------------------------------------------------| +| DeepSeek-V2-Lite | 16B / 2.4B | `deepseek_v2_lite.yaml` | +| DeepSeek-V2 | 236B / 21B | `deepseek_v2.yaml` | +| DeepSeek-V3 | 671B / 37B | `deepseek_v3.yaml` | +| MoE-1T | 1T / 44B | `moe_1T.yaml` | +| MoE-2T | 2T / 80B | `moe_2T.yaml` | + +Ready-to-run pretrain scripts live in `examples/moe_package/`, e.g.: + +- `examples/moe_package/run_deepseek_v2_lite_pretrain_mi355x.sh` +- `examples/moe_package/run_deepseek_v2_pretrain_mi355x.sh` +- `examples/moe_package/run_deepseek_v3_pretrain_mi355x.sh` + +Each example script is a convenience wrapper: it sets environment + parallelism and selects an experiment YAML under `examples/moe_package/configs/`, then launches training. You can run the same training directly with the unified CLI, passing the experiment YAML with `--config` and the MoE feature toggles (Section 4) as overrides: + +```bash +# DeepSeek-V2-Lite baseline + DeepEP + sync-free + loss fusion + manual GC, via primus-cli +export ENABLE_NUMA_BINDING=1 HSA_KERNARG_POOL_SIZE=12582912 # feature 6 (env, not CLI flags) +./runner/primus-cli direct -- train pretrain \ + --config examples/moe_package/configs/MI355X/deepseek_v2_lite-pretrain-baseline.yaml \ + --enable_primus_turbo True \ + --use_turbo_deepep True --turbo_deepep_num_cu 64 --moe_router_dtype fp32 \ + --turbo_sync_free_moe_stage 1 \ + --cross_entropy_fusion_impl te --cross_entropy_loss_fusion True \ + --manual_gc True --manual_gc_interval 1 +``` + +Use `./runner/primus-cli slurm srun -N -- train pretrain --config ...` for multi-node. The feature tables below list the exact flags so you can compose your own command. (Optimizations that are environment variables—NUMA binding, `HSA_KERNARG_POOL_SIZE`, UCCL-EP—are exported before the command rather than passed as `--flags`.) + +--- + +## 3. Profiling and analysis workflow + +Diagnose before optimizing. The recommended order: + +1. **Torch Profiler**—capture operator times, memory, and GPU utilization. Enable through the Megatron profiling flags (`--profile`, `--use_pytorch_profiler`, `--profile_step_start`, `--profile_step_end`, `--disable_profiler_activity_cpu`). Load the trace in [Perfetto](https://ui.perfetto.dev/) to inspect CPU/GPU overlap, launch delays, and idle gaps. See [Profiling & observability](./profiling-and-observability.md). +2. **TraceLens**—AMD's automated trace analyzer for hierarchical breakdowns, roofline/efficiency, communication-vs-sync separation, and trace diffing. Wired into Primus via `generate_tracelens_report` / `mlflow_upload_tracelens_report` (see the profiling guide). +3. **Memory projection**—model VRAM across params, gradients, activations, and optimizer state *before* launching, via `./primus-cli direct -- projection memory --config .yaml`. See [Projection](../02-user-guide/projection.md). +4. **Pipeline visualization**—dump pipeline schedule data (`--dump_pp_data true`) and render stage utilization with `tools/visualization/pp_vis/vis.py` to find bubbles and stage imbalance. + +--- + +## 4. Primus MoE optimizations + +The `examples/moe_package/run_*` scripts expose these as composable "MoE features." The table maps each feature to the **actual `--flags` (or environment variables)** you pass to `train pretrain`—the same toggles the example scripts set. + +| Feature | Flags (real keys) | What it does | +|---------|-------------------|--------------| +| Turbo attention | `--enable_primus_turbo True --use_turbo_attention True` | Optimized attention kernels (Primus-Turbo). | +| Turbo grouped GEMM | `--enable_primus_turbo True --use_turbo_grouped_gemm True` | Fused CK grouped GEMM processes all experts in one launch instead of multi-stream. | +| Loss fusion | `--cross_entropy_fusion_impl te --cross_entropy_loss_fusion True` | Fuses large-vocab loss into one kernel to cut memory + launch overhead. | +| DeepEP acceleration | `--enable_primus_turbo True --use_turbo_deepep True --turbo_deepep_num_cu 64 --turbo_deepep_use_comm_stream False --moe_shared_expert_overlap False --moe_router_dtype fp32` | GPU-side index calc + sync-free dispatch to cut redundant cross-node A2A traffic. | +| Sync-free MoE | `--enable_primus_turbo True --turbo_sync_free_moe_stage ` | Removes CPU D2H syncs across Router → Dispatcher → Permutation → GroupMLP. | +| NUMA binding | `export ENABLE_NUMA_BINDING=1` | Pins each GPU process to its NUMA socket for better memory bandwidth/stability. | +| HIP kernarg pool | `export HSA_KERNARG_POOL_SIZE=12582912` | Enlarges the kernel-argument pool (12 MB) to avoid launch stalls under many small kernels. | +| Manual GC | `--manual_gc True --manual_gc_interval 1` | Periodic host GC to remove iteration-time jitter on long runs. | +| UCCL-EP | `export USING_UEP=1` | Use the UCCL transport for DeepEP dispatch/combine (sets `PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND=DEEP_EP` + UCCL network env). Requires the `uccl` and `deep_ep` packages. | + +> **Grouped GEMM keys.** Enable Turbo grouped GEMM for MoE with `use_turbo_grouped_gemm` (`--use_turbo_grouped_gemm True`). The older `use_turbo_grouped_mlp` alias has been **removed**—passing it now raises an assertion error (`use_turbo_grouped_mlp has been removed; please use use_turbo_grouped_gemm instead`). +> +> **Legacy path.** The legacy multi-stream grouped GEMM path is selected with `--moe_use_legacy_grouped_gemm True` (the scripts default `LEGACY_GG=True`). Turbo grouped GEMM is **incompatible** with the legacy path—set `--moe_use_legacy_grouped_gemm False` whenever `use_turbo_grouped_gemm` is enabled (Primus raises an error otherwise). + +### Sync-free MoE stages + +`turbo_sync_free_moe_stage` is a single knob with four levels (`0`–`3`, validated in `primus/backends/megatron/patches/args/rocm_arg_validation.py`). Each stage **auto-enables** a set of fusion flags: + +| Level | Auto-enabled flags | Behavior | +|-------|--------------------|----------| +| `0` (default) | — | Disabled—standard baseline. | +| `1` | `moe_use_fused_router_with_aux_score`, `moe_permute_fusion` | Sync-free **Router** + **Permutation** fusion. | +| `2` | stage 1 + `use_turbo_deepep`, `use_turbo_grouped_gemm` | Adds sync-free **DeepEP** dispatch and **Turbo grouped GEMM**. | +| `3` | stage 2 + `use_turbo_fused_act_with_probs` | Full sync-free pipeline (adds fused activation). Per the blog this **uses significantly more GPU memory**—only enable with headroom. | + +Requirements (enforced at startup): + +- All stages require `--enable_primus_turbo True`. +- Stages `2` and `3` require Turbo grouped GEMM and are therefore **incompatible with `--moe_use_legacy_grouped_gemm True`**. + +Practical guidance from the run scripts: + +- **MI355X:** `--turbo_sync_free_moe_stage 1` (compatible with the default legacy grouped GEMM path, since stage 1 does not enable Turbo grouped GEMM). +- **MI300X / MI325X:** the example scripts suggest stage `2` (with `--moe_shared_expert_overlap False --moe_router_dtype fp32`). Because stage 2 auto-enables Turbo grouped GEMM, you must also set `--moe_use_legacy_grouped_gemm False`. + +### Scheduling and memory features + +- **1F1B A2A overlap**—interleaves micro-batch N's expert communication with micro-batch N-1's backward compute on top of interleaved-1F1B pipeline parallelism, hiding A2A behind compute while preserving the bubble rate and roughly the same peak memory. +- **Arbitrary pipeline partition**—manual stage layout instead of automatic even splits, to balance per-stage memory/compute. Use the Megatron-core `--pipeline_model_parallel_layout` flag (as the DeepSeek-V3 script does) or the Primus `decoder_pipeline_manual_split_list` config key (`primus/configs/modules/megatron/primus_megatron_module.yaml`). +- **Selective layer recompute**—recompute specific transformer layers with `--recompute_layer_ids 0,1,2,3` (keep `RECOMPUTE_LAYERS=0` so this is the only recompute control), or full block recompute via `--recompute_granularity full --recompute_method block --recompute_num_layers N`. +- **MoE expert-parallel comm overlap**—`overlap_moe_expert_parallel_comm: true` (`trainer_base.yaml`). + +See [Performance tuning](./performance-tuning.md) for the full Primus-Turbo flag reference. + +--- + +## 5. Model-specific tuning + +### DeepSeek-V2-Lite (16B / 2.4B, 27 layers) + +A compute/memory-efficient variant ideal for high-throughput pretraining. AMD Instinct's large HBM (192 GB on MI300X, 288 GB on MI355X) lets you push **micro-batch size (MBS)** high to maximize throughput. + +Recommended optimization stack (matches `run_deepseek_v2_lite_pretrain_mi355x.sh`, where `MoE_Features=(3 4 5 6 7 8)`): + +1. Manual GC for stable iteration time. +2. Loss fusion for the large vocabulary. +3. DeepEP for A2A. +4. Sync-free mode (stage 1 on MI355X) to remove D2H syncs. +5. NUMA binding for CPU affinity (`ENABLE_NUMA_BINDING=1` + `HSA_KERNARG_POOL_SIZE`). +6. MBS scaling using the memory freed by the above (the blog reports peak memory dropping from ~99.8% to ~84.3% at MBS=12, enabling MBS=14). + +The script's default `MoE_Features=(3 4 5 6 7 8)` also enables feature `8` = UCCL-EP (see the feature table above). Default parallelism in the script: `TP=1 ETP=1 PP=1 EP=8 CP=1`, `MBS=14 GBS=896 SEQ=4096`. + +### DeepSeek-V2 (236B / 21B, 60 layers) + +Scale up with parallelism for max throughput across nodes. Recommended stack: + +1. Manual GC, 2) Loss fusion, 3) DeepEP, 4) NUMA binding, 5) Sync-free mode, plus **interleaved pipeline parallelism (VPP)** to cut the pipeline bubble ratio. Enabling VPP (`--num_virtual_stages_per_pipeline_rank > 1`) also improves sync-free mode effectiveness. + +Default parallelism in the script: `TP=1 ETP=1 PP=4 VPP=5 EP=8 CP=1` (interleaved PP), `SEQ=4096`. + +### 1T+ parameter models (MoE-1T / MoE-2T, 96 layers) + +Ultra-large training combines every advanced technique. Use **memory projection first**—at this scale **activations dominate memory**, not parameters/optimizer state. + +Findings from the blog's projections (768–1024 GPUs): + +- **Context parallelism (CP2)** roughly halves activation memory (~76 GB/GPU saved for 1T, ~131 GB/GPU for 2T)—the most effective single lever. +- **Increasing EP** (8→16) barely reduces memory but adds A2A time. +- **Increasing PP** (24→48) doesn't materially cut memory and raises pipeline bubbles + activation memory. + +Suggested configs: + +| Model | MI300X | MI355X | +|-------|--------|--------| +| MoE-1T | PP24 EP8 CP2 | PP24 EP8 (no checkpointing) | +| MoE-2T | PP24 EP16 CP2 | PP24 EP8 CP2 (benefits from larger DP) | + +**Pipeline bubble at scale.** When the global batch is constrained, gradient-accumulation (GA) per iteration drops and the bubble ratio rises. Interleaved PP (VPP) mitigates this: + +$$ +\text{bubble ratio} = \frac{PP-1}{(PP-1) + GA \times VPP} +$$ + +For PP=16, GA=16: VPP=1 gives ~48% bubble; VPP=6 gives ~14%—a large efficiency win verified on a 64-node setup. + +**Inter-node dispatch.** Profiling on a 2T/1024-GPU run showed A2A consuming **25–30%** of step time; DeepEP delivered roughly **1.05×–7.66×** end-to-end speedup over plain A2A and kept EP scaling nearly flat. + +--- + +## 6. Quick checklist + +1. **Profile first**—Torch Profiler + TraceLens; project memory before launching ultra-large runs. +2. **Turn on Turbo**—`--enable_primus_turbo True`, then grouped GEMM, attention, DeepEP as needed (requires the external `primus_turbo` package). +3. **Kill CPU syncs**—`--turbo_sync_free_moe_stage` (1 on MI355X, 2 on MI300X/MI325X). +4. **Stabilize + bind**—`--manual_gc True`, `ENABLE_NUMA_BINDING=1`, `HSA_KERNARG_POOL_SIZE=12582912`. +5. **Scale memory headroom into throughput**—raise MBS; use CP2 + selective recompute for 1T+; use VPP to cut pipeline bubbles. + +--- + +## Related documentation + +- [Performance tuning](./performance-tuning.md)—Primus-Turbo flags, HipBLASLt, precision, recompute. +- [Parallelism strategies](./parallelism-strategies.md) and [Parallelism configuration](./parallelism-configuration.md)—EP, PP, CP, VPP. +- [Collective operations](./collective-operations.md)—A2A and DeepEP context. +- [Profiling & observability](./profiling-and-observability.md) and [Projection](../02-user-guide/projection.md). +- Source blog: `examples/moe_package/README.md`. diff --git a/docs/04-technical-guides/multi-node-networking.md b/docs/04-technical-guides/multi-node-networking.md new file mode 100644 index 000000000..62a5357a3 --- /dev/null +++ b/docs/04-technical-guides/multi-node-networking.md @@ -0,0 +1,203 @@ +# Multi-node networking guide + +Multi-node training depends on **high-bandwidth, low-latency** communication between GPUs. On AMD systems, **RCCL** (ROCm Collective Communications Library) provides GPU collectives with an API aligned to **NCCL**, so most **NCCL-prefixed** environment variables apply to RCCL as well. + +This guide summarizes how Primus configures networking, how **InfiniBand**, **RoCE**, and **AINIC (AMD AI NIC)** fit in, and how to validate and troubleshoot cluster connectivity. + +**Primary sources in this repository** + +| Topic | File | +|-------|------| +| Default NCCL/RCCL and socket setup | `runner/helpers/envs/base_env.sh` | +| IB HCA detection | `runner/helpers/envs/get_nccl_ib_hca.sh` | +| Socket / interface detection | `runner/helpers/envs/get_ip_interface.sh` | +| AINIC hook (container/CLI integration) | `runner/helpers/hooks/03_enable_ainic.sh` | +| AINIC CLI defaults | `runner/use_ainic.yaml` | +| ANP / `NCCL_NET_PLUGIN` example | `examples/run_pretrain.sh` | + +--- + +## 1. Overview + +- **Goal:** Keep gradient and parameter exchanges from becoming the bottleneck when scaling across nodes. +- **Stack:** PyTorch distributed uses the ROCm **NCCL** backend name in many configs; the implementation is **RCCL** on AMD GPUs. +- **Transports:** Common fabrics include **InfiniBand (IB)**, **RoCE** (RDMA over Converged Ethernet), and **AINIC** on supported AMD platforms. Primus scripts set or detect **HCAs**, **socket interfaces**, and optional **AINIC** tuning. + +--- + +## 2. InfiniBand configuration + +These variables are standard in NCCL/RCCL deployments. Primus seeds several from `runner/helpers/envs/base_env.sh` when that script is sourced. + +| Variable | Role | +|----------|------| +| `NCCL_IB_HCA` | Selects **InfiniBand Host Channel Adapters** (device:port list). | +| `NCCL_IB_GID_INDEX` | **GID index** for the active port (RoCE and IB differ; see vendor docs). | +| `NCCL_IB_TC` | **Traffic class** for InfiniBand. | +| `NCCL_IB_FIFO_TC` | Traffic class for FIFO traffic. | +| `NCCL_IB_RETRY_CNT` | Retry count for IB operations (tune with vendor guidance). | +| `NCCL_IB_TIMEOUT` | Timeout for IB operations. | +| `NCCL_IB_QPS_PER_CONNECTION` | Queue pairs per connection. | +| `NCCL_NET_GDR_LEVEL` | **GPUDirect RDMA** level for NIC/GPU transfers. | +| `NCCL_DMABUF_ENABLE` | Use **DMA-BUF** path where supported. | + +### Auto-detection in Primus + +If `NCCL_IB_HCA` is **unset**, `base_env.sh` runs `runner/helpers/envs/get_nccl_ib_hca.sh`, which enumerates `/sys/class/infiniband/`, skips bonded/storage-style devices, and builds a comma-separated `device:port` list for `NCCL_IB_HCA`. + +Default in `base_env.sh`: + +```bash +export NCCL_IB_GID_INDEX=${NCCL_IB_GID_INDEX:-3} +``` + +AINIC-oriented configs often override `NCCL_IB_GID_INDEX` to `1` (see `runner/use_ainic.yaml` and `03_enable_ainic.sh`). + +--- + +## 3. RoCE (RDMA over converged ethernet) + +RoCE reuses much of the **IB verb** stack; the same **`NCCL_IB_*`** knobs apply. + +| Variable | Typical use | +|----------|-------------| +| `NCCL_IB_ROCE_VERSION_NUM` | RoCE version (commonly **2** for RoCE v2). | + +GID selection (`NCCL_IB_GID_INDEX`) and traffic classes (`NCCL_IB_TC`, `NCCL_IB_FIFO_TC`) remain important on RoCE fabrics. Follow your network team’s mapping (often **GID index 1** for RoCE v2 vs **3** for some IB fabrics—your site may differ). + +--- + +## 4. AINIC (AMD AI NIC) + +**AINIC** refers to AMD’s AI-optimized NIC path (for example, the **AMD Pensando™ Pollara 400 AI NIC**) used in some clusters. Enabling it is a combination of **environment**, **container image**, and **device pass-through**. + +### Enable AINIC + +- Set **`USING_AINIC=1`**. The hook `runner/helpers/hooks/03_enable_ainic.sh` runs when this is set and exports AINIC-related variables back to the caller (`env.VAR=VALUE` lines). +- Use container images built for AINIC when required by your site. Examples in this repository use tags such as `docker.io/tasimage/primus:-ainic` (see `examples/customer_package/` and `.github/workflows/ci.yaml`). Match the image to your ROCm and ANP bundle. + +### `runner/use_ainic.yaml` + +Primus CLI system defaults for AINIC-oriented runs include: + +- Container **`device`** mounts: `/dev/kfd`, `/dev/dri`, `/dev/infiniband` (required for GPU and IB access in the container). +- Environment entries such as `USING_AINIC=1`, `NCCL_PXN_DISABLE=0`, and `NCCL_IB_GID_INDEX=1`. + +Adjust **`NCCL_IB_GID_INDEX`** and **`container.options.image`** to match your cluster; comments in `runner/use_ainic.yaml` call this out explicitly. + +### AINIC hook + +**`runner/helpers/hooks/03_enable_ainic.sh`** is the supported hook path: it sets ANP/RCCL/MPI home directories, IB QoS, RoCE version, P2P channel counts, GDR flush behavior, `LD_LIBRARY_PATH` (including `libibverbs` and RCCL/ANP/MPI build paths), and related flags. Default `NCCL_IB_FIFO_TC` in the hook is **192**; align this value with your fabric. + +### RCCL network plugin (ANP) + +For ANP-based networking, clusters often set **`NCCL_NET_PLUGIN`** to **`librccl-anp.so`** when that library is present under `ANP_HOME_DIR`, falling back to `librccl-net.so` otherwise—see the logic in `examples/run_pretrain.sh`. This complements the library paths from `03_enable_ainic.sh`. + +### Variables commonly set for AINIC + +From `03_enable_ainic.sh` (non-exhaustive): + +| Variable | Purpose | +|----------|---------| +| `ANP_HOME_DIR`, `RCCL_HOME_DIR`, `MPI_HOME_DIR` | Install roots for ANP, RCCL, and Open MPI. | +| `NCCL_IB_TC`, `NCCL_IB_FIFO_TC` | Traffic classes for IB/RoCE. | +| `NCCL_IB_GID_INDEX` | Often **1** for AINIC-oriented configs in Primus examples. | +| `NCCL_IB_ROCE_VERSION_NUM` | RoCE v2. | +| `RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING` | Stricter GDR flush ordering (set to `0` in these scripts). | +| `LD_LIBRARY_PATH` | Prepends `libibverbs`, RCCL, ANP, and MPI library paths. | + +--- + +## 5. Socket configuration + +CPU-side and fallback socket traffic uses interface selection: + +| Variable | Role | +|----------|------| +| `NCCL_SOCKET_IFNAME` | Interface name or pattern for NCCL socket transport (e.g. `eth0`, or `^docker0,lo` to **exclude** virtual interfaces). | +| `GLOO_SOCKET_IFNAME` | Interface for **Gloo** process groups (CPU barriers and related). | + +**Primus behavior:** `base_env.sh` sets `IP_INTERFACE` via `runner/helpers/envs/get_ip_interface.sh` (fallback: first address from `hostname -I`). Both `NCCL_SOCKET_IFNAME` and `GLOO_SOCKET_IFNAME` default to **`IP_INTERFACE`** when unset. + +**Requirement:** All nodes must agree on a **reachable** address family and interface choice; mismatched bindings are a frequent source of hangs. + +--- + +## 6. PCIe cross-NIC (PXN) + +| Variable | Default in `base_env.sh` | Meaning | +|----------|--------------------------|---------| +| `NCCL_PXN_DISABLE` | `1` | **PXN disabled** by default (saves GPU memory per comment in `base_env.sh`). | + +When **`NCCL_PXN_DISABLE=0`**, **PCIe cross-NIC** is enabled: GPUs may use NICs attached to **other** PCIe switches, which can improve **multi-rail** bandwidth at the cost of **higher GPU memory** use. `runner/use_ainic.yaml` sets `NCCL_PXN_DISABLE=0` for AINIC-oriented runs. + +--- + +## 7. Network diagnostics + +### Preflight + +```bash +primus-cli direct -- preflight --network +``` + +For multi-node (Slurm example): + +```bash +primus-cli slurm srun -N 4 -- preflight --host --gpu --network +``` + +See `docs/02-user-guide/preflight.md` for flags, output locations (`output/preflight` by default), and interpretation. + +Set **`PRIMUS_EXPECT_IB=1`** when InfiniBand is **required** for validation; preflight uses this in `primus/tools/preflight/network/network_standard.py`. + +### RCCL benchmark + +```bash +primus-cli slurm srun -N 4 -- benchmark rccl --op all_reduce --min-bytes 1M --max-bytes 128M +``` + +This exercises collective bandwidth and latency across a message-size sweep. See `docs/02-user-guide/benchmarking.md` and `primus/tools/benchmark/rccl_bench_args.py` for options (dtypes, operations, output files). + +### Verbose RCCL logs + +```bash +export NCCL_DEBUG=INFO +``` + +Use for short, controlled runs; **TRACE** can be extremely verbose. + +--- + +## 8. Multi-node setup checklist + +- **ROCm version** matches across all nodes (driver and container image). +- **`NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME`** (or auto-detected `IP_INTERFACE`) identify the **same logical network** on every node. +- **InfiniBand or RoCE** is up (`ibstat`, `/dev/infiniband`, kernel modules such as `ib_core` / `mlx5_core` as appropriate). +- **Firewall** allows ports required by your launcher and collective tests (`MASTER_ADDR` / `MASTER_PORT` reachable). +- **`MASTER_ADDR`** resolves and is reachable from **all** nodes. +- **`GPUS_PER_NODE`** matches physical GPUs per node. +- **Containers** mount `/dev/kfd`, `/dev/dri`, and `/dev/infiniband` when using IB/RoCE/AINIC (see `runner/use_ainic.yaml`). + +--- + +## 9. Troubleshooting network issues + +| Symptom | What to check | +|---------|----------------| +| **Timeout or hang** at init | `MASTER_ADDR` / `MASTER_PORT`, firewall, VPN, wrong `NCCL_SOCKET_IFNAME`, or inconsistent interface across nodes. | +| **Slow** collectives | IB vs Ethernet path, `NCCL_NET_GDR_LEVEL`, fabric errors, or contention; compare **`benchmark rccl`** to baseline. | +| **IB not detected** | `/dev/infiniband` missing, modules not loaded, or wrong container devices. | +| **Wrong interface** | Restrict with `NCCL_SOCKET_IFNAME=^docker0,lo` (exclude loopback and Docker bridges). | +| **GID / RoCE issues** | `NCCL_IB_GID_INDEX` vs site documentation; RoCE v2 settings (`NCCL_IB_ROCE_VERSION_NUM`). | + +For a consolidated list of `NCCL_*` / `RCCL_*` variables, see `docs/03-configuration-reference/environment-variables.md` and the upstream [RCCL environment variables](https://rocm.docs.amd.com/projects/rccl/en/develop/api-reference/env-variables.html) documentation. + +--- + +## Related documentation + +- [Preflight diagnostics](../02-user-guide/preflight.md) +- [Benchmark suite](../02-user-guide/benchmarking.md) +- [NCCL/RCCL collective operations](./collective-operations.md) +- [Environment variables](../03-configuration-reference/environment-variables.md) diff --git a/docs/README_NATIVE_SFT_LORA_EN.md b/docs/04-technical-guides/native-sft-lora.md similarity index 95% rename from docs/README_NATIVE_SFT_LORA_EN.md rename to docs/04-technical-guides/native-sft-lora.md index 9ff918858..ec160614b 100644 --- a/docs/README_NATIVE_SFT_LORA_EN.md +++ b/docs/04-technical-guides/native-sft-lora.md @@ -1,4 +1,4 @@ -# Primus Native SFT LoRA — Quick Start +# Primus native SFT LoRA—quick start > **Branch**: `feat/megatron/support-sft-native` (PR701) > **Backend**: Megatron-LM **native** (no Megatron-Bridge runtime dependency) @@ -38,13 +38,13 @@ Entry point: `primus/backends/megatron/megatron_sft_trainer.py` (`MegatronSFTTra --- -## 2. Runtime Environment +## 2. Runtime environment ### 2.1 Docker container | Container | Image | Notes | |---|---|---| -| **`sft_primus_0507_native`** | `rocm/primus:v26.2` | Recommended; verified | +| **`sft_primus_0507_native`** | `rocm/primus:v26.3` | Recommended; verified | Container mounts (set once at container start): @@ -63,13 +63,13 @@ export EXP_NAME="llama2_70b_native_$(date +%Y%m%d_%H%M%S)" ``` Set automatically inside the container by `examples/run_pretrain.sh` (you don't need to touch these): -- `TRITON_CACHE_DIR`, `MIOPEN_USER_DB_PATH`, `PRIMUS_CACHE_ROOT` — persistent JIT cache -- `NCCL_*` / `RCCL_*` — communication tuning -- `HSA_*` / `GPU_MAX_HW_QUEUES` — AMD GPU performance tuning +- `TRITON_CACHE_DIR`, `MIOPEN_USER_DB_PATH`, `PRIMUS_CACHE_ROOT`—persistent JIT cache +- `NCCL_*` / `RCCL_*`—communication tuning +- `HSA_*` / `GPU_MAX_HW_QUEUES`—AMD GPU performance tuning --- -## 3. Launch Commands (verified) +## 3. Launch commands (verified) ### 3.1 BF16 / FP8 (existing yaml configs, ready to run) @@ -130,7 +130,7 @@ The plumbing is already in place: Hard constraints: -1. **TransformerEngine ≥ 2.7.0.dev0** required (the `rocm/primus:v26.2` image already satisfies this) +1. **TransformerEngine ≥ 2.7.0.dev0** required (the `rocm/primus:v26.3` image already satisfies this) 2. **FP4 and FP8 are mutually exclusive**: `args.fp4 and args.fp8` raises in Megatron (`arguments.py:885-887`) 3. **`fp4_param` must be paired with `fp4`**: enabling `fp4_param` alone raises (`arguments.py:889-891`) @@ -310,7 +310,7 @@ modules: # ===================================================================== enable_primus_turbo: true use_turbo_attention: false - use_turbo_grouped_mlp: false + use_turbo_grouped_gemm: false use_turbo_rms_norm: false # ---------- Cross-entropy fusion ---- @@ -403,11 +403,11 @@ grep -E "throughput per GPU" "$RANK0" | head -5 ### Q1: `--fp4-format requires Transformer Engine >= 2.7.0.dev0` -Upgrade TE inside the container, or switch to image `rocm/primus:v26.2`+. +Upgrade TE inside the container, or switch to image `rocm/primus:v26.3`+. ### Q2: `--fp4-format and --fp8-format cannot be used simultaneously` -Leftover `fp8: hybrid` / `fp8: e4m3` in the yaml — must be removed. +Leftover `fp8: hybrid` / `fp8: e4m3` in the yaml—must be removed. ### Q3: `--fp4-param-gather must be used together with --fp4-format` @@ -437,9 +437,10 @@ done --- -## 6. References / Further Reading +## 6. References / further reading -- **PR #701** — Full implementation of this native SFT stack: +- **Post-training overview**: [Post-Training (SFT / LoRA / DPO)](../02-user-guide/posttraining.md)—how this native SFT LoRA path fits into the broader fine-tuning workflow. +- **PR #701**—Full implementation of this native SFT stack: https://github.com/AMD-AGI/Primus/pull/701 - **Megatron-LM FP4 design**: `third_party/Megatron-LM/megatron/core/fp4_utils.py` + @@ -454,6 +455,6 @@ done ## 7. Maintainers -- @wenxie-amd — PR #701 main author -- @Xiaoming-AMD — co-author (trainer + dataset core) -- @botaohu001 — packing / mlperf-aligned recipe / diagnostic tools +- @wenxie-amd—PR #701 main author +- @Xiaoming-AMD—co-author (trainer + dataset core) +- @botaohu001—packing / mlperf-aligned recipe / diagnostic tools diff --git a/docs/04-technical-guides/parallelism-configuration.md b/docs/04-technical-guides/parallelism-configuration.md new file mode 100644 index 000000000..d0041054d --- /dev/null +++ b/docs/04-technical-guides/parallelism-configuration.md @@ -0,0 +1,255 @@ +# Parallelism configuration guide + +Primus is a YAML-driven training framework for AMD GPUs. Megatron-LM, TorchTitan, and MaxText each expose parallelism through different configuration namespaces. This guide explains how to set parallelism and batch-related parameters, how global batch size relates to micro batch size and data parallel width, and how to choose a parallel strategy for common model sizes. + +Default values cited below come from Primus module presets: + +- Megatron trainer: `primus/configs/modules/megatron/trainer_base.yaml` +- Megatron model (tensor/pipeline/expert/context parallel): `primus/configs/models/megatron/language_model.yaml` +- TorchTitan: `primus/configs/modules/torchtitan/pre_trainer.yaml` + +Experiment YAMLs in `examples/` often override these defaults for specific models and hardware. + +--- + +## 1. Megatron parallelism configuration + +Model-parallel degrees live on the **model** config (for example under `model:` in your experiment YAML, merged from `language_model.yaml`). Training batch and overlap settings live on the **trainer** module (`trainer_base.yaml`). + +### Core parallel degrees + +| Parameter | Default (Primus `language_model.yaml`) | Description | +|-----------|------------------------------------------|-------------| +| `tensor_model_parallel_size` | `1` | Tensor parallelism (TP): shards attention and MLP across this many GPUs. | +| `pipeline_model_parallel_size` | `1` | Pipeline parallelism (PP): number of pipeline stages. | +| `expert_model_parallel_size` | `1` | Expert parallelism (EP) for MoE: shards experts across this many GPUs. | +| `context_parallel_size` | `1` | Context parallelism (CP) for long sequences. | +| `sequence_parallel` | `true` | Sequence parallelism (SP); typically used with TP greater than 1. | + +### Virtual pipeline (VPP) and pipeline communication + +| Parameter | Description | +|-----------|-------------| +| `virtual_pipeline_model_parallel_size` | Interleaved pipeline depth (null disables VPP). | +| `num_layers_per_virtual_pipeline_stage` | Layers per virtual stage when using VPP. | +| `overlap_p2p_comm` | Overlap pipeline P2P with compute (default `true` in `trainer_base.yaml`). | + +### Optimizer, FSDP, and overlap (trainer module) + +| Parameter | Default (`trainer_base.yaml`) | Description | +|-----------|-------------------------------|-------------| +| `use_distributed_optimizer` | `false` | ZeRO-1 style optimizer state sharding when enabled. | +| `use_torch_fsdp2` | `false` | Full FSDP2 integration. | +| `overlap_grad_reduce` | `false` | Overlap gradient all-reduce with backward. | +| `overlap_param_gather` | `false` | Overlap parameter gathering with forward. | + +Set these to `true` in your experiment when you want communication/compute overlap; many production configs enable `use_distributed_optimizer` and overlap flags for large runs. + +### Data parallel size (implicit) + +For Megatron, data parallel size is not a single YAML key; it is implied by the world size and the product of parallel degrees: + +\[ +\text{DP} = \frac{\text{world\_size}}{\text{TP} \times \text{PP} \times \text{EP}} +\] + +(Adjust if you also use context parallelism or other groupings; your job’s process layout must match the configured degrees.) + +### Batch parameters + +| Parameter | Default (`trainer_base.yaml`) | Description | +|-----------|-------------------------------|-------------| +| `micro_batch_size` | `2` | Micro batch size per data-parallel rank (MBS). | +| `global_batch_size` | `128` | Target global batch size (GBS) across the data parallel group. | + +Megatron derives **gradient accumulation** from `global_batch_size`, `micro_batch_size`, and the effective data parallel size so that: + +\[ +\text{GBS} = \text{MBS} \times \text{DP} \times \text{gradient\_accumulation\_steps} +\] + +Equivalently: + +\[ +\text{gradient\_accumulation\_steps} = \frac{\text{GBS}}{\text{MBS} \times \text{DP}} +\] + +You normally set `global_batch_size` and `micro_batch_size` in YAML; Megatron computes the number of accumulation steps automatically. + +--- + +## 2. TorchTitan parallelism configuration + +TorchTitan parallelism is grouped under the `parallelism:` key in the TorchTitan module (see `primus/configs/modules/torchtitan/pre_trainer.yaml`). + +### `parallelism.*` parameters + +| Key | Default | Description | +|-----|---------|-------------| +| `parallelism.tensor_parallel_degree` | `1` | Tensor parallelism degree. | +| `parallelism.pipeline_parallel_degree` | `1` | Pipeline parallelism degree. | +| `parallelism.data_parallel_shard_degree` | `-1` | FSDP shard degree; `-1` lets the framework choose. | +| `parallelism.data_parallel_replicate_degree` | `1` | DDP-style replication degree. | +| `parallelism.expert_parallel_degree` | `1` | Expert parallelism for MoE. | +| `parallelism.context_parallel_degree` | `1` | Context parallelism. | +| `parallelism.fsdp_reshard_after_forward` | `default` | FSDP reshard policy (`default` uses TorchTitan’s default behavior). | +| `parallelism.enable_async_tensor_parallel` | `false` | Async tensor-parallel communication. | +| `parallelism.pipeline_parallel_schedule` | `1F1B` | Pipeline schedule (for example `1F1B`). | +| `parallelism.pipeline_parallel_microbatch_size` | `1` | Microbatch size for pipeline stages. | + +### Batch parameters under `training.*` + +| Key | Default | Description | +|-----|---------|-------------| +| `training.global_batch_size` | `-1` | Global batch size; `-1` typically means unset or derived. | +| `training.local_batch_size` | `8` | Per-rank local (micro) batch size. | + +### Global batch relationship + +For TorchTitan, a useful relationship when using replicate and shard degrees explicitly is: + +\[ +\text{global\_batch\_size} \approx \text{local\_batch\_size} \times \text{data\_parallel\_replicate\_degree} \times \text{data\_parallel\_shard\_degree} +\] + +Exact semantics follow TorchTitan’s distributed layout; set `training.global_batch_size` and parallelism degrees consistently with your launcher’s world size. + +--- + +## 3. MaxText parallelism configuration + +MaxText (JAX) uses a **device mesh** with **ICI** (intra-node / “in-cluster interconnect”) and **DCN** (inter-node / “data center network”) axes for parallelism. Defaults and parameter names come from upstream MaxText, for example `third_party/maxtext/src/MaxText/configs/base.yml`, not from Primus presets alone. + +### Common parallelism keys (from `base.yml`) + +Examples include: + +- `ici_tensor_parallelism`—tensor parallelism within a node +- `ici_fsdp_parallelism`—FSDP-style sharding on ICI (default `-1` for auto in many layouts) +- `dcn_data_parallelism`—data parallelism across nodes (default `-1` for auto) +- `dcn_fsdp_parallelism`—FSDP across DCN + +### Batch sizing + +- `per_device_batch_size`—primary knob for per-device batch (see `base.yml`). + +Consult MaxText’s mesh documentation and your chosen model YAML for valid combinations of ICI/DCN axes. + +--- + +## 4. Batch size relationships + +### Megatron-style identity + +\[ +\text{GBS} = \text{MBS} \times \text{DP} \times \text{grad\_accum} +\] + +\[ +\text{DP} = \frac{\text{world\_size}}{\text{TP} \times \text{PP} \times \text{EP}} +\] + +(Subject to your exact parallel groups; CP and custom layouts can introduce additional groups.) + +### How GBS, MBS, and DP interact + +| Goal | What to change | +|------|----------------| +| Increase global batch without more per-GPU memory | Increase `gradient_accumulation_steps` (Megatron) or increase accumulation / GBS while keeping MBS fixed. | +| Increase throughput per step | Increase `micro_batch_size` if memory allows; may require lowering accumulation to keep GBS fixed. | +| Scale to more GPUs | Increase world size; often increase DP; keep GBS stable by adjusting accumulation. | + +### Memory and convergence + +| Factor | Effect | +|--------|--------| +| **MBS** | Strongly affects per-GPU activation memory; larger MBS often improves GPU utilization but can OOM. | +| **GBS** | Affects effective noise in the gradient and optimal learning rate scaling; many recipes scale LR with GBS. | + +**Practical recommendation:** start with `micro_batch_size` of `1` or `2`, verify stability and memory. Increase `global_batch_size` (via accumulation or more DP ranks) gradually while monitoring loss and adjusting learning rate per your recipe. + +### Example numeric table (Megatron-style) + +Assume TP=1, PP=1, EP=1, so DP equals world size. + +| World size (DP) | MBS | Grad accum | GBS | +|-----------------|-----|------------|-----| +| 8 | 1 | 16 | 128 | +| 8 | 2 | 8 | 128 | +| 16 | 1 | 8 | 128 | +| 16 | 2 | 4 | 128 | + +--- + +## 5. Decision guide: Choosing parallelism + +| Situation | Suggested direction | +|-----------|----------------------| +| Model fits on **one GPU** | Use DP and/or FSDP only; TP=1, PP=1. | +| Model fits on **one node** but not one GPU | **TP** within the node; **DP** across any remaining replicas. | +| Model needs **multiple nodes** | **TP** within node where possible; **PP** across nodes for very large depth; **DP** for remaining width. | +| **MoE** | Add **EP**; align expert count and routing with `expert_model_parallel_size` / `parallelism.expert_parallel_degree`. | +| **Very long sequences** | Increase **CP** (`context_parallel_size` / `context_parallel_degree`) as supported by the backend. | + +### Example configurations (illustrative) + +These are representative topologies; always validate with your checkpoint format, memory profile, and hardware interconnect. + +| Profile | GPUs | TP | PP | EP | DP (illustrative) | +|---------|------|----|----|----|---------------------| +| ~7B | 8 | 1 | 1 | 1 | 8 | +| ~70B | 64 | 8 | 2 | 1 | 4 | +| Large MoE (~671B class) | many | 8 | 4 | 8 | remainder | + +--- + +## 6. Common parallelism recipes (YAML snippets) + +### Megatron: 8-GPU data parallel only + +```yaml +# model (or merged language_model section) +tensor_model_parallel_size: 1 +pipeline_model_parallel_size: 1 +expert_model_parallel_size: 1 +context_parallel_size: 1 +sequence_parallel: false + +# trainer +micro_batch_size: 2 +global_batch_size: 128 +``` + +### Megatron: Tensor + pipeline + data parallel + +```yaml +tensor_model_parallel_size: 8 +pipeline_model_parallel_size: 2 +expert_model_parallel_size: 1 +context_parallel_size: 1 +sequence_parallel: true + +micro_batch_size: 1 +global_batch_size: 512 +``` + +### TorchTitan: TP + PP with explicit schedule + +```yaml +parallelism: + tensor_parallel_degree: 4 + pipeline_parallel_degree: 2 + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + expert_parallel_degree: 1 + context_parallel_degree: 1 + pipeline_parallel_schedule: 1F1B + pipeline_parallel_microbatch_size: 1 + enable_async_tensor_parallel: false + +training: + global_batch_size: 256 + local_batch_size: 4 +``` + +For full worked examples, see `examples/megatron/configs/` and `examples/torchtitan/configs/` under your target hardware (for example `MI300X/`). diff --git a/docs/04-technical-guides/parallelism-strategies.md b/docs/04-technical-guides/parallelism-strategies.md new file mode 100644 index 000000000..83b028f6e --- /dev/null +++ b/docs/04-technical-guides/parallelism-strategies.md @@ -0,0 +1,361 @@ +# Parallelism strategies for distributed training + +This guide explains the parallelism dimensions used when training large foundation models on AMD GPUs with Primus. It moves from basic data parallelism to advanced combinations of tensor, pipeline, context, and expert parallelism, including how Primus exposes these options through Megatron-LM and TorchTitan. + +For Megatron YAML flags and environment tuning, see [Megatron parameters](../03-configuration-reference/megatron-parameters.md) and [Environment variables](../03-configuration-reference/environment-variables.md). + +--- + +## 1. Introduction + +### Why parallelism is needed + +Modern foundation models often exceed the memory of a single accelerator: parameters, activations, optimizer states, and KV caches cannot all reside on one device at useful batch sizes. Even when a model *fits*, training throughput may be too low without scaling across many GPUs. Parallelism splits the problem along several independent **dimensions** so that: + +- **Memory** is shared across devices (sharding, pipeline stages, sequence splits). +- **Compute** is scaled by processing more data in parallel or by overlapping communication with computation. + +### Overview of parallelism dimensions + +| Dimension | What is split | Primary goal | +|-----------|----------------|--------------| +| **Data parallelism (DP)** | Input batches | Throughput; same model on each GPU | +| **FSDP / ZeRO** | Parameters, gradients, optimizer (by stage) | Memory; keep DP semantics | +| **Tensor parallelism (TP)** | Individual weight matrices / matmuls | Memory per layer; needs fast links | +| **Sequence parallelism (SP)** | Sequence in non-TP regions | Activation memory with TP | +| **Pipeline parallelism (PP)** | Layer groups across stages | Memory; depth-wise split | +| **Context parallelism (CP)** | Sequence for attention (e.g. ring) | Very long contexts | +| **Expert parallelism (EP)** | MoE experts across devices | Memory and compute for MoE | + +These can be **combined**. The product of parallel degrees must match how processes are laid out on the cluster (see [Section 9](#9-combining-parallelism-strategies)). + +--- + +## 2. Data parallelism (DP) + +In **classic data parallelism**, every GPU holds a **full copy** of the model. Each rank receives a **different mini-batch** of data. After the backward pass, **gradients are synchronized** so that all ranks apply the same update. + +``` + Batch shard 0 Batch shard 1 Batch shard 2 Batch shard 3 + | | | | + v v v v + +--------+ +--------+ +--------+ +--------+ + | GPU 0 | | GPU 1 | | GPU 2 | | GPU 3 | + | full | | full | | full | | full | + | model | | model | | model | | model | + +--------+ +--------+ +--------+ +--------+ + | | | | + +------------------+------------------+------------------+ + | + AllReduce(gradients) + | + v + Same weights on all ranks after optimizer step +``` + +**Properties** + +- Simple to reason about and widely supported. +- Requires the **full model, activations for one micro-batch, and optimizer state** to fit in **one GPU’s memory** (unless combined with other strategies). + +**Effective batch size** + +For a single update that aggregates over data-parallel ranks and gradient accumulation: + +\[ +\text{effective\_batch\_size} = \text{micro\_batch\_size} \times \text{num\_GPUs}_{\text{DP}} \times \text{gradient\_accumulation\_steps} +\] + +Here `num_GPUs_DP` is the **data-parallel group size** (not always the same as `world_size` when TP/PP/EP are also used). + +--- + +## 3. Fully sharded data parallel (FSDP / ZeRO) + +**ZeRO** (Zero Redundancy Optimizer) reduces redundant storage by **sharding** optimizer states, gradients, and/or parameters across data-parallel ranks. + +| ZeRO stage | Sharded | Idea | +|------------|---------|------| +| **Stage 1** | Optimizer states | Each rank keeps only \(1/N\) of optimizer tensors | +| **Stage 2** | + Gradients | Gradients are sharded; reduced where needed | +| **Stage 3** | + Parameters | Each rank holds \(1/N\) of parameters; gather before use | + +**FSDP** (Fully Sharded Data Parallel) in PyTorch is the common **implementation** of sharded data parallel training; in the Megatron ecosystem, **ZeRO-3-style** behavior is often discussed alongside **FSDP** for full parameter sharding. + +**Typical execution pattern (conceptual)** + +1. **Forward:** **AllGather** (or equivalent) to materialize parameters needed for the current layer/batch on each rank. +2. **Backward:** **ReduceScatter** (or equivalent) to write shard-sized gradient pieces back to ranks. + +**Memory intuition** + +If replicated training used \(M\) memory per rank for parameters+gradients+optimizer, **ideal** full sharding across \(N\) ranks approaches **\(M/N\)** for the sharded pieces (plus buffers and fragmentation). Moving from **full replication** to **\(1/N\)** sharding for those tensors saves roughly **\((N-1)/N\)** of that component—for **8 GPUs**, about **87.5%** of the replicated footprint for the sharded tensors. + +### In Primus + +| Backend | Configuration | +|---------|----------------| +| **Megatron-LM** | `use_distributed_optimizer: true` enables the **distributed optimizer** (ZeRO-1–style optimizer sharding in Megatron). For **PyTorch FSDP2**, set `use_torch_fsdp2: true` (see Megatron constraints: FSDP2 and distributed optimizer are not used together). | +| **TorchTitan** | `data_parallel_shard_degree` controls how ranks participate in **FSDP-style** sharding (see TorchTitan job config; `-1` often means auto). | + +Exact interactions with checkpoint formats and DDP are documented in [Megatron parameters](../03-configuration-reference/megatron-parameters.md). + +--- + +## 4. Tensor parallelism (TP) + +**Tensor parallelism** splits **individual layers** (usually linear / attention projections) across GPUs so **no single GPU stores the full weight matrix** for that layer. + +### Column-parallel vs row-parallel + +Consider a linear layer \(Y = X W\) with weight matrix \(W\). **Column-parallel** splits \(W\) **along the output dimension** (columns). **Row-parallel** splits \(W\) **along the input dimension** (rows) and **splits \(X\)** so each rank’s matmul dimensions match. + +**Column-parallel linear**—each rank holds **disjoint columns** of \(W\); each rank's output is a **disjoint column shard** (half the width for 2-way TP). To recover the full-width tensor the shards are **concatenated (All-Gather along the output dim)**—this is only done when the full tensor is actually needed: + +``` + SAME full X replicated on each TP rank + | + +-----------------+-----------------+ + | | + v v + Rank 0: X @ W[:,0:h/2] Rank 1: X @ W[:,h/2:h] + | | + v v + partial Y_0 partial Y_1 + (narrow) (narrow) + | | + +-----------------+-----------------+ + | + All-Gather (concatenate) on output dim + (only when the full tensor is needed; with + gather_output=False the output stays column- + sharded and feeds the next layer with no comm) + | + v + full-width Y (concatenation of shards) +``` + +**Row-parallel linear**—each rank holds **disjoint rows** of \(W\); **input \(X\)** is **split** along the **input feature** dimension so each rank computes part of the reduction: + +``` + Rank 0: X_0 @ W[0:r/2,:] ----+ + +-- AllReduce --> Y + Rank 1: X_1 @ W[r/2:r,:] ----+ + (X split along features) (partial sums add to full Y) +``` + +Typical **transformer block** pattern: **column-parallel** for the first projection—its column-sharded output is fed directly into the next layer **without communication**—then **row-parallel** for the second projection, which performs the single **AllReduce** that reconstructs the full output. Column-parallel itself only communicates when `gather_output=True`. + +**Communication** + +- Often **AllReduce** of partial outputs, or **ReduceScatter** + **AllGather** sequences depending on implementation and **sequence parallelism** (see next section). + +**When to use** + +- Best **within a node** (NVLink / high-bandwidth GPU–GPU paths). Multi-node TP is possible but latency-sensitive. + +### In Primus + +| Backend | Parameter | +|---------|-----------| +| Megatron-LM | `tensor_model_parallel_size` | +| TorchTitan | `parallelism.tensor_parallel_degree` | + +--- + +## 5. Sequence parallelism (SP) + +**Sequence parallelism** extends TP by splitting the **sequence dimension** in regions that are not covered by tensor-parallel matmuls—commonly **LayerNorm**, **dropout**, and sometimes **residual** paths—so **activation memory** scales better when **TP > 1**. + +**Interaction with TP** + +- After a **column-parallel** region, partial activations can be **ReduceScatter**d along the sequence. +- Before a **row-parallel** region, activations may be **AllGather**d along the sequence. + +So SP trades **extra collectives** for **lower per-rank activation footprint** on long sequences. + +### In Primus + +| Backend | Parameter | +|---------|-----------| +| Megatron-LM | `sequence_parallel: true` (used with TP) | +| TorchTitan | Sequence-parallel behavior is integrated with TP/parallelization pipelines in supported models | + +--- + +## 6. Pipeline parallelism (PP) + +**Pipeline parallelism** assigns **disjoint subsets of layers** to **stages** on different devices. Activations (and gradients) move **between stages** with **point-to-point** communication. + +``` + Microbatch 1: Stage0 -> Stage1 -> Stage2 -> Stage3 + Microbatch 2: Stage0 -> Stage1 -> Stage2 -> Stage3 + ... +``` + +### Pipeline bubbles + +If a stage waits for input while other stages compute, **idle time** appears (**pipeline bubble**). Schedulers reduce bubbles by overlapping forwards and backwards across microbatches. + +**Common schedules** + +| Schedule | Idea | +|----------|------| +| **1F1B** | One forward, one backward; classic **warmup / steady / cooldown** phases | +| **1F1B interleaved (VPP)** | **Virtual pipeline** stages: multiple chunks per device to improve utilization | +| **Zero-bubble (ZB)** | Reorders / splits backward so **forward and backward** hide each other better; may separate **input-gradient** vs **weight-gradient** phases | +| **V-Schedule / V-Half / V-Min** | Variants reducing bubbles further (names vary by codebase) | +| **DualPipe** | Bidirectional pipeline scheduling (e.g. DeepSeek-style) to overlap forward/backward paths | + +**Bubble rate** + +\[ +\text{bubble\_rate} = \frac{\text{idle time}}{\text{total time}} +\] + +Lower is better; large **microbatch counts** and better schedules reduce bubble overhead. + +### In Primus (Megatron) + +| Parameter | Role | +|-----------|------| +| `pipeline_model_parallel_size` | Number of pipeline stages | +| `patch_zero_bubble` | Enable Primus/Megatron **zero-bubble** pipeline patches | +| `patch_primus_pipeline` | Use Primus pipeline implementation for schedule logic | +| `pp_algorithm` | e.g. `1f1b`, `1f1b-interleaved`, `zero-bubble`, `zero-bubble-heuristic`, `zbv-formatted`, `v-half`, `v-min` | + +See `primus/configs/modules/megatron/primus_pipeline.yaml` and `zero_bubble.yaml` in the repo for defaults. + +### In Primus (TorchTitan) + +| Parameter | Role | +|-----------|------| +| `parallelism.pipeline_parallel_degree` | Pipeline depth | +| `parallelism.pipeline_parallel_schedule` | e.g. `1F1B`, `Interleaved1F1B`, `GPipe`, zero-bubble variants where supported | + +--- + +## 7. Context parallelism (CP) + +**Context parallelism** splits the **sequence length** across devices for long-context training. A common pattern is **ring attention**: each rank holds a **chunk** of queries/keys/values and participates in a **ring** of message passing so attention covers the full sequence without centralizing all activations on one GPU. + +**Use cases** + +- Long documents, 32K–128K+ tokens, where **per-layer activation memory** and **attention compute** must be distributed. + +### In Primus + +| Backend | Parameter | +|---------|-----------| +| Megatron-LM | `context_parallel_size` | +| TorchTitan | `parallelism.context_parallel_degree` | + +--- + +## 8. Expert parallelism (EP) + +**Mixture-of-Experts (MoE)** models route each token to a small subset of **experts**. **Expert parallelism** assigns **different experts** to **different GPUs** so expert weights are not duplicated on every device. + +**Communication** + +- **AllToAll** (or equivalent) is typical: **dispatch** tokens to expert ranks and **combine** expert outputs back. + +**Expert tensor parallelism (ETP)** + +- Experts can be further **tensor-parallel** within a subset of GPUs, analogous to TP for dense layers. + +### In Primus + +| Backend | Parameter | +|---------|-----------| +| Megatron-LM | `expert_model_parallel_size` | +| TorchTitan | `parallelism.expert_parallel_degree`, `parallelism.expert_tensor_parallel_degree` | + +--- + +## 9. Combining parallelism strategies + +### Common pattern + +- **TP within a node** (fast interconnect). +- **PP across nodes or across groups** when layers do not fit on one device. +- **DP / FSDP** for scaling batch size and sharding optimizer state or parameters. + +### GPU count (simplified) + +For dense models (ignoring CP and detailed MoE layout): + +\[ +\text{world\_size} \approx \text{TP} \times \text{PP} \times \text{DP} +\] + +For MoE-heavy setups, you often see: + +\[ +\text{world\_size} \approx \text{TP} \times \text{PP} \times \text{EP} \times \text{DP} +\] + +**Context parallelism** introduces another multiplicative factor in layouts where CP ranks are part of the global mesh (exact rank ordering is implementation-specific). + +### Memory vs communication + +- **More TP** → smaller matrices per GPU but **more frequent** collectives within layers. +- **More PP** → less memory per stage but **pipeline bubbles** and **latency** between stages. +- **More DP/FSDP** → better throughput scaling if communication is not saturated. + +### Example configurations (illustrative) + +| Scenario | TP | PP | DP / notes | +|----------|----|----|------------| +| ~7B on 8 GPUs | 1 | 1 | 8-way DP (or FSDP) | +| ~70B on 64 GPUs (8 nodes × 8) | 8 | 2 | 4-way DP | +| Large MoE (e.g. 671B-class) on 256 GPUs | 8 | 4 | EP 8 (example; real jobs vary widely) | + +Always validate against **memory profiling**, **checkpoint sharding**, and **network** on your cluster. + +--- + +## 10. Batch size relationships + +Let: + +- \(B_{\text{micro}}\) = micro-batch size per forward/backward **per data-parallel rank** (per step inside accumulation), +- \(D\) = **data parallel size** (ranks that share the same model split for DP), +- \(G\) = **gradient accumulation** steps, +- \(B_{\text{global}}\) = **global batch size** across all DP ranks for one optimizer update. + +Then: + +\[ +B_{\text{global}} = B_{\text{micro}} \times D \times G +\] + +**Data parallel size** from world size (when using TP, PP, EP): + +\[ +D = \frac{\text{world\_size}}{\text{TP} \times \text{PP} \times \text{EP}} +\] + +(If **context parallelism** is present, the denominator must include **CP** in the same way your trainer defines the mesh.) + +Solve for accumulation: + +\[ +G = \frac{B_{\text{global}}}{B_{\text{micro}} \times D} +\] + +**Practical notes** + +- **Micro batch** drives **per-GPU activation memory** (often linearly in sequence length for attention). +- **Global batch** affects **convergence** and learning dynamics; scaling laws often refer to global batch. +- **Gradient accumulation** increases **time per optimizer step** but **reduces memory** by using smaller \(B_{\text{micro}}\). + +Megatron-specific names for batch arguments appear in [Megatron parameters](../03-configuration-reference/megatron-parameters.md). + +--- + +## Related documentation + +- [NCCL/RCCL collective operations guide](./collective-operations.md)—which collectives each strategy uses. +- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) +- [Environment variables](../03-configuration-reference/environment-variables.md) diff --git a/docs/04-technical-guides/performance-tuning.md b/docs/04-technical-guides/performance-tuning.md new file mode 100644 index 000000000..b9a5c9a9a --- /dev/null +++ b/docs/04-technical-guides/performance-tuning.md @@ -0,0 +1,242 @@ +# Performance tuning guide + +This guide covers AMD-focused performance work in Primus: HipBLASLt autotuning for GEMMs, **Primus-Turbo** optional kernels, mixed precision, activation recomputation, communication overlap, memory settings, and MoE-specific flags. It references Primus examples and Megatron module YAMLs. + +--- + +## 1. HipBLASLt autotuning + +Transformer Engine and GEMM-heavy training benefit from HipBLASLt kernel selection. Primus integrates a **three-stage** workflow controlled by `PRIMUS_HIPBLASLT_TUNING_STAGE` (see `examples/README.md` and `examples/run_pretrain.sh`). + +> **Activate tuning first.** The stage variable is only honored when the master switch `PRIMUS_HIPBLASLT_TUNING=1` is set (and `PRIMUS_DETERMINISTIC` is not `1`). Without `PRIMUS_HIPBLASLT_TUNING=1`, both `run_pretrain.sh` and the CLI hook `runner/helpers/hooks/train/pretrain/prepare_experiment.sh` skip tuning entirely and force `TE_HIPBLASLT_TUNING_RUN_COUNT=0` / `TE_HIPBLASLT_TUNING_ALGO_COUNT=0`. Export `PRIMUS_HIPBLASLT_TUNING=1` alongside the stage in every command below. + +### Stage 0 (default) + +No tuning: + +```bash +export PRIMUS_HIPBLASLT_TUNING_STAGE=0 # default +``` + +### Stage 1: Dump GEMM shapes + +Run a **short** training job so shapes are collected during real forward/backward passes. Reduce `train_iters` (or equivalent) for faster shape collection. + +```bash +export PRIMUS_HIPBLASLT_TUNING=1 +export PRIMUS_HIPBLASLT_TUNING_STAGE=1 +./runner/primus-cli direct -- train pretrain \ + --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +``` + +Output layout (from `examples/README.md`): + +- `./output/tune_hipblaslt/${PRIMUS_MODEL}/gemm_shape` + +### Stage 2: Offline tuning + +Runs offline tuning from dumped shapes (often 10–30 minutes depending on model and shapes): + +```bash +export PRIMUS_HIPBLASLT_TUNING=1 +export PRIMUS_HIPBLASLT_TUNING_STAGE=2 +./runner/primus-cli direct -- train pretrain \ + --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +``` + +Expected output: + +- `./output/tune_hipblaslt/${PRIMUS_MODEL}/gemm_tune/tune_hipblas_gemm_results.txt` + +### Stage 3: Train with tuned kernels + +Point the runtime at the tuned override file: + +```bash +export PRIMUS_HIPBLASLT_TUNING=1 +export PRIMUS_HIPBLASLT_TUNING_STAGE=3 +export HIPBLASLT_TUNING_OVERRIDE_FILE=/path/to/tune_hipblas_gemm_results.txt +./runner/primus-cli direct -- train pretrain \ + --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml +``` + +### Related environment variables + +| Variable | Role | +|----------|------| +| `TE_HIPBLASLT_TUNING_ALGO_COUNT` | Breadth of algorithm search for TE HipBLASLt tuning (see `examples/run_pretrain.sh` defaults). | +| `TE_HIPBLASLT_TUNING_RUN_COUNT` | Number of benchmark runs per shape during TE tuning. | +| `TE_HIPBLASLT_TUNING_ALGO_FILE` | Optional algorithm file for TE tuning flows. | +| `TE_HIPBLASLT_TUNING` | When set, interacts with deterministic mode; avoid conflicting settings with shape dump (see script comments in `examples/run_pretrain.sh`). | +| `HIPBLASLT_TUNING_OVERRIDE_FILE` | Override file for stage 3 training. | + +### Standalone offline tool + +For manual HipBLASLt bench workflows, see `examples/offline_tune/offline_tune_gemm.py` and `examples/offline_tune/README.md` (hipblaslt-bench integration and `HIPBLASLT_TUNING_OVERRIDE_FILE` usage). + +--- + +## 2. Primus-Turbo optimization + +**Primus-Turbo** is a separate package of optimized AMD GPU kernels used by Primus Megatron and TorchTitan integrations. It is controlled by the master flag `enable_primus_turbo` in Megatron configs (`primus/configs/modules/megatron/primus_turbo.yaml` extends into trainer/model as needed). **You must install the external `primus_turbo` package** for these paths to be available. + +### Master flag (Megatron) + +```yaml +enable_primus_turbo: true +``` + +### Feature flags (Megatron) + +Defaults in `primus/configs/modules/megatron/primus_turbo.yaml` are mostly `false` until enabled. + +| Flag | Purpose | +|------|---------| +| `use_turbo_attention` | Optimized attention kernels. | +| `use_turbo_parallel_linear` | Optimized tensor-parallel linear layers. | +| `use_turbo_grouped_gemm` | Optimized grouped GEMM for MoE. | +| `use_turbo_grouped_mlp` | Removed—use `use_turbo_grouped_gemm` (passing this key now raises an error). | +| `use_turbo_rms_norm` | Optimized RMSNorm. | +| `moe_use_fused_router_with_aux_score` | Fused MoE router (requires Primus-Turbo backend; see [Backend Patch Notes](../06-developer-guide/backend-patch-notes.md)). | +| `use_turbo_deepep` | DeepEP token dispatcher; set with `enable_primus_turbo: true`. | +| `turbo_deepep_num_cu` | Compute units for DeepEP (patch notes suggest practices such as 64 or 80 for EP8, 32 for EP16–64). | +| `turbo_sync_free_moe_stage` | Sync-free MoE stages (`0`–`3`; `0` disables, stage `2` recommended for performance per patch notes). See [MoE training deep-dive](./moe-training.md). | +| `use_turbo_fused_act_with_probs` | Fused activation with probabilities to reduce redundant work. | + +### Feature flags (TorchTitan) + +TorchTitan presets include `primus_turbo` in `primus/configs/modules/torchtitan/pre_trainer.yaml` + +Example keys: + +```yaml +primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + use_turbo_async_tp: true + use_turbo_float8_linear: true + use_turbo_grouped_mm: false +``` + +### Documentation + +Extended Megatron arguments and Turbo-related behavior are summarized in [Backend Patch Notes](../06-developer-guide/backend-patch-notes.md). + +--- + +## 3. Mixed precision training + +### Megatron (`trainer_base.yaml` patterns) + +| Setting | Description | +|---------|-------------| +| `bf16: true` | BFloat16 training (default `true` in `trainer_base.yaml`). | +| `fp16: false` | FP16 training (optional). | +| `fp8` | FP8 recipe control (`null` / recipes such as delayed scaling in upstream Megatron). | +| `fp8_recipe`, `fp8_margin`, `fp8_interval` | FP8 scaling behavior. | +| `fp4`, `fp4_recipe` | Experimental FP4 paths. | +| `first_last_layers_bf16: true` | Keep first/last layers in BF16 for stability (`num_layers_at_start_in_bf16`, `num_layers_at_end_in_bf16` fine-tune). | + +### TorchTitan + +| Setting | Location | +|---------|----------| +| `training.mixed_precision_param: bfloat16` | `pre_trainer.yaml` default | +| `training.mixed_precision_reduce: float32` | Reduce precision | +| FP8 / quantization | `quantize.linear.float8.*` in `primus/configs/modules/torchtitan/quantize.yaml` | + +### Loss fusion (Megatron model) + +From `primus/configs/models/megatron/language_model.yaml`: + +- Default: `cross_entropy_loss_fusion: false` with `cross_entropy_fusion_impl: "native"`. +- To use Transformer Engine fused cross entropy where supported, enable the fusion explicitly and set `cross_entropy_fusion_impl: "te"`. + +--- + +## 4. Activation recomputation + +### Megatron + +| Parameter | Typical values | Notes | +|-----------|----------------|-------| +| `recompute_granularity` | `full`, `selective` | `full` recomputes more; max memory savings. | +| `recompute_method` | `uniform`, `block` | How recomputation is distributed. | +| `recompute_num_layers` | integer | Layers to recompute when using selective/uniform strategies. | +| `recompute_layer_ids` | list or null | Primus extension: **global** layer indices from `0` to `num_layers - 1` (the patch resolves block-local indices to global ids via `layer_offset`). Use with `recompute_granularity: full` and supported recompute methods. | + +### TorchTitan + +| Parameter | Location | +|-----------|----------| +| `activation_checkpoint.mode` | `none` in default `pre_trainer.yaml` | +| `activation_checkpoint.selective_ac_option` | Selective AC options | + +--- + +## 5. Communication overlap + +### Megatron + +From `trainer_base.yaml` and model settings: + +| Setting | Purpose | +|---------|---------| +| `overlap_grad_reduce` | Overlap gradient reduction with backward. | +| `overlap_param_gather` | Overlap parameter gather with forward. | +| `overlap_p2p_comm` | Pipeline P2P overlap. | +| `async_tensor_model_parallel_allreduce` | Async TP all-reduce (model config). | + +### TorchTitan + +| Setting | Purpose | +|---------|---------| +| `parallelism.enable_async_tensor_parallel: true` | Async tensor parallelism. | + +### Environment + +`CUDA_DEVICE_MAX_CONNECTIONS=1` is commonly required for **correct** overlap behavior in TP/PP stacks (see `docs/03-configuration-reference/environment-variables.md` and Megatron tests). Primus launch scripts or your cluster setup may set this. + +--- + +## 6. Memory optimization + +### Megatron + +| Parameter | Purpose | +|-----------|---------| +| `optimizer_cpu_offload: true` | Offload optimizer state to CPU. | +| `optimizer_offload_fraction` | Fraction to offload (`1.0` in `trainer_base.yaml`). | +| `use_distributed_optimizer: true` | Shards optimizer state across DP ranks (when enabled). | +| `empty_unused_memory_level` | Aggressive emptying of unused memory (`0` default). | +| `global_batch_size` + `micro_batch_size` | Increase global batch via **gradient accumulation** without increasing per-step activation memory. | + +### TorchTitan + +| Parameter | Purpose | +|-----------|---------| +| `training.enable_cpu_offload` | CPU offload path in `pre_trainer.yaml` | + +--- + +## 7. MoE-specific optimization + +### Megatron (model + turbo) + +| Setting | Purpose | +|---------|---------| +| `moe_permute_fusion: true` | Fuse permutation / unpermutation (`patch-notes.md`). | +| `moe_use_fused_router_with_aux_score: true` | Fused router + aux loss (Primus-Turbo). | +| `use_turbo_deepep: true` | DeepEP dispatcher (`enable_primus_turbo` must be true). | +| `turbo_sync_free_moe_stage: 2` | Recommended stage for sync-free MoE (per patch notes). | +| `overlap_moe_expert_parallel_comm: true` | Overlap expert parallel communication (`trainer_base.yaml`). | + +--- + +## Quick checklist + +1. **GEMMs:** run HipBLASLt stages 1–3 or use `offline_tune_gemm.py` for custom workflows. +2. **Kernels:** enable `primus_turbo` after installing `primus_turbo`; turn on attention/MoE flags as needed. +3. **Precision:** BF16 by default; add FP8/FP4 only with recipe testing. +4. **Memory:** recomputation + distributed optimizer + CPU offload + accumulation before buying more GPUs. +5. **MoE:** fusion + DeepEP + sync-free stages + EP comm overlap when supported. diff --git a/docs/04-technical-guides/profiling-and-observability.md b/docs/04-technical-guides/profiling-and-observability.md new file mode 100644 index 000000000..00040b8f5 --- /dev/null +++ b/docs/04-technical-guides/profiling-and-observability.md @@ -0,0 +1,170 @@ +# Profiling and observability + +This guide covers how to capture and analyze performance data in Primus: the PyTorch/Kineto profiler, GPU memory snapshots, AMD's TraceLens trace analysis, ROCm memory sampling, memory/performance projection, and pipeline-schedule visualization. Parameters are grounded in `primus/configs/modules/megatron/`, `primus/configs/modules/torchtitan/pre_trainer.yaml`, and `primus/configs/modules/maxtext/pre_trainer.yaml`. + +--- + +## 1. Torch profiler (Megatron) + +Megatron training integrates the PyTorch profiler. Defaults live in `primus/configs/modules/megatron/trainer_base.yaml` and `primus_megatron_module.yaml`. + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `profile` | `false` | Master switch for profiling. | +| `use_pytorch_profiler` | `false` | Use the PyTorch (Kineto) profiler path. | +| `profile_ranks` | `[0]` | Which global ranks to profile. | +| `profile_step_start` | `10` | First step to capture. | +| `profile_step_end` | `12` | Last step to capture. | +| `disable_profiler_activity_cpu` | `false` | Drop CPU-side activity to shrink traces (GPU-only trace). | +| `torch_profiler_record_shapes` | `true` | Record tensor shapes per op. | +| `torch_profiler_with_stack` | `true` | Record Python/C++ stacks (larger traces). | +| `torch_profiler_use_gzip` | `false` | Gzip the exported trace. | + +Enable via CLI overrides on `train pretrain`: + +```bash +./runner/primus-cli direct -- train pretrain \ + --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml \ + --profile True \ + --use_pytorch_profiler True \ + --profile_step_start 5 \ + --profile_step_end 6 \ + --disable_profiler_activity_cpu False +``` + +Keep the capture window **short** (a few steps after warmup)—traces grow quickly, especially with `with_stack` and CPU activity enabled. Load the resulting trace in [Perfetto](https://ui.perfetto.dev/) to inspect CPU/GPU overlap, kernel launch delays, and idle gaps. + +--- + +## 2. GPU memory profiling (Megatron) + +Two complementary mechanisms: + +**Memory history snapshot** (`trainer_base.yaml`): + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `record_memory_history` | `false` | Record the CUDA/HIP allocator history for snapshot analysis. | +| `memory_snapshot_path` | `snapshot.pickle` | Output path for the allocator snapshot. | + +Load the pickle with PyTorch's memory visualizer to find fragmentation and peak allocations. + +**ROCm memory sampling** (`primus_megatron_module.yaml`): + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `use_rocm_mem_info` | `false` | When `true`, collect ROCm memory info via `rocm-smi` **every** iteration. | +| `use_rocm_mem_info_iters` | `[1, 2]` | When `use_rocm_mem_info=false`, only sample at these iterations. | + +Also relevant: `log_memory_to_tensorboard` (`trainer_base.yaml`) writes memory metrics to TensorBoard. + +--- + +## 3. TraceLens automated trace analysis (Megatron) + +[TraceLens](https://github.com/AMD-AGI/TraceLens) turns raw profiler traces into hierarchical breakdowns (roofline/efficiency, compute-vs-memory bound kernels, communication-vs-sync separation, trace diffing). Primus can generate and optionally upload these reports. Configured in `primus/configs/modules/megatron/primus_megatron_module.yaml`: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `generate_tracelens_report` | `false` | Generate TraceLens reports locally (auto-enabled when upload is on). | +| `mlflow_upload_tracelens_report` | `false` | Upload reports to MLflow (auto-enables generation, profiling, tensorboard). | +| `mlflow_tracelens_ranks` | `null` | Ranks to analyze (`null` = all; e.g. `[0, 8]` for one rank/node). | +| `mlflow_tracelens_output_format` | `xlsx` | `xlsx` (fastest), `csv`, or `all`. | +| `mlflow_tracelens_cleanup_after_upload` | `false` | Delete local reports after upload to save disk. | +| `mlflow_tracelens_auto_install` | `true` | Auto-install TraceLens if missing (set `false` to disable). | + +Related profiler/log uploads: `mlflow_upload_traces` (upload raw trace files) and `mlflow_upload_logs` (upload training logs). See [Logging & experiment tracking](./logging-and-experiment-tracking.md) for MLflow setup. + +--- + +## 4. Performance metrics to MLflow (Megatron) + +`mlflow_upload_performance_metrics: false` (`primus_megatron_module.yaml`) enables a comprehensive scaling-test metric set when turned on (implicitly enabling throughput calculation): + +- `perf/throughput_tflops_per_gpu`, `perf/tps_tokens_per_sec_per_gpu`, `perf/iteration_time_ms` +- `perf/{rocm,hip}_current_mem_gb`, `perf/{rocm,hip}_mem_utilization_pct` +- `perf/gpu_utilization_pct_rank{N}`, `perf/gpu_utilization_pct_avg` + +> GPU utilization collection uses an `all_gather` every `log_interval`, which synchronizes ranks—keep this in mind for throughput-sensitive runs. + +--- + +## 5. Profiling (TorchTitan) + +Configured under `profiling:` in `primus/configs/modules/torchtitan/pre_trainer.yaml`: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `enable_profiling` | `false` | Enable the Torch profiler. | +| `profile_freq` | `10` | Capture every N steps. | +| `save_traces_folder` | `profile_traces` | Output folder for traces. | +| `enable_memory_snapshot` | `false` | Capture GPU memory snapshots. | +| `save_memory_snapshot_folder` | `memory_snapshot` | Output folder for snapshots. | + +Communication tracing is configured under `comm:` (`trace_buf_size`, `save_traces_folder: comm_traces`). + +--- + +## 6. Profiling (MaxText) + +Configured in `primus/configs/modules/maxtext/pre_trainer.yaml`: + +| Parameter | Default | Purpose | +|-----------|---------|---------| +| `profiler` | `xplane` | Profiler backend (XPlane traces). | +| `skip_first_n_steps_for_profiler` | `3` | Warmup steps to skip before capture. | +| `profiler_steps` | `1` | Number of steps to capture. | + +--- + +## 7. Memory and performance projection + +Project resource usage **before** launching, without consuming a full cluster. Exposed through the Primus CLI `projection` subcommand (`primus/cli/subcommands/projection.py`). + +```bash +# Memory projection from an experiment config +./primus-cli direct -- projection memory --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml + +# Performance projection (single-node benchmarking) +./primus-cli direct -- projection performance --config .yaml + +# Performance projection scaled to N nodes (simulation) +./primus-cli direct -- projection performance --config .yaml --target-nodes 4 +``` + +Memory projection breaks VRAM down across parameters, gradients, activations, optimizer states, and mixed-precision overhead—invaluable for MoE/ultra-large models where activations dominate. See [Projection](../02-user-guide/projection.md) for the full reference. + +--- + +## 8. Pipeline schedule visualization + +Diagnose pipeline bubbles and stage imbalance with the built-in tool `tools/visualization/pp_vis/`. + +1. Dump per-rank schedule data during training with `--dump_pp_data true` (Megatron flag `dump_pp_data`, `primus_megatron_module.yaml`). Output lands under `output/pp_data/` (`config.json`, `pp_rank_*.json`). +2. Install the tool deps and run the viewer: + +```bash +pip install -r tools/visualization/pp_vis/requirements.txt +python tools/visualization/pp_vis/vis.py # open http://127.0.0.1:8988 +``` + +Configure `task_list` in `vis.py` to point at your dumped `log_path` and the iterations to render. The tool can also visualize the PP simulator output (see `tools/visualization/pp_vis/README.md`). + +--- + +## 9. Recommended workflow + +1. **Project first**—run `projection memory` to confirm the config fits before booking GPUs. +2. **Capture a short trace**—a few steps after warmup with `profile` + `use_pytorch_profiler`. +3. **Inspect**—Perfetto for the timeline; TraceLens for automated hierarchical analysis and trace diffs. +4. **Check memory**—`record_memory_history` / ROCm sampling for fragmentation and peaks. +5. **Check pipelines**—`dump_pp_data` + `pp_vis` for bubbles and stage imbalance. + +--- + +## Related documentation + +- [Logging & experiment tracking](./logging-and-experiment-tracking.md)—WandB / TensorBoard / MLflow setup. +- [MoE training deep-dive](./moe-training.md)—applying this workflow to sparse models. +- [Performance tuning](./performance-tuning.md)—what to change once you've found the bottleneck. +- [Projection](../02-user-guide/projection.md) and [Monitoring and logging](../05-operations/monitoring-logging.md). diff --git a/docs/05-operations/README.md b/docs/05-operations/README.md new file mode 100644 index 000000000..6b9c1e016 --- /dev/null +++ b/docs/05-operations/README.md @@ -0,0 +1,12 @@ +# Operations + +Production deployment and operational guidance. + +- [Deployment](deployment.md): container, Slurm, and Kubernetes deployment +- [Monitoring and logging](monitoring-logging.md): WandB, TensorBoard, MLflow, Primus logging +- [Troubleshooting](troubleshooting.md): common failures, diagnostics, and fixes +- [Security](security.md): secrets handling, container security, dependencies + +--- + +[← Documentation home](../README.md) diff --git a/docs/05-operations/deployment.md b/docs/05-operations/deployment.md new file mode 100644 index 000000000..4c6797390 --- /dev/null +++ b/docs/05-operations/deployment.md @@ -0,0 +1,237 @@ +# Deployment guide + +This guide describes how to deploy Primus training across **container**, **direct (bare metal)**, and **Slurm** environments using the unified `primus-cli` launcher. For environment variable semantics, see [Environment variables](../03-configuration-reference/environment-variables.md). For YAML hierarchy and precedence, see [Configuration system](../02-user-guide/configuration-system.md). + +--- + +## 1. Deployment overview + +Primus supports three deployment modes: + +| Mode | Description | Typical use | +|------|-------------|-------------| +| **Container** | Docker/Podman with ROCm-capable GPU devices and capabilities | Recommended default; reproducible images | +| **Direct** | Runs on the current host (or inside an existing container) | Local debugging, single-node, clusters with ROCm on nodes | +| **Slurm** | Wraps `srun`/`sbatch` and launches per-node entry scripts | Multi-node clusters with Slurm | + +**Container image:** `docker.io/rocm/primus:v26.3` (default in `runner/.primus.yaml`). For clusters using **AINIC**, use `runner/use_ainic.yaml` and tune the image and NCCL-related variables (for example `USING_AINIC`, `NCCL_IB_GID_INDEX`) to match your fabric. + +**Prerequisites (baseline):** + +- **AMD ROCm** >= 7.0 on the host (or in the image when using containers) +- **Docker** or **Podman** >= 24.0 when using container mode +- **AMD Instinct** GPUs and working ROCm stack (`rocm-smi` should report devices) + +--- + +## 2. Container deployment + +### 2.1 Pull the image + +```bash +docker pull docker.io/rocm/primus:v26.3 +``` + +The default `container.options.image` in `runner/.primus.yaml` is `rocm/primus:v26.3` (equivalent to `docker.io/rocm/primus:v26.3` when the registry is omitted). + +### 2.2 Required device mounts + +System defaults (`runner/.primus.yaml`, `container.options.device`) pass each path as `--device` to the runtime: + +| Device | Purpose | +|--------|---------| +| `/dev/kfd` | Kernel Fusion Driver (ROCm core) | +| `/dev/dri` | Direct Rendering Infrastructure (GPU access) | +| `/dev/infiniband` | InfiniBand character devices (multi-node / RDMA) | + +### 2.3 Required capabilities + +Defaults (`container.options.cap-add`): + +| Capability | Purpose | +|------------|---------| +| `SYS_PTRACE` | Debugging and profiling tools | +| `CAP_SYS_ADMIN` | Administrative operations required by some ROCm/GPU workflows | + +### 2.4 Container runtime options + +Defaults in `runner/.primus.yaml` include: + +| Option | Value | +|--------|--------| +| `ipc` | `host` | +| `network` | `host` | +| `privileged` | `true` | +| `security-opt` | `seccomp=unconfined` | +| `group-add` | `video` | + +`primus-cli-container.sh` always mounts the **Primus repository root** into the container at the same path (`-v $PRIMUS_PATH:$PRIMUS_PATH`). Mount additional paths for **datasets**, **model weights**, and **outputs** with `--volume` (or `container.options.volume` in YAML). + +### 2.5 Environment passthrough + +`container.options.env` lists names that are forwarded into the **inner** `primus-cli` invocation as `--env` when set on the host (see `runner/.primus.yaml`). Examples include: + +`MASTER_ADDR`, `MASTER_PORT`, `NNODES`, `NODE_RANK`, `GPUS_PER_NODE`, `DOCKER_IMAGE`, `HF_TOKEN`, `WANDB_API_KEY`, `ENABLE_NUMA_BINDING`, `USING_AINIC`, and NCCL/GLOO socket and IB-related variables (`NCCL_IB_HCA`, `NCCL_SOCKET_IFNAME`, `GLOO_SOCKET_IFNAME`, `NCCL_IB_GID_INDEX`, and others). + +Additionally, `primus-cli-container.sh` auto-forwards any environment variable whose name starts with `PRIMUS_`, `NCCL_`, `RCCL_`, `GLOO_`, `IONIC_`, or `HIPBLASLT_` when present on the host. + +### 2.6 Single-node example + +```bash +./primus-cli container --volume /data:/data -- train pretrain --config /data/exp.yaml +``` + +### 2.7 Multi-node container deployment + +Set `MASTER_ADDR`, `MASTER_PORT`, `NNODES`, `NODE_RANK`, and `GPUS_PER_NODE` on each node (Slurm or your orchestrator sets these; see `runner/primus-cli-slurm-entry.sh`). Example pattern when launching manually: + +```bash +export MASTER_ADDR= +export MASTER_PORT=1234 +export NNODES=4 +export NODE_RANK=<0-based index for this node> +export GPUS_PER_NODE=8 + +./primus-cli container -- train pretrain --config /path/to/config.yaml +``` + +Use `--clean` before launch to remove existing containers (`primus-cli-container.sh`). + +--- + +## 3. Slurm deployment + +### 3.1 `srun` (interactive or blocking) + +```bash +./primus-cli slurm srun -N -p -- train pretrain --config +``` + +The Slurm entry script invokes the container launcher on each allocated node. Set the image through `runner/.primus.yaml`, a custom launcher config file, or site policy; the default is `rocm/primus:v26.3`. + +### 3.2 `sbatch` (batch jobs) + +```bash +./primus-cli slurm sbatch -N -p --time --job-name -o -- \ + train pretrain --config +``` + +Add `-e ` if you want separate stderr. + +### 3.3 Slurm-to-Primus environment mapping + +`runner/primus-cli-slurm-entry.sh` sets: + +| Variable | Source (typical) | +|----------|-------------------| +| `NNODES` | `SLURM_NNODES`, or `SLURM_JOB_NUM_NODES`, or existing `NNODES` | +| `NODE_RANK` | `SLURM_NODEID`, or `SLURM_PROCID`, or existing `NODE_RANK` | +| `GPUS_PER_NODE` | Default `8` if unset | +| `MASTER_ADDR` | First host in `SLURM_NODELIST` if unset | +| `MASTER_PORT` | Default `1234` if unset | + +The entry script then exports `MASTER_ADDR`, `MASTER_PORT`, `NNODES`, `NODE_RANK`, and `GPUS_PER_NODE` into the container launcher. + +### 3.4 Slurm YAML defaults (`runner/.primus.yaml`) + +| Key | Default | +|-----|---------| +| `slurm.nodes` | `1` | +| `slurm.gpus_per_node` | `8` | +| `slurm.time` | `"4:00:00"` | +| `slurm.partition` | (commented; set per site) | + +CLI Slurm flags override YAML when both are specified (see `runner/primus-cli-slurm.sh`). + +### 3.5 Entry after the first `--` + +Production examples pass the Primus Python command after the Slurm `--` separator, for example: + +```bash +./primus-cli slurm srun -N 4 -p gpu -- train pretrain --config exp.yaml +``` + +The shipped `primus-cli-slurm-entry.sh` invokes **`primus-cli-container.sh`** with distributed variables set from Slurm. Container options should come from launcher configuration instead of a literal `container` token in the inner command. For **bare-metal** nodes without Docker, run `primus-cli direct` under your allocation and ensure the same distributed variables and ROCm layout as in [Multi-node configuration](#5-multi-node-configuration). + +--- + +## 4. Kubernetes deployment + +Kubernetes integration is **not** shipped as a Helm chart or operator in this repository. The repo includes **`examples/run_k8s_pretrain.sh`**, a client script that talks to a Kubernetes **API** to create and manage training workloads (image default `docker.io/rocm/primus:v26.3`). + +Use that script as a reference for your platform; adapt networking, storage, and scheduling to your cluster policies. + +--- + +## 5. Multi-node configuration + +Required variables for distributed training: + +| Variable | Role | +|----------|------| +| `MASTER_ADDR` | Hostname or IP of rank-0 process | +| `MASTER_PORT` | TCP port for the process group rendezvous | +| `NNODES` | Number of nodes | +| `NODE_RANK` | Zero-based index of this node | +| `GPUS_PER_NODE` | GPUs per node used by `torchrun` | + +**Flow:** User or Slurm sets the environment → `primus-cli` and `primus-cli-direct.sh` load GPU and comm settings → **`torchrun`** launches `primus/cli/main.py` with the distributed topology. + +**Defaults from `runner/.primus.yaml` (`direct` section):** + +| Key | Default | +|-----|---------| +| `direct.master_port` | `1234` | +| `direct.gpus_per_node` | `8` | +| `direct.nnodes` | `1` | +| `direct.master_addr` | `"localhost"` | + +--- + +## 6. Startup and shutdown + +**Lifecycle (high level):** + +1. Parse CLI and load YAML (`--config` chain: see [Configuration system](../02-user-guide/configuration-system.md)). +2. Load environment (GPU detection, hooks, patches in `primus-cli-direct.sh`). +3. Launch training via **`torchrun`** into the Python CLI. + +**Verification:** + +- `--dry-run` prints the command that would run without executing (supported in container and Slurm scripts). +- `--debug` sets `PRIMUS_LOG_LEVEL=DEBUG` for verbose launcher and shell logging. + +**Shutdown:** + +- Normal completion or **Ctrl+C** terminates the training process. +- In container mode, **`--clean`** removes existing containers before launch (`primus-cli-container.sh`). + +**Timeouts (config):** + +| Backend | Parameter | Location | +|---------|-----------|----------| +| Megatron | `distributed_timeout_minutes` | `primus/configs/modules/megatron/trainer_base.yaml` (default `10`) | +| TorchTitan | `comm.init_timeout_seconds` | `primus/configs/modules/torchtitan/pre_trainer.yaml` (default `300`) | + +--- + +## 7. Production checklist + +| Item | Action | +|------|--------| +| ROCm drivers | Install and verify with `rocm-smi` | +| Container image | Pulled and aligned with host ROCm expectations | +| Network | Run `preflight --network` (see [Preflight](../02-user-guide/preflight.md)) | +| Shared data | Paths visible and consistent on all nodes | +| Hugging Face | Set `HF_TOKEN` if using gated models | +| Checkpoints | Save directory on shared or replicated storage with sufficient space | +| Monitoring | Configure WandB or TensorBoard (see [Monitoring and Logging](./monitoring-logging.md)) | +| Resources | Slurm time limits, partitions, and GPU counts match your YAML and hardware | + +--- + +## Related documentation + +- [CLI reference](../02-user-guide/cli-reference.md) +- [Troubleshooting](./troubleshooting.md) +- [Multi-node networking](../04-technical-guides/multi-node-networking.md) diff --git a/docs/05-operations/monitoring-logging.md b/docs/05-operations/monitoring-logging.md new file mode 100644 index 000000000..9e93e01a3 --- /dev/null +++ b/docs/05-operations/monitoring-logging.md @@ -0,0 +1,222 @@ +# Monitoring and logging + +This page summarizes how Primus configures application logging, experiment tracking (Weights and Biases, TensorBoard, MLflow), training metrics, profilers, ROCm memory probes, and how to capture a reproducible configuration snapshot. + +--- + +## 1. Primus logging system + +Primus uses **loguru** for structured logging. Initialization wires **file sinks** (per log level) and a **stderr** sink, binds experiment and distributed context (`team`, `user`, `exp`, `module_name`, `node_ip`, `rank`, `world_size`), and installs an **intercept handler** so legacy `logging` output from frameworks such as Megatron is forwarded to loguru with consistent formatting. + +**Rank-aware behavior** + +- Worker processes write under `{exp_root}/logs/{module_name}/rank-{rank}/` with separate rotated files for `debug`, `info`, `warning`, and `error` (subject to `file_sink_level`). +- The launcher **master** process can use `logs/master/` when the master logger is configured with `is_head=True`. + +**Levels from module configuration** (`primus/configs/modules/module_base.yaml`) + +| Parameter | Default | Role | +|-----------|---------|------| +| `sink_level` | `null` | If set, overrides both file and stderr sink levels. | +| `file_sink_level` | `DEBUG` | Minimum level for file sinks when `sink_level` is unset. | +| `stderr_sink_level` | `INFO` | Minimum level for stderr when `sink_level` is unset. | + +`init_worker_logger` in `primus/core/runtime/logging.py` reads `sink_level`, `file_sink_level`, and `stderr_sink_level` from the merged module config. The Megatron trainer maps `stderr_sink_level` to Megatron’s numeric `logging_level` (deprecated `logging_level` in `trainer_base.yaml` is replaced by this mapping). + +**Shell / runner environment** (see `docs/03-configuration-reference/environment-variables.md`) + +| Variable | Purpose | +|----------|---------| +| `PRIMUS_LOG_LEVEL` | Runner verbosity: `DEBUG`, `INFO`, `WARN`, `ERROR` (default `INFO`). | +| `PRIMUS_LOG_TIMESTAMP` | `1` enables timestamps on runner logs; `0` disables. | +| `PRIMUS_LOG_COLOR` | `1` enables ANSI colors when appropriate; often `0` in non-TTY contexts. | + +**CLI** + +- `primus-cli --debug` sets `PRIMUS_LOG_LEVEL=DEBUG` so launcher and shell logging are verbose (see `docs/02-user-guide/cli-reference.md`). + +--- + +## 2. Weights and biases + +### Megatron (`primus/configs/modules/megatron/trainer_base.yaml`, `primus_megatron_module.yaml`) + +Defaults in `primus_megatron_module.yaml` disable WandB; trainer fields in `trainer_base.yaml` supply names and paths when enabled. + +| Parameter | Default (module / trainer) | Description | +|-----------|----------------------------|-------------| +| `disable_wandb` | `true` (`primus_megatron_module.yaml`) | Master switch; when `false`, Primus sets paths and default project/run names from experiment metadata. | +| `wandb_project` | `null` | If unset when WandB is enabled, defaults to `{work_group}_{user_name}`. | +| `wandb_exp_name` | `null` | If unset, defaults to `exp_name`. | +| `wandb_entity` | `null` | Optional WandB entity/team. | +| `wandb_save_dir` | `null` | Deprecated in favor of `{exp_root}`; artifacts use `{exp_root}/wandb`. | + +**Environment** + +- `WANDB_API_KEY` is **required** when WandB is enabled; Primus emits a warning if it is missing (`primus/backends/megatron/patches/args/wandb_config_patches.py`). + +### TorchTitan (`primus/configs/modules/torchtitan/pre_trainer.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `metrics.enable_wandb` | `false` | Enables WandB in the TorchTitan metrics stack. | + +When enabled, `primus/backends/torchtitan/patches/wandb_patches.py` can set `WANDB_PROJECT` and `WANDB_RUN_NAME` from Primus experiment metadata if unset. Use `WANDB_API_KEY` for authentication. + +--- + +## 3. TensorBoard + +### Megatron + +**Module toggles** (`primus_megatron_module.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `disable_tensorboard` | `true` | When `false`, TensorBoard output is placed under `{exp_root}/tensorboard` (Primus overrides deprecated `tensorboard_dir` with this path). | + +**Trainer** (`trainer_base.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `tensorboard_log_interval` | `1` | Steps between TensorBoard writes. | +| `tensorboard_queue_size` | `1000` | Event file queue size. | +| `log_timers_to_tensorboard` | `false` | Log timer stats. | +| `log_batch_size_to_tensorboard` | `false` | Log batch size. | +| `log_learning_rate_to_tensorboard` | `true` | Log learning rate. | +| `log_validation_ppl_to_tensorboard` | `false` | Log validation perplexity. | +| `log_memory_to_tensorboard` | `false` | Log memory stats. | +| `log_world_size_to_tensorboard` | `false` | Log world size. | +| `log_loss_scale_to_tensorboard` | `true` | Log loss scale. | +| `tensorboard_dir` | `null` | Deprecated; Primus sets the directory under `exp_root`. | + +**Note:** Enabling Megatron **profiling** (`profile: true`) forces `disable_tensorboard` off in `update_primus_config` so TensorBoard is available for profile-related views. + +### TorchTitan (`pre_trainer.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `metrics.enable_tensorboard` | `false` | Enables TensorBoard logging. | +| `metrics.save_tb_folder` | `tb` | Subfolder name (typically under the job dump directory in TorchTitan layouts). | + +**Launch TensorBoard locally** + +```bash +tensorboard --logdir +``` + +Point `` at the Megatron `tensorboard` directory under the experiment root, or at the TorchTitan metrics folder that contains the `save_tb_folder` subtree. + +--- + +## 4. MLflow + +MLflow integration is **Megatron-only** in the paths described here. + +**Module** (`primus_megatron_module.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `disable_mlflow` | `true` | When `false`, MLflow run setup runs on the **last** global rank (`world_size - 1`). | +| `mlflow_run_name` | `null` | If unset when enabled, defaults to `{work_group}_{user_name}`. | +| `mlflow_experiment_name` | `null` | Passed to `mlflow.set_experiment` when set. | + +**Startup behavior** (`primus/backends/megatron/training/global_vars.py`) + +- Logs training `args` as parameters. +- Logs filtered environment variables with an `env__` prefix. +- Collects git metadata, sets MLflow source tags, and writes `system/git_metadata.json` as a run artifact. + +**Environment** (typical Databricks / hosted tracking) + +| Variable | Role | +|----------|------| +| `DATABRICKS_HOST` | Checked by the Megatron trainer when MLflow is enabled; a warning is printed if unset. | +| `DATABRICKS_TOKEN` | Authentication for Databricks-hosted tracking (see environment reference). | +| `MLFLOW_TRACKING_URI` | Tracking server URI; optional depending on deployment. | + +--- + +## 5. Training metrics + +### Megatron (`trainer_base.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `log_interval` | `100` | Steps between standard training log lines. | +| `log_throughput` | `false` | Log throughput metrics. | +| `log_avg_skip_iterations` | `2` | Skip initial iterations when computing averages. | +| `log_avg_reset_interval` | `10` | Interval for resetting running averages. | +| `log_params_norm` | `false` | Log parameter norm. | +| `log_num_zeros_in_grad` | `false` | Log count of zero gradients. | +| `log_progress` | `false` | Progress-style logging. | +| `timing_log_level` | `0` | Timing log verbosity. | +| `timing_log_option` | `minmax` | Timing aggregation option. | + +### TorchTitan (`pre_trainer.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `metrics.log_freq` | `10` | Metric logging frequency (steps). | +| `metrics.disable_color_printing` | `false` | Disable colored console metrics. | +| `metrics.save_for_all_ranks` | `false` | Save metrics from every rank vs. reduced ranks. | + +--- + +## 6. Profiling + +### Megatron (`trainer_base.yaml`, `primus_megatron_module.yaml`) + +| Parameter | Source | Default | Description | +|-----------|--------|---------|-------------| +| `profile` | `trainer_base.yaml` | `false` | Enables Megatron profiling path; also forces TensorBoard on when `true`. | +| `use_pytorch_profiler` | `trainer_base.yaml` | `false` | Use PyTorch profiler integration. | +| `profile_ranks` | `trainer_base.yaml` | `[0]` | Ranks to profile. | +| `profile_step_start` | `trainer_base.yaml` | `10` | First step to profile. | +| `profile_step_end` | `trainer_base.yaml` | `12` | Last step to profile. | +| `record_memory_history` | `trainer_base.yaml` | `false` | Record memory history. | +| `memory_snapshot_path` | `trainer_base.yaml` | `snapshot.pickle` | Memory snapshot file name. | +| `disable_profiler_activity_cpu` | `primus_megatron_module.yaml` | `false` | Disable CPU activities in the profiler. | +| `torch_profiler_record_shapes` | `primus_megatron_module.yaml` | `true` | Record tensor shapes. | +| `torch_profiler_with_stack` | `primus_megatron_module.yaml` | `true` | Capture Python stacks. | +| `torch_profiler_use_gzip` | `primus_megatron_module.yaml` | `false` | Gzip profiler traces. | + +### TorchTitan (`pre_trainer.yaml`) + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `profiling.enable_profiling` | `false` | Master profiling toggle. | +| `profiling.profile_freq` | `10` | How often to capture traces. | +| `profiling.enable_memory_snapshot` | `false` | Enable memory snapshots. | +| `profiling.save_memory_snapshot_folder` | `memory_snapshot` | Output folder for snapshots. | +| `profiling.save_traces_folder` | `profile_traces` | Folder for profiler traces. | + +--- + +## 7. ROCm memory monitoring + +Configured in `primus/configs/modules/megatron/primus_megatron_module.yaml` and applied in the Megatron trainer when logging throughput. + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `use_rocm_mem_info` | `false` | When `true`, collect ROCm memory information via `rocm-smi` on **every** iteration that hits the throughput logging branch. | +| `use_rocm_mem_info_iters` | `[1, 2]` | When `use_rocm_mem_info` is `false`, `rocm-smi` runs only on these iteration numbers (same branch). | + +Collection is evaluated where `log_throughput` drives the extended iteration log (see `primus/backends/megatron/patches/training_log/print_rank_last_patches.py`): enable `log_throughput` in `trainer_base.yaml` (or overrides) when you need ROCm memory lines in the training log. + +--- + +## 8. Experiment snapshots + +**On disk (every run)** + +- **Experiment root**: `{workspace}/{work_group}/{user_name}/{exp_name}` is created at config load time (`PrimusConfig`). +- **Per-rank logs**: `{exp_root}/logs/{module_name}/rank-{rank}/` with rotated level-specific files. +- **Checkpoints**: Megatron uses `{exp_root}/checkpoints` (trainer sets `save` to this path). +- **TensorBoard / WandB**: Under `exp_root` as described above when those features are enabled. + +**MLflow** (Megatron, when enabled): Parameters, environment snapshot, and git metadata artifact provide a structured record of the run configuration and repository state. + +**Resolved configuration** + +The launcher and parser accept `--export_config`, but the default core training path (`primus/cli/subcommands/train.py` into `PrimusRuntime`) does not currently write a resolved YAML file. Archive the submitted experiment YAML, any referenced presets, launcher config, and runtime logs with each run. Treat resolved-config export as a legacy or future capability unless your deployment has implemented it on the core runtime path. diff --git a/docs/05-operations/security.md b/docs/05-operations/security.md new file mode 100644 index 000000000..c51d362d0 --- /dev/null +++ b/docs/05-operations/security.md @@ -0,0 +1,154 @@ +# Security considerations + +This document describes security-relevant properties of Primus as a **YAML-driven training framework** for AMD GPUs (ROCm, RCCL, containers). It is intended for operators, platform engineers, and security reviewers. It does not replace organizational policies, threat models, or vendor hardening guides. + +**Related documentation:** [Environment variables](../03-configuration-reference/environment-variables.md), [Installation](../01-getting-started/installation.md), [CLI reference](../02-user-guide/cli-reference.md). + +--- + +## 1. Overview + +| Aspect | Description | +|--------|-------------| +| Role | Primus orchestrates distributed **training** jobs; it is **not** a general user-facing network service. | +| Authentication / authorization | **No built-in** authentication, authorization, or multi-tenant isolation in Primus itself. | +| Responsibility | **Security posture is determined by** the scheduler, container runtime, network, storage, identity systems, and operational practices of the deployment environment. | + +Treat Primus like privileged infrastructure software: run it on appropriately isolated hosts and networks, and govern secrets and data the same way you would for large-scale ML training elsewhere. + +--- + +## 2. Secrets management + +Secrets are commonly passed as **environment variables** consumed by Primus, launchers, or third-party libraries. + +| Variable (examples) | Typical use | +|---------------------|-------------| +| `HF_TOKEN` | Hugging Face token for **gated models** and authenticated downloads. | +| `WANDB_API_KEY` | Weights & Biases API key for experiment logging. | +| `DATABRICKS_HOST` / `DATABRICKS_TOKEN` | Databricks or MLflow-related credentials when those integrations are used. | + +**Practices** + +| Practice | Detail | +|----------|--------| +| Do not hardcode secrets | Avoid putting tokens or passwords directly in YAML, shell history, or committed scripts. | +| Prefer indirection | Use **`${VAR}`** substitution in configs to reference environment-injected values rather than literals. | +| Slurm | Use **`--export`** deliberately; prefer site-specific **secret injection** or **credential helpers** where available. | +| Containers | Pass secrets with **`--env`** or via **`runner/.primus.yaml`** env forwarding—never bake them into images. | +| Rotation | Rotate API keys and tokens on a schedule and after personnel or scope changes. | + +A broader catalog of variables appears in [Environment variables](../03-configuration-reference/environment-variables.md). + +--- + +## 3. Container security + +Primus-oriented container runs often require **elevated access** so ROCm, profilers, and high-performance networking behave correctly. + +**Common high-privilege options** + +| Option | Typical purpose | +|--------|-----------------| +| `--privileged true` | Broad device access (often required for ROCm workflows on some setups). | +| `--cap-add SYS_PTRACE` | Debugging and profiling tooling. | +| `--cap-add CAP_SYS_ADMIN` | Administrative operations expected by parts of the ROCm/tooling stack. | +| `--security-opt seccomp=unconfined` | Relaxes seccomp constraints for compatibility with drivers and tools. | +| `--ipc host` | Shared memory semantics for large tensors and collectives. | +| `--network host` | Host networking—frequently used for **multi-node RCCL** performance and simplicity. | + +**Device access (examples)** + +| Device | Role | +|--------|------| +| `/dev/kfd` | ROCm kernel interface. | +| `/dev/dri` | GPU render nodes. | +| `/dev/infiniband` | InfiniBand character devices when using IB. | + +**Risks** + +| Risk | Why it matters | +|------|----------------| +| Privileged containers | Substantial **host** access; container escape or compromise has high impact. | +| Host networking | Exposes the container to the **host’s network namespace**; services may bind broadly. | +| Shared IPC | Potential for **cross-process interference** or information leakage if workloads share hosts improperly. | + +**Mitigations** + +| Mitigation | Detail | +|------------|--------| +| Dedicated training nodes | Run training on **isolated** machines rather than mixed with user-facing services. | +| Network controls | Apply **firewall rules** and **segmentation** so only required ports and peers are reachable. | +| Trusted images | Pull from **trusted registries**, pin digests, and verify image provenance. | +| Monitoring | Track **CPU, memory, GPU, and network** usage; alert on anomalous processes or egress. | + +--- + +## 4. Third-party dependencies + +Primus integrates **third-party submodules** and Python packages; each carries its own license and maintenance cadence. + +**Representative submodules** + +| Component | Notes (non-exhaustive) | +|-----------|-------------------------| +| Megatron-LM | MIT License; NVIDIA upstream. | +| TorchTitan | Apache 2.0; PyTorch / Meta ecosystem. | +| MaxText | Apache 2.0; Google upstream. | +| Megatron-Bridge | NVIDIA NeMo ecosystem. | +| Emerging-Optimizers | NVIDIA NeMo ecosystem. | +| HummingbirdXT | AMD AGI ecosystem. | + +**Python dependencies** + +Runtime tooling often includes packages such as **loguru**, **wandb**, **nltk**, **matplotlib**, **mlflow**, and others as declared in project requirements—verify the canonical list in the repository’s `requirements.txt` (or lockfile) for your revision. + +**Recommendations** + +| Recommendation | Rationale | +|----------------|-----------| +| Pin versions | Reproducible builds and controlled upgrade paths. | +| Update submodules | Security and correctness fixes flow from upstream projects. | +| Monitor advisories | Subscribe to upstream security notices for frameworks you enable. | + +--- + +## 5. Network security + +| Property | Detail | +|----------|--------| +| RCCL / NCCL traffic | **Not encrypted** at the application layer; assumes a **trusted network path**. | +| Coordination | **`MASTER_ADDR`** and **`MASTER_PORT`** should reside on a **private** or otherwise **trusted** segment. | +| InfiniBand | Often on a **dedicated fabric**; still treat adjacent compromised hosts as in-scope for lateral movement. | +| TLS / mTLS | Primus does **not** provide TLS or mTLS for inter-node training traffic by default. | + +For physical and logical networking topics, see [Multi-node networking](../04-technical-guides/multi-node-networking.md). + +--- + +## 6. Data security + +| Asset | Consideration | +|-------|-----------------| +| Training data | May include **PII**, licensed corpora, or export-controlled material—classify and restrict accordingly. | +| Checkpoints | Contain **full model state**; treat as sensitive intellectual property. | +| Storage permissions | Use **least privilege** on shared filesystems and object stores. | +| `HF_TOKEN` | Grants access to **gated** Hugging Face assets—protect like any other long-lived credential. | + +Checkpoint formats and operational practices are described in [Checkpoint management](../04-technical-guides/checkpoint-management.md). + +--- + +## 7. What is not verified + +The following items reflect **typical gaps** in public-facing evidence for many research and infrastructure codebases; confirm against your organization’s audits and CI for your fork and deployment. + +| Topic | Status (evidence-based caveat) | +|-------|--------------------------------| +| Independent security audit | **No** comprehensive third-party audit of this codebase is asserted here. | +| CI secrets scanning | **No** guarantee of automated secret detection in CI unless your pipeline adds it. | +| Dependency vulnerability scanning | **No** guarantee of continuous SCA unless your pipeline adds it. | +| Container images | Images may contain **unpatched** OS or Python packages—scan and rebuild on a schedule. | +| RCCL / NCCL traffic | **Not** encrypted or mutually authenticated by default; rely on network trust boundaries. | + +Use this section as a checklist for **your** production controls: add scanning, signing, policy-as-code, and periodic reviews appropriate to your threat model. diff --git a/docs/05-operations/troubleshooting.md b/docs/05-operations/troubleshooting.md new file mode 100644 index 000000000..1def25e8b --- /dev/null +++ b/docs/05-operations/troubleshooting.md @@ -0,0 +1,193 @@ +# Troubleshooting guide + +This guide is the primary reference for diagnosing and resolving common failures when running Primus on AMD GPUs (ROCm, RCCL, Docker). It complements the [CLI reference](../02-user-guide/cli-reference.md), [Preflight](../02-user-guide/preflight.md), [Benchmarking](../02-user-guide/benchmarking.md), and [Configuration system](../02-user-guide/configuration-system.md) documentation. + +--- + +## 1. Diagnostic tools + +Use these tools before scaling a job or when a failure is hard to localize. + +| Tool | Purpose | +|------|---------| +| `primus-cli --debug` | Enables verbose logging in `primus-cli` for command construction, delegation, and runtime details. | +| `primus-cli --dry-run` | Prints the commands Primus would run without executing them—useful to verify wrappers, paths, and MPI/launcher wiring. | +| `--export_config ` | Parsed by the training config parser, but not written by the current default `PrimusRuntime` path. Treat it as legacy/future functionality unless your deployment has implemented it. | +| `NCCL_DEBUG=INFO` | Surfaces detailed RCCL/NCCL connection and collective logs (set in the job environment). | +| `primus-cli direct -- preflight --host --gpu --network` | Fast host, GPU, and network validation. See [Preflight](../02-user-guide/preflight.md). | +| `PRIMUS_PATCHES=none` | Disables Primus patches to the selected backend to isolate whether a failure is Primus-specific or upstream. | + +**Examples** + +```bash +# Verbose CLI + dry run (no training executed) +primus-cli --debug --dry-run direct -- train pretrain --config path/to/config.yaml + +# Preflight: environment validation only +primus-cli direct -- preflight --host --gpu --network + +# RCCL/NCCL verbose logs in the training process environment +export NCCL_DEBUG=INFO + +# Isolate backend vs. Primus patch layer +export PRIMUS_PATCHES=none +``` + +For end-to-end checks including optional performance probes, see [Preflight](../02-user-guide/preflight.md) (`preflight --perf-test`). + +--- + +## 2. Out of memory (OOM) errors + +**Symptoms:** HIP/CUDA OOM messages, worker processes killed by the OOM killer, or abrupt exits during forward/backward. + +**Primary levers and mitigations** + +| Approach | What to change | +|----------|------------------| +| Reduce per-device activation memory | Lower **`micro_batch_size`** (often the first knob). | +| Shard weights/activations across devices | Increase **tensor parallelism** (`tensor_model_parallel_size` or `parallelism.tensor_parallel_degree`, depending on backend). See [Parallelism configuration](../04-technical-guides/parallelism-configuration.md). | +| Trade compute for memory | **Activation recomputation:** `recompute_granularity: full`, `recompute_method: uniform`, `recompute_num_layers: `. | +| Shard optimizer state | **`use_distributed_optimizer: true`** (Megatron-style stacks). | +| FSDP / sharded data parallel | Megatron: **`use_torch_fsdp2: true`**. TorchTitan: increase **`data_parallel_shard_degree`**. | +| CPU offload | **`optimizer_cpu_offload: true`** where supported. | +| Sequence length | Reduce **maximum sequence length** if the workload allows. | +| MoE / scratch memory | Set **`HSA_NO_SCRATCH_RECLAIM=1`** to reduce scratch-memory conflicts on some MoE workloads. | +| Plan before you run | **`primus-cli ... -- projection memory --config `**—see [Projection](../02-user-guide/projection.md). | + +**Related references:** [Parallelism strategies](../04-technical-guides/parallelism-strategies.md), [Performance tuning](../04-technical-guides/performance-tuning.md), [Megatron parameters](../03-configuration-reference/megatron-parameters.md), [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md). + +--- + +## 3. Distributed communication failures + +**Symptoms:** Hangs at initialization, NCCL/RCCL timeouts, connection errors, or inconsistent ranks. + +| Cause | What to verify / fix | +|-------|----------------------| +| Wrong network interface | Set `NCCL_SOCKET_IFNAME` to the correct interface; exclude virtual interfaces, e.g. `^docker0,lo`. | +| `MASTER_ADDR` unreachable | From every node, resolve DNS/IP consistently; verify firewalls and routing. | +| Port already in use | Change **`MASTER_PORT`** to a free port on all nodes. | +| InfiniBand not used or missing | Confirm `/dev/infiniband` exists where expected; run `ibstat`; ensure IB kernel modules are loaded. See [Multi-node networking](../04-technical-guides/multi-node-networking.md). | +| Timeout too aggressive | Megatron: increase **`distributed_timeout_minutes`**. TorchTitan: increase **`comm.init_timeout_seconds`**. | +| Mismatched world size | Align **`NNODES`**, **`GPUS_PER_NODE`**, and launcher settings across **all** nodes. | + +**Debugging** + +```bash +export NCCL_DEBUG=INFO +``` + +**Validation** + +```bash +primus-cli direct -- preflight --network +primus-cli direct -- benchmark rccl --op all_reduce +``` + +Collective behavior and RCCL roles are summarized in [Collective operations](../04-technical-guides/collective-operations.md). + +--- + +## 4. Container issues + +**Symptoms:** Permission denied on devices, GPUs not visible inside the container, immediate exit, or RCCL failures only under Docker. + +| Symptom | Typical cause | Mitigation | +|---------|-----------------|------------| +| GPU not visible | Devices not passed through | Ensure **`--device /dev/kfd`** and **`--device /dev/dri`** (Primus container mode typically sets these). | +| Permission denied on GPU | Group membership / permissions | Add **`--group-add video`**; ensure the user has **video/render** access on the host. | +| InfiniBand missing in container | Device not mounted | Add **`--device /dev/infiniband`** (and related uverbs devices as required by your site). | +| Debugger/profiler failures | Missing capabilities | **`--cap-add SYS_PTRACE`** and **`--cap-add CAP_SYS_ADMIN`** are required for many ROCm tooling paths. | +| Driver/library mismatch | Image vs. host ROCm | Match **container image ROCm** to **host ROCm driver** version. | +| Data or code not found | Bind mounts | Use **`--volume /host/path:/container/path`** for datasets and workspace. | +| Env vars missing in container | Forwarding | Check **`runner/.primus.yaml`** `container.options.env` for auto-forwarded variables; add extras with **`--env KEY=VALUE`**. | + +Installation and container-oriented setup are covered in [Installation](../01-getting-started/installation.md). + +--- + +## 5. Configuration errors + +**Symptoms:** YAML parse failures, "unknown key" or type errors, silent wrong behavior after edits. + +| Issue | Resolution | +|-------|------------| +| Unset `${VAR}` | `${VAR}` with no default fails if `VAR` is unset. Use **`${VAR:default}`** or **export** the variable before launch. | +| Broken `extends:` chain | Confirm every referenced file exists and paths resolve relative to the expected directory. | +| Wrong parameter name | Cross-check backend docs: [Megatron](../03-configuration-reference/megatron-parameters.md), [TorchTitan](../03-configuration-reference/torchtitan-parameters.md), [MaxText](../03-configuration-reference/maxtext-parameters.md), [Megatron Bridge](../03-configuration-reference/megatron-bridge-parameters.md). | +| Override not applied | Review merge order: **CLI > experiment overrides > module preset with model preset additions**; duplicate top-level module keys win over model preset keys. See [Configuration system](../02-user-guide/configuration-system.md). | +| Effective config unknown | Use **`--dry-run`** to inspect the launch command and manually trace the experiment YAML, module preset, model preset, and overrides. Resolved-config export is not currently written by the default core runtime. | + +--- + +## 6. Backend-specific issues + +### Megatron + +| Issue | Mitigation | +|-------|------------| +| Suspected Primus patch interaction | `export PRIMUS_PATCHES=none` and retry with vanilla Megatron behavior. | +| Custom kernel compile failures | `disable_compile_dependencies: true` skips custom kernel compilation where applicable. | +| Wrong third-party path | Set **`BACKEND_PATH`** to override third-party resolution. | + +### TorchTitan + +| Issue | Mitigation | +|-------|------------| +| Submodule drift | `git submodule update --recursive` so `third_party/torchtitan` matches the Primus revision you run. | +| `torch.compile` instability | `export TORCH_COMPILE_DISABLE=1` or set **`compile.enable: false`** in config. | + +### MaxText (JAX) + +| Issue | Mitigation | +|-------|------------| +| JAX / jaxlib vs ROCm | Verify JAX and jaxlib builds match your ROCm stack. | +| XLA memory pressure | Tune **`XLA_PYTHON_CLIENT_MEM_FRACTION`** to cap client-side allocator use. | + +--- + +## 7. Performance issues + +**Symptoms:** Low tokens/sec, long iteration time, or poor scaling versus expectations. + +**Diagnosis** + +| Step | Command / action | +|------|-------------------| +| Compute sanity | `primus-cli direct -- benchmark gemm` | +| Interconnect | `primus-cli direct -- benchmark rccl --op all_reduce` | +| Broader probe | `primus-cli direct -- preflight --perf-test` (see [Preflight](../02-user-guide/preflight.md)) | + +**Common fixes** + +| Area | Action | +|------|--------| +| GEMM / kernels | Enable **HipBLASLt tuning** (multi-stage workflow—see [Performance tuning](../04-technical-guides/performance-tuning.md)). | +| Primus stack | Enable **`enable_primus_turbo: true`** where supported. | +| Communication overlap | **`overlap_grad_reduce: true`**, **`overlap_param_gather: true`** (when applicable to your backend). | +| MoE / scratch | Confirm **`HSA_NO_SCRATCH_RECLAIM=1`** when recommended for your model class. | +| FP8 | Enable FP8 when hardware and backend support it. | + +--- + +## 8. Data issues + +| Symptom | Checks | +|---------|--------| +| Mock data works; real data fails | **`data_path`** format: Megatron typically expects **`.bin` / `.idx`** pairs. See [Data preparation](../04-technical-guides/data-preparation.md). | +| Tokenizer errors | **`tokenizer_type`** and **`tokenizer_model`** must match (e.g., Hugging Face tokenizer ID for `HuggingFaceTokenizer`). | +| Hugging Face download failures | Set **`HF_TOKEN`** for gated models; verify outbound network and cache directories. | + +--- + +## 9. Known limitations + +| Area | Note | +|------|------| +| MaxText | Parameter completeness depends on upstream MaxText **`base.yml`**; some keys are inherited from upstream defaults. | +| Megatron Bridge | Recipe parameters may be loaded dynamically; not every key appears in static reference tables. | +| HummingbirdXT | Less mature than other backends; expect sharper edges in configs and tooling. | +| Primus-Turbo | Requires a **separate installation** step; not always present by default. | + +For terminology, see the [Glossary](../01-getting-started/glossary.md). For checkpoint-related failures, see [Checkpoint management](../04-technical-guides/checkpoint-management.md). diff --git a/docs/06-developer-guide/README.md b/docs/06-developer-guide/README.md new file mode 100644 index 000000000..a8c21c5a6 --- /dev/null +++ b/docs/06-developer-guide/README.md @@ -0,0 +1,17 @@ +# Developer guide + +For contributors and maintainers. + +- [Architecture](architecture.md): system design, runtime, backends, patch system +- [Contributing](contributing.md): development setup, code style, PR process +- [Testing](testing.md): test types, running tests, CI pipeline +- [Extending backends](extending-backends.md): adding new training backends +- [Adding models](adding-models.md): adding model configurations per backend +- [Model support matrix](model-support-matrix.md): supported models per backend and GPU +- [CLI architecture](cli-architecture.md): CLI internals: subcommand discovery, dispatch, and launch wrappers +- [Backend patch notes](backend-patch-notes.md): Primus-specific backend arguments and the files they patch +- [Tooling](tooling.md): auxiliary analysis, benchmarking, visualization, and diagnostics tools + +--- + +[← Documentation home](../README.md) diff --git a/docs/06-developer-guide/adding-models.md b/docs/06-developer-guide/adding-models.md new file mode 100644 index 000000000..c7f3f09da --- /dev/null +++ b/docs/06-developer-guide/adding-models.md @@ -0,0 +1,379 @@ +# Adding model configurations + +This guide explains how to add **model configuration YAML** for each Primus training backend. Model presets live under `primus/configs/models//` and are referenced from **experiment** YAML under `examples//configs/...`. Backend-specific parameter references: + +- [Megatron parameters](../03-configuration-reference/megatron-parameters.md) +- [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md) +- [MaxText parameters](../03-configuration-reference/maxtext-parameters.md) +- [Megatron Bridge parameters](../03-configuration-reference/megatron-bridge-parameters.md) + +--- + +## Overview: Three-layer configuration + +For each backend, Primus composes configuration in three layers: + +1. **Experiment config** (entry point): `examples//configs//.yaml`—selects `framework`, `config` (module preset), `model` (model preset), and `overrides`. +2. **Module config** (trainer defaults): `primus/configs/modules//.yaml`—training loop defaults, logging, optimizer blocks, and backend-specific knobs. +3. **Model config** (architecture and assets): `primus/configs/models//.yaml`—architecture fields and tokenizer or Hugging Face paths, shaped differently per backend (see sections below). + +At runtime, `modules..model: .yaml` resolves to `primus/configs/models//.yaml` and is merged into module parameters before the backend adapter converts them. + +--- + +## Adding a Megatron model + +### How Megatron configs are wired + +1. **Experiment config** (entry point): + + ```yaml + # examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml + modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml # module-level trainer config + + # model to run + model: llama3.1_8B.yaml # model config name + ``` + +2. **Module config** (trainer-level defaults): `primus/configs/modules/megatron/pre_trainer.yaml`—extends shared bases and sets Megatron training defaults. + +3. **Model config** (architecture + tokenizer): + + ```yaml + # primus/configs/models/megatron/llama3.1_8B.yaml + extends: + - llama3_8B.yaml + + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: meta-llama/Llama-3.1-8B + + max_position_embeddings: 131072 + ``` + +At runtime, `modules.pre_trainer.model: llama3.1_8B.yaml` resolves to `primus/configs/models/megatron/llama3.1_8B.yaml`. The `extends` chain pulls in parent files (for example `llama3_8B.yaml` → `llama3_base.yaml` → `llama_base.yaml`). + +### Files you typically add + +| Artifact | Purpose | +| -------- | ------- | +| **Model preset** (required) | New YAML under `primus/configs/models/megatron/`—architecture, tokenizer, and optional `extends`. | +| **Experiment config** (required) | New or copied YAML under `examples/megatron/configs/MI300X/` or `MI355X/`—points `model:` at your preset and sets `overrides` (batch size, precision, parallelism, `mock_data`, and so on). | +| **Module preset** (optional) | Only if you need trainer defaults that differ from `pre_trainer.yaml`—new file under `primus/configs/modules/megatron/` and reference it as `config:` in the experiment. | + +### Example: TinyLlama 1.1B from Hugging Face + +Assume Hugging Face repo `TinyLlama/TinyLlama-1.1B-Chat-v1.0` is not yet represented in Primus. You can add a local model preset (not necessarily committed upstream) as follows. + +**1. Decide the architecture** + +Because TinyLlama is not shipped as a Megatron preset in Primus, you can: + +- **Option A (recommended):** extend `language_model.yaml` and set all architecture fields explicitly. +- **Option B:** extend the closest existing model (for example a LLaMA-style preset) and override differing fields. + +**2. Map Hugging Face `config.json` to Megatron keys** + +Read from the Hugging Face model (typically `config.json` or the model card): + +| Hugging Face / concept | Megatron YAML (typical keys) | +| ---------------------- | ---------------------------- | +| `hidden_size` | `hidden_size` | +| `intermediate_size` | `ffn_hidden_size` | +| `num_attention_heads` | `num_attention_heads` | +| `num_hidden_layers` | `num_layers` | +| `num_key_value_heads` | Use with `num_attention_heads` to set `num_query_groups` (often `num_attention_heads / num_key_value_heads`) | +| `max_position_embeddings` | `max_position_embeddings` | + +**3. Create `tinyllama_1.1B.yaml`** + +Path: `primus/configs/models/megatron/tinyllama_1.1B.yaml` + +```yaml +extends: + - language_model.yaml # generic Megatron language model base + +tokenizer_type: HuggingFaceTokenizer +tokenizer_model: TinyLlama/TinyLlama-1.1B-Chat-v1.0 + +hidden_size: 2048 +ffn_hidden_size: 5632 # intermediate_size in HF config.json +num_attention_heads: 32 +num_layers: 22 # num_hidden_layers in HF config.json +num_query_groups: 8 # e.g. 32 / 4 if HF has 4 KV heads + +max_position_embeddings: 2048 +position_embedding_type: rope +``` + +**4. Point an experiment at the new model** + +Copy an existing experiment (for example `examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml`) and set `model:` to your preset. Use **mock data** first for a quick sanity check: + +```yaml +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:tinyllama_1.1B-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + model: tinyllama_1.1B.yaml + overrides: + save: null + disable_last_saving: true + stderr_sink_level: DEBUG + + mock_data: true + train_iters: 50 + micro_batch_size: 2 + global_batch_size: 128 + + seq_length: 2048 +``` + +**5. Run verification** + +```bash +./primus-cli direct -- \ + train pretrain \ + --config examples/megatron/configs/MI300X/tinyllama_1.1B-pretrain.yaml +``` + +Confirm in logs that `framework` is `megatron`, the resolved model file is `tinyllama_1.1B.yaml`, and the tokenizer matches your preset. + +### Megatron checklist + +- [ ] Choose an appropriate base under `primus/configs/models/megatron/` (`language_model.yaml` or a close LLaMA-style model). +- [ ] Set `tokenizer_type`, `tokenizer_model`, and architecture fields aligned with Hugging Face. +- [ ] Add or update an experiment YAML under `examples/megatron/configs/...` with `model: .yaml`. +- [ ] Run `./primus-cli direct -- train pretrain --config ...` to validate resolution and a short run. + +--- + +## Adding a TorchTitan model + +### How TorchTitan configs are wired in Primus + +1. **Experiment config**: + + ```yaml + # examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml + modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: llama3.1_8B.yaml + overrides: + training: + local_batch_size: 4 + seq_len: 8192 + mock_data: false + steps: 50 + ``` + +2. **Module config**: `primus/configs/modules/torchtitan/pre_trainer.yaml`—training defaults, quantization fragments, and TorchTitan-oriented structure. + +3. **Model config**—`job` and `model` sections consumed by the TorchTitan launcher: + + ```yaml + # primus/configs/models/torchtitan/llama3.1_8B.yaml + job: + dump_folder: "./outputs" + description: "Llama 3.1 8B training" + + model: + name: "llama3" + flavor: "8B" + hf_assets_path: "meta-llama/Llama-3.1-8B" + converters: + - primus_turbo + ``` + +At runtime, `modules.pre_trainer.model: llama3.1_8B.yaml` resolves to `primus/configs/models/torchtitan/llama3.1_8B.yaml`. The launcher uses `job` and `model` to wire the PyTorch model and training loop. + +### Mapping from Hugging Face to TorchTitan + +You need: + +- **Model family** (`model.name`): must match a family implemented in TorchTitan (for example `llama3`, `qwen3`, `deepseek_v3`). +- **Flavor** (`model.flavor`): a size key defined in TorchTitan code (for example `8B`, `70B`, `1.7b`)—see `third_party/torchtitan/torchtitan/models//`. +- **Hugging Face assets** (`model.hf_assets_path`): repository used to load weights and tokenizer. + +**Important limitations** + +- TorchTitan can only train models that are **implemented in the TorchTitan codebase**. The YAML under `primus/configs/models/torchtitan/` does **not** define new architectures; it selects and configures existing `*ModelArgs` entries. +- If a family or flavor is missing in TorchTitan, you cannot enable it with YAML alone—extend TorchTitan first, then add a Primus preset. + +### Example pattern: Qwen3 8B preset + +Qwen3 8B already exists in this repository as a TorchTitan preset and example. Use it as a pattern when adding a different TorchTitan model or flavor that is implemented upstream but not yet represented in Primus. + +**Existing file:** `primus/configs/models/torchtitan/qwen3_8b.yaml` + +```yaml +job: + dump_folder: "./outputs" + description: "Qwen 3 8B training" + +model: + name: "qwen3" + flavor: "8B" + hf_assets_path: "Qwen/Qwen3-8B" + converters: + - primus_turbo +``` + +**Field meanings:** + +- **`job.dump_folder`**—where TorchTitan writes logs and checkpoints for the job. +- **`job.description`**—free-form description shown in logs and metadata. +- **`model.name` / `model.flavor`**—the TorchTitan family and size key; both must exist in the TorchTitan code. +- **`model.hf_assets_path`**—Hugging Face repository used to load weights and tokenizer. +- **`model.converters`**—extra TorchTitan converters; `primus_turbo` is the default used in the Primus examples. + +The architecture itself lives in TorchTitan code, not in this YAML. For example, Qwen3 8B is declared in `third_party/torchtitan/torchtitan/models/qwen3/__init__.py`: + +```python +"8B": Qwen3ModelArgs( + vocab_size=151936, + max_seq_len=4096, + head_dim=128, + dim=4096, + n_layers=36, + n_heads=32, + n_kv_heads=8, + qk_norm=True, + hidden_dim=12288, + rope_theta=1000000, +), +``` + +The Primus preset only selects and configures such a definition. + +**Experiment snippet** (copy from an existing TorchTitan example and change `model:`): + +```yaml +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + model: qwen3_8b.yaml + overrides: + training: + local_batch_size: 4 + seq_len: 4096 + mock_data: true + steps: 50 +``` + +**Run:** + +```bash +./primus-cli direct -- \ + train pretrain \ + --config examples/torchtitan/configs/MI300X/qwen3_8B-pretrain.yaml +``` + +For a new model, create a new preset and example path that matches the upstream TorchTitan family/flavor you are adding. + +### TorchTitan checklist + +- [ ] Define `job` (for example `dump_folder`, `description`) and `model` (`name`, `flavor`, `hf_assets_path`, `converters`). +- [ ] Add an experiment under `examples/torchtitan/configs/...` referencing `model: .yaml`. +- [ ] Run `./primus-cli direct -- train pretrain --config ...` for a short job. + +--- + +## Adding a MaxText model + +MaxText (JAX) model presets in Primus are intentionally thin: they set **`model_name`** and **`tokenizer_path`** (and extend `model_base.yaml`) so MaxText can load its own architecture tables when available. + +**Typical model preset** + +Path pattern: `primus/configs/models/maxtext/.yaml` + +```yaml +extends: + - model_base.yaml + +model_name: "llama3-8b" +tokenizer_path: "meta-llama/Meta-Llama-3-8B" +``` + +Comments in `primus/configs/models/maxtext/model_base.yaml` explain that architecture parameters are resolved from MaxText’s `configs/models/.yml` when present, or from Primus overrides as appropriate. + +**Experiment wiring** + +Experiments reference the preset the same way as other backends, for example: + +```yaml +modules: + pre_trainer: + framework: maxtext + config: pre_trainer.yaml + model: llama3_8B.yaml +``` + +**Supported architectures** + +For the authoritative list of model names and architectures MaxText supports, see the [MaxText](https://github.com/AI-Hypercomputer/maxtext) repository and upstream documentation. Primus examples under `examples/maxtext/configs/MI300X/` and `MI355X/` illustrate which presets are exercised in this tree. + +--- + +## Adding a Megatron Bridge model (post-training) + +Megatron Bridge model presets are small YAML files that select a **recipe**, **flavor**, and **Hugging Face path**, plus optional dataset blocks. + +**Example preset** (`primus/configs/models/megatron_bridge/qwen3_8b.yaml`): + +```yaml +recipe: qwen.qwen3 +flavor: qwen3_8b_finetune_config +hf_path: Qwen/Qwen3-8B + +dataset: + dataset_name: "rajpurkar/squad" +``` + +| Field | Role | +| ----- | ---- | +| `recipe` | Logical recipe module (for example `qwen.qwen3`, `llama.llama3`). | +| `flavor` | Named configuration within the recipe (for example `qwen3_8b_finetune_config`). | +| `hf_path` | Hugging Face model id for weights and tokenizer. | + +**Experiment** (pattern from `examples/megatron_bridge/configs/`): + +```yaml +modules: + post_trainer: + framework: megatron_bridge + config: sft_trainer.yaml + model: qwen3_8b.yaml + overrides: + precision_config: bf16_mixed + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 +``` + +See [Megatron Bridge parameters](../03-configuration-reference/megatron-bridge-parameters.md) for the full override surface. + +--- + +## Testing new models + +Use a staged approach so failures are easy to localize. + +| Stage | Goal | Typical settings | +| ----- | ---- | ---------------- | +| **Mock or synthetic data** | Validate config resolution, tokenizer, and a few steps without real datasets. | Megatron: `mock_data: true`. TorchTitan: `training.mock_data: true`. MaxText: `dataset_type: "synthetic"` in overrides where applicable. Keep `train_iters` / `steps` small. | +| **Single-GPU** | Confirm numerics and memory before scaling. | Set tensor and pipeline parallelism to 1 in overrides; use one process / one device per your launcher docs. | +| **Multi-GPU** | Match production parallelism. | Set `tensor_model_parallel_size`, `pipeline_model_parallel_size`, expert / context parallel sizes, or MaxText `ici_*` / `dcn_*` fields as required by the model size and hardware. | + +Cross-check [Parallelism configuration](../04-technical-guides/parallelism-configuration.md) and [Model support matrix](./model-support-matrix.md) when moving from single-device to multi-device runs. diff --git a/docs/06-developer-guide/architecture.md b/docs/06-developer-guide/architecture.md new file mode 100644 index 000000000..a82520b5e --- /dev/null +++ b/docs/06-developer-guide/architecture.md @@ -0,0 +1,119 @@ +# Architecture overview + +This document describes how the Primus training framework is structured: CLI and configuration, the core runtime orchestrator, backend adapters, trainer lifecycle, and the patch system. + +## 1. System overview + +Primus is organized into three conceptual layers: runtime launch (how processes and GPUs are started), hooks and patches (environment and in-process adjustments), and task execution (CLI subcommands that drive training and utilities). + +``` +┌─────────────────────────────────────────────────────┐ +│ Runtime Layer (runner/) │ +│ direct | container | slurm │ +│ GPU detection, env setup, distributed launch │ +├─────────────────────────────────────────────────────┤ +│ Hook / Patch System │ +│ runner/helpers/hooks/ | primus/core/patches/ │ +│ Pre/post processing, runtime monkey-patches │ +├─────────────────────────────────────────────────────┤ +│ Task Execution Layer │ +│ primus/cli/subcommands/ │ +│ train | benchmark | preflight | projection │ +└─────────────────────────────────────────────────────┘ +``` + +The repository also provides shell entrypoints under `runner/` (for example `primus-cli-direct.sh`, `primus-cli-container.sh`, `primus-cli-slurm.sh`) that prepare the environment and invoke the Python CLI. + +## 2. CLI and plugin system + +- **Entry point:** `primus/cli/main.py` is the unified CLI entry. It discovers subcommand modules under `primus/cli/subcommands/` with `pkgutil.walk_packages`, skipping modules whose leaf name starts with `_`. +- **Registration contract:** Each subcommand module exposes `register_subcommand(subparsers)` and must return the configured parser. The parser must call `set_defaults(func=run)` so `main()` can dispatch to the handler. +- **Parsing:** The CLI uses the standard library `argparse` only (no Click or Typer). +- **Unknown arguments:** `main()` calls `parse_known_args()`. For selected subcommands (`train`, `projection`, `preflight`), trailing tokens are passed through to the handler as overrides; for other commands, unknown arguments are rejected. + +## 3. Configuration pipeline + +Configuration flows from experiment YAML to a resolved structure consumed by the runtime. + +1. **CLI** parses `--config` / `--exp` (required for train flows) pointing at an experiment YAML file. +2. **`load_primus_config()`** (used by `PrimusRuntime`) delegates to **`PrimusParser.parse()`** in `primus/core/launcher/parser.py`. The parser loads the experiment file via **`yaml_utils.parse_yaml_to_namespace()`**, which uses **`primus/core/config/yaml_loader.parse_yaml()`** for `${VAR}` / `${VAR:default}` substitution and `extends:` inheritance with deep merge. +3. **Per trainer module** (names containing `trainer`, for example `pre_trainer`): + - **`PresetLoader.load()`** loads the module preset from `primus/configs/modules//.yaml`. + - **`PresetLoader.load()`** loads the model preset from `primus/configs/models//.yaml`. + - Each preset is loaded through the same YAML pipeline (env substitution and `extends:` chains). +4. **`parse_platform()`** merges platform settings from `primus/configs/platforms/` (defaulting to `platform_azure.yaml` when the experiment omits `platform`). +5. **CLI overrides:** For `primus train`, `main()` passes `unknown_args` into the train handler. `PrimusRuntime` parses them with `parse_cli_overrides()` and **deep-merges** them into `module_config.params`. +6. **Result:** A resolved configuration where each module exposes a **`params`** namespace (`SimpleNamespace`) for training parameters, produced by `_normalize_module_for_runtime()` in `primus/core/config/primus_config.py`. + +The object returned from `load_primus_config()` is a lightweight `SimpleNamespace` (not `PrimusConfig`), with `modules` as a **list** of module configs, each tagged with a `.name` field. + +## 4. Core runtime (PrimusRuntime) + +`primus/core/runtime/train_runtime.py` defines **`PrimusRuntime`**, the main orchestrator for the new core training path. Execution for a single module follows this flow: + +1. **`load_primus_config()`** loads and validates the experiment; **`get_module_config()`** selects the requested module (for example `pre_trainer` or `post_trainer`). +2. **`_apply_overrides()`** merges CLI overrides into `module_config.params`. +3. **`_initialize_environment()`** ensures the data directory exists and calls **`setup_training_env()`** (Hugging Face cache and related setup). +4. **`_initialize_distributed_context()`** reads torchrun-style rank and master information via **`get_torchrun_env()`**. +5. **`_initialize_logging()`** initializes worker logging. +6. **`BackendRegistry.get_adapter(framework)`** resolves the **`BackendAdapter`** (lazy-importing `primus.backends.` if needed). +7. **`adapter.setup_backend_path()`** inserts the backend tree on `sys.path`. Resolution order: CLI `--backend_path`, then the `BACKEND_PATH` env var, then the default—`third_party/` under the repo root, followed by `$PRIMUS_THIRDPARTY_DIR` or `~/.cache/Primus/third_party` (the `primus-cli deps sync` location). +8. **`adapter.prepare_backend()`** runs backend setup hooks (via **`BackendRegistry.run_setup()`** by default). +9. **`adapter.convert_config(module_config.params)`** produces **`backend_args`** for the trainer. +10. **`run_patches(phase="build_args", ...)`** runs registered patches; backend version detection runs when patches first need it (**`adapter.detect_backend_version()`** via **`_get_backend_version()`**). +11. **`merge_namespace()`** merges `backend_args` into `module_config.params` (backend wins on conflicts); **`adapter.load_trainer_class(stage)`** resolves the trainer class (default stage `pretrain`). +12. **`TrainerClass(backend_args=backend_args)`** constructs the trainer. +13. **`run_patches(phase="setup")`** then **`trainer.setup()`**. +14. **`trainer.init()`**. +15. **`run_patches(phase="before_train")`** then **`trainer.train()`** then **`run_patches(phase="after_train")`** then **`trainer.cleanup()`**. +16. On failure, **`_safe_cleanup()`** calls **`trainer.cleanup(on_error=True)`** when possible. + +## 5. Backend system + +- **`BackendAdapter`** (`primus/core/backend/backend_adapter.py`) is the abstract integration surface. Subclasses implement **`convert_config()`**, **`load_trainer_class()`**, and **`detect_backend_version()`**. Shared behavior includes **`setup_backend_path()`** and a default **`prepare_backend()`** that runs registered setup hooks. +- **`BackendRegistry`** (`primus/core/backend/backend_registry.py`) maps backend names to adapter classes, supports **lazy import** of `primus.backends.`, and stores optional **setup hooks** per backend. +- **Registered adapters** (via each backend package’s `__init__.py` calling **`BackendRegistry.register_adapter()`**): **`megatron`**, **`torchtitan`**, **`maxtext`**, **`megatron_bridge`**, **`hummingbirdxt`**. +- Backend code lives under **`primus/backends//`**. Importing the package registers the adapter and any trainers or hooks that package defines. + +## 6. Trainer lifecycle + +- **`BaseTrainer`** (`primus/core/trainer/base_trainer.py`) defines the lifecycle. **`setup()`**, **`init()`**, and **`train()`** are **abstract** (subclasses must implement them); **`cleanup(on_error=False)`** is **optional**—it ships a default (no-op) implementation that subclasses may override. The constructor stores **`backend_args`** and reads distributed settings from **`get_torchrun_env()`**. +- Concrete trainers (for example Megatron or TorchTitan pretrain classes) subclass **`BaseTrainer`** and implement the abstract methods. +- **`PrimusRuntime`** drives **`setup` → `init` → `train` → `cleanup`**, with patch phases **`build_args`** (before the trainer is created), **`setup`**, and **`before_train`/`after_train`** (around `train`). No patch phase runs around `cleanup`—`after_train` fires before `cleanup` (see §4). + +## 7. Patch system + +- **`PatchRegistry`** (`primus/core/patches/patch_registry.py`) stores **`FunctionPatch`** objects keyed by backend and phase, with wildcard buckets (`None`) for patches that apply broadly. +- The **`@register_patch`** decorator registers a patch with metadata (priority, optional version patterns, tags). +- **`run_patches()`** (`primus/core/patches/patch_runner.py`) collects applicable patches, filters by **`PatchContext`**, sorts by **priority**, and runs handlers. It accepts an optional **`enabled_ids`** list; if omitted, behavior is controlled by **`PRIMUS_PATCHES`**: + - unset or **`all`**: all patches + - **`none`**: disable all + - comma-separated IDs: only those patches +- **Phases** used by the core runtime include **`build_args`**, **`setup`**, **`before_train`**, and **`after_train`**. +- Patch implementations are typically colocated with backends under **`primus/backends//patches/`**. + +## 8. Legacy runtime + +The legacy pretrain path—previously selected with **`PRIMUS_TRAIN_RUNTIME=legacy`** and backed by the **`primus/modules/`** stack (**`BaseModule`**-style composition)—has been **removed**. `primus/modules/` no longer contains any source code, and **`primus/cli/subcommands/train.py`** no longer reads `PRIMUS_TRAIN_RUNTIME` or resolves a legacy-vs-core runtime. + +All training now runs exclusively through the **core runtime**: both `primus train pretrain` and `primus train posttrain` construct a **`PrimusRuntime`** (**`primus/core/runtime/train_runtime.py`**). **`primus/pretrain.py`** now only provides shared backend-path / environment helpers (for example **`setup_backend_path()`**) used by the training and projection entry points; it no longer defines a `launch_pretrain_from_cli()` legacy launcher. + +## 9. Key source files + +| Path | Role | +|------|------| +| `primus/cli/main.py` | CLI entry, subcommand discovery, dispatch | +| `primus/cli/subcommands/train.py` | `train` subcommand; chooses core vs legacy pretrain; `posttrain` via `PrimusRuntime` | +| `primus/core/launcher/parser.py` | **`PrimusParser`**: experiment, platform, and module preset loading | +| `primus/core/config/preset_loader.py` | **`PresetLoader`**: load framework presets from `primus/configs/` | +| `primus/core/config/yaml_loader.py` | YAML load with env substitution and `extends` | +| `primus/core/config/primus_config.py` | **`load_primus_config()`**, **`get_module_config()`**, module normalization | +| `primus/core/runtime/train_runtime.py` | **`PrimusRuntime`**, **`TrainContext`** | +| `primus/core/backend/backend_adapter.py` | **`BackendAdapter`** ABC | +| `primus/core/backend/backend_registry.py` | **`BackendRegistry`** | +| `primus/core/trainer/base_trainer.py` | **`BaseTrainer`** ABC | +| `primus/core/patches/patch_registry.py` | **`PatchRegistry`**, **`@register_patch`** | +| `primus/core/patches/patch_runner.py` | **`run_patches()`**, **`PRIMUS_PATCHES`** parsing | +| `runner/primus-cli-*.sh` | Shell wrappers for direct, container, and Slurm launch | + +For a deep dive on the CLI internals (subcommand discovery, dispatch, and the launch wrappers), see [CLI Architecture](cli-architecture.md). For day-to-day contribution workflows (style, tests, CI), see [Contributing Guide](contributing.md) and [Testing Guide](testing.md). diff --git a/docs/06-developer-guide/backend-patch-notes.md b/docs/06-developer-guide/backend-patch-notes.md new file mode 100644 index 000000000..8d2d9e4a3 --- /dev/null +++ b/docs/06-developer-guide/backend-patch-notes.md @@ -0,0 +1,142 @@ +# Backend patch notes + +Primus integrates several large-model backends (Megatron-LM, TorchTitan, JAX MaxText, …) and applies a lightweight patch layer to keep configuration flags consistent with the Primus CLI. This page captures those backend-specific switches so they live alongside the rest of the documentation (instead of the historical `primus/README_patch.md` file). + +## How to read these notes + +- Start with the **Base Module Parameters** table below—every backend module inherits these knobs. +- Jump to the backend-specific section for details on extra CLI/config options and links to the patched source files. +- When editing configs or CLI presets, cross-reference the [Primus CLI Reference](../02-user-guide/cli-reference.md) so command examples and backend parameters stay in sync. + +## Supported models + +This section lists, at a high level, the model families Primus currently targets on each backend. For more details and configuration examples, refer to the backend-specific patch notes below. + +### Megatron-LM + +- **LLaMA family**: LLaMA2, LLaMA3, LLaMA3.1, LLaMA3.3, LLaMA4 (various sizes from 7B up to 405B+) +- **DeepSeek family**: DeepSeek-V2 (lite/base/full) and DeepSeek-V3 +- **MoE / Mixtral**: Mixtral-8x7B / 8x22B, large MoE configs (515B, 1T, 2T, 4T) and DeepSeek-style MoE +- **Qwen family**: Qwen2.5 (7B/72B) and Qwen3 (8B/30B/235B variants) +- **Other GPT-style models**: Grok1/2, GPT-OSS 20B and generic `language_model.yaml` + +### TorchTitan + +- **LLaMA family**: LLaMA3, LLaMA3.1, LLaMA3.3 (various sizes, including FP8 variants) +- **DeepSeek family**: DeepSeek-V3 (16B and 671B, FP8 and BF16 configs) +- **Qwen family**: Qwen3 small/medium models (0.6B, 1.7B, 32B) + +### JAX MaxText + +- **LLaMA family**: LLaMA2 (7B/70B), LLaMA3 (8B/70B), LLaMA3.3 (70B) +- **DeepSeek family**: DeepSeek-V2 16B +- **MoE / Mixtral**: Mixtral-8x7B +- **Other models**: Grok1 and additional MaxText-supported transformers (see MaxText docs for the full list) + +## Base module parameters + +All modules inherit the options defined in [`primus/configs/modules/module_base.yaml`](https://github.com/AMD-AGI/Primus/blob/main/primus/configs/modules/module_base.yaml): + +| Argument Name | Default Value | Description | +| ------------------- | ------------- | ------------------------------------------------------------------------------------------ | +| `trainable` | `false` | Whether the module participates in training. | +| `sink_level` | `null` | Global sink level for logging. Overrides `file_sink_level` and `stderr_sink_level` if set. | +| `file_sink_level` | `DEBUG` | Logging level for file sink (log files). | +| `stderr_sink_level` | `INFO` | Logging level for stderr/console output. | + +### Backend index + +- [Megatron-LM patch notes](#megatron-lm-patch-notes) +- [TorchTitan patch notes](#torchtitan-patch-notes) +- [JAX MaxText patch notes](#jax-maxtext-patch-notes) + +--- + +## Megatron-LM patch notes + +Primus keeps a curated patch layer on top of upstream Megatron-LM so CLI presets and configs can expose additional controls. Use this section with the [Base Module Parameters](#base-module-parameters) above for shared module parameters, and the [Primus CLI Reference](../02-user-guide/cli-reference.md) for CLI/config usage patterns. + +> ℹ️ The **Version** column maps to Primus internal patch tags (v0.x.y) so you know when a flag landed. + +### 1. Module-level parameters + +These arguments are introduced in the Megatron module logic (e.g., training loop, logging, resume logic). They are defined via patching and can be configured to control training behavior and logging utilities. + +| New Argument | Default Value | Version | Description | Patched Files | Notes | +| ------------------------------------ | ------------- | ------- | ---------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------ | +| `disable_tensorboard` | `true` | v0.1.0 | Whether to disable TensorBoard. Set to `false` if you want to enable profiling or torch trace. | NA | Required for timeline and performance debugging. | +| `disable_wandb` | `true` | v0.1.0 | Whether to disable Weights & Biases logging. | NA | Useful for internal benchmarking. | +| `disable_compile_dependencies` | `true` | v0.1.0 | Disables Megatron’s custom kernel compilation. Most ops are already covered by TE. | NA | Avoids redundant compilation steps. | +| `auto_continue_train` | `false` | v0.1.0 | Automatically resume training from the latest checkpoint if found in the `--save` path. | NA | Simplifies job restarts. | +| `disable_last_saving` | `false` | v0.1.0 | Skip saving the final checkpoint at the last iteration. | NA | Useful for profiling or benchmarking runs. | +| `no_fp8_weight_transpose_cache` | `false` | v0.2.0 | Disable the FP8 weight transpose cache to save memory. | `megatron.core.extensions.transformer_engine.TELinear`, `megatron.core.extensions.transformer_engine.TELayerNormColumnParallelLinear`, `megatron.core.extensions.transformer_engine.TEDelayedScaling` | May affect performance but reduce memory use. | +| `decoder_pipeline_manual_split_list` | `null` | v0.2.0 | Enable manual pipeline split in (interleaved) 1F1B pipeline parallelism. | `megatron.core.transformer.transformer_block.get_num_layers_to_build`, `megatron.core.transformer.transformer_layer.get_transformer_layer_offset` | Deprecated. Use `pipeline_model_parallel_layout` instead. | +| `pp_warmup` | `false` | v0.2.0 | Add fwd/bwd warmup to save iter1's time when pp degree is large. | NA | Can save much time for pipeline debug. | +| `dump_pp_data` | `false` | v0.2.0 | Enable dumping pp schedule data for visualization. | `megatron.core.pipeline_parallel.schedules.forward_step`, `megatron.core.pipeline_parallel.schedules.backward_step`, `megatron.core.pipeline_parallel.schedules.forward_backward_pipelining_with_interleaving`, `megatron.core.pipeline_parallel.schedules.forward_backward_pipelining_without_interleaving` | Useful for pipeline schedule visualization. | +| `disable_profiler_activity_cpu` | `false` | v0.2.0 | Disable CPU activity in torch profiling. | NA | If you only want to trace CUDA kernels and get a smaller trace JSON file, you can enable this option. However, if you plan to run with TraceLen, please do not enable it. more torch profiler args:
`torch_profiler_record_shapes: true`,
`torch_profiler_with_stack: true`,
`torch_profiler_use_gzip: true` | +| `use_rocm_mem_info` | `false` | v0.2.0 | Logging ROCm memory information in Megatron-LM Trainer | NA | If `use_rocm_mem_info = True`, ROCm memory information will be collected with `rocm-smi` at every iteration. | +| `use_rocm_mem_info_iters` | `[1,2]` | v0.2.0 | Logging ROCm memory information in Megatron-LM Trainer for some iterations | NA | If `use_rocm_mem_info = False`, ROCm memory information will be collected at the iterations specified in `use_rocm_mem_info_iters`. | +| `patch_zero_bubble` | `false` | v0.2.0 | Using Zero-Bubble pipeline parallism | `megatron.core.optimizer.ChainedOptimizer`, `megatron.core.pipeline_parallel.get_forward_backward_func`, `megatron.core.tensor_parallel.layers.LinearWithGradAccumulationAndAsyncCommunication`, `megatron.core.parallel_stat.default_embedding_ranks`, `megatron.core.parallel_stat.is_pipeline_last_stage`, `megatron.core.parallel_stat.is_rank_in_embedding_group`, `megatron.core.distributed.finalize_model_grads`, `megatron.core.transformer.transformer_layer.get_transformer_layer_offset` | If `patch_zero_bubble = True`, Zero bubble pipeline parallism will be enable to use. See more detail at [ZeroBubble User Guide](../../primus/backends/megatron/core/pipeline_parallel/zerobubble/README.md) | +| `disable_mlflow` | `true` | v0.3.0 | Track model development using MLflow | NA | Envs:
`export DATABRICKS_TOKEN=your_token`
`export DATABRICKS_HOST=your_host`
`export MLFLOW_TRACKING_URI=databricks`
`export MLFLOW_REGISTRY_URI=databricks-uc`
Arguments:
`mlflow_run_name: null`,
`mlflow_experiment_name: null` | +| `recompute_layer_ids` | `null` | v0.4.0 | Specify the exact IDs of layers to recompute, enabling more flexible memory reduction | NA | Using `recompute_layer_ids=[layer_id_0, layer_id_1,...]` together with `recompute_granularity=full`, where layer_id ranges from 0 to num_layers - 1. | +| `dataloader_mp_context` | `null` | v0.5.0 | Set `DataLoader.multiprocessing_context` to avoid SIGSEGV caused by fork()-hostile native state (RDMA MRs, HIP runtime, IPC handles). | `torch.utils.data.DataLoader.__init__` | Accepted values: `"forkserver"`, `"spawn"`, `"fork"`, or `null` (keep PyTorch default). Only takes effect when `num_workers > 0` and no explicit `multiprocessing_context` is passed. | + +--- + +### 2. Model-definition parameters + +These arguments affect the internal architecture or layer implementations. They are patched into the model construction logic and used for tuning or debugging specific variants. + +| New Argument | Default Value | Version | Description | Patched Files | Notes | +| ----------------------------------- | ------------- | ------- | ------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------- | +| `disable_primus_topk_router` | `false` | v0.1.0 | Disable PrimusTopkRouter and use TopkRouter implemented by megatron. | `megatron.core.transformer.moe.router.TopKRouter` | Used to debug internal. | +| `moe_router_force_load_balancing` | `false` | v0.1.0 | Force token redistribution in MoE to achieve load balance across experts. | `megatron.core.transformer.moe.router.TopKRouter` | Use to debug MoE imbalance issues. | +| `use_deprecated_20241209_moe_layer` | `false` | v0.1.0 | Enable legacy MoE implementation for debugging/perf comparison. | `megatron.core.transformer.moe.moe_layer.MoELayer`, `megatron.core.transformer.moe.moe_layer.MoESubmodules`, `megatron.core.transformer.moe.experts.GroupedMLP`, `megatron.core.transformer.moe.experts.SequentialMLP`, `megatron.core.transformer.moe.experts.TEGroupedMLP`, `megatron.core.transformer.moe.router.TopKRouter` | Deprecated, used for internal testing only. | +| `moe_permute_fusion` | `false` | v0.1.0 | Permutation and unpermutation fusion. | `megatron.core.extensions.transformer_engine`, `megatron.core.transformer.moe.moe_utils` | Fuse permutation and unpermutation in moe layer. | +| `moe_use_fused_router_with_aux_score` | `false` | v0.2.0 | Fused router topk and calculation of moe aux loss score. Need Primus turbo backend | `megatron.core.transformer.moe.router.TopKRouter` | Used to reduce launch overhead of the small kernels in router. | + +### 3. Primus-Turbo related options + +| New Argument | Default Value | Version | Description | Patched Files | Notes | +| ----------------------------------- | ------------- | ------- | ------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------- | +| `use_turbo_gemm` | `false` | v0.8.0 | Use Primus-Turbo linear modules (`PrimusTurboLinear`, `PrimusTurboColumnParallelLinear`, `PrimusTurboRowParallelLinear`, `PrimusTurboLayerNormColumnParallelLinear`) in place of the TE linear modules. | `megatron.core.extensions.transformer_engine.TELinear`, `megatron.core.extensions.transformer_engine.TEColumnParallelLinear`, `megatron.core.extensions.transformer_engine.TERowParallelLinear`, `megatron.core.extensions.transformer_engine.TELayerNormColumnParallelLinear` | Accelerates dense GEMMs. Supports FP8 recipes (`tensorwise`, `blockwise`, `mxfp8`) and FP4 recipe (`mxfp4`). Replaces the deprecated `use_turbo_parallel_linear`. **Please set `enable_primus_turbo=True` first.** | +| `use_turbo_grouped_gemm` | `false` | v0.8.0 | Use Primus-Turbo grouped GEMM (`PrimusGroupedMLP` with `PrimusTurboColumnParallelGroupedLinear` / `PrimusTurboRowParallelGroupedLinear`) for MoE experts in place of `TEGroupedMLP`. | `megatron.core.transformer.moe.experts.TEGroupedMLP`, `megatron.core.extensions.transformer_engine.TEColumnParallelGroupedLinear`, `megatron.core.extensions.transformer_engine.TERowParallelGroupedLinear` | Accelerates MoE grouped GEMMs. Incompatible with `moe_use_legacy_grouped_gemm=True`. Required by Sync-Free MoE stage 2/3. Replaces the deprecated `use_turbo_grouped_mlp`. **Please set `enable_primus_turbo=True` first.** | +| `use_turbo_permute_padding` | `false` | v0.8.0 | Pad tokens of every experts to 16 or 32 multiple to reduce d2h and h2d. | `megatron.core.transformer.moe.token_dispatcher.MoEFlexTokenDispatcher`, `megatron.core.transformer.moe.experts.TEGroupedMLP` | Only effective under FP8/FP4 with `use_turbo_deepep=True`. Pad multiple is 16/32 (FP8, depending on recipe) or 32 (FP4). **Please set `enable_primus_turbo=True` first.** | +| `use_turbo_deepep` | `false` | v0.4.0 | Use Primus-turbo `DeepEPTokenDispatcher`. | `megatron.core.transformer.moe.token_dispatcher.MoEFlexTokenDispatcher` | Used Primus-Turbo DeepEP to accelerate MoE token dispatcher. **You must both set`enable_primus_turbo=True` and `use_turbo_deepep=True` to enable this function.** | +| `turbo_deepep_num_cu` | `32` | v0.4.0 | Set the number of CUs to use for Primus-Turbo DeepEP. | | 64 or 80 for ep8, 32 for ep16-64 is best practice. | +| `turbo_deepep_use_comm_stream` | `false` | v0.4.0 | Primus-Turbo DeepEP will use an internal stream to dispatch/combine when enabled, default used `current_stream` | | **Please both set`enable_primus_turbo=True` and `use_turbo_deepep=True` first** +| `turbo_sync_free_moe_stage` | `0` | v0.4.0 | Primus Sync-Free MoE has 4 stages. See [RFC: Primus-Megatron SyncFree MoE](https://github.com/AMD-AGI/Primus/issues/203) for more details. | | stage 2 is recommended for better performance. **Please set`enable_primus_turbo=True` first** | + +--- + +## TorchTitan patch notes + +TorchTitan integration uses the same Primus configuration surface (CLI flags + YAML) but exposes a few extra knobs via patches. Pair this with the [Base Module Parameters](#base-module-parameters) above for shared module parameters. + +| New Argument | Default Value | Version | Description | Patched Files | Notes | +| ------------ | ------------- | ------- | ----------- | ------------- | ----- | +| `primus_turbo.enable_embedding_autocast` | `true` | v0.4.0 | Automatically casts `nn.Embedding` outputs to the AMP dtype (e.g., bf16) during training so downstream layers stay in sync. | (Primus TorchTitan patch set) | Disable only if you manage casting manually. | + +--- + +## JAX MaxText patch notes + +Primus integrates JAX MaxText as a backend for running LLaMA and related transformer models on AMD GPUs. At the moment, Primus does not apply any additional patch-layer arguments on top of MaxText—the MaxText configuration surface (YAML + CLI) is used as-is. + +Use this section together with: + +- The [Base Module Parameters](#base-module-parameters) and [Supported Models](#supported-models) above for a high-level model overview +- The [Primus CLI Reference](../02-user-guide/cli-reference.md) for Primus CLI usage patterns +- The official MaxText documentation for the full set of MaxText-specific arguments + +### MaxText-specific notes + +- Primus currently wires MaxText via `primus/configs/models/maxtext` and `primus/configs/modules/maxtext`. +- Model families currently exercised in examples include: + - LLaMA2 7B/70B + - LLaMA3 8B/70B + - LLaMA3.3 70B + - DeepSeek-V2 16B + - Mixtral-8x7B + - Grok1 +- There are no extra Primus-only flags for MaxText yet; as we add MaxText-specific patches (e.g., ROCm optimizations, logging helpers), they will be documented in tables here in the same style as the Megatron-LM and TorchTitan patch notes. diff --git a/docs/cli/CLI-ARCHITECTURE.md b/docs/06-developer-guide/cli-architecture.md similarity index 93% rename from docs/cli/CLI-ARCHITECTURE.md rename to docs/06-developer-guide/cli-architecture.md index 944988aa3..8dd86e78f 100644 --- a/docs/cli/CLI-ARCHITECTURE.md +++ b/docs/06-developer-guide/cli-architecture.md @@ -1,4 +1,4 @@ -# 🚀 From Chaos to Order: Building a Unified Entry Point for AMD GPU LLM Training +# 🚀 From chaos to order: Building a unified entry point for AMD GPU LLM training > ⚠️ **NOTE**: This is a draft version and not the final release. > @@ -8,7 +8,7 @@ --- -## 📖 The Beginning: Pain Points in Training Workflows +## 📖 The beginning: Pain points in training workflows Imagine this scenario: @@ -41,7 +41,7 @@ The traditional approach is to use a large number of Bash scripts to handle thes --- -## 💡 Design Philosophy: One Command, Done +## 💡 Design philosophy: One command, done Our core philosophy is simple: **One command, from environment configuration to training launch, fully automated.** @@ -50,7 +50,7 @@ Our core philosophy is simple: **One command, from environment configuration to primus-cli direct -- train pretrain --config deepseek_v2.yaml ``` -### 🏗️ Three-Layer Architecture Design +### 🏗️ Three-layer architecture design Primus CLI adopts a clear **three-layer structure + plugin system**: @@ -74,7 +74,7 @@ Primus CLI adopts a clear **three-layer structure + plugin system**: └─────────────────────────────────────────────────────┘ ``` -### 🎯 Four Design Goals +### 🎯 Four design goals | Goal | Implementation | User Benefits | |------|---------------|---------------| @@ -85,9 +85,9 @@ Primus CLI adopts a clear **three-layer structure + plugin system**: --- -## 🔍 Deep Dive: Architecture Dissection +## 🔍 Deep dive: Architecture dissection -### ⚙️ Layer 1: Intelligent Runtime Abstraction +### ⚙️ Layer 1: Intelligent runtime abstraction Different scenarios require different runtime environments, but users shouldn't have to worry about these details. Primus CLI provides three seamlessly switchable runtime modes: @@ -114,7 +114,7 @@ primus-cli slurm srun -N 8 -- benchmark gemm -M 4096 --- -### 🔁 Layer 2: Hook & Patch System +### 🔁 Layer 2: Hook and patch system Training is more than just running a Python script. You might need to: - 🗂️ Preprocess datasets before training @@ -147,7 +147,7 @@ This is especially useful when you need to quickly apply temporary fixes or make --- -### 🧩 Layer 3: Task Execution Layer +### 🧩 Layer 3: Task execution layer This layer is responsible for executing specific business logic—training, testing, environment checks, and other actual tasks. Remember we said "zero-intrusion extension"? How is this achieved? @@ -185,15 +185,15 @@ This plugin-based design allows Primus CLI to quickly respond to new requirement --- -## 🌐 The Magic Behind: Intelligent Environment Detection +## 🌐 The magic behind: Intelligent environment detection This is probably the most "black tech" part of Primus CLI. -### Problem: Different GPUs Need Different Configurations +### Problem: Different GPUs need different configurations AMD's GPU family is rich: MI300X, MI250X, MI210... Each GPU has its optimal ROCm configuration and environment variable settings. The traditional approach is to let users manually select configurations, but this is both error-prone and insufficiently automated. -### Solution: Three-Step Auto-Configuration +### Solution: Three-step auto-configuration **Step 1: Load Common Environment** @@ -226,7 +226,7 @@ Now, `MI300X.sh` can contain all best practices for this GPU model: **Users don't need to worry about these details at all - everything is automatic.** -### Real-World Example +### Real-world example ```bash # On MI300X cluster @@ -242,7 +242,7 @@ primus-cli direct -- train pretrain --config config.yaml --- -## 🧪 Foundation of Scientific Experiments: Reproducibility +## 🧪 Foundation of scientific experiments: Reproducibility In machine learning research, reproducibility is crucial. But reality is harsh: @@ -250,7 +250,7 @@ In machine learning research, reproducibility is crucial. But reality is harsh: Does this sound familiar? Primus CLI completely solves this problem with an **automated snapshot mechanism**. -### Auto-Record Everything +### Auto-record everything Every time training starts, Primus CLI automatically saves the complete runtime context: @@ -270,7 +270,7 @@ output/exp_2025_11_10_134522/ └── metadata.json # Runtime metadata ``` -### One-Click Reproduction +### One-click reproduction Three months later, when you want to reproduce this experiment: @@ -285,7 +285,7 @@ Primus CLI will automatically: 3. Verify GPU and system environment 4. Start training (if environment is compatible) -### Real-World Value +### Real-world value | Scenario | Traditional Approach | Using Primus CLI | |----------|---------------------|------------------| @@ -297,11 +297,11 @@ Primus CLI will automatically: --- -## 📊 Real-World Case: From Development to Production +## 📊 Real-world case: From development to production Let's see how Primus CLI simplifies the entire workflow through a real scenario. -### Scenario: Training DeepSeek-V2 Model +### Scenario: Training DeepSeek-V2 model **Step 1: Local Development & Validation** 🖥️ @@ -346,7 +346,7 @@ primus-cli slurm sbatch \ -- train pretrain --config configs/deepseek_v2_prod.yaml ``` -### Key Insight +### Key insight Notice? **From development to production, the core command structure remains unchanged**: ``` @@ -357,7 +357,7 @@ Only the runtime environment (`direct` → `container` → `slurm`) changes - ev --- -## 🎯 Core Advantages Summary +## 🎯 Core advantages summary After the detailed introduction above, let's summarize the core value Primus CLI brings: @@ -373,11 +373,11 @@ After the detailed introduction above, let's summarize the core value Primus CLI --- -## 🛣️ Future Roadmap +## 🛣️ Future roadmap Primus CLI continues to evolve, and our near-term plans include: -### Short-Term Goals (2025) +### Short-term goals (2025) - 🎯 **Python Hook API**: Support writing Hooks in Python scripts for more flexible extension capabilities - 🎯 **Intelligent Preflight**: Auto-check GPU health, network topology, InfiniBand connectivity before launch - 🎯 **Configuration Template System**: Built-in best practice config templates for common models @@ -386,12 +386,12 @@ Primus CLI continues to evolve, and our near-term plans include: - 🎯 **Extended Framework Support**: Improve support for more training frameworks like TorchTitan, JAX/Flax - 🎯 **CI/CD Integration**: Provide standardized testing and validation workflows, support automated regression testing -### Long-Term Vision +### Long-term vision - 🌟 Become the **standard training entry point for the ROCm ecosystem** --- -## 🎓 Summary: The Power of One Command +## 🎓 Summary: The power of one command Back to the question at the beginning: How do we make large model training go from complex to simple? @@ -415,9 +415,10 @@ primus-cli direct -- train pretrain --config deepseek_v2.yaml --- -## 📚 Learn More +## 📚 Learn more -- 📖 **User Guide**: [PRIMUS-CLI-GUIDE.md](./PRIMUS-CLI-GUIDE.md) +- 📖 **CLI Reference (user guide)**: [cli-reference.md](../02-user-guide/cli-reference.md) +- 🏛 **System Architecture**: [architecture.md](./architecture.md) - 🔧 **Quick Start**: `primus-cli --help` - 💬 **Issue Reporting**: GitHub Issues - 🌐 **ROCm Ecosystem**: [rocm.github.io](https://rocm.github.io) diff --git a/docs/06-developer-guide/contributing.md b/docs/06-developer-guide/contributing.md new file mode 100644 index 000000000..b6e52b656 --- /dev/null +++ b/docs/06-developer-guide/contributing.md @@ -0,0 +1,142 @@ +# Contributing guide + +This guide summarizes how to set up a development environment, follow project conventions, run checks locally, and align with the CI pipeline. For test commands and layout, see [Testing Guide](testing.md). The repository root [CONTRIBUTING.md](../../CONTRIBUTING.md) repeats branch naming, commit style, and pull request steps in short form. + +## 1. Development setup + +1. **Clone the repository** (include submodules): + + ```bash + git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git + cd Primus + ``` + +2. **Install Python dependencies:** + + ```bash + pip install -r requirements.txt + ``` + +3. **Install pre-commit hooks** (recommended): + + ```bash + pip install pre-commit + pre-commit install + ``` + +4. **Optional—JAX / MaxText work:** + + ```bash + pip install -r requirements-jax.txt + ``` + +5. **Quick verification** (from the repository root, with Primus on your `PATH` or via the bundled launcher): + + ```bash + ./primus-cli direct -- benchmark gemm --M 4096 --N 4096 --K 4096 + ``` + +## 2. Code style + +Configuration lives in `.pre-commit-config.yaml`. Hooks run automatically on `git commit` after `pre-commit install`. + +| Tool | Version | Purpose | +|------|---------|---------| +| **black** | 24.8.0 | Python formatter, line length 110 | +| **isort** | 5.13.2 | Import sorting, profile `black` | +| **autoflake** | 2.3.1 | Removes unused imports and variables (see hook args for star imports and `__init__`) | +| **shellcheck** | 0.10.0.1 (shellcheck-py) | Shell script analysis | +| **pre-commit-hooks** | v4.0.1 | `trailing-whitespace`, `end-of-file-fixer`, `check-yaml`, `check-added-large-files`, `check-merge-conflict` | + +Manual one-off runs (repository root): + +```bash +black --line-length=110 . +isort --profile black . +autoflake --remove-all-unused-imports --remove-unused-variables --expand-star-imports --ignore-init-module-imports --recursive --in-place . +``` + +CI runs `pre-commit run --all-files --show-diff-on-failure`, so lint behavior follows `.pre-commit-config.yaml` rather than a separate hand-written list of formatter commands. + +## 3. Branch naming convention + +Format: + +```text +// +``` + +**Types:** `feat`, `fix`, `docs`, `refactor`, `test`, `chore`, `ci` + +**Scope (optional):** `engine`, `model`, `scheduler`, `docs`, `tests`, `config`, or another short area name. + +**Examples:** + +- `feat/model/implement-moe-routing` +- `fix/engine/init-error` + +## 4. Commit convention + +Follow [Conventional Commits](https://www.conventionalcommits.org/): + +```text +(): +``` + +**Examples:** + +- `feat(model): add MOE routing functionality` +- `fix(engine): resolve initialization error` + +## 5. Testing requirements + +Before opening a pull request, run the following from the repository root: + +- **Shell integration tests:** + + ```bash + bash ./tests/runner/run_all_tests.sh + ``` + +- **Python unit tests:** + + ```bash + pytest tests/unit_tests/ --maxfail=1 -s + ``` + +- **Backend / trainer tests** (GPU, datasets, and sometimes Hugging Face tokens): run the relevant file under `tests/trainer/` when your change touches that backend. See [Testing Guide](testing.md). + +- **Pre-commit on all files:** + + ```bash + pre-commit run --all-files + ``` + +## 6. Pull request process + +1. Fork the repository (unless you have write access and use a feature branch). +2. Create a branch that follows the naming convention above. +3. Implement changes and commit using the commit message convention. +4. Run tests and pre-commit locally. +5. Push and open a pull request with a clear description. +6. Reference related issues when applicable. +7. Request reviewers. +8. Address review feedback. +9. Ensure CI passes (lint and unit tests on the paths your PR triggers). + +## 7. CI pipeline + +The workflow **`.github/workflows/ci.yaml`** defines how changes are validated. + +**Triggers:** `workflow_dispatch`, pushes to `main`, tags matching `v*`, and pull requests. + +**Jobs (high level):** + +- **`code-lint`:** Ubuntu, Python 3.12—installs `pre-commit` and runs `pre-commit run --all-files --show-diff-on-failure`, so the checks (black, isort, autoflake, shellcheck, and the `pre-commit-hooks` set) follow `.pre-commit-config.yaml` exactly. +- **`build-docker`:** Builds and pushes Docker images (depends on `code-lint`). +- **`run-unittest-torch`:** Self-hosted GPU runner—installs dependencies (including Primus-Turbo and AITER as defined in the workflow), runs `bash ./tests/runner/run_all_tests.sh`, `pytest` on `tests/unit_tests/` (with a few deselected tests), then Megatron and TorchTitan trainer tests with `DATA_PATH`, `MASTER_PORT`, `HSA_NO_SCRATCH_RECLAIM=1`, and `HF_TOKEN` where required. +- **`run-unittest-jax`:** JAX runner—installs `requirements-jax.txt`, runs shell tests and `python ./tests/run_unit_tests.py --jax` with CI-specific environment variables. + +Lint checks mirror the pre-commit stack. Trainer jobs require GPU resources and shared secrets (for example `HF_TOKEN`) in the hosted environment. + +For a focused description of local vs CI test commands, see [Testing Guide](testing.md). diff --git a/docs/06-developer-guide/extending-backends.md b/docs/06-developer-guide/extending-backends.md new file mode 100644 index 000000000..cc42151fc --- /dev/null +++ b/docs/06-developer-guide/extending-backends.md @@ -0,0 +1,419 @@ +# Extending backends + +This guide explains how to add a **new training backend** to Primus using the current runtime architecture. It complements the high-level picture in [Primus overview](../01-getting-started/overview.md): adapters sit under the unified CLI and configuration system ([Configuration system](../02-user-guide/configuration-system.md)), and each backend plugs in through the same lifecycle and hook points as Megatron-LM, TorchTitan, MaxText, and the other integrated stacks. + +The runtime is built around: + +- **`BackendAdapter`** – integrates a backend framework +- **`BackendRegistry`** – discovers and instantiates adapters +- **`BaseTrainer`** – defines the minimal training lifecycle that all backends follow +- **`PrimusRuntime`** – orchestrates config loading, environment setup, patches, adapter, and trainer + +The examples below use a minimal **`dummy`** backend as a template. The dummy files are illustrative and are not checked into this repository; existing backends such as Megatron, TorchTitan, MaxText, Megatron Bridge, and HummingbirdXT show the production pattern. + +--- + +## What happens when you run Primus? + +When you run: + +```bash +primus train pretrain --config +``` + +the runtime (`PrimusRuntime`) does roughly: + +1. Load the experiment config—`load_primus_config()` returns a lightweight `SimpleNamespace` (not a `PrimusConfig`)—and select the `module_config` +2. Apply CLI overrides to `module_config.params` +3. Initialize environment (HF, logging, distributed environment, data directory) +4. Resolve backend adapter via `BackendRegistry.get_adapter(framework)` +5. Call `adapter.setup_backend_path(...)` to put the backend on `sys.path` +6. Call `adapter.prepare_backend(module_config)` (usually runs backend setup hooks) +7. Build backend arguments: + + ```python + backend_args = adapter.convert_config(module_config.params) + # run "build_args" patches and merge back into module_config.params + ``` + +8. Load and construct the trainer: + + ```python + TrainerClass = adapter.load_trainer_class(stage=module_config.params.stage or "pretrain") + trainer = TrainerClass(backend_args=backend_args) + ``` + +9. Execute the trainer lifecycle (with patches around it). Backend version detection is lazy during patch handling through `adapter.detect_backend_version()` rather than a separate pre-trainer step: + + ```python + # PrimusRuntime (the "build_args" patches from step 7 already ran + # before the trainer was constructed): + run_patches(phase="setup", backend_args=backend_args) + trainer.setup() + + trainer.init() + + run_patches(phase="before_train", backend_args=backend_args) + trainer.train() + run_patches(phase="after_train", backend_args=backend_args) + + trainer.cleanup() + ``` + +So a complete backend must provide: + +- An **adapter** subclassing `BackendAdapter` +- A **trainer** subclassing `BaseTrainer` and implementing `setup`, `init`, and `train` (and optionally overriding `cleanup`, which has a default no-op implementation) +- A small `primus.backends..__init__` that calls `BackendRegistry.register_adapter(...)` + +--- + +## Minimal backend layout + +Create a new backend folder under `primus/backends/`: + +```text +primus/backends/dummy/ +├── __init__.py +├── dummy_adapter.py +└── dummy_pretrain_trainer.py +``` + +This mirrors the pattern used by existing backends (for example Megatron, TorchTitan). + +--- + +## Implement the adapter (`BackendAdapter`) + +**File**: `primus/backends/dummy/dummy_adapter.py` + +```python +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +from primus.core.backend.backend_adapter import BackendAdapter +from primus.core.backend.backend_registry import BackendRegistry +from primus.core.utils.module_utils import log_rank_0 + + +class DummyAdapter(BackendAdapter): + """Minimal adapter for a 'dummy' backend.""" + + def __init__(self, framework: str = "dummy"): + super().__init__(framework) + + def setup_backend_path(self, backend_path=None) -> str: + """ + Dummy backend lives inside the Primus tree (no third_party submodule), + so we don't need to modify sys.path or resolve any external path here. + + For real backends that live under third_party/, you can rely on + the default implementation in BackendAdapter instead. + """ + log_rank_0("[Primus:DummyAdapter] setup_backend_path: no-op for in-tree dummy backend") + return "" + + def convert_config(self, params: Any) -> Any: + """ + Convert Primus module params → backend-specific args. + + For a real backend you would build a structured args object. Here we + just wrap the incoming params in a SimpleNamespace. + """ + if isinstance(params, dict): + backend_args = SimpleNamespace(**params) + else: + backend_args = params + log_rank_0("[Primus:DummyAdapter] Converted Primus params -> dummy backend_args") + return backend_args + + def detect_backend_version(self) -> str: + """Return a version string used by patch filtering.""" + return "dummy-0.1" + + def load_trainer_class(self, stage: str = "pretrain"): + """Return the Trainer class for the specified training stage.""" + from primus.backends.dummy.dummy_pretrain_trainer import DummyPretrainTrainer + + log_rank_0("[Primus:DummyAdapter] Loaded trainer class: DummyPretrainTrainer") + return DummyPretrainTrainer +``` + +Key points: + +- Since the dummy backend is implemented directly under `primus.backends.dummy` (not in `third_party/`), it overrides `setup_backend_path()` as a **no-op** so that the default third-party path resolution is skipped. +- `convert_config()` returns whatever your trainer expects as `backend_args`. +- `load_trainer_class()` imports and returns `DummyPretrainTrainer` directly (similar to `MegatronAdapter`), without going through a registry lookup. + +--- + +## Implement a runnable trainer (`BaseTrainer`) + +**File**: `primus/backends/dummy/dummy_pretrain_trainer.py` + +```python +from typing import Any + +from primus.core.trainer.base_trainer import BaseTrainer +from primus.core.utils.module_utils import log_rank_0 + + +class DummyPretrainTrainer(BaseTrainer): + """Minimal runnable trainer for the dummy backend.""" + + def __init__(self, backend_args: Any): + # BaseTrainer stores backend_args and reads torchrun env (rank, world_size, etc.) + super().__init__(backend_args=backend_args) + self._initialized = False + + def setup(self): + # Optional pre-init setup (e.g., logging, sanity checks) + log_rank_0(f"[DummyPretrainTrainer] setup() on rank={self.rank}") + + def init(self): + # Build your model / optimizer / dataloader here in a real backend. + log_rank_0("[DummyPretrainTrainer] init()") + self._initialized = True + + def train(self): + if not self._initialized: + raise RuntimeError("DummyPretrainTrainer.init() must be called before train().") + + log_rank_0("[DummyPretrainTrainer] train()") + # Example: access a custom param (e.g., 'hello') from backend_args. + hello_value = getattr(self.backend_args, "hello", "") + log_rank_0(f"[DummyPretrainTrainer] hello={hello_value}") + # Real training loop would go here. + log_rank_0("[DummyPretrainTrainer] training finished successfully.") + + def cleanup(self, on_error: bool = False): + # Optional cleanup logic (close files, finalize logging, etc.) + status = "error" if on_error else "success" + log_rank_0(f"[DummyPretrainTrainer] cleanup(on_error={status})") +``` + +Why this matches the core architecture: + +- `BaseTrainer.__init__` reads distributed environment from `get_torchrun_env()`. +- `PrimusRuntime` drives the lifecycle: `setup` → `init` → `train` → `cleanup` and runs patch phases around these steps. +- Your trainer only needs to implement `setup`, `init`, `train`, and `cleanup` using `backend_args` and the resolved environment information. + +--- + +## Register the adapter in `BackendRegistry` + +**File**: `primus/backends/dummy/__init__.py` + +```python +from primus.backends.dummy.dummy_adapter import DummyAdapter +from primus.core.backend.backend_registry import BackendRegistry + + +# Register adapter (backend name → adapter class) +BackendRegistry.register_adapter("dummy", DummyAdapter) +``` + +At runtime, when `framework: dummy` is requested: + +- `BackendRegistry.get_adapter("dummy")` lazily imports `primus.backends.dummy` (this file), which calls `register_adapter("dummy", DummyAdapter)`. +- The adapter instance is created and used by `PrimusRuntime` to set up the backend path, run setup hooks, build `backend_args`, and load and construct the trainer. + +--- + +## Minimal config example + +Create an experiment YAML (simplified; full template in the next section): + +```yaml +modules: + pre_trainer: + framework: dummy + config: dummy_trainer.yaml + model: dummy_8B.yaml +``` + +Run: + +```bash +./primus-cli direct -- train pretrain --config examples/dummy/configs/dummy_8B-pretrain.yaml +``` + +Because this dummy backend is an in-tree template and `setup_backend_path()` is a no-op, you should see logs similar to: + +- `[Primus:DummyAdapter] setup_backend_path: no-op for in-tree dummy backend` +- `[Primus:DummyAdapter] Converted Primus params -> dummy backend_args` +- `[DummyPretrainTrainer] setup()` +- `[DummyPretrainTrainer] init()` +- `[DummyPretrainTrainer] train()` + +--- + +## Example end-to-end YAML configs + +This template mirrors the Megatron pattern. Create these files only when you are actually adding a dummy backend for local development or tests: + +- The **top-level experiment config** lives under `examples//configs/...` +- The **module config** is resolved from `primus/configs/modules/{framework}/` +- The **model config** is resolved from `primus/configs/models/{framework}/` + +### Top-level experiment config + +**File 1**: `examples/dummy/configs/dummy_8B-pretrain.yaml` + +```yaml +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:dummy_8B-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: dummy + config: dummy_trainer.yaml + + # model to run + model: dummy_8B.yaml + + overrides: + # log / debug + stderr_sink_level: DEBUG + + # example training overrides (merged into module params) + train_iters: 100 + global_batch_size: 32 + micro_batch_size: 4 + seq_length: 1024 + hello: world +``` + +### Module-level trainer config + +**File 2**: `primus/configs/modules/dummy/dummy_trainer.yaml` + +```yaml +extends: + - trainer_base.yaml # optional, if you have a shared base; otherwise omit + +train_iters: 1000 +global_batch_size: 16 +micro_batch_size: 1 +seq_length: 512 + +log_interval: 1 +save_interval: 100 +``` + +This file defines **default training hyperparameters** for the `dummy` backend. Fields under `modules.pre_trainer.overrides` in the top-level config are deep-merged on top of these defaults. + +### Model-level config + +**File 3**: `primus/configs/models/dummy/dummy_8B.yaml` + +```yaml +extends: [] + +model_name: dummy_8B +vocab_size: 32000 +hidden_size: 4096 +num_layers: 32 +num_attention_heads: 32 +``` + +This file plays the same role as Megatron model configs under `primus/configs/models/megatron/`. It is loaded via: + +- `modules.pre_trainer.model: dummy_8B.yaml` +- resolved as `primus/configs/models/{framework}/dummy_8B.yaml` + +### Running the example + +Run: + +```bash +./primus-cli direct -- train pretrain --config examples/dummy/configs/dummy_8B-pretrain.yaml +``` + +Primus will: + +- load `examples/dummy/configs/dummy_8B-pretrain.yaml` +- resolve `modules.pre_trainer.config` → `primus/configs/modules/dummy/dummy_trainer.yaml` +- resolve `modules.pre_trainer.model` → `primus/configs/models/dummy/dummy_8B.yaml` +- build `module_config.params` from these sources plus `overrides` +- call `DummyAdapter.convert_config(params)` to build `backend_args` +- construct `DummyPretrainTrainer(backend_args=...)` +- execute `setup` → `init` → `train` → `cleanup`. + +For adding model YAML for existing backends (Megatron, TorchTitan, and others), see [Adding model configurations](./adding-models.md). + +--- + +## Checklist for a complete backend + +Use this as a quick checklist when adding a new backend: + +- [ ] Adapter subclass of `BackendAdapter` implements: + - `load_trainer_class(stage: str)` + - `convert_config(params)` + - `detect_backend_version()` + - (optionally) overrides `prepare_backend()` / `third_party_dir_name` +- [ ] Trainer subclass of `BaseTrainer` implements: + - `setup()`, `init()`, `train()`, and optional `cleanup(on_error: bool)` +- [ ] `BackendRegistry.register_adapter(backend, AdapterClass)` is called in `primus.backends..__init__` +- [ ] At least one unit test is added under `tests/unit_tests/backends/` + +Once these are in place, your backend is fully integrated into the Primus runtime and follows the same lifecycle and patch phases as the built-in backends. + +--- + +## Advanced: Backend-specific setup with train hooks + +For more advanced scenarios (for example installing extra Python packages or configuring backend-specific environment variables at runtime), you can use **train hooks** under `runner/helpers/hooks`. + +- **Hook locations for training**: + - Global hooks (run for all commands): `runner/helpers/hooks/*.sh` and `runner/helpers/hooks/*.py`. These are discovered with `find ... -maxdepth 1 \( -name "*.sh" -o -name "*.py" \)` and executed in **lexicographical order** of their filenames (see `runner/helpers/execute_hooks.sh`). + - Command-specific hooks: `runner/helpers/hooks/train/pretrain/*.sh|*.py` (and `.../posttrain/...`), discovered and ordered the same way. For pretrain, this directory contains the dispatcher `prepare_experiment.sh`. + - Per-framework hooks: `runner/helpers/hooks/train/pretrain//` and `runner/helpers/hooks/train/posttrain//`, where `` is `megatron`, `torchtitan`, `dummy`, and so on. These are **not** run directly by `execute_hooks`; instead `prepare_experiment.sh` detects the framework from the experiment config, runs that framework folder's `*.sh` files in lexicographical order, and then invokes the framework's `prepare.py` dispatcher. + +When you run: + +```bash +./primus-cli direct -- train pretrain --config +``` + +Primus will: + +- Call `execute_hooks train pretrain ...`, which: + - Runs global hooks under `runner/helpers/hooks/` (lexicographical order) + - Then runs command-specific hooks under `runner/helpers/hooks/train/pretrain/`, including `prepare_experiment.sh` + - `prepare_experiment.sh` resolves the framework from the config and runs the per-framework hooks under `runner/helpers/hooks/train/pretrain//` (its `*.sh` files in lexicographical order, then `prepare.py`) + +Each hook script can **emit control lines on stdout** that Primus parses (the framework hooks' stdout is captured through `prepare_experiment.sh`): + +- **`env.*=value` → environment variables** + + ```bash + # inside runner/helpers/hooks/train/pretrain//.sh + echo "env.MY_BACKEND_FLAG=1" # becomes: export MY_BACKEND_FLAG=1 + echo "env.PYTHONPATH=/opt/mylib:$PYTHONPATH" + ``` + + These are exported into the environment of the `primus-cli direct` process, so downstream backend code and trainers see them. + +- **`extra.*=value` → extra CLI arguments** + + ```bash + # inside the same hook + echo "extra.backend_path=/opt/my-backend" # becomes: --backend_path /opt/my-backend + echo "extra.train_data_path=/my/data" # becomes: --train_data_path /my/data + ``` + + These `extra.*` pairs are appended to the Primus CLI invocation as `-- ` after hook execution. + +Typical pattern to install or configure packages for a backend: + +- Add a script under `runner/helpers/hooks/train/pretrain//-setup.sh` (use a numeric prefix such as `000-` or `010-` to control ordering). +- In that script: + - Optionally run `python -m pip install ...` or other setup commands. + - Emit `env.*=...` lines to export any required environment variables. + - Emit `extra.*=...` lines if you need to pass additional CLI arguments (for example `backend_path`) into the Primus runtime for this run. diff --git a/docs/06-developer-guide/model-support-matrix.md b/docs/06-developer-guide/model-support-matrix.md new file mode 100644 index 000000000..f8f2ac545 --- /dev/null +++ b/docs/06-developer-guide/model-support-matrix.md @@ -0,0 +1,210 @@ +# Model support matrix + +This document summarizes which model families Primus targets per backend and lists representative checked-in model presets and example experiment YAML under the repository. It distinguishes **curated examples** from **theoretical** support (a preset or upstream stack may exist without a matching `examples/` entry). Use the filesystem under `primus/configs/models/` and `examples/*/configs/` as the authoritative live inventory. + +For how to add presets, see [Adding model configurations](./adding-models.md). Backend parameter references: [Megatron](../03-configuration-reference/megatron-parameters.md), [TorchTitan](../03-configuration-reference/torchtitan-parameters.md), [MaxText](../03-configuration-reference/maxtext-parameters.md), [Megatron Bridge](../03-configuration-reference/megatron-bridge-parameters.md). + +--- + +## Overview: Supported model families (high level) + +The following aligns with the backend overview and the configs present in this tree. + +| Backend | Model families (documentation / stack scope) | +| ------- | ---------------------------------------------- | +| **Megatron-LM** | LLaMA2 / LLaMA3 / LLaMA3.1 / LLaMA3.3 / LLaMA4 (sizes from small to 405B+), DeepSeek-V2 (including lite) and DeepSeek-V3, Mixtral MoE and large MoE recipe YAML, Qwen2.5 and Qwen3 (dense and MoE), Grok, GPT-OSS (20B / 120B), GLM, Kimi K2, LFM2, MiniMax, Zebra LLaMA, Mamba, and generic `language_model.yaml` bases. | +| **TorchTitan** | LLaMA3 family (including 3.1), LLaMA4 examples, DeepSeek-V3 examples, and Qwen3 examples including 0.6B, 1.7B, 4B, 8B, 14B, and 32B variants where present. Additional presets exist under `primus/configs/models/torchtitan/` without being exhaustively listed here. | +| **MaxText (JAX)** | LLaMA2 / LLaMA3 / LLaMA3.3, DeepSeek-V2 16B, Mixtral-8x7B, Grok1, Qwen3 14B / 30B-A3B (per presets and examples). Broader coverage may exist in upstream MaxText; see [MaxText](https://github.com/AI-Hypercomputer/maxtext). | +| **Megatron Bridge** | Qwen3 pretraining and post-training examples, plus post-training examples for Zebra LLaMA and Mamba where present. LLaMA 3.1 70B Bridge examples appear under MI355X. | +| **HummingbirdXT** | Registered backend with a post-training trainer and one checked-in example; user-facing support level still needs maintainer confirmation. | + +**Interpretation:** “Supported” in upstream code can exceed what this repository ships as YAML. Rows below reference representative files that exist under `primus/configs/models/` and `examples/`; they should not be treated as a complete generated inventory. + +--- + +## Megatron model configs + +Model presets live in `primus/configs/models/megatron/`. Example experiments that reference those presets appear under `examples/megatron/configs/MI300X/`, `MI325X/`, and `MI355X/`. + +For **TorchTitan**, the MI300X, MI325X, and MI355X example directories carry the same model set (21 configs each). For **Megatron**, MI300X and MI325X are nearly identical **except** that MI325X omits `qwen3_5_35B_A3B` (BF16 and FP8)—so MI300X has 70 example configs while MI325X has 68—and **MI355X** is a superset (99 configs; it adds models such as `glm5`, `gpt_oss_120B`, `kimi_k2`, `lfm2_8B_A1B`, and `minimax_m2.5`). Each row's SKU list below reflects exactly which SKUs ship a curated example (see, for example, `qwen3_5_35B_A3B`, which is MI300X/MI355X only). + +| Model name (file) | Preset path | Role | Example experiment dirs | Precision in examples | +| ----------------- | ----------- | ---- | ----------------------- | ---------------------- | +| `deepseek_v2.yaml` | `primus/configs/models/megatron/deepseek_v2.yaml` | Dense model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `deepseek_v2_base.yaml` | `primus/configs/models/megatron/deepseek_v2_base.yaml` | Base fragment (`extends` only) | — | — | +| `deepseek_v2_lite.yaml` | `primus/configs/models/megatron/deepseek_v2_lite.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `deepseek_v3.yaml` | `primus/configs/models/megatron/deepseek_v3.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `deepseek_v3_base.yaml` | `primus/configs/models/megatron/deepseek_v3_base.yaml` | Base fragment | — | — | +| `glm4_7.yaml` | `primus/configs/models/megatron/glm4_7.yaml` | Model preset | No curated example in this repo | — | +| `glm5.yaml` | `primus/configs/models/megatron/glm5.yaml` | Model preset | MI355X | BF16, FP8 | +| `gpt_oss_20B.yaml` | `primus/configs/models/megatron/gpt_oss_20B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `gpt_oss_120B.yaml` | `primus/configs/models/megatron/gpt_oss_120B.yaml` | Model preset | MI355X | BF16, FP8 | +| `grok1.yaml` | `primus/configs/models/megatron/grok1.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `grok2.yaml` | `primus/configs/models/megatron/grok2.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `grok_base.yaml` | `primus/configs/models/megatron/grok_base.yaml` | Base fragment | — | — | +| `hybrid_model_base.yaml` | `primus/configs/models/megatron/hybrid_model_base.yaml` | Base fragment | — | — | +| `kimi_k2.yaml` | `primus/configs/models/megatron/kimi_k2.yaml` | MoE model preset | MI355X | BF16, FP8 | +| `language_model.yaml` | `primus/configs/models/megatron/language_model.yaml` | Generic Megatron LM defaults | Used via `extends` | — | +| `lfm2_8B_A1B.yaml` | `primus/configs/models/megatron/lfm2_8B_A1B.yaml` | MoE model preset | MI355X | BF16, FP8 | +| `lfm_base.yaml` | `primus/configs/models/megatron/lfm_base.yaml` | Base fragment | — | — | +| `llama2_7B.yaml` | `primus/configs/models/megatron/llama2_7B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama2_13B.yaml` | `primus/configs/models/megatron/llama2_13B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama2_70B.yaml` | `primus/configs/models/megatron/llama2_70B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama2_base.yaml` | `primus/configs/models/megatron/llama2_base.yaml` | Base fragment | — | — | +| `llama_base.yaml` | `primus/configs/models/megatron/llama_base.yaml` | Base fragment | — | — | +| `llama3_8B.yaml` | `primus/configs/models/megatron/llama3_8B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama3_70B.yaml` | `primus/configs/models/megatron/llama3_70B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama3_base.yaml` | `primus/configs/models/megatron/llama3_base.yaml` | Base fragment | — | — | +| `llama3.1_8B.yaml` | `primus/configs/models/megatron/llama3.1_8B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama3.1_70B.yaml` | `primus/configs/models/megatron/llama3.1_70B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama3.1_405B.yaml` | `primus/configs/models/megatron/llama3.1_405B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama3.2_1B.yaml` | `primus/configs/models/megatron/llama3.2_1B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama3.2_3B.yaml` | `primus/configs/models/megatron/llama3.2_3B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama3.3_70B.yaml` | `primus/configs/models/megatron/llama3.3_70B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama4_17B128E.yaml` | `primus/configs/models/megatron/llama4_17B128E.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama4_17B16E.yaml` | `primus/configs/models/megatron/llama4_17B16E.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `llama4_base.yaml` | `primus/configs/models/megatron/llama4_base.yaml` | Base fragment | — | — | +| `mamba_370M.yaml` | `primus/configs/models/megatron/mamba_370M.yaml` | Model preset | MI300X, MI325X, MI355X | Set in experiment overrides | +| `mamba_base.yaml` | `primus/configs/models/megatron/mamba_base.yaml` | Base fragment | — | — | +| `minimax_m2.5.yaml` | `primus/configs/models/megatron/minimax_m2.5.yaml` | MoE model preset | MI355X | BF16, FP8 | +| `mixtral_8x7B_v0.1.yaml` | `primus/configs/models/megatron/mixtral_8x7B_v0.1.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `mixtral_8x22B_v0.1.yaml` | `primus/configs/models/megatron/mixtral_8x22B_v0.1.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `mixtral_base.yaml` | `primus/configs/models/megatron/mixtral_base.yaml` | Base fragment | — | — | +| `moe_515B.yaml` | `primus/configs/models/megatron/moe_515B.yaml` | Large MoE template | No curated example in this repo | — | +| `moe_1T.yaml` | `primus/configs/models/megatron/moe_1T.yaml` | Large MoE template | No curated example in this repo | — | +| `moe_2T.yaml` | `primus/configs/models/megatron/moe_2T.yaml` | Large MoE template | No curated example in this repo | — | +| `moe_4T.yaml` | `primus/configs/models/megatron/moe_4T.yaml` | Large MoE template | No curated example in this repo | — | +| `moe_proxy_single_node.yaml` | `primus/configs/models/megatron/moe_proxy_single_node.yaml` | MoE proxy / test template | No curated example in this repo | — | +| `primus_megatron_model.yaml` | `primus/configs/models/megatron/primus_megatron_model.yaml` | Primus Megatron root defaults | Used via `extends` | — | +| `qwen2.5_3B.yaml` | `primus/configs/models/megatron/qwen2.5_3B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `qwen2.5_7B.yaml` | `primus/configs/models/megatron/qwen2.5_7B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `qwen2.5_14B.yaml` | `primus/configs/models/megatron/qwen2.5_14B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `qwen2.5_32B.yaml` | `primus/configs/models/megatron/qwen2.5_32B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `qwen2.5_72B.yaml` | `primus/configs/models/megatron/qwen2.5_72B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `qwen2.5_base.yaml` | `primus/configs/models/megatron/qwen2.5_base.yaml` | Base fragment | — | — | +| `qwen3_4B.yaml` | `primus/configs/models/megatron/qwen3_4B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `qwen3_8B.yaml` | `primus/configs/models/megatron/qwen3_8B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `qwen3_14B.yaml` | `primus/configs/models/megatron/qwen3_14B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `qwen3_32B.yaml` | `primus/configs/models/megatron/qwen3_32B.yaml` | Model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `qwen3_30B_A3B.yaml` | `primus/configs/models/megatron/qwen3_30B_A3B.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `qwen3_5_35B_A3B.yaml` | `primus/configs/models/megatron/qwen3_5_35B_A3B.yaml` | MoE model preset | MI300X, MI355X | BF16, FP8 | +| `qwen3_235B_A22B.yaml` | `primus/configs/models/megatron/qwen3_235B_A22B.yaml` | MoE model preset | MI300X, MI325X, MI355X | BF16, FP8 | +| `zebra_llama_1B.yaml` | `primus/configs/models/megatron/zebra_llama_1B.yaml` | Model preset | MI300X, MI325X, MI355X | Set in experiment overrides | +| `zebra_llama_3B.yaml` | `primus/configs/models/megatron/zebra_llama_3B.yaml` | Model preset | MI300X, MI325X, MI355X | Set in experiment overrides | +| `zebra_llama_8B.yaml` | `primus/configs/models/megatron/zebra_llama_8B.yaml` | Model preset | MI300X, MI325X, MI355X | Set in experiment overrides | + +**Parallelism:** Tensor, pipeline, and expert parallel sizes are **not** fixed in model presets; they are set in experiment `overrides` (for example `tensor_model_parallel_size`, `pipeline_model_parallel_size`, `expert_model_parallel_size`). MoE presets such as `qwen3_235B_A22B.yaml` typically require non-default expert parallelism in real runs—see the matching experiment YAML. + +--- + +## TorchTitan model configs + +Presets: `primus/configs/models/torchtitan/`. Examples: `examples/torchtitan/configs/MI300X/`, `MI325X/`, and `MI355X/`. + +| Model name (file) | Preset path | Example experiment dirs | Precision in examples | +| ----------------- | ----------- | ------------------------- | ---------------------- | +| `deepseek_v3_16b.yaml` | `primus/configs/models/torchtitan/deepseek_v3_16b.yaml` | MI300X, MI325X, MI355X | BF16 | +| `deepseek_v3_16b-fp8.yaml` | `primus/configs/models/torchtitan/deepseek_v3_16b-fp8.yaml` | MI300X, MI325X, MI355X | FP8 | +| `deepseek_v3_236b.yaml` | `primus/configs/models/torchtitan/deepseek_v3_236b.yaml` | MI300X, MI325X, MI355X | BF16 | +| `deepseek_v3_236b-fp8.yaml` | `primus/configs/models/torchtitan/deepseek_v3_236b-fp8.yaml` | MI300X, MI325X, MI355X | FP8 | +| `deepseek_v3_671b.yaml` | `primus/configs/models/torchtitan/deepseek_v3_671b.yaml` | MI300X, MI325X, MI355X | (see experiment) | +| `deepseek_v3_671b-fp8.yaml` | `primus/configs/models/torchtitan/deepseek_v3_671b-fp8.yaml` | Preset only; stock examples use `deepseek_v3_671b.yaml` | — | +| `llama3_8B.yaml` | `primus/configs/models/torchtitan/llama3_8B.yaml` | No example in this repo | — | +| `llama3_8B-fp8.yaml` | `primus/configs/models/torchtitan/llama3_8B-fp8.yaml` | No example in this repo | — | +| `llama3_70B.yaml` | `primus/configs/models/torchtitan/llama3_70B.yaml` | No example in this repo | — | +| `llama3_70B-fp8.yaml` | `primus/configs/models/torchtitan/llama3_70B-fp8.yaml` | No example in this repo | — | +| `llama3.1_8B.yaml` | `primus/configs/models/torchtitan/llama3.1_8B.yaml` | MI300X, MI325X, MI355X | BF16 | +| `llama3.1_8B-fp8.yaml` | `primus/configs/models/torchtitan/llama3.1_8B-fp8.yaml` | MI300X, MI325X, MI355X | FP8 | +| `llama3.1_70B.yaml` | `primus/configs/models/torchtitan/llama3.1_70B.yaml` | MI300X, MI325X, MI355X | BF16 | +| `llama3.1_70B-fp8.yaml` | `primus/configs/models/torchtitan/llama3.1_70B-fp8.yaml` | MI300X, MI325X, MI355X | FP8 | +| `llama3.1_405B.yaml` | `primus/configs/models/torchtitan/llama3.1_405B.yaml` | MI300X, MI325X, MI355X | BF16 | +| `llama3.1_405B-fp8.yaml` | `primus/configs/models/torchtitan/llama3.1_405B-fp8.yaml` | MI300X, MI325X, MI355X | FP8 | +| `llama3.2_1B.yaml` | `primus/configs/models/torchtitan/llama3.2_1B.yaml` | No example in this repo | — | +| `llama3.3_70B.yaml` | `primus/configs/models/torchtitan/llama3.3_70B.yaml` | No example in this repo | — | +| `llama3.3_70B-fp8.yaml` | `primus/configs/models/torchtitan/llama3.3_70B-fp8.yaml` | No example in this repo | — | +| `llama4_17Bx128E.yaml` | `primus/configs/models/torchtitan/llama4_17Bx128E.yaml` | MoE; MI300X, MI325X, MI355X | BF16 | +| `llama4_17Bx128E-fp8.yaml` | `primus/configs/models/torchtitan/llama4_17Bx128E-fp8.yaml` | MoE; MI300X, MI325X, MI355X | FP8 | +| `llama4_17Bx16E.yaml` | `primus/configs/models/torchtitan/llama4_17Bx16E.yaml` | MoE; MI300X, MI325X, MI355X | BF16 | +| `llama4_17Bx16E-fp8.yaml` | `primus/configs/models/torchtitan/llama4_17Bx16E-fp8.yaml` | MoE; MI300X, MI325X, MI355X | FP8 | +| `qwen3_0.6b.yaml` | `primus/configs/models/torchtitan/qwen3_0.6b.yaml` | MI300X, MI325X, MI355X | (see experiment) | +| `qwen3_1.7b.yaml` | `primus/configs/models/torchtitan/qwen3_1.7b.yaml` | MI300X, MI325X, MI355X | (see experiment) | +| `qwen3_4b.yaml` | `primus/configs/models/torchtitan/qwen3_4b.yaml` | MI300X, MI325X, MI355X | (see experiment) | +| `qwen3_8b.yaml` | `primus/configs/models/torchtitan/qwen3_8b.yaml` | MI300X, MI325X, MI355X | (see experiment) | +| `qwen3_14b.yaml` | `primus/configs/models/torchtitan/qwen3_14b.yaml` | MI300X, MI325X, MI355X | (see experiment) | +| `qwen3_32b.yaml` | `primus/configs/models/torchtitan/qwen3_32b.yaml` | MI300X, MI325X, MI355X | (see experiment) | + +**Parallelism:** Controlled by TorchTitan launch configuration and Primus module overrides (see TorchTitan patch notes and [TorchTitan parameters](../03-configuration-reference/torchtitan-parameters.md)); not embedded in the small `job` / `model` preset alone. + +--- + +## MaxText model configs + +Presets: `primus/configs/models/maxtext/`. Examples: `examples/maxtext/configs/MI300X/` and `examples/maxtext/configs/MI355X/`. + +| Model name (file) | Preset path | Example experiment dirs | +| ----------------- | ----------- | ------------------------ | +| `deepseek_v2_16B.yaml` | `primus/configs/models/maxtext/deepseek_v2_16B.yaml` | MI300X, MI355X | +| `grok1.yaml` | `primus/configs/models/maxtext/grok1.yaml` | MI300X | +| `llama2_7B.yaml` | `primus/configs/models/maxtext/llama2_7B.yaml` | MI300X, MI355X | +| `llama2_70B.yaml` | `primus/configs/models/maxtext/llama2_70B.yaml` | MI300X, MI355X | +| `llama3_8B.yaml` | `primus/configs/models/maxtext/llama3_8B.yaml` | MI300X, MI355X | +| `llama3_70B.yaml` | `primus/configs/models/maxtext/llama3_70B.yaml` | MI300X, MI355X | +| `llama3.1_405B.yaml` | `primus/configs/models/maxtext/llama3.1_405B.yaml` | MI355X | +| `llama3.3_70B.yaml` | `primus/configs/models/maxtext/llama3.3_70B.yaml` | MI300X, MI355X | +| `mixtral_8x7B.yaml` | `primus/configs/models/maxtext/mixtral_8x7B.yaml` | MI300X, MI355X | +| `qwen3_14B.yaml` | `primus/configs/models/maxtext/qwen3_14B.yaml` | MI300X, MI355X | +| `qwen3_30B_A3B.yaml` | `primus/configs/models/maxtext/qwen3_30B_A3B.yaml` | MI300X, MI355X | +| `model_base.yaml` | `primus/configs/models/maxtext/model_base.yaml` | Extended by other presets (not a standalone run) | + +**Parallelism:** JAX / MaxText sharding is configured in experiment overrides (for example `ici_fsdp_parallelism`, `ici_data_parallelism`, `dcn_*` in sample experiments). See [MaxText parameters](../03-configuration-reference/maxtext-parameters.md). + +--- + +## Megatron Bridge model configs + +Presets: `primus/configs/models/megatron_bridge/`. Examples: `examples/megatron_bridge/configs/MI300X/` and `examples/megatron_bridge/configs/MI355X/`. + +| Model name (file) | Preset path | Recipe / flavor (from preset) | Example experiment dirs | +| ----------------- | ----------- | ----------------------------- | ------------------------ | +| `qwen3_8b.yaml` | `primus/configs/models/megatron_bridge/qwen3_8b.yaml` | `qwen.qwen3` / `qwen3_8b_finetune_config` | MI300X pretrain, MI355X posttrain | +| `qwen3_32b.yaml` | `primus/configs/models/megatron_bridge/qwen3_32b.yaml` | `qwen.qwen3` / `qwen3_32b_finetune_config` | MI300X, MI355X | +| `llama31_70b.yaml` | `primus/configs/models/megatron_bridge/llama31_70b.yaml` | `llama.llama3` / `llama31_70b_finetune_config` | MI355X | +| `zebra_llama_1B.yaml`, `zebra_llama_3B.yaml`, `zebra_llama_8B.yaml` | `primus/configs/models/megatron_bridge/` | Zebra LLaMA presets | MI300X posttrain | +| `mamba_370M.yaml` | `primus/configs/models/megatron_bridge/mamba_370M.yaml` | Mamba preset | MI300X posttrain | + +Example filenames include `*_pretrain.yaml`, `*_sft_posttrain.yaml`, and `*_lora_posttrain.yaml`; precision such as `bf16_mixed` is set in experiment `overrides`. + +--- + +## Hardware compatibility (example directories) + +Curated example layouts under `examples/` use GPU SKU subdirectories. As of this document: + +| GPU SKU | `examples/megatron/configs/` | `examples/torchtitan/configs/` | `examples/maxtext/configs/` | `examples/megatron_bridge/configs/` | +| ------- | ---------------------------- | ------------------------------ | ---------------------------- | ----------------------------------- | +| **MI300X** | Yes | Yes | Yes | Yes | +| **MI355X** | Yes | Yes | Yes | Yes | +| **MI325X** | Yes | Yes | No | No | + +Megatron and TorchTitan ship MI325X example directories in addition to MI300X and MI355X examples. MaxText includes MI300X and MI355X examples, including MI355X-only entries such as `llama3.1_405B-pretrain.yaml`. Megatron Bridge MI300X examples include Qwen3 8B and 32B pretraining plus Qwen3 32B, Zebra LLaMA, and Mamba post-training examples; LLaMA 3.1 70B Bridge examples appear under MI355X. + +Absence of a SKU directory for a given backend does **not** imply the backend cannot run there; it means this tree does not currently provide a checked-in example path to copy from. + +--- + +## Model architecture reference (Megatron presets) + +Values below come from `primus/configs/models/megatron/` presets (merged through `extends`). **Vocabulary size** is usually defined by the tokenizer / Hugging Face config, not duplicated in every YAML; **context** is `max_position_embeddings` where set in the chain. Use this table as a quick reference for common sizes—not an exhaustive spec of every parameter. + +| Model family | Example preset | Hidden size | Layers | Attention heads | KV heads (GQA) | Max position (context) | +| ------------ | -------------- | ----------- | ------ | ----------------- | ---------------- | ------------------------ | +| LLaMA 2 7B | `llama2_7B.yaml` | 4096 | 32 | 32 | 32 (no GQA) | From `llama2_base` / tokenizer | +| LLaMA 3 8B | `llama3_8B.yaml` | 4096 | 32 | 32 | 8 | 8192 (`llama3_base`) | +| LLaMA 3 70B | `llama3_70B.yaml` | 8192 | 80 | 64 | 8 | 8192 | +| LLaMA 3.1 405B | `llama3.1_405B.yaml` | 16384 | 126 | 128 | 8 | 8192 | +| Qwen3 8B | `qwen3_8B.yaml` | 4096 | 36 | 32 | 8 | 131072 (`qwen2.5_base` chain) | +| Mixtral 8x7B | `mixtral_8x7B_v0.1.yaml` | 4096 | 32 | 32 | — | 4096 | +| DeepSeek-V3 (MoE) | `deepseek_v3.yaml` | 7168 | 61 | 128 (MLA) | — | See preset / HF | +| Mamba 370M | `mamba_370M.yaml` | (Mamba stack) | — | — | — | — | + +For MoE and hybrid architectures (LLaMA 4, Qwen3-MoE, large `moe_*.yaml` templates), refer to the full YAML and upstream model cards; headline dimensions alone do not capture expert layout or MLA. diff --git a/docs/06-developer-guide/testing.md b/docs/06-developer-guide/testing.md new file mode 100644 index 000000000..9c80203b7 --- /dev/null +++ b/docs/06-developer-guide/testing.md @@ -0,0 +1,122 @@ +# Testing guide + +This guide describes where tests live, how to run them locally, and how they map to CI. For coding standards and PR workflow, see [Contributing Guide](contributing.md). The canonical CI definition is `.github/workflows/ci.yaml`. + +## 1. Test organization + +Layout (simplified from the repository root): + +```text +Primus/ +├── runner/ +│ └── lib/ +│ └── common.sh # Shared logging/helpers sourced by the shell test runner +├── tests/ +│ ├── runner/ # Shell integration tests +│ │ ├── run_all_tests.sh # Master shell test runner +│ │ ├── lib/ # test_common.sh, test_config.sh, test_validation.sh +│ │ ├── helpers/ # Hook and env tests +│ │ ├── test_primus_cli.sh +│ │ ├── test_primus_cli_direct.sh +│ │ ├── test_primus_cli_container.sh +│ │ └── test_primus_cli_slurm.sh +│ ├── unit_tests/ # Python unit tests (pytest) +│ │ ├── agents/ # Tuning-agent tests +│ │ ├── backends/ +│ │ ├── ci/ +│ │ ├── cli/ +│ │ ├── core/ # config, backend, launcher, patches, projection, pipeline_parallel, runtime, trainer, utils +│ │ ├── megatron/ # Megatron-specific unit tests +│ │ ├── modules/ +│ │ └── tools/ +│ ├── trainer/ # Integration tests (typically need GPU + data) +│ │ ├── test_megatron_trainer.py +│ │ ├── test_torchtitan_trainer.py +│ │ └── test_maxtext_trainer.py +│ ├── scripts/ # CI unit/integration launch scripts and UT patches +│ ├── utils.py # Shared test utilities +│ └── run_unit_tests.py # Optional orchestrator (walks tests/, see below) +``` + +`tests/runner/run_all_tests.sh` sources shared helpers from **`runner/lib/common.sh`** at the repository root (not under `tests/runner/`). + +## 2. Running tests + +**Shell integration tests** (CLI behavior, config loading, hooks, environment): + +```bash +bash ./tests/runner/run_all_tests.sh +``` + +**Python unit tests:** + +```bash +pytest tests/unit_tests/ --maxfail=1 -s +``` + +**Trainer integration tests** (GPU and data; may require Hugging Face access): + +```bash +# Megatron +DATA_PATH= pytest tests/trainer/test_megatron_trainer.py -s + +# TorchTitan +DATA_PATH= pytest tests/trainer/test_torchtitan_trainer.py -s + +# MaxText (JAX) — often run via the orchestrator in CI +python ./tests/run_unit_tests.py --jax +``` + +`tests/run_unit_tests.py` walks **`tests/`** and runs every `test_*.py` it finds, except for **`tests/trainer/test_maxtext_trainer.py`** in the default mode (that file is only selected when **`--jax`** is set). That means the default orchestrator run includes **`tests/unit_tests/`** and **`tests/trainer/`** (and any other matching tests), which is broader than `pytest tests/unit_tests/` alone. + +**Orchestrator (default—all discovered tests except MaxText trainer):** + +```bash +python ./tests/run_unit_tests.py +``` + +**Orchestrator (JAX / MaxText trainer only):** + +```bash +python ./tests/run_unit_tests.py --jax +``` + +## 3. Test types + +- **Shell tests:** Exercise runner scripts, CLI wiring, configuration loading, hook execution, and environment setup. Implemented as bash scripts under `tests/runner/` and orchestrated by `run_all_tests.sh`. +- **Unit tests:** Cover configuration parsing, preset loading, CLI behavior, patch registration, adapters, and other library logic under `tests/unit_tests/`. +- **Trainer tests:** End-to-end training against real backends; require AMD GPUs and appropriate data paths (and sometimes tokens). See `.github/workflows/ci.yaml` for CI values such as `DATA_PATH`, `MASTER_PORT`, and `HSA_NO_SCRATCH_RECLAIM`. + +## 4. Writing new tests + +- **Pytest:** Add files named `test_*.py` under `tests/unit_tests/`, following existing patterns and reusing fixtures from `conftest.py` where present. +- **Shell:** Add scripts under `tests/runner/` or extend `tests/runner/run_all_tests.sh` to invoke new suites, consistent with existing `test_primus_cli*.sh` scripts. +- **Backends:** Prefer `tests/unit_tests/backends//` for adapter-focused tests. + +## 5. CI pipeline details + +From `.github/workflows/ci.yaml`: + +- **`code-lint`:** Python 3.12 on GitHub-hosted runners. Runs `pre-commit run --all-files --show-diff-on-failure`, so checks follow `.pre-commit-config.yaml`. +- **`dependency-review`:** Runs `actions/dependency-review-action` on pull requests to flag dependency changes. +- **`run-unittest-torch`:** Self-hosted GPU runner. Installs `requirements.txt`, runs `bash ./tests/runner/run_all_tests.sh`, then `pytest tests/unit_tests/` under coverage (`--cov=primus --cov-report=term-missing`) with specific tests `--deselect`ed (currently some `megatron/cco` TP-overlap and `megatron/transformer/moe` dispatcher cases—see the workflow file). Trainer steps set `MASTER_PORT`, `DATA_PATH`, `HSA_NO_SCRATCH_RECLAIM=1`, and `HF_TOKEN` for Megatron and TorchTitan trainer tests. A follow-up **coverage** step combines unit and E2E coverage. +- **`run-unittest-jax`:** JAX runner. Installs `requirements-jax.txt`, runs the same shell test script, then `python ./tests/run_unit_tests.py --jax` with CI environment variables (for example `JAX_SKIP_UT=1` and `DATA_PATH` as defined in the workflow). + +The **`build-docker`** job builds images after lint passes; unit test jobs depend on **`code-lint`**, not on **`build-docker`**. + +## 6. Pre-commit hooks + +Install once per clone: + +```bash +pip install pre-commit +pre-commit install +``` + +Run manually on the whole tree: + +```bash +pre-commit run --all-files +``` + +Hooks include: `trailing-whitespace`, `end-of-file-fixer`, `check-yaml`, `check-added-large-files`, `check-merge-conflict`, `isort`, `autoflake`, `black`, and `shellcheck` (as configured in `.pre-commit-config.yaml`). These align with the **`code-lint`** job in CI; see [Contributing Guide](contributing.md) for manual equivalents. diff --git a/docs/06-developer-guide/tooling.md b/docs/06-developer-guide/tooling.md new file mode 100644 index 000000000..982755a41 --- /dev/null +++ b/docs/06-developer-guide/tooling.md @@ -0,0 +1,24 @@ +# Tooling + +Primus ships a set of auxiliary tools for analysis, benchmarking, visualization, installation, and diagnostics. They live under [`tools/`](../../tools/README.md) in the repository; each tool keeps its own README with detailed usage instructions. This page is a lightweight index so the tooling is discoverable from the documentation. + +## Available tools + +| Tool | Directory | What it does | Docs | +|------|-----------|--------------|------| +| **IRLens** | `tools/IRLens/` | Parses XLA HLO text dumps and prints an execution skeleton with control flow, separating communication vs compute ops. | [README](../../tools/IRLens/README.md) | +| **model_stats** | `tools/model_stats/` | Generates charts from the model config registry under `primus/configs/models`. | [README](../../tools/model_stats/README.md) | +| **Pipeline Visualization** | `tools/visualization/pp_vis/` | Visualizes pipeline-parallelism schedules from dumped data or PP-simulator JSON via a local web UI. | [README](../../tools/visualization/pp_vis/README.md) | +| **Auto Benchmark** | `tools/auto_benchmark/` | Interactive benchmark menu for Megatron/TorchTitan on MI300X/MI355X with metrics collection. | [README](../../tools/auto_benchmark/Primus_Auto_Benchmark_README.md) | +| **Backend Gap Report / Engineering Dashboard** | `tools/backend_gap_report/` | Generation and publishing toolchain for the shared Primus engineering dashboard and backend-gap reports. | [README](../../tools/backend_gap_report/README.md) | +| **Installation (venv)** | `tools/installation/` | Reproduces the Primus training Docker environment in a Python virtual environment (no Docker, no sudo). | [README](../../tools/installation/README.md) | +| **Daily Report** | `tools/daily/` | Benchmark summary CSV generation used by CI workflows. | — | +| **Docker Helpers** | `tools/docker/` | Container startup and proxy scripts. | — | +| **Profile Trace** | `tools/profile_trace/` | Trace-file merging utility. | — | + +## Related documentation + +- [Primus tools](../02-user-guide/primus-tools.md)—the full catalog of Primus tools (CLI, tuning agent, ecosystem) with how-to starting points. +- [Tools overview](../../tools/README.md)—the top-level index maintained alongside the code. +- [Profiling and observability](../04-technical-guides/profiling-and-observability.md)—how these tools fit into performance analysis. +- [Benchmarking](../02-user-guide/benchmarking.md)—running the benchmark suites the tools summarize. diff --git a/docs/README.md b/docs/README.md index eb68c4fa5..bebc45bb5 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,72 +1,146 @@ -# Primus Documentation +# Primus documentation -Welcome to the Primus documentation! This guide will help you get started with training large-scale foundation models on AMD GPUs. +Production documentation for **Primus**, a large-scale foundation model training framework for AMD GPUs. -## 📚 Documentation Structure +--- + +## Choose your starting point + +| I am a... | Start here | +|-----------|------------| +| **New user** | [Getting started](./01-getting-started/overview.md) | +| **User** running training jobs | [User guide](./02-user-guide/pretraining.md) | +| **User** writing YAML configurations | [Configuration reference](./03-configuration-reference/megatron-parameters.md) | +| **Engineer** tuning performance | [Technical guides](./04-technical-guides/performance-tuning.md) | +| **Operator** deploying to production | [Operations](./05-operations/deployment.md) | +| **Contributor** to the codebase | [Developer guide](./06-developer-guide/architecture.md) | + +--- + +## Documentation structure + +### [Getting started](./01-getting-started/) + +Start here if you are new to Primus. -### 🚀 Getting Started +- [Project overview](./01-getting-started/overview.md): what Primus does, who it is for, key capabilities +- [Installation guide](./01-getting-started/installation.md): prerequisites, Docker/bare-metal/Slurm setup +- [Quickstart](./01-getting-started/quickstart.md): first training run in 5 minutes +- [Glossary](./01-getting-started/glossary.md): terms, acronyms, and domain concepts -Start here if you're new to Primus: +### [User guide](./02-user-guide/) -- **[Quick Start Guide](./quickstart.md)** - Get up and running in 5 minutes -- **[CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md)** - Complete command-line reference -- **[CLI Architecture](./cli/CLI-ARCHITECTURE.md)** - Design philosophy and deep dive +Core workflows and day-to-day usage. -### 📖 User Guides +- [CLI reference](./02-user-guide/cli-reference.md): `primus-cli` modes, flags, and subcommands +- [Configuration system](./02-user-guide/configuration-system.md): YAML configuration model, presets, overrides, inheritance +- [Pretraining](./02-user-guide/pretraining.md): pretraining **concepts**: backends, YAML structure, parallelism, configuration inventory +- [Backend training recipes](./02-user-guide/training-recipes.md): pretraining **commands**: copy-paste, GPU-arch-specific run commands +- [Post-training](./02-user-guide/posttraining.md): SFT and LoRA fine-tuning via Megatron Bridge +- [Benchmarking](./02-user-guide/benchmarking.md): GEMM, RCCL, and dense-GEMM benchmark suites +- [Preflight](./02-user-guide/preflight.md): cluster diagnostics and environment validation +- [Projection](./02-user-guide/projection.md): memory and performance projection tools +- [Tuning agent](./02-user-guide/tuning-agent.md): LLM-driven search for an optimal training configuration (uses projection as an oracle) +- [Primus tools](./02-user-guide/primus-tools.md): catalog of all Primus tools and ecosystem projects with how-to starting points -Guides for common workflows and features: +### [Configuration reference](./03-configuration-reference/) -- **[Configuration Guide](./configuration.md)** - YAML/TOML configuration, recommended patterns, and examples -- **[Slurm & Container Usage](./slurm-container.md)** - Distributed training and containerization workflows -- **[Experiment Management](./experiments.md)** - Organizing and tracking your training runs +Parameter references for Primus presets, backend-facing keys, and commonly used environment variables. -### 🔧 Technical References +- [Megatron parameters](./03-configuration-reference/megatron-parameters.md): Megatron-LM backend YAML parameters and Primus overrides +- [TorchTitan parameters](./03-configuration-reference/torchtitan-parameters.md): Primus TorchTitan preset keys and common JobConfig fields +- [MaxText parameters](./03-configuration-reference/maxtext-parameters.md): Primus MaxText overlay defaults and common fields +- [Megatron Bridge parameters](./03-configuration-reference/megatron-bridge-parameters.md): Megatron Bridge recipe, SFT, and pretraining fields surfaced through Primus +- [Environment variables](./03-configuration-reference/environment-variables.md): practical reference for commonly encountered environment variables -In-depth technical documentation: +### [Technical guides](./04-technical-guides/) -- **[Post-Training Guide](./posttraining.md)** - Fine-tuning with SFT and LoRA using Primus CLI -- **[Native SFT & LoRA Quick Start](./README_NATIVE_SFT_LORA_EN.md)** - Megatron-native SFT/LoRA launch guide (BF16/FP8/FP4), no Megatron-Bridge runtime dependency -- **[Performance Projection](./projection.md)** - Project training performance and memory to multi-node configurations -- **[Tuning Agent](./tuning_agent.md)** - LLM-driven search for an optimal training config — parallelism plus batching, schedule, memory, MoE-comm, and precision knobs (drives the projection tool as an oracle) -- **[Preflight](./preflight.md)** - Cluster diagnostics (host/GPU/network info + perf tests) -- **[Benchmark Suite](./benchmark.md)** - GEMM, RCCL, end-to-end benchmarks and profiling -- **[Supported Models](./backends/overview.md#supported-models)** - Supported LLM architectures and feature compatibility matrix -- **[Advanced Features](./advanced.md)** - Mixed precision, parallelism strategies, optimization techniques -- **[Backend Patch Notes](./backends/overview.md)** - Primus-specific arguments for Megatron, TorchTitan, etc. -- **[Backend Extension Guide](./backends/extending-backends.md)** - How to add a new backend using the current adapter/trainer architecture - - **[Megatron Model Extension Guide](./backends/adding-megatron-models.md)** - How to add a new Megatron model config - - **[TorchTitan Model Extension Guide](./backends/adding-torchtitan-models.md)** - How to add a new TorchTitan model config -- **[Flux Diffusion Models](./backends/megatron/diffusion/README.md)** - Flux diffusion model architecture, training, and API reference -- **[FP8 Training Guide](./backends/megatron/diffusion/fp8_training.md)** - FP8 precision training on AMD MI300X/MI355X: configuration, benchmarks, and tuning +Deep technical topics for advanced users. -### 💡 Help & Support +- [Parallelism strategies](./04-technical-guides/parallelism-strategies.md): DP, TP, PP, SP, CP, EP, FSDP explained +- [Parallelism configuration](./04-technical-guides/parallelism-configuration.md): per-backend parallelism setup and batch size relationships +- [Collective operations](./04-technical-guides/collective-operations.md): NCCL/RCCL operations and their role in each parallelism strategy +- [Performance tuning](./04-technical-guides/performance-tuning.md): HipBLASLt, Primus-Turbo, FP8, MoE optimization +- [MoE training deep-dive](./04-technical-guides/moe-training.md): bottlenecks and Primus-Turbo optimizations for Mixture-of-Experts models +- [Data preparation](./04-technical-guides/data-preparation.md): tokenization, data formats, mock data +- [Checkpoint management](./04-technical-guides/checkpoint-management.md): formats, save/load, distributed checkpointing +- [Multi-node networking](./04-technical-guides/multi-node-networking.md): InfiniBand, RoCE, AINIC configuration +- [Profiling and observability](./04-technical-guides/profiling-and-observability.md): Torch profiler, TraceLens, memory snapshots, projection, pp_vis +- [Logging and experiment tracking](./04-technical-guides/logging-and-experiment-tracking.md): TensorBoard, WandB, MLflow setup per backend +- [Fault tolerance and elastic training](./04-technical-guides/fault-tolerance-and-elastic-training.md): graceful exit, auto-resume, in-process restart, torchft +- [Determinism and reproducibility](./04-technical-guides/determinism-and-reproducibility.md): deterministic mode, seeds, trade-offs +- [Diffusion models](./04-technical-guides/diffusion-models/README.md): Flux diffusion architecture, data pipeline, and FP8 / MXFP4 training +- [Native SFT and LoRA](./04-technical-guides/native-sft-lora.md): Megatron-native SFT/LoRA runbook (BF16 / FP8 / FP4), no Megatron-Bridge dependency -Get help and find answers: +### [Operations](./05-operations/) -- **[FAQ](./faq.md)** - Frequently asked questions and troubleshooting -- **[Examples](../examples/README.md)** - Real-world training examples and templates -- **[Preflight Tool](../primus/tools/preflight/README.md)** - Cluster sanity checker to verify environment readiness +Production deployment and operational guidance. -## 🎯 Quick Navigation by Use Case +- [Deployment](./05-operations/deployment.md): container, Slurm, and Kubernetes deployment +- [Monitoring and logging](./05-operations/monitoring-logging.md): WandB, TensorBoard, MLflow, Primus logging +- [Troubleshooting](./05-operations/troubleshooting.md): common failures, diagnostics, and fixes +- [Security](./05-operations/security.md): secrets handling, container security, dependencies + +### [Developer guide](./06-developer-guide/) + +For contributors and maintainers. + +- [Architecture](./06-developer-guide/architecture.md): system design, runtime, backends, patch system +- [Contributing](./06-developer-guide/contributing.md): development setup, code style, PR process +- [Testing](./06-developer-guide/testing.md): test types, running tests, CI pipeline +- [Extending backends](./06-developer-guide/extending-backends.md): adding new training backends +- [Adding models](./06-developer-guide/adding-models.md): adding model configurations per backend +- [Model support matrix](./06-developer-guide/model-support-matrix.md): supported models per backend and GPU +- [CLI architecture](./06-developer-guide/cli-architecture.md): CLI internals: subcommand discovery, dispatch, and launch wrappers +- [Backend patch notes](./06-developer-guide/backend-patch-notes.md): Primus-specific backend arguments and the files they patch +- [Tooling](./06-developer-guide/tooling.md): auxiliary analysis, benchmarking, visualization, and diagnostics tools + +--- + +## Common use cases ### I want to... -- **Train a model locally** → [Quick Start](./quickstart.md) + [CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md) -- **Run distributed training on Slurm** → [Slurm & Container Usage](./slurm-container.md) -- **Configure my training run** → [Configuration Guide](./configuration.md) -- **Project performance to multi-node** → [Performance Projection](./projection.md) -- **Auto-tune my training config (parallelism + knobs)** → [Tuning Agent](./tuning_agent.md) -- **Benchmark performance** → [Benchmark Suite](./benchmark.md) -- **Understand the CLI design** → [CLI Architecture](./cli/CLI-ARCHITECTURE.md) -- **Troubleshoot issues** → [FAQ](./faq.md) +| Goal | Document | +|------|----------| +| Understand what Primus is | [Overview](./01-getting-started/overview.md) | +| Browse all Primus tools | [Primus tools](./02-user-guide/primus-tools.md) | +| Install Primus | [Installation](./01-getting-started/installation.md) | +| Run my first training | [Quickstart](./01-getting-started/quickstart.md) | +| Get an exact run command for my model/GPU | [Backend training recipes](./02-user-guide/training-recipes.md) | +| Write a training YAML configuration | [Configuration system](./02-user-guide/configuration-system.md) | +| Look up a Megatron parameter | [Megatron parameters](./03-configuration-reference/megatron-parameters.md) | +| Look up a TorchTitan parameter | [TorchTitan parameters](./03-configuration-reference/torchtitan-parameters.md) | +| Look up an environment variable | [Environment variables](./03-configuration-reference/environment-variables.md) | +| Understand parallelism strategies | [Parallelism strategies](./04-technical-guides/parallelism-strategies.md) | +| Configure parallelism for my model | [Parallelism configuration](./04-technical-guides/parallelism-configuration.md) | +| Tune training performance | [Performance tuning](./04-technical-guides/performance-tuning.md) | +| Train a Mixture-of-Experts model | [MoE training deep-dive](./04-technical-guides/moe-training.md) | +| Train a diffusion (Flux) model | [Diffusion models](./04-technical-guides/diffusion-models/README.md) | +| Fine-tune with native SFT / LoRA | [Native SFT and LoRA](./04-technical-guides/native-sft-lora.md) | +| Auto-tune my training configuration | [Tuning agent](./02-user-guide/tuning-agent.md) | +| Profile a training run | [Profiling and observability](./04-technical-guides/profiling-and-observability.md) | +| Track experiments (WandB/MLflow/TensorBoard) | [Logging and experiment tracking](./04-technical-guides/logging-and-experiment-tracking.md) | +| Survive node failures on long runs | [Fault tolerance and elastic training](./04-technical-guides/fault-tolerance-and-elastic-training.md) | +| Reproduce results bit-for-bit | [Determinism and reproducibility](./04-technical-guides/determinism-and-reproducibility.md) | +| Prepare training data | [Data preparation](./04-technical-guides/data-preparation.md) | +| Deploy to a Slurm cluster | [Deployment](./05-operations/deployment.md) | +| Debug a training failure | [Troubleshooting](./05-operations/troubleshooting.md) | +| Contribute to Primus | [Contributing](./06-developer-guide/contributing.md) | +| Understand the code architecture | [Architecture](./06-developer-guide/architecture.md) | +| Add a new training backend | [Extending backends](./06-developer-guide/extending-backends.md) | + +--- -## 🔗 External Resources +## External resources -- [Primus-Turbo](https://github.com/AMD-AGI/Primus-Turbo) - High-performance operators & modules -- [Primus-SaFE](https://github.com/AMD-AGI/Primus-SaFE) - Stability & platform layer -- [AMD ROCm Documentation](https://rocm.docs.amd.com/) -- [TorchTitan Documentation](https://github.com/pytorch/torchtitan) +- [Primus-Turbo](https://github.com/AMD-AGI/Primus-Turbo): high-performance operators and kernels +- [Primus-SaFE](https://github.com/AMD-AGI/Primus-SaFE): external stability/platform layer; this repository does not include a production integration guide +- [AMD ROCm documentation](https://rocm.docs.amd.com/) +- [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) +- [TorchTitan](https://github.com/pytorch/torchtitan) +- [MaxText](https://github.com/AI-Hypercomputer/maxtext) --- -**Need help?** Check the [FAQ](./faq.md) or open an issue on [GitHub](https://github.com/AMD-AGI/Primus/issues). +**Need help?** Open an issue on [GitHub](https://github.com/AMD-AGI/Primus/issues). diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 000000000..b2818f364 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,37 @@ +# Configuration file for the Sphinx documentation builder. +# +# This file only contains a selection of the most common options. For a full +# list see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +import re +from pathlib import Path + + +def _get_version(): + init = Path(__file__).parent.parent / "primus" / "__init__.py" + match = re.search(r'^__version__\s*=\s*["\']([^"\']+)["\']', init.read_text(), re.MULTILINE) + if not match: + raise ValueError("Could not find __version__ in primus/__init__.py") + return match.group(1) + + +# Project info +version = _get_version() +release = version +project = f"AMD Primus {version}" +author = "Advanced Micro Devices, Inc." +copyright = "Copyright (c) %Y Advanced Micro Devices, Inc. All rights reserved." + +# Theme-related configs +html_theme = "rocm_docs_theme" +html_theme_options = { + "flavor": "ai-ecosystem", + "link_main_doc": False, +} +html_title = project + +# Sphinx extension-related configs +extensions = ["rocm_docs"] +external_toc_path = "./sphinx/_toc.yml" +external_projects_current_project = "primus" diff --git a/docs/license.md b/docs/license.md new file mode 100644 index 000000000..1f8761f24 --- /dev/null +++ b/docs/license.md @@ -0,0 +1,4 @@ +# License + +```{include} ../LICENSE +``` diff --git a/docs/sphinx/_toc.yml.in b/docs/sphinx/_toc.yml.in new file mode 100644 index 000000000..455628414 --- /dev/null +++ b/docs/sphinx/_toc.yml.in @@ -0,0 +1,139 @@ +# Variables of the form ${} are substituted, currently the following +# list is supported: +# - ${branch} (or {branch}) the name of the current branch +# - ${url} (or {url}) github url of the current project +# - ${project:} base url of the documentation of +# based on intersphinx_mapping. +# These comments will also be removed. +defaults: + numbered: false +root: 01-getting-started/overview.md +subtrees: + - caption: Getting started + entries: + - file: 01-getting-started/quickstart.md + title: Quickstart + - file: 01-getting-started/installation.md + title: Installation and setup + - file: 01-getting-started/glossary.md + title: Glossary + + - caption: User guide + entries: + - file: 02-user-guide/pretraining.md + title: Pretraining workflows + - file: 02-user-guide/posttraining.md + title: Post-training workflows + - file: 02-user-guide/training-recipes.md + title: Backend training recipes + - file: 02-user-guide/cli-reference.md + title: CLI reference + - file: 02-user-guide/preflight.md + title: Preflight diagnostics + - file: 02-user-guide/configuration-system.md + title: Configuration system + - file: 02-user-guide/projection.md + title: Memory and performance projection + - file: 02-user-guide/benchmarking.md + title: Benchmark suite + - file: 02-user-guide/primus-tools.md + title: Primus tools + - file: 02-user-guide/tuning-agent.md + title: Tuning agent + + - caption: Configuration reference + entries: + - file: 03-configuration-reference/megatron-parameters.md + title: Megatron backend + - file: 03-configuration-reference/torchtitan-parameters.md + title: TorchTitan backend + - file: 03-configuration-reference/megatron-bridge-parameters.md + title: Megatron Bridge backend + - file: 03-configuration-reference/maxtext-parameters.md + title: MaxText backend + - file: 03-configuration-reference/environment-variables.md + title: Environment variables + + - caption: Technical guides + entries: + - file: 04-technical-guides/parallelism-strategies.md + title: Parallelism strategies + - file: 04-technical-guides/parallelism-configuration.md + title: Parallelism configuration + - file: 04-technical-guides/multi-node-networking.md + title: Multi-node networking + - file: 04-technical-guides/collective-operations.md + title: Collective operations (NCCL/RCCL) + - file: 04-technical-guides/data-preparation.md + title: Data preparation + - file: 04-technical-guides/checkpoint-management.md + title: Checkpoint management + - file: 04-technical-guides/moe-training.md + title: MoE training + - file: 04-technical-guides/performance-tuning.md + title: Performance tuning + - file: 04-technical-guides/profiling-and-observability.md + title: Profiling and observability + - file: 04-technical-guides/logging-and-experiment-tracking.md + title: Logging and experiment tracking + - file: 04-technical-guides/fault-tolerance-and-elastic-training.md + title: Fault tolerance and elastic training + - file: 04-technical-guides/determinism-and-reproducibility.md + title: Determinism and reproducibility + - file: 04-technical-guides/native-sft-lora.md + title: Native SFT LoRA + - file: 04-technical-guides/diffusion-models/architecture_overview.md + title: Diffusion model architecture + subtrees: + - entries: + - file: 04-technical-guides/diffusion-models/flux_architecture.md + title: Flux architecture + - file: 04-technical-guides/diffusion-models/fp8_training.md + title: FP8 training + - file: 04-technical-guides/diffusion-models/mxfp4_training.md + title: MXFP4 training + - file: 04-technical-guides/diffusion-models/data_preprocessing.md + title: Data preprocessing + - file: 04-technical-guides/diffusion-models/energon_integration.md + title: Energon integration + - file: 04-technical-guides/diffusion-models/adding_new_models.md + title: Adding new models + - file: 04-technical-guides/diffusion-models/api_reference.md + title: API reference + + - caption: Operations + entries: + - file: 05-operations/deployment.md + title: Deployment + - file: 05-operations/monitoring-logging.md + title: Monitoring and logging + - file: 05-operations/troubleshooting.md + title: Troubleshooting + - file: 05-operations/security.md + title: Security considerations + + - caption: Developer guide + entries: + - file: 06-developer-guide/architecture.md + title: Architecture overview + - file: 06-developer-guide/cli-architecture.md + title: CLI architecture + - file: 06-developer-guide/contributing.md + title: Contributing + - file: 06-developer-guide/adding-models.md + title: Adding model configurations + - file: 06-developer-guide/extending-backends.md + title: Extending backends + - file: 06-developer-guide/model-support-matrix.md + title: Model support matrix + - file: 06-developer-guide/testing.md + title: Testing + - file: 06-developer-guide/tooling.md + title: Tooling + - file: 06-developer-guide/backend-patch-notes.md + title: Backend patch notes + + - caption: About + entries: + - file: license.md + title: License diff --git a/docs/sphinx/requirements.in b/docs/sphinx/requirements.in new file mode 100644 index 000000000..7dfacf6a0 --- /dev/null +++ b/docs/sphinx/requirements.in @@ -0,0 +1 @@ +rocm-docs-core @ git+https://github.com/ROCm/rocm-docs-core.git@ai-ecosystem-theme diff --git a/docs/sphinx/requirements.txt b/docs/sphinx/requirements.txt new file mode 100644 index 000000000..2a0741d17 --- /dev/null +++ b/docs/sphinx/requirements.txt @@ -0,0 +1,279 @@ +# +# This file is autogenerated by pip-compile with Python 3.12 +# by the following command: +# +# pip-compile production_docs/sphinx/requirements.in +# +accessible-pygments==0.0.5 + # via pydata-sphinx-theme +alabaster==1.0.0 + # via sphinx +asttokens==3.0.2 + # via stack-data +attrs==26.1.0 + # via + # jsonschema + # jupyter-cache + # referencing +babel==2.18.0 + # via + # pydata-sphinx-theme + # sphinx +beautifulsoup4==4.15.0 + # via pydata-sphinx-theme +breathe==4.36.0 + # via rocm-docs-core +certifi==2026.6.17 + # via requests +cffi==2.1.0 + # via + # cryptography + # pynacl +charset-normalizer==3.4.9 + # via requests +click==8.4.2 + # via + # jupyter-cache + # sphinx-external-toc +comm==0.2.3 + # via ipykernel +cryptography==49.0.0 + # via pyjwt +debugpy==1.8.21 + # via ipykernel +decorator==5.3.1 + # via ipython +docutils==0.22.4 + # via + # myst-parser + # pydata-sphinx-theme + # sphinx +executing==2.2.1 + # via stack-data +fastjsonschema==2.21.2 + # via + # nbformat + # rocm-docs-core +gitdb==4.0.12 + # via gitpython +gitpython==3.1.51 + # via rocm-docs-core +greenlet==3.5.3 + # via sqlalchemy +idna==3.18 + # via requests +imagesize==2.0.0 + # via sphinx +importlib-metadata==9.0.0 + # via + # jupyter-cache + # myst-nb +ipykernel==7.3.0 + # via myst-nb +ipython==9.15.0 + # via + # ipykernel + # myst-nb +ipython-pygments-lexers==1.1.1 + # via ipython +jedi==0.20.0 + # via ipython +jinja2==3.1.6 + # via + # myst-parser + # sphinx +jsonschema==4.26.0 + # via nbformat +jsonschema-specifications==2025.9.1 + # via jsonschema +jupyter-cache==1.0.1 + # via myst-nb +jupyter-client==8.9.1 + # via + # ipykernel + # nbclient +jupyter-core==5.9.1 + # via + # ipykernel + # jupyter-client + # nbclient + # nbformat +markdown-it-py==4.2.0 + # via + # mdit-py-plugins + # myst-parser +markupsafe==3.0.3 + # via jinja2 +matplotlib-inline==0.2.2 + # via + # ipykernel + # ipython +mdit-py-plugins==0.6.1 + # via myst-parser +mdurl==0.1.2 + # via markdown-it-py +myst-nb==1.4.0 + # via rocm-docs-core +myst-parser==5.1.0 + # via myst-nb +nbclient==0.11.0 + # via + # jupyter-cache + # myst-nb +nbformat==5.10.4 + # via + # jupyter-cache + # myst-nb + # nbclient +nest-asyncio2==1.7.2 + # via ipykernel +packaging==26.2 + # via + # ipykernel + # pydata-sphinx-theme + # sphinx +parso==0.8.7 + # via jedi +pexpect==4.9.0 + # via ipython +platformdirs==4.10.0 + # via jupyter-core +prompt-toolkit==3.0.52 + # via ipython +psutil==7.2.2 + # via + # ipykernel + # ipython +ptyprocess==0.7.0 + # via pexpect +pure-eval==0.2.3 + # via stack-data +pycparser==3.0 + # via cffi +pydata-sphinx-theme==0.15.4 + # via + # rocm-docs-core + # sphinx-book-theme +pygithub==2.9.1 + # via rocm-docs-core +pygments==2.20.0 + # via + # accessible-pygments + # ipython + # ipython-pygments-lexers + # pydata-sphinx-theme + # sphinx +pyjwt[crypto]==2.13.0 + # via pygithub +pynacl==1.6.2 + # via pygithub +python-dateutil==2.9.0.post0 + # via jupyter-client +pyyaml==6.0.3 + # via + # jupyter-cache + # myst-nb + # myst-parser + # rocm-docs-core + # sphinx-external-toc +pyzmq==27.1.0 + # via + # ipykernel + # jupyter-client +referencing==0.37.0 + # via + # jsonschema + # jsonschema-specifications +requests==2.34.2 + # via + # pygithub + # sphinx +rocm-docs-core @ git+https://github.com/ROCm/rocm-docs-core.git@ai-ecosystem-theme + # via -r production_docs/sphinx/requirements.in +roman-numerals==4.1.0 + # via sphinx +rpds-py==2026.6.3 + # via + # jsonschema + # referencing +six==1.17.0 + # via python-dateutil +smmap==5.0.3 + # via gitdb +snowballstemmer==3.1.1 + # via sphinx +soupsieve==2.8.4 + # via beautifulsoup4 +sphinx==9.1.0 + # via + # breathe + # myst-nb + # myst-parser + # pydata-sphinx-theme + # rocm-docs-core + # sphinx-book-theme + # sphinx-copybutton + # sphinx-design + # sphinx-external-toc + # sphinx-multitoc-numbering + # sphinx-notfound-page +sphinx-book-theme==1.1.4 + # via rocm-docs-core +sphinx-copybutton==0.5.2 + # via rocm-docs-core +sphinx-design==0.7.0 + # via rocm-docs-core +sphinx-external-toc==1.1.0 + # via rocm-docs-core +sphinx-multitoc-numbering==0.1.3 + # via sphinx-external-toc +sphinx-notfound-page==1.1.0 + # via rocm-docs-core +sphinxcontrib-applehelp==2.0.0 + # via sphinx +sphinxcontrib-devhelp==2.0.0 + # via sphinx +sphinxcontrib-htmlhelp==2.1.0 + # via sphinx +sphinxcontrib-jsmath==1.0.1 + # via sphinx +sphinxcontrib-qthelp==2.0.0 + # via sphinx +sphinxcontrib-serializinghtml==2.0.0 + # via sphinx +sqlalchemy==2.0.51 + # via jupyter-cache +stack-data==0.6.3 + # via ipython +tabulate==0.10.0 + # via jupyter-cache +tornado==6.5.7 + # via + # ipykernel + # jupyter-client +traitlets==5.15.1 + # via + # ipykernel + # ipython + # jupyter-client + # jupyter-core + # matplotlib-inline + # nbclient + # nbformat +typing-extensions==4.16.0 + # via + # beautifulsoup4 + # jupyter-client + # myst-nb + # pydata-sphinx-theme + # pygithub + # referencing + # sqlalchemy +urllib3==2.7.0 + # via + # pygithub + # requests +wcwidth==0.8.2 + # via prompt-toolkit +zipp==4.1.0 + # via importlib-metadata diff --git a/docs_deprecated/README.md b/docs_deprecated/README.md new file mode 100644 index 000000000..d2659d274 --- /dev/null +++ b/docs_deprecated/README.md @@ -0,0 +1,82 @@ +# Primus Documentation + +Welcome to the Primus documentation! This guide will help you get started with training large-scale foundation models on AMD GPUs. + +> **Comprehensive Documentation**: For the complete production documentation set, see [`docs/`](../docs/README.md). It includes configuration references, parallelism guides, environment variable documentation, and more. + +## Documentation Structure + +### Getting Started + +Start here if you're new to Primus: + +- **[Quick Start Guide](./quickstart.md)** - Get up and running in 5 minutes +- **[CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md)** - Complete command-line reference +- **[CLI Architecture](../docs/06-developer-guide/cli-architecture.md)** - Design philosophy and deep dive + +### User Guides + +Guides for common workflows and features: + +- **[Configuration System](../docs/02-user-guide/configuration-system.md)** - YAML configuration, presets, overrides, and inheritance +- **[Deployment Guide](../docs/05-operations/deployment.md)** - Container, Slurm, and Kubernetes deployment + +### Technical References + +In-depth technical documentation: + +- **[Post-Training Guide](./posttraining.md)** - Fine-tuning with SFT and LoRA using Primus CLI +- **[Native SFT & LoRA Quick Start](../docs/04-technical-guides/native-sft-lora.md)** - Megatron-native SFT/LoRA launch guide (BF16/FP8/FP4), no Megatron-Bridge runtime dependency +- **[Performance Projection](./projection.md)** - Project training performance and memory to multi-node configurations +- **[Tuning Agent](../docs/02-user-guide/tuning-agent.md)** - LLM-driven search for an optimal training config — parallelism plus batching, schedule, memory, MoE-comm, and precision knobs (drives the projection tool as an oracle) +- **[Preflight](./preflight.md)** - Cluster diagnostics (host/GPU/network info + perf tests) +- **[Benchmark Suite](./benchmark.md)** - GEMM, RCCL, end-to-end benchmarks and profiling +- **[Supported Models](./backends/overview.md#supported-models)** - Supported LLM architectures and feature compatibility matrix +- **[Backend Patch Notes](./backends/overview.md)** - Primus-specific arguments for Megatron, TorchTitan, etc. +- **[Backend Extension Guide](./backends/extending-backends.md)** - How to add a new backend using the current adapter/trainer architecture + - **[Megatron Model Extension Guide](./backends/adding-megatron-models.md)** - How to add a new Megatron model config + - **[TorchTitan Model Extension Guide](./backends/adding-torchtitan-models.md)** - How to add a new TorchTitan model config +- **[Flux Diffusion Models](../docs/04-technical-guides/diffusion-models/README.md)** - Flux diffusion model architecture, training, and API reference +- **[FP8 Training Guide](../docs/04-technical-guides/diffusion-models/fp8_training.md)** - FP8 precision training on AMD MI300X/MI355X: configuration, benchmarks, and tuning + +### Production Documentation + +For comprehensive coverage, see the [Production Documentation](../docs/README.md): + +- **[Configuration References](../docs/03-configuration-reference/megatron-parameters.md)** - Per-backend YAML parameter documentation +- **[Environment Variables](../docs/03-configuration-reference/environment-variables.md)** - Complete environment variable reference +- **[Parallelism Strategies](../docs/04-technical-guides/parallelism-strategies.md)** - Distributed training parallelism explained +- **[Performance Tuning](../docs/04-technical-guides/performance-tuning.md)** - HipBLASLt, Primus-Turbo, FP8, MoE optimization +- **[Troubleshooting](../docs/05-operations/troubleshooting.md)** - Common issues and solutions +- **[Architecture](../docs/06-developer-guide/architecture.md)** - System design and code architecture + +### Help and Support + +- **[Troubleshooting Guide](../docs/05-operations/troubleshooting.md)** - Common issues and solutions +- **[Examples](../examples/README.md)** - Real-world training examples and templates +- **[Preflight Tool](../primus/tools/preflight/README.md)** - Cluster sanity checker to verify environment readiness + +## Quick Navigation by Use Case + +### I want to... + +- **Train a model locally** → [Quick Start](./quickstart.md) + [CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md) +- **Run distributed training on Slurm** → [Deployment Guide](../docs/05-operations/deployment.md) +- **Configure my training run** → [Configuration System](../docs/02-user-guide/configuration-system.md) +- **Look up YAML parameters** → [Configuration References](../docs/03-configuration-reference/megatron-parameters.md) +- **Project performance to multi-node** → [Performance Projection](./projection.md) +- **Auto-tune my training config (parallelism + knobs)** → [Tuning Agent](../docs/02-user-guide/tuning-agent.md) +- **Benchmark performance** → [Benchmark Suite](./benchmark.md) +- **Understand the CLI design** → [CLI Architecture](../docs/06-developer-guide/cli-architecture.md) +- **Troubleshoot issues** → [Troubleshooting](../docs/05-operations/troubleshooting.md) + +## External Resources + +- [Primus-Turbo](https://github.com/AMD-AGI/Primus-Turbo) - High-performance operators and modules +- [Primus-SaFE](https://github.com/AMD-AGI/Primus-SaFE) - Stability and platform layer +- [AMD ROCm Documentation](https://rocm.docs.amd.com/) +- [TorchTitan Documentation](https://github.com/pytorch/torchtitan) + +--- + +**Need help?** Check the [FAQ](./faq.md) or open an issue on [GitHub](https://github.com/AMD-AGI/Primus/issues). diff --git a/docs/backend-gap/dashboard-data/reports/megatron-upstream-main-2026-04-30.json b/docs_deprecated/backend-gap/dashboard-data/reports/megatron-upstream-main-2026-04-30.json similarity index 100% rename from docs/backend-gap/dashboard-data/reports/megatron-upstream-main-2026-04-30.json rename to docs_deprecated/backend-gap/dashboard-data/reports/megatron-upstream-main-2026-04-30.json diff --git a/docs/backend-gap/dashboard-data/reports/torchtitan-upstream-main-2026-04-21.json b/docs_deprecated/backend-gap/dashboard-data/reports/torchtitan-upstream-main-2026-04-21.json similarity index 100% rename from docs/backend-gap/dashboard-data/reports/torchtitan-upstream-main-2026-04-21.json rename to docs_deprecated/backend-gap/dashboard-data/reports/torchtitan-upstream-main-2026-04-21.json diff --git a/docs/backend-gap/reports/megatron/upstream-main/report.md b/docs_deprecated/backend-gap/reports/megatron/upstream-main/report.md similarity index 100% rename from docs/backend-gap/reports/megatron/upstream-main/report.md rename to docs_deprecated/backend-gap/reports/megatron/upstream-main/report.md diff --git a/docs/backend-gap/reports/megatron/upstream-main/summary.md b/docs_deprecated/backend-gap/reports/megatron/upstream-main/summary.md similarity index 100% rename from docs/backend-gap/reports/megatron/upstream-main/summary.md rename to docs_deprecated/backend-gap/reports/megatron/upstream-main/summary.md diff --git a/docs/backend-gap/reports/torchtitan/upstream-main/report.md b/docs_deprecated/backend-gap/reports/torchtitan/upstream-main/report.md similarity index 100% rename from docs/backend-gap/reports/torchtitan/upstream-main/report.md rename to docs_deprecated/backend-gap/reports/torchtitan/upstream-main/report.md diff --git a/docs/backend-gap/reports/torchtitan/upstream-main/summary.md b/docs_deprecated/backend-gap/reports/torchtitan/upstream-main/summary.md similarity index 100% rename from docs/backend-gap/reports/torchtitan/upstream-main/summary.md rename to docs_deprecated/backend-gap/reports/torchtitan/upstream-main/summary.md diff --git a/docs/backends/adding-megatron-models.md b/docs_deprecated/backends/adding-megatron-models.md similarity index 100% rename from docs/backends/adding-megatron-models.md rename to docs_deprecated/backends/adding-megatron-models.md diff --git a/docs/backends/adding-torchtitan-models.md b/docs_deprecated/backends/adding-torchtitan-models.md similarity index 100% rename from docs/backends/adding-torchtitan-models.md rename to docs_deprecated/backends/adding-torchtitan-models.md diff --git a/docs/backends/extending-backends.md b/docs_deprecated/backends/extending-backends.md similarity index 100% rename from docs/backends/extending-backends.md rename to docs_deprecated/backends/extending-backends.md diff --git a/docs/backends/maxtext/patch-notes.md b/docs_deprecated/backends/maxtext/patch-notes.md similarity index 100% rename from docs/backends/maxtext/patch-notes.md rename to docs_deprecated/backends/maxtext/patch-notes.md diff --git a/docs/backends/megatron/patch-notes.md b/docs_deprecated/backends/megatron/patch-notes.md similarity index 100% rename from docs/backends/megatron/patch-notes.md rename to docs_deprecated/backends/megatron/patch-notes.md diff --git a/docs/backends/overview.md b/docs_deprecated/backends/overview.md similarity index 100% rename from docs/backends/overview.md rename to docs_deprecated/backends/overview.md diff --git a/docs/backends/torchtitan/patch-notes.md b/docs_deprecated/backends/torchtitan/patch-notes.md similarity index 100% rename from docs/backends/torchtitan/patch-notes.md rename to docs_deprecated/backends/torchtitan/patch-notes.md diff --git a/docs/benchmark.md b/docs_deprecated/benchmark.md similarity index 100% rename from docs/benchmark.md rename to docs_deprecated/benchmark.md diff --git a/docs/cli/PRIMUS-CLI-GUIDE.md b/docs_deprecated/cli/PRIMUS-CLI-GUIDE.md similarity index 99% rename from docs/cli/PRIMUS-CLI-GUIDE.md rename to docs_deprecated/cli/PRIMUS-CLI-GUIDE.md index c48915e6f..c884a7e89 100644 --- a/docs/cli/PRIMUS-CLI-GUIDE.md +++ b/docs_deprecated/cli/PRIMUS-CLI-GUIDE.md @@ -982,7 +982,7 @@ export PRIMUS_LOG_LEVEL=DEBUG ## Reference Resources ### Related Documentation -- [CLI Architecture](./CLI-ARCHITECTURE.md) - Primus CLI architecture deep dive +- [CLI Architecture](../../docs/06-developer-guide/cli-architecture.md) - Primus CLI architecture deep dive - [Main Documentation](../README.md) - Complete Primus documentation index - [.primus.yaml](../../runner/.primus.yaml) - Default configuration example diff --git a/docs/cli/README.md b/docs_deprecated/cli/README.md similarity index 95% rename from docs/cli/README.md rename to docs_deprecated/cli/README.md index bd5061462..a7e322a46 100644 --- a/docs/cli/README.md +++ b/docs_deprecated/cli/README.md @@ -13,7 +13,7 @@ The Primus CLI provides a unified command-line interface for training, benchmark - Configuration files and options - Best practices and troubleshooting -- **[Architecture Deep Dive](./CLI-ARCHITECTURE.md)** +- **[Architecture Deep Dive](../../docs/06-developer-guide/cli-architecture.md)** - Design philosophy and principles - Three-layer architecture explained - Plugin system and extensibility @@ -70,7 +70,7 @@ primus-cli direct -- data diffusion-ingest \ ## 📖 Learn More - For detailed usage instructions, see the [User Guide](./PRIMUS-CLI-GUIDE.md) -- For architecture and design details, see [Architecture Deep Dive](./CLI-ARCHITECTURE.md) +- For architecture and design details, see [Architecture Deep Dive](../../docs/06-developer-guide/cli-architecture.md) - For the main Primus documentation, see [Primus README](../../README.md) ## 🔗 Related Documentation diff --git a/docs/install-on-host.md b/docs_deprecated/install-on-host.md similarity index 100% rename from docs/install-on-host.md rename to docs_deprecated/install-on-host.md diff --git a/docs/posttraining.md b/docs_deprecated/posttraining.md similarity index 100% rename from docs/posttraining.md rename to docs_deprecated/posttraining.md diff --git a/docs/preflight.md b/docs_deprecated/preflight.md similarity index 100% rename from docs/preflight.md rename to docs_deprecated/preflight.md diff --git a/docs/projection.md b/docs_deprecated/projection.md similarity index 99% rename from docs/projection.md rename to docs_deprecated/projection.md index 02596c6b3..3588c5efb 100644 --- a/docs/projection.md +++ b/docs_deprecated/projection.md @@ -626,7 +626,7 @@ Splitting the residual into these two terms is more robust than the older single With `--memory-mode both`, a side-by-side table shows the simulate vs. benchmark per-component deltas. A small delta (within ~10–20%) means the analytical model and the residual term are well-calibrated; a large positive delta (simulate ≫ benchmark) usually means simulate is over-counting unsharded components, while a large negative delta means the bench captured overhead the analytical model under-estimates. -> The benchmark-based memory projection is what the [Tuning Agent](tuning_agent.md) uses for OOM-accurate feasibility filtering when a GPU is available, so its `tokens/s/GPU` rankings never include configs that would OOM on the real cluster. +> The benchmark-based memory projection is what the [Tuning Agent](../docs/02-user-guide/tuning-agent.md) uses for OOM-accurate feasibility filtering when a GPU is available, so its `tokens/s/GPU` rankings never include configs that would OOM on the real cluster. --- diff --git a/docs/quickstart.md b/docs_deprecated/quickstart.md similarity index 96% rename from docs/quickstart.md rename to docs_deprecated/quickstart.md index 4d6dd1e6c..6aa149624 100644 --- a/docs/quickstart.md +++ b/docs_deprecated/quickstart.md @@ -87,7 +87,7 @@ primus-cli [options] [mode-args] -- [command] **Learn More:** - [CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md) - Complete reference -- [CLI Architecture](./cli/CLI-ARCHITECTURE.md) - Design deep dive +- [CLI Architecture](../docs/06-developer-guide/cli-architecture.md) - Design deep dive - [Configuration Guide](./configuration.md) - YAML configuration - [Examples](../examples/README.md) - Real-world templates diff --git a/docs/tech_blogs/primus_cli_unified_entry_rocm.md b/docs_deprecated/tech_blogs/primus_cli_unified_entry_rocm.md similarity index 100% rename from docs/tech_blogs/primus_cli_unified_entry_rocm.md rename to docs_deprecated/tech_blogs/primus_cli_unified_entry_rocm.md diff --git a/docs/tech_blogs/primus_pipeline/imgs/actual-perf.png b/docs_deprecated/tech_blogs/primus_pipeline/imgs/actual-perf.png similarity index 100% rename from docs/tech_blogs/primus_pipeline/imgs/actual-perf.png rename to docs_deprecated/tech_blogs/primus_pipeline/imgs/actual-perf.png diff --git a/docs/tech_blogs/primus_pipeline/imgs/llama2-7B-perf.png b/docs_deprecated/tech_blogs/primus_pipeline/imgs/llama2-7B-perf.png similarity index 100% rename from docs/tech_blogs/primus_pipeline/imgs/llama2-7B-perf.png rename to docs_deprecated/tech_blogs/primus_pipeline/imgs/llama2-7B-perf.png diff --git a/docs/tech_blogs/primus_pipeline/imgs/qwen-235B-perf.png b/docs_deprecated/tech_blogs/primus_pipeline/imgs/qwen-235B-perf.png similarity index 100% rename from docs/tech_blogs/primus_pipeline/imgs/qwen-235B-perf.png rename to docs_deprecated/tech_blogs/primus_pipeline/imgs/qwen-235B-perf.png diff --git a/docs/tech_blogs/primus_pipeline/imgs/simulation.png b/docs_deprecated/tech_blogs/primus_pipeline/imgs/simulation.png similarity index 100% rename from docs/tech_blogs/primus_pipeline/imgs/simulation.png rename to docs_deprecated/tech_blogs/primus_pipeline/imgs/simulation.png diff --git a/docs/tech_blogs/primus_pipeline/imgs/simulator_shell.png b/docs_deprecated/tech_blogs/primus_pipeline/imgs/simulator_shell.png similarity index 100% rename from docs/tech_blogs/primus_pipeline/imgs/simulator_shell.png rename to docs_deprecated/tech_blogs/primus_pipeline/imgs/simulator_shell.png diff --git a/docs/tech_blogs/primus_pipeline/primus_pipeline.md b/docs_deprecated/tech_blogs/primus_pipeline/primus_pipeline.md similarity index 100% rename from docs/tech_blogs/primus_pipeline/primus_pipeline.md rename to docs_deprecated/tech_blogs/primus_pipeline/primus_pipeline.md diff --git a/docs/tech_blogs/projection/projection.md b/docs_deprecated/tech_blogs/projection/projection.md similarity index 100% rename from docs/tech_blogs/projection/projection.md rename to docs_deprecated/tech_blogs/projection/projection.md diff --git a/docs/weekly_reports/2026-W17-primus-weekly.md b/docs_deprecated/weekly_reports/2026-W17-primus-weekly.md similarity index 100% rename from docs/weekly_reports/2026-W17-primus-weekly.md rename to docs_deprecated/weekly_reports/2026-W17-primus-weekly.md diff --git a/docs/weekly_reports/2026-W18-primus-weekly.md b/docs_deprecated/weekly_reports/2026-W18-primus-weekly.md similarity index 100% rename from docs/weekly_reports/2026-W18-primus-weekly.md rename to docs_deprecated/weekly_reports/2026-W18-primus-weekly.md diff --git a/docs/weekly_reports/2026-W19-primus-weekly.md b/docs_deprecated/weekly_reports/2026-W19-primus-weekly.md similarity index 100% rename from docs/weekly_reports/2026-W19-primus-weekly.md rename to docs_deprecated/weekly_reports/2026-W19-primus-weekly.md diff --git a/docs/weekly_reports/dashboard-data/reports/2026-W17.json b/docs_deprecated/weekly_reports/dashboard-data/reports/2026-W17.json similarity index 100% rename from docs/weekly_reports/dashboard-data/reports/2026-W17.json rename to docs_deprecated/weekly_reports/dashboard-data/reports/2026-W17.json diff --git a/docs/weekly_reports/dashboard-data/reports/2026-W18.json b/docs_deprecated/weekly_reports/dashboard-data/reports/2026-W18.json similarity index 100% rename from docs/weekly_reports/dashboard-data/reports/2026-W18.json rename to docs_deprecated/weekly_reports/dashboard-data/reports/2026-W18.json diff --git a/docs/weekly_reports/dashboard-data/reports/2026-W19.json b/docs_deprecated/weekly_reports/dashboard-data/reports/2026-W19.json similarity index 100% rename from docs/weekly_reports/dashboard-data/reports/2026-W19.json rename to docs_deprecated/weekly_reports/dashboard-data/reports/2026-W19.json diff --git a/examples/README.md b/examples/README.md index 7067527c8..1a2aec904 100644 --- a/examples/README.md +++ b/examples/README.md @@ -212,8 +212,8 @@ The following models are supported out of the box via provided configuration fil - **Flux** - Flow-based diffusion model for text-to-image generation - Training guide: [examples/megatron/diffusion/README.md](megatron/diffusion/README.md) (Flux 535M and 12B) - - Architecture & developer docs: [docs/backends/megatron/diffusion/README.md](../docs/backends/megatron/diffusion/README.md) - - FP8 training: [docs/backends/megatron/diffusion/fp8_training.md](../docs/backends/megatron/diffusion/fp8_training.md) + - Architecture & developer docs: [docs/04-technical-guides/diffusion-models/README.md](../docs/04-technical-guides/diffusion-models/README.md) + - FP8 training: [docs/04-technical-guides/diffusion-models/fp8_training.md](../docs/04-technical-guides/diffusion-models/fp8_training.md) --- diff --git a/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml b/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml index 83eda2b58..fa4aa8da7 100644 --- a/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml +++ b/examples/megatron/configs/MI300X/diffusion/flux_535m_pretrain_fp8.yaml @@ -171,7 +171,7 @@ modules: # - Training time: Minutes # # Validation Checklist: -# [ ] Setup FP8 environment (see docs/backends/megatron/diffusion/fp8_training.md) +# [ ] Setup FP8 environment (see docs/04-technical-guides/diffusion-models/fp8_training.md) # [ ] Verify TE FP8 support is available # [ ] Run this config to validate FP8 training # [ ] Check logs for NaN/Inf (should be none) @@ -196,4 +196,4 @@ modules: # - If slow: Verify ROCm FP8 tensor cores are being used # - If unstable: Try fp8_wgrad: false in model config # -# For more information: See docs/backends/megatron/diffusion/fp8_training.md +# For more information: See docs/04-technical-guides/diffusion-models/fp8_training.md diff --git a/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml b/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml index 5db547f21..1ec0e6433 100644 --- a/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml +++ b/examples/megatron/configs/MI355X/diffusion/flux_535m_pretrain_fp8.yaml @@ -171,7 +171,7 @@ modules: # - Training time: Minutes # # Validation Checklist: -# [ ] Setup FP8 environment (see docs/backends/megatron/diffusion/fp8_training.md) +# [ ] Setup FP8 environment (see docs/04-technical-guides/diffusion-models/fp8_training.md) # [ ] Verify TE FP8 support is available # [ ] Run this config to validate FP8 training # [ ] Check logs for NaN/Inf (should be none) @@ -196,4 +196,4 @@ modules: # - If slow: Verify ROCm FP8 tensor cores are being used # - If unstable: Try fp8_wgrad: false in model config # -# For more information: See docs/backends/megatron/diffusion/fp8_training.md +# For more information: See docs/04-technical-guides/diffusion-models/fp8_training.md diff --git a/examples/megatron/diffusion/README.md b/examples/megatron/diffusion/README.md index d03f7babd..bfc6ada27 100644 --- a/examples/megatron/diffusion/README.md +++ b/examples/megatron/diffusion/README.md @@ -4,10 +4,10 @@ Training examples for Flux diffusion models with Primus-Megatron on AMD GPUs. ## Related Documentation -- **Architecture & Developer Guide:** [docs/backends/megatron/diffusion/README.md](../../../docs/backends/megatron/diffusion/README.md) -- **API Reference:** [docs/backends/megatron/diffusion/api_reference.md](../../../docs/backends/megatron/diffusion/api_reference.md) -- **FP8 Training Guide:** [docs/backends/megatron/diffusion/fp8_training.md](../../../docs/backends/megatron/diffusion/fp8_training.md) -- **MXFP4 Training Guide:** [docs/backends/megatron/diffusion/mxfp4_training.md](../../../docs/backends/megatron/diffusion/mxfp4_training.md) +- **Architecture & Developer Guide:** [docs/04-technical-guides/diffusion-models/README.md](../../../docs/04-technical-guides/diffusion-models/README.md) +- **API Reference:** [docs/04-technical-guides/diffusion-models/api_reference.md](../../../docs/04-technical-guides/diffusion-models/api_reference.md) +- **FP8 Training Guide:** [docs/04-technical-guides/diffusion-models/fp8_training.md](../../../docs/04-technical-guides/diffusion-models/fp8_training.md) +- **MXFP4 Training Guide:** [docs/04-technical-guides/diffusion-models/mxfp4_training.md](../../../docs/04-technical-guides/diffusion-models/mxfp4_training.md) - **Dataset Preparation:** [primus/configs/data/megatron/diffusion/README.md](../../../primus/configs/data/megatron/diffusion/README.md) - **Tests:** [tests/unit_tests/backends/megatron/diffusion/](../../../tests/unit_tests/backends/megatron/diffusion/) @@ -178,7 +178,7 @@ bash examples/run_slurm_pretrain.sh | Flux 12B | BF16 | ~40-50GB | 1 | 1.0x | | Flux 12B | FP8 | ~20-25GB | 2 | 1.5-2x | -For configuration details, tuning recipes, benchmarks, and troubleshooting, see the [FP8 Training Guide](../../../docs/backends/megatron/diffusion/fp8_training.md). +For configuration details, tuning recipes, benchmarks, and troubleshooting, see the [FP8 Training Guide](../../../docs/04-technical-guides/diffusion-models/fp8_training.md). --- @@ -198,7 +198,7 @@ AITER_LOG_TUNED_CONFIG=1 \ bash examples/run_pretrain.sh ``` -For configuration knobs, backend-selector semantics, tuned-GEMM verification, and troubleshooting, see the [MXFP4 Training Guide](../../../docs/backends/megatron/diffusion/mxfp4_training.md). +For configuration knobs, backend-selector semantics, tuned-GEMM verification, and troubleshooting, see the [MXFP4 Training Guide](../../../docs/04-technical-guides/diffusion-models/mxfp4_training.md). --- diff --git a/primus/agents/tuning_agent/README.md b/primus/agents/tuning_agent/README.md index e4df0a730..ed7543714 100644 --- a/primus/agents/tuning_agent/README.md +++ b/primus/agents/tuning_agent/README.md @@ -6,7 +6,7 @@ provides two estimates — **memory** and **performance** — each of which runs **benchmark-anchored by default** (measure what fits on a sub-node run, scale the rest analytically) with a fully analytical **no-GPU `simulate`** fallback. -See [`docs/tuning_agent.md`](../../../docs/tuning_agent.md) — the full +See [`docs/02-user-guide/tuning-agent.md`](../../../docs/02-user-guide/tuning-agent.md) — the full user/operator guide (modes, configuration, troubleshooting, worked example) plus the design write-up, paper-ready problem statement, and the list of deferred future features (cluster-spec retrieval, persistent memory cache, diff --git a/primus/configs/data/megatron/diffusion/README.md b/primus/configs/data/megatron/diffusion/README.md index a674ba69c..827f1420b 100644 --- a/primus/configs/data/megatron/diffusion/README.md +++ b/primus/configs/data/megatron/diffusion/README.md @@ -46,7 +46,7 @@ When using `primus-cli direct --` commands inside containers, understanding Dock ### Default Docker Setup -From [`tools/docker/start_container.sh`](../../../../tools/docker/start_container.sh): +From [`tools/docker/start_container.sh`](../../../../../tools/docker/start_container.sh): ```bash DATA_PATH=${DATA_PATH:-"${PRIMUS_PATH}/data"} # Default: ./data relative to repo # Mounted as: -v "${DATA_PATH}:${DATA_PATH}" @@ -580,15 +580,15 @@ model: ## Additional Resources ### Documentation -- **Training Guide**: [`examples/megatron/diffusion/README.md`](../../../../examples/megatron/diffusion/README.md) -- **Energon Integration**: [`docs/backends/megatron/diffusion/energon_integration.md`](../../../../docs/backends/megatron/diffusion/energon_integration.md) -- **FP8 Training Guide**: [`docs/backends/megatron/diffusion/fp8_training.md`](../../../../docs/backends/megatron/diffusion/fp8_training.md) +- **Training Guide**: [`examples/megatron/diffusion/README.md`](../../../../../examples/megatron/diffusion/README.md) +- **Energon Integration**: [`docs/04-technical-guides/diffusion-models/energon_integration.md`](../../../../../docs/04-technical-guides/diffusion-models/energon_integration.md) +- **FP8 Training Guide**: [`docs/04-technical-guides/diffusion-models/fp8_training.md`](../../../../../docs/04-technical-guides/diffusion-models/fp8_training.md) ### Related Configs -- **Encoder Configs**: [`primus/configs/models/megatron/diffusion/encoders.yaml`](../../models/megatron/diffusion/encoders.yaml) -- **Model Architecture**: [`primus/configs/models/megatron/diffusion/`](../../models/megatron/diffusion/) -- **Training Configs (MI300X)**: [`examples/megatron/configs/MI300X/diffusion/`](../../../../examples/megatron/configs/MI300X/diffusion/) -- **Training Configs (MI355X)**: [`examples/megatron/configs/MI355X/diffusion/`](../../../../examples/megatron/configs/MI355X/diffusion/) +- **Encoder Configs**: [`primus/configs/models/megatron/diffusion/encoders.yaml`](../../../models/megatron/diffusion/encoders.yaml) +- **Model Architecture**: [`primus/configs/models/megatron/diffusion/`](../../../models/megatron/diffusion/) +- **Training Configs (MI300X)**: [`examples/megatron/configs/MI300X/diffusion/`](../../../../../examples/megatron/configs/MI300X/diffusion/) +- **Training Configs (MI355X)**: [`examples/megatron/configs/MI355X/diffusion/`](../../../../../examples/megatron/configs/MI355X/diffusion/) --- diff --git a/primus/configs/data/megatron/diffusion/templates/metadataset.yaml b/primus/configs/data/megatron/diffusion/templates/metadataset.yaml index b6563b368..8e41e70b3 100644 --- a/primus/configs/data/megatron/diffusion/templates/metadataset.yaml +++ b/primus/configs/data/megatron/diffusion/templates/metadataset.yaml @@ -77,5 +77,5 @@ max_samples_per_sequence: 100 # For more information, see: # - primus/configs/data/megatron/diffusion/README.md # - examples/megatron/diffusion/README.md -# - docs/backends/megatron/diffusion/energon_integration.md +# - docs/04-technical-guides/diffusion-models/energon_integration.md # diff --git a/primus/configs/models/megatron/diffusion/encoders.yaml b/primus/configs/models/megatron/diffusion/encoders.yaml index 43d771fdb..6d5f8632d 100644 --- a/primus/configs/models/megatron/diffusion/encoders.yaml +++ b/primus/configs/models/megatron/diffusion/encoders.yaml @@ -92,7 +92,7 @@ clip: # - Encoders above ARE loaded during training # - Images/text are encoded on-the-fly (slower but saves disk space) # -# See: docs/backends/megatron/diffusion/data_preprocessing.md +# See: docs/04-technical-guides/diffusion-models/data_preprocessing.md # ============================================================================ # ============================================================================ diff --git a/primus/configs/modules/megatron/primus_megatron_module.yaml b/primus/configs/modules/megatron/primus_megatron_module.yaml index dce09dbbe..838ca07a4 100644 --- a/primus/configs/modules/megatron/primus_megatron_module.yaml +++ b/primus/configs/modules/megatron/primus_megatron_module.yaml @@ -92,7 +92,7 @@ pp_warmup: false # set to true to decrease iter-1 time when using pp dump_pp_data: false # recompute -recompute_layer_ids: null #int list，id srange from 0 to (num_layers_per_pp_stage - 1) +recompute_layer_ids: null # int list; global layer ids, range from 0 to (num_layers - 1) # dataloader dataloader_mp_context: null # "forkserver" | "spawn" | "fork" | null diff --git a/tests/README.md b/tests/README.md index e69de29bb..bad41361f 100644 --- a/tests/README.md +++ b/tests/README.md @@ -0,0 +1,49 @@ +# Primus Tests + +This directory contains the test suite for Primus. + +## Test Structure + +``` +tests/ +├── runner/ # Shell integration tests for primus-cli +│ ├── run_all_tests.sh # Master test runner +│ ├── lib/ # Library function tests +│ ├── helpers/ # Hook and environment tests +│ └── test_primus_cli*.sh # CLI mode tests +├── unit_tests/ # Python unit tests (pytest) +│ ├── agents/ # Tuning-agent tests +│ ├── backends/ # Backend-specific tests (megatron, torchtitan, maxtext, ...) +│ ├── ci/ # CI helper tests +│ ├── cli/ # CLI tests +│ ├── core/ # Core library tests (config/, patches/, backend/, launcher/, projection/, pipeline_parallel/, runtime/, trainer/, utils/) +│ ├── megatron/ # Megatron-specific unit tests +│ ├── modules/ # Module/trainer tests +│ └── tools/ # Tooling tests +├── trainer/ # Integration tests (require GPU) +│ ├── test_megatron_trainer.py +│ ├── test_torchtitan_trainer.py +│ └── test_maxtext_trainer.py +├── scripts/ # CI unit/integration launch scripts and UT patches +├── utils.py # Shared test utilities +├── conftest.py # Shared pytest fixtures +└── run_unit_tests.py # Python test orchestrator (walks tests/) +``` + +> **Note:** `config/` and `patches/` live under `unit_tests/core/` (i.e. `tests/unit_tests/core/config/` and `tests/unit_tests/core/patches/`), not directly under `unit_tests/`. + +## Running Tests + +```bash +# Shell integration tests +bash ./tests/runner/run_all_tests.sh + +# Python unit tests +pytest tests/unit_tests/ --maxfail=1 -s + +# All tests via orchestrator +python ./tests/run_unit_tests.py # Torch backends +python ./tests/run_unit_tests.py --jax # JAX/MaxText backend +``` + +For comprehensive testing documentation, see the [Testing Guide](../docs/06-developer-guide/testing.md). diff --git a/tools/README.md b/tools/README.md index e69de29bb..495effa93 100644 --- a/tools/README.md +++ b/tools/README.md @@ -0,0 +1,24 @@ +# Primus Tools + +Auxiliary tools for analysis, benchmarking, visualization, and diagnostics. + +## Available Tools + +| Tool | Directory | Description | +|------|-----------|-------------| +| **IRLens** | `tools/IRLens/` | Parses XLA HLO text dumps; prints execution skeleton with control flow, communication vs compute ops | +| **model_stats** | `tools/model_stats/` | Generates charts from the model config registry under `primus/configs/models` | +| **Pipeline Visualization** | `tools/visualization/pp_vis/` | Visualizes pipeline parallelism schedules from dumped data or PP simulator JSON via a local web UI | +| **Auto Benchmark** | `tools/auto_benchmark/` | Interactive benchmark menu for Megatron/TorchTitan on MI300X/MI355X with metrics collection | +| **Daily Report** | `tools/daily/` | Benchmark summary CSV generation used by CI workflows | +| **Docker Helpers** | `tools/docker/` | Container startup and proxy scripts | +| **Profile Trace** | `tools/profile_trace/` | Trace file merging utility | + +## Per-Tool Documentation + +Each tool has its own README with usage instructions: + +- [IRLens README](./IRLens/README.md) +- [model_stats README](./model_stats/README.md) +- [Pipeline Visualization README](./visualization/pp_vis/README.md) +- [Auto Benchmark README](./auto_benchmark/Primus_Auto_Benchmark_README.md) From e2f24878f9eeb3067c1dd8bfcff9d6b0d875658e Mon Sep 17 00:00:00 2001 From: vidushi8 Date: Tue, 14 Jul 2026 20:08:20 -0700 Subject: [PATCH 027/127] Add MLperf examples for llama3.1 8b and gpt oss 20B (#854) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds MLPerf-style pretraining examples for **Llama 3.1 8B** and **GPT-OSS 20B** on AMD MI355X, covering low-precision recipes (MXFP4 for Llama, FP8 for GPT-OSS), plus the supporting Megatron/TE patches needed to run them end-to-end. ## What's included ### Llama 3.1 8B (MI355X, MXFP4/FP4) - `examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml` — FP4 pretrain config - `examples/mlperf/llama3.1_8b/config_MI355X_1x8x1.sh` — single-node 8-GPU launch config (TP1/PP1/EP1) - `examples/mlperf/llama3.1_8b/run_and_time.sh` — MLPerf run+timing wrapper - `examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv` — tuned GEMM configs - `examples/mlperf/llama3.1_8b/README.md` ### GPT-OSS 20B (MI355X, FP8) - `examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml` — FP8 pretrain config - `examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh` — launch config (gbs 32) - `examples/mlperf/gpt_oss_20b/run_and_time.sh` - `examples/mlperf/gpt_oss_20b/tune_gemm_results.txt` - `examples/mlperf/gpt_oss_20b/README.md` ### Framework patches - `primus/backends/megatron/te_patches/fused_bias_swiglu_patches.py` — fused bias-SwiGLU support - `primus/backends/megatron/patches/validation_data_sampling_patches.py` — validation data sampling - `primus/backends/megatron/training/evaluator.py` — evaluator updates - `primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py` — trainer tweaks ## Notes - Paths in the launch/run scripts are relative to `${PRIMUS_PATH}`. - Passes pre-commit (`shellcheck`, YAML lint, `thirdparty-lock`). - Stats: 14 files, +1166 / −5. ## Test plan - [x] `bash examples/mlperf/llama3.1_8b/run_and_time.sh` completes on 1×MI355X (8 GPU), MXFP4 - [x] `bash examples/mlperf/gpt_oss_20b/run_and_time.sh` completes on 1×MI355X (8 GPU), FP8 - [x] Loss/throughput sanity-checked - [train.mlperfpretrain.llama3.1-8b.log](https://github.com/user-attachments/files/29991139/train.mlperfpretrain.llama3.1-8b.log) ; [train.mlperfpretrain.gptoss.log](https://github.com/user-attachments/files/29991140/train.mlperfpretrain.gptoss.log) - [x] Docker image - `unifiedtrainingdockers.azurecr.io/utd/ci:primus_the_rock_ci_2e8745e_20260709` - [x] For GPT OSS 20 FP8 use atomic FP32 `NVTE_CK_IS_V3_ATOMIC_FP32=1` to avoid NaN and `use_turbo_grouped_gemm: false` ## Results (1×MI355X, 8 GPU, GBS=32) | Model | Precision | Target eval loss | Samples to target | Final eval loss | Time to train | Throughput | |---|---|---|---|---|---|---| | Llama 3.1 8B | MXFP4 | ≤ 3.30 | 196,608 | 3.2807 | 102.42 min (6,145.34 s) | 31.99 samples/s | | GPT-OSS 20B | FP8 | ≤ 3.34 | 233,472 | 3.3297 | 153.33 min (9,199.68 s) | 25.38 samples/s | Metrics from MLPerf MLLOG `run_start`→`run_stop` (both runs converged, `status: success`). "Samples to target" = `samples_count` at the first eval where `eval_accuracy ≤ MLLOG_TARGET_EVAL_LOSS`. Throughput = `overall_throughput` (samples/s). Target loss set via `MLLOG_TARGET_EVAL_LOSS` (3.3 / 3.34). --------- Co-authored-by: Vidushi Goyal --- examples/mlperf/gpt_oss_20b/README.md | 64 +++++ .../config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh | 182 ++++++++++++++ .../gpt_oss_20B-FP8-mlperf-pretrain.yaml | 222 ++++++++++++++++++ examples/mlperf/gpt_oss_20b/run_and_time.sh | 60 +++++ .../mlperf/gpt_oss_20b/tune_gemm_results.txt | 25 ++ examples/mlperf/llama3.1_8b/README.md | 53 +++++ .../mlperf/llama3.1_8b/a4w4_tuned_gemms.csv | 12 + .../mlperf/llama3.1_8b/config_MI355X_1x8x1.sh | 115 +++++++++ .../MI355X/llama3.1_8B-pretrain-FP4.yaml | 113 +++++++++ examples/mlperf/llama3.1_8b/run_and_time.sh | 60 +++++ .../mlperf/mlperf_pretrain_trainer.py | 8 +- .../te_patches/fused_bias_swiglu_patches.py | 60 +++++ .../validation_data_sampling_patches.py | 168 +++++++++++++ .../backends/megatron/training/evaluator.py | 29 ++- 14 files changed, 1166 insertions(+), 5 deletions(-) create mode 100644 examples/mlperf/gpt_oss_20b/README.md create mode 100644 examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh create mode 100644 examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml create mode 100755 examples/mlperf/gpt_oss_20b/run_and_time.sh create mode 100644 examples/mlperf/gpt_oss_20b/tune_gemm_results.txt create mode 100644 examples/mlperf/llama3.1_8b/README.md create mode 100644 examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv create mode 100755 examples/mlperf/llama3.1_8b/config_MI355X_1x8x1.sh create mode 100644 examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml create mode 100755 examples/mlperf/llama3.1_8b/run_and_time.sh create mode 100644 primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py create mode 100644 primus/backends/megatron/patches/validation_data_sampling_patches.py diff --git a/examples/mlperf/gpt_oss_20b/README.md b/examples/mlperf/gpt_oss_20b/README.md new file mode 100644 index 000000000..8f29f31a0 --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/README.md @@ -0,0 +1,64 @@ +# GPT-OSS-20B Pretraining Benchmark + +GPT-OSS 20B (Mixture of Experts) + + +## Setup + +### Start Docker Image + +```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 +``` + + +### 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: + +```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 +``` + +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` + +The data directory is approximately **80 GB** and model directory is approximately **30 GB**. + +### How to run + +```bash +export HF_TOKEN= +source config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh +bash run_and_time.sh +``` +## Notes + +- `log_interval: 99999999` suppresses regular Primus logs 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 new file mode 100644 index 000000000..b3c64add3 --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh @@ -0,0 +1,182 @@ +#!/bin/bash +# ============================================================================= +# MLPerf GPT-OSS-20B Configuration for MI355X (1 node, 8 GPUs) +# ============================================================================= + +# ----------------------------------------------------------------------------- +# System Configuration +# ----------------------------------------------------------------------------- +export DGXSYSTEM=MI355X_1x8x1 +export GPUS_PER_NODE=8 +export NNODES=1 +export NODE_RANK=0 +export MASTER_ADDR=localhost +export MASTER_PORT=29501 + +# ----------------------------------------------------------------------------- +# Paths +# ----------------------------------------------------------------------------- +export PRIMUS_PATH=/workspace/Primus +export PYTHONPATH="${PRIMUS_PATH}:${PRIMUS_PATH}/third_party/Megatron-LM:${PYTHONPATH}" +export EXP=${PRIMUS_PATH}/examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml +export DATA_PATH=/data + +# ----------------------------------------------------------------------------- +# Training Hyperparameters +# ----------------------------------------------------------------------------- +export PRIMUS_MICRO_BATCH_SIZE=4 +export PRIMUS_GLOBAL_BATCH_SIZE=32 +export EVAL_ITERS=$((1024 / PRIMUS_GLOBAL_BATCH_SIZE)) # MLPerf closed: eval_iters * GBS = 1024 eval samples +export PRIMUS_LR=8.0e-4 +export PRIMUS_MIN_LR=8.0e-5 # Set to 10% of max LR +export PRIMUS_TRAIN_ITERS=1200000 +export PRIMUS_LR_WARMUP_ITERS=128 +export PRIMUS_LR_DECAY_ITERS=$((PRIMUS_TRAIN_ITERS-PRIMUS_LR_WARMUP_ITERS)) # 1200000 - 128 = 1199872 +# export SEED=30279 + +# Evaluation frequency (sample-based, adjusts automatically with GBS) +export EVAL_SAMPLES_INTERVAL=12288 # Evaluate every 12,288 samples +export PRIMUS_EVAL_INTERVAL=$((EVAL_SAMPLES_INTERVAL / PRIMUS_GLOBAL_BATCH_SIZE)) # Auto-computed + +# ----------------------------------------------------------------------------- +# Parallelism +# ----------------------------------------------------------------------------- +export PRIMUS_TP=1 +export PRIMUS_PP=1 +export PRIMUS_EP=1 + +# ----------------------------------------------------------------------------- +# Primus Configuration +# ----------------------------------------------------------------------------- +export PRIMUS_TURBO_GROUPED_GEMM_BACKEND=TRITON +export PRIMUS_GRAD_REDUCE_IN_BF16=true +export USE_TURBO_RMS_NORM=true + +# ----------------------------------------------------------------------------- +# ROCm / System Runtime +# ----------------------------------------------------------------------------- +export GPU_MAX_HW_QUEUES=2 +export HIP_FORCE_DEV_KERNARG=1 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export HSA_KERNARG_POOL_SIZE=12582912 +export TORCH_NCCL_HIGH_PRIORITY=1 +export ENABLE_NUMA_BINDING=1 +export PYTORCH_ALLOC_CONF=expandable_segments:True +export HSA_NO_SCRATCH_RECLAIM=1 +export HSA_ENABLE_SDMA=1 +export HSA_ENABLE_INTERRUPT=0 +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export OMP_NUM_THREADS=1 +export PYTHONWARNINGS=ignore +export TOKENIZERS_PARALLELISM=false + +# ----------------------------------------------------------------------------- +# RCCL / NCCL Tuning +# ----------------------------------------------------------------------------- +export NCCL_MIN_P2P_NCHANNELS=32 +export NCCL_MIN_CTAS=32 +export NCCL_NCHANNELS_PER_NET_PEER=32 +export NCCL_NVLS_ENABLE=0 +export NCCL_CHECKS_DISABLE=1 + +# ----------------------------------------------------------------------------- +# hipBLASLt +# ----------------------------------------------------------------------------- +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 + +# ----------------------------------------------------------------------------- +# NVTE — FP8 & Cast Transpose +# ----------------------------------------------------------------------------- +export NVTE_ROCM_ENABLE_MXFP8=0 +export NVTE_USE_CAST_TRANSPOSE_TRITON=0 +export NVTE_USE_OPTIMIZED_HIPIFIED_CAST_TRANSPOSE=1 + +# ----------------------------------------------------------------------------- +# NVTE — FMHA / CK Backend +# ----------------------------------------------------------------------------- +export NVTE_FLASH_ATTN=0 # Disable FlashAttention so FusedAttention (CK/ASM) is used +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. + +# ----------------------------------------------------------------------------- +# NVTE — Debug +# ----------------------------------------------------------------------------- +export NVTE_DEBUG=0 +export NVTE_DEBUG_LEVEL=0 +export NVTE_LOG_FUSED_ATTN_CONFIG=0 +export NVTE_LOG_CK_CONFIG=0 +export CK_FUSED_ATTN_LOG_CONFIG=0 +# export NVTE_FMHA_DEBUG=1 # keep commented; debug-only knob + +# ----------------------------------------------------------------------------- +# MLPerf Logging +# ----------------------------------------------------------------------------- +export LOG_INTERVAL=999999 +export MLLOG_TRAIN_LOSS_LOG_FREQ=0 +export MLLOG_TARGET_EVAL_LOSS=3.34 +export MLLOG_OUTPUT_FILE=/results/mlperf_logging.out +export MLLOG_SAVE_TO_FILE=0 +export MLLOG_SUBMISSION_BENCHMARK=gpt_oss_20b +export MLLOG_SUBMISSION_DIVISION=closed +export MLLOG_SUBMISSION_ORG=AMD +export MLLOG_SUBMISSION_PLATFORM=MI355X + +export MLLOG_TENSOR_PARALLELISM=1 +export MLLOG_PIPELINE_PARALLELISM=1 +export MLLOG_CONTEXT_PARALLELISM=1 +export MLLOG_EXPERT_PARALLELISM=1 +export MLLOG_MICRO_BATCH_SIZE=4 +MLLOG_CONFIG_FILENAME=$(basename "${BASH_SOURCE[0]}") +export MLLOG_CONFIG_FILENAME +export MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR='fp8' + +# ----------------------------------------------------------------------------- +# Synthetic Warmup (kernel pre-compilation) +# ----------------------------------------------------------------------------- +export SYNTH_WARMUP_STEPS=3 + +# ----------------------------------------------------------------------------- +# MoE Token Dispatcher +# ----------------------------------------------------------------------------- +# 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. +export MOE_SKIP_IDENTITY_SORT=1 + +# ----------------------------------------------------------------------------- +# DDP Parameter All-Gather (SDMA) +# ----------------------------------------------------------------------------- +export ENABLE_SDMA_ALLGATHER=1 +# 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 + +# ----------------------------------------------------------------------------- +# Run-log verbosity +# ----------------------------------------------------------------------------- +# MLPerf run-log verbosity. Default 0 keeps only :::MLLOG + ``run_and_time.sh`` +# banners on stdout. Set to 1 to restore the full framework output (Primus +# loguru banners / Megatron / TE / aiter / Gloo / torchrun / hipify / ...) +# when debugging. See src/_log_suppression.py for the full strategy. +export MLPERF_VERBOSE_LOGS=${MLPERF_VERBOSE_LOGS:-0} + +# fused rms and swiglu no cat +export PRIMUS_FUSED_RESIDUAL_NORM=1 +export PRIMUS_MOE_SWIGLU_NOCAT=1 +export MLLOG_BLOCK_TPUT_LOG=0 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 new file mode 100644 index 000000000..0e714040b --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml @@ -0,0 +1,222 @@ +work_group: ${TEAM:amd} +user_name: ${USER:root} +exp_name: ${EXP_NAME:gpt_oss_20b} +workspace: ./output + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: ${PRIMUS_MODEL:gpt_oss_20B}.yaml + overrides: + + # Activate the migrated MLPerf pretrain trainer (mllog + MLPerf hooks). + stage: mlperf_pretrain + + # tokenizer + tokenizer_type: Llama3Tokenizer + tokenizer_model: ${MODEL:meta-llama/Llama-3.1-8B} + + # model + num_layers: 24 + hidden_size: 2880 + ffn_hidden_size: 2880 + num_attention_heads: 64 + num_query_groups: 8 # Group Query Attention (GQA) - matches HF num_key_value_heads + num_experts: 32 + activation_func: swiglu # SiLU activation (matches HF hidden_act: "silu") + + # rotary + position_embedding_type: rope + rotary_base: 150000 + + # mixed-precision + attention_softmax_in_fp32: false + grad_reduce_in_bf16: ${PRIMUS_GRAD_REDUCE_IN_BF16:true} + + # log + wandb_project: "Primus_GPT_OSS_20B" + stderr_sink_level: ERROR + log_interval: 999999 + + # debug + # moe_router_force_load_balancing: true + # log_avg_skip_iterations: 2 + # log_avg_reset_interval: 50 + + # profile + profile: ${PRIMUS_PROFILE:false} + 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] + + # enable fp8 training + fp8: e4m3 + fp8_recipe: tensorwise + clip_grad: 1.0 # Gradient clipping (already default, but explicit) + check_for_nan_in_loss_and_grad: false + + # hyper parameters + train_iters: ${PRIMUS_TRAIN_ITERS:1200000} + micro_batch_size: ${PRIMUS_MICRO_BATCH_SIZE:2} + global_batch_size: ${PRIMUS_GLOBAL_BATCH_SIZE:16} + seq_length: ${PRIMUS_SEQ_LENGTH:8192} + max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:131072} + seed: ${SEED:1234} # Random seed for reproducibility + lr: ${PRIMUS_LR:8.0e-4} # Reduced from 8e-4 for FP8 stability + min_lr: ${PRIMUS_MIN_LR:8.0e-5} # Set to 10% of max LR + lr_warmup_iters: ${PRIMUS_LR_WARMUP_ITERS:128} + lr_decay_iters: ${PRIMUS_LR_DECAY_ITERS:1199872} + lr_decay_style: cosine + weight_decay: 0.1 + optimizer: adam + use_distributed_optimizer: true # use distributed optimizer + adam_beta1: 0.9 + adam_beta2: 0.95 + adam_eps: 1.0e-5 + eod_mask_loss: true + init_method_std: 0.008 + norm_epsilon: 1.0e-6 + layernorm_epsilon: 1.0e-05 # RMSNorm epsilon (matches HF rms_norm_eps) + + # Dropout (disabled for training) + hidden_dropout: 0.0 + attention_dropout: 0.0 + + # parallel + tensor_model_parallel_size: ${PRIMUS_TP:1} + pipeline_model_parallel_size: ${PRIMUS_PP:1} + expert_model_parallel_size: ${PRIMUS_EP:8} + overlap_grad_reduce: true + overlap_param_gather: true + ddp_num_buckets: 8 + ddp_average_in_collective: true + + # data + mock_data: false + num_workers: ${PRIMUS_NUM_WORKERS:0} + train_data_path: "10 /data/c4-train.en_6_text_document" + valid_data_path: "/data/c4-validation-91205-samples.en_text_document" + test_data_path: "/data/c4-validation-91205-samples.en_text_document" + # Avoid copying a dense (B, 1, S, S) CPU attention mask every step. + # TE receives causal/sliding-window metadata from attn_mask_type + window_size. + # Use numeric 0/1 because Primus env expansion only type-casts numbers. + create_attention_mask_in_dataloader: ${PRIMUS_CREATE_ATTENTION_MASK_IN_DATALOADER:0} + + # fusion + moe_permute_fusion: true + gradient_accumulation_fusion: true + moe_use_legacy_grouped_gemm: false # Sync-Free MoE stage 2 or 3 require PrimusTurboGroupedMLP, please set `moe_use_legacy_grouped_gemm=True + moe_use_fused_router_with_aux_score: true + multi_latent_attention: false # Flag config.ENABLE_EXPERIMENTAL not enabled + apply_rope_fusion: true + + + # sliding window attention (GPT-OSS-20B model definition; matches HF sliding_window: 128) + # use_turbo_attention is false so non-turbo attention (which supports sliding window) is used. + # Pattern: alternating sliding_attention (1) and full_attention (0) for 24 layers + # window_size must be a tuple (left_window, right_window) for Transformer Engine + # For causal attention: left = past tokens, right = 0 (no future tokens) + # HF sliding_window: 128 means 128 past tokens, so use (128, 0) + window_size: [128, 0] # Left window: 128 past tokens, Right: 0 (causal) + window_attn_skip_freq: [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0] + + # MoE settings + moe_apply_probs_on_input: false + moe_aux_loss_coeff: 0.0 #0.9 + moe_deepep_num_sms: 20 + moe_enable_deepep: false + moe_expert_capacity_factor: null + moe_extended_tp: false + moe_ffn_hidden_size: 2880 + moe_flex_dispatcher_backend: deepep + moe_grouped_gemm: true + moe_hybridep_num_sms: 16 + moe_input_jitter_eps: null + moe_latent_size: null + moe_layer_freq: 1 + moe_layer_recompute: false + moe_pad_expert_input_to_capacity: false + moe_per_layer_logging: false + moe_router_bias_update_rate: 0.001 + moe_router_dtype: fp32 # DeepEP only supports float32 probs + moe_router_enable_expert_bias: false + moe_router_force_load_balancing: false + moe_router_fusion: true + moe_router_group_topk: null + moe_router_load_balancing_type: none + moe_router_num_groups: null + moe_router_padding_for_fp8: false + moe_router_padding_for_quantization: false + moe_router_pre_softmax: false + moe_router_score_function: softmax + moe_router_topk: 4 + moe_router_topk_limited_devices: null + moe_router_topk_scaling_factor: null + moe_shared_expert_gate: false + moe_shared_expert_intermediate_size: null + moe_shared_expert_overlap: false + moe_token_dispatcher_type: alltoall + moe_token_drop_policy: probs + moe_token_dropping: false + moe_z_loss_coeff: null + + # ckpt + finetune: false + auto_continue_train: false + load: null + no_load_optim: null + no_load_rng: null + save: null + save_interval: 100000 + no_save_optim: null + no_save_rng: null + disable_last_saving: true + exit_on_missing_checkpoint: false + ckpt_format: torch + eval_iters: ${EVAL_ITERS:64} # eval_samples = eval_iters * GBS = 1024; set EVAL_ITERS in config shell (1024/GBS). + eval_interval: ${PRIMUS_EVAL_INTERVAL:768} + + # Turbo + enable_primus_turbo: true + use_turbo_attention: false + use_turbo_grouped_gemm: false + use_turbo_rms_norm: ${USE_TURBO_RMS_NORM: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 + + # 64 or 80 for ep8, 32 for ep16-64 is best practice + turbo_deepep_num_cu: 64 + turbo_deepep_use_comm_stream: false + + # sync-free moe support stage 0-3, 0 means not use sync-free moe + # stage 3 is completely no gpu-cpu sync in MoE, but cost more memory + # stage 2 is recommended for better performance + turbo_sync_free_moe_stage: 0 + + # Cross entropy flags + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true + + # tensorboard logging, set 'disable_tensorboard: false' to enable tensorboard logging + disable_tensorboard: true + tensorboard_dir: /workspace/code/tensorboard + tensorboard_log_interval: 1 + tensorboard_queue_size: 1000 + log_timers_to_tensorboard: true + log_batch_size_to_tensorboard: true + log_learning_rate_to_tensorboard: true + log_validation_ppl_to_tensorboard: true + log_memory_to_tensorboard: true + log_world_size_to_tensorboard: true + log_loss_scale_to_tensorboard: true diff --git a/examples/mlperf/gpt_oss_20b/run_and_time.sh b/examples/mlperf/gpt_oss_20b/run_and_time.sh new file mode 100755 index 000000000..fc07e6fe1 --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/run_and_time.sh @@ -0,0 +1,60 @@ +#!/bin/bash + +set -e + +# Create results directory +mkdir -p /results + +cd "${PRIMUS_PATH}/examples/mlperf/gpt_oss_20b" + +# 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. +if [[ -n "${SLURM_NNODES:-}" && "${SLURM_NNODES}" -gt 1 ]]; then + NNODES="${SLURM_NNODES}" + NODE_RANK="${SLURM_NODEID:-0}" +fi + +echo "============================================" +echo "MLPerf GPT-OSS-20B Training" +echo "============================================" +echo "Config: ${EXP}" +echo "Data: ${DATA_PATH}" +echo "GPUs: ${GPUS_PER_NODE}" +echo "Nodes: ${NNODES}" +echo "Rank: ${NODE_RANK}" +echo "Master: ${MASTER_ADDR}:${MASTER_PORT}" +echo "============================================" + +# Start timing +start=$(date +%s) +start_fmt=$(date +%Y-%m-%d\ %r) +echo "STARTING TIMING RUN AT $start_fmt" + +# Launch through Primus CLI and keep the real exit code even though output is +# piped through tee. +set +e +"${PRIMUS_PATH}/primus-cli" direct -- \ + train pretrain \ + --config "${EXP}" \ + 2>&1 | tee train.mlperfpretrain.exp.log +ret_code=${PIPESTATUS[0]} +set -e + +# End timing +end=$(date +%s) +end_fmt=$(date +%Y-%m-%d\ %r) +echo "ENDING TIMING RUN AT $end_fmt" + +# Report result +result=$(( end - start )) +result_name="GPT_OSS_20B" +echo "RESULT,$result_name,,$result,AMD,$start_fmt" + +if [[ $ret_code != 0 ]]; then + echo "Training failed with exit code: $ret_code" + exit "$ret_code" +fi + +exit 0 diff --git a/examples/mlperf/gpt_oss_20b/tune_gemm_results.txt b/examples/mlperf/gpt_oss_20b/tune_gemm_results.txt new file mode 100644 index 000000000..9d987b63d --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/tune_gemm_results.txt @@ -0,0 +1,25 @@ +Git Version: de5c1aebb6-dirty + 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,151492,4511.27,39.8689,303082,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,81491.9,4800.79,74.1153,311532,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,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.69535e+06,1086.96,286.825,306398,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,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.59496e+06,788.341,297.921,306251,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,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.42444e+06,1094.09,318.875,306237,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,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.39063e+06,1024.47,404.231,306248,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,64447.9,3796.71,93.716,311699,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,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.81411e+06,1006.88,343.401,306396,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,1,1,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.96197e+06,832.489,326.259,306396,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.32172e+06,474.576,18315.1,300658,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.45958e+06,524.074,16585.3,302286,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.31028e+06,544.947,18475.1,300155,gfx950:sramecc+:xnack-,256 diff --git a/examples/mlperf/llama3.1_8b/README.md b/examples/mlperf/llama3.1_8b/README.md new file mode 100644 index 000000000..4a9289941 --- /dev/null +++ b/examples/mlperf/llama3.1_8b/README.md @@ -0,0 +1,53 @@ +# LLama3.1 8B MLPerf Pretraining + +MLPerf-compliant LLama3.1 8B pretraining using Primus + +## Setup + +### Start Docker Image + +```bash +export MLPERF_PAT= +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 +``` + + +### Configuration + +- **Model**: LLama3.1 8B (4096 hidden, 32 layers, 32 attention heads) +- **Training**: 1.2M iterations, GBS=32, MBS=2, LR=8e-4 +- **Precision**: MXFP4 +- **Data**: C4 dataset (tokenized) + +## Key Files + +- `configs/MI355X/llama3.1_8B-pretrain-FP4.yaml` - Model and training config + - Update `train_data_path` and `train_data_path` to your local downloaded location +- `config_MI355X_1x8x1.sh` - System config and env vars + - Update `PRIMUS_PATH` to clone Primus Repo + - Update `EXP`to `/examples/mlperf/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml` +- `run_and_time.sh` - Run script + +### Data + +Download preprocessed C4 dataset: + +```bash +mkdir -p /data/mlperf_llama31_8b +cd /data/mlperf_llama31_8b +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 +``` + +### How to run + +```bash +export HF_TOKEN= +source config_MI355X_1x8x1.sh +bash run_and_time.sh +``` +## Notes + +- `log_interval: 99999999` suppresses regular Primus logs diff --git a/examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv b/examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv new file mode 100644 index 000000000..5cf68d18b --- /dev/null +++ b/examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv @@ -0,0 +1,12 @@ +cu_num,M,N,K,kernelId,splitK,kernelName,errRatio +256,4096,4096,16384,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,6144,4096,16384,50,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_192x256E,0.0 +256,16384,4096,4096,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,16384,4096,6144,45,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_128x512E,0.0 +256,16384,4096,14336,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,16384,4096,28672,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,28672,4096,16384,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,16384,6144,4096,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,4096,14336,16384,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,16384,14336,4096,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,16384,28672,4096,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 diff --git a/examples/mlperf/llama3.1_8b/config_MI355X_1x8x1.sh b/examples/mlperf/llama3.1_8b/config_MI355X_1x8x1.sh new file mode 100755 index 000000000..f56e999d1 --- /dev/null +++ b/examples/mlperf/llama3.1_8b/config_MI355X_1x8x1.sh @@ -0,0 +1,115 @@ +#!/bin/bash +# MLPerf LLama3.1 8B Configuration for MI355X (1x8x1) + +export DGXSYSTEM=MI355X_1x8x1 +export GPUS_PER_NODE=8 +export NNODES=1 +export NODE_RANK=0 +export MASTER_ADDR=localhost +export MASTER_PORT=29502 + +export PRIMUS_PATH=/workspace/Primus +export PRIMUS_MLPERF=1 +export PYTHONPATH="${PRIMUS_PATH}:${PRIMUS_PATH}/third_party/Megatron-LM:${PYTHONPATH}" +export EXP=${PRIMUS_PATH}/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml +export DATA_PATH=/data + +export PRIMUS_MICRO_BATCH_SIZE=2 +export PRIMUS_GLOBAL_BATCH_SIZE=32 +export PRIMUS_LR=8e-4 +export PRIMUS_MIN_LR=8e-5 +export PRIMUS_TRAIN_ITERS=1200000 +export PRIMUS_LR_WARMUP_ITERS=64 +export EVAL_SAMPLES_INTERVAL=12288 +export PRIMUS_EVAL_INTERVAL=$((EVAL_SAMPLES_INTERVAL / PRIMUS_GLOBAL_BATCH_SIZE)) # Auto-computed + +export HSA_ENABLE_INTERRUPT=0 +export HSA_TOOLS_LIB=/opt/rocm/lib/libroctracer64.so +export PRIMUS_APPLY_ROPE_FUSION=True +export PRIMUS_FP8_RECIPE=hybrid + +export HSA_NO_SCRATCH_RECLAIM=1 +export HSA_ENABLE_SDMA=1 +export GPU_MAX_HW_QUEUES=2 +export CUDA_DEVICE_MAX_CONNECTIONS=1 + +export NVTE_FUSED_ATTN=1 +export NVTE_FUSED_ATTN_CK=1 +export NVTE_FUSED_ATTN_AOTRITON=1 +export NVTE_CK_USES_FWD_V3=1 +export NVTE_CK_USES_BWD_V3=1 +export NVTE_CK_IS_V3_ATOMIC_FP32=0 +export NVTE_USE_AITER_ROPE=1 +export NVTE_MXFP4_USE_HADAMARD=1 +export NVTE_FLASH_ATTN=0 +export NVTE_FUSED_ATTN=1 +export NVTE_USE_CAST_TRANSPOSE_TRITON=0 +export NVTE_ASYNC_AMAX_REDUCTION=1 +export NVTE_DP_AMAX_REDUCE_INTERVAL=0 +export NVTE_USE_RMSNORM_TRITON=0 +export NVTE_LOG_CK_CONFIG=0 +export NVTE_LOG_FUSED_ATTN_CONFIG=0 +export USE_TE_SWIGLU=1 + +export ENABLE_TRANSPOSE_CACHE=1 +export CK_FUSED_ATTN_LOG_CONFIG=0 +export CHECK_FOR_NAN_IN_GRAD=0 + +export TOKENIZERS_PARALLELISM=false +export NCCL_CHECKS_DISABLE=1 +export TORCH_NCCL_HIGH_PRIORITY=1 + +export AITER_CONFIG_GEMM_A4W4=${PRIMUS_PATH}/examples/mlperf/llama3.1_8b/a4w4_tuned_gemms.csv +export AITER_LOG_TUNED_CONFIG=0 +export LOG_AITER_GEMMS=0 +export AITER_LOG_LEVEL=ERROR +export AITER_LOG_MORE=0 + +export ENABLE_MLLOG=1 +export MLLOG_OUTPUT_FILE=/results/mlperf_logging.out +export MLLOG_TRAIN_LOSS_LOG_FREQ=0 +export MLLOG_TARGET_EVAL_LOSS=3.3 +export TARGET_EVAL_LOSS=3.3 +export MLLOG_SUBMISSION_BENCHMARK=llama31_8b +export MLLOG_SUBMISSION_DIVISION=closed +export MLLOG_SUBMISSION_ORG=AMD +export MLLOG_SUBMISSION_PLATFORM=MI355X +export MLLOG_TENSOR_PARALLELISM=1 +export MLLOG_PIPELINE_PARALLELISM=1 +export MLLOG_CONTEXT_PARALLELISM=1 +export MLLOG_EXPERT_PARALLELISM=1 +export MLLOG_MICRO_BATCH_SIZE=2 +MLLOG_CONFIG_FILENAME=$(basename "${BASH_SOURCE[0]}") +export MLLOG_CONFIG_FILENAME +export MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR='mxfp4' + +export FP4=true +export FP4_RECIPE=mxfp4 + +export USE_HIPBLASLT=1 +export TORCH_BLAS_PREFER_HIPBLASLT=1 + +export SYNTH_WARMUP_STEPS=20 +export SYNTH_WARMUP_VALID_STEPS=10 +export WARMUP_RECIPE=fp8_hybrid +export SYNTH_WARMUP_EMPTY_CACHE=1 +export MLPERF_VERBOSE_LOGS=${MLPERF_VERBOSE_LOGS:-0} + +export NCCL_MIN_P2P_NCHANNELS=32 +export NCCL_MIN_CTAS=32 +export NCCL_NCHANNELS_PER_NET_PEER=32 +export NCCL_NVLS_ENABLE=0 +export NCCL_P2P_LEVEL=5 +export NCCL_SINGLE_RING_THRESHOLD=0 +export NCCL_BUFFSIZE=2097152 + +export TP_COMM_OVERLAP=False +export MC_TP_OVERLAP_AG=False +export MC_TP_OVERLAP_RS=False +export MC_TP_OVERLAP_RS_DGRAD=False + +export HIP_FORCE_DEV_KERNARG=1 +export HSA_ENABLE_SDMA_OPTIMIZATIONS=1 +export PYTORCH_ROC_ALLOC_CONF=expandable_segments:True +export HIP_API_BLOCKING=0 +export PYTORCH_NO_CUDA_MEMORY_CACHING=0 diff --git a/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml b/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml new file mode 100644 index 000000000..cdc2efe8e --- /dev/null +++ b/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml @@ -0,0 +1,113 @@ +work_group: ${TEAM:amd} +user_name: ${USER:root} +exp_name: ${EXP_NAME:llama3.1_8B-pretrain-mlperf} +workspace: ./output + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + # model to run + model: llama3.1_8B.yaml + overrides: + # --- Logging Config --- + stage: mlperf_pretrain + disable_wandb: true + disable_tensorboard: true + stderr_sink_level: ERROR + log_interval: 999999 + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + eval_iters: ${PRIMUS_EVAL_ITERS:32} # 32 * GBS = 1024 eval samples + eval_interval: ${PRIMUS_EVAL_INTERVAL:10} # 10 * GBS = 320 eval samples perform evaluation. + + # --- Training Config --- + train_iters: ${PRIMUS_TRAIN_ITERS:200} + micro_batch_size: ${PRIMUS_MICRO_BATCH_SIZE:2} # grad_acc = global_batch_size / (micro_batch_size * num_gpus) = 32 / (2 * 8) = 2 + global_batch_size: ${PRIMUS_GLOBAL_BATCH_SIZE:32} + seq_length: 8192 + max_position_embeddings: 8192 + seed: ${SEED:1234} + lr: ${PRIMUS_LR:0.0008} # 8e-4 + min_lr: ${PRIMUS_MIN_LR:0.00008} # 10% of lr + lr_warmup_iters: ${PRIMUS_LR_WARMUP_ITERS:64} + lr_decay_iters: 1199936 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: false + init_method_std: 0.02 + norm_epsilon: 1.0e-5 + adam_eps: 1.0e-5 + check_for_nan_in_loss_and_grad: false # default true + check_for_spiky_loss: false # default false, but setting it here explicitly + check_for_large_grads: false # default false, but setting it here explicitly + + # --- Model Parallel Config --- + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + expert_model_parallel_size: 1 + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: true + masked_softmax_fusion: true + ddp_average_in_collective: true + fused_single_qkv_rope: true + + # --- Data Config --- + mock_data: false + train_data_path: "/data/c4-train.en_6_text_document" + valid_data_path: "/data/c4-validation-91205-samples.en_text_document" + test_data_path: null + mmap_bin_files: true + data_cache_path: null + + # --- Profiling Config --- + profile: false + use_pytorch_profiler: true + profile_ranks: [0] # Only profile rank 0 to save disk space + profile_step_start: 8 # Start after warmup (step 8) + profile_step_end: 13 # Profile 5 iterations (8-12) + disable_profiler_activity_cpu: false # GPU kernels only (smaller files) + torch_profiler_record_shapes: false # Disable for smaller traces + torch_profiler_with_stack: false # Disable for smaller traces + torch_profiler_use_gzip: true # Compress output + + # --- Checkpointing Config --- + finetune: false + auto_continue_train: false + load: null + no_load_optim: null + no_load_rng: null + save: null + save_interval: 2000000 + no_save_optim: null + no_save_rng: null + disable_last_saving: true + ckpt_format: torch + + # --- FSDP Config --- + use_torch_fsdp2: false + use_distributed_optimizer: true # this is needed for fsdp2 + + # Cross entropy flags + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true + + # --- Mixed Precision Config --- + fp4: e2m1 + fp4_recipe: mxfp4 + accumulate_allreduce_grads_in_fp32: false + grad_reduce_in_bf16: true + attention_softmax_in_fp32: false + + + # --- Primus Turbo Config --- + enable_primus_turbo: false + use_turbo_attention: false + use_turbo_parallel_linear: false # can't use together with delayed recipe + use_turbo_grouped_mlp: false + moe_use_fused_router_with_aux_score: false + enable_turbo_attention_float8 : false diff --git a/examples/mlperf/llama3.1_8b/run_and_time.sh b/examples/mlperf/llama3.1_8b/run_and_time.sh new file mode 100755 index 000000000..d9c11f583 --- /dev/null +++ b/examples/mlperf/llama3.1_8b/run_and_time.sh @@ -0,0 +1,60 @@ +#!/bin/bash + +set -e + +# Create results directory +mkdir -p /results + +cd "${PRIMUS_PATH}/examples/mlperf/llama3.1_8b" + +# 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. +if [[ -n "${SLURM_NNODES:-}" && "${SLURM_NNODES}" -gt 1 ]]; then + NNODES="${SLURM_NNODES}" + NODE_RANK="${SLURM_NODEID:-0}" +fi + +echo "============================================" +echo "MLPerf LLama3.1 8B Training" +echo "============================================" +echo "Config: ${EXP}" +echo "Data: ${DATA_PATH}" +echo "GPUs: ${GPUS_PER_NODE}" +echo "Nodes: ${NNODES}" +echo "Rank: ${NODE_RANK}" +echo "Master: ${MASTER_ADDR}:${MASTER_PORT}" +echo "============================================" + +# Start timing +start=$(date +%s) +start_fmt=$(date +%Y-%m-%d\ %r) +echo "STARTING TIMING RUN AT $start_fmt" + +# Launch through Primus CLI and keep the real exit code even though output is +# piped through tee. +set +e +"${PRIMUS_PATH}/primus-cli" direct -- \ + train pretrain \ + --config "${EXP}" \ + 2>&1 | tee train.mlperfpretrain.exp.log +ret_code=${PIPESTATUS[0]} +set -e + +# End timing +end=$(date +%s) +end_fmt=$(date +%Y-%m-%d\ %r) +echo "ENDING TIMING RUN AT $end_fmt" + +# Report result +result=$(( end - start )) +result_name="LLAMA3.1_8B" +echo "RESULT,$result_name,,$result,AMD,$start_fmt" + +if [[ $ret_code != 0 ]]; then + echo "Training failed with exit code: $ret_code" + exit "$ret_code" +fi + +exit 0 diff --git a/primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py b/primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py index 360013b2b..7d6194f01 100644 --- a/primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py +++ b/primus/backends/megatron/mlperf/mlperf_pretrain_trainer.py @@ -28,6 +28,7 @@ import os import time +from typing import Any from primus.backends.megatron.megatron_pretrain_trainer import MegatronPretrainTrainer from primus.backends.megatron.mlperf.mlperf_logger import MLPerfLogger, ThroughputTimer @@ -54,8 +55,11 @@ def _get_arg(args, kwargs, index, name): class MLPerfMegatronPretrainTrainer(MegatronPretrainTrainer): """MegatronPretrainTrainer with MLPerf (mllog) logging.""" - def __init__(self, backend_args): - super().__init__(backend_args) + def __init__(self, backend_args: Any = None, **kwargs): + # The core runtime instantiates every trainer with BaseModule-style + # context kwargs (module_name, primus_config, module_rank, ...). Accept + # and forward them so BaseTrainer can filter them cooperatively. + super().__init__(backend_args=backend_args, **kwargs) self.mllogger = MLPerfLogger() self.throughput_timer = None diff --git a/primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py b/primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py new file mode 100644 index 000000000..14e775197 --- /dev/null +++ b/primus/backends/megatron/patches/te_patches/fused_bias_swiglu_patches.py @@ -0,0 +1,60 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Transformer Engine Fused Bias SwiGLU Patches + +Patches SwiGLUFunction to use TE's fused swiglu/dswiglu kernels when +USE_TE_SWIGLU=1 is set, providing better performance on ROCm GPUs. +""" + +import os + +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +@register_patch( + "megatron.te.fused_bias_swiglu", + backend="megatron", + phase="before_train", + description="Use TE fused swiglu/dswiglu kernels in SwiGLUFunction forward/backward", + condition=lambda ctx: os.getenv("USE_TE_SWIGLU", "0") == "1", +) +def patch_te_fused_bias_swiglu(ctx: PatchContext): + """ + Patch SwiGLUFunction to use Transformer Engine's fused swiglu/dswiglu + C++ extensions in forward and backward passes. + + Activated when USE_TE_SWIGLU=1 is set in the environment. + """ + from megatron.core.fusions.fused_bias_swiglu import SwiGLUFunction + from transformer_engine.pytorch.cpp_extensions import dswiglu as te_dswiglu + from transformer_engine.pytorch.cpp_extensions import swiglu as te_swiglu + + @staticmethod + def new_forward(ctx, input, fp8_input_store, cpu_offload_input): + input_for_backward = input.to(__import__("torch").float8_e4m3fn) if fp8_input_store else input + if cpu_offload_input: + input_for_backward.activation_offloading = True + ctx.save_for_backward(input_for_backward) + ctx.ori_input_dtype = input.dtype + ctx.fp8_input_store = fp8_input_store + return te_swiglu(input, None) + + @staticmethod + def new_backward(ctx, grad_output): + input = ctx.saved_tensors[0] + input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input + return te_dswiglu(grad_output, input, None), None, None + + SwiGLUFunction.forward = new_forward + SwiGLUFunction.backward = new_backward + + log_rank_0( + "[Patch:megatron.te.fused_bias_swiglu] Patched SwiGLUFunction " + "to use TE fused swiglu/dswiglu kernels (USE_TE_SWIGLU=1)" + ) diff --git a/primus/backends/megatron/patches/validation_data_sampling_patches.py b/primus/backends/megatron/patches/validation_data_sampling_patches.py new file mode 100644 index 000000000..ca3e5c13f --- /dev/null +++ b/primus/backends/megatron/patches/validation_data_sampling_patches.py @@ -0,0 +1,168 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Megatron Validation Data Sampling Patches + +Patches Megatron's validation data loading to use a fixed, reproducible +subset of validation samples (eval_iters * global_batch_size) starting +from offset 0, matching MLPerf evaluation protocol. + +Without this patch, Megatron: + 1. Over-allocates validation samples proportional to total training steps + 2. Advances through validation data using consumed_valid_samples offset, + making eval loss non-reproducible across runs +""" + +import os + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +def _is_mlperf_enabled(ctx): + return os.getenv("PRIMUS_MLPERF", "0") == "1" and getattr(get_args(ctx), "eval_iters", 0) > 0 + + +@register_patch( + "megatron.training.validation_num_samples", + backend="megatron", + phase="before_train", + description=( + "Fix validation sample count to eval_iters * global_batch_size " "instead of scaling with train_iters" + ), + condition=_is_mlperf_enabled, +) +def patch_validation_num_samples(ctx: PatchContext): + """ + Patch get_train_valid_test_num_samples to allocate only + eval_iters * global_batch_size validation samples instead of + (train_iters // eval_interval + 1) * eval_iters * global_batch_size. + """ + import megatron.training.training as training_module + + def patched_get_train_valid_test_num_samples(): + from megatron.training import get_args + + args = get_args() + + if args.train_samples: + train_samples = args.train_samples + else: + train_samples = args.train_iters * args.global_batch_size + + if args.full_validation: + eval_samples = None + else: + eval_samples = args.eval_iters * args.global_batch_size + + test_samples = args.eval_iters * args.global_batch_size + + if hasattr(args, "phase_transition_iterations") and args.phase_transition_iterations: + phase_transition_samples = ( + [0] + + [t * args.global_batch_size for t in args.phase_transition_iterations] + + [args.train_samples] + ) + current_sample = args.iteration * args.global_batch_size + for i in range(len(phase_transition_samples) - 1): + if phase_transition_samples[i] <= current_sample < phase_transition_samples[i + 1]: + train_samples = phase_transition_samples[i + 1] - phase_transition_samples[i] + break + + return (train_samples, eval_samples, test_samples) + + training_module.get_train_valid_test_num_samples = patched_get_train_valid_test_num_samples + log_rank_0( + "[Patch:megatron.training.validation_num_samples] " + "Patched get_train_valid_test_num_samples: eval_samples = eval_iters * gbs" + ) + + +@register_patch( + "megatron.training.validation_data_loader", + backend="megatron", + phase="before_train", + description=( + "Always build validation dataloader with consumed_samples=0 and " + "cap total_samples at eval_iters * gbs for reproducible evaluation" + ), + condition=_is_mlperf_enabled, +) +def patch_validation_data_loader(ctx: PatchContext): + """ + Patch build_pretraining_data_loader to detect validation datasets and + force consumed_samples=0 with total_samples capped at eval_iters * gbs. + + Also patches build_train_valid_test_data_loaders to always pass + consumed_samples=0 for validation dataloaders (instead of + args.consumed_valid_samples). + """ + import torch.utils.data + from megatron.core import mpu + from megatron.core.datasets.utils import Split + + try: + from megatron.training.datasets import data_samplers as samplers_module + from megatron.training.datasets.data_samplers import MegatronPretrainingSampler + except ImportError: + from megatron.legacy.data import data_samplers as samplers_module + from megatron.legacy.data.data_samplers import MegatronPretrainingSampler + + orig_build_loader = samplers_module.build_pretraining_data_loader + + def patched_build_pretraining_data_loader(dataset, consumed_samples, name=""): + """Replacement that forces validation to always start from sample 0 + with a fixed sample count.""" + if dataset is None: + return None + + from megatron.training import get_args + + args = get_args() + + if hasattr(dataset, "split"): + split = dataset.split + elif hasattr(dataset, "index_split"): + split = dataset.index_split + else: + split = None + + is_validation = (split == Split.valid) or (name == "validation") + + if is_validation: + eval_samples = args.eval_iters * args.global_batch_size + total_samples = min(len(dataset), eval_samples) + batch_sampler = MegatronPretrainingSampler( + total_samples=total_samples, + consumed_samples=0, + micro_batch_size=args.micro_batch_size, + data_parallel_rank=mpu.get_data_parallel_rank(), + data_parallel_size=mpu.get_data_parallel_world_size(), + ) + return torch.utils.data.DataLoader( + dataset, + batch_sampler=batch_sampler, + num_workers=args.num_workers, + pin_memory=True, + persistent_workers=True if args.num_workers > 0 else False, + ) + + return orig_build_loader(dataset, consumed_samples) + + # Replace in both the data_samplers module and training module so all + # call sites (including build_train_valid_test_data_loaders) use our version + samplers_module.build_pretraining_data_loader = patched_build_pretraining_data_loader + + import megatron.training.training as training_module + + training_module.build_pretraining_data_loader = patched_build_pretraining_data_loader + + log_rank_0( + "[Patch:megatron.training.validation_data_loader] " + "Patched build_pretraining_data_loader: validation always uses " + "consumed_samples=0, total_samples capped at eval_iters * gbs" + ) diff --git a/primus/backends/megatron/training/evaluator.py b/primus/backends/megatron/training/evaluator.py index 06071b46f..68da99d69 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -7,6 +7,7 @@ import time import torch +from megatron.core import parallel_state from megatron.core.full_cuda_graph import FullCudaGraphWrapper from megatron.core.num_microbatches_calculator import get_num_microbatches from megatron.core.pipeline_parallel import get_forward_backward_func @@ -140,10 +141,32 @@ def primus_evaluate( total_loss_dict = {} if is_pipeline_stage_containing_loss(): for key in total_loss_numerators.keys(): - if total_loss_denominators[key] > 0: - total_loss_dict[key] = total_loss_numerators[key] / total_loss_denominators[key] + # Reduce numerator/denominator across data-parallel ranks so the + # validation loss is a TRUE global average, identical on every rank. + # Without this, args._eval_val_loss stays a per-rank local value, and + # the target-eval-loss early stop (mlperf_pretrain_trainer.py) is then + # evaluated inconsistently: near the target one rank's local loss can + # dip <= target and exit train() alone while the others keep training, + # desyncing collectives (grad-norm all-reduce) -> NCCL hang at ~172k. + reduced = torch.tensor( + [float(total_loss_numerators[key]), float(total_loss_denominators[key])], + dtype=torch.float64, + device="cuda", + ) + torch.distributed.all_reduce( + reduced, + op=torch.distributed.ReduceOp.SUM, + group=parallel_state.get_data_parallel_group(), + ) + # Keep the result as a 0-dim tensor: downstream Megatron code + # (evaluate_and_print_results) and mlperf logging call .item() on it. + if reduced[1].item() > 0: + total_loss_dict[key] = (reduced[0] / reduced[1]).to(torch.float32) else: - total_loss_dict[key] = 0.0 + total_loss_dict[key] = torch.zeros((), dtype=torch.float32, device="cuda") + if "lm loss" in total_loss_dict: + val = total_loss_dict["lm loss"] + args._eval_val_loss = val.item() if hasattr(val, "item") else float(val) collected_non_loss_data = None if non_loss_data_func is not None: From 04e9a35f71f4fc48816f17d711cf7cacd4d4d603 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Wed, 15 Jul 2026 08:48:15 +0300 Subject: [PATCH 028/127] feat(flux): Flux HF->Primus checkpoint conversion tools (#818) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/flux` — review after it. ## What this changes Standalone tooling — the HF→Primus Flux checkpoint converter and the empty-encoding generator, plus the converter test. The converter module itself ships in the Flux model PR; this PR is just the CLI tools that use it. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/flux`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion/test_flux_checkpoint_converter.py`. Validated locally on an AMD GPU container: 5 passed. ## Files 3 (checkpoint converter tool, empty-encoding generator + test). Co-authored-by: Flux Split Trial Co-authored-by: luiza-amd --- .../test_flux_checkpoint_converter.py | 315 ++++++++++++++++++ .../convert_flux_hf_to_primus.py | 154 +++++++++ tools/generate_empty_encodings.py | 149 +++++++++ 3 files changed, 618 insertions(+) create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_flux_checkpoint_converter.py create mode 100755 tools/checkpoint_conversion/convert_flux_hf_to_primus.py create mode 100644 tools/generate_empty_encodings.py diff --git a/tests/unit_tests/backends/megatron/diffusion/test_flux_checkpoint_converter.py b/tests/unit_tests/backends/megatron/diffusion/test_flux_checkpoint_converter.py new file mode 100644 index 000000000..19cda2faf --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_flux_checkpoint_converter.py @@ -0,0 +1,315 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for Flux checkpoint converter. + +Tests QKV fusion, key mapping, and checkpoint conversion logic. +""" + +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux import FluxConfig +from tests.utils import PrimusUT + +# Note: We use lazy imports for checkpoint_converter functions to avoid +# triggering parent package imports that require Megatron. Functions are +# imported inside test methods when needed. + + +class TestQKVWeightFusion(PrimusUT): + """Test QKV weight fusion for GQA.""" + + def test_qkv_weight_fusion_shape(self): + """Test that QKV fusion produces correct output shape.""" + # Lazy import to avoid parent package Megatron dependency + from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import ( + _fuse_qkv_weights, + ) + + # Use Flux 12B config: hidden_size=3072, num_heads=24, head_size=128 + config = FluxConfig.flux_12b() + + hidden_size = config.hidden_size + q_weight = torch.randn(hidden_size, hidden_size) + k_weight = torch.randn(hidden_size, hidden_size) + v_weight = torch.randn(hidden_size, hidden_size) + + # Fuse + qkv_weight = _fuse_qkv_weights(config, q_weight, k_weight, v_weight) + + # Expected shape: [head_size * (num_heads + 2*num_query_groups), hidden_size] + # For Flux: num_query_groups = num_heads = 24 + # So: [128 * (24 + 2*24), 3072] = [128 * 72, 3072] = [9216, 3072] + num_heads = config.num_attention_heads + num_query_groups = getattr(config, "num_query_groups", num_heads) + head_size = hidden_size // num_heads + expected_shape = (head_size * (num_heads + 2 * num_query_groups), hidden_size) + + assert qkv_weight.shape == expected_shape, f"Expected {expected_shape}, got {qkv_weight.shape}" + + def test_qkv_weight_fusion_interleaving(self): + """Test that QKV weights are interleaved correctly per group.""" + # Lazy import to avoid parent package Megatron dependency + from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import ( + _fuse_qkv_weights, + ) + + # Use Flux 535M config with simple values for easier validation + # Note: Flux 535M has hidden_size=3072, num_heads=24, but we adjust expectations + # for the interleaving test which needs smaller values + config = FluxConfig.flux_535m() + # Override for test simplicity - fusion algorithm works with any config + config.hidden_size = 64 + config.num_attention_heads = 4 + config.num_query_groups = 4 + + hidden_size = 64 + head_size = 16 + + # Create identifiable weights + q_weight = torch.ones(hidden_size, hidden_size) * 1.0 + k_weight = torch.ones(hidden_size, hidden_size) * 2.0 + v_weight = torch.ones(hidden_size, hidden_size) * 3.0 + + qkv_weight = _fuse_qkv_weights(config, q_weight, k_weight, v_weight) + + # Reshape to verify interleaving: [heads_per_group + 2, num_groups, head_size, hidden_size] + # For this config: heads_per_group=1, num_groups=4 + # So we expect pattern: [Q0, K0, V0, Q1, K1, V1, Q2, K2, V2, Q3, K3, V3] + qkv_reshaped = qkv_weight.reshape(4, 3, head_size, hidden_size) + + # Check first group + assert torch.allclose(qkv_reshaped[0, 0], torch.ones(head_size, hidden_size) * 1.0), "Q0 mismatch" + assert torch.allclose(qkv_reshaped[0, 1], torch.ones(head_size, hidden_size) * 2.0), "K0 mismatch" + assert torch.allclose(qkv_reshaped[0, 2], torch.ones(head_size, hidden_size) * 3.0), "V0 mismatch" + + +class TestQKVBiasFusion(PrimusUT): + """Test QKV bias fusion for GQA.""" + + def test_qkv_bias_fusion_shape(self): + """Test that QKV bias fusion produces correct output shape.""" + # Lazy import to avoid parent package Megatron dependency + from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import ( + _fuse_qkv_bias, + ) + + # Use Flux 12B config + config = FluxConfig.flux_12b() + + hidden_size = config.hidden_size + q_bias = torch.randn(hidden_size) + k_bias = torch.randn(hidden_size) + v_bias = torch.randn(hidden_size) + + qkv_bias = _fuse_qkv_bias(config, q_bias, k_bias, v_bias) + + # Expected shape: [head_size * (num_heads + 2*num_query_groups)] + num_heads = config.num_attention_heads + num_query_groups = getattr(config, "num_query_groups", num_heads) + head_size = hidden_size // num_heads + expected_shape = (head_size * (num_heads + 2 * num_query_groups),) + + assert qkv_bias.shape == expected_shape, f"Expected {expected_shape}, got {qkv_bias.shape}" + + def test_qkv_bias_fusion_interleaving(self): + """Test that QKV biases are interleaved correctly.""" + # Lazy import to avoid parent package Megatron dependency + from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import ( + _fuse_qkv_bias, + ) + + # Use Flux 535M config with simple values for easier validation + config = FluxConfig.flux_535m() + # Override for test simplicity - fusion algorithm works with any config + config.hidden_size = 64 + config.num_attention_heads = 4 + config.num_query_groups = 4 + + # Create identifiable biases + q_bias = torch.ones(64) * 1.0 + k_bias = torch.ones(64) * 2.0 + v_bias = torch.ones(64) * 3.0 + + qkv_bias = _fuse_qkv_bias(config, q_bias, k_bias, v_bias) + + # Reshape to verify: [num_groups, heads_per_group + 2, head_size] + qkv_reshaped = qkv_bias.reshape(4, 3, 16) + + # Check first group + assert torch.allclose(qkv_reshaped[0, 0], torch.ones(16) * 1.0) + assert torch.allclose(qkv_reshaped[0, 1], torch.ones(16) * 2.0) + assert torch.allclose(qkv_reshaped[0, 2], torch.ones(16) * 3.0) + + +class TestCheckpointConversion: + """Test end-to-end checkpoint conversion (with mock data).""" + + def test_mock_conversion(self, tmp_path): + """Test conversion with mock HF checkpoint.""" + from safetensors.torch import save_file as save_safetensors + + from primus.backends.megatron.core.models.diffusion.flux.checkpoint_converter import ( + convert_hf_checkpoint, + ) + + # Create minimal mock HF checkpoint for Flux 535M (1 joint + 1 single layer) + config = FluxConfig.flux_535m() + hidden_size = config.hidden_size + + mock_state_dict = {} + + # Double block 0 + mock_state_dict["transformer_blocks.0.norm1.linear.weight"] = torch.randn(hidden_size, hidden_size) + mock_state_dict["transformer_blocks.0.norm1.linear.bias"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.to_q.weight"] = torch.randn(hidden_size, hidden_size) + mock_state_dict["transformer_blocks.0.attn.to_q.bias"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.to_k.weight"] = torch.randn(hidden_size, hidden_size) + mock_state_dict["transformer_blocks.0.attn.to_k.bias"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.to_v.weight"] = torch.randn(hidden_size, hidden_size) + mock_state_dict["transformer_blocks.0.attn.to_v.bias"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.to_out.0.weight"] = torch.randn(hidden_size, hidden_size) + mock_state_dict["transformer_blocks.0.attn.to_out.0.bias"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.norm_q.weight"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.norm_k.weight"] = torch.randn(hidden_size) + + # Added attention + mock_state_dict["transformer_blocks.0.attn.add_q_proj.weight"] = torch.randn(hidden_size, hidden_size) + mock_state_dict["transformer_blocks.0.attn.add_q_proj.bias"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.add_k_proj.weight"] = torch.randn(hidden_size, hidden_size) + mock_state_dict["transformer_blocks.0.attn.add_k_proj.bias"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.add_v_proj.weight"] = torch.randn(hidden_size, hidden_size) + mock_state_dict["transformer_blocks.0.attn.add_v_proj.bias"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.to_add_out.weight"] = torch.randn(hidden_size, hidden_size) + mock_state_dict["transformer_blocks.0.attn.to_add_out.bias"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.norm_added_q.weight"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.attn.norm_added_k.weight"] = torch.randn(hidden_size) + mock_state_dict["transformer_blocks.0.norm1_context.linear.weight"] = torch.randn( + hidden_size, hidden_size + ) + mock_state_dict["transformer_blocks.0.norm1_context.linear.bias"] = torch.randn(hidden_size) + + # MLP + mock_state_dict["transformer_blocks.0.ff.net.0.proj.weight"] = torch.randn( + hidden_size * 4, hidden_size + ) + mock_state_dict["transformer_blocks.0.ff.net.0.proj.bias"] = torch.randn(hidden_size * 4) + mock_state_dict["transformer_blocks.0.ff.net.2.weight"] = torch.randn(hidden_size, hidden_size * 4) + mock_state_dict["transformer_blocks.0.ff.net.2.bias"] = torch.randn(hidden_size) + + # Context MLP + mock_state_dict["transformer_blocks.0.ff_context.net.0.proj.weight"] = torch.randn( + hidden_size * 4, hidden_size + ) + mock_state_dict["transformer_blocks.0.ff_context.net.0.proj.bias"] = torch.randn(hidden_size * 4) + mock_state_dict["transformer_blocks.0.ff_context.net.2.weight"] = torch.randn( + hidden_size, hidden_size * 4 + ) + mock_state_dict["transformer_blocks.0.ff_context.net.2.bias"] = torch.randn(hidden_size) + + # Single block 0 + mock_state_dict["single_transformer_blocks.0.norm.linear.weight"] = torch.randn( + hidden_size, hidden_size + ) + mock_state_dict["single_transformer_blocks.0.norm.linear.bias"] = torch.randn(hidden_size) + mock_state_dict["single_transformer_blocks.0.attn.to_q.weight"] = torch.randn( + hidden_size, hidden_size + ) + mock_state_dict["single_transformer_blocks.0.attn.to_q.bias"] = torch.randn(hidden_size) + mock_state_dict["single_transformer_blocks.0.attn.to_k.weight"] = torch.randn( + hidden_size, hidden_size + ) + mock_state_dict["single_transformer_blocks.0.attn.to_k.bias"] = torch.randn(hidden_size) + mock_state_dict["single_transformer_blocks.0.attn.to_v.weight"] = torch.randn( + hidden_size, hidden_size + ) + mock_state_dict["single_transformer_blocks.0.attn.to_v.bias"] = torch.randn(hidden_size) + mock_state_dict["single_transformer_blocks.0.attn.norm_q.weight"] = torch.randn(hidden_size) + mock_state_dict["single_transformer_blocks.0.attn.norm_k.weight"] = torch.randn(hidden_size) + mock_state_dict["single_transformer_blocks.0.proj_mlp.weight"] = torch.randn( + hidden_size * 4, hidden_size + ) + mock_state_dict["single_transformer_blocks.0.proj_mlp.bias"] = torch.randn(hidden_size * 4) + mock_state_dict["single_transformer_blocks.0.proj_out.weight"] = torch.randn( + hidden_size, hidden_size * 2 + ) + mock_state_dict["single_transformer_blocks.0.proj_out.bias"] = torch.randn(hidden_size) + + # Root level + mock_state_dict["x_embedder.weight"] = torch.randn(hidden_size, 64) + mock_state_dict["x_embedder.bias"] = torch.randn(hidden_size) + mock_state_dict["context_embedder.weight"] = torch.randn(hidden_size, 4096) + mock_state_dict["context_embedder.bias"] = torch.randn(hidden_size) + mock_state_dict["time_text_embed.timestep_embedder.linear_1.weight"] = torch.randn(hidden_size, 256) + mock_state_dict["time_text_embed.timestep_embedder.linear_1.bias"] = torch.randn(hidden_size) + mock_state_dict["time_text_embed.timestep_embedder.linear_2.weight"] = torch.randn( + hidden_size, hidden_size + ) + mock_state_dict["time_text_embed.timestep_embedder.linear_2.bias"] = torch.randn(hidden_size) + mock_state_dict["time_text_embed.text_embedder.linear_1.weight"] = torch.randn(hidden_size, 768) + mock_state_dict["time_text_embed.text_embedder.linear_1.bias"] = torch.randn(hidden_size) + mock_state_dict["time_text_embed.text_embedder.linear_2.weight"] = torch.randn( + hidden_size, hidden_size + ) + mock_state_dict["time_text_embed.text_embedder.linear_2.bias"] = torch.randn(hidden_size) + mock_state_dict["time_text_embed.guidance_embedder.linear_1.weight"] = torch.randn(hidden_size, 256) + mock_state_dict["time_text_embed.guidance_embedder.linear_1.bias"] = torch.randn(hidden_size) + mock_state_dict["time_text_embed.guidance_embedder.linear_2.weight"] = torch.randn( + hidden_size, hidden_size + ) + mock_state_dict["time_text_embed.guidance_embedder.linear_2.bias"] = torch.randn(hidden_size) + mock_state_dict["norm_out.linear.weight"] = torch.randn(hidden_size, hidden_size) + mock_state_dict["norm_out.linear.bias"] = torch.randn(hidden_size) + mock_state_dict["proj_out.weight"] = torch.randn(64, hidden_size) + mock_state_dict["proj_out.bias"] = torch.randn(64) + + # Save mock checkpoint + checkpoint_path = tmp_path / "mock_flux.safetensors" + save_safetensors(mock_state_dict, str(checkpoint_path)) + + # Convert + primus_state_dict = convert_hf_checkpoint( + checkpoint_path, + flux_config=config, + save_to=None, + ) + + # Verify conversion + assert len(primus_state_dict) > 0, "Converted state dict is empty" + + # Check that QKV was fused (with new TransformerBlock naming) + assert "transformer.layers.0.self_attention.linear_qkv.weight" in primus_state_dict + assert "transformer.layers.0.self_attention.linear_qkv.bias" in primus_state_dict + assert "transformer.layers.1.self_attention.linear_qkv.weight" in primus_state_dict + + # Check that proj_out was split for single blocks + assert "transformer.layers.1.self_attention.linear_proj.weight" in primus_state_dict + assert "transformer.layers.1.mlp.linear_fc2.weight" in primus_state_dict + + # Verify key mapping worked (with new TransformerBlock naming) + assert "transformer.layers.0.adaln.adaLN_modulation.1.weight" in primus_state_dict + # Note: Layer 1 is FluxSingleTransformerBlock which has different adaln structure + + # norm_out scale/shift swap (the converter's only non-trivial root-level + # math): HF Diffusers stores the modulation as [SCALE; SHIFT] but + # Primus/BFL native expects [SHIFT; SCALE], so the two halves must be + # swapped along dim 0 -- a plain key-rename/passthrough would be a bug. + orig_w = mock_state_dict["norm_out.linear.weight"] + orig_b = mock_state_dict["norm_out.linear.bias"] + half_w = orig_w.shape[0] // 2 + half_b = orig_b.shape[0] // 2 + expected_w = torch.cat([orig_w[half_w:, :], orig_w[:half_w, :]], dim=0) + expected_b = torch.cat([orig_b[half_b:], orig_b[:half_b]], dim=0) + + conv_w = primus_state_dict["norm_out.adaLN_modulation.1.weight"] + conv_b = primus_state_dict["norm_out.adaLN_modulation.1.bias"] + assert torch.equal(conv_w, expected_w), "norm_out weight halves were not swapped" + assert torch.equal(conv_b, expected_b), "norm_out bias halves were not swapped" + # Positive control: the swap must actually reorder, not pass through. + assert not torch.equal(conv_w, orig_w), "norm_out weight passed through without the scale/shift swap" + assert not torch.equal(conv_b, orig_b), "norm_out bias passed through without the scale/shift swap" diff --git a/tools/checkpoint_conversion/convert_flux_hf_to_primus.py b/tools/checkpoint_conversion/convert_flux_hf_to_primus.py new file mode 100755 index 000000000..b73b95b9a --- /dev/null +++ b/tools/checkpoint_conversion/convert_flux_hf_to_primus.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Convert HuggingFace Flux checkpoint to Primus format. + +This tool converts HuggingFace Diffusers Flux transformer checkpoints to +Primus/Megatron-Core compatible format. It handles: +- QKV weight fusion for grouped-query attention (GQA) +- Key mapping from HF to Primus naming conventions +- Single block proj_out splitting +- Multi-file safetensors loading + +Usage: + # Convert FLUX.1-dev checkpoint + python tools/checkpoint_conversion/convert_flux_hf_to_primus.py \\ + --input black-forest-labs/FLUX.1-dev \\ + --output checkpoints/primus_flux_12b.safetensors \\ + --variant flux_12b + + # Convert with custom architecture + python tools/checkpoint_conversion/convert_flux_hf_to_primus.py \\ + --input path/to/checkpoint \\ + --output checkpoints/primus_custom.safetensors \\ + --variant custom \\ + --num-joint-layers 10 \\ + --num-single-layers 20 + +Example: + $ python tools/checkpoint_conversion/convert_flux_hf_to_primus.py \\ + --input black-forest-labs/FLUX.1-dev \\ + --output checkpoints/primus_flux_12b.safetensors \\ + --variant flux_12b + + Converting flux_12b checkpoint + Input: black-forest-labs/FLUX.1-dev + Output: checkpoints/primus_flux_12b.safetensors + Architecture: 19 joint + 38 single layers + Loading HuggingFace checkpoint from: black-forest-labs/FLUX.1-dev + ... + Conversion complete! +""" + +import argparse +import sys +from pathlib import Path + +# Add primus to path +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +from primus.backends.megatron.core.models.diffusion.flux import ( + FluxConfig, + convert_hf_checkpoint, +) + + +def main(): + parser = argparse.ArgumentParser( + description="Convert HuggingFace Flux checkpoint to Primus format", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Convert FLUX.1-dev (12B parameters) + %(prog)s --input black-forest-labs/FLUX.1-dev \\ + --output checkpoints/primus_flux_12b.safetensors \\ + --variant flux_12b + + # Convert local checkpoint directory + %(prog)s --input /path/to/flux/checkpoint \\ + --output primus_flux.safetensors \\ + --variant flux_12b + + # Convert with custom architecture + %(prog)s --input /path/to/checkpoint \\ + --output primus_custom.safetensors \\ + --variant custom \\ + --num-joint-layers 10 \\ + --num-single-layers 20 + """, + ) + + parser.add_argument( + "--input", + required=True, + help="Path to HF checkpoint (file, directory, or HF model ID like 'black-forest-labs/FLUX.1-dev')", + ) + parser.add_argument("--output", required=True, help="Output path for Primus checkpoint (.safetensors)") + parser.add_argument( + "--variant", + choices=["flux_535m", "flux_12b", "custom"], + default="flux_12b", + help="Flux variant (determines architecture). Default: flux_12b", + ) + parser.add_argument( + "--num-joint-layers", type=int, help="Number of joint layers (only with --variant custom)" + ) + parser.add_argument( + "--num-single-layers", type=int, help="Number of single layers (only with --variant custom)" + ) + + args = parser.parse_args() + + # Validate custom variant arguments + if args.variant == "custom": + if not args.num_joint_layers or not args.num_single_layers: + parser.error("--num-joint-layers and --num-single-layers are required with --variant custom") + else: + if args.num_joint_layers or args.num_single_layers: + parser.error("--num-joint-layers and --num-single-layers can only be used with --variant custom") + + # Create config based on variant + if args.variant == "flux_535m": + config = FluxConfig.flux_535m() + elif args.variant == "flux_12b": + config = FluxConfig.flux_12b() + else: # custom + config = FluxConfig(num_joint_layers=args.num_joint_layers, num_single_layers=args.num_single_layers) + + # Print conversion info + print("=" * 80) + print(f"Converting {args.variant} checkpoint") + print(f" Input: {args.input}") + print(f" Output: {args.output}") + print(f" Architecture: {config.num_joint_layers} joint + {config.num_single_layers} single layers") + print("=" * 80) + print() + + # Convert checkpoint + try: + convert_hf_checkpoint( + checkpoint_path=args.input, + flux_config=config, + save_to=args.output, + ) + + print() + print("=" * 80) + print("✓ Conversion complete!") + print(f"Primus checkpoint saved to: {args.output}") + print("=" * 80) + + return 0 + + except Exception as e: + print() + print("=" * 80) + print(f"✗ Conversion failed: {e}") + print("=" * 80) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/generate_empty_encodings.py b/tools/generate_empty_encodings.py new file mode 100644 index 000000000..10e55781d --- /dev/null +++ b/tools/generate_empty_encodings.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Generate empty T5 and CLIP encodings for CFG dropout. + +Runs T5-XXL and CLIP-L on the empty string ("") and saves the resulting +embeddings as .npy files. These are loaded by FluxPretrainTrainer at init +to replace text embeddings during classifier-free guidance dropout. + +Using real model outputs (instead of torch.randn) is critical for training +convergence — see NVIDIA MLPerf Training v5.1 reference. + +Usage: + python tools/generate_empty_encodings.py \ + --output_dir /path/to/empty_encodings \ + --t5_model google/t5-v1_1-xxl \ + --clip_model openai/clip-vit-large-patch14 \ + --t5_max_length 256 + +Output files: + t5_empty.npy - shape (1, seq_len, 4096) + clip_empty.npy - shape (1, 768) +""" + +import argparse +import os + +import numpy as np +import torch + + +def generate_t5_empty(model_name: str, max_length: int, device: str) -> np.ndarray: + """Generate T5-XXL encoding for empty string.""" + from transformers import T5EncoderModel, T5Tokenizer + + print(f"Loading T5 tokenizer: {model_name}") + tokenizer = T5Tokenizer.from_pretrained(model_name) + + print(f"Loading T5 model: {model_name}") + model = T5EncoderModel.from_pretrained(model_name, torch_dtype=torch.float32) + model = model.to(device).eval() + + with torch.no_grad(): + inputs = tokenizer( + "", + max_length=max_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ).to(device) + + outputs = model(**inputs) + # outputs.last_hidden_state: (1, seq_len, 4096) + embeddings = outputs.last_hidden_state.cpu().numpy() + + print(f"T5 empty encoding shape: {embeddings.shape}") + del model + torch.cuda.empty_cache() + return embeddings + + +def generate_clip_empty(model_name: str, device: str) -> np.ndarray: + """Generate CLIP-L pooled encoding for empty string.""" + from transformers import CLIPTextModel, CLIPTokenizer + + print(f"Loading CLIP tokenizer: {model_name}") + tokenizer = CLIPTokenizer.from_pretrained(model_name) + + print(f"Loading CLIP model: {model_name}") + model = CLIPTextModel.from_pretrained(model_name, torch_dtype=torch.float32) + model = model.to(device).eval() + + with torch.no_grad(): + inputs = tokenizer( + "", + max_length=tokenizer.model_max_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ).to(device) + + outputs = model(**inputs) + # outputs.pooler_output: (1, 768) + pooled = outputs.pooler_output.cpu().numpy() + + print(f"CLIP empty encoding shape: {pooled.shape}") + del model + torch.cuda.empty_cache() + return pooled + + +def main(): + parser = argparse.ArgumentParser(description="Generate empty T5/CLIP encodings for CFG dropout") + parser.add_argument( + "--output_dir", + type=str, + required=True, + help="Directory to save t5_empty.npy and clip_empty.npy", + ) + parser.add_argument( + "--t5_model", + type=str, + default="google/t5-v1_1-xxl", + help="T5 model name or path (default: google/t5-v1_1-xxl)", + ) + parser.add_argument( + "--clip_model", + type=str, + default="openai/clip-vit-large-patch14", + help="CLIP model name or path (default: openai/clip-vit-large-patch14)", + ) + parser.add_argument( + "--t5_max_length", + type=int, + default=256, + help="Max sequence length for T5 encoding (default: 256 for schnell)", + ) + parser.add_argument( + "--device", + type=str, + default="cuda" if torch.cuda.is_available() else "cpu", + help="Device to run models on (default: cuda if available)", + ) + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + + t5_path = os.path.join(args.output_dir, "t5_empty.npy") + clip_path = os.path.join(args.output_dir, "clip_empty.npy") + + t5_embeddings = generate_t5_empty(args.t5_model, args.t5_max_length, args.device) + np.save(t5_path, t5_embeddings) + print(f"Saved T5 empty encodings to: {t5_path}") + + clip_embeddings = generate_clip_empty(args.clip_model, args.device) + np.save(clip_path, clip_embeddings) + print(f"Saved CLIP empty encodings to: {clip_path}") + + print("\nDone! Add to your YAML config:") + print(f" empty_encodings_path: {args.output_dir}") + + +if __name__ == "__main__": + main() From 82a3c5d5421638730436fc5d46adbc1ccce140e8 Mon Sep 17 00:00:00 2001 From: Kailash Gogineni Date: Tue, 14 Jul 2026 22:53:58 -0700 Subject: [PATCH 029/127] [Model][Hardware][AMD] Add MLPerf Training 6.0 Llama2-70B LoRA post-training example on MI355X (#877) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds an end-to-end MLPerf Training 6.0 **Llama2-70B LoRA** post-training example targeting **AMD MI355X** (8× GPU, 1 node), driven through Megatron-Bridge and `primus-cli`. - Dataset: [GovReport](https://gov-report-data.github.io/) (SCROLLS `gov_report`), packed to **8192** tokens. - Model: **meta-llama/Llama-2-70b-hf** with LoRA (rank 16, alpha 32). - Precision: **MXFP4** + BF16, with **FP8 delayed scaling** after healing at step 340. - Quality target: eval loss **< 0.925**. ## Changes ### Example (`examples/mlperf/llama2_70b/`) - `README.md` — full run instructions (container launch, data/checkpoint paths, config reference). - `config_MI355X_1x8x1.sh` — MLPerf 6.0 MI355X env (MXFP4, AITER, NCCL, MLLOG, 550 iters, lr=0.0006). - `configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml` — post-train overrides. - `a4w4_tuned_gemms.csv` — tuned AITER A4W4 GEMM configs. - `run_and_time.sh` / `run_in_container.sh` — one-shot MLPerf runners via `primus-cli`. ### Megatron-Bridge runtime patches (`primus/backends/megatron_bridge/patches/mlperf_llama2_70b/`) Applied only when the MLPerf Llama2-70B run is selected (see `conditions.py`), replacing prior git patches to `third_party/Megatron-Bridge`: - `megatron_patches.py` — MXFP4 recipe + phase tracking, optional TE SwiGLU (`USE_TE_SWIGLU=1`). - `bridge_patches.py` — data loaders, deterministic eval reset, SFT attention-mask cache, NeMo-style step timing. - `lora.py` — NeMo-stable LoRA (`use_te_fused_lora=False`). - `resettable_data_iterator.py` — deterministic validation iterator. ### MLPerf recipe (`primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/`) - `llama2_custom.py` — `llama2_70b_lora_mxfp4_config` recipe + custom training/eval loop. - `pre_quantize_mxfp4.py` — `PRE_QUANTIZED_MODEL=True` pre-quantization (FP8 stash on CPU, MXFP4 weight swap). - `mxfp4_healing.py` — MXFP4→FP8 healing at `HEALING_ITER=340`. - `nemo_loss.py` — NeMo-equivalent `MaskedTokenLossReduction`. - `_log_suppression.py` — non-MLLOG log suppression for clean submission logs. ### Core / plumbing - `config_utils.py` — recipe resolution now supports direct custom module paths (e.g. `primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.llama2_custom`) with a fallback. - Nested config override logic (`_apply_nested_overrides`) moved from the pretrain trainer to the shared base trainer and wired into the post-train trainer; `runtime_config_update` now invoked before finetune. - `transformer_engine_spec_provider.py` — optional eager-attention fallback. - `cli/main.py` — prefer the git checkout over an installed wheel for in-tree Primus modules; make MLPerf log suppression import optional. - `train_runtime.py` — guard against empty `data_path`. ### Hooks (`runner/helpers/hooks/train/posttrain/megatron_bridge/`) - `00_install_requirements.sh` — container-safe pip cache path; pin `fsspec`. - `01_convert_checkpoints.sh` — container-aware data root / HF cache resolution; HF→Megatron conversion with correct attention env handling. - `02_prepare_mlperf_dataset.sh` (new) — download/convert SCROLLS gov-report and build packed `.npy` + metadata. ### Top-level dataset utilities - `download_dataset.py`, `convert_dataset.py`, `create_metadata.py`, `hash.py` — dataset download, packing, metadata, and integrity hashing. ## MLPerf configuration | Parameter | Value | |-----------|-------| | `train_iters` | 550 | | `global_batch_size` | 8 | | `micro_batch_size` | 1 | | `seq_length` | 8192 | | `lr` | 0.0006 | | `eval_interval` / `eval_iters` | 48 / 24 | | Parallelism | TP=1, PP=1, CP=1 (8 GPUs data parallel) | | Quality target | eval loss < 0.925 | ## Results Measured on **MI355X 1×8×1** (8 GPUs, 1 node) via `examples/mlperf/llama2_70b/run_and_time.sh`. Log: `logs/log_20260715_014540.txt`. | Metric | Value | |--------|-------| | **Time to train** | **517.1 s (8.62 min)** training-loop wall time | | **Final iteration** | 384 / 550 (early exit — quality target reached) | | **Final eval loss** | **0.9242** (target: < 0.925) ✓ | | **Final eval PPL** | 2.52 | | **Consumed train samples** | 3,072 | | **MXFP4→FP8 healing** | Applied at step 340 (`DelayedScaling`, 320 FP8 weights restored) | ### Throughput Steady-state training throughput from `logs/log_20260715_014540.txt` (NeMo-style train-step wall clock, `global_batch_size=8`, `seq_length=8192`): | Phase | Steps | Step time | Model TFLOP/s/GPU | Model TFLOP/s (8× GPU) | Tokens/s/GPU | Tokens/s (8× GPU) | Samples/s | |-------|-------|-----------|-------------------|------------------------|--------------|-------------------|-----------| | Warmup | 10 | 1.30 s | 2,797 | 22,376 | 6,292 | 50,336 | 6.15 | | **MXFP4** | 20–339 | **1.15 s** | **~3,165** | **~25,320** | **~7,120** | **~56,960** | **~6.96** | | **FP8** (post-healing) | 350–380 | **1.43 s** | **~2,540** | **~20,320** | **~5,715** | **~45,720** | **~5.59** | | Eval (iter 384) | — | 13.2 s | — | — | — | — | **5.40** | Cluster tokens/s = `tokens/s/GPU × 8`. Samples/s = `global_batch_size / step_time`. **Eval loss progression** (validation at eval intervals): | Iteration | Eval loss | |-----------|-----------| | 192 | 0.9685 | | 240 | 0.9597 | | 288 | 0.9514 | | 336 | 0.9491 | | **384** | **0.9242** | Training exited early at iteration 384 when eval loss dropped below the MLPerf quality target (0.925). --- ⚡ *Prepared with [PR Pundit](https://github.com/AMD-AGI/pr-pundit) — AMD OSS Agent* --------- Co-authored-by: vidushi8 Co-authored-by: Cursor Co-authored-by: shekhar Co-authored-by: Xiaoming-AMD --- .gitignore | 1 + examples/mlperf/llama2_70b/README.md | 190 ++ .../mlperf/llama2_70b/a4w4_tuned_gemms.csv | 10 + .../mlperf/llama2_70b/config_MI355X_1x8x1.sh | 165 ++ .../llama2_70b_lora_mlperf_posttrain.yaml | 60 + examples/mlperf/llama2_70b/run_and_time.sh | 53 + .../mlperf/llama2_70b/run_in_container.sh | 19 + .../transformer_engine_spec_provider.py | 6 +- primus/backends/megatron/patches/__init__.py | 3 + primus/backends/megatron/sft/dataset.py | 7 +- .../megatron/sft/mlperf_packed_dataset.py | 2 +- .../backends/megatron_bridge/config_utils.py | 30 +- .../megatron_bridge_base_trainer.py | 36 + .../megatron_bridge_posttrain_trainer.py | 4 + .../megatron_bridge_pretrain_trainer.py | 39 - .../patches/mlperf_llama2_70b/__init__.py | 7 + .../mlperf_llama2_70b/bridge_patches.py | 487 ++++ .../patches/mlperf_llama2_70b/conditions.py | 43 + .../patches/mlperf_llama2_70b/lora.py | 188 ++ .../mlperf_llama2_70b/megatron_patches.py | 193 ++ .../resettable_data_iterator.py | 35 + .../recipes/mlperf_llama2_70b/__init__.py | 7 + .../mlperf_llama2_70b/_log_suppression.py | 224 ++ .../mlperf_llama2_70b/convert_dataset.py | 58 + .../mlperf_llama2_70b/create_metadata.py | 40 + .../recipes/mlperf_llama2_70b/dataset_hash.py | 51 + .../mlperf_llama2_70b/download_dataset.py | 51 + .../mlperf_llama2_70b/llama2_custom.py | 2056 +++++++++++++++++ .../mlperf_llama2_70b/mxfp4_healing.py | 515 +++++ .../recipes/mlperf_llama2_70b/nemo_loss.py | 591 +++++ .../mlperf_llama2_70b/pre_quantize_mxfp4.py | 464 ++++ primus/cli/main.py | 32 +- .../llama2_70b_lora_mxfp4.yaml | 4 + primus/core/runtime/train_runtime.py | 5 +- .../00_install_requirements.sh | 12 +- .../megatron_bridge/01_convert_checkpoints.sh | 104 +- .../02_prepare_mlperf_dataset.sh | 98 + 37 files changed, 5824 insertions(+), 66 deletions(-) create mode 100644 examples/mlperf/llama2_70b/README.md create mode 100644 examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv create mode 100755 examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh create mode 100644 examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml create mode 100755 examples/mlperf/llama2_70b/run_and_time.sh create mode 100755 examples/mlperf/llama2_70b/run_in_container.sh create mode 100644 primus/backends/megatron_bridge/patches/mlperf_llama2_70b/__init__.py create mode 100644 primus/backends/megatron_bridge/patches/mlperf_llama2_70b/bridge_patches.py create mode 100644 primus/backends/megatron_bridge/patches/mlperf_llama2_70b/conditions.py create mode 100644 primus/backends/megatron_bridge/patches/mlperf_llama2_70b/lora.py create mode 100644 primus/backends/megatron_bridge/patches/mlperf_llama2_70b/megatron_patches.py create mode 100644 primus/backends/megatron_bridge/patches/mlperf_llama2_70b/resettable_data_iterator.py create mode 100644 primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/__init__.py create mode 100644 primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/_log_suppression.py create mode 100644 primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py create mode 100644 primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/create_metadata.py create mode 100644 primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py create mode 100644 primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py create mode 100644 primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/llama2_custom.py create mode 100644 primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/mxfp4_healing.py create mode 100644 primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/nemo_loss.py create mode 100644 primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/pre_quantize_mxfp4.py create mode 100644 primus/configs/models/megatron_bridge/llama2_70b_lora_mxfp4.yaml create mode 100755 runner/helpers/hooks/train/posttrain/megatron_bridge/02_prepare_mlperf_dataset.sh diff --git a/.gitignore b/.gitignore index 1d04b58e5..59e120234 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ dist/ *.egg-info *~ logs +.pip_cache .vscode local/ .gitmodules diff --git a/examples/mlperf/llama2_70b/README.md b/examples/mlperf/llama2_70b/README.md new file mode 100644 index 000000000..4c67ecff7 --- /dev/null +++ b/examples/mlperf/llama2_70b/README.md @@ -0,0 +1,190 @@ +# Llama2-70B LoRA MLPerf on MI355X (Primus) + +MLPerf Training 6.0 Llama2-70B LoRA on **MI355X** (8× GPU, 1 node) via Megatron-Bridge and `primus-cli`. + +Dataset: [GovReport](https://gov-report-data.github.io/) (SCROLLS `gov_report`), packed to **8192** tokens. +Model: **meta-llama/Llama-2-70b-hf** with LoRA (rank 16, alpha 32). +Precision: **MXFP4** + BF16; **FP8 delayed scaling** after healing at step 340. + +## Key files + +- `configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml` — post-train overrides +- `config_MI355X_1x8x1.sh` — system config and env vars (set `PRIMUS_PATH` to your Primus clone) +- `run_and_time.sh` — one-shot MLPerf run via `primus-cli` +- `a4w4_tuned_gemms.csv` — tuned AITER A4W4 GEMM configs + +--- + +## Prerequisites + +- 8× MI355X GPUs on one node +- Hugging Face access to `meta-llama/Llama-2-70b-hf` (`HF_TOKEN`) +- ~300 GB disk for packed data + Megatron checkpoint +- Docker with ROCm (`/dev/kfd`, `/dev/dri`) + +--- + +## 1. Launch container + +```bash +docker pull rocm/primus:v26.4 + +docker run -it \ + --device=/dev/kfd \ + --device=/dev/dri \ + --security-opt seccomp=unconfined \ + --group-add 44 \ + --group-add 109 \ + --cap-add=SYS_PTRACE \ + --ipc=host \ + --shm-size=32g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --memory=0 \ + --memory-swap=0 \ + --privileged \ + --ulimit nofile=65535:65535 \ + -v /home/kgoginen@amd.com/Primus:/workspace/Primus \ + rocm/primus:v26.4 +``` + +Change the `-v` host path to your Primus checkout. Repo is at `/workspace/Primus` inside the container. + +Optional: mount a data volume if data/checkpoints live on the host: + +```bash + -v /path/on/host/data:/data \ +``` + +--- + +## 2. Set data and checkpoint paths + +Inside the container: + +```bash +cd /workspace/Primus + +export HF_TOKEN=hf_... # required on first run (hooks download model + dataset) + +# Packed GovReport .npy files (train.npy, validation.npy, packed_metadata.jsonl) +export PACKED_DATA_DIR=/data + +# Megatron checkpoint root (must contain latest_train_state.pt, not iter_0000000/) +export PRETRAINED_CHECKPOINT=/data/megatron_checkpoints/Llama-2-70b-hf +``` + +Hooks create these under `/data` on first run if missing. Point `PACKED_DATA_DIR` and `PRETRAINED_CHECKPOINT` at existing paths to skip re-download. + +--- + +## 3. Run training + +```bash +bash examples/mlperf/llama2_70b/run_and_time.sh +``` + +Hooks run automatically: pip deps → dataset → HF→Megatron checkpoint. + +Equivalent: + +```bash +source examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh +./runner/primus-cli direct train posttrain \ + --config examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml +``` + +--- + +## 4. MLPerf experiment configuration + +### Config files + +| File | Role | +|------|------| +| `examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml` | Post-train overrides | +| `examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh` | MLPerf env (MXFP4, AITER, NCCL, MLLOG) | +| `examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv` | Tuned AITER A4W4 GEMM configs | +| `primus/configs/models/megatron_bridge/llama2_70b_lora_mxfp4.yaml` | Model recipe | +| `primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/llama2_custom.py` | `llama2_70b_lora_mxfp4_config` | + +### Training schedule + +| Parameter | Value | +|-----------|-------| +| `train_iters` | 550 | +| `global_batch_size` | 8 | +| `micro_batch_size` | 1 | +| `seq_length` | 8192 | +| `lr` | 0.0006 | +| `eval_interval` / `eval_iters` | 48 / 24 | +| Quality target | eval loss **< 0.925** | + +### Precision + +**MXFP4 (steps 0–339):** `fp4=mxfp4`, `fp8=None`, `PRE_QUANTIZED_MODEL=True`, fused attention, AITER A4W4 GEMMs (`a4w4_tuned_gemms.csv`). + +**FP8 healing (step 340+):** `HEALING_ITER=340`, delayed scaling via `FP8_*` env vars in `config_MI355X_1x8x1.sh`. + +### LoRA + +Targets `linear_qkv`, `linear_proj` (dim 16, alpha 32). `stable_lora_with_te_op_fuser=True` (unfused `LoRALinear` adapters). + +### Parallelism + +TP=1, PP=1, CP=1, 8 GPUs data parallel. + +### MLPerf overrides (Primus-side, no third_party git patches) + +Runtime patches under `primus/backends/megatron_bridge/patches/mlperf_llama2_70b/` apply only when +`llama2_70b_lora_mxfp4` / `llama2_70b_lora_mlperf_posttrain.yaml` is selected. + +Recipe code lives under `primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/`. + +| File | Role | +|------|------| +| `lora.py` | NeMo-stable LoRA (`use_te_fused_lora=False`) | +| `resettable_data_iterator.py` | Deterministic validation iterator | +| `bridge_patches.py` | Data loaders, eval reset, SFT mask cache, NeMo timing | +| `megatron_patches.py` | MXFP4 recipe + optional TE SwiGLU | +| `conditions.py` | Scopes patches to MLPerf Llama2-70B only | + +One-time cleanup if you previously applied git patches to submodules: + +```bash +git -C third_party/Megatron-Bridge checkout -- . +git -C third_party/Megatron-Bridge/3rdparty/Megatron-LM checkout -- . +``` + +--- + +## 5. Logging + +Bring-up defaults (`config_MI355X_1x8x1.sh`): `log_interval=10`, `PRIMUS_LOG_GPU_MEM=1`, `VERBOSE_TRAINING_LOG=1`. + +MLPerf submission (quiet): + +```bash +export PRIMUS_LOG_GPU_MEM=0 +export VERBOSE_TRAINING_LOG=0 +# yaml: log_interval: 99999, stderr_sink_level: ERROR +``` + +### Common issues + +| Symptom | Fix | +|---------|-----| +| NCCL hang, 0% GPU | `NCCL_IB_DISABLE=1` (default in config) | +| Invalid pretrained checkpoint | Point at checkpoint **root**, not `iter_0000000` | +| Long silence at start | Pre-quantize + warmup + AITER JIT (normal) | + +--- + +## 6. Optional overrides + +```bash +export PRIMUS_TRAIN_ITERS=550 +export SEED=1234 +export SYNTH_WARMUP_STEPS=0 +export NCCL_IB_DISABLE=0 # if RDMA works on your system +``` diff --git a/examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv b/examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv new file mode 100644 index 000000000..8c4d18c33 --- /dev/null +++ b/examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv @@ -0,0 +1,10 @@ +cu_num,M,N,K,kernelId,splitK,kernelName,errRatio +256,8192,10240,8192,45,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_128x512E,0.0 +256,8192,8192,8192,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,8192,57344,8192,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,8192,8192,28672,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,8192,8192,10240,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,8192,8192,57344,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,8192,28672,8192,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,10240,8192,8192,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 +256,57344,8192,8192,54,0,_ZN5aiter42f4gemm_bf16_per1x32Fp4_BpreShuffle_256x256E,0.0 diff --git a/examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh b/examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh new file mode 100755 index 000000000..59448a069 --- /dev/null +++ b/examples/mlperf/llama2_70b/config_MI355X_1x8x1.sh @@ -0,0 +1,165 @@ +#!/bin/bash +# MLPerf 6.0 environment for Llama2-70B LoRA on MI355X (8 GPUs, 1 node). +# Source before run_and_time.sh or primus-cli direct train posttrain. + +export DGXSYSTEM=MI355X_1x8x1 +export GPUS_PER_NODE=8 +export NNODES=1 +export NODE_RANK=0 +export MASTER_ADDR=localhost +export MASTER_PORT=29502 + +# MLPerf timed runs set SEED=$RANDOM per trial; default here for single-shot primus-cli. +export SEED="${SEED:-$RANDOM}" + +export PRIMUS_PATH="${PRIMUS_PATH:-/workspace/Primus}" +export PYTHONPATH="${PRIMUS_PATH}:${PRIMUS_PATH}/third_party/Megatron-Bridge:${PYTHONPATH:-}" +export EXP="${EXP:-${PRIMUS_PATH}/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml}" +export DATA_PATH="${DATA_PATH:-/data}" + +export PACKED_TRAIN_DATA_PATH="${DATA_PATH}/train.npy" +export PACKED_VAL_DATA_PATH="${DATA_PATH}/validation.npy" +export PACKED_METADATA_PATH="${DATA_PATH}/packed_metadata.jsonl" +export PACKED_DATA_DIR="${DATA_PATH}" + +export PRETRAINED_CHECKPOINT="${PRETRAINED_CHECKPOINT:-/data/megatron_checkpoints/Llama-2-70b-hf}" + +export LR=0.0006 + +export HSA_NO_SCRATCH_RECLAIM=1 +export HSA_ENABLE_SDMA=1 +export HSA_ENABLE_INTERRUPT=0 +export GPU_MAX_HW_QUEUES=2 +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export TORCH_NCCL_HIGH_PRIORITY=1 +export NCCL_CHECKS_DISABLE=1 +# Single-node: use GPU P2P, not ionic IB. Broken libibverbs ABI (ionic kernel +# mismatch) commonly hangs NCCL here with 0% GPU util. MLPerf systems with +# working RDMA can override: NCCL_IB_DISABLE=0 +export NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-1}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-lo}" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-lo}" +export OMP_NUM_THREADS=1 + +export NVTE_USE_AITER_ROPE=1 +# Fused attention for MXFP4 training (must match recipe attention_backend=fused). +export NVTE_FLASH_ATTN=0 +export NVTE_FUSED_ATTN=1 +export NVTE_UNFUSED_ATTN=0 +export NVTE_FUSED_ATTN_CK=1 +export NVTE_FUSED_ATTN_AOTRITON=1 +export NVTE_CK_USES_FWD_V3=1 +export NVTE_CK_USES_BWD_V3=1 +export NVTE_CK_IS_V3_ATOMIC_FP32=0 +export NVTE_RS_STRIDED_ATOMIC=2 +export NVTE_FP8_DPA_BWD=1 +export NVTE_USE_HIPBLASLT=1 +export NVTE_USE_CAST_TRANSPOSE_TRITON=0 +export NVTE_USE_OPTIMIZED_HIPIFIED_CAST_TRANSPOSE=0 +export NVTE_USE_RMSNORM_TRITON=0 +export USE_TE_SWIGLU=1 +export ENABLE_TRANSPOSE_CACHE=0 +export NVTE_MXFP4_USE_HADAMARD=${NVTE_MXFP4_USE_HADAMARD:-1} +export NVTE_DEBUG=0 +export NVTE_DEBUG_LEVEL=0 + +export HEALING_ITER=${HEALING_ITER:-340} +export HEALING_PRECISION=${HEALING_PRECISION:-FP8_DS} +export PRE_QUANTIZED_MODEL=${PRE_QUANTIZED_MODEL:-True} +export NCCL_MIN_P2P_NCHANNELS=${NCCL_MIN_P2P_NCHANNELS:-32} +export NCCL_MIN_CTAS=${NCCL_MIN_CTAS:-32} +export NCCL_NCHANNELS_PER_NET_PEER=${NCCL_NCHANNELS_PER_NET_PEER:-32} +export NCCL_NVLS_ENABLE=${NCCL_NVLS_ENABLE:-0} + +export MEGATRON_BRIDGE_LOGGING_LEVEL=50 +export PYTHONWARNINGS=ignore +export PRIMUS_LOG_LEVEL=ERROR + +# Print rank-0 GPU memory (allocated/reserved/peak + torch memory_stats) every log_interval. +export PRIMUS_LOG_GPU_MEM=${PRIMUS_LOG_GPU_MEM:-1} +# Megatron iteration / TFLOP / loss lines use print_rank_0 (always on). Primus log_rank_0 +# helpers need this for recipe-internal banners during bring-up. +export VERBOSE_TRAINING_LOG=${VERBOSE_TRAINING_LOG:-1} + +export SYNTH_WARMUP_STEPS=5 +export SYNTH_WARMUP_VALID_STEPS=5 + +export ENABLE_MLLOG=1 +export MLLOG_OUTPUT_FILE=/results/mlperf_logging.out +export MLLOG_TRAIN_LOSS_LOG_FREQ=0 +export MLLOG_TARGET_EVAL_LOSS=0.925 +export MLLOG_SUBMISSION_BENCHMARK=llama2_70b_lora +export MLLOG_SUBMISSION_DIVISION=closed +export MLLOG_SUBMISSION_ORG=AMD +export MLLOG_SUBMISSION_PLATFORM=MI355X + +export MLLOG_TENSOR_PARALLELISM=1 +export MLLOG_PIPELINE_PARALLELISM=1 +export MLLOG_CONTEXT_PARALLELISM=1 +export MLLOG_EXPERT_PARALLELISM=1 +export MLLOG_MICRO_BATCH_SIZE=1 +MLLOG_CONFIG_FILENAME=$(basename "${BASH_SOURCE[0]}") +export MLLOG_CONFIG_FILENAME +export MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR=mxfp4 + +export TP_COMM_OVERLAP=False +export MC_TP_OVERLAP_AG=False +export MC_TP_OVERLAP_RS=False +export MC_TP_OVERLAP_RS_DGRAD=False + +export CUBLAS_FORCE_XMMA_KERNEL_INIT=DEVICE + +export LORA_A2A=1 +export POSSIBLE_USER_WARNINGS=0 +export CUDNN_FRONTEND_ATTN_DP_WORKSPACE_LIMIT=0 + +export TP=1 +export PP=1 +export CP=1 +export SP=False +export VBOOST_VALUE=1 +export MBS=1 +export MINIBS=1 +export SKIP_EVALS=3 +export VAL_CHECK_INTERVAL=384 +export HYDRA_FULL_ERROR=1 + +export FP8_DPA=0 +# FP8 env flags below apply to healing (HEALING_ITER=340) and TE delayed scaling, +# not Megatron model_cfg.fp8 during the MXFP4 phase (recipe sets fp8=None, fp4=mxfp4). +export FP8=True +export FP8_AMAX_ALGO=most_recent +export FP8_REDUCE_AMAX=False +export FP8_AMAX_HISTORY=4 +export FP8_ACTIVATION=True + +export FUSED_SOFTMAX=0 +export RMSNORM_CAST=0 + +export PT_TENSOR_VALIDATION=0 +export PROFILE_RPD=0 + +export USE_HIPBLASLT=1 +export TORCH_BLAS_PREFER_HIPBLASLT=1 + +export LOGGING_INTERVAL=5000 +# FP4 weights (MXFP4 e2m1 linear layers); distinct from FP8_* healing flags above. +export FP4=True +export FP4_RECIPE=mxfp4 +export MAX_STEPS=550 +export NEXP=1 + +export LOAD_CKPT=True +export MCORE_CUDA_GRAPH=False +export RESET_CG_AFTER_HEALING=False + +export RECOMPUTE_GRANULARITY=null +export RECOMPUTE_METHOD=null +export RECOMPUTE_NUM_LAYERS=null + +export FP8_ACT=0 +export AITER_CONFIG_GEMM_A4W4="${PRIMUS_PATH}/examples/mlperf/llama2_70b/a4w4_tuned_gemms.csv" +export AITER_LOG_TUNED_CONFIG=0 +export NVTE_FP4_LOG_GEMM_SHAPES=0 +export AITER_LOG_LEVEL=ERROR +export AITER_LOG_MORE=0 diff --git a/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml b/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml new file mode 100644 index 000000000..7a929a43e --- /dev/null +++ b/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml @@ -0,0 +1,60 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:llama2_70b_lora_mlperf_posttrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + post_trainer: + framework: megatron_bridge + config: sft_trainer.yaml + + # MXFP4 LoRA recipe (primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.llama2_70b_lora_mxfp4_config) + model: llama2_70b_lora_mxfp4.yaml + + overrides: + pretrained_checkpoint: ${PRETRAINED_CHECKPOINT:/data/megatron_checkpoints/Llama-2-70b-hf} + stderr_sink_level: DEBUG + log_interval: 10 + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + pipeline_dtype: null + virtual_pipeline_model_parallel_size: null + context_parallel_size: 1 + sequence_parallel: false + use_megatron_fsdp: false + + peft: lora + dataset_type: mlperf_dataset + packed_sequence: true + packed_train_data_path: ${PACKED_DATA_DIR:/data}/train.npy + packed_val_data_path: ${PACKED_DATA_DIR:/data}/validation.npy + packed_metadata_path: ${PACKED_DATA_DIR:/data}/packed_metadata.jsonl + + train_iters: ${PRIMUS_TRAIN_ITERS:550} + global_batch_size: ${PRIMUS_GLOBAL_BATCH_SIZE:8} + micro_batch_size: ${PRIMUS_MICRO_BATCH_SIZE:1} + seq_length: 8192 + eval_interval: ${PRIMUS_EVAL_INTERVAL:48} + eval_iters: ${PRIMUS_EVAL_ITERS:24} + save_interval: null + + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-8 + weight_decay: 0.0001 + min_lr: 0.0 + lr: ${LR:0.0006} + lr_warmup_iters: ${PRIMUS_LR_WARMUP_ITERS:0} + lr_decay_iters: 550 + clip_grad: 0.3 + + comm_overlap_config: null + + seed: ${SEED:1234} + + enable_primus_turbo: false + use_turbo_attention: false + use_turbo_rms_norm: false + use_turbo_parallel_linear: false + + check_for_nan_in_loss: false diff --git a/examples/mlperf/llama2_70b/run_and_time.sh b/examples/mlperf/llama2_70b/run_and_time.sh new file mode 100755 index 000000000..e99711c69 --- /dev/null +++ b/examples/mlperf/llama2_70b/run_and_time.sh @@ -0,0 +1,53 @@ +#!/bin/bash +# One-shot MLPerf Llama2-70B LoRA via primus-cli (dataset + checkpoint hooks). +# +# MLPerf overrides live in primus/backends/megatron_bridge/patches/mlperf_llama2_70b/ and are +# applied automatically via @register_patch when llama2_70b_lora_mxfp4 is selected. +# +# Hooks (megatron_bridge) run automatically before training: +# 00_install_requirements.sh — pip deps +# 01_convert_checkpoints.sh — HF → Megatron checkpoint (needs HF_TOKEN) +# 02_prepare_mlperf_dataset.sh — SCROLLS gov-report .npy + metadata (needs HF_TOKEN) +# +# Usage (inside Primus container, repo root): +# export HF_TOKEN=... +# bash examples/mlperf/llama2_70b/run_and_time.sh +# +# Optional: +# PACKED_DATA_DIR=/data/mlperf_llama2 +# MLLOG_VERBOSE_LOGS=1 +# PRIMUS_LOG_GPU_MEM=1 # GPU mem every log_interval (default on in config_MI355X) +# SEED=1234 + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PRIMUS_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +cd "${PRIMUS_ROOT}" + +DATA_ROOT="${PACKED_DATA_DIR:-${DATA_PATH:-${PRIMUS_ROOT}/data/mlperf_llama2}}" +export PACKED_DATA_DIR="${DATA_ROOT}" +export DATA_PATH="${DATA_ROOT}" +export HF_HOME="${HF_HOME:-${DATA_ROOT}/.cache/huggingface}" +mkdir -p "${DATA_ROOT}" "${HF_HOME}" + +if [[ -z "${HF_TOKEN:-}" ]]; then + echo "[ERROR] HF_TOKEN is required (meta-llama/Llama-2-70b-hf + MLPerf dataset hub access)." >&2 + exit 1 +fi +export HF_TOKEN + +export SEED="${SEED:-$RANDOM}" + +# Exact MLPerf 6.0 MI355X env (MXFP4, AITER, NCCL, MLLOG, 550 iters, lr=0.0006, ...) +# shellcheck disable=SC1091 +source "${SCRIPT_DIR}/config_MI355X_1x8x1.sh" + +CONFIG="${CONFIG:-${EXP}}" + +echo "[INFO] Primus root: ${PRIMUS_ROOT}" +echo "[INFO] Data root: ${DATA_ROOT}" +echo "[INFO] HF cache: ${HF_HOME}" +echo "[INFO] Training config: ${CONFIG}" + +exec ./runner/primus-cli direct train posttrain --config "${CONFIG}" "$@" diff --git a/examples/mlperf/llama2_70b/run_in_container.sh b/examples/mlperf/llama2_70b/run_in_container.sh new file mode 100755 index 000000000..a3989b082 --- /dev/null +++ b/examples/mlperf/llama2_70b/run_in_container.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Run Llama 2 70B LoRA MLPerf (Megatron-Bridge, MI355X) via primus-cli. +# Dataset and model download are handled by posttrain hooks automatically. +# MLPerf overrides are applied at runtime from primus/backends/megatron_bridge/patches/mlperf_llama2_70b/. +# +# Usage (inside container): +# export HF_TOKEN="hf_..." +# bash /workspace/Primus/examples/mlperf/llama2_70b/run_in_container.sh +# +# Optional overrides: +# PACKED_DATA_DIR=/path/to/data +# MLLOG_VERBOSE_LOGS=1 + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +export PRIMUS_ROOT="${PRIMUS_ROOT:-$(cd "${SCRIPT_DIR}/../../.." && pwd)}" + +exec bash "${SCRIPT_DIR}/run_and_time.sh" "$@" 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 9b762b721..f78172ba8 100644 --- a/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py +++ b/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py @@ -21,6 +21,7 @@ from megatron.core.fusions.fused_layer_norm import FusedLayerNorm from megatron.core.models.backends import BackendSpecProvider from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear +from megatron.core.transformer.dot_product_attention import DotProductAttention from megatron.core.transformer.mlp import MLPSubmodules from megatron.core.transformer.moe.experts import ( SequentialMLP, @@ -104,8 +105,9 @@ def __init__( class PrimusTurboSpecProvider(BackendSpecProvider): """A protocol for providing the submodules used in Spec building.""" - def __init__(self): + def __init__(self, fallback_to_eager_attn: bool = False): self.cfg = get_primus_args() + self.fallback_to_eager_attn = fallback_to_eager_attn def linear(self) -> type: """Which linear module TE backend uses""" @@ -156,6 +158,8 @@ def layer_norm(self, rms_norm: bool = False, for_qk: bool = False) -> type: def core_attention(self) -> type: """Which module to use for attention""" + if self.fallback_to_eager_attn: + return DotProductAttention return ( _require_primus_turbo(PrimusTurboAttention, "attention") if self.cfg.use_turbo_attention diff --git a/primus/backends/megatron/patches/__init__.py b/primus/backends/megatron/patches/__init__.py index dd51eaf5e..e269a2a68 100644 --- a/primus/backends/megatron/patches/__init__.py +++ b/primus/backends/megatron/patches/__init__.py @@ -63,3 +63,6 @@ def _auto_import_patch_modules() -> None: # Eagerly import all patch modules on package import so patches are registered # before any backend-specific logic runs. _auto_import_patch_modules() + +# MLPerf Llama2-70B LoRA Megatron-LM overrides (MXFP4 recipe, optional TE SwiGLU). +import primus.backends.megatron_bridge.patches.mlperf_llama2_70b.megatron_patches # noqa: F401, E402 diff --git a/primus/backends/megatron/sft/dataset.py b/primus/backends/megatron/sft/dataset.py index a21b9c23a..cdd9f78d7 100644 --- a/primus/backends/megatron/sft/dataset.py +++ b/primus/backends/megatron/sft/dataset.py @@ -121,9 +121,10 @@ def build_train_valid_test_datasets( HF / jsonl tokenize+pack pipeline and route to ``MlperfPackedDataset``. This lets a Native SFT run consume the exact byte-identical packs produced by the upstream mlperf - ``download_dataset.py + convert_dataset.py + create_metadata.py`` - pipeline (used by ``examples/megatron_bridge/configs/MI355X/ - llama2_70b_lora_posttrain.yaml``). + ``primus.backends.megatron_bridge.recipes.mlperf_llama2_70b`` dataset utilities + (``download_dataset.py + convert_dataset.py + create_metadata.py``) + pipeline (used by ``examples/mlperf/llama2_70b/configs/MI355X/ + ``llama2_70b_lora_mlperf_posttrain.yaml``). """ from primus.backends.megatron.sft.mlperf_packed_dataset import ( build_mlperf_packed_datasets, diff --git a/primus/backends/megatron/sft/mlperf_packed_dataset.py b/primus/backends/megatron/sft/mlperf_packed_dataset.py index 5a0876d3a..18d59de14 100644 --- a/primus/backends/megatron/sft/mlperf_packed_dataset.py +++ b/primus/backends/megatron/sft/mlperf_packed_dataset.py @@ -9,7 +9,7 @@ This module lets the Megatron-native SFT backend consume the exact same ``train.npy`` / ``validation.npy`` / ``packed_metadata.jsonl`` artefacts that the upstream mlperf LLama-2-70B PEFT recipe ships through Megatron-Bridge -(``examples/megatron_bridge/configs/MI355X/llama2_70b_lora_posttrain.yaml``). +(``examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml``). Why a separate dataset class? ----------------------------- diff --git a/primus/backends/megatron_bridge/config_utils.py b/primus/backends/megatron_bridge/config_utils.py index 31c8b0090..e0eb0ee36 100644 --- a/primus/backends/megatron_bridge/config_utils.py +++ b/primus/backends/megatron_bridge/config_utils.py @@ -329,8 +329,10 @@ def _resolve_recipe(recipe: str, flavor: str): Resolve a recipe module and function by searching multiple namespaces. Search order: - 1. primus.backends.megatron_bridge.recipes.{recipe} (Primus-side extensions) - 2. megatron.bridge.recipes.{recipe} (upstream Megatron-Bridge) + 1. Direct custom module path (e.g., primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.llama2_custom) + 2. primus.backends.megatron_bridge.recipes.{recipe} (Primus-side extensions) + 3. megatron.bridge.recipes.{recipe} (upstream Megatron-Bridge) + 4. recipe as a direct Python module path (fallback) Returns: Tuple of (module, full_module_path) for the first namespace that @@ -339,6 +341,18 @@ def _resolve_recipe(recipe: str, flavor: str): Raises: AssertionError if the recipe cannot be found in any namespace. """ + custom_prefixes = ("primus.backends.megatron_bridge.recipes.", "primus.") + if any(recipe.startswith(prefix) for prefix in custom_prefixes): + try: + module = importlib.import_module(recipe) + except ImportError as e: + assert False, f"Recipe loading failed: Cannot import custom recipe '{recipe}': {e}" + assert hasattr( + module, flavor + ), f"Recipe loading failed: Function '{flavor}' not found in custom recipe '{recipe}'" + log_rank_0(f" ℹ️ Using custom recipe from: {recipe}") + return module, recipe + search_prefixes = [ "primus.backends.megatron_bridge.recipes", "megatron.bridge.recipes", @@ -353,8 +367,20 @@ def _resolve_recipe(recipe: str, flavor: str): if hasattr(module, flavor): return module, full_module_path + # Fallback: try recipe as a direct Python module path + if "." in recipe and not recipe.startswith(("/", ".")): + try: + module = importlib.import_module(recipe) + if hasattr(module, flavor): + log_rank_0(f" ℹ️ Using custom recipe from: {recipe}") + return module, recipe + except ImportError: + pass + # Build a helpful error message listing all paths that were tried. tried = [f"{p}.{recipe}" for p in search_prefixes] + if "." in recipe: + tried.append(recipe) assert False, f"Recipe loading failed: Function '{flavor}' not found. " f"Searched modules: {tried}" diff --git a/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py index f2943e381..545519161 100644 --- a/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py +++ b/primus/backends/megatron_bridge/megatron_bridge_base_trainer.py @@ -104,3 +104,39 @@ def detect_megatron_version(cls) -> str: "Please ensure Megatron-LM is properly installed and " "megatron.core.package_info is available." ) from e + + def _apply_nested_overrides(self) -> None: + """Apply flat backend_args overrides to nested ConfigContainer fields. + + ConfigContainer uses nested dataclasses (train, logger, checkpoint, etc.) + that cannot be reached by the flat _merge_dict_to_dataclass pass in + load_recipe_config. This bridges user-facing YAML keys (e.g. + ``log_interval: 99999``) to their nested targets. + """ + args = self.backend_args + cfg = self.cfg_container + + if hasattr(args, "log_interval"): + val = getattr(args, "log_interval") + if val is not None: + cfg.logger.log_interval = int(val) + log_rank_0(f" ↳ Override logger.log_interval = {cfg.logger.log_interval}") + + for key in ("eval_interval", "eval_iters"): + if hasattr(args, key): + val = getattr(args, key) + if val is not None: + setattr(cfg.train, key, int(val)) + log_rank_0(f" ↳ Override train.{key} = {val}") + + if hasattr(args, "save_interval"): + val = getattr(args, "save_interval") + if val is not None: + cfg.checkpoint.save_interval = int(val) + log_rank_0(f" ↳ Override checkpoint.save_interval = {val}") + + if hasattr(args, "skip_save") and args.skip_save: + cfg.checkpoint.save_interval = 0 + cfg.checkpoint.save = None + log_rank_0(" ↳ Override checkpoint.save_interval = 0 (skip periodic save)") + log_rank_0(" ↳ Override checkpoint.save = None (skip final save)") diff --git a/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py index dd9ed4a9e..924607935 100644 --- a/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py +++ b/primus/backends/megatron_bridge/megatron_bridge_posttrain_trainer.py @@ -79,6 +79,7 @@ def init(self): log_rank_0("Initializing Megatron-Bridge post-training components...") self.cfg_container = load_recipe_config(self.backend_args) + self._apply_nested_overrides() log_rank_0("Post-training initialization completed") @@ -100,10 +101,13 @@ def train(self): try: # Execute post-training based on configuration + from megatron.bridge.training.config import runtime_config_update from megatron.bridge.training.finetune import finetune from megatron.bridge.training.vlm_step import forward_step # log_rank_0(f"ConfigContainer: {self.cfg_container}") + runtime_config_update(self.cfg_container) + log_dict_aligned("ConfigContainer", self.cfg_container.to_dict()) finetune(self.cfg_container, forward_step_func=forward_step) diff --git a/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py b/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py index aa76e8daa..e88a4bcc4 100644 --- a/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py +++ b/primus/backends/megatron_bridge/megatron_bridge_pretrain_trainer.py @@ -72,45 +72,6 @@ def init(self): log_rank_0("Pre-training initialization completed") - def _apply_nested_overrides(self): - """Apply backend_args overrides to nested ConfigContainer fields. - - ConfigContainer uses nested dataclasses (train, logger, checkpoint, etc.) - that cannot be reached by the flat _merge_dict_to_dataclass pass. - This method bridges user-facing flat keys to their nested targets. - """ - args = self.backend_args - cfg = self.cfg_container - - # logger.* - if hasattr(args, "log_interval"): - val = getattr(args, "log_interval") - if val is not None: - cfg.logger.log_interval = int(val) - log_rank_0(f" ↳ Override logger.log_interval = {cfg.logger.log_interval}") - - # train.* - for key in ("eval_interval", "eval_iters"): - if hasattr(args, key): - val = getattr(args, key) - if val is not None: - setattr(cfg.train, key, int(val)) - log_rank_0(f" ↳ Override train.{key} = {val}") - - # checkpoint.* - if hasattr(args, "save_interval"): - val = getattr(args, "save_interval") - if val is not None: - cfg.checkpoint.save_interval = int(val) - log_rank_0(f" ↳ Override checkpoint.save_interval = {val}") - - # Skip final/in-training save entirely. - if hasattr(args, "skip_save") and args.skip_save: - cfg.checkpoint.save_interval = 0 - cfg.checkpoint.save = None - log_rank_0(" ↳ Override checkpoint.save_interval = 0 (skip periodic save)") - log_rank_0(" ↳ Override checkpoint.save = None (skip final save)") - def train(self): """ Execute Megatron-Bridge pre-training. diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/__init__.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/__init__.py new file mode 100644 index 000000000..b0d190202 --- /dev/null +++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Llama2-70B LoRA MLPerf overrides under ``megatron_bridge/patches/mlperf_llama2_70b/``.""" diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/bridge_patches.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/bridge_patches.py new file mode 100644 index 000000000..e42e91514 --- /dev/null +++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/bridge_patches.py @@ -0,0 +1,487 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Runtime Megatron-Bridge patches for MLPerf Llama2-70B LoRA. + +Replaces the former git patches under ``third_party/Megatron-Bridge`` without +writing to ``third_party/Megatron-Bridge``. +""" + +from __future__ import annotations + +import functools +import inspect +from typing import Any, Optional + +import torch + +from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.conditions import ( + is_llama2_70b_mlperf, +) +from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.resettable_data_iterator import ( + ResettableDataIterator, +) +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_rank_0 + +_PATCHED_ATTR = "_primus_mlperf_llama2_70b_patched" + + +def _mark_patched(obj: Any) -> None: + setattr(obj, _PATCHED_ATTR, True) + + +def _already_patched(obj: Any) -> bool: + return bool(getattr(obj, _PATCHED_ATTR, False)) + + +@register_patch( + "mlperf_llama2_70b.bridge.data_patches", + backend="megatron_bridge", + phase="before_train", + condition=is_llama2_70b_mlperf, + description="MLPerf data/eval/sampler overrides for Megatron-Bridge", +) +def patch_bridge_data(ctx: PatchContext) -> None: + import megatron.bridge.data.loaders as loaders + import megatron.bridge.data.samplers as samplers + import megatron.bridge.training.eval as bridge_eval + from megatron.core.rerun_state_machine import RerunDataIterator + + loaders.ResettableDataIterator = ResettableDataIterator + + import megatron.bridge.peft.lora as bridge_lora + + from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.lora import ( + LoRA, + VLMLoRA, + ) + + bridge_lora.LoRA = LoRA + bridge_lora.VLMLoRA = VLMLoRA + + if not _already_patched(loaders.get_train_valid_test_num_samples): + orig_num_samples = loaders.get_train_valid_test_num_samples + + @functools.wraps(orig_num_samples) + def _mlperf_get_train_valid_test_num_samples(cfg): + if cfg.train.train_samples is not None: + train_samples = cfg.train.train_samples + else: + train_samples = cfg.train.train_iters * cfg.train.global_batch_size + eval_iters = cfg.train.eval_iters + test_iters = cfg.train.eval_iters + return ( + train_samples, + eval_iters * cfg.train.global_batch_size, + test_iters * cfg.train.global_batch_size, + ) + + loaders.get_train_valid_test_num_samples = _mlperf_get_train_valid_test_num_samples + _mark_patched(_mlperf_get_train_valid_test_num_samples) + + if not _already_patched(samplers.build_pretraining_data_loader): + orig_build_loader = samplers.build_pretraining_data_loader + + @functools.wraps(orig_build_loader) + def _mlperf_build_pretraining_data_loader( + dataset, + consumed_samples, + dataloader_type, + micro_batch_size, + num_workers, + data_sharding, + worker_init_fn=None, + collate_fn=None, + pin_memory=True, + persistent_workers=False, + data_parallel_rank=0, + data_parallel_size=1, + drop_last=True, + global_batch_size=None, + eval_iters: Optional[int] = None, + name: str = "", + ): + if dataset is None: + return None + + if name == "validation": + if eval_iters is None or global_batch_size is None: + raise RuntimeError( + "eval_iters and global_batch_size must be provided when creating " + "a validation dataloader for MLPerf Llama2." + ) + eval_samples = eval_iters * global_batch_size + total_samples = min(len(dataset), eval_samples) + batch_sampler = samplers.MegatronPretrainingSampler( + total_samples=total_samples, + consumed_samples=0, + micro_batch_size=micro_batch_size, + data_parallel_rank=data_parallel_rank, + data_parallel_size=data_parallel_size, + ) + elif dataloader_type == "single": + batch_sampler = samplers.MegatronPretrainingSampler( + total_samples=len(dataset), + consumed_samples=consumed_samples, + micro_batch_size=micro_batch_size, + data_parallel_rank=data_parallel_rank, + data_parallel_size=data_parallel_size, + drop_last=drop_last, + ) + elif dataloader_type == "cyclic": + batch_sampler = samplers.MegatronPretrainingRandomSampler( + dataset, + total_samples=len(dataset), + consumed_samples=consumed_samples, + micro_batch_size=micro_batch_size, + data_parallel_rank=data_parallel_rank, + data_parallel_size=data_parallel_size, + data_sharding=data_sharding, + ) + elif dataloader_type == "batch": + if global_batch_size is None: + raise RuntimeError( + "global_batch_size must be provided when using dataloader_type='batch'." + ) + batch_sampler = samplers.MegatronPretrainingBatchSampler( + total_samples=len(dataset), + consumed_samples=consumed_samples, + micro_batch_size=micro_batch_size, + global_batch_size=global_batch_size, + data_parallel_rank=data_parallel_rank, + data_parallel_size=data_parallel_size, + drop_last=drop_last, + pad_samples_to_global_batch_size=not drop_last, + ) + elif dataloader_type == "external": + return dataset + else: + raise Exception(f"Unsupported dataloader_type: {dataloader_type}") + + return torch.utils.data.DataLoader( + dataset, + batch_sampler=batch_sampler, + num_workers=num_workers, + pin_memory=pin_memory, + persistent_workers=persistent_workers, + worker_init_fn=worker_init_fn, + collate_fn=collate_fn, + ) + + samplers.build_pretraining_data_loader = _mlperf_build_pretraining_data_loader + _mark_patched(_mlperf_build_pretraining_data_loader) + + if not _already_patched(loaders.build_train_valid_test_data_loaders): + orig_build_loaders = loaders.build_train_valid_test_data_loaders + + @functools.wraps(orig_build_loaders) + def _mlperf_build_train_valid_test_data_loaders( + cfg, train_state, build_train_valid_test_datasets_provider + ): + (train_dataloader, valid_dataloader, test_dataloader) = (None, None, None) + loaders.print_rank_0("> building train, validation, and test datasets ...") + train_ds, valid_ds, test_ds = loaders.build_train_valid_test_datasets( + cfg=cfg, + build_train_valid_test_datasets_provider=build_train_valid_test_datasets_provider, + ) + exit_signal = cfg.train.exit_signal + + from megatron.bridge.training.utils.sig_utils import ( + DistributedSignalHandler, + ) + + def worker_init_fn(_): + DistributedSignalHandler(exit_signal).__enter__() + + maybe_worker_init_fn = worker_init_fn if cfg.train.exit_signal_handler_for_dataloader else None + + train_dataloader = samplers.build_pretraining_data_loader( + train_ds, + train_state.consumed_train_samples, + cfg.dataset.dataloader_type, + cfg.train.micro_batch_size, + cfg.dataset.num_workers, + cfg.dataset.data_sharding, + worker_init_fn=maybe_worker_init_fn, + collate_fn=train_ds.collate_fn if hasattr(train_ds, "collate_fn") else None, + pin_memory=cfg.dataset.pin_memory, + persistent_workers=cfg.dataset.persistent_workers, + data_parallel_rank=loaders.mpu.get_data_parallel_rank(), + data_parallel_size=loaders.mpu.get_data_parallel_world_size(), + global_batch_size=cfg.train.global_batch_size, + name="train", + ) + if cfg.train.skip_train and cfg.train.eval_iters > 0: + valid_dataloader = samplers.build_pretraining_data_loader( + valid_ds, + 0, + cfg.dataset.dataloader_type, + cfg.train.micro_batch_size, + cfg.dataset.num_workers, + cfg.dataset.data_sharding, + worker_init_fn=maybe_worker_init_fn, + collate_fn=valid_ds.collate_fn if hasattr(valid_ds, "collate_fn") else None, + pin_memory=cfg.dataset.pin_memory, + persistent_workers=cfg.dataset.persistent_workers, + data_parallel_rank=loaders.mpu.get_data_parallel_rank(), + data_parallel_size=loaders.mpu.get_data_parallel_world_size(), + global_batch_size=cfg.train.global_batch_size, + name="validation", + eval_iters=cfg.train.eval_iters, + ) + elif cfg.train.eval_iters > 0: + val_dataloader_type = ( + "cyclic" + if isinstance(cfg.dataset, loaders.GPTDatasetConfig) + else cfg.dataset.dataloader_type + ) + valid_dataloader = samplers.build_pretraining_data_loader( + valid_ds, + train_state.consumed_valid_samples, + val_dataloader_type, + cfg.train.micro_batch_size, + cfg.dataset.num_workers, + cfg.dataset.data_sharding, + worker_init_fn=maybe_worker_init_fn, + collate_fn=valid_ds.collate_fn if hasattr(valid_ds, "collate_fn") else None, + pin_memory=cfg.dataset.pin_memory, + persistent_workers=cfg.dataset.persistent_workers, + data_parallel_rank=loaders.mpu.get_data_parallel_rank(), + data_parallel_size=loaders.mpu.get_data_parallel_world_size(), + global_batch_size=cfg.train.global_batch_size, + name="validation", + eval_iters=cfg.train.eval_iters, + ) + + if cfg.train.eval_iters > 0: + test_dataloader = samplers.build_pretraining_data_loader( + test_ds, + 0, + cfg.dataset.dataloader_type, + cfg.train.micro_batch_size, + cfg.dataset.num_workers, + cfg.dataset.data_sharding, + worker_init_fn=maybe_worker_init_fn, + collate_fn=test_ds.collate_fn if hasattr(test_ds, "collate_fn") else None, + pin_memory=cfg.dataset.pin_memory, + persistent_workers=cfg.dataset.persistent_workers, + data_parallel_rank=loaders.mpu.get_data_parallel_rank(), + data_parallel_size=loaders.mpu.get_data_parallel_world_size(), + global_batch_size=cfg.train.global_batch_size, + name="test", + ) + + do_train = train_dataloader is not None and cfg.train.train_iters > 0 + do_valid = valid_dataloader is not None and cfg.train.eval_iters > 0 + do_test = test_dataloader is not None and cfg.train.eval_iters > 0 + flags = torch.tensor( + [int(do_train), int(do_valid), int(do_test)], dtype=torch.long, device="cuda" + ) + torch.distributed.broadcast(flags, 0) + train_state.do_train = flags[0].item() + train_state.do_valid = flags[1].item() + train_state.do_test = flags[2].item() + return train_dataloader, valid_dataloader, test_dataloader + + loaders.build_train_valid_test_data_loaders = _mlperf_build_train_valid_test_data_loaders + _mark_patched(_mlperf_build_train_valid_test_data_loaders) + + if not _already_patched(loaders.build_train_valid_test_data_iterators): + orig_build_iterators = loaders.build_train_valid_test_data_iterators + + @functools.wraps(orig_build_iterators) + def _mlperf_build_train_valid_test_data_iterators( + cfg, train_state, build_train_valid_test_datasets_provider + ): + train_dataloader, valid_dataloader, test_dataloader = loaders.build_train_valid_test_data_loaders( + cfg=cfg, + train_state=train_state, + build_train_valid_test_datasets_provider=build_train_valid_test_datasets_provider, + ) + dl_type = cfg.dataset.dataloader_type + assert dl_type in ["single", "cyclic", "batch", "external"] + + def _get_iterator(dataloader_type, dataloader): + if dataloader_type == "single": + return RerunDataIterator(iter(dataloader)) + if dataloader_type in ("cyclic", "batch"): + return RerunDataIterator(iter(loaders.cyclic_iter(dataloader))) + if dataloader_type == "external": + if isinstance(dataloader, list): + return [RerunDataIterator(d) for d in dataloader] + return RerunDataIterator(dataloader) + raise RuntimeError("unexpected dataloader type") + + train_data_iterator = ( + _get_iterator(dl_type, train_dataloader) if train_dataloader is not None else None + ) + if valid_dataloader is not None: + valid_data_iterator = RerunDataIterator(ResettableDataIterator(valid_dataloader)) + else: + valid_data_iterator = None + test_data_iterator = ( + _get_iterator(dl_type, test_dataloader) if test_dataloader is not None else None + ) + return train_data_iterator, valid_data_iterator, test_data_iterator + + loaders.build_train_valid_test_data_iterators = _mlperf_build_train_valid_test_data_iterators + _mark_patched(_mlperf_build_train_valid_test_data_iterators) + + def _reset_data_iterator(data_iterator): + if data_iterator is None: + return + if isinstance(data_iterator, list): + for it in data_iterator: + _reset_data_iterator(it) + return + if isinstance(data_iterator, RerunDataIterator): + inner = data_iterator.iterable + if isinstance(inner, ResettableDataIterator): + inner.reset() + data_iterator.saved_microbatches.clear() + data_iterator.replaying = False + data_iterator.replay_pos = 0 + elif isinstance(data_iterator, ResettableDataIterator): + data_iterator.reset() + + if not _already_patched(bridge_eval.evaluate): + orig_evaluate = bridge_eval.evaluate + + @functools.wraps(orig_evaluate) + def _mlperf_evaluate(*args, **kwargs): + if len(args) >= 3: + _reset_data_iterator(args[2]) + else: + _reset_data_iterator(kwargs.get("data_iterator")) + return orig_evaluate(*args, **kwargs) + + bridge_eval.evaluate = _mlperf_evaluate + _mark_patched(_mlperf_evaluate) + + log_rank_0("[Patch:mlperf_llama2_70b.bridge.data_patches] Megatron-Bridge data patches applied") + + +@register_patch( + "mlperf_llama2_70b.bridge.sft_attention_mask", + backend="megatron_bridge", + phase="before_train", + condition=is_llama2_70b_mlperf, + description="Cache causal attention masks in GPTSFTDataset for MLPerf steady-state SFT", +) +def patch_sft_attention_mask(ctx: PatchContext) -> None: + from megatron.bridge.data.datasets import sft as sft_mod + + if _already_patched(sft_mod.GPTSFTDataset._create_attention_mask): + return + + orig_create_mask = sft_mod.GPTSFTDataset._create_attention_mask + + @functools.wraps(orig_create_mask) + def _cached_create_attention_mask(self, max_length): + cache = getattr(self, "_attention_mask_cache", None) + if cache is None: + cache = {} + object.__setattr__(self, "_attention_mask_cache", cache) + cached = cache.get(max_length) + if cached is not None: + return cached + attention_mask = torch.tril(torch.ones((max_length, max_length))).unsqueeze(0) + attention_mask = attention_mask < 0.5 + cache[max_length] = attention_mask + return attention_mask + + sft_mod.GPTSFTDataset._create_attention_mask = _cached_create_attention_mask + _mark_patched(_cached_create_attention_mask) + log_rank_0("[Patch:mlperf_llama2_70b.bridge.sft_attention_mask] SFT attention-mask cache applied") + + +@register_patch( + "mlperf_llama2_70b.bridge.training_log_nemo", + backend="megatron_bridge", + phase="before_train", + condition=is_llama2_70b_mlperf, + description="NeMo-style train_step timing support in Megatron-Bridge training_log", +) +def patch_training_log_nemo(ctx: PatchContext) -> None: + import megatron.bridge.training.utils.train_utils as train_utils + from megatron.bridge.training.utils import flop_utils + from megatron.bridge.utils.common_utils import get_world_size_safe, print_rank_0 + + original_training_log = train_utils.training_log + if original_training_log is None or _already_patched(original_training_log): + return + + sig = inspect.signature(original_training_log) + if "nemo_elapsed_time_per_iter_sec" in sig.parameters: + log_rank_0( + "[Patch:mlperf_llama2_70b.bridge.training_log_nemo] " + "training_log already supports NeMo timing; skipping wrapper" + ) + return + + @functools.wraps(original_training_log) + def _training_log_with_nemo( + loss_dict, + total_loss_dict, + learning_rate, + decoupled_learning_rate, + loss_scale, + report_memory_flag, + skipped_iter, + grad_norm, + params_norm, + num_zeros_in_grad, + config, + global_state, + history_wct, + model, + log_max_attention_logit=None, + nemo_elapsed_time_per_iter_sec=None, + ): + result = original_training_log( + loss_dict, + total_loss_dict, + learning_rate, + decoupled_learning_rate, + loss_scale, + report_memory_flag, + skipped_iter, + grad_norm, + params_norm, + num_zeros_in_grad, + config, + global_state, + history_wct, + model, + log_max_attention_logit, + ) + + if ( + nemo_elapsed_time_per_iter_sec is not None + and global_state.train_state.step % config.logger.log_interval == 0 + ): + batch_size = config.train.global_batch_size + num_flops = flop_utils.num_floating_point_operations(config, batch_size) + per_gpu_tf = num_flops / nemo_elapsed_time_per_iter_sec / get_world_size_safe() / 1e12 + elapsed = global_state.timers("interval-time").elapsed(barrier=True) + interval_ms = (elapsed / max(config.logger.log_interval, 1)) * 1000.0 + print_rank_0( + f"[TFLOP/s basis: NeMo-style train_step wall clock] " + f"Step time: {nemo_elapsed_time_per_iter_sec:.2f}s | " + f"model TFLOP/s/GPU: {per_gpu_tf:.1f}" + ) + print_rank_0( + f" NeMo-style time/iter (ms): {nemo_elapsed_time_per_iter_sec * 1000.0:.1f} |" + f" Megatron-Bridge interval-time/iter (ms): {interval_ms:.1f} |" + ) + + return result + + train_utils.training_log = _training_log_with_nemo + _mark_patched(_training_log_with_nemo) + log_rank_0("[Patch:mlperf_llama2_70b.bridge.training_log_nemo] NeMo timing wrapper applied") diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/conditions.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/conditions.py new file mode 100644 index 000000000..15909073c --- /dev/null +++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/conditions.py @@ -0,0 +1,43 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +from primus.core.patches import PatchContext, get_args + +_MLPERF_MODEL_TOKEN = "llama2_70b_lora_mxfp4" +_MLPERF_FLAVOR_TOKEN = "llama2_70b_lora_mxfp4_config" +_MLPERF_CONFIG_TOKEN = "llama2_70b_lora_mlperf_posttrain" + + +def is_llama2_70b_mlperf(ctx: PatchContext) -> bool: + """Return True when the active run targets MLPerf Llama2-70B LoRA.""" + model_name = str(ctx.model_name or "") + if _MLPERF_MODEL_TOKEN in model_name: + return True + + try: + args = get_args(ctx) + except AssertionError: + args = None + + if args is not None: + model = str(getattr(args, "model", "") or "") + if _MLPERF_MODEL_TOKEN in model: + return True + + recipe = str(getattr(args, "recipe", "") or "") + flavor = str(getattr(args, "flavor", "") or "") + if _MLPERF_FLAVOR_TOKEN in flavor and ( + "llama2_custom" in recipe or "recipes.mlperf_llama2_70b" in recipe + ): + return True + + primus_config = ctx.extra.get("primus_config") + if primus_config is not None: + config_file = str(getattr(primus_config, "config_file", "") or "") + if _MLPERF_CONFIG_TOKEN in config_file: + return True + + return False diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/lora.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/lora.py new file mode 100644 index 000000000..e89677678 --- /dev/null +++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/lora.py @@ -0,0 +1,188 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +NeMo-stable LoRA for MLPerf Llama2-70B. + +Default ``use_te_fused_lora=False`` keeps unfused :class:`LoRALinear` adapters so +loss matches the MLPerf reference path. +""" + +import logging +from dataclasses import dataclass, field +from typing import List, Literal, Optional + +import torch +import torch.nn as nn +import transformer_engine.pytorch as te +from megatron.bridge.peft.base import PEFT +from megatron.bridge.peft.lora_layers import ( + LinearAdapter, + LoRALinear, + LoRATopKRouter, + TEFusedLoRALinear, + TELinearAdapter, + patch_linear_module, +) +from megatron.bridge.peft.module_matcher import ModuleMatcher +from megatron.bridge.peft.utils import ( + ParallelLinearAdapter, + get_adapter_attributes_from_linear, + is_expert_linear, + wildcard_match, +) +from megatron.core import parallel_state +from megatron.core.transformer.moe.router import TopKRouter +from megatron.core.utils import unwrap_model + +logger = logging.getLogger(__name__) + +try: + import bitsandbytes + + HAVE_BNB = True +except ImportError: + HAVE_BNB = False + + +def _te_fused_lora_allowed_for_module( + full_name: str, + module_name: Optional[str], + include_modules: Optional[List[str]], + exclude_modules: List[str], +) -> bool: + """Return True if this FQN may use :class:`TEFusedLoRALinear`.""" + for pattern in exclude_modules: + if module_name == pattern or wildcard_match(pattern, full_name): + return False + if include_modules is None: + return True + for pattern in include_modules: + if module_name == pattern or wildcard_match(pattern, full_name): + return True + return False + + +@dataclass +class LoRA(PEFT, ModuleMatcher): + """LoRA with explicit control over TE fused adapters (MLPerf defaults to unfused).""" + + target_modules: List[str] = field( + default_factory=lambda: ["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"] + ) + dim: int = 32 + alpha: int = 32 + dropout: float = 0.0 + dropout_position: Literal["pre", "post"] = "pre" + lora_A_init_method: str = "xavier" + lora_B_init_method: str = "zero" + a2a_experimental: bool = False + lora_dtype: torch.dtype = None + use_te_fused_lora: bool = False + te_fused_lora_include_modules: Optional[List[str]] = None + te_fused_lora_exclude_modules: List[str] = field(default_factory=list) + + def transform( + self, module: nn.Module, name: Optional[str] = None, prefix: Optional[str] = None + ) -> nn.Module: + adapter_types = (LinearAdapter, LoRALinear, LoRATopKRouter, TELinearAdapter) + if isinstance(module, adapter_types): + return module + + if (ans := self.match(module, name, prefix)) is not None: + (match, full_name) = ans + if isinstance(module, nn.Linear) or (module.__class__ == te.Linear): + if hasattr(module.weight.data, "_local_tensor") or ( + HAVE_BNB + and getattr(module, "quant_state", None) is not None + and module.quant_state.__class__ == bitsandbytes.functional.QuantState + ): + lora_cls = patch_linear_module + elif module.__class__ == te.Linear: + lora_cls = TELinearAdapter + else: + lora_cls = LinearAdapter + + return lora_cls( + module, + dim=self.dim, + alpha=self.alpha, + dropout=self.dropout, + lora_A_init_method=self.lora_A_init_method, + lora_dtype=self.lora_dtype, + ) + + is_expert = is_expert_linear(full_name) + attrs = get_adapter_attributes_from_linear(module, is_expert=is_expert) + + enable_op_fuser = ( + self.use_te_fused_lora + and hasattr(module, "config") + and getattr(module.config, "use_transformer_engine_op_fuser", False) + and parallel_state.get_tensor_model_parallel_world_size() == 1 + and _te_fused_lora_allowed_for_module( + full_name, + name, + self.te_fused_lora_include_modules, + self.te_fused_lora_exclude_modules, + ) + ) + + logging.info(f"Adding lora to: {full_name}") + adapter = ParallelLinearAdapter( + attrs.in_features, + attrs.out_features, + self.dim, + base_linear_name=full_name, + activation="identity", + norm_type=None, + column_init_method=self.lora_A_init_method, + row_init_method=self.lora_B_init_method, + gather_output=False, + input_is_parallel=attrs.input_is_parallel, + dropout=self.dropout, + dropout_position=self.dropout_position, + model_parallel_config=getattr(module, "config", None), + alpha=self.alpha, + is_expert=is_expert, + a2a_experimental=self.a2a_experimental, + disable_tensor_parallel_comm=attrs.disable_tensor_parallel_comm, + disable_sequence_parallel_comm=attrs.disable_sequence_parallel_comm, + base_linear_is_parallel=attrs.base_linear_is_parallel, + ) + if isinstance(module, TopKRouter): + return LoRATopKRouter(module, adapter) + if enable_op_fuser: + return TEFusedLoRALinear(module, adapter) + return LoRALinear(module, adapter) + return module + + +@dataclass +class VLMLoRA(LoRA): + """VLM LoRA variant (re-exported for parity with upstream Megatron-Bridge).""" + + freeze_vision_model: bool = True + freeze_vision_projection: bool = True + freeze_language_model: bool = True + + def freeze_model(self, model: nn.Module, training: bool = True) -> None: + modules_to_freeze = [] + model = unwrap_model(model)[0] + if hasattr(model, "llava_model"): + model = model.llava_model + + if self.freeze_vision_model and hasattr(model, "vision_model"): + modules_to_freeze.append(model.vision_model) + if self.freeze_vision_projection and hasattr(model, "vision_projection"): + modules_to_freeze.append(model.vision_projection) + if self.freeze_language_model and hasattr(model, "language_model"): + modules_to_freeze.append(model.language_model) + + for module in modules_to_freeze: + module.eval() + for param in module.parameters(): + param.requires_grad = training is False diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/megatron_patches.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/megatron_patches.py new file mode 100644 index 000000000..997968c1d --- /dev/null +++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/megatron_patches.py @@ -0,0 +1,193 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Runtime Megatron-LM patches for MLPerf Llama2-70B LoRA MXFP4. + +Important: never delegate to Primus-Turbo ``fp4_utils.get_fp4_recipe`` here — +that function can return a ``(recipe, reason)`` tuple, which breaks TE when it +expects a recipe object with ``.mxfp4()`` / ``.mxfp8()`` methods. +""" + +from __future__ import annotations + +import enum +import functools +import os +from typing import Any + +from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.conditions import ( + is_llama2_70b_mlperf, +) +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_rank_0 + +_PATCHED_ATTR = "_primus_mlperf_llama2_70b_megatron_patched" +_mxfp4_phase = True + + +def _mark_patched(obj: Any) -> None: + setattr(obj, _PATCHED_ATTR, True) + + +def _already_patched(obj: Any) -> bool: + return bool(getattr(obj, _PATCHED_ATTR, False)) + + +def is_mxfp4_phase() -> bool: + return _mxfp4_phase + + +def set_mxfp4_phase(active: bool) -> None: + global _mxfp4_phase + _mxfp4_phase = active + + +def _upstream_get_fp4_recipe_handles_mxfp4() -> bool: + """True only when Megatron's get_fp4_recipe already builds MXFP4BlockScaling.""" + try: + import inspect + + from megatron.core import fp4_utils + + source = inspect.getsource(fp4_utils.get_fp4_recipe) + return "MXFP4BlockScaling" in source or "Fp4Recipe.mxfp4" in source + except Exception: + return False + + +def _build_mxfp4_get_fp4_recipe(orig_get_fp4_recipe): + import transformer_engine.common.recipe + from megatron.core.enums import Fp4Recipe + from megatron.core.fp8_utils import _get_custom_recipe + from megatron.core.transformer.transformer_config import TransformerConfig + from megatron.core.utils import is_te_min_version + + @functools.wraps(orig_get_fp4_recipe) + def _mlperf_get_fp4_recipe(config: TransformerConfig): + if config.fp4_recipe == Fp4Recipe.nvfp4: + if not is_te_min_version("2.7.0.dev0"): + raise ValueError("NVFP4BlockScaling requires TransformerEngine >= 2.7.0.dev0.") + fp4_recipe = transformer_engine.common.recipe.NVFP4BlockScaling() + elif config.fp4_recipe == Fp4Recipe.mxfp4: + if not is_te_min_version("2.8.0"): + raise ValueError("MXFP4BlockScaling requires TransformerEngine >= 2.8.0.") + fp4_recipe = transformer_engine.common.recipe.MXFP4BlockScaling() + fp4_recipe.use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1" + elif config.fp4_recipe == Fp4Recipe.custom: + fp4_recipe = _get_custom_recipe(config.fp4_quantizer_factory) + else: + raise ValueError(f"Unsupported FP4 recipe: {config.fp4_recipe}. Supported: nvfp4, mxfp4, custom.") + return fp4_recipe + + return _mlperf_get_fp4_recipe + + +@register_patch( + "mlperf_llama2_70b.megatron.fp4_mxfp4", + backend="megatron", + phase="before_train", + condition=is_llama2_70b_mlperf, + description="MXFP4 recipe + phase tracking for MLPerf Llama2 (single recipe object)", +) +def patch_fp4_mxfp4(ctx: PatchContext) -> None: + from megatron.core import enums, fp4_utils + + recipe_already_handles_mxfp4 = _upstream_get_fp4_recipe_handles_mxfp4() + + if not hasattr(enums.Fp4Recipe, "mxfp4"): + + class _MlperfFp4Recipe(str, enum.Enum): + nvfp4 = "nvfp4" + mxfp4 = "mxfp4" + custom = "custom" + + enums.Fp4Recipe = _MlperfFp4Recipe + log_rank_0("[Patch:mlperf_llama2_70b.megatron.fp4_mxfp4] Added Fp4Recipe.mxfp4") + + fp4_utils.is_mxfp4_phase = is_mxfp4_phase + fp4_utils.set_mxfp4_phase = set_mxfp4_phase + fp4_utils._mxfp4_phase = _mxfp4_phase + + if recipe_already_handles_mxfp4: + log_rank_0( + "[Patch:mlperf_llama2_70b.megatron.fp4_mxfp4] " + "Upstream get_fp4_recipe already supports mxfp4; phase tracking only" + ) + return + + if _already_patched(fp4_utils.get_fp4_recipe): + log_rank_0( + "[Patch:mlperf_llama2_70b.megatron.fp4_mxfp4] " "get_fp4_recipe already patched for MLPerf mxfp4" + ) + return + + fp4_utils.get_fp4_recipe = _build_mxfp4_get_fp4_recipe(fp4_utils.get_fp4_recipe) + _mark_patched(fp4_utils.get_fp4_recipe) + + log_rank_0("[Patch:mlperf_llama2_70b.megatron.fp4_mxfp4] MXFP4 get_fp4_recipe patch applied") + + +@register_patch( + "mlperf_llama2_70b.megatron.te_swiglu", + backend="megatron", + phase="before_train", + condition=is_llama2_70b_mlperf, + description="Optional TE SwiGLU path when USE_TE_SWIGLU=1", +) +def patch_te_swiglu(ctx: PatchContext) -> None: + if os.getenv("USE_TE_SWIGLU", "0") != "1": + log_rank_0("[Patch:mlperf_llama2_70b.megatron.te_swiglu] USE_TE_SWIGLU!=1; skipping") + return + + import transformer_engine.common.recipe + import transformer_engine.pytorch as te + import transformer_engine_torch as tex + from megatron.core.fusions import fused_bias_swiglu as swiglu_mod + + if _already_patched(swiglu_mod.SwiGLUFunction.forward): + return + + swiglu_mod.SwiGLUFunction.forward + + @staticmethod + @swiglu_mod.nvtx_decorator() + def _te_swiglu_forward(ctx, input_tensor, fp8_input_store, cpu_offload_input): + import torch + + ctx.fp8_input_store = fp8_input_store + ctx.ori_input_dtype = input_tensor.dtype + input_for_backward = input_tensor.to(torch.float8_e4m3fn) if fp8_input_store else input_tensor + if cpu_offload_input: + input_for_backward.activation_offloading = True + ctx.save_for_backward(input_for_backward) + + swiglu_op = te.ops.SwiGLU() + if fp8_input_store: + recipe = transformer_engine.common.recipe.DelayedScaling( + fp8_format=transformer_engine.common.recipe.Format.E4M3, + amax_history_len=8, + amax_compute_algo="max", + margin=2, + interval=1, + ) + with te.fp8_autocast(enabled=True, fp8_recipe=recipe): + return swiglu_op(input_tensor) + return swiglu_op(input_tensor) + + @staticmethod + @swiglu_mod.nvtx_decorator() + def _te_swiglu_backward(ctx, grad_output): + pass + + input_tensor = ctx.saved_tensors[0] + input_tensor = input_tensor.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input_tensor + return tex.dswiglu(grad_output, input_tensor, None), None, None + + swiglu_mod.SwiGLUFunction.forward = _te_swiglu_forward + swiglu_mod.SwiGLUFunction.backward = _te_swiglu_backward + _mark_patched(_te_swiglu_forward) + log_rank_0("[Patch:mlperf_llama2_70b.megatron.te_swiglu] TE SwiGLU patches applied") diff --git a/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/resettable_data_iterator.py b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/resettable_data_iterator.py new file mode 100644 index 000000000..a5773379f --- /dev/null +++ b/primus/backends/megatron_bridge/patches/mlperf_llama2_70b/resettable_data_iterator.py @@ -0,0 +1,35 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Resettable validation iterator for deterministic MLPerf eval.""" + + +class ResettableDataIterator: + """Iterator wrapper that restarts from the beginning of the dataloader on reset(). + + Unlike cyclic_iter which continuously cycles, this iterator explicitly resets + to produce the exact same sequence of batches on each reset(). This guarantees + deterministic validation: every evaluation pass sees identical data in identical + order, regardless of how many evaluations have occurred. + """ + + def __init__(self, dataloader): + self._dataloader = dataloader + self._iterator = iter(dataloader) + + def __iter__(self): + return self + + def __next__(self): + try: + return next(self._iterator) + except StopIteration: + self._iterator = iter(self._dataloader) + return next(self._iterator) + + def reset(self): + """Reset to the beginning of the dataloader for deterministic iteration.""" + self._iterator = iter(self._dataloader) diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/__init__.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/__init__.py new file mode 100644 index 000000000..31519705e --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""MLPerf Llama2-70B LoRA recipe and supporting modules.""" diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/_log_suppression.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/_log_suppression.py new file mode 100644 index 000000000..25db669ce --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/_log_suppression.py @@ -0,0 +1,224 @@ +# --------------------------------------------------------------------------- +# Non-MLLOG log suppression for Primus Llama2 SFT training pipeline. +# +# Adapted from small_llm_pretraining/primus/src/_log_suppression.py for use +# with the Llama2 70B LoRA SFT recipe (Megatron-Bridge + TransformerEngine +# + AITER). +# +# The MLPerf reference logs (":::MLLOG ...") plus training timing/result +# banners are the only lines we want on stdout. All other framework output +# (Megatron deprecations, TE warnings, AITER/hipblaslt/gloo C++ writes, +# torch.distributed noise, Primus internal INFO, ...) is suppressed by +# default to keep CI and submission logs clean. +# +# Control: +# MLLOG_VERBOSE_LOGS=1 -> restore the full verbose output +# MLLOG_VERBOSE_LOGS=0 -> quiet mode (default): only MLLOG + training +# timing/result lines are emitted. +# +# Strategy: +# 1. Export env vars that noisy libraries honour (AITER_LOG_LEVEL, +# PYTHONWARNINGS, TRANSFORMERS_VERBOSITY, ...). +# 2. Raise Python ``logging`` levels for the logger names emitting the +# noise (megatron.*, transformer_engine.*, torch.distributed, ...). +# 3. Install an FD-level line filter on stdout so native C++ writes +# (hipModuleLoad, hipblaslt latency, gloo peer-connect) are caught. +# Stderr is redirected to /dev/null in quiet mode. +# +# This module should be imported as early as possible in the recipe to +# suppress logs emitted during model construction and training. +# --------------------------------------------------------------------------- +import logging as _logging +import os as _os + +VERBOSE_LOGS = _os.environ.get("MLLOG_VERBOSE_LOGS", "0") == "1" + +QUIET_LOGGER_NAMES = ( + "megatron", + "megatron.core.utils", + "megatron.core.rerun_state_machine", + "transformer_engine", + "transformer_engine.aiter_rope", + "torch.distributed", + "torch.distributed.c10d_logger", + "primus", + "primus.cli", + "primus.backends", + "primus.modules", +) + + +def reapply_quiet_logger_levels() -> None: + """Raise levels on noisy Python loggers. Safe to call multiple times.""" + for name in QUIET_LOGGER_NAMES: + _logging.getLogger(name).setLevel(_logging.ERROR) + + +def _configure_env_and_loggers() -> None: + """Silence non-MLLOG sources controllable via env vars or logging.""" + _os.environ.setdefault("AITER_LOG_LEVEL", "ERROR") + _os.environ.setdefault("AITER_LOG_MORE", "0") + _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error") + _os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") + _os.environ.setdefault("PYTHONWARNINGS", "ignore") + + reapply_quiet_logger_levels() + + +def _install_fd_level_fallback_filter() -> None: + """Last-resort line filter for logs that bypass Python ``logging``. + + * FD 2 (stderr) -> left untouched unless MLLOG_SUPPRESS_STDERR=1 (/dev/null). + * FD 1 (stdout) -> pipe + reader thread that applies regex suppression. + ``:::MLLOG`` lines always pass through unconditionally. + """ + import re + import sys + import threading + + ansi_re = re.compile(r"\x1b\[[0-9;]*[ -/]*[@-~]") + + suppress_patterns = tuple( + re.compile(p) + for p in ( + # aiter C++ hipModuleLoad / hipModuleGetFunction banners + r"\[aiter\] hipModule", + r"\[aiter\] hipModuleGetFunction:", + r"^\[aiter\] import \[", + # hipblaslt latency warnings (ROCm) + r"Warning: Latency not found for MI_M=", + r"Returning latency value of 32 \(really slow\)", + r"Warning: Stream-K Data Parallel does not support GSU", + r"^libibverbs: Warning:", + r"^\[WARNING\] Field \"", + r"`torch_dtype` is deprecated", + r"^(?:\s*(?:BFloat8Float8_fnuz|Float8_fnuz|,\s*(?:MI_[MNK]|mi_input_type)=\d*|\d+)\s*)+\.?\s*$", + # Gloo C++ peer-connect banners + r"\[Gloo\] Rank \d+ is connected to ", + r"^Expected number of connected peer ranks is\s*:", + # PyTorch AccumulateGrad stream mismatch warning + r"AccumulateGrad node's stream does not match", + r"^\s*Variable\._execution_engine\.run_backward\(", + # torch.distributed c10d barrier warning + r"barrier\(\): using the device under current context", + # hipify preprocessed banner + r"^(?:\x1b\[[0-9;]*m)?Successfully preprocessed all matching files\.", + # Primus runner INFO/DEBUG banners (coloured) + r"^\[0;34m\[.*\]\s*\[INFO\]", + r"^\[0;32m\[.*\]\s*\[(INFO|SUCCESS)\]", + r"^\[1;33m\[.*\]\s*\[WARN\]", + r"^\[DEBUG\]", + # Primus rank-0 coloured INFO lines + r"^\[\[32m\d{8}\s+\d{2}:\d{2}:\d{2}\[0m\].*\[INFO\]", + # [MLPerf Train] framework startup banners + r"^\[MLPerf Train\]", + # Created path banners from multi-rank output directory creation + r"^Created path:", + # Orphan fragments from interleaved C++ writes + r"^\s*Success\s*$", + r"^\s*\d+\s*$", + # torchrun OMP_NUM_THREADS banner + r"Setting OMP_NUM_THREADS environment variable for each process", + r"^\*{5,}$", + r"^W\d{4}\s+\d{2}:\d{2}:\d{2}\.\d+\s+\d+\s+torch/distributed/run\.py", + # pip install noise + r"^\[notice\] A new release of pip", + r"^WARNING: The directory .* pip_cache", + r"^Requirement already satisfied:", + r"^Downloading\s", + r"^Installing collected packages:", + r"^Successfully installed\s", + r"^ERROR: pip's dependency resolver", + # Primus hook execution banners + r"^\[fix_aiter_asm_dir\]", + r"^\[cast_transpose_mxfp4", + r"^\[mxfp4-", + r"^\[te_fp4_debug\]", + r"^\[aiter-rope\]", + r"^\[megatron-te-swiglu\]", + r"^patching file\s", + r"^\[OK\]\s", + r"^\[\+\]\s", + # env.VAR= debug lines + r"^env\.\w+=", + r"^extra\.\w+=", + ) + ) + + def _should_suppress(line: str) -> bool: + stripped = ansi_re.sub("", line) + if ":::MLLOG" in stripped: + return False + if "RESULT," in stripped: + return False + if not stripped.strip(): + return True + for pat in suppress_patterns: + if pat.search(stripped): + return True + return False + + def _start_reader(read_fd: int, out_fd: int) -> None: + def _run() -> None: + buf = b"" + try: + while True: + chunk = _os.read(read_fd, 4096) + if not chunk: + break + buf += chunk + while b"\n" in buf: + raw, buf = buf.split(b"\n", 1) + line = raw.decode("utf-8", errors="replace") + if not _should_suppress(line): + _os.write(out_fd, raw + b"\n") + except Exception: + pass + finally: + if buf: + line = buf.decode("utf-8", errors="replace") + if not _should_suppress(line): + try: + _os.write(out_fd, buf) + except Exception: + pass + + threading.Thread(target=_run, daemon=True).start() + + sys.stdout.flush() + sys.stderr.flush() + + orig_stdout_fd = _os.dup(1) + + if _os.environ.get("MLLOG_SUPPRESS_STDERR", "0") == "1": + devnull_fd = _os.open(_os.devnull, _os.O_WRONLY) + _os.dup2(devnull_fd, 2) + _os.close(devnull_fd) + + stdout_r, stdout_w = _os.pipe() + _os.dup2(stdout_w, 1) + _os.close(stdout_w) + + _start_reader(stdout_r, orig_stdout_fd) + + sys.stdout = _os.fdopen(1, "w", buffering=1, closefd=False) + sys.stderr = _os.fdopen(2, "w", buffering=1, closefd=False) + + +_INSTALLED = False + + +def install() -> None: + """Install quiet-mode suppression exactly once per process.""" + global _INSTALLED + if _INSTALLED: + return + _INSTALLED = True + if VERBOSE_LOGS: + return + _configure_env_and_loggers() + _install_fd_level_fallback_filter() + + +# Auto-install on import. +install() diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py new file mode 100644 index 000000000..2a8f96c0f --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py @@ -0,0 +1,58 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import subprocess +import sys +from pathlib import Path + +_RECIPE_DIR = Path(__file__).resolve().parent +if str(_RECIPE_DIR) not in sys.path: + sys.path.insert(0, str(_RECIPE_DIR)) + +import numpy as np +import pandas as pd +from dataset_hash import hash_directory + + +def convert(data_dir, split): + df = pd.read_parquet(f"{data_dir}/{split}-00000-of-00001.parquet") + transformed_data = df.apply(lambda row: transform_row(row), axis=1).tolist() + np.save(f"{data_dir}/{split}", transformed_data) + + +def transform_row(row): + return { + "input_ids": row["input_ids"], + "loss_mask": [int(x != -100) for x in row["labels"]], + "seq_start_id": [0], + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Convert dataset script") + parser.add_argument("--data_dir", type=str, default="/data", help="The directory of the data files") + args = parser.parse_args() + + convert(args.data_dir, "train") + convert(args.data_dir, "validation") + + subprocess.run( + f"rm {args.data_dir}/train-00000-of-00001.parquet {args.data_dir}/validation-00000-of-00001.parquet", + shell=True, + executable="/bin/bash", + check=True, + ) + directory_hash = hash_directory(args.data_dir) + print(f"Succesfully converted dataset with hash {directory_hash}") diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/create_metadata.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/create_metadata.py new file mode 100644 index 000000000..3b3e193aa --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/create_metadata.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python3 +"""Create minimal metadata file for manually packed sequences.""" + +import argparse +import json +from pathlib import Path + + +def create_metadata(seq_length: int, output_path: str): + """Create minimal metadata for packed sequences. + + Args: + seq_length: The sequence length of your packed data + output_path: Path to save the metadata JSON file + """ + metadata = [ + { + "max_samples_per_bin": 1, + "dataset_max_seqlen": seq_length, + "min_packed_seqlen": seq_length, + } + ] + + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + + with open(output_path, "w") as f: + json.dump(metadata, f, indent=2) + + print(f"✓ Created metadata file: {output_path}") + print(" - max_samples_per_bin: 1") + print(f" - dataset_max_seqlen: {seq_length}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Create packed metadata for MLPerf Llama2-70B") + parser.add_argument("seq_length", type=int, help="Packed sequence length (e.g. 8192)") + parser.add_argument("output_path", type=str, help="Path to packed_metadata.jsonl") + args = parser.parse_args() + + create_metadata(args.seq_length, args.output_path) diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py new file mode 100644 index 000000000..7f53eecdc --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py @@ -0,0 +1,51 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import hashlib +import os +from concurrent.futures import ThreadPoolExecutor + + +def hash_file_md5(file_path, chunk_size=4194304): # Default chunk size 4MB. + """Hashes a single file using MD5 and returns the hex digest.""" + md5_hash = hashlib.md5() + try: + with open(file_path, "rb") as f: + while chunk := f.read(chunk_size): + md5_hash.update(chunk) + except FileNotFoundError: + return None + return md5_hash.hexdigest() + + +def hash_directory(path): + """Hashes all files in a directory in parallel using MD5.""" + hashes = [] + with ThreadPoolExecutor() as executor: + futures = [ + executor.submit(hash_file_md5, os.path.join(root, file)) + for root, dirs, files in os.walk(path) + for file in files + ] + for future in futures: + file_hash = future.result() + if file_hash: + hashes.append(file_hash) + + # Combine the individual file hashes into a single hash + combined_hash = hashlib.md5() + for file_hash in sorted(hashes): # Sort to maintain consistency + combined_hash.update(file_hash.encode()) + + return combined_hash.hexdigest() diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py new file mode 100644 index 000000000..0d885e07b --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py @@ -0,0 +1,51 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import subprocess +import sys +from pathlib import Path + +_RECIPE_DIR = Path(__file__).resolve().parent +if str(_RECIPE_DIR) not in sys.path: + sys.path.insert(0, str(_RECIPE_DIR)) + +from dataset_hash import hash_directory +from huggingface_hub import snapshot_download + +parser = argparse.ArgumentParser() +parser.add_argument("--data_dir", default="/data", type=str, help="Path to the dataset location") +args = parser.parse_args() + +snapshot_download( + "regisss/scrolls_gov_report_preprocessed_mlperf_2", + revision="21ff1233ee3e87bc780ab719c755170148aba1cb", + allow_patterns="*.parquet", + local_dir=args.data_dir, + local_dir_use_symlinks=False, + max_workers=16, + repo_type="dataset", +) +subprocess.run( + f"mv {args.data_dir}/data/* {args.data_dir}/ && find {args.data_dir} -mindepth 1 ! -name '*.parquet' -exec rm -rf {{}} +", + shell=True, + executable="/bin/bash", + check=True, +) + +directory_hash = hash_directory(args.data_dir) +assert ( + directory_hash == "682a5f40b790a56751bf8303554efc08" +), f"Expected hash 682a5f40b790a56751bf8303554efc08, but got {directory_hash}" +print(f"Succesfully downloaded and verified dataset with hash {directory_hash}") diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/llama2_custom.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/llama2_custom.py new file mode 100644 index 000000000..650b1e981 --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/llama2_custom.py @@ -0,0 +1,2056 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Custom Llama2 recipe for Primus. + +This is a custom recipe based on Megatron-Bridge's llama2.py recipe, +but placed in Primus for easier customization and extension. +""" + +import gc +import os +import sys +import time +from collections import deque +from datetime import timedelta +from typing import Any, Callable, List, Optional, Union + +import torch +from megatron.core.full_cuda_graph import FullCudaGraphWrapper +from megatron.core.num_microbatches_calculator import ( + get_current_global_batch_size, + get_current_running_global_batch_size, + get_num_microbatches, +) +from megatron.core.optimizer import MegatronOptimizer +from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler +from megatron.core.parallel_state import update_pg_timeout +from megatron.core.pipeline_parallel import get_forward_backward_func +from megatron.core.pipeline_parallel.utils import is_pp_last_stage +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.rerun_state_machine import ( + RerunDataIterator, + RerunMode, + get_rerun_state_machine, +) +from megatron.core.transformer import MegatronModule +from megatron.core.transformer.cuda_graphs import TECudaGraphHelper +from megatron.core.transformer.enums import AttnBackend +from megatron.core.utils import check_param_hashes_across_dp_replicas, get_model_config +from typing_extensions import TypedDict, Unpack + +from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b import ( # noqa: F401, E402 + _log_suppression, +) +from primus.core.utils.module_utils import log_rank_0 as _orig_log_rank_0 +from primus.core.utils.module_utils import log_rank_last as _orig_log_rank_last + +_log_suppression.reapply_quiet_logger_levels() + +_verbose_logging = os.environ.get("VERBOSE_TRAINING_LOG", "0") == "1" +if _verbose_logging: + log_rank_0 = _orig_log_rank_0 + log_rank_last = _orig_log_rank_last +else: + + def log_rank_0(*args, **kwargs): + pass + + def log_rank_last(*args, **kwargs): + pass + + +from megatron.bridge.utils import common_utils + +# Megatron-Bridge training_log prints iteration / TFLOP/s / loss via print_rank_0 +# (plain stdout). Only silence Primus log_rank_0 helpers — not Megatron metrics. +_megatron_print_rank_0 = common_utils.print_rank_0 +_megatron_print_rank_last = common_utils.print_rank_last +_megatron_warn_rank_0 = common_utils.warn_rank_0 +common_utils.print_rank_0 = _megatron_print_rank_0 +common_utils.print_rank_last = _megatron_print_rank_last +if _verbose_logging: + common_utils.warn_rank_0 = _megatron_warn_rank_0 +else: + common_utils.warn_rank_0 = log_rank_0 + + +def _truthy_env(name: str, default: bool = False) -> bool: + raw = os.getenv(name) + if raw is None: + return default + return raw.strip().lower() in ("1", "true", "yes", "on") + + +def _log_training_gpu_mem(tag: str, memory_keys=None) -> None: + """Rank-0 CUDA memory snapshot when ``PRIMUS_LOG_GPU_MEM=1``.""" + if not _truthy_env("PRIMUS_LOG_GPU_MEM", default=False): + return + try: + from megatron.bridge.training.utils.train_utils import report_memory + + if not torch.cuda.is_available(): + _orig_log_rank_0(f"[GPU mem] {tag} | CUDA not available") + return + torch.cuda.synchronize() + dev = torch.cuda.current_device() + alloc = torch.cuda.memory_allocated(dev) / (1024**3) + reserved = torch.cuda.memory_reserved(dev) / (1024**3) + peak_alloc = torch.cuda.max_memory_allocated(dev) / (1024**3) + peak_reserved = torch.cuda.max_memory_reserved(dev) / (1024**3) + _orig_log_rank_0( + f"[GPU mem] {tag} | " + f"allocated={alloc:.2f} GiB reserved={reserved:.2f} GiB " + f"max_alloc={peak_alloc:.2f} GiB max_reserved={peak_reserved:.2f} GiB" + ) + mem = report_memory(memory_keys) + if mem: + detail = " | ".join(f"{k}={v}" for k, v in sorted(mem.items())) + _orig_log_rank_0(f"[GPU mem detail] {tag} | {detail}") + except Exception as exc: + _orig_log_rank_0(f"[GPU mem] {tag} | failed ({type(exc).__name__}: {exc})") + + +from megatron.bridge import AutoBridge +from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs +from megatron.bridge.data.finetuning import prepare_finetuning_batch +from megatron.bridge.data.iterator_utils import make_data_iterator_list +from megatron.bridge.recipes.utils.finetune_utils import default_squad_config +from megatron.bridge.recipes.utils.tokenizer_utils import ( + DEFAULT_NULL_TOKENIZER_VOCAB_SIZE, +) +from megatron.bridge.training import fault_tolerance +from megatron.bridge.training.checkpointing import maybe_finalize_async_save +from megatron.bridge.training.comm_overlap import CommOverlapConfig +from megatron.bridge.training.config import ( + CheckpointConfig, + ConfigContainer, + DistributedDataParallelConfig, + FinetuningDatasetConfig, + LoggerConfig, + ProfilingConfig, + RerunStateMachineConfig, + RNGConfig, + StragglerDetectionConfig, + TensorInspectConfig, + TokenizerConfig, + TrainingConfig, +) +from megatron.bridge.training.forward_step_func_types import ForwardStepCallable +from megatron.bridge.training.mixed_precision import ( + MixedPrecisionConfig, + bf16_mixed, + register, +) +from megatron.bridge.training.nvrx_straggler import safe_shutdown_nvrx_straggler_manager +from megatron.bridge.training.profiling import ( + handle_profiling_step, + handle_profiling_stop, + initialize_pytorch_profiler, + should_profile_rank, +) +from megatron.bridge.training.state import GlobalState +from megatron.bridge.training.tensor_inspect import ( + tensor_inspect_end_if_enabled, + tensor_inspect_step_if_enabled, +) +from megatron.bridge.training.utils import flop_utils +from megatron.bridge.training.utils import train_utils as _megatron_train_utils +from megatron.bridge.training.utils.pg_utils import get_pg_collection +from megatron.bridge.training.utils.train_utils import ( + calc_params_l2_norm, + prepare_forward_step_func, + training_log, +) + +from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.lora import LoRA +from primus.backends.megatron_bridge.patches.mlperf_llama2_70b.resettable_data_iterator import ( + ResettableDataIterator, +) + +# train_utils binds print_rank_* at import time; rebind so training_log always uses Primus loggers even if +# this module was imported after an earlier train_utils load. +_megatron_train_utils.print_rank_0 = log_rank_0 +_megatron_train_utils.print_rank_last = log_rank_last +from megatron.bridge.training.train import ( + _delete_cuda_graphs, + disable_forward_pre_hook, + enable_forward_pre_hook, + maybe_check_weight_hash_across_dp_replicas, + maybe_report_stragglers, + maybe_run_manual_gc, + maybe_synchronize_training_step, + should_disable_forward_pre_hook, + train_step, +) +from megatron.bridge.utils.common_utils import is_last_rank + +# Importing this module installs a monkey-patch that swaps Megatron-Bridge's +from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b import ( # noqa: F401 + nemo_loss as _nemo_loss, +) + +MLPERF_TARGET_LOSS = 0.925 + +# Sticky flag: once ``evaluate_and_print_results_custom`` observes an lm-loss +# value strictly below ``MLPERF_TARGET_LOSS`` we flip this to True. Subsequent +# calls to ``evaluate_and_print_results_custom`` / ``warmup_eval`` become +# no-ops that return ``(should_exit=True, None)`` without running another full +# validation pass. Prevents redundant evals from a resumed loop or any outer +# re-entry path after the target is met. +_TARGET_LOSS_REACHED: bool = False + +# --------------------------------------------------------------------------- +# MLPerf logging singleton (initialised lazily in the recipe config function) +# --------------------------------------------------------------------------- +_sft_logger = None + + +def _get_sft_logger(): + return _sft_logger + + +@register +def bf16_with_fp8_hybrid() -> MixedPrecisionConfig: + """Create a MixedPrecisionConfig for mixed precision training using BF16 with MXFP8. + + Returns: + MixedPrecisionConfig: Configuration for BF16 with MXFP8 mixed precision training + """ + cfg = bf16_mixed() + cfg.fp8 = "hybrid" + cfg.fp8_recipe = "delayed" + cfg.fp8_amax_history_len = 4 + cfg.fp8_amax_compute_algo = "most_recent" + cfg.fp8_param_gather = True + return cfg + + +@register +def bf16_with_mxfp4_mixed() -> MixedPrecisionConfig: + """BF16 + MXFP4 (e2m1) mixed precision, registered here so it resolves when + ``precision_config="bf16_with_mxfp4_mixed"`` is passed from YAML before + ``runtime_config_update`` runs. + """ + cfg = bf16_mixed() + cfg.fp8 = None + cfg.fp4 = "e2m1" + cfg.fp4_recipe = "mxfp4" + cfg.fp8_recipe = "delayed" + cfg.fp8_amax_history_len = 4 + cfg.fp8_amax_compute_algo = "most_recent" + cfg.fp8_reduce_amax = False + cfg.fp8_interval = 1 + cfg.fp8_margin = 0 + cfg.fp8_dot_product_attention = False + cfg.fp8_param_gather = False + cfg.grad_reduce_in_fp32 = False + return cfg + + +class Timer: + def __init__(self, gbs): + self.start_time = None + self.stop_time = None + self.elapsed_time = 0 + self.samples = 0 + self.gbs = gbs + self.consumed_samples = 0 + + def start(self): + self.start_time = time.time() + + def stop(self): + self.stop_time = time.time() + self.samples += self.gbs + self.consumed_samples += self.gbs + self.elapsed_time += self.stop_time - self.start_time + + def get_throughput(self): + throughput = self.samples / self.elapsed_time + self.samples = 0 + self.elapsed_time = 0 + return throughput + + +class Llama2CustomKwargs(TypedDict, total=False): + """Typed options accepted by custom Llama2 recipe helper functions.""" + + # Core identifiers + hf_path: str + dir: Optional[str] + name: str + # Dataset configuration + data_paths: Optional[List[str]] + data_args_path: Optional[str] + train_data_path: Optional[List[str]] + valid_data_path: Optional[List[str]] + test_data_path: Optional[List[str]] + per_split_data_args_path: Optional[str] + mock: bool + # Model configuration + tensor_model_parallel_size: int + pipeline_model_parallel_size: int + pipeline_dtype: Optional[torch.dtype] + virtual_pipeline_model_parallel_size: Optional[int] + context_parallel_size: int + sequence_parallel: bool + use_megatron_fsdp: bool + # Training hyperparameters + train_iters: int + global_batch_size: int + micro_batch_size: int + seq_length: int + lr: float + min_lr: float + lr_warmup_iters: int + lr_decay_iters: Optional[int] + eval_interval: int + save_interval: int + use_null_tokenizer: bool + # Precision / overlap configs + precision_config: Optional[Union[MixedPrecisionConfig, str]] + comm_overlap_config: Optional[CommOverlapConfig] + adam_beta1: float = 0.9 + adam_beta2: float = 0.99 + adam_eps: float = 1e-8 + weight_decay: float = 0.0001 + eval_iters: int = 22 + clip_grad: float = 0.3 + pretrained_checkpoint: str | None + packed_sequence: bool + packed_train_data_path: str | None + packed_val_data_path: str | None + packed_metadata_path: str | None + dataset_type: str + seed: int + check_for_nan_in_loss: bool + te_fused_lora_include_modules: Optional[List[str]] + te_fused_lora_exclude_modules: Optional[List[str]] + # Optional MXFP4-phase activation recompute (left None outside the MXFP4 recipe). + recompute_granularity: Optional[str] + recompute_method: Optional[str] + recompute_num_layers: Optional[int] + # TE attention backend ("flash", "fused", "unfused", "local", "auto"). None keeps Megatron's default. + attention_backend: Optional[str] + + +def llama2_70b_lora_config(**user_kwargs: Unpack[Llama2CustomKwargs]) -> ConfigContainer: + """ + Return a custom pre-training config for Llama-2 70B. + + This is a custom variant that can be modified without changing Megatron-Bridge code. + See `_llama2_lora` for the full list of parameters. + """ + recommended_kwargs: Llama2CustomKwargs = { + "hf_path": "meta-llama/Llama-2-70b-hf", + "tensor_model_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "train_iters": 1000, + "global_batch_size": 8, + "micro_batch_size": 1, + "eval_interval": 48, + "eval_iters": 22, + "adam_beta1": 0.9, + "adam_beta2": 0.999, + "adam_eps": 1e-8, + "weight_decay": 0.0001, + "clip_grad": 0.3, + } + # Combine defaults with user kwargs; user values take precedence. + combined_kwargs: Llama2CustomKwargs = {**recommended_kwargs, **user_kwargs} + cfg = _llama2_lora(**combined_kwargs) + + # --- MLPerf logging: init phase --- + if os.getenv("ENABLE_MLLOG", "0") == "1": + global _sft_logger + try: + from primus_mllog import MLPerfSFTLogger + + kw = combined_kwargs + gbs = kw.get("global_batch_size", 8) + mbs = kw.get("micro_batch_size", 1) + _sft_logger = MLPerfSFTLogger( + global_batch_size=gbs, + micro_batch_size=mbs, + ) + _sft_logger.log_cache_clear_and_init_start() + + data_root = os.getenv("DATA_PATH", "/data") + world_size = int(os.environ.get("WORLD_SIZE", "1")) + tp = kw.get("tensor_model_parallel_size", 1) + pp = kw.get("pipeline_model_parallel_size", 1) + dp_size = world_size // (tp * pp) + init_cfg = MLPerfSFTLogger.extract_sft_configs( + train_gbs=gbs, + train_mbs=mbs, + train_iters=kw.get("train_iters", 1000), + eval_iters=kw.get("eval_iters", 22), + seq_length=kw.get("seq_length", 8192), + seed=kw.get("seed", int(os.getenv("SEED", "1234"))), + lr=kw.get("lr", float(os.getenv("LR", "0.0004"))), + weight_decay=kw.get("weight_decay", 0.0001), + clip_grad=kw.get("clip_grad", 0.3), + lr_warmup_iters=kw.get("lr_warmup_iters", 0), + adam_beta1=kw.get("adam_beta1", 0.9), + adam_beta2=kw.get("adam_beta2", 0.999), + adam_eps=kw.get("adam_eps", 1e-8), + lora_rank=16, + lora_alpha=32, + data_root=data_root, + data_parallel_size=dp_size, + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + context_parallel_size=int(os.getenv("MLLOG_CONTEXT_PARALLELISM", "1")), + config_filename=os.getenv("MLLOG_CONFIG_FILENAME", ""), + lowest_numerical_precision_linear=os.getenv( + "MLLOG_LOWEST_NUMERICAL_PRECISION_LINEAR", "mxfp4" + ), + ) + _sft_logger.log_init_params(init_cfg) + except ImportError: + _orig_log_rank_0("primus_mllog not installed — MLPerf logging disabled") + _sft_logger = None + except Exception as exc: + _orig_log_rank_0(f"MLPerf logging init failed ({type(exc).__name__}: {exc}) — disabled") + _sft_logger = None + + return cfg + + +def llama2_70b_lora_mxfp4_config(**user_kwargs: Unpack[Llama2CustomKwargs]) -> ConfigContainer: + """Llama-2 70B LoRA with MXFP4 mixed precision (``bf16_with_mxfp4_mixed``). + + Matches MLPerf 6.0 NeMo MI355X FP4 defaults: full/block recompute over 8 + layers; FusedAttention backend (CK / AOTriton on ROCm) instead of + FlashAttention (which is not validated for the MXFP4 path on MI355X). + All defaults are user-overridable via ``user_kwargs``. + """ + mxfp4_defaults: Llama2CustomKwargs = { + "precision_config": "bf16_with_mxfp4_mixed", + "recompute_granularity": None, + "recompute_method": None, + "recompute_num_layers": None, + "attention_backend": "fused", + } + combined_kwargs: Llama2CustomKwargs = {**mxfp4_defaults, **user_kwargs} + return llama2_70b_lora_config(**combined_kwargs) + + +def _llama2_lora( + hf_path: str, + dir: Optional[str] = None, + name: str = "default", + # Model configuration + tensor_model_parallel_size: int = 1, + pipeline_model_parallel_size: int = 1, + virtual_pipeline_model_parallel_size: Optional[int] = None, + context_parallel_size: int = 1, + sequence_parallel: bool = False, + # Training hyperparameters + train_iters: int = 1000, + global_batch_size: int = 8, + micro_batch_size: int = 1, + seq_length: int = 8192, + lr: float = 4e-4, + min_lr: float = 0.0, + lr_warmup_iters: int = 0, + lr_decay_iters: Optional[int] = None, + eval_interval: int = 48, + eval_iters: int = 22, + use_null_tokenizer: bool = False, + pretrained_checkpoint: str | None = None, + packed_sequence: bool = False, + packed_train_data_path: str | None = None, + packed_val_data_path: str | None = None, + packed_metadata_path: str | None = None, + dataset_type: str = "mlperf_dataset", + adam_beta1: float = 0.9, + adam_beta2: float = 0.999, + adam_eps: float = 1e-8, + weight_decay: float = 0.0001, + clip_grad: float = 0.3, + # Precision recipe + precision_config: Optional[Union[MixedPrecisionConfig, str]] = "bf16_with_mxfp4_mixed", + comm_overlap_config: Optional[CommOverlapConfig] = None, + seed: int = 1234, + check_for_nan_in_loss: bool = False, + te_fused_lora_include_modules: Optional[List[str]] = None, + te_fused_lora_exclude_modules: Optional[List[str]] = None, + recompute_granularity: Optional[str] = None, + recompute_method: Optional[str] = None, + recompute_num_layers: Optional[int] = None, + attention_backend: Optional[str] = None, +) -> ConfigContainer: + """ + Create a custom pre-training configuration for Llama2 models. + + This is based on Megatron-Bridge's llama2 recipe but can be customized + for Primus-specific needs. + + Args: + hf_path (str): HuggingFace model path (e.g., "meta-llama/Llama-2-70b-hf"). + dir (Optional[str]): Base directory for saving logs and checkpoints. + name (str): Name of the pre-training run. + data_paths (Optional[List[str]]): List of paths to dataset files. If None, mock data will be used. + data_args_path (Optional[str]): Path to file containing data arguments. + train_data_path (Optional[List[str]]): List of training data paths. + valid_data_path (Optional[List[str]]): List of validation data paths. + test_data_path (Optional[List[str]]): List of test data paths. + per_split_data_args_path (Optional[str]): Path to JSON file with per-split data configuration. + mock (bool): Whether to use mock data. If True, ignores data_paths. + tensor_model_parallel_size (int): Degree of tensor model parallelism. + pipeline_model_parallel_size (int): Degree of pipeline model parallelism. + pipeline_dtype (Optional[torch.dtype]): Data type for pipeline parallelism. + virtual_pipeline_model_parallel_size (Optional[int]): Size of virtual pipeline parallelism. + context_parallel_size (int): Degree of context parallelism to be passed to model_config. + sequence_parallel (bool): Whether to use sequence parallelism. + use_megatron_fsdp (bool): Whether to use Megatron FSDP. + train_iters (int): Total number of training iterations. + global_batch_size (int): Global batch size for training. + micro_batch_size (int): Micro batch size for training. + seq_length (int): Sequence length for training data. + lr (float): Learning rate. + min_lr (float): Minimum learning rate for cosine decay. + lr_warmup_iters (int): Number of warmup iterations for the learning rate. + lr_decay_iters (Optional[int]): Number of iterations over which to decay the LR. + eval_interval (int): Evaluation interval. + save_interval (int): Save interval. + precision_config (Optional[Union[MixedPrecisionConfig, str]]): Precision configuration for the model. + comm_overlap_config (Optional[CommOverlapConfig]): Communication overlap configuration for the model. + adam_beta1 (float): Beta1 parameter for Adam optimizer. + adam_beta2 (float): Beta2 parameter for Adam optimizer. + adam_eps (float): Epsilon parameter for Adam optimizer. + weight_decay (float): Weight decay parameter for Adam optimizer. + eval_iters (int): Number of iterations to run for evaluation validation/test for. + pretrained_checkpoint (str | None): Path to pretrained checkpoint to load. + peft (str | PEFT | None): PEFT configuration (e.g., "lora" or LoRA object). + packed_sequence (bool): Whether to use packed sequences. + packed_train_data_path (str | None): Path to packed training data. + packed_val_data_path (str | None): Path to packed validation data. + packed_metadata_path (str | None): Path to packed metadata. + dataset_type (str): Dataset type to use. Either "squad" (default) or "mlperf_dataset". + use_transformer_engine_op_fuser (bool): If True, set ``model_cfg.use_transformer_engine_op_fuser`` + (TE op-fuser path on the backbone, e.g. fused MLP). Set False if LM/validation loss diverges + from a known-good run. + stable_lora_with_te_op_fuser (bool): Single Primus knob for the **stable LoRA + op fuser** combo. + If True (default): backbone follows ``use_transformer_engine_op_fuser``, but LoRA always + uses unfused :class:`LoRALinear` (``use_te_fused_lora=False``) so loss matches the safe path. + If False: **legacy** behavior — LoRA uses ``use_te_fused_lora = use_transformer_engine_op_fuser`` + (when TP=1, fused :class:`TEFusedLoRALinear` tracks backbone op fuser, as in older Bridge). + te_fused_lora_include_modules (Optional[List[str]]): Passed to :class:`LoRA`; only applies when + fused LoRA is enabled (``stable_lora_with_te_op_fuser=False`` and backbone op fuser on). + te_fused_lora_exclude_modules (Optional[List[str]]): Passed to :class:`LoRA`; same scope as include. + + Returns: + ConfigContainer: Configuration for pre-training. + """ + from transformers import AutoConfig + + config = AutoConfig.from_pretrained("meta-llama/Llama-2-70b-hf") + bridge = AutoBridge.from_hf_config(config) + model_cfg = bridge.to_megatron_provider(load_weights=False) # GPTProvider + model_cfg.tensor_model_parallel_size = tensor_model_parallel_size + model_cfg.pipeline_model_parallel_size = pipeline_model_parallel_size + model_cfg.pipeline_dtype = None + model_cfg.virtual_pipeline_model_parallel_size = virtual_pipeline_model_parallel_size + model_cfg.context_parallel_size = context_parallel_size + model_cfg.sequence_parallel = sequence_parallel + model_cfg.seq_length = seq_length + model_cfg.perform_initialization = True + model_cfg.fp16 = False + model_cfg.bf16 = True + model_cfg.params_dtype = torch.bfloat16 + model_cfg.autocast_dtype = torch.bfloat16 + model_cfg.pipeline_dtype = torch.bfloat16 + # Fusions / CE: TE parallel cross-entropy; grad-acc fusion (not used with Megatron FSDP). + model_cfg.cross_entropy_loss_fusion = False + model_cfg.cross_entropy_fusion_impl = "native" + model_cfg.gradient_accumulation_fusion = False + model_cfg.bias_dropout_fusion = True + model_cfg.fused_single_qkv_rope = False + model_cfg.apply_rope_fusion = True + model_cfg.use_transformer_engine_op_fuser = False + # Activation offload hint for TP paths (distinct from cpu_offloading / cpu_offloading_num_layers). + model_cfg.cpu_offloading_activations = True + + # MXFP4 weights (fp4/e2m1); fp8=None during MXFP4 phase. FP8_* env vars in + # config_MI355X_1x8x1.sh configure healing (step 340) and TE delayed scaling, + # not Megatron model_cfg.fp8 (which stays None until healing). + model_cfg.fp8 = None + model_cfg.fp4 = "e2m1" + model_cfg.fp4_recipe = "mxfp4" + model_cfg.fp8_param_gather = False + model_cfg.grad_reduce_in_fp32 = False + + # Used when cuda_graph_impl is not "none" (harmless when graphs are disabled). + model_cfg.cuda_graph_retain_backward_graph = True + model_cfg.cuda_graph_use_single_mempool = True + model_cfg.fp8_recipe = "delayed" + model_cfg.fp8_amax_history_len = 4 + model_cfg.fp8_amax_compute_algo = "most_recent" + model_cfg.fp8_dot_product_attention = False + model_cfg.disable_parameter_transpose_cache = False + model_cfg.fine_grained_activation_offloading = False + model_cfg.use_transformer_engine_full_layer_spec = ( + False # Doesn't work beacuse of RMSNorm is not supported in FusedLayerNorm + ) + model_cfg.cpu_offloading = False + model_cfg.cpu_offloading_num_layers = 0 + model_cfg.empty_unused_memory_level = ( + 0 # 0: No empty, 1: Empty at end of eval, 2: Empty at end of eval and train. + ) + # Optional MXFP4-style activation recompute (left untouched when kwargs are None). + if recompute_granularity is not None: + model_cfg.recompute_granularity = recompute_granularity + if recompute_method is not None: + model_cfg.recompute_method = recompute_method + if recompute_num_layers is not None: + model_cfg.recompute_num_layers = recompute_num_layers + # Pin TE's attention backend ("fused" forces CK / AOTriton on ROCm for the MXFP4 path). + # Leave as None to keep Megatron's "auto" default (which picks FlashAttention on ROCm). + if attention_backend is not None: + try: + model_cfg.attention_backend = AttnBackend[attention_backend] + except KeyError as e: + raise ValueError( + f"Unknown attention_backend {attention_backend!r}; expected one of " + f"{[b.name for b in AttnBackend]}." + ) from e + # Disable attention QK clipping / max-logit scans in the optimizer path (extra GPU work per step). + if hasattr(model_cfg, "qk_clip"): + model_cfg.qk_clip = False + if hasattr(model_cfg, "log_max_attention_logit"): + model_cfg.log_max_attention_logit = False + + from megatron.bridge.training.config import OptimizerConfig, SchedulerConfig + + opt_config = OptimizerConfig( + optimizer="adam", + lr=lr, + min_lr=min_lr, + clip_grad=clip_grad, + weight_decay=weight_decay, + adam_beta1=adam_beta1, + adam_beta2=adam_beta2, + adam_eps=adam_eps, + bf16=True, + params_dtype=torch.bfloat16, + use_distributed_optimizer=True, + overlap_param_gather_with_optimizer_step=True, + barrier_with_L1_time=False, + log_num_zeros_in_grad=False, + ) + + scheduler = SchedulerConfig( + lr_decay_style="cosine", + lr_decay_iters=lr_decay_iters, # same as max_steps in nemo + lr_warmup_iters=lr_warmup_iters, # value 0 same as nemo + lr_warmup_init=0.0, + start_weight_decay=weight_decay, + end_weight_decay=weight_decay, + weight_decay_incr_style="constant", + override_opt_param_scheduler=True, + ) + + peft_config = LoRA( + dim=16, + alpha=32, + dropout=0.1, + dropout_position="pre", + lora_A_init_method="xavier", + lora_B_init_method="zero", + a2a_experimental=True, + target_modules=["linear_qkv", "linear_proj"], + use_te_fused_lora=False, + te_fused_lora_include_modules=te_fused_lora_include_modules, + te_fused_lora_exclude_modules=( + te_fused_lora_exclude_modules if te_fused_lora_exclude_modules is not None else [] + ), + ) + + # Dataset configuration - switch between squad and mlperf_dataset + if dataset_type == "squad": + dataset_cfg = default_squad_config(seq_length, packed_sequence) + elif dataset_type == "mlperf_dataset": + if packed_sequence: + packed_sequence_specs = PackedSequenceSpecs( + packed_sequence_size=seq_length, + tokenizer_model_name=hf_path, + packed_train_data_path=packed_train_data_path or "/data/train.npy", + packed_val_data_path=packed_val_data_path or "/data/validation.npy", + packed_metadata_path=packed_metadata_path or "/data/packed_metadata.jsonl", + ) + else: + packed_sequence_specs = None + dataset_cfg = FinetuningDatasetConfig( + dataset_root="/data", + seq_length=seq_length, + seed=seed, + packed_sequence_specs=packed_sequence_specs, + data_sharding=True, + dataloader_type="batch", + num_workers=1, + do_test=False, + do_validation=True, + dataset_kwargs={"return_cu_seqlen": False}, + ) + else: + raise ValueError(f"Unknown dataset_type: {dataset_type!r}. Expected 'squad' or 'mlperf_dataset'.") + + dataset_cfg.num_workers = 0 + dataset_cfg.memmap_workers = 1 # needs to be 1>0 + dataset_cfg.pin_memory = True + dataset_cfg.persistent_workers = False + dataset_cfg.dataloader_type = "batch" + + # Config Container + cfg = ConfigContainer( + model=model_cfg, + train=TrainingConfig( + train_iters=train_iters, + eval_interval=eval_interval, + eval_iters=eval_iters, + global_batch_size=global_batch_size, + micro_batch_size=micro_batch_size, + # Manual GC aligns collections across ranks but adds periodic host pauses; use default GC for best step time. + manual_gc=False, + manual_gc_interval=0, + manual_gc_eval=False, + empty_unused_memory_level=0, # 0: No empty, 1: Empty at end of eval, 2: Empty at end of eval and train. + train_sync_interval=None, + check_weight_hash_across_dp_replicas_interval=None, + ), + optimizer=opt_config, + scheduler=scheduler, + ddp=DistributedDataParallelConfig( + # Per-step NaN grad scan + sync; disable for throughput when loss NaN checks are off. + check_for_nan_in_grad=False, + grad_reduce_in_fp32=False, + overlap_grad_reduce=False, + overlap_param_gather=False, + average_in_collective=False, + use_distributed_optimizer=True, + # gradient_reduce_div_fusion=True, + # pad_buckets_for_high_nccl_busbw=True, + use_megatron_fsdp=False, + keep_fp8_transpose_cache=( + os.getenv("ENABLE_TRANSPOSE_CACHE", "").strip().lower() in ("1", "true", "yes", "on") + ), + fp8_param_gather=False, + ), + dataset=dataset_cfg, + logger=LoggerConfig( + log_interval=10, + tensorboard_dir=None, + # Per-step / log-interval overhead toggles (keep off for throughput). + log_params_norm=False, + log_throughput=True, + log_energy=False, + log_progress=False, + timing_log_level=0, + log_loss_scale_to_tensorboard=False, + log_timers_to_tensorboard=False, + log_throughput_to_tensorboard=False, + log_validation_ppl_to_tensorboard=False, + log_memory_to_tensorboard=False, + log_runtime_to_tensorboard=False, + log_world_size_to_tensorboard=False, + log_l2_norm_grad_to_tensorboard=False, + wandb_project=None, + wandb_exp_name=None, + wandb_save_dir=None, + wandb_entity=None, + ), + tokenizer=TokenizerConfig( + tokenizer_type="NullTokenizer" if use_null_tokenizer else "HuggingFaceTokenizer", + tokenizer_model=hf_path if not use_null_tokenizer else None, + vocab_size=DEFAULT_NULL_TOKENIZER_VOCAB_SIZE if use_null_tokenizer else None, + ), + checkpoint=CheckpointConfig( + save_interval=None, + save=None, + load=None, + save_optim=False, + save_rng=False, + pretrained_checkpoint=pretrained_checkpoint, + ckpt_format="torch_dist", + replication_factor=0, + most_recent_k=0, + finetune=True, + load_main_params_from_ckpt=False, + load_optim=False, + load_rng=False, + ), + rng=RNGConfig(seed=seed, te_rng_tracker=True), + rerun_state_machine=RerunStateMachineConfig( + check_for_nan_in_loss=check_for_nan_in_loss, + ), + straggler=StragglerDetectionConfig(log_straggler=False), + tensor_inspect=TensorInspectConfig(enabled=False), + comm_overlap=comm_overlap_config, + mixed_precision=precision_config, + peft=peft_config, + profiling=ProfilingConfig( + use_pytorch_profiler=False, + profile_step_start=140, + profile_step_end=144, + profile_ranks=list(range(8)), # 1 node × 8 GPUs + record_shapes=False, + nvtx_ranges=False, + ), + ) + + if cfg.comm_overlap is None: + cfg.comm_overlap = CommOverlapConfig( + tp_comm_overlap=False, + ) + + return cfg + + +def _reset_data_iterator(data_iterator): + """Reset data iterator to the beginning for deterministic evaluation. + Traverses through RerunDataIterator wrappers to find and reset any + ResettableDataIterator, ensuring evaluation always starts from the + same data point. + """ + if data_iterator is None: + return + if isinstance(data_iterator, list): + for it in data_iterator: + _reset_data_iterator(it) + return + if isinstance(data_iterator, RerunDataIterator): + inner = data_iterator.iterable + if isinstance(inner, ResettableDataIterator): + inner.reset() + data_iterator.saved_microbatches.clear() + data_iterator.replaying = False + data_iterator.replay_pos = 0 + elif isinstance(data_iterator, ResettableDataIterator): + data_iterator.reset() + + +def evaluate( + state: GlobalState, + forward_step_func: ForwardStepCallable, + data_iterator: Optional[Union[RerunDataIterator, list[RerunDataIterator]]], + model: list[MegatronModule], + process_non_loss_data_func: Optional[Callable], + config: ConfigContainer, + verbose: bool = False, + non_loss_data_func: Optional[Callable] = None, +) -> tuple[Optional[dict[str, torch.Tensor]], Optional[Any], bool]: + """Evaluation function (from eval_m.py). + Validation loss aggregation matches NeMo: mean of per-microbatch means + (same as NeMo MaskedTokenLossReduction with validation_step=True, val_drop_last=True). + """ + _reset_data_iterator(data_iterator) + wrapped_forward_step = prepare_forward_step_func(forward_step_func, state) + timers = state.timers + timers("evaluate", log_level=0).start(barrier=True) + for model_module in model: + model_module.eval() + pg_collection = get_pg_collection(model) + rerun_state_machine = get_rerun_state_machine() + rerun_mode = rerun_state_machine.get_mode() + rerun_state_machine.set_mode(RerunMode.DISABLED) + total_loss_dict = {} + eval_batch_size = state.cfg.train.global_batch_size + eval_num_microbatches = eval_batch_size // ( + state.cfg.train.micro_batch_size * state.cfg.data_parallel_size + ) + with torch.no_grad(): + if verbose: + log_rank_0(f"Evaluating on {state.cfg.train.eval_iters * eval_batch_size} samples") + if ( + state.cfg.model.cuda_graph_impl == "local" + and "full_iteration" in state.cfg.model.cuda_graph_scope + ): + forward_backward_func = FullCudaGraphWrapper( + get_forward_backward_func(), + cuda_graph_warmup_steps=state.cfg.model.cuda_graph_warmup_steps, + ) + else: + forward_backward_func = get_forward_backward_func() + iteration = 0 + while iteration < state.cfg.train.eval_iters: + iteration += 1 + if verbose: + log_rank_0(f"Evaluating iter {iteration}/{state.cfg.train.eval_iters}") + seq_length = state.cfg.model.seq_length + eval_data_iterator = data_iterator + if state.cfg.dataset.dataloader_type == "batch": + eval_microbatch_iterator, seq_length = prepare_finetuning_batch( + data_iterator=data_iterator, + num_microbatches=eval_num_microbatches, + default_seq_length=state.cfg.model.seq_length, + seq_key="tokens", + ) + eval_data_iterator = make_data_iterator_list( + model=model, + data_iterator=eval_microbatch_iterator, + ) + config.timers = None + fault_tolerance.on_eval_step_start(state) + loss_dicts = forward_backward_func( + forward_step_func=wrapped_forward_step, + data_iterator=eval_data_iterator, + model=model, + num_microbatches=eval_num_microbatches, + seq_length=seq_length, + micro_batch_size=state.cfg.train.micro_batch_size, + forward_only=True, + ) + fault_tolerance.on_eval_step_end(state) + config.timers = state.timers + if is_pp_last_stage(pg_collection.pp): + for key in loss_dicts[0].keys(): + if key not in total_loss_dict: + total_loss_dict[key] = torch.zeros(2, dtype=torch.float32, device="cuda") + val = [x[key].reshape(-1) for x in loss_dicts] + if val[0].numel() == 2: + # NeMo MLPerf-equivalent: accumulate raw [sum, count] across micros & eval iters. + # Per-micro [sum, count] is already DP/CP-reduced inside + # nemo_loss.MaskedTokenLossReduction.forward, so no extra all_reduce here. + per_iter = torch.vstack([v.float() for v in val]).sum( + dim=0 + ) # [Σ_micro sum, Σ_micro count] + total_loss_dict[key] += per_iter + elif val[0].numel() == 1: + # legacy single-scalar branch + micro_sum = torch.stack([v.float() for v in val], dim=0).sum() + total_loss_dict[key][0] += micro_sum + total_loss_dict[key][1] += float(len(loss_dicts)) + else: + raise ValueError(f"Invalid value shape: {val[0].shape} for key {key}") + + state.train_state.consumed_valid_samples += eval_batch_size + if state.cfg.train.exit_duration_in_mins: + train_time = (time.time() - state.start_time) / 60.0 + done_cuda = torch.tensor( + [train_time > state.cfg.train.exit_duration_in_mins], + dtype=torch.int, + device="cuda", + ) + torch.distributed.all_reduce(done_cuda, op=torch.distributed.ReduceOp.MAX) + done = bool(done_cuda.item()) + if done: + rerun_state_machine.set_mode(rerun_mode) + log_rank_0("Exiting during evaluation, timelimit reached") + return None, None, True + + collected_non_loss_data = None + if non_loss_data_func is not None: + collected_non_loss_data = non_loss_data_func(model) + elif process_non_loss_data_func is not None and is_last_rank(): + non_loss_data_iterator = data_iterator + non_loss_seq_length = state.cfg.model.seq_length + if state.cfg.dataset.dataloader_type == "batch": + non_loss_microbatch_iterator, non_loss_seq_length = prepare_finetuning_batch( + data_iterator=data_iterator, + num_microbatches=get_num_microbatches(), + default_seq_length=state.cfg.model.seq_length, + seq_key="tokens", + ) + non_loss_data_iterator = make_data_iterator_list( + model=model, + data_iterator=non_loss_microbatch_iterator, + ) + collected_non_loss_data = forward_backward_func( + forward_step_func=wrapped_forward_step, + data_iterator=non_loss_data_iterator, + model=model, + num_microbatches=get_num_microbatches(), + seq_length=non_loss_seq_length, + micro_batch_size=state.cfg.train.micro_batch_size, + forward_only=True, + collect_non_loss_data=True, + ) + for model_module in model: + model_module.train() + for key in total_loss_dict: + val_loss_sum, val_loss_count = total_loss_dict[key] + if val_loss_count > 0: + total_loss_dict[key] = val_loss_sum / val_loss_count + else: + total_loss_dict[key] = val_loss_sum + timers("evaluate").stop() + timers.log(["evaluate"]) + rerun_state_machine.set_mode(rerun_mode) + return total_loss_dict, collected_non_loss_data, False + + +from megatron.bridge.training import eval + + +def evaluate_and_print_results_custom( + state: GlobalState, + prefix: str, + forward_step_func: ForwardStepCallable, + data_iterator: Optional[Union[RerunDataIterator, list[RerunDataIterator]]], + model: list[MegatronModule], + config: ConfigContainer, + verbose: bool = False, + write_to_tensorboard: bool = False, + process_non_loss_data_func: Optional[Callable] = None, + non_loss_data_func: Optional[Callable] = None, + throughput: float = 0.0, # samples/sec +) -> tuple: + """Helper function to evaluate and dump results on screen. + + Args: + state (GlobalState): The global state object. + prefix (str): Prefix for logging evaluation results. + forward_step_func (Callable): The function that performs a forward step. + data_iterator (Optional[Union[RerunDataIterator, list[RerunDataIterator]]]): Iterator over evaluation data. + model (list[MegatronModule]): list of model chunks. + config (ConfigContainer): Configuration container (potentially redundant). + verbose (bool, optional): Whether to print evaluation progress. Defaults to False. + write_to_tensorboard (bool, optional): Whether to write results to TensorBoard. Defaults to False. + process_non_loss_data_func (Optional[Callable], optional): Function to process non-loss data. Defaults to None. + non_loss_data_func (Optional[Callable], optional): Function to compute non-loss data. Defaults to None. + """ + global _TARGET_LOSS_REACHED + if _TARGET_LOSS_REACHED: + # Target already hit on a previous pass; skip this entire eval. + # Returning should_exit=True keeps the outer while-loop on its exit path. + _orig_log_rank_0( + f"Skipping evaluation at {prefix}: target loss < " f"{MLPERF_TARGET_LOSS} already reached." + ) + return True, None + + log_rank_0(f"Evaluating and printing results at {prefix}") + should_exit = False + + def is_last_rank(): + return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1) + + import math + + if write_to_tensorboard: + writer = state.tensorboard_logger + else: + writer = None + + wandb_writer = state.wandb_logger + + eval_start = time.time() + total_loss_dict, collected_non_loss_data, timelimit = evaluate( + state, + forward_step_func, + data_iterator, + model, + process_non_loss_data_func, + config, + verbose, + non_loss_data_func, + ) + eval_duration = time.time() - eval_start + # Timelimit hit during evaluation + if timelimit: + return False, None + string = f" validation loss at {prefix} | " + for key in total_loss_dict: + string += "{} value: {:.6E} | ".format(key, total_loss_dict[key].item()) + ppl = math.exp(min(20, total_loss_dict[key].item())) + string += "{} PPL: {:.6E} | ".format(key, ppl) + if writer: + writer.add_scalar( + "{} validation".format(key), total_loss_dict[key].item(), state.train_state.step + ) + writer.add_scalar( + "{} validation vs samples".format(key), + total_loss_dict[key].item(), + state.train_state.consumed_train_samples, + ) + if state.cfg.logger.log_validation_ppl_to_tensorboard: + writer.add_scalar("{} validation ppl".format(key), ppl, state.train_state.step) + writer.add_scalar( + "{} validation ppl vs samples".format(key), ppl, state.train_state.consumed_train_samples + ) + + if wandb_writer and is_last_rank(): + wandb_writer.log( + {"{} validation".format(key): total_loss_dict[key].item()}, state.train_state.step + ) + if state.cfg.logger.log_validation_ppl_to_tensorboard: + wandb_writer.log({"{} validation ppl".format(key): ppl}, state.train_state.step) + + if process_non_loss_data_func is not None and writer and is_last_rank(): + process_non_loss_data_func(collected_non_loss_data, state.train_state.step, writer) + + string += "throughput: {:.6E} | ".format(throughput) + string += "eval duration: {:.6E} | ".format(eval_duration) + if writer: + writer.add_scalar("throughput samples/sec", throughput, state.train_state.step) + writer.add_scalar( + "throughput samples/sec vs samples", throughput, state.train_state.consumed_train_samples + ) + + if wandb_writer and is_last_rank(): + wandb_writer.log({"throughput samples/sec": throughput}, state.train_state.step) + + length = len(string) + 1 + # Match training logs: emit on rank 0 so validation lines show up in typical Primus rank-0 log streams. + log_rank_0("-" * length) + log_rank_0(string) + log_rank_0("-" * length) + # Guard against non-PP-last ranks (or otherwise empty total_loss_dict) where + # "lm loss" may not be populated; only run the MLPerf early-exit check when + # the key is present. + eval_loss_value = None + if total_loss_dict and "lm loss" in total_loss_dict: + eval_loss_value = total_loss_dict["lm loss"].item() + if eval_loss_value < MLPERF_TARGET_LOSS: + should_exit = True + _TARGET_LOSS_REACHED = True + log_rank_0(f"Validation loss is less than {MLPERF_TARGET_LOSS}, exiting training") + return should_exit, eval_loss_value + + +eval.evaluate_and_print_results = evaluate_and_print_results_custom + + +def warmup_eval( + state: GlobalState, + forward_step_func: ForwardStepCallable, + data_iterator: Optional[Union[RerunDataIterator, list[RerunDataIterator]]], + model: list[MegatronModule], + config: ConfigContainer, + verbose: bool = False, + process_non_loss_data_func: Optional[Callable] = None, + non_loss_data_func: Optional[Callable] = None, + num_warmup_iters: int = 10, +) -> None: + log_rank_0(f"Starting warmup eval...") + for i in range(num_warmup_iters): + log_rank_0(f"Warmup eval iteration {i} running...") + evaluate_and_print_results_custom( + state, + f"warmup iteration {i}", + forward_step_func, + data_iterator, + model, + config, + verbose=verbose, + write_to_tensorboard=False, + process_non_loss_data_func=process_non_loss_data_func, + non_loss_data_func=non_loss_data_func, + ) + log_rank_0(f"Warmup eval iteration {i} completed") + log_rank_0(f"Warmup eval completed") + + +class _SyntheticSFTDataIterator: + """Infinite iterator yielding synthetic finetuning batches for warmup. + + Produces random token tensors matching the SFT packed-sequence shape expected + by prepare_finetuning_batch / forward_backward_func. + """ + + def __init__(self, seq_length: int, micro_batch_size: int, vocab_size: int = 32000): + self._seq_length = seq_length + self._mbs = micro_batch_size + self._vocab_size = vocab_size + + def __iter__(self): + return self + + def __next__(self): + sl = self._seq_length + mbs = self._mbs + tokens = torch.randint(0, self._vocab_size, (mbs, sl), dtype=torch.int64, device="cuda") + labels = torch.randint(0, self._vocab_size, (mbs, sl), dtype=torch.int64, device="cuda") + loss_mask = torch.ones(mbs, sl, dtype=torch.float32, device="cuda") + position_ids = torch.arange(sl, dtype=torch.int64, device="cuda").unsqueeze(0).expand(mbs, -1) + return {"tokens": tokens, "labels": labels, "loss_mask": loss_mask, "position_ids": position_ids} + + +def run_synthetic_warmup( + forward_step_func: ForwardStepCallable, + model: list[MegatronModule], + optimizer: MegatronOptimizer, + scheduler: OptimizerParamScheduler, + global_state: GlobalState, + pg_collection: ProcessGroupCollection, +) -> None: + """Run synthetic warmup steps to pre-compile JIT kernels before RUN_START. + + Mirrors NeMo's warmup behavior: runs N training (fwd+bwd+opt) steps and + M validation (fwd-only) steps with random data, then restores all state so + that real training starts from a clean slate. + + Controlled by environment variables: + SYNTH_WARMUP_STEPS - number of training warmup steps (default 5, 0 disables) + SYNTH_WARMUP_VALID_STEPS - number of validation warmup steps (default 5, 0 disables) + """ + warmup_train_steps = int(os.getenv("SYNTH_WARMUP_STEPS", "5")) + warmup_valid_steps = int(os.getenv("SYNTH_WARMUP_VALID_STEPS", "5")) + + if warmup_train_steps <= 0 and warmup_valid_steps <= 0: + return + + config = global_state.cfg + seq_length = config.model.seq_length + mbs = config.train.micro_batch_size + gbs = config.train.global_batch_size + dp_size = getattr(config, "data_parallel_size", 1) + num_microbatches = gbs // (mbs * dp_size) + + synth_iter = _SyntheticSFTDataIterator(seq_length, mbs) + wrapped_forward_step = prepare_forward_step_func(forward_step_func, global_state) + forward_backward_func = get_forward_backward_func() + + # --- Save model parameters --- + models = model if isinstance(model, (list, tuple)) else [model] + saved_params = {} + for m in models: + for name, p in m.named_parameters(): + if p.requires_grad: + saved_params[(id(m), name)] = p.data.to("cpu", copy=True) + + # --- Save optimizer state (neuter to prevent real updates) --- + saved_opt_states = [] + inner_opts = [] + if hasattr(optimizer, "chained_optimizers"): + for sub in optimizer.chained_optimizers: + inner_opts.append(getattr(sub, "optimizer", sub)) + else: + inner_opts.append(getattr(optimizer, "optimizer", optimizer)) + + for inner in inner_opts: + saved = [] + for group in inner.param_groups: + state = {} + for key in ("betas", "weight_decay", "bias_correction"): + if key in group: + state[key] = group[key] + saved.append(state) + if "betas" in group: + group["betas"] = [1.0, 1.0] + if "weight_decay" in group: + group["weight_decay"] = 0.0 + if "bias_correction" in group: + group["bias_correction"] = False + saved_opt_states.append(saved) + + # --- Save LR scheduler state --- + saved_sched = {} + if scheduler is not None: + for k in ("num_steps", "num_floating_point_operations_so_far"): + if hasattr(scheduler, k): + saved_sched[k] = getattr(scheduler, k) + + # --- Training warmup steps (fwd + bwd + optimizer) --- + if warmup_train_steps > 0: + for m in models: + m.train() + for step in range(1, warmup_train_steps + 1): + train_step( + wrapped_forward_step, + synth_iter, + model, + optimizer, + scheduler, + global_state, + pg_collection, + forward_backward_func, + ) + torch.cuda.synchronize() + + # --- Validation warmup steps (forward-only) --- + if warmup_valid_steps > 0: + for m in models: + m.eval() + with torch.no_grad(): + for step in range(1, warmup_valid_steps + 1): + forward_backward_func( + forward_step_func=wrapped_forward_step, + data_iterator=synth_iter, + model=model, + num_microbatches=num_microbatches, + seq_length=seq_length, + micro_batch_size=mbs, + forward_only=True, + ) + torch.cuda.synchronize() + for m in models: + m.train() + + # --- Restore model parameters --- + for m in models: + for name, p in m.named_parameters(): + key = (id(m), name) + if key in saved_params: + p.data.copy_(saved_params[key].to(p.device)) + del saved_params + + # --- Restore optimizer state --- + for inner, saved in zip(inner_opts, saved_opt_states): + for group, state in zip(inner.param_groups, saved): + for key, val in state.items(): + group[key] = val + + if hasattr(optimizer, "reload_model_params"): + optimizer.reload_model_params() + + # --- Zero optimizer state tensors (exp_avg / exp_avg_sq) --- + for inner in inner_opts: + for param_states in inner.state.values(): + for k, v in param_states.items(): + if isinstance(v, torch.Tensor) and v.is_floating_point(): + v.zero_() + + # --- Restore LR scheduler and re-sync param_groups['lr'] --- + if scheduler is not None: + for k, v in saved_sched.items(): + setattr(scheduler, k, v) + scheduler.step(0) + + # --- Reset FP8 state --- + for m in models: + for module in m.modules(): + if hasattr(module, "fp8_initialized"): + module.fp8_initialized = False + if hasattr(module, "reset_fp8_meta_tensors"): + try: + module.reset_fp8_meta_tensors() + except Exception: + pass + + # --- Seed FP8 amax_history to prevent scale=inf after reset --- + for m in models: + for module in m.modules(): + fp8_meta = getattr(module, "fp8_meta", None) + if fp8_meta is None or not isinstance(fp8_meta, dict): + continue + for key in ("scaling_fwd", "scaling_bwd"): + if key not in fp8_meta: + continue + tensor_meta = fp8_meta[key] + if hasattr(tensor_meta, "amax_history"): + tensor_meta.amax_history.fill_(1.0) + + # --- Reset FP4/MXFP4 state (block scales, quantizer caches) --- + for m in models: + for module in m.modules(): + is_fp4 = bool(getattr(module, "fp4", False)) or hasattr(module, "fp4_initialized") + if not is_fp4: + fp8_meta = getattr(module, "fp8_meta", None) + if isinstance(fp8_meta, dict): + for key in ("scaling_fwd", "scaling_bwd"): + tm = fp8_meta.get(key) + if tm is not None and ("FP4" in type(tm).__name__ or "Fp4" in type(tm).__name__): + is_fp4 = True + break + if not is_fp4: + continue + if hasattr(module, "fp4_initialized"): + module.fp4_initialized = False + if hasattr(module, "fp8_initialized"): + module.fp8_initialized = False + if hasattr(module, "reset_fp4_meta_tensors"): + try: + module.reset_fp4_meta_tensors() + except Exception: + pass + elif hasattr(module, "reset_fp8_meta_tensors"): + try: + module.reset_fp8_meta_tensors() + except Exception: + pass + + # --- Clear gradients --- + for m in models: + if hasattr(m, "zero_grad_buffer"): + m.zero_grad_buffer() + m.zero_grad(set_to_none=True) + + # --- Reset consumed samples / step counter back to 0 --- + global_state.train_state.step = 0 + global_state.train_state.consumed_train_samples = 0 + + # --- Clear CUDA cache (single clear after warmup, before RUN_START) --- + torch.cuda.synchronize() + torch.cuda.empty_cache() + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + +def megatron_bridge_train_override( + forward_step_func: ForwardStepCallable, + model: list[MegatronModule], + optimizer: MegatronOptimizer, + scheduler: OptimizerParamScheduler, + train_data_iterator: Optional[Union[RerunDataIterator, list[RerunDataIterator]]], + valid_data_iterator: Optional[Union[RerunDataIterator, list[RerunDataIterator]]], + global_state: GlobalState, + checkpointing_context: dict[str, Any], + pg_collection: ProcessGroupCollection, + process_non_loss_data_func: Optional[Callable] = None, + non_loss_data_func: Optional[Callable] = None, +) -> None: + """Main training loop. + + Handles the overall training process, including the iteration loop, + calling train_step, evaluation, checkpointing, logging, and exit conditions. + + Args: + forward_step_func: Callable that executes a single forward step. + model: list of model chunks (potentially wrapped in DDP). + optimizer: The optimizer instance. + scheduler: The learning rate scheduler instance. + train_data_iterator: Iterator for the training dataset. + valid_data_iterator: Iterator for the validation dataset. + global_state: The GlobalState object holding various training states. + checkpointing_context: Context dictionary for checkpointing. + process_non_loss_data_func: Optional function to process non-loss data during evaluation. + non_loss_data_func: Optional function to compute non-loss data during evaluation. + + Warnings: + This is an experimental API and is subject to change in backwards + incompatible ways without notice. + """ + + # import ctypes + # import gc + # libc = ctypes.CDLL("libc.so.6") + # gc.collect() + # libc.malloc_trim(0) + + config: ConfigContainer = global_state.cfg + model_config = get_model_config(model[0]) + train_config = config.train + timers = global_state.timers + straggler_timer = global_state.straggler_timer + energy_monitor = global_state.energy_monitor + + # Prepare forward_step_func (check signature and inject state if needed). + # This is done once to prevent creating new partial objects every iteration. + # + # Note on reference semantics: + # - functools.partial stores a reference to global_state, not a copy + # - When global_state.train_state.step changes, the partial sees the updated value + # - This is safe because GlobalState is a mutable object passed by reference + # + # For functors (classes with __call__ defined): + # - For functors: partial(functor_instance, state) still allows functor's internal state to work + # - inspect.signature() properly inspects the __call__ method of functors + wrapped_forward_step_func = prepare_forward_step_func(forward_step_func, global_state) + + # Turn on training mode which enables dropout. + for model_module in model: + model_module.train() + log_rank_0(f"Model module: {model_module}") + + # Tracking loss. + total_loss_dict = {} + + # Make sure rerun_state_machine has the right iteration loaded from checkpoint. + rerun_state_machine = get_rerun_state_machine() + if rerun_state_machine.current_iteration != global_state.train_state.step: + log_rank_0(f"Setting rerun_state_machine.current_iteration to {global_state.train_state.step}...") + rerun_state_machine.current_iteration = global_state.train_state.step + + global_state.train_state.floating_point_operations_so_far + num_floating_point_operations_since_last_log_event = 0.0 + + if energy_monitor is not None: + energy_monitor.setup() + energy_monitor.resume() + + # interval-time is started/stopped once per iteration around train_step only (see training loop). + # A single long-running interval-time includes logging, hooks, and other post-step work, which + # inflates "elapsed time per iteration (ms)" vs profiler / train-step wall time. + report_memory_flag = True + pre_hook_enabled = False + should_exit = False + exit_code = 0 + + if train_config.manual_gc: + # Disable the default garbage collector and perform the collection manually. + # This is to align the timing of garbage collection across ranks. + assert ( + train_config.manual_gc_interval >= 0 + ), "Manual garbage collection interval should be larger than or equal to 0" + gc.disable() + gc.collect() + + if config.straggler and config.straggler.log_straggler: + world = torch.distributed.get_world_size() + rank = torch.distributed.get_rank() + mmcnt = config.straggler.straggler_minmax_count + straggler_timer.configure( + world, + rank, + mmcnt=mmcnt, + enabled=not config.straggler.disable_straggler_on_startup, + port=config.straggler.straggler_ctrlr_port, + ) + + # Initialize NVRx straggler detection if enabled + nvrx_straggler_manager = global_state.nvrx_straggler_manager + if nvrx_straggler_manager is not None: + try: + # Initialize the straggler detector first + nvrx_straggler_manager.initialize() + # Wrap the train_step function for monitoring + # Note: The nvidia-resiliency-ext library will monitor the actual train_step calls + nvrx_straggler_manager.wrap_train_step_function(train_step) + except Exception as e: + log_rank_0(f"Failed to initialize NVRx straggler detection: {e}") + # Set to None to disable further checks + global_state._nvrx_straggler_manager = None + + get_num_microbatches() + eval_duration = 0.0 + eval_iterations = 0 + + prof = None + nsys_nvtx_context = None # NVTX context for nsys profiling + prof_config = config.profiling + if prof_config and should_profile_rank(prof_config, torch.distributed.get_rank()): + if prof_config.use_pytorch_profiler: + trace_dir = config.logger.tensorboard_dir or os.path.join(os.getcwd(), "torch_profiler_traces") + prof = initialize_pytorch_profiler(prof_config, trace_dir) + prof.start() + + # Initialize RPD profiler if enabled + rpd = None + rpd_status = None + profiler_type = os.getenv("PROFILER", "") + rpd_warmup_steps = int(os.getenv("RPD_WARMUP_STEPS", "0")) + rpd_active_steps = int(os.getenv("RPD_ACTIVE_STEPS", "100")) + if profiler_type == "rpd": + try: + from rpdTracerControl import rpdTracerControl + + rank = torch.distributed.get_rank() + rpd_filename = os.getenv("RPD_TRACE_FILENAME", f"trace.rpd") + rpdTracerControl.setFilename(name=rpd_filename, append=True) + rpd = rpdTracerControl() + log_rank_0( + f"RPD profiler initialized. Will start at step {rpd_warmup_steps} and stop at step {rpd_warmup_steps + rpd_active_steps}" + ) + except Exception as e: + log_rank_0(f"Failed to initialize RPD profiler: {e}") + rpd = None + else: + log_rank_0(f"### Profiler type is {profiler_type}") + + # Megatron FSDP and FSDP2 does not have this hook + should_toggle_forward_pre_hook = should_disable_forward_pre_hook( + config.ddp.use_megatron_fsdp, + config.optimizer.use_distributed_optimizer, + config.ddp.overlap_param_gather, + ) + # Disable forward pre-hook to start training to ensure that errors in checkpoint loading + # or random initialization don't propagate to all ranks in first all-gather (which is a + # no-op if things work correctly). + if should_toggle_forward_pre_hook: + disable_forward_pre_hook(model, param_sync=False) + # Also remove param_sync_func temporarily so that sync calls made in + # `forward_backward_func` are no-ops. + param_sync_func = model_config.param_sync_func + model_config.param_sync_func = None + pre_hook_enabled = False + # Also, check weight hash across DP replicas to be very pedantic. + if train_config.check_weight_hash_across_dp_replicas_interval is not None: + assert check_param_hashes_across_dp_replicas( + model, cross_check=True + ), "Parameter hashes not matching across DP replicas" + torch.distributed.barrier() + log_rank_0(f">>> Weight hashes match after {global_state.train_state.step} iterations...") + + # Capture CUDA Graphs. + cuda_graph_helper = None + if model_config.cuda_graph_impl == "transformer_engine": + cuda_graph_helper = TECudaGraphHelper( + model=model, + config=model_config, + seq_length=config.model.seq_length, + micro_batch_size=config.train.micro_batch_size, + optimizers=[optimizer], + ) + + # Track train step elapsed time for throughput logging + history_wct = None + if config.logger.log_throughput_to_tensorboard: + history_wct = deque(maxlen=config.logger.throughput_window_size + 1) + + # Wrap forward_backward_func for Full iteration CUDA graph + forward_backward_func = get_forward_backward_func() + if config.model.cuda_graph_impl == "local" and "full_iteration" in config.model.cuda_graph_scope: + forward_backward_func = FullCudaGraphWrapper( + forward_backward_func, cuda_graph_warmup_steps=config.model.cuda_graph_warmup_steps + ) + + start_iteration = global_state.train_state.step + log_rank_0(f"Starting training loop at iteration {start_iteration}") + + dp_size = pg_collection.dp.size() + batch_size = dp_size * train_config.micro_batch_size * get_num_microbatches() + timer = Timer(batch_size) + + # Synthetic warmup: pre-compile JIT kernels before measured training begins. + run_synthetic_warmup( + forward_step_func, + model, + optimizer, + scheduler, + global_state, + pg_collection, + ) + _log_training_gpu_mem("after synthetic warmup", config.logger.memory_keys) + + sft_logger = _get_sft_logger() + + # MLPerf logging: transition from init to training (after warmup) + if sft_logger is not None: + sft_logger.log_init_stop_run_start() + + # Wall-clock for the training loop only (first through last train step; warmup eval disabled). + training_wall_start = time.perf_counter() + # Skip the first N interval evals (e.g. 48, 96, 144 when eval_interval=48); run from step 4*interval (192). + eval_skip_first_n = 3 + + # One-shot MXFP4 healing env banner (no-op unless HEALING_ITER > 0 and + # MXFP4_HEALING_PHASE_LOG is on). Safe to call even when the healing + # module isn't imported — gracefully degrades. + try: + from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.mxfp4_healing import ( + log_healing_env_banner_once, + ) + + log_healing_env_banner_once() + except Exception: # noqa: BLE001 + pass + + # NeMo / MLPerf reference (MI355X implementation): CustomCallback uses time.time() in Lightning + # on_train_batch_start / on_train_batch_end — i.e. wall time over training_step only, no cuda.synchronize. + # We mirror that by timing megatron.train_step only, then averaging over logger.log_interval like interval-time. + nemo_style_iter_seconds_accum = 0.0 + nemo_style_iter_count = 0 + + # Run training iterations till done. + while global_state.train_state.step < train_config.train_iters: + # Handle RPD profiling start + if rpd and global_state.train_state.step >= rpd_warmup_steps and not rpd_status: + log_rank_0(f"Starting RPD profiling at step {global_state.train_state.step}") + rpd.start() + rpd_status = "running" + + # Handle RPD profiling stop + if ( + rpd + and rpd_status == "running" + and global_state.train_state.step >= rpd_warmup_steps + rpd_active_steps + ): + log_rank_0(f"Stopping RPD profiling at step {global_state.train_state.step}") + rpd.stop() + rpd_status = "finished" + + # Handle profiling for this step + nvtx_ctx = handle_profiling_step( + prof_config, + global_state.train_state.step, + torch.distributed.get_rank(), + prof, + ) + if nvtx_ctx is not None: + nsys_nvtx_context = nvtx_ctx + + # Update the timeout for all process groups after initialization + # We update the timeout after the first successful iteration, + # which takes longer than others usually + if global_state.train_state.step == start_iteration + 1: + distributed_timeout_seconds_after_init = ( + global_state.cfg.dist.distributed_timeout_seconds_after_init + ) + if distributed_timeout_seconds_after_init is not None: + update_pg_timeout(timedelta(seconds=distributed_timeout_seconds_after_init)) + + # Capture CUDA Graphs after warmup. + # if ( + # model_config.cuda_graph_impl == "transformer_engine" + # and cuda_graph_helper is not None + # and not cuda_graph_helper.graphs_created() + # and global_state.train_state.step - start_iteration == model_config.cuda_graph_warmup_steps + # ): + # if model_config.cuda_graph_warmup_steps > 0 and should_toggle_forward_pre_hook: + # disable_forward_pre_hook(model, param_sync=False) + # cuda_graph_helper.create_cudagraphs() + # if model_config.cuda_graph_warmup_steps > 0 and should_toggle_forward_pre_hook: + # enable_forward_pre_hook(model) + # cuda_graph_helper.cuda_graph_set_manual_hooks() + + # Run training step. + timers("interval-time", log_level=0).start(barrier=False) + timer.start() + fault_tolerance.on_training_step_start(global_state) + _nemo_t0 = time.time() + ( + loss_dict, + skipped_iter, + should_checkpoint, + should_exit, + exit_code, + grad_norm, + num_zeros_in_grad, + log_max_attention_logit, + ) = train_step( + wrapped_forward_step_func, + train_data_iterator, + model, + optimizer, + scheduler, + global_state, + pg_collection, + forward_backward_func, + ) + nemo_style_iter_seconds_accum += time.time() - _nemo_t0 + nemo_style_iter_count += 1 + fault_tolerance.on_training_step_end(global_state) + timer.stop() + timers("interval-time").stop() + + # Advance NVIDIA DLFw Inspect step if enabled + tensor_inspect_step_if_enabled(config.tensor_inspect) + + if config.logger.log_throughput_to_tensorboard: + history_wct.append(time.time() - global_state.start_time) + if should_exit: + break + + # Enable forward pre-hooks after first set of forward and backward passes. + # When running in fp16, skip all NaN iterations until steady-state loss scaling value + # is reached. + if global_state.train_state.step == start_iteration: + if skipped_iter: + # Only enable forward pre-hook after a training step has successfully run. Relevant + # for fp16 codepath where first XX iterations are skipped until steady-state loss + # scale value is reached. + start_iteration = global_state.train_state.step + 1 + else: + # Enable forward pre-hook after training step has successfully run. All subsequent + # forward passes will use the forward pre-hook / `param_sync_func` in + # `forward_backward_func`. + if should_toggle_forward_pre_hook: + enable_forward_pre_hook(model) + model_config.param_sync_func = param_sync_func + pre_hook_enabled = True + # Set the manual hooks here since it's not set right after the capturing. + if ( + model_config.cuda_graph_impl == "transformer_engine" + and model_config.cuda_graph_warmup_steps == 0 + ): + assert cuda_graph_helper.graphs_created(), "CUDA Graphs should have been created." + cuda_graph_helper.cuda_graph_set_manual_hooks() + + global_state.train_state.step += 1 + + # MXFP4 healing: when train_state.step + 1 == HEALING_ITER, restore FP8 + # weights from the CPU stash (see ``mxfp4_healing``) and + # switch ``megatron.core.fp4_utils`` out of the MXFP4 phase. No-op when + # HEALING_ITER == 0 (default) so BF16/MXFP8 submission runs are + # unaffected. + try: + from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b import ( + mxfp4_healing as _mxh, + ) + + if _mxh.healing_iter() > 0: + _mxh.apply_healing_after_step(model, model_config, global_state.train_state.step) + _mxh.log_training_step_phase(global_state.train_state.step) + except Exception as _heal_err: # noqa: BLE001 + _orig_log_rank_0( + f"[mxfp4_healing] Failed at step={global_state.train_state.step}: " + f"{type(_heal_err).__name__}: {_heal_err}" + ) + raise + + # If fsdp_manual_registration is enabled, manually register FSDP communication buffers after one training step. + # if global_state.train_state.step == start_iteration + 1 and config.ddp.use_megatron_fsdp: + # _maybe_register_fsdp_buffers(config, model) + + global_state.train_state.consumed_train_samples += batch_size + num_skipped_samples_in_batch = ( + get_current_global_batch_size() - get_current_running_global_batch_size() + ) + if train_config.decrease_batch_size_if_needed: + assert num_skipped_samples_in_batch >= 0 + else: + assert num_skipped_samples_in_batch == 0 + global_state.train_state.skipped_train_samples += num_skipped_samples_in_batch + num_floating_point_operations_in_batch = flop_utils.num_floating_point_operations(config, batch_size) + global_state.train_state.floating_point_operations_so_far += num_floating_point_operations_in_batch + global_state.train_state.floating_point_operations_so_far + num_floating_point_operations_since_last_log_event += num_floating_point_operations_in_batch + + # Logging. + if hasattr(optimizer, "is_stub_optimizer") and not optimizer.is_stub_optimizer: + loss_scale = optimizer.get_loss_scale().item() + else: + loss_scale = 1.0 + params_norm = None + + if config.logger.log_params_norm: + params_norm = calc_params_l2_norm( + model, model_config, use_megatron_fsdp=config.dist.use_megatron_fsdp + ) + learning_rate = None + decoupled_learning_rate = None + for param_group in optimizer.param_groups: + if len(param_group) == 0: + continue + if param_group["is_decoupled_lr"]: + decoupled_learning_rate = param_group["lr"] + else: + learning_rate = param_group["lr"] + nemo_elapsed_time_per_iter_sec = None + if ( + config.logger.log_interval + and global_state.train_state.step % config.logger.log_interval == 0 + and nemo_style_iter_count > 0 + ): + nemo_elapsed_time_per_iter_sec = nemo_style_iter_seconds_accum / nemo_style_iter_count + report_memory_flag = training_log( + loss_dict, + total_loss_dict, + learning_rate, + decoupled_learning_rate, + loss_scale, + report_memory_flag, + skipped_iter, + grad_norm, + params_norm, + num_zeros_in_grad, + config, + global_state, + history_wct, + model, + log_max_attention_logit, + nemo_elapsed_time_per_iter_sec=nemo_elapsed_time_per_iter_sec, + ) + if config.logger.log_interval and global_state.train_state.step % config.logger.log_interval == 0: + _log_training_gpu_mem( + f"step {global_state.train_state.step}", + config.logger.memory_keys, + ) + if nemo_elapsed_time_per_iter_sec is not None: + nemo_style_iter_seconds_accum = 0.0 + nemo_style_iter_count = 0 + + # MLPerf logging: per-step train loss + if sft_logger is not None and not skipped_iter and loss_dict: + sft_logger.on_train_step( + step=global_state.train_state.step, + loss_dict=loss_dict, + lr=learning_rate, + consumed_samples=global_state.train_state.consumed_train_samples, + ) + + if ( + global_state.train_state.do_valid + and train_config.eval_interval + and global_state.train_state.step % train_config.eval_interval == 0 + and global_state.train_state.step > eval_skip_first_n * train_config.eval_interval + ): + if energy_monitor is not None: + energy_monitor.pause() + if should_toggle_forward_pre_hook: + disable_forward_pre_hook(model) + pre_hook_enabled = False + if train_config.manual_gc and train_config.manual_gc_eval: + # Collect all objects. + gc.collect() + prefix = f"iteration {global_state.train_state.step}" + timers("eval-time", log_level=0).start(barrier=True) + assert ( + timer.consumed_samples == global_state.train_state.consumed_train_samples + ), "Timer and global_state sample mismatch" + + if sft_logger is not None: + sft_logger.on_eval_start(global_state.train_state.consumed_train_samples) + + should_exit, eval_loss_value = evaluate_and_print_results_custom( + global_state, + prefix, + forward_step_func, + valid_data_iterator, + model, + model_config, + verbose=False, + write_to_tensorboard=False, + process_non_loss_data_func=process_non_loss_data_func, + non_loss_data_func=non_loss_data_func, + throughput=timer.get_throughput(), + ) + + if sft_logger is not None and eval_loss_value is not None: + target_hit = sft_logger.on_eval_end( + global_state.train_state.consumed_train_samples, + eval_loss_value, + ) + if target_hit: + should_exit = True + if should_exit: + exit_code = 0 + eval_duration += timers("eval-time").elapsed() + eval_iterations += train_config.eval_iters + timers("eval-time").stop() + + if train_config.manual_gc and train_config.manual_gc_eval: + # Collect only the objects created and used in evaluation. + gc.collect(generation=0) + if should_toggle_forward_pre_hook: + enable_forward_pre_hook(model) + pre_hook_enabled = True + if energy_monitor is not None: + energy_monitor.resume() + + # Miscellaneous post-training-step functions (e.g., FT heartbeats, GC). + # Some of these only happen at specific iterations. + maybe_synchronize_training_step(config.train.train_sync_interval, global_state.train_state.step) + num_floating_point_operations_since_last_log_event = maybe_report_stragglers( + config.logger.log_interval, + bool(getattr(config.straggler, "log_straggler", False)), + straggler_timer, + global_state.train_state.step, + num_floating_point_operations_since_last_log_event, + ) + maybe_check_weight_hash_across_dp_replicas( + model, + config.train.check_weight_hash_across_dp_replicas_interval, + global_state.train_state.step, + should_toggle_forward_pre_hook, + ) + handle_profiling_stop( + config.profiling, + global_state.train_state.step, + torch.distributed.get_rank(), + prof, + nsys_nvtx_context, + ) + maybe_run_manual_gc( + config.train.manual_gc, + config.train.manual_gc_interval, + global_state.train_state.step, + ) + if should_exit: + break + + training_wall_elapsed_s = time.perf_counter() - training_wall_start + log_rank_0( + f"Training loop finished: wall time {training_wall_elapsed_s:.2f} s " + f"({training_wall_elapsed_s / 60.0:.2f} min); " + f"final_iteration={global_state.train_state.step}; " + f"consumed_train_samples={global_state.train_state.consumed_train_samples}" + ) + + # MLPerf logging: end of training + if sft_logger is not None: + sft_logger.log_run_stop(global_state.train_state.consumed_train_samples) + + _delete_cuda_graphs(cuda_graph_helper) + + # Stop RPD profiler if still running + if rpd and rpd_status == "running": + log_rank_0(f"Stopping RPD profiling at training end (step {global_state.train_state.step})") + rpd.stop() + rpd_status = "finished" + + # Flush TensorBoard writer if present (TB/WandB disabled in LoggerConfig for this recipe). + writer = global_state.tensorboard_logger + if writer: + writer.flush() + + # Close out pre-hooks if using distributed optimizer and overlapped param gather. + if pre_hook_enabled: + disable_forward_pre_hook(model) + + # This will finalize all unfinalized async request and terminate + # a persistent async worker if persistent ckpt worker is enabled + fault_tolerance.on_checkpointing_start(global_state) + maybe_finalize_async_save( + global_state=global_state, ckpt_cfg=config.checkpoint, blocking=True, terminate=True + ) + fault_tolerance.on_checkpointing_end(global_state=global_state, is_async_finalization=True) + + # Shutdown NVRx straggler detection if enabled + safe_shutdown_nvrx_straggler_manager(global_state.nvrx_straggler_manager) + + if energy_monitor is not None: + energy_monitor.lap() + total_energy = energy_monitor.get_total() + log_rank_0(f"Total training energy (GPU): {total_energy / 1e6} MJ") + energy_monitor.shutdown() + + # If any exit conditions (signal handler, duration, iterations) have been reached, exit. + if should_exit: + # Stop RPD profiler if still running before exit + if rpd and rpd_status == "running": + log_rank_0(f"Stopping RPD profiling before exit at step {global_state.train_state.step}") + rpd.stop() + rpd_status = "finished" + + # Close NVIDIA DLFw Inspect if enabled + tensor_inspect_end_if_enabled(config.tensor_inspect) + maybe_finalize_async_save( + global_state=global_state, ckpt_cfg=config.checkpoint, blocking=True, terminate=True + ) + wandb_writer = global_state.wandb_logger + if wandb_writer: + wandb_writer.finish() + fault_tolerance.shutdown(global_state) + if exit_code != 0: + sys.exit(exit_code) + + # Close NVIDIA DLFw Inspect at clean finish + tensor_inspect_end_if_enabled(config.tensor_inspect) + + +# --------------------------------------------------------------------------- +# Install ``megatron_bridge_train_override`` on the megatron-bridge training +# modules. When ``PRE_QUANTIZED_MODEL`` is enabled, the override is first +# wrapped by ``install_pre_quantize_wrap`` so that the very first entry into +# the training loop pre-quantizes all TE Linear / LayerNormLinear weights to +# MXFP4 (and stashes FP8 copies on CPU for healing) before delegating to the +# override. When disabled (default), the override is installed directly (no +# wrap, zero runtime cost). +# --------------------------------------------------------------------------- +from megatron.bridge.training import train as _mb_train_mod + +from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.pre_quantize_mxfp4 import ( + install_pre_quantize_wrap, +) + +_installed_train = install_pre_quantize_wrap(megatron_bridge_train_override) +setattr(megatron_bridge_train_override, "_primus_llama2_custom_train_override", True) + +_mb_train_mod.train = _installed_train + +try: + from megatron.bridge.training import pretrain as _mb_pretrain_mod + + _mb_pretrain_mod.train = _installed_train +except ImportError: + pass + +try: + from megatron.bridge.training import finetune as _mb_finetune_mod + + if hasattr(_mb_finetune_mod, "train"): + _mb_finetune_mod.train = _installed_train +except ImportError: + pass + +# Keep the historical name `train` bound to the training module so any +# lingering `train.train = ...` expectations elsewhere resolve. +train = _mb_train_mod diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/mxfp4_healing.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/mxfp4_healing.py new file mode 100644 index 000000000..8bba1ee67 --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/mxfp4_healing.py @@ -0,0 +1,515 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""NeMo-parity MXFP4 -> FP8 healing for Primus / Megatron-Bridge. + +Reference: ``mlperf_code_llama2_70b_0430/src/callbacks/custom_callbacks.py`` +``CustomCallback`` (``_healing_setup``, ``_set_quantized_params_cpu``, +``_reset_full_iteration_cuda_graphs``, ``on_train_batch_end`` healing +branch). This module mirrors that logic byte-for-byte; only the trigger +surface differs: + +* NeMo's pre-quantize fires from ``on_train_start``; healing fires from + ``on_train_batch_end`` when ``trainer.global_step + 1 == healing_iter``. +* Primus's pre-quantize fires from a megatron-bridge ``train()`` wrapper + (``pre_quantize_mxfp4``); healing fires when the recipe + calls ``apply_healing_after_step(model, model_config, train_state.step)`` + immediately after ``train_state.step += 1`` -- the same arithmetic + identity as NeMo. + +All NeMo-side state lives on ``CustomCallback`` (``self.fp8_cpu_params``, +``self.healing_lambda``); Primus has no callback instance, so we use module +state. Configuration that NeMo reads from OmegaConf ``cfg.model.*`` is +read here from environment variables with NeMo's same defaults +(``HEALING_ITER`` / ``HEALING_PRECISION`` / ``ENABLE_TRANSPOSE_CACHE`` / +``RESET_CG_AFTER_HEALING`` / ``FIRST_LAST_LAYERS_BF16`` / +``NUM_LAYERS_AT_{START,END}_IN_BF16`` / ``STORE_GPU``). +""" + +from __future__ import annotations + +import os +from typing import Any, List + +import torch + +from primus.core.utils.module_utils import log_rank_0 + +# --------------------------------------------------------------------------- +# Module state. +# --------------------------------------------------------------------------- + +# Layered FP8 (E4M3) CPU stash, populated from outside via +# ``set_fp8_cpu_params`` (the bridge from ``pre_quantize_mxfp4._FP8_CPU_PARAMS``) +# in the same shape as NeMo's ``CustomCallback.fp8_cpu_params``: +# ``List[List[Float8Tensor]]`` indexed first by ``decoder.layers`` index and +# then by intra-layer iteration over TE Linear / LayerNormLinear modules. +_FP8_CPU_PARAMS: List[List[Any]] = [] +_HEALING_APPLIED: bool = False + + +# --------------------------------------------------------------------------- +# Env helpers (NeMo ``custom_callbacks._env_enabled`` parity). +# --------------------------------------------------------------------------- + + +def _env_enabled(name: str) -> bool: + """NeMo ``custom_callbacks.py:37-43``: missing/unrecognized = False.""" + val = os.getenv(name, "").strip().lower() + if val in {"1", "true", "yes", "on"}: + return True + return False + + +def healing_iter() -> int: + """``cfg.model.healing_iter`` analog (0 = healing disabled).""" + raw = os.getenv("HEALING_ITER", "0") + try: + return max(0, int(raw.strip())) + except ValueError: + return 0 + + +def healing_precision() -> str: + """``cfg.model.healing_precision`` analog: ``FP8_DS`` (default) or ``MXFP8``.""" + return os.getenv("HEALING_PRECISION", "FP8_DS").strip().upper() + + +# --------------------------------------------------------------------------- +# Model traversal helpers. +# --------------------------------------------------------------------------- + + +def _extract_module(model): + """NeMo ``custom_callbacks.py:92-97``. Adapted: Primus passes + ``List[ModelChunk]``; pick chunk 0 then strip ``.module``. + """ + m = model + if isinstance(m, (list, tuple)): + m = m[0] + while hasattr(m, "module"): + m = m.module + return m + + +def _all_modules(model): + chunks = list(model) if isinstance(model, (list, tuple)) else [model] + for ch in chunks: + yield from ch.modules() + + +def _te_linear_types(): + import transformer_engine.pytorch as te + + types = (te.Linear, te.LayerNormLinear) + if hasattr(te, "LayerNormMLP"): + types = types + (te.LayerNormMLP,) + return types + + +# --------------------------------------------------------------------------- +# NeMo ``_healing_setup`` / ``healing_lambda``. +# --------------------------------------------------------------------------- + + +def _build_healing_lambda(model_config: Any): + """NeMo ``CustomCallback._healing_setup`` (``custom_callbacks.py:176-187``). + + Reads ``fp8_amax_history_len`` / ``fp8_amax_compute_algo`` / + ``fp8_reduce_amax`` / ``fp8_dot_product_attention`` from the live + ``TransformerConfig`` (NeMo reads them from OmegaConf ``cfg.model``; + Megatron-Bridge stores the same fields directly on ``model_config``). + NeMo defaults: ``amax_history_len=4``, ``amax_compute_algo="most_recent"``, + ``reduce_amax=False``, ``fp8_dpa=False``. + """ + import transformer_engine.common.recipe as _recipe + + prec = healing_precision() + if prec == "FP8_DS": + amax_len = int(getattr(model_config, "fp8_amax_history_len", 4) or 4) + amax_algo = str(getattr(model_config, "fp8_amax_compute_algo", "most_recent") or "most_recent") + reduce_amax = bool(getattr(model_config, "fp8_reduce_amax", False)) + fp8_dpa = bool(getattr(model_config, "fp8_dot_product_attention", False)) + + log_rank_0( + "[mxfp4_healing] Healing recipe = DelayedScaling(" + f"amax_history_len={amax_len}, " + f"amax_compute_algo={amax_algo!r}, " + f"reduce_amax={reduce_amax}, " + f"fp8_dpa={fp8_dpa})" + ) + + def _delayed(_config: Any): + return _recipe.DelayedScaling( + amax_history_len=amax_len, + amax_compute_algo=amax_algo, + reduce_amax=reduce_amax, + fp8_dpa=fp8_dpa, + ) + + return _delayed + + if prec == "MXFP8": + log_rank_0("[mxfp4_healing] Healing recipe = MXFP8BlockScaling() (TE defaults)") + + def _mxfp8(_config: Any): + return _recipe.MXFP8BlockScaling() + + return _mxfp8 + + raise ValueError(f"Unsupported HEALING_PRECISION={prec!r} (expected FP8_DS or MXFP8)") + + +# --------------------------------------------------------------------------- +# NeMo ``_set_quantized_params_cpu`` (custom_callbacks.py:293-355). +# --------------------------------------------------------------------------- + + +def _set_quantized_params_cpu(model: Any, cpu_params: List[List[Any]], qtype: str) -> int: + """Restore CPU-stashed quantized weights to GPU and reinstall on TE modules. + + Byte-equivalent to NeMo ``CustomCallback._set_quantized_params_cpu`` + (``custom_callbacks.py:293-355``). Iterates ``decoder.layers`` in the same + order as ``_get_quantized_params_cpu`` (which built ``cpu_params``), so the + positional index ``cpu_params[layer_idx][param_idx]`` (consumed via + ``.pop(0)``) lines up with the ``param_idx``-th TE Linear / + LayerNormLinear / LayerNormMLP module encountered inside + ``layers[layer_idx].named_modules()``. + + NeMo's two-pass design (set ``module._parameters['weight']=None`` for all + target modules first, then assign the restored weight) is preserved: + decoupling the unlink from the install lets the caching allocator release + the old MXFP4 weight storage (~75 GiB) before the FP8 GPU side is wired + in (~64 GiB), preventing both sets resident at peak. + """ + extracted = _extract_module(model) + layers = extracted.decoder.layers + layer_count = len(layers) + dev = torch.cuda.current_device() + + first_last = _env_enabled("FIRST_LAST_LAYERS_BF16") + n_start = int(os.getenv("NUM_LAYERS_AT_START_IN_BF16", "0") or "0") + n_end = int(os.getenv("NUM_LAYERS_AT_END_IN_BF16", "0") or "0") + store_gpu = _env_enabled("STORE_GPU") + enable_tc = _env_enabled("ENABLE_TRANSPOSE_CACHE") + use_fuser = bool(getattr(extracted.config, "use_transformer_engine_op_fuser", False)) + + te_types = _te_linear_types() + + torch.cuda.synchronize() + torch.cuda.empty_cache() + + xfer_stream = torch.cuda.Stream() + n_restored = 0 + + for layer_idx, layer in enumerate(layers): + if first_last: + if layer_idx < n_start or layer_idx >= layer_count - n_end: + continue + + target_modules = [] + for _name, module in layer.named_modules(): + if not isinstance(module, te_types): + continue + if not hasattr(module, "weight"): + continue + target_modules.append(module) + + with torch.no_grad(): + for module in target_modules: + module._parameters["weight"] = None + + for module in target_modules: + with torch.no_grad(): + weight = cpu_params[layer_idx].pop(0) + + if not store_gpu: + with torch.cuda.stream(xfer_stream): + if qtype == "FP8_DS": + weight._data = weight._data.to(dev, non_blocking=True) + if weight._transpose is not None: + if enable_tc: + weight._transpose = weight._transpose.to(dev, non_blocking=True) + else: + weight._transpose = None + elif qtype == "MXFP4": + weight._rowwise_data = weight._rowwise_data.to(dev, non_blocking=True) + if weight._columnwise_data is not None: + weight._columnwise_data = weight._columnwise_data.to(dev, non_blocking=True) + else: + raise ValueError(f"Unsupported quantization type: {qtype}") + + weight.requires_grad = False + module._parameters["weight"] = weight + n_restored += 1 + + if use_fuser: + mlp = getattr(layer, "mlp", None) + if mlp is not None and hasattr(mlp, "_make_fused_impl"): + mlp._fused_impl = (mlp._make_fused_impl(),) + sa = getattr(layer, "self_attention", None) + if sa is not None: + for attr in ("linear_proj", "linear_qkv"): + mod = getattr(sa, attr, None) + if mod is not None and hasattr(mod, "_make_fused_branches"): + mod._fused_branches = mod._make_fused_branches() + + xfer_stream.synchronize() + return n_restored + + +# --------------------------------------------------------------------------- +# NeMo ``custom_llama.reset_fp8_state`` (custom_llama.py:77-92). +# --------------------------------------------------------------------------- + + +def _reset_fp8_state(model: Any) -> None: + """Force TE to re-initialize ``fp8_meta`` under the new recipe on next forward. + + Without this, modules keep MXFP4 ``RecipeState`` while the recipe and + weights changed under them; the next GEMM dereferences scale/amax + pointers from the wrong RecipeState type. Mirrors NeMo's + ``custom_llama.reset_fp8_state``: MX recipes (MXFP4/MXFP8) use + ``RecipeState`` without ``.scale``, and TE's ``reset_fp8_meta_tensors`` + assumes DelayedScaling-style state and raises against MX, so we detect + the MX case by inspecting ``fp8_meta`` and skip the reset for those + modules. + """ + + def reset_fp8(m: Any) -> None: + if not hasattr(m, "fp8_initialized"): + return + m.fp8_initialized = False + # MX recipes (MXFP4/MXFP8) use RecipeState without `.scale`; TE's + # reset_fp8_meta_tensors assumes DelayedScaling-style state and raises. + fp8_meta = getattr(m, "fp8_meta", None) or {} + for key in ("scaling_fwd", "scaling_bwd"): + state = fp8_meta.get(key) + if state is not None and not hasattr(state, "scale"): + return + m.reset_fp8_meta_tensors() + + for m in _all_modules(model): + reset_fp8(m) + + +# --------------------------------------------------------------------------- +# NeMo ``_reset_full_iteration_cuda_graphs`` (custom_callbacks.py:148-174). +# --------------------------------------------------------------------------- + + +def _reset_full_iteration_cuda_graphs() -> None: + """Drop Megatron full-iteration CUDA graph captures before the recipe swap. + + After MXFP4 -> FP8 weight swap and recipe change, replaying old graphs + is invalid. ``RESET_CG_AFTER_HEALING`` (NeMo: ``cfg.model.reset_cg_after_healing``) + additionally resets per-stage warmup step counters so + ``FullCudaGraphWrapper`` re-runs ``cuda_graph_warmup_steps`` before + re-capture. + """ + try: + from megatron.core.full_cuda_graph import FullCudaGraphWrapper + except ImportError: + return + + FullCudaGraphWrapper.cuda_graph["training"] = None + FullCudaGraphWrapper.cuda_graph["validation"] = None + FullCudaGraphWrapper.result["training"] = None + FullCudaGraphWrapper.result["validation"] = None + + if _env_enabled("RESET_CG_AFTER_HEALING"): + FullCudaGraphWrapper.curr_iteration["training"] = 0 + FullCudaGraphWrapper.curr_iteration["validation"] = 0 + + +# --------------------------------------------------------------------------- +# Public stash bridge. +# --------------------------------------------------------------------------- + + +def set_fp8_cpu_params(layered: List[List[Any]]) -> None: + """Hand off NeMo-style ``List[List[Float8Tensor]]`` from pre-quantize. + + Mirrors the assignment ``self.fp8_cpu_params = self._get_quantized_params_cpu(...)`` + in NeMo ``CustomCallback._pre_quantize_model``. Storing a reference (not + a copy) means ``_set_quantized_params_cpu``'s ``.pop(0)`` mutates the same + list as the caller, identical to NeMo. + """ + global _FP8_CPU_PARAMS, _HEALING_APPLIED + _FP8_CPU_PARAMS = layered + _HEALING_APPLIED = False + + +# --------------------------------------------------------------------------- +# NeMo ``on_train_batch_end`` healing branch (custom_callbacks.py:620-677). +# --------------------------------------------------------------------------- + + +def apply_healing_after_step(model: Any, model_config: Any, train_state_step: int) -> None: + """Run NeMo's ``on_train_batch_end`` healing branch. + + Trigger condition mirrors NeMo: + ``if trainer.global_step + 1 == self.cfg.model.healing_iter:`` + Primus's ``train_state.step`` is incremented before this call (in the + recipe's ``megatron_bridge_train_override``), so it carries the same + value as NeMo's ``trainer.global_step``. + + Steps below are 1:1 with the NeMo branch. + """ + global _HEALING_APPLIED + + hi = healing_iter() + if hi <= 0 or _HEALING_APPLIED: + return + if train_state_step + 1 != hi: + return + + if not _FP8_CPU_PARAMS: + log_rank_0( + "[mxfp4_healing] HEALING_ITER is set but the FP8 CPU stash is empty; " + "either PRE_QUANTIZED_MODEL=True is missing or pre-quantize did not run. " + "Skipping healing." + ) + _HEALING_APPLIED = True + return + + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + dist.barrier() + + # 1. NeMo: ``self._reset_full_iteration_cuda_graphs(trainer)``. + _reset_full_iteration_cuda_graphs() + + # 2. NeMo: ``self._set_quantized_params_cpu(trainer.model, self.fp8_cpu_params, self.healing_precision)``. + n_restored = _set_quantized_params_cpu(model, _FP8_CPU_PARAMS, healing_precision()) + torch.cuda.synchronize() + torch.cuda.empty_cache() + + # 3. NeMo: patch ``megatron.core.fp4_utils`` to the healing recipe and clear ``_mxfp4_phase``. + import megatron.core.fp4_utils as fp4u + + fp4u.get_fp4_recipe = _build_healing_lambda(model_config) + fp4u._mxfp4_phase = False + + # 4. NeMo: when ``ENABLE_TRANSPOSE_CACHE`` is off, tell every TE Linear to stop + # keeping a transpose cache (custom_callbacks.py:637-642). + if not _env_enabled("ENABLE_TRANSPOSE_CACHE"): + te_types = _te_linear_types() + for m in _all_modules(model): + if isinstance(m, te_types): + m.keep_fp8_weight_transpose_cache = False + + # 5. NeMo: drop activation recompute on the live TransformerConfig + # (custom_callbacks.py:644-648). + extracted = _extract_module(model) + cfg = extracted.config + if cfg.recompute_granularity is not None: + cfg.recompute_granularity = None + cfg.recompute_num_layers = None + cfg.recompute_method = None + + # 6. NeMo: clear MXFP4 caches on every module (custom_callbacks.py:650-656). + for m in _all_modules(model): + if hasattr(m, "_mxfp4_weight_cache"): + del m._mxfp4_weight_cache + if hasattr(m, "_mxfp4_persist_columnwise"): + del m._mxfp4_persist_columnwise + + # 6b. Primus delta vs NeMo: clear ``save_original_input`` on every TE module. + # + # Megatron-LM shipped in Primus's tree (third_party/Megatron-Bridge/3rdparty/ + # Megatron-LM) sets ``module.save_original_input = True`` on every TE + # ``linear_proj`` (and some ``mlp.linear_fc1``) when ``config.fp4`` is + # active -- it's an MX-only optimization that lets TE re-quantize at + # backward time from the saved BF16 input (legal under MXFP4BlockScaling + # because the per-block scale is reproducible from the BF16 input). + # ``DelayedScaling`` derives its FP8 scale from amax history that is + # updated each step, so the saved BF16 input cannot reproduce the + # forward FP8 input. TE guards this combination with a hard error:: + # + # RuntimeError: DelayedScaling recipe is not supported with save_original_input + # + # ...which fires on the very next forward after we patch + # ``get_fp4_recipe -> FP8_DS``. NeMo MLPerf doesn't hit this because + # their (older) Megatron-LM lacks the FP4 branch that sets the flag in + # the first place. We reset it defensively to mirror NeMo's behavior. + n_save_orig_cleared = 0 + for m in _all_modules(model): + if getattr(m, "save_original_input", False): + try: + m.save_original_input = False + n_save_orig_cleared += 1 + except Exception: # noqa: BLE001 + pass + + # 7. NeMo: ``reset_fp8_state(trainer.lightning_module)`` (custom_callbacks.py:668). + _reset_fp8_state(model) + + torch.cuda.synchronize() + torch.cuda.empty_cache() + + _HEALING_APPLIED = True + + log_rank_0( + f"[mxfp4_healing] healing applied at train_state.step={train_state_step} " + f"(step+1==HEALING_ITER={hi}); recipe={healing_precision()}; " + f"restored {n_restored} FP8 weights; " + f"cleared save_original_input on {n_save_orig_cleared} TE modules." + ) + + +# --------------------------------------------------------------------------- +# Status helpers + test hygiene. +# --------------------------------------------------------------------------- + + +def is_fp8_healing_phase() -> bool: + """True after ``apply_healing_after_step`` has fired (TE weights are FP8).""" + return _HEALING_APPLIED + + +def reset_healing_state() -> None: + """Test / multi-run hygiene: forget the stash and the healed flag.""" + global _FP8_CPU_PARAMS, _HEALING_APPLIED + _FP8_CPU_PARAMS = [] + _HEALING_APPLIED = False + + +# --------------------------------------------------------------------------- +# Backwards-compatible shims for existing callers. +# +# These keep ``pre_quantize_mxfp4`` and ``llama2_custom`` importable without +# edits. The new +# code is ``set_fp8_cpu_params`` (above); the legacy names below just +# delegate / no-op. +# --------------------------------------------------------------------------- + + +def _set_ordered_fp8_stash_from_layered(model: Any, fp8_cpu_params: List[List[Any]]) -> None: + """Legacy alias for ``set_fp8_cpu_params``. ``model`` is unused; we keep + the layered stash exactly as NeMo does. + """ + del model # unused; layered stash is enough (NeMo parity) + set_fp8_cpu_params(fp8_cpu_params) + + +def log_healing_env_banner_once() -> None: + """No-op (NeMo doesn't have this; logging removed from main flow).""" + return + + +def log_training_step_phase(train_state_step: int) -> None: + """No-op (NeMo doesn't have this; logging removed from main flow).""" + del train_state_step + return + + +def run_fp8_warmup_for_kernel_jit(*_args: Any, **_kwargs: Any) -> None: + """No-op. NeMo's ``_pre_ttt_fp8_warmup`` call site is commented out + (custom_callbacks.py:509-510), so the warmup is dead in NeMo too. + """ + return diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/nemo_loss.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/nemo_loss.py new file mode 100644 index 000000000..d8a0881b7 --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/nemo_loss.py @@ -0,0 +1,591 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +NeMo-equivalent training/validation loss for Primus (exact port of +``MaskedTokenLossReduction.forward + reduce``). + +Background +---------- +NeMo MLPerf (`mlperf_code_llama2_70b_0430/src/custom_llama.py`) overrides +``training_loss_reduction`` / ``validation_loss_reduction`` on the model +with ``MaskedTokenLossReduction``. The class has **two distinct paths** +depending on its constructor flags: + +1. ``validation_step=False`` (training) **or** ``val_drop_last=True``: + * per microbatch: per-sample token-mean + ``loss_for_ub[i] = sum(losses[i] * mask[i]) / sum(mask[i])`` + * across microbatches: ``concat(per_sample_means).mean()`` +2. ``validation_step=True`` AND ``val_drop_last=False`` (NeMo's default for + eval; see ``CustomLlamaModel.validation_loss_reduction``): + * per microbatch: per-sample token-mean (same as path 1) **plus** a + 0/1 indicator ``num_valid_tokens_in_ub = (num_valid_tokens > 0)``, + stacked to produce ``[loss_for_ub, valid_indicator]``. + * across microbatches: ``vstack(stacked).sum(dim=0)`` → returns + ``[total_loss_sum, total_valid_sample_count]``. + * the eval-time CP all-reduce on ``loss_mask`` / ``loss_for_ub`` is + gated by ``disabled_cp_for_eval`` (i.e., CP is **off** for eval). + +This is **not** equivalent to Megatron-Bridge's default +``masked_next_token_loss`` which computes a *global* token-mean across all +tokens in all samples. When samples have very different numbers of +unmasked tokens (typical for SFT/LoRA runs), the two objectives differ +and produce different gradient directions. + +Implementation +-------------- +We expose a NeMo-equivalent per-microbatch loss function that returns the +3-tuple Megatron-Core's pipeline schedule expects:: + + (loss, num_tokens, {"lm loss": [loss_sum, count]}) + +For both paths, ``loss = loss_for_ub.sum()`` (sum over the microbatch's +per-sample means). The choice of ``num_tokens`` and the ``[sum, count]`` +reporting tensor depends on the path: + +* **Train / val_drop_last=True path**: ``num_tokens = num_samples = B``. + After Megatron-Bridge's ``output /= clamp(num_tokens, 1); output /= + num_microbatches`` and DP+microbatch ``[sum, count]`` aggregation, the + effective backward target and reported loss are both + ``mean over (microbatch, sample) of per_sample_mean`` -- byte-identical + to NeMo's ``concat(per_sample_means).mean()`` when ``B`` is uniform. + +* **val_drop_last=False path** (eval, NeMo default): ``num_tokens = + num_samples = sum_i (num_valid_tokens_i > 0)``. Empty-mask samples + (those whose entire sequence is masked) are excluded from the + denominator -- mirroring NeMo's ``num_valid_tokens_in_ub`` indicator. + Because ``loss_for_ub`` is already zeroed for those samples, the + numerator (sum) is unaffected. Megatron-Bridge's eval reduction does + ``total_loss_sum / total_valid_count`` across DP+CP+microbatch, which + matches NeMo's ``vstack(...).sum(dim=0)`` semantics. + +Train vs. eval detection +~~~~~~~~~~~~~~~~~~~~~~~~ +Megatron-Bridge has a single ``forward_step`` for both train and eval and +wraps the eval pass in ``with torch.no_grad():`` (see +``megatron/bridge/training/eval.py:118``). We detect eval at call time +with ``torch.is_grad_enabled()`` -- a Megatron-Bridge-specific but stable +signal. + +CP gating (``cp_eval`` / ``disabled_cp_for_eval``) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +NeMo's ``cp_eval`` parameter (``CustomLlamaModel.config.cp_eval``) toggles +``disabled_cp_for_eval``: when CP is **disabled** for eval, the CP +all-reduces on ``loss_mask`` and ``loss_for_ub`` must be **skipped** (the +data isn't actually sharded across CP ranks even though +``parallel_state.get_context_parallel_world_size() > 1``). We expose this +through env var ``PRIMUS_NEMO_LOSS_DISABLED_CP_FOR_EVAL`` (default ``0``, +i.e. CP-on-for-eval). + +``val_drop_last`` is exposed via ``PRIMUS_NEMO_LOSS_VAL_DROP_LAST`` +(default ``0``, matching NeMo's eval setup ``val_drop_last=False``). + +Reporting parity +---------------- +Megatron-Bridge's training & eval aggregation computes ``loss_sum / +count`` over (microbatch, DP-rank, CP-rank). With our return shape: + +* Train / val_drop_last=True: + ``sum_(rank, microbatch) sum_samples(per_sample_mean) / + sum_(rank, microbatch) B = mean per_sample_mean`` +* val_drop_last=False: + ``sum_(rank, microbatch) sum_samples(per_sample_mean) / + sum_(rank, microbatch) count_valid_samples + = total_loss_sum / total_valid_sample_count`` (matches NeMo). + +Activation +---------- +Auto-installs at import time. Disable by setting ``PRIMUS_NEMO_LOSS=0``. + +Env vars +~~~~~~~~ +* ``PRIMUS_NEMO_LOSS`` (default ``1``): master switch. +* ``PRIMUS_NEMO_LOSS_VAL_DROP_LAST`` (default ``0``): if ``1``, + treat eval like train (NeMo's ``val_drop_last=True``). +* ``PRIMUS_NEMO_LOSS_DISABLED_CP_FOR_EVAL`` (default ``0``): if ``1``, + skip the CP all-reduces during eval (NeMo's ``cp_eval`` set). +""" + +from __future__ import annotations + +import os +from functools import partial +from typing import Dict, List, Optional, Tuple, Union + +import torch +from megatron.core import parallel_state +from megatron.core.rerun_state_machine import get_rerun_state_machine + +from primus.core.utils.module_utils import log_rank_0 + + +def _safe_log_rank_0(msg: str) -> None: + """Log on rank 0; fall back to print if the global logger is not ready yet.""" + try: + log_rank_0(msg) + except AttributeError: + if os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0")) in ("0", "", None): + print(msg, flush=True) + + +# Same value Megatron-Bridge uses in ``masked_next_token_loss``. +_SPIKY_LOSS_FACTOR: int = 10 + +# Module-level flag so callers can verify the patch landed. +_INSTALLED: bool = False + + +def _enabled() -> bool: + """Return True when the NeMo-equivalent loss should be active.""" + flag = os.environ.get("PRIMUS_NEMO_LOSS", "1").strip().lower() + return flag not in ("0", "false", "no", "off") + + +def _env_bool(name: str, default: bool = False) -> bool: + """Parse a boolean env var (``1/true/yes/on`` -> True; anything else -> False).""" + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in ("1", "true", "yes", "on") + + +def _is_validation_step() -> bool: + """Detect whether we're inside a validation forward pass. + + Megatron-Bridge wraps the entire eval loop in + ``with torch.no_grad():`` (``megatron/bridge/training/eval.py:118``) + and reuses the same ``forward_step`` for both train and eval, so + ``torch.is_grad_enabled()`` is a stable, framework-internal signal + that distinguishes the two. + """ + return not torch.is_grad_enabled() + + +def _val_drop_last() -> bool: + """``val_drop_last`` flag, mirrored to NeMo's ``MaskedTokenLossReduction``. + + NeMo's ``CustomLlamaModel.validation_loss_reduction`` constructs the + eval reduction with ``val_drop_last=False`` (drop_last=False so all + eval samples count). Default here matches that. Set + ``PRIMUS_NEMO_LOSS_VAL_DROP_LAST=1`` to fall back to the train-style + averaging (``concat(per_sample_means).mean()`` over all microbatches). + """ + return _env_bool("PRIMUS_NEMO_LOSS_VAL_DROP_LAST", default=False) + + +def _disabled_cp_for_eval() -> bool: + """``disabled_cp_for_eval`` flag, mirrored to NeMo's ``cp_eval`` parameter. + + When CP is disabled for the eval forward pass (i.e., the model + doesn't actually shard across CP ranks during eval, even though + ``parallel_state.get_context_parallel_world_size() > 1``), the CP + all-reduces on ``loss_mask`` and ``loss_for_ub`` MUST be skipped -- + otherwise we double-count along the CP dimension. NeMo encodes this + via ``cp_eval is not None`` in ``MaskedTokenLossReduction.__init__``. + Default here is ``False`` (CP-on-for-eval); set + ``PRIMUS_NEMO_LOSS_DISABLED_CP_FOR_EVAL=1`` if your recipe disables + CP during eval. + """ + return _env_bool("PRIMUS_NEMO_LOSS_DISABLED_CP_FOR_EVAL", default=False) + + +def _cp_eval_value() -> Optional[int]: + """Return NeMo's ``cp_eval`` integer value for the validation reduction. + + NeMo's ``CustomLlamaModel.validation_loss_reduction`` passes + ``cp_eval=self.config.cp_eval``, which is the *integer* + ``cfg.model.eval_cp`` (or ``None``). The class only checks + ``cp_eval is not None`` to set ``disabled_cp_for_eval``, so the + integer value itself never affects the loss math -- but we keep the + int-or-None contract to remain byte-identical to NeMo. + + Resolution order: + 1. ``PRIMUS_NEMO_LOSS_CP_EVAL`` (preferred). Set to an integer + (e.g. ``1``) to mark "CP is disabled for eval". Empty / unset + leaves it as ``None``. + 2. Fallback: ``PRIMUS_NEMO_LOSS_DISABLED_CP_FOR_EVAL`` (boolean). + If set truthy, returns ``1``; else ``None``. Provided for + backwards compat with the function-based env var. + """ + raw = os.environ.get("PRIMUS_NEMO_LOSS_CP_EVAL") + if raw is not None and raw.strip(): + try: + return int(raw) + except ValueError: + return 1 if raw.strip().lower() in ("1", "true", "yes", "on") else None + return 1 if _disabled_cp_for_eval() else None + + +# ============================================================================ +# Byte-identical port of NeMo's ``MaskedTokenLossReduction`` class. +# +# Source: +# mlperf_code_llama2_70b_0430/src/custom_llama.py:160-221 +# Original imports replaced 1:1: +# nemo.lightning.megatron_parallel.MegatronLossReduction -> stub below +# megatron.core.parallel_state -> identical +# torch -> identical +# +# The class body is reproduced as-is so the per-microbatch loss math, the +# CP / DP all-reduce gating, and the train vs. val branching are +# guaranteed identical to NeMo. The Megatron-Bridge wiring (which +# expects a (loss, num_tokens, dict) 3-tuple instead of NeMo's +# (loss_for_ub, dict)) lives in ``nemo_masked_next_token_loss`` further +# down -- it instantiates the right train / val singleton, calls +# ``forward(batch, per_token_losses)``, and repackages the return. +# ============================================================================ + + +class MegatronLossReduction: + """Stub of ``nemo.lightning.megatron_parallel.MegatronLossReduction``. + + The NeMo base class is just an interface marker (its ``forward`` / + ``reduce`` are abstract). We reproduce it as an empty class so the + ported ``MaskedTokenLossReduction`` keeps the same MRO for any + callers that introspect via ``isinstance``. + """ + + +class MaskedTokenLossReduction(MegatronLossReduction): + """Byte-identical port of NeMo's ``MaskedTokenLossReduction``. + + Identical to ``custom_llama.MaskedTokenLossReduction`` (lines 160-221 + of ``mlperf_code_llama2_70b_0430/src/custom_llama.py``): + + * ``__init__`` signature, attribute names, and ``disabled_cp_for_eval`` + derivation are unchanged. + * ``forward(batch, per_token_losses)`` returns NeMo's 2-tuple + ``(loss_for_ub, dict)`` where the dict is either ``{"avg": + reduced_loss}`` (train / ``val_drop_last=True``) or + ``{"loss_sum_and_ub_size": <(B, 2) all-reduced tensor>}`` + (validation with ``val_drop_last=False``). + * ``reduce(losses_reduced_per_micro_batch)`` aggregates microbatch + outputs identically (``concat(...).mean()`` for ``"avg"``, + ``vstack(...).sum(dim=0)`` for ``"loss_sum_and_ub_size"``). + + To use this class with Megatron-Bridge's function-based loss API + (which expects ``(loss, num_tokens, {"lm loss": [sum, count]})``), + go through ``nemo_masked_next_token_loss`` -- it builds an instance + (one singleton for train, one for val), calls ``forward``, and + repackages the return. + """ + + def __init__( + self, + validation_step: bool = False, + val_drop_last: bool = True, + cp_eval: Optional[int] = None, + ) -> None: + super().__init__() + self.validation_step = validation_step + self.train_step = not validation_step + self.val_drop_last = val_drop_last + self.disabled_cp_for_eval = cp_eval is not None + + def forward( + self, + batch: Dict[str, torch.Tensor], + per_token_losses: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + if isinstance(per_token_losses, tuple): + per_token_losses, loss_mask = per_token_losses + batch["loss_mask"] = loss_mask + masked_losses = per_token_losses * batch["loss_mask"] + + cp_size = parallel_state.get_context_parallel_world_size() + if cp_size > 1 and (self.train_step or not self.disabled_cp_for_eval): + torch.distributed.all_reduce( + batch["loss_mask"], group=parallel_state.get_context_parallel_group() + ) + + num_valid_tokens = batch["loss_mask"].sum(1) + loss_for_ub = torch.sum(masked_losses, dim=1) / num_valid_tokens + loss_for_ub = torch.where(num_valid_tokens == 0, torch.zeros_like(loss_for_ub), loss_for_ub) + + if cp_size > 1 and (self.train_step or not self.disabled_cp_for_eval): + torch.distributed.all_reduce(loss_for_ub, group=parallel_state.get_context_parallel_group()) + + if self.validation_step and not self.val_drop_last: + num_valid_tokens_in_ub = (num_valid_tokens > 0).long() + loss_sum_and_ub_size_all_gpu = torch.stack( + [loss_for_ub.clone().detach(), num_valid_tokens_in_ub.clone().detach()], + dim=1, + ) + if self.disabled_cp_for_eval: + torch.distributed.all_reduce(loss_sum_and_ub_size_all_gpu) + else: + torch.distributed.all_reduce( + loss_sum_and_ub_size_all_gpu, + group=parallel_state.get_data_parallel_group(), + ) + return loss_for_ub, {"loss_sum_and_ub_size": loss_sum_and_ub_size_all_gpu} + + reduced_loss = loss_for_ub + return loss_for_ub, {"avg": reduced_loss} + + def reduce(self, losses_reduced_per_micro_batch: List[Dict[str, torch.Tensor]]) -> torch.Tensor: + if losses_reduced_per_micro_batch: + if "avg" in losses_reduced_per_micro_batch[0]: + loss_tensors_list = [loss_reduced["avg"] for loss_reduced in losses_reduced_per_micro_batch] + loss_tensor = torch.concat(loss_tensors_list) + return loss_tensor.mean() + + loss_sum_tensors_list: List[torch.Tensor] = [ + loss_sum["loss_sum_and_ub_size"] for loss_sum in losses_reduced_per_micro_batch + ] + loss_sum = ( + torch.vstack(loss_sum_tensors_list).sum(dim=0) + if len(loss_sum_tensors_list) > 0 + else torch.tensor([0.0, 0.0], device=torch.cuda.current_device()) + ) + return loss_sum + + return torch.tensor(0.0, device=torch.cuda.current_device()) + + +# ---------------------------------------------------------------------------- +# Lazy singletons mirroring NeMo's ``CustomLlamaModel`` setup: +# +# self._training_loss_reduction = MaskedTokenLossReduction() +# self._validation_loss_reduction = MaskedTokenLossReduction( +# validation_step=True, +# val_drop_last=False, +# cp_eval=self.config.cp_eval, +# ) +# +# (See ``mlperf_code_llama2_70b_0430/src/custom_llama.py:143-157``.) +# We construct lazily because env vars may not be set at import time. +# ---------------------------------------------------------------------------- + +_TRAIN_REDUCTION: Optional[MaskedTokenLossReduction] = None +_VAL_REDUCTION: Optional[MaskedTokenLossReduction] = None + + +def _get_train_reduction() -> MaskedTokenLossReduction: + """Return the per-process training ``MaskedTokenLossReduction`` singleton.""" + global _TRAIN_REDUCTION + if _TRAIN_REDUCTION is None: + _TRAIN_REDUCTION = MaskedTokenLossReduction() + return _TRAIN_REDUCTION + + +def _get_val_reduction() -> MaskedTokenLossReduction: + """Return the per-process validation ``MaskedTokenLossReduction`` singleton.""" + global _VAL_REDUCTION + if _VAL_REDUCTION is None: + _VAL_REDUCTION = MaskedTokenLossReduction( + validation_step=True, + val_drop_last=_val_drop_last(), + cp_eval=_cp_eval_value(), + ) + return _VAL_REDUCTION + + +def reset_loss_reduction_singletons() -> None: + """Drop cached singletons so subsequent calls re-read env vars. + + Useful for tests or for recipes that reconfigure ``cp_eval`` / + ``val_drop_last`` after import time. + """ + global _TRAIN_REDUCTION, _VAL_REDUCTION + _TRAIN_REDUCTION = None + _VAL_REDUCTION = None + + +def nemo_masked_next_token_loss( + loss_mask: torch.Tensor, + output_tensor: Union[torch.Tensor, Tuple[torch.Tensor, ...]], + check_for_nan_in_loss: bool = True, + check_for_spiky_loss: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: + """Megatron-Bridge adapter for the byte-identical ``MaskedTokenLossReduction``. + + All loss math (CP / DP all-reduces, val_drop_last branching, + train vs. eval gating) lives in the class above and is identical + to NeMo. This adapter does three things: + + 1. Picks the right per-process singleton (train vs. val) based on + ``torch.is_grad_enabled()`` -- Megatron-Bridge wraps eval in + ``with torch.no_grad():`` (``eval.py:118``) and reuses the same + ``forward_step`` for both, so this is the only stable signal. + 2. Calls ``MaskedTokenLossReduction.forward(batch, per_token_losses)`` + with a ``batch = {"loss_mask": loss_mask}`` shim. The class + mutates ``batch["loss_mask"]`` in-place via the CP all-reduce -- + same as NeMo, so we accept the side effect. + 3. Repackages NeMo's 2-tuple ``(loss_for_ub, dict)`` return into the + Megatron-Core schedule's expected 3-tuple + ``(loss, num_tokens, {"lm loss": [sum, count]})``: + + * ``"avg"`` branch (train / ``val_drop_last=True``): + - ``loss = loss_for_ub.sum()``, ``num_tokens = B``. + Megatron-Core computes ``loss / clamp(num_tokens, 1) / + num_microbatches`` for backward, giving + ``mean over (microbatch, sample) of per_sample_mean`` -- + byte-identical to NeMo's ``concat(...).mean()`` reduction + when ``B`` is uniform across microbatches and DP ranks. + * ``"loss_sum_and_ub_size"`` branch (val w/ ``val_drop_last=False``): + - The class already DP-all-reduced the (B, 2) tensor inside + ``forward``. We sum dim=0 (matches NeMo's ``vstack(...).sum(dim=0)`` + reduce) to get ``[total_loss_sum, total_valid_count]`` for + this microbatch, then return it as the reporting tensor. + Megatron-Bridge's eval reduction then sums it across + microbatches and DP+CP (``eval.py:183``); since both + entries scale by the same factor in any further reduction, + the final ratio ``loss_sum / count`` is identical to + NeMo's reported eval loss. + + NaN / Inf / spike validation runs on the local per-sample sum so it + triggers regardless of which branch the class took. + """ + if isinstance(output_tensor, tuple): + per_token_losses, dynamic_mask = output_tensor + loss_mask = dynamic_mask + else: + per_token_losses = output_tensor + + per_token_losses = per_token_losses.float() + loss_mask = loss_mask.float() + + batch_size = loss_mask.shape[0] + per_token_losses = per_token_losses.view(batch_size, -1) + loss_mask = loss_mask.view(batch_size, -1) + + reduction = _get_val_reduction() if _is_validation_step() else _get_train_reduction() + + batch: Dict[str, torch.Tensor] = {"loss_mask": loss_mask} + loss_for_ub, reduced_dict = reduction.forward(batch, per_token_losses) + + loss_sum = loss_for_ub.sum() + + rerun_state_machine = get_rerun_state_machine() + if check_for_nan_in_loss: + rerun_state_machine.validate_result( + result=loss_sum, + rejection_func=torch.isnan, + message="found NaN in local forward loss calculation", + tolerance=0.0, + fatal=True, + ) + rerun_state_machine.validate_result( + result=loss_sum, + rejection_func=torch.isinf, + message="found Inf in local forward loss calculation", + tolerance=0.0, + fatal=True, + ) + if check_for_spiky_loss: + rerun_state_machine.validate_result( + result=loss_sum, + rejection_func=partial( + rerun_state_machine.is_unexpectedly_large, + threshold=_SPIKY_LOSS_FACTOR, + context="loss", + ), + message="Spiky loss", + tolerance=0.0, + fatal=False, + ) + + if "loss_sum_and_ub_size" in reduced_dict: + agg = reduced_dict["loss_sum_and_ub_size"].sum(dim=0).float() + agg_loss_sum = agg[0] + agg_count = agg[1] + num_samples = agg_count.to(torch.int) + reporting_loss = torch.cat( + [ + agg_loss_sum.clone().detach().view(1), + agg_count.clone().detach().view(1), + ] + ) + return (agg_loss_sum, num_samples, {"lm loss": reporting_loss}) + + num_samples = torch.tensor(loss_for_ub.numel(), dtype=torch.int, device=loss_for_ub.device) + reporting_loss = torch.cat( + [ + loss_sum.clone().detach().view(1), + num_samples.detach().view(1).float(), + ] + ) + return (loss_sum, num_samples, {"lm loss": reporting_loss}) + + +def create_nemo_masked_next_token_loss_function( + loss_mask: torch.Tensor, + check_for_nan_in_loss: bool, + check_for_spiky_loss: bool, +) -> partial: + """Factory matching Megatron-Bridge's ``create_masked_next_token_loss_function``.""" + return partial( + nemo_masked_next_token_loss, + loss_mask, + check_for_nan_in_loss=check_for_nan_in_loss, + check_for_spiky_loss=check_for_spiky_loss, + ) + + +def install_nemo_loss_if_enabled() -> bool: + """Monkey-patch Megatron-Bridge step modules to use the NeMo-equivalent loss. + + Affected bindings: + + * ``megatron.bridge.training.vlm_step._create_loss_function`` -> + ``create_nemo_masked_next_token_loss_function`` (used by + ``MegatronBridgePosttrainTrainer.train()`` -> ``finetune(..., + forward_step_func=vlm_step.forward_step)``). + * ``megatron.bridge.training.gpt_step.masked_next_token_loss`` -> + ``nemo_masked_next_token_loss`` (used by other gpt_step-based recipes + and any test callers). + + Idempotent. Gated by env var ``PRIMUS_NEMO_LOSS`` (default: enabled; + disable with ``PRIMUS_NEMO_LOSS=0``). + + Returns True when patching took effect, False if disabled or already + installed. + """ + global _INSTALLED + if _INSTALLED: + return False + if not _enabled(): + _safe_log_rank_0( + "[primus.nemo_loss] PRIMUS_NEMO_LOSS disabled; using Megatron-Bridge " + "default global-token-mean loss" + ) + return False + + patched_any = False + + try: + from megatron.bridge.training import vlm_step as _vlm + + _vlm._create_loss_function = create_nemo_masked_next_token_loss_function + _safe_log_rank_0( + "[primus.nemo_loss] Patched megatron.bridge.training.vlm_step._create_loss_function " + "-> NeMo MaskedTokenLossReduction-equivalent (per-sample token-mean)" + ) + patched_any = True + except Exception as e: # pragma: no cover - defensive + _safe_log_rank_0(f"[primus.nemo_loss] Skipping vlm_step patch: {e!r}") + + try: + from megatron.bridge.training import gpt_step as _gpt + + _gpt.masked_next_token_loss = nemo_masked_next_token_loss + _safe_log_rank_0( + "[primus.nemo_loss] Patched megatron.bridge.training.gpt_step.masked_next_token_loss " + "-> NeMo per-sample-mean variant" + ) + patched_any = True + except Exception as e: # pragma: no cover - defensive + _safe_log_rank_0(f"[primus.nemo_loss] Skipping gpt_step patch: {e!r}") + + _INSTALLED = patched_any + return patched_any + + +# Auto-install on import. Importing this module from ``llama2_custom.py`` +# (or any recipe entry point) is sufficient to activate NeMo-equivalent +# loss reporting + gradient semantics for the entire training run. +install_nemo_loss_if_enabled() diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/pre_quantize_mxfp4.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/pre_quantize_mxfp4.py new file mode 100644 index 000000000..71978483d --- /dev/null +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/pre_quantize_mxfp4.py @@ -0,0 +1,464 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +PRE_QUANTIZED_MODEL=True parity with NeMo MLPerf 6.0 MI355X FP4. + +Recipe-local pre-quantization wiring. The three byte-equivalent NeMo +reference functions (``_extract_module``, ``_get_quantized_params_cpu``, +``_pre_quantize_model``) are ports of: + + ``mlperf-training-6-0/llama2_sft/nemo/src/callbacks/custom_callbacks.py`` + lines 57-62 (``_extract_module``), 122-149 (``_pre_quantize_model``), + 151-212 (``_get_quantized_params_cpu``). + +Adaptations vs NeMo (NOT behavioural changes): + +* NeMo reads ``self.first_last_layers_bf16``, + ``self.num_layers_at_start_in_bf16``, ``self.num_layers_at_end_in_bf16``, + ``self.store_quantized_params_on_gpu``, ``self.fp8_quantizer``, + ``self.mxfp4_quantizer`` from a ``CustomCallback`` instance bound to + the Lightning trainer. Primus has no such instance, so these are read + from environment variables with NeMo's *default* values + (``False`` / ``0`` / ``False``) and the FP8/MXFP4 quantizers are + module-locals built inside ``_pre_quantize_model``. +* NeMo passes ``trainer.model`` (single Lightning module, already + unwrapped from MegatronParallel by PL). Primus passes a + ``List[ModelChunk]`` from megatron-bridge; ``_extract_module`` strips + the list wrapper before stripping ``.module`` chains, so the input + semantically matches NeMo's. + +Integration: ``install_pre_quantize_wrap(orig_train)`` returns a drop-in +wrapper around ``megatron.bridge.training.train.train`` that, on the +first call, stashes FP8 weights on CPU, swaps live weights to MXFP4, and +bridges the stash to ``mxfp4_healing._ORDERED_FP8_STASH`` +before delegating to ``orig_train``. When ``PRE_QUANTIZED_MODEL`` is not +enabled the wrap is a no-op (it just returns ``orig_train`` unchanged). + +Historical note: this code used to live under +``primus.backends.megatron.patches.te_patches.pre_quantize_mxfp4_patches`` +as a ``@register_patch(phase="before_train")`` module that fired during +``MegatronBridgeBaseTrainer.__init__``. Because the patch ran *before* +the recipe was imported, a separate wrap-aware setter mechanism +(``_primus_set_orig_train``) was needed to inject the recipe's own +``megatron_bridge_train_override`` into the already-installed wrapper. +By moving the install into the recipe itself we can wrap the override +directly at recipe import time, deleting that plumbing entirely. +""" + +from __future__ import annotations + +import logging +import os +from typing import Any, Callable, List + +from primus.core.utils.module_utils import log_rank_0 + +_LOG = logging.getLogger(__name__) + +_TAG = "[pre_quantize_mxfp4]" + +# Layered FP8 CPU stash, populated by ``_pre_quantize_model``. +# Format mirrors NeMo's ``self.fp8_cpu_params``: +# List[List[Float8Tensor]] +# indexed by ``decoder.layers`` index, then by intra-layer iteration +# order over ``isinstance(m, (te.Linear, te.LayerNormLinear))`` modules. +# Same ``Float8Tensor`` objects are also referenced from +# ``mxfp4_healing._ORDERED_FP8_STASH`` (flat +# ``(module, fp8_tensor)`` view) for restore-time bookkeeping. +_FP8_CPU_PARAMS: List[List[Any]] = [] + + +def _truthy_env(name: str, default: bool = False) -> bool: + raw = os.getenv(name) + if raw is None: + return default + return raw.strip().lower() in ("1", "true", "yes", "on") + + +def is_pre_quantized_enabled() -> bool: + """Single switch for the pre-quantize path (``PRE_QUANTIZED_MODEL``).""" + return _truthy_env("PRE_QUANTIZED_MODEL", default=False) + + +# --------------------------------------------------------------------------- +# NeMo MLPerf 6.0 MI355X FP4 byte-equivalent pre-quantization functions. +# +# Source of truth: the actual NeMo container code, NOT the older host-tree +# mirror. Inside the running ``nemo_mxfp4_lora`` Docker image: +# /workspace/code/src/callbacks/custom_callbacks.py +# - _extract_module : lines 57-62 +# - _pre_quantize_model : lines 122-148 +# - _get_quantized_params_cpu : lines 150-213 +# The host-tree file under ``mlperf-training-6-0/`` is 25 lines shorter +# (660 vs 685) and out of date in three places: the MXFP4Quantizer kwargs, +# the ``_columnwise_data`` assertion, and the CPU-pin/GPU-restore +# columnwise-None guard. We mirror the *container* code. +# --------------------------------------------------------------------------- + + +def _extract_module(model): + """Unwrap the model from MegatronParallel / DDP wrappers to get the GPT module. + + NeMo (``custom_callbacks.py:57-62``):: + + m = model + while hasattr(m, "module"): + m = m.module + return m + + Adaptation: Primus passes ``model`` as ``List[ModelChunk]`` from + megatron-bridge (one chunk per pipeline stage). Pick chunk 0 then + follow ``.module`` exactly as NeMo does on the unwrapped trainer + model. + """ + m = model + if isinstance(m, (list, tuple)): + m = m[0] + while hasattr(m, "module"): + m = m.module + return m + + +def _pre_quantize_model(model): + """Pre-quantize model: store FP8 clones on CPU for healing, replace weights with MXFP4. + + Byte-equivalent to NeMo + ``custom_callbacks.CustomCallback._pre_quantize_model`` + (``mlperf-training-6-0/llama2_sft/nemo/src/callbacks/custom_callbacks.py:122-149``). + """ + import torch + import transformer_engine_torch as tex + from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer + from transformer_engine.pytorch.tensor.mxfp4_tensor import MXFP4Quantizer + + global _FP8_CPU_PARAMS + + device = next(_extract_module(model).parameters()).device + + # Clone current FP8 weights to CPU for healing. + # Must happen BEFORE MXFP4 replacement since we reference the live FP8 params. + # + # ``columnwise=`` mirrors NeMo's + # ``custom_callbacks.py::_pre_quantize_model`` (line 205). When the env + # is set, ``Float8Quantizer`` allocates a columnwise (transpose) buffer + # alongside the rowwise data so the post-healing FP8 GEMMs can reuse a + # precomputed transpose. When it is off, no transpose is kept -- and + # ``_restore_fp8_weights_to_gpu`` / ``_pre_ttt_fp8_warmup`` additionally + # set ``keep_fp8_weight_transpose_cache=False`` on each TE Linear so TE + # recomputes the transpose on demand instead of dereferencing a cache + # entry that was never populated. + # + # Default ``False`` matches NeMo's ``_env_enabled`` helper (returns False + # for missing / empty env). The Primus MLPerf-parity shell + # (``setup_llama2_70b_lora_fp4_training.sh``) exports + # ``ENABLE_TRANSPOSE_CACHE=1`` so the reference run gets the fast path. + _enable_transpose_cache = _truthy_env("ENABLE_TRANSPOSE_CACHE", default=False) + fp8_quantizer = Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device=device), + amax=torch.zeros(1, dtype=torch.float32, device=device), + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=_enable_transpose_cache, + ) + _FP8_CPU_PARAMS = _get_quantized_params_cpu(model, fp8_quantizer, "FP8_DS") + + _use_hadamard = os.environ.get("NVTE_MXFP4_USE_HADAMARD", "0") == "1" + + mxfp4_quantizer = MXFP4Quantizer( + rowwise=True, + columnwise=True, + with_gemm_swizzled_scales=True, + shuffle_rowwise_data=True, + shuffle_columnwise_data=True, + use_hadamard=_use_hadamard, + ) + _get_quantized_params_cpu(model, mxfp4_quantizer, "MXFP4", replace=True) + + torch.cuda.empty_cache() + + +def _get_quantized_params_cpu(model, quantizer, qtype: str, replace: bool = False) -> List: + """Byte-equivalent to NeMo + ``custom_callbacks.CustomCallback._get_quantized_params_cpu`` + (``mlperf-training-6-0/llama2_sft/nemo/src/callbacks/custom_callbacks.py:151-212``). + + Adaptation: + + * NeMo reads ``self.first_last_layers_bf16`` / ``self.num_layers_at_start_in_bf16`` + / ``self.num_layers_at_end_in_bf16`` / ``self.store_quantized_params_on_gpu`` + from the OmegaConf ``cfg.model`` (env-defaulted to ``False`` / + ``0`` / ``0`` / ``False`` in NeMo's MI355X FP4 config). Primus + reads the same env vars (``FIRST_LAST_LAYERS_BF16``, + ``NUM_LAYERS_AT_START_IN_BF16``, ``NUM_LAYERS_AT_END_IN_BF16``, + ``STORE_GPU``) with NeMo's same defaults. NeMo's MI355X reference + run does not set any of them so the effective behaviour is + identical. + """ + import torch + import transformer_engine.pytorch as te + import transformer_engine_torch as tex + from transformer_engine.pytorch.quantized_tensor import QuantizedTensor + from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor + + first_last_layers_bf16 = _truthy_env("FIRST_LAST_LAYERS_BF16", default=False) + num_layers_at_start_in_bf16 = int(os.getenv("NUM_LAYERS_AT_START_IN_BF16", "0") or "0") + num_layers_at_end_in_bf16 = int(os.getenv("NUM_LAYERS_AT_END_IN_BF16", "0") or "0") + store_quantized_params_on_gpu = _truthy_env("STORE_GPU", default=False) + + extracted_module = _extract_module(model) + layers = extracted_module.decoder.layers + layer_count = len(layers) + + quantized_params = [] + for layer_idx, layer in enumerate(layers): + if first_last_layers_bf16: + if ( + layer_idx < num_layers_at_start_in_bf16 + or layer_idx >= layer_count - num_layers_at_end_in_bf16 + ): + quantized_params.append([]) + continue + + quantized_layer_params = [] + for name, module in layer.named_modules(): + if not isinstance(module, (te.Linear, te.LayerNormLinear)): + continue + if not hasattr(module, "weight"): + continue + param = module.weight + with torch.no_grad(): + if qtype == "MXFP4": + if isinstance(param, QuantizedTensor): + bf16_param = param.dequantize().detach() + else: + bf16_param = param.detach().to(torch.bfloat16) + bf16_param = bf16_param.contiguous() + qparam = quantizer(bf16_param) + del bf16_param + assert qparam._rowwise_data is not None, "No rowwise data." + elif qtype == "FP8_DS": + fp8_data = param.data.to(torch.float8_e4m3fn).view(torch.uint8) + scale_inv = torch.ones(1, dtype=torch.float32, device=param.device) + qparam = Float8Tensor( + shape=param.shape, + dtype=torch.bfloat16, + data=fp8_data, + fp8_scale_inv=scale_inv, + fp8_dtype=tex.DType.kFloat8E4M3, + quantizer=quantizer, + ) + assert qparam._data is not None, "No data." + else: + raise ValueError(f"Unsupported quantization type: {qtype}") + + if replace: + qparam.requires_grad = False + module._parameters["weight"] = qparam + else: + qparam = qparam.clone() + if not store_quantized_params_on_gpu: + if qtype == "MXFP4": + qparam._rowwise_data = qparam._rowwise_data.cpu().pin_memory() + if qparam._columnwise_data is not None: + qparam._columnwise_data = qparam._columnwise_data.cpu().pin_memory() + elif qtype == "FP8_DS": + qparam._data = qparam._data.cpu().pin_memory() + if qparam._transpose is not None: + qparam._transpose = qparam._transpose.cpu().pin_memory() + quantized_layer_params.append(qparam) + quantized_params.append(quantized_layer_params) + + return quantized_params + + +# --------------------------------------------------------------------------- +# Memory diagnostics (Primus-only; not part of NeMo's pre-quantize logic). +# --------------------------------------------------------------------------- + + +def _log_gpu_mem(tag: str) -> None: + """Print rank-0 GPU mem checkpoint (allocated / reserved / max-allocated, GiB). + + Enabled when ``PRIMUS_LOG_GPU_MEM=1`` or ``MXFP4_HEALING_DEBUG=1``. + """ + if not ( + _truthy_env("PRIMUS_LOG_GPU_MEM", default=False) or _truthy_env("MXFP4_HEALING_DEBUG", default=False) + ): + return + try: + import torch + + if not torch.cuda.is_available(): + return + dev = torch.cuda.current_device() + alloc = torch.cuda.memory_allocated(dev) / (1024**3) + reserved = torch.cuda.memory_reserved(dev) / (1024**3) + max_alloc = torch.cuda.max_memory_allocated(dev) / (1024**3) + max_reserved = torch.cuda.max_memory_reserved(dev) / (1024**3) + log_rank_0( + f"{_TAG}[mem] {tag:<48s} " + f"allocated={alloc:7.2f} GiB | reserved={reserved:7.2f} GiB | " + f"max_alloc={max_alloc:7.2f} GiB | max_reserved={max_reserved:7.2f} GiB" + ) + except Exception as exc: # noqa: BLE001 + log_rank_0(f"{_TAG}[mem] {tag}: failed to read " f"({type(exc).__name__}: {exc})") + + +def _empty_cache_and_collect(tag: str) -> None: + """Run gc.collect() + torch.cuda.empty_cache() to release any cached blocks + holding onto the freed BF16 storage. Prints a mem checkpoint after.""" + try: + import gc + + import torch + + gc.collect() + if torch.cuda.is_available(): + torch.cuda.synchronize() + torch.cuda.empty_cache() + _log_gpu_mem(tag) + except Exception as exc: # noqa: BLE001 + log_rank_0(f"{_TAG}[mem] {tag}: empty_cache failed " f"({type(exc).__name__}: {exc})") + + +# --------------------------------------------------------------------------- +# train() wrapper -- invokes _pre_quantize_model once before the first iter. +# --------------------------------------------------------------------------- + + +def _make_pre_quantizing_train(orig_train: Callable) -> Callable: + """Return a thin wrapper around megatron-bridge's ``train`` that runs + pre-quantization once on the first call, then delegates to ``orig_train``. + + ``orig_train`` is closed over directly; since this wrapper is installed + from the recipe (after the recipe-level ``megatron_bridge_train_override`` + is defined), there is no need for the setter/getter swap dance the + old ``@register_patch`` version had. + """ + + if getattr(orig_train, "_primus_pre_quantize_wrapped", False): + return orig_train # idempotent + + state = {"done": False} + + def _wrapped_train(forward_step_func, model, *args, **kwargs): + if not state["done"]: + try: + try: + from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.mxfp4_healing import ( + log_healing_env_banner_once, + ) + + log_healing_env_banner_once() + except Exception: # noqa: BLE001 + pass + log_rank_0( + f"{_TAG} PRE_QUANTIZED_MODEL=True -> running " + "NeMo-equivalent _pre_quantize_model " + "(FP8 stash on CPU + MXFP4 swap)." + ) + + # Memory diagnostic: confirms whether the BF16 storage is freed + # after the in-place MXFP4 swap. If `allocated` does NOT drop by + # roughly half between "before swap" and "after swap (post empty_cache)", + # something outside this loop (DDP grad bucket, distopt mapping, TE + # FP8 weight cache, etc.) is still holding the old BF16 storage. + _log_gpu_mem("before _pre_quantize_model") + + # NeMo MLPerf 6.0 MI355X FP4 reference (custom_callbacks.py:122-149). + # Byte-equivalent: clones FP8 (E4M3) weights to CPU for healing, + # then replaces every TE Linear/LayerNormLinear weight with an + # MXFP4Tensor (rowwise+columnwise data populated). HEALING_ITER + # gating is enforced by _ORDERED_FP8_STASH consumers; here we + # *always* stash so PRE_QUANTIZED_MODEL=True and HEALING_ITER>0 + # have identical behaviour to NeMo (which always stashes when + # cfg.model.pre_quantized_model is True). + _pre_quantize_model(model) + + _log_gpu_mem("after _pre_quantize_model (BEFORE empty_cache)") + _empty_cache_and_collect("after _pre_quantize_model (AFTER empty_cache)") + + # Bridge: populate mxfp4_healing._ORDERED_FP8_STASH + # with flat (module, fp8_cpu_tensor) pairs that share storage + # with _FP8_CPU_PARAMS, so existing healing-side code + # (restore-to-GPU, refcount audit, FP8 warmup) keeps working. + try: + from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.mxfp4_healing import ( + _set_ordered_fp8_stash_from_layered, + ) + + _set_ordered_fp8_stash_from_layered(model, _FP8_CPU_PARAMS) + except Exception as bridge_err: # noqa: BLE001 + log_rank_0( + f"{_TAG} FP8 stash bridge to mxfp4_healing failed " + f"({type(bridge_err).__name__}: {bridge_err}); " + f"healing/restore will not work." + ) + + # NeMo MLPerf MI355X parity: optional FP8 warmup (off by default). + # Honors MXFP4_HEALING_FP8_WARMUP=1 to enable. + try: + from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.mxfp4_healing import ( + healing_iter as _healing_iter, + ) + from primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.mxfp4_healing import ( + run_fp8_warmup_for_kernel_jit, + ) + + if _healing_iter() > 0: + _optimizer = args[0] if len(args) > 0 else kwargs.get("optimizer") + _state = args[4] if len(args) > 4 else kwargs.get("state") + _model_cfg = None + try: + chunks = list(model) if isinstance(model, (list, tuple)) else [model] + _model_cfg = getattr(chunks[0], "config", None) + except Exception: # noqa: BLE001 + pass + if _model_cfg is None and _state is not None: + _model_cfg = getattr(getattr(_state, "cfg", None), "model", None) + + if _model_cfg is None: + log_rank_0( + f"{_TAG} FP8 warmup skipped: could not resolve " + "model_config from train() args (no model.config " + "and no state.cfg.model)." + ) + else: + run_fp8_warmup_for_kernel_jit( + model=model, + model_config=_model_cfg, + optimizer=_optimizer, + forward_step_func=forward_step_func, + state=_state, + ) + except Exception as warm_err: # noqa: BLE001 + log_rank_0( + f"{_TAG} FP8 warmup raised " + f"({type(warm_err).__name__}: {warm_err}); " + f"continuing to iter 1 without warmup." + ) + finally: + # Even on failure, mark as done so we don't re-attempt every + # call in case of caller retry. The exception will propagate. + state["done"] = True + return orig_train(forward_step_func, model, *args, **kwargs) + + _wrapped_train._primus_pre_quantize_wrapped = True # type: ignore[attr-defined] + return _wrapped_train + + +def install_pre_quantize_wrap(orig_train: Callable) -> Callable: + """Recipe-facing entry point. + + If ``PRE_QUANTIZED_MODEL`` is enabled, returns a pre-quantizing + wrapper around ``orig_train``; otherwise returns ``orig_train`` + unchanged. Idempotent: re-wrapping an already-wrapped callable is + a no-op. + """ + if not is_pre_quantized_enabled(): + return orig_train + return _make_pre_quantizing_train(orig_train) diff --git a/primus/cli/main.py b/primus/cli/main.py index 26606fa8d..b4b0b49a7 100644 --- a/primus/cli/main.py +++ b/primus/cli/main.py @@ -8,14 +8,34 @@ # aiter, ...) that may print/log at import time. Importing this module only # triggers the light ``primus/__init__.py`` and installs the FD-level filter # when ``PRIMUS_LOG_SUPPRESSION=1`` is set; it is a complete no-op otherwise. -import primus.mlperf_log_suppression # noqa: F401 # isort: skip +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _prefer_checkout_primus_on_sys_path() -> None: + """Prefer the git checkout over an installed wheel for in-tree Primus modules.""" + if not (_REPO_ROOT / "primus" / "mlperf_log_suppression.py").is_file(): + return + root = str(_REPO_ROOT) + if root in sys.path: + sys.path.remove(root) + sys.path.insert(0, root) + + +_prefer_checkout_primus_on_sys_path() + +try: + import primus.mlperf_log_suppression # noqa: F401 # isort: skip +except ModuleNotFoundError: + # Optional in-tree module; recipe-local _log_suppression covers MLPerf runs. + pass import argparse import importlib import pkgutil -import sys import traceback -from pathlib import Path from typing import Callable, Dict, Iterable, Optional, Set SUBCOMMAND_PACKAGE = "primus.cli.subcommands" @@ -26,11 +46,7 @@ def _ensure_project_root_on_path() -> None: Allow running `python primus/cli/main.py` from the repo root without requiring an installed package. """ - if __package__: - return - project_root = Path(__file__).resolve().parents[2] - if str(project_root) not in sys.path: - sys.path.insert(0, str(project_root)) + _prefer_checkout_primus_on_sys_path() def _iter_subcommand_modules() -> Iterable[str]: diff --git a/primus/configs/models/megatron_bridge/llama2_70b_lora_mxfp4.yaml b/primus/configs/models/megatron_bridge/llama2_70b_lora_mxfp4.yaml new file mode 100644 index 000000000..a2690cbf0 --- /dev/null +++ b/primus/configs/models/megatron_bridge/llama2_70b_lora_mxfp4.yaml @@ -0,0 +1,4 @@ +# MXFP4 LoRA recipe for MLPerf Llama2-70B (MI355X) +recipe: primus.backends.megatron_bridge.recipes.mlperf_llama2_70b.llama2_custom +flavor: llama2_70b_lora_mxfp4_config +hf_path: meta-llama/Llama-2-70b-hf diff --git a/primus/core/runtime/train_runtime.py b/primus/core/runtime/train_runtime.py index 61c8b7479..f314e956a 100644 --- a/primus/core/runtime/train_runtime.py +++ b/primus/core/runtime/train_runtime.py @@ -269,9 +269,12 @@ def _initialize_configuration( # Initialize TrainContext based on raw configuration (before CLI overrides). # Use primus_config_obj (PrimusConfig) for BaseModule compatibility + _data_path = getattr(self.args, "data_path", None) + if not _data_path: + _data_path = "./data" self.ctx = TrainContext( config_path=cfg_path, - data_path=Path(getattr(self.args, "data_path", "./data")), + data_path=Path(_data_path), module_name=module_name, primus_config=primus_config_obj, # Use PrimusConfig object, not SimpleNamespace module_config=module_cfg, diff --git a/runner/helpers/hooks/train/posttrain/megatron_bridge/00_install_requirements.sh b/runner/helpers/hooks/train/posttrain/megatron_bridge/00_install_requirements.sh index 302def462..b62fba157 100755 --- a/runner/helpers/hooks/train/posttrain/megatron_bridge/00_install_requirements.sh +++ b/runner/helpers/hooks/train/posttrain/megatron_bridge/00_install_requirements.sh @@ -10,10 +10,11 @@ set -euo pipefail echo "[+] Installing Megatron-Bridge dependencies..." SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -# Set up pip cache directory -PRIMUS_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" -DATA_PATH="${DATA_PATH:-${PRIMUS_ROOT}/data}" -PIP_CACHE_DIR="${PIP_CACHE_DIR:-${DATA_PATH}/pip_cache}" +# Pip cache must live on a path that exists inside THIS environment (e.g. Docker). Schedulers often set +# DATA_PATH to a *host* path (e.g. /home/.../data/mlperf_llama2) for dataset mounts; that path is not +# writable or may not exist in the container, so do not derive PIP_CACHE_DIR from DATA_PATH by default. +# Keep the cache outside the repo so it is never accidentally committed. +PIP_CACHE_DIR="${PIP_CACHE_DIR:-/tmp/primus-cache/pip}" echo "[INFO] Using pip cache: ${PIP_CACHE_DIR}" mkdir -p "${PIP_CACHE_DIR}" @@ -26,4 +27,7 @@ pip install --cache-dir="${PIP_CACHE_DIR}" -U "datasets>=2.14.0" pip install --cache-dir="${PIP_CACHE_DIR}" -r "${SCRIPT_DIR}/requirements-megatron-bridge.txt" +# datasets 5.x requires fsspec<=2026.4.0; megatron-bridge deps may upgrade it. +pip install --cache-dir="${PIP_CACHE_DIR}" 'fsspec>=2023.1.0,<=2026.4.0' + echo "[OK] Megatron-Bridge dependencies installed" diff --git a/runner/helpers/hooks/train/posttrain/megatron_bridge/01_convert_checkpoints.sh b/runner/helpers/hooks/train/posttrain/megatron_bridge/01_convert_checkpoints.sh index 4698c3042..632cb1036 100755 --- a/runner/helpers/hooks/train/posttrain/megatron_bridge/01_convert_checkpoints.sh +++ b/runner/helpers/hooks/train/posttrain/megatron_bridge/01_convert_checkpoints.sh @@ -25,7 +25,8 @@ if [[ -z "$CONFIG_FILE" ]]; then fi # Parse the complete config with all extends and nested configs -PRIMUS_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../../../.." && pwd)" +# Repo root: megatron_bridge -> posttrain -> train -> hooks -> helpers -> runner -> Primus (6 levels) +PRIMUS_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../../../../" && pwd)" cd "$PRIMUS_ROOT" # Convert CONFIG_FILE to absolute path @@ -78,10 +79,72 @@ if [[ -z "$HF_PATH" ]]; then exit 0 fi -# Set paths -DATA_PATH="${DATA_PATH:-${PRIMUS_ROOT}/data}" -HF_CACHE="${HF_HOME:-${DATA_PATH}/huggingface}/hub" -MEGATRON_PATH="${DATA_PATH}/megatron_checkpoints/$(basename "${HF_PATH}")" +# Data root: DATA_PATH is often a *host* path. The path may appear in the container (bind mount +# metadata) but still be unusable. Pick the first root where we can actually mkdir. +resolve_effective_data() { + local c testdir + local -a candidates=() + [[ -n "${MOUNT_DATA_PATH:-}" ]] && candidates+=("${MOUNT_DATA_PATH}") + candidates+=("/data") + [[ -n "${DATA_PATH:-}" ]] && candidates+=("${DATA_PATH}") + candidates+=("${PRIMUS_ROOT}/data") + + for c in "${candidates[@]}"; do + [[ -z "$c" ]] && continue + testdir="${c}/.primus_megatron_write_test_$$" + if mkdir -p "${testdir}" 2>/dev/null; then + rmdir "${testdir}" 2>/dev/null || rm -rf "${testdir}" 2>/dev/null || true + printf '%s' "$c" + return 0 + fi + done + mkdir -p "${PRIMUS_ROOT}/data" + printf '%s' "${PRIMUS_ROOT}/data" +} +EFFECTIVE_DATA="$(resolve_effective_data)" +LOG_INFO_RANK0 "Checkpoint/HF data root: ${EFFECTIVE_DATA}" + +# Hugging Face cache dirs: schedulers often export HF_HOME (or HF_HUB_CACHE / TRANSFORMERS_CACHE) +# to host paths that do not exist inside the container. Override when we cannot mkdir there. +HF_FALLBACK="${EFFECTIVE_DATA}/huggingface" +if [[ -n "${HF_HOME:-}" ]]; then + if mkdir -p "${HF_HOME}/hub/.primus_hf_write_test_$$" 2>/dev/null; then + rm -rf "${HF_HOME}/hub/.primus_hf_write_test_$$" 2>/dev/null || true + else + LOG_INFO_RANK0 "HF_HOME=${HF_HOME} is not usable here; using ${HF_FALLBACK}" + export HF_HOME="${HF_FALLBACK}" + fi +else + export HF_HOME="${HF_FALLBACK}" +fi +mkdir -p "${HF_HOME}/hub" + +if [[ -n "${HF_HUB_CACHE:-}" ]]; then + if mkdir -p "${HF_HUB_CACHE}/.primus_hf_write_test_$$" 2>/dev/null; then + rm -rf "${HF_HUB_CACHE}/.primus_hf_write_test_$$" 2>/dev/null || true + else + LOG_INFO_RANK0 "HF_HUB_CACHE=${HF_HUB_CACHE} is not usable here; using ${HF_HOME}/hub" + export HF_HUB_CACHE="${HF_HOME}/hub" + fi +fi + +if [[ -n "${TRANSFORMERS_CACHE:-}" ]]; then + if mkdir -p "${TRANSFORMERS_CACHE}/.primus_hf_write_test_$$" 2>/dev/null; then + rm -rf "${TRANSFORMERS_CACHE}/.primus_hf_write_test_$$" 2>/dev/null || true + else + LOG_INFO_RANK0 "TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE} is not usable here; using ${HF_HOME}/hub" + export TRANSFORMERS_CACHE="${HF_HOME}/hub" + fi +fi + +HF_CACHE="${HF_HUB_CACHE:-${HF_HOME}/hub}" + +# So the training process sees the same paths (execute_hooks exports env.* from hook stdout) +echo "env.HF_HOME=${HF_HOME}" +[[ -n "${HF_HUB_CACHE:-}" ]] && echo "env.HF_HUB_CACHE=${HF_HUB_CACHE}" +[[ -n "${TRANSFORMERS_CACHE:-}" ]] && echo "env.TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE}" + +MEGATRON_PATH="${EFFECTIVE_DATA}/megatron_checkpoints/$(basename "${HF_PATH}")" LOG_INFO_RANK0 "HF Model: ${HF_PATH}" LOG_INFO_RANK0 "HF Cache: ${HF_CACHE}" @@ -95,10 +158,22 @@ else LOG_INFO_RANK0 "HF checkpoint will be downloaded from ${HF_PATH}" fi +resolve_pretrained_checkpoint() { + local base="$1" + # Megatron-Bridge PEFT validation (checkpoint_exists) and load both expect the + # checkpoint *root* (latest_train_state.pt / latest_checkpointed_iteration.txt + # live here; weights are under iter_*). Do not point at iter_0000000. + printf '%s' "${base}" +} + # Check if Megatron checkpoint already exists if [[ -d "$MEGATRON_PATH" ]]; then LOG_INFO_RANK0 "Megatron checkpoint already exists at ${MEGATRON_PATH}, skipping conversion" - echo "extra.pretrained_checkpoint=${MEGATRON_PATH}" + CKPT_PATH="$(resolve_pretrained_checkpoint "${MEGATRON_PATH}")" + echo "extra.pretrained_checkpoint=${CKPT_PATH}" + echo "env.NVTE_FLASH_ATTN=0" + echo "env.NVTE_FUSED_ATTN=1" + echo "env.NVTE_UNFUSED_ATTN=0" exit 0 fi @@ -118,10 +193,23 @@ if [[ "$NODE_RANK" == "0" ]]; then # Set up Python path for Megatron-Bridge export PYTHONPATH="${PRIMUS_ROOT}/third_party/Megatron-Bridge/src:${PRIMUS_ROOT}/third_party/Megatron-Bridge/3rdparty/Megatron-LM:${PYTHONPATH:-}" + # HF→Megatron conversion uses attention_backend=auto; clear pre-set TE + # env vars from the container image (e.g. NVTE_FLASH_ATTN=0) so Megatron + # can configure them for the conversion pass. + unset NVTE_FLASH_ATTN NVTE_FUSED_ATTN NVTE_UNFUSED_ATTN + python3 third_party/Megatron-Bridge/examples/conversion/convert_checkpoints.py import \ --hf-model "${HF_PATH}" \ --megatron-path "${MEGATRON_PATH}" + # Restore MLPerf fused-attention env for the training run. + export NVTE_FLASH_ATTN=0 + export NVTE_FUSED_ATTN=1 + export NVTE_UNFUSED_ATTN=0 + echo "env.NVTE_FLASH_ATTN=0" + echo "env.NVTE_FUSED_ATTN=1" + echo "env.NVTE_UNFUSED_ATTN=0" + # Create done file and remove lock touch "$DONE_FILE" rm -f "$LOCK_FILE" @@ -153,4 +241,6 @@ else echo "[OK] [RANK ${NODE_RANK}] Checkpoint ready at ${MEGATRON_PATH}" fi -echo "extra.pretrained_checkpoint=${MEGATRON_PATH}" +CKPT_PATH="$(resolve_pretrained_checkpoint "${MEGATRON_PATH}")" +echo "extra.pretrained_checkpoint=${CKPT_PATH}" +exit 0 diff --git a/runner/helpers/hooks/train/posttrain/megatron_bridge/02_prepare_mlperf_dataset.sh b/runner/helpers/hooks/train/posttrain/megatron_bridge/02_prepare_mlperf_dataset.sh new file mode 100755 index 000000000..ca92cc1b6 --- /dev/null +++ b/runner/helpers/hooks/train/posttrain/megatron_bridge/02_prepare_mlperf_dataset.sh @@ -0,0 +1,98 @@ +#!/bin/bash +############################################################################### +# Download SCROLLS gov-report MLPerf dataset and build packed .npy + metadata +# when post_trainer uses dataset_type=mlperf_dataset. +############################################################################### +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck disable=SC1091 +source "${SCRIPT_DIR}/../../../../../lib/common.sh" || { + echo "[ERROR] Failed to load common library" >&2 + exit 1 +} + +PRIMUS_ROOT="$(cd "${SCRIPT_DIR}/../../../../../../" && pwd)" +cd "${PRIMUS_ROOT}" + +CONFIG_FILE="" +for ((i=1; i<=$#; i++)); do + if [[ "${!i}" == "--config" ]]; then + j=$((i+1)) + CONFIG_FILE="${!j}" + break + fi +done + +if [[ -z "$CONFIG_FILE" ]]; then + exit 0 +fi + +if [[ ! "$CONFIG_FILE" = /* ]]; then + CONFIG_FILE="${PRIMUS_ROOT}/${CONFIG_FILE#./}" +fi + +read -r DATA_DIR SEQ_LENGTH <<< "$(python3 -c " +import os +import sys +sys.path.insert(0, '${PRIMUS_ROOT}') +from pathlib import Path +from primus.core.config.primus_config import load_primus_config, get_module_config + +cfg = load_primus_config(Path('${CONFIG_FILE}'), None) +post = get_module_config(cfg, 'post_trainer') +if post is None: + sys.exit(0) +params = getattr(post, 'params', None) +if params is None or getattr(params, 'dataset_type', '') != 'mlperf_dataset': + sys.exit(0) + +train_path = getattr(params, 'packed_train_data_path', None) or os.environ.get('PACKED_DATA_DIR', '/data') +train_path = os.path.expandvars(str(train_path)) +data_dir = os.path.dirname(train_path) if train_path.endswith('.npy') else train_path +seq_length = int(getattr(params, 'seq_length', 8192) or 8192) +print(data_dir, seq_length) +" 2>/dev/null || true)" + +if [[ -z "${DATA_DIR}" ]]; then + exit 0 +fi + +SEQ_LENGTH="${SEQ_LENGTH:-8192}" +mkdir -p "${DATA_DIR}" + +LOG_INFO_RANK0 "[mlperf-data] Data root: ${DATA_DIR} (seq_length=${SEQ_LENGTH})" +echo "env.PACKED_DATA_DIR=${DATA_DIR}" +echo "env.DATA_PATH=${DATA_DIR}" + +if [[ -f "${DATA_DIR}/train.npy" && -f "${DATA_DIR}/validation.npy" && -f "${DATA_DIR}/packed_metadata.jsonl" ]]; then + LOG_INFO_RANK0 "[mlperf-data] Dataset already present; skipping download" + exit 0 +fi + +if [[ -z "${HF_TOKEN:-}" && ( ! -f "${DATA_DIR}/train.npy" || ! -f "${DATA_DIR}/validation.npy" ) ]]; then + LOG_ERROR_RANK0 "[mlperf-data] HF_TOKEN is required to download regisss/scrolls_gov_report_preprocessed_mlperf_2" + exit 1 +fi + +MLPERF_RECIPE_DIR="${PRIMUS_ROOT}/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b" + +if [[ ! -f "${DATA_DIR}/train.npy" || ! -f "${DATA_DIR}/validation.npy" ]]; then + LOG_INFO_RANK0 "[mlperf-data] Downloading and converting MLPerf dataset..." + python3 "${MLPERF_RECIPE_DIR}/download_dataset.py" --data_dir "${DATA_DIR}" + python3 "${MLPERF_RECIPE_DIR}/convert_dataset.py" --data_dir "${DATA_DIR}" +fi + +if [[ ! -f "${DATA_DIR}/packed_metadata.jsonl" ]]; then + LOG_INFO_RANK0 "[mlperf-data] Creating packed_metadata.jsonl..." + python3 "${MLPERF_RECIPE_DIR}/create_metadata.py" "${SEQ_LENGTH}" "${DATA_DIR}/packed_metadata.jsonl" +fi + +for f in train.npy validation.npy packed_metadata.jsonl; do + if [[ ! -f "${DATA_DIR}/${f}" ]]; then + LOG_ERROR_RANK0 "[mlperf-data] Expected file missing: ${DATA_DIR}/${f}" + exit 1 + fi +done + +LOG_SUCCESS_RANK0 "[mlperf-data] MLPerf dataset ready under ${DATA_DIR}" From 87f57877a1f4fd5195b35b29f4350333ed96cd56 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Wed, 15 Jul 2026 15:16:41 +0300 Subject: [PATCH 030/127] feat(flux): diffusion data preprocessing pipelines + data CLI (#817) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/data` — review after it. ## What this changes The offline data-preprocessing layer (auth / download / finalize / validate plus the raw/ingest/encoded pipeline stages), wired into a new `primus data` CLI subcommand. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); builds on `feat/flux/data` (uses its providers and inherits the energon/webdataset deps). ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion/data/preprocessing tests/unit_tests/cli/test_data_config.py`. Validated locally on an AMD GPU container: 63 passed. ## Files 20 (preprocessing pipelines, data CLI subcommand + tests). Co-authored-by: Flux Split Trial --- .../data/diffusion/preprocessing/__init__.py | 9 + .../data/diffusion/preprocessing/auth.py | 174 ++++ .../data/diffusion/preprocessing/download.py | 156 +++ .../data/diffusion/preprocessing/finalize.py | 230 +++++ .../preprocessing/pipelines/__init__.py | 21 + .../diffusion/preprocessing/pipelines/base.py | 70 ++ .../preprocessing/pipelines/encoded.py | 608 +++++++++++ .../preprocessing/pipelines/ingest.py | 311 ++++++ .../diffusion/preprocessing/pipelines/raw.py | 234 +++++ .../data/diffusion/preprocessing/utils.py | 546 ++++++++++ .../data/diffusion/preprocessing/validate.py | 302 ++++++ primus/cli/main.py | 5 +- primus/cli/subcommands/data.py | 950 ++++++++++++++++++ .../diffusion/data/preprocessing/__init__.py | 2 + .../data/preprocessing/test_download.py | 251 +++++ .../data/preprocessing/test_finalize.py | 121 +++ .../data/preprocessing/test_ingest.py | 330 ++++++ .../data/preprocessing/test_utils.py | 164 +++ .../data/preprocessing/test_validate.py | 194 ++++ tests/unit_tests/cli/test_data_config.py | 253 +++++ 20 files changed, 4930 insertions(+), 1 deletion(-) create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/__init__.py create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/auth.py create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/download.py create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/finalize.py create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/pipelines/__init__.py create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/pipelines/base.py create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/pipelines/encoded.py create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/pipelines/raw.py create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/utils.py create mode 100644 primus/backends/megatron/data/diffusion/preprocessing/validate.py create mode 100644 primus/cli/subcommands/data.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/preprocessing/__init__.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_download.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_finalize.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_utils.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_validate.py create mode 100644 tests/unit_tests/cli/test_data_config.py diff --git a/primus/backends/megatron/data/diffusion/preprocessing/__init__.py b/primus/backends/megatron/data/diffusion/preprocessing/__init__.py new file mode 100644 index 000000000..21f7888d0 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Data preprocessing for diffusion models. + +Pipeline implementations live in the ``pipelines`` subpackage; shared helpers +live in ``utils``. Import those submodules directly. +""" diff --git a/primus/backends/megatron/data/diffusion/preprocessing/auth.py b/primus/backends/megatron/data/diffusion/preprocessing/auth.py new file mode 100644 index 000000000..8cb574f08 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/auth.py @@ -0,0 +1,174 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Authentication utilities for HuggingFace dataset access. + +Provides secure token loading from files with permission checks. +""" + +import logging +import os +import stat +from pathlib import Path +from typing import Optional + +logger = logging.getLogger(__name__) + + +class HFAuthError(Exception): + """Exception raised for HuggingFace authentication errors.""" + + +def check_file_permissions(file_path: Path) -> bool: + """ + Check if file has secure permissions (600 or 400). + + Args: + file_path: Path to token file + + Returns: + True if permissions are secure, False otherwise + """ + try: + file_stat = os.stat(file_path) + mode = stat.S_IMODE(file_stat.st_mode) + + # Check if file is readable by others or group + if mode & (stat.S_IRGRP | stat.S_IROTH): + return False + + # Check if file is writable by others or group + if mode & (stat.S_IWGRP | stat.S_IWOTH): + return False + + return True + except OSError as e: + logger.error(f"Failed to check file permissions: {e}") + return False + + +def load_token_from_file(token_file: str) -> str: + """ + Load HuggingFace token from file with security checks. + + Args: + token_file: Path to file containing HuggingFace token + + Returns: + Token string + + Raises: + HFAuthError: If file has insecure permissions or cannot be read + """ + token_path = Path(token_file).expanduser().resolve() + + # Check if file exists + if not token_path.exists(): + raise HFAuthError( + f"Token file not found: {token_path}\n" f"Please create the file or check the path." + ) + + # Check if it's a file (not directory) + if not token_path.is_file(): + raise HFAuthError(f"Token path is not a file: {token_path}") + + # Check file permissions + if not check_file_permissions(token_path): + raise HFAuthError( + f"Token file has insecure permissions: {token_path}\n" + f"Please set secure permissions:\n" + f" chmod 600 {token_path}\n" + f"Current permissions allow read/write by group or others." + ) + + # Read token + try: + with open(token_path, "r") as f: + token = f.read().strip() + + if not token: + raise HFAuthError(f"Token file is empty: {token_path}") + + # Basic validation (HF tokens start with 'hf_') + if not token.startswith("hf_"): + logger.warning( + f"Token from {token_path} doesn't start with 'hf_' - " + f"this may not be a valid HuggingFace token" + ) + + logger.info(f"Loaded HuggingFace token from {token_path}") + return token + + except IOError as e: + raise HFAuthError(f"Failed to read token file {token_path}: {e}") + + +def setup_hf_authentication(token_file: Optional[str] = None, use_env: bool = True) -> Optional[str]: + """ + Setup HuggingFace authentication with multiple fallback options. + + Priority: + 1. Token from file (if token_file provided) + 2. Token from HF_TOKEN environment variable (if use_env=True) + 3. Token from HF CLI login (~/.cache/huggingface/token) + 4. No authentication (public datasets only) + + Args: + token_file: Optional path to token file + use_env: Whether to check HF_TOKEN environment variable + + Returns: + Token string if found, None otherwise + + Side effects: + Sets HF_TOKEN environment variable if token is found + """ + # Priority 1: Token file + if token_file: + try: + token = load_token_from_file(token_file) + os.environ["HF_TOKEN"] = token + logger.info("Using HuggingFace token from file") + return token + except HFAuthError as e: + logger.error(str(e)) + raise + + # Priority 2: Environment variable + if use_env and "HF_TOKEN" in os.environ: + token = os.environ["HF_TOKEN"] + if token: + logger.info("Using HuggingFace token from HF_TOKEN environment variable") + return token + + # Priority 3: HF CLI login + hf_cache_token = Path.home() / ".cache" / "huggingface" / "token" + if hf_cache_token.exists(): + if not check_file_permissions(hf_cache_token): + logger.warning( + f"HuggingFace CLI token file has insecure permissions: {hf_cache_token}. " + "Skipping it; run `chmod 600` on the file to use it." + ) + else: + try: + with open(hf_cache_token, "r") as f: + token = f.read().strip() + if token: + os.environ["HF_TOKEN"] = token + logger.info("Using HuggingFace token from CLI login (~/.cache/huggingface/token)") + return token + except IOError as e: + logger.debug(f"Could not read HF CLI token: {e}") + + # No authentication found + logger.info("No HuggingFace authentication found. " "Only public datasets will be accessible.") + return None + + +__all__ = [ + "HFAuthError", + "load_token_from_file", + "setup_hf_authentication", + "check_file_permissions", +] diff --git a/primus/backends/megatron/data/diffusion/preprocessing/download.py b/primus/backends/megatron/data/diffusion/preprocessing/download.py new file mode 100644 index 000000000..c58c2814e --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/download.py @@ -0,0 +1,156 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Reusable download utilities for dataset preparation pipelines. + +Provides HTTP download with exponential backoff, MD5 verification, +and MLCommons R2 manifest resolution (.uri / .md5 protocol). +Uses urllib.request (stdlib) -- no external HTTP dependencies needed. +""" + +import hashlib +import logging +import random +import shutil +import time +import urllib.error +import urllib.request +from http.client import HTTPResponse +from pathlib import Path +from typing import List, Optional, Tuple + +logger = logging.getLogger(__name__) + +_DOWNLOAD_TIMEOUT = 300 # 5 min per file +_MAX_RETRIES = 5 +_BASE_DELAY = 1.0 +_MAX_DELAY = 60.0 +_USER_AGENT = "Wget/1.21" + + +class _MD5MismatchError(Exception): + """Raised when a downloaded file's MD5 doesn't match the expected value.""" + + +def _is_retryable(status_code: int) -> bool: + return status_code == 429 or 500 <= status_code < 600 + + +def download_with_backoff( + url: str, + dest: Path, + expected_md5: Optional[str] = None, + max_retries: int = _MAX_RETRIES, + base_delay: float = _BASE_DELAY, + timeout: int = _DOWNLOAD_TIMEOUT, +) -> None: + """Download a file via HTTP with exponential backoff on 429/5xx. + + Streams the response to disk to avoid holding large files in memory. + Verifies MD5 checksum if provided. + """ + dest.parent.mkdir(parents=True, exist_ok=True) + + last_error: Optional[Exception] = None + for attempt in range(max_retries + 1): + try: + req = urllib.request.Request(url, headers={"User-Agent": _USER_AGENT}) + resp: HTTPResponse = urllib.request.urlopen(req, timeout=timeout) + with open(dest, "wb") as f: + shutil.copyfileobj(resp, f) + + if expected_md5: + actual_md5 = hashlib.md5(dest.read_bytes()).hexdigest() + if actual_md5 != expected_md5: + dest.unlink(missing_ok=True) + raise _MD5MismatchError( + f"MD5 mismatch for {dest.name}: " f"expected {expected_md5}, got {actual_md5}" + ) + return + + except urllib.error.HTTPError as e: + last_error = e + if not _is_retryable(e.code) or attempt == max_retries: + dest.unlink(missing_ok=True) + raise RuntimeError( + f"HTTP {e.code} downloading {url} " f"(attempt {attempt + 1}/{max_retries + 1})" + ) from e + delay = min(base_delay * (2**attempt) + random.uniform(0, 1), _MAX_DELAY) + logger.warning(f" HTTP {e.code} on {url}, retry {attempt + 1}/{max_retries} " f"in {delay:.1f}s") + time.sleep(delay) + + except (urllib.error.URLError, TimeoutError, OSError) as e: + last_error = e + if attempt == max_retries: + dest.unlink(missing_ok=True) + raise RuntimeError(f"Download failed for {url} after {max_retries + 1} attempts: {e}") from e + delay = min(base_delay * (2**attempt) + random.uniform(0, 1), _MAX_DELAY) + logger.warning( + f" Network error on {url}: {e}, retry {attempt + 1}/{max_retries} " f"in {delay:.1f}s" + ) + time.sleep(delay) + + except _MD5MismatchError as e: + last_error = e + if attempt == max_retries: + dest.unlink(missing_ok=True) + raise RuntimeError(str(e)) from e + delay = min(base_delay * (2**attempt) + random.uniform(0, 1), _MAX_DELAY) + logger.warning(f" {e}, retry {attempt + 1}/{max_retries} in {delay:.1f}s") + time.sleep(delay) + + dest.unlink(missing_ok=True) + raise RuntimeError(f"Download failed for {url}: {last_error}") + + +def fetch_url_text(url: str, timeout: int = 30) -> str: + """Fetch a small text resource (manifest, etc.) via HTTP.""" + req = urllib.request.Request(url, headers={"User-Agent": _USER_AGENT}) + resp = urllib.request.urlopen(req, timeout=timeout) + return resp.read().decode("utf-8") + + +def parse_md5_manifest( + manifest_text: str, + suffix_filter: Optional[str] = None, +) -> List[Tuple[str, str]]: + """Parse an MLCommons .md5 manifest into (md5, filename) pairs. + + Args: + manifest_text: Raw text content of the .md5 file. + suffix_filter: If provided, only include files ending with this + suffix (e.g. ".arrow"). None means include all files. + + Returns: + Sorted list of (md5, filename) tuples, sorted by filename + for deterministic ordering. + """ + entries = [] + for line in manifest_text.strip().splitlines(): + parts = line.strip().split(None, 1) + if len(parts) != 2: + continue + md5, fname = parts + if suffix_filter is not None and not fname.endswith(suffix_filter): + continue + entries.append((md5, fname)) + entries.sort(key=lambda x: x[1]) + return entries + + +def fetch_manifest(manifest_url: str) -> Tuple[str, List[Tuple[str, str]]]: + """Fetch .uri and .md5 manifests, return (base_url, [(md5, filename)]). + + The manifest_url should end with '.uri' or '.md5'. The function derives + the complementary URL by replacing the suffix. + """ + uri_url = manifest_url.replace(".md5", ".uri") + md5_url = manifest_url.replace(".uri", ".md5") + + base_url = fetch_url_text(uri_url).strip() + md5_text = fetch_url_text(md5_url) + entries = parse_md5_manifest(md5_text) + + logger.info(f"Manifest: {len(entries)} files, base URL: {base_url}") + return base_url, entries diff --git a/primus/backends/megatron/data/diffusion/preprocessing/finalize.py b/primus/backends/megatron/data/diffusion/preprocessing/finalize.py new file mode 100644 index 000000000..a2c88fbc2 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/finalize.py @@ -0,0 +1,230 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Automatic Energon dataset finalization for Primus diffusion pipelines. + +This module automates the post-processing steps required after data preprocessing: +1. Create .nv-meta/dataset.yaml with correct sample type configuration +2. Run ``energon prepare`` non-interactively to index the dataset +3. Verify the dataset is ready for training + +Supports two indexing strategies: +- **Ratio-based** (default): splits shards by train/val/test ratio. + Works with both the programmatic ``BaseWebdatasetFactory`` API and + the ``energon prepare`` subprocess fallback. +- **Pattern-based**: assigns shards to splits by regex patterns on + relative paths (e.g. ``"train/.*"``). Requires the programmatic + API (``megatron-energon >= 7.x``). + +Usage: + from primus.backends.megatron.data.diffusion.preprocessing.finalize import finalize_energon_dataset + + # Ratio-based (original behavior) + finalize_energon_dataset( + output_dir='/workspace/Primus/data/encoded_pokemon', + train_split=1.0, + encoding='preencoded', + ) + + # Pattern-based (MLPerf layout with separate train/ and val/ dirs) + finalize_energon_dataset( + output_dir='/workspace/Primus/data/mlperf_flux1', + encoding='preencoded_numpy', + split_parts_patterns=[("train", "train/.*"), ("val", "val/.*")], + ) +""" + +import logging +import subprocess +from pathlib import Path +from typing import List, Literal, Optional, Tuple + +logger = logging.getLogger(__name__) + +_ENCODING_TYPE = Literal["preencoded", "preencoded_numpy", "raw"] + + +def _write_dataset_yaml(output_path: Path, encoding: str) -> Path: + """Write .nv-meta/dataset.yaml with the correct subflavor.""" + meta_dir = output_path / ".nv-meta" + meta_dir.mkdir(exist_ok=True) + + dataset_yaml_path = meta_dir / "dataset.yaml" + dataset_yaml_path.write_text( + f"__module__: megatron.energon\n" + f"__class__: CrudeWebdataset\n" + f"subflavors:\n" + f" encoding: {encoding}\n" + ) + logger.info(f"✓ Created {dataset_yaml_path}") + return dataset_yaml_path + + +def _prepare_programmatic( + output_path: Path, + num_workers: int, + train_split: float, + split_parts_patterns: Optional[List[Tuple[str, str]]], +) -> None: + """Index dataset using BaseWebdatasetFactory.prepare_dataset().""" + from megatron.energon import BaseWebdatasetFactory + + shard_paths = sorted(str(p.relative_to(output_path)) for p in output_path.glob("**/*.tar")) + if not shard_paths: + raise FileNotFoundError(f"No .tar shards found under {output_path}") + + kwargs: dict = { + "parent_path": str(output_path), + "paths": shard_paths, + "workers": num_workers, + } + + if split_parts_patterns: + kwargs["split_parts_patterns"] = split_parts_patterns + else: + val_split = (1.0 - train_split) / 2 + test_split = (1.0 - train_split) / 2 + kwargs["split_parts_ratio"] = [ + ("train", train_split), + ("val", val_split), + ("test", test_split), + ] + + logger.info(f" Indexing with BaseWebdatasetFactory ({len(shard_paths)} shards)") + BaseWebdatasetFactory.prepare_dataset(**kwargs) + logger.info("✓ Energon indexing complete (programmatic API)") + + +def _prepare_subprocess( + output_path: Path, + num_workers: int, + train_split: float, + split_parts_patterns: Optional[List[Tuple[str, str]]], +) -> None: + """Fallback: index dataset using ``energon prepare`` subprocess. + + NOTE: The subprocess path only supports ratio-based splitting. + Pattern-based ``split_parts_patterns`` requires the programmatic API. + """ + if split_parts_patterns: + raise RuntimeError( + "Pattern-based split_parts_patterns requires " + "megatron.energon.BaseWebdatasetFactory (programmatic API). " + "Install megatron-energon >= 7.x or use ratio-based splitting." + ) + + val_split = (1.0 - train_split) / 2 + test_split = (1.0 - train_split) / 2 + split_input = f"{train_split}, {val_split}, {test_split}\nn\n" + + logger.info(f" Split ratios: train={train_split:.2f}, " f"val={val_split:.2f}, test={test_split:.2f}") + logger.info(" Running energon prepare subprocess...") + + try: + process = subprocess.Popen( + [ + "energon", + "prepare", + str(output_path), + "--num-workers", + str(num_workers), + ], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + + stdout, _ = process.communicate(input=split_input, timeout=300) + + if process.returncode != 0: + logger.error("energon prepare failed:") + logger.error(stdout) + raise RuntimeError(f"energon prepare exited with code {process.returncode}") + + for line in stdout.splitlines(): + if any(skip in line for skip in ["libibverbs", "Warning:", "UserWarning"]): + continue + if any(key in line for key in ["Indexing", "Done", "Found", "samples", "shards"]): + logger.info(f" {line.strip()}") + + logger.info("✓ Energon indexing complete (subprocess)") + + except subprocess.TimeoutExpired: + process.kill() + raise RuntimeError( + "energon prepare timed out after 5 minutes. " + "This may indicate a problem with the dataset or system resources." + ) + except FileNotFoundError: + raise RuntimeError( + "energon command not found. Make sure megatron-energon is installed:\n" + " pip install megatron-energon" + ) + + +def finalize_energon_dataset( + output_dir: str, + train_split: float = 1.0, + encoding: _ENCODING_TYPE = "preencoded", + num_workers: int = 8, + split_parts_patterns: Optional[List[Tuple[str, str]]] = None, +) -> None: + """ + Finalize Energon WebDataset by creating dataset.yaml and running energon prepare. + + This automates the manual steps typically required after data preprocessing: + 1. Create .nv-meta/dataset.yaml with correct subflavor configuration + 2. Run energon indexing (programmatic API with subprocess fallback) + 3. Verify the dataset is ready for training + + Args: + output_dir: Path to dataset directory containing tar files. + train_split: Fraction for training (rest split evenly between + val/test). Ignored when split_parts_patterns is provided. + encoding: Dataset encoding mode. + num_workers: Number of workers for energon prepare (default: 8). + split_parts_patterns: List of (split_name, pattern) tuples for + pattern-based splitting, e.g. + ``[("train", "train/.*"), ("val", "val/.*")]``. + When provided, train_split is ignored. + + Raises: + RuntimeError: If energon prepare fails or energon command not found + FileNotFoundError: If output_dir doesn't exist or has no tar files + """ + output_path = Path(output_dir) + + if not output_path.exists(): + raise FileNotFoundError(f"Output directory not found: {output_dir}") + + tar_files = list(output_path.glob("**/*.tar")) + if not tar_files: + raise FileNotFoundError( + f"No .tar files found in {output_dir}. " "Make sure data preprocessing completed successfully." + ) + + logger.info("=" * 80) + logger.info("Finalizing Energon dataset for training...") + logger.info("=" * 80) + logger.info(f"Found {len(tar_files)} shard(s) to index") + + _write_dataset_yaml(output_path, encoding) + + logger.info("Running energon prepare (this may take a few minutes)...") + + try: + _prepare_programmatic(output_path, num_workers, train_split, split_parts_patterns) + except ImportError: + logger.info("BaseWebdatasetFactory not available, falling back to subprocess") + _prepare_subprocess(output_path, num_workers, train_split, split_parts_patterns) + + from .validate import validate_energon_dataset + + validate_energon_dataset(output_dir, encoding=encoding) + + logger.info("=" * 80) + logger.info(f"✓ Dataset finalized: {output_dir}") + logger.info(f" To use in training, set: dataset_path: {output_dir}") + logger.info("=" * 80) diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/__init__.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/__init__.py new file mode 100644 index 000000000..15280d896 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Dataset preparation pipelines for Megatron diffusion models. + +Provides reusable pipeline classes for creating Energon WebDatasets +from various input sources (HuggingFace, directories, WebDatasets). +""" + +from .base import DatasetPipeline +from .encoded import EncodedDatasetPipeline +from .ingest import StreamingIngestPipeline +from .raw import RawDatasetPipeline + +__all__ = [ + "DatasetPipeline", + "RawDatasetPipeline", + "EncodedDatasetPipeline", + "StreamingIngestPipeline", +] diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/base.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/base.py new file mode 100644 index 000000000..65d579cea --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/base.py @@ -0,0 +1,70 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +"""Abstract base class for dataset preparation pipelines.""" + +import logging +from abc import ABC, abstractmethod +from typing import Any, Dict + +from ..utils import load_from_directory, load_from_huggingface, load_from_webdataset + +logger = logging.getLogger(__name__) + + +class DatasetPipeline(ABC): + """Abstract base for all dataset preparation pipelines. + + Subclasses implement run() which executes the full pipeline + and returns a dict of statistics (samples_processed, shards_written, etc.). + """ + + # Whether HuggingFace sources are loaded in streaming mode. Subclasses that + # need the full dataset materialized (e.g. for distributed total-count + # splitting) set this to False. + HF_STREAMING: bool = True + + @abstractmethod + def run(self, **kwargs) -> Dict[str, Any]: ... + + def load_data(self, **source_kwargs): + """Load data based on ``self.source_type``. + + Args: + **source_kwargs: Source-specific arguments + - directory: input_dir + - huggingface: hf_dataset, hf_split, hf_data_files, image/caption keys + - webdataset: input_path + + Returns: + Iterator over samples with 'image' and 'caption' keys. + """ + if self.source_type == "directory": + logger.info(f"Loading from directory: {source_kwargs['input_dir']}") + return load_from_directory(source_kwargs["input_dir"]) + elif self.source_type == "huggingface": + hf_split = source_kwargs.get("hf_split", "train") + hf_data_files = source_kwargs.get("hf_data_files") + if hf_data_files: + logger.info( + f"Loading from HuggingFace: {source_kwargs['hf_dataset']} " + f"(split: {hf_split}, data_files: {hf_data_files})" + ) + else: + logger.info(f"Loading from HuggingFace: {source_kwargs['hf_dataset']} (split: {hf_split})") + + return load_from_huggingface( + source_kwargs["hf_dataset"], + split=hf_split, + streaming=self.HF_STREAMING, + data_files=hf_data_files, + image_key=source_kwargs.get("image_key"), + caption_key=source_kwargs.get("caption_key"), + image_keys=source_kwargs.get("image_keys"), + caption_keys=source_kwargs.get("caption_keys"), + ) + elif self.source_type == "webdataset": + logger.info(f"Loading from WebDataset: {source_kwargs['input_path']}") + return load_from_webdataset(source_kwargs["input_path"]) + else: + raise ValueError(f"Unknown source type: {self.source_type}") diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/encoded.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/encoded.py new file mode 100644 index 000000000..ba4d171ba --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/encoded.py @@ -0,0 +1,608 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Pre-encoded dataset preparation pipeline for Megatron diffusion models. + +Creates pre-encoded Energon WebDatasets with VAE/T5/CLIP encoded tensors +for faster training (no on-the-fly encoding overhead). +""" + +import logging +import time +from pathlib import Path +from typing import Optional + +import numpy as np +import torch +from tqdm import tqdm + +from primus.backends.megatron.data.diffusion.encoders import get_encoder +from primus.backends.megatron.data.diffusion.encoders.config import ( + CLIPLConfig, + T5XXLConfig, + VAEConfig, +) + +from ..utils import ( + get_distributed_info, + preprocess_image, + save_to_webdataset, + split_work_for_rank, +) + +logger = logging.getLogger(__name__) + + +def set_reproducibility(seed: int = 42) -> None: + """Set global seeds and cuDNN flags for reproducible encoding. + + Called explicitly from the pipeline entry point rather than at import time, + so importing this module does not mutate global torch/cuDNN state for + unrelated callers. + """ + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +from .base import DatasetPipeline + + +class EncodedDatasetPipeline(DatasetPipeline): + """ + Pipeline for creating pre-encoded Energon WebDataset with VAE/T5/CLIP tensors. + + This pipeline pre-encodes images with VAE and text with T5/CLIP, creating larger + datasets that enable faster training (no encoding overhead). + + Args: + source_type: Type of input source ('directory', 'huggingface', 'webdataset') + output_dir: Output directory for Energon WebDataset + model_path: Pretrained model path (HF or local, default: black-forest-labs/FLUX.1-dev) + vae_path: Custom VAE path (overrides model_path) + t5_path: Custom T5 path (overrides model_path) + clip_path: Custom CLIP path (overrides model_path) + precision: Model precision ('bf16', 'fp16', 'fp32', default: 'bf16') + device: Device for encoding (default: 'cuda') + batch_size: Encoding batch size (default: 8) + t5_max_length: T5 max sequence length (default: 512, use 256 for schnell) + image_size: Target image size (default: 1024) + variable_size: If True, resize to nearest multiple of 16 instead of fixed size (default: False) + center_crop: Whether to center crop images (default: True) + max_size: Maximum dimension when variable_size is True (default: 1024) + shard_size: Samples per shard (default: 1000) + max_samples: Maximum samples to process (default: None for all) + compress: Whether to compress tar files with gzip (default: False) + hf_token_file: Path to HuggingFace token file (default: None) + vae_latent_mode: 'presampled' stores only latents; 'resample' additionally + stores mean and logvar for training-time reparameterization (default: 'presampled') + + Example: + >>> pipeline = EncodedDatasetPipeline( + ... source_type='huggingface', + ... output_dir='/data/encoded_pokemon', + ... model_path='black-forest-labs/FLUX.1-dev', + ... batch_size=8, + ... image_size=1024, + ... ) + >>> results = pipeline.run( + ... hf_dataset='diffusers/pokemon-gpt4-captions', + ... hf_split='train' + ... ) + >>> print(f"Processed {results['samples_processed']} samples") + """ + + # Encoded pipeline materializes the full HF dataset (non-streaming) so the + # distributed loader can split by total count. + HF_STREAMING: bool = False + + def __init__( + self, + source_type: str, + output_dir: str, + model_path: str = "black-forest-labs/FLUX.1-dev", + vae_path: Optional[str] = None, + t5_path: Optional[str] = None, + clip_path: Optional[str] = None, + precision: str = "bf16", + device: str = "cuda", + batch_size: int = 8, + t5_max_length: int = 512, + variable_size: bool = False, + image_size: int = 1024, + center_crop: bool = True, + max_size: int = 1024, + shard_size: int = 1000, + max_samples: Optional[int] = None, + compress: bool = False, + hf_token_file: Optional[str] = None, + vae_latent_mode: str = "presampled", + ): + if vae_latent_mode not in ("presampled", "resample"): + raise ValueError(f"vae_latent_mode must be 'presampled' or 'resample', got '{vae_latent_mode}'") + + self.source_type = source_type + self.output_dir = Path(output_dir) + self.model_path = model_path + self.vae_path = vae_path + self.t5_path = t5_path + self.clip_path = clip_path + self.precision = precision + self.device = device + self.batch_size = batch_size + self.t5_max_length = t5_max_length + self.variable_size = variable_size + self.image_size = image_size + self.center_crop = center_crop + self.max_size = max_size + self.shard_size = shard_size + self.max_samples = max_samples + self.compress = compress + self.vae_latent_mode = vae_latent_mode + + # Setup HF authentication if token file provided + if hf_token_file: + from ..auth import HFAuthError, setup_hf_authentication + + try: + setup_hf_authentication(token_file=hf_token_file) + except HFAuthError as e: + raise ValueError(f"HuggingFace authentication failed: {e}") + + # Get distributed info + self.rank, self.world_size = get_distributed_info() + + logger.info(f"Initialized EncodedDatasetPipeline (rank {self.rank}/{self.world_size})") + logger.info(f"Model: {self.model_path}") + logger.info(f"Precision: {self.precision}") + logger.info(f"Batch size: {self.batch_size}") + logger.info(f"T5 max sequence length: {self.t5_max_length}") + logger.info(f"VAE latent mode: {self.vae_latent_mode}") + + # Load encoders + self.vae, self.t5, self.clip = self._load_encoders() + + @staticmethod + def _raise_encoder_auth_error(encoder_name: str, model_path: str, original_error: Exception): + """Raise a clear error when encoder download fails due to authentication.""" + err_str = str(original_error).lower() + is_auth = any( + kw in err_str + for kw in [ + "token", + "permission", + "private repository", + "gated", + "401", + "403", + "authentication", + "login", + ] + ) + if is_auth: + raise RuntimeError( + f"Failed to download {encoder_name} encoder from '{model_path}'.\n" + f"This model likely requires HuggingFace authentication.\n\n" + f"To fix, provide a token using one of:\n" + f" 1. --hf-token-file /path/to/.hf_token\n" + f" 2. export HF_TOKEN=hf_xxx\n" + f" 3. huggingface-cli login\n" + ) from original_error + raise RuntimeError( + f"Failed to load {encoder_name} encoder from '{model_path}': {original_error}" + ) from original_error + + def _load_encoders(self): + """Load VAE, T5, and CLIP encoders.""" + logger.info("Loading encoders...") + + # Determine if we're using FLUX model (needs subfolders) + vae_model_path = self.vae_path or self.model_path + t5_model_path = self.t5_path or self.model_path + clip_model_path = self.clip_path or self.model_path + + is_flux = "FLUX" in self.model_path or "flux" in self.model_path.lower() + + # Create encoder configs with subfolders for FLUX models + vae_config = VAEConfig( + type="autoencoder_kl", + model_path=vae_model_path, + subfolder="vae" if (is_flux and not self.vae_path) else None, + precision=self.precision, + device=self.device, + ) + + t5_config = T5XXLConfig( + type="t5_xxl", + model_path=t5_model_path, + subfolder="text_encoder_2" if (is_flux and not self.t5_path) else None, + tokenizer_subfolder="tokenizer_2" if (is_flux and not self.t5_path) else None, + max_length=self.t5_max_length, + precision=self.precision, + device=self.device, + ) + + clip_config = CLIPLConfig( + type="clip_l", + model_path=clip_model_path, + subfolder="text_encoder" if (is_flux and not self.clip_path) else None, + tokenizer_subfolder="tokenizer" if (is_flux and not self.clip_path) else None, + precision=self.precision, + device=self.device, + ) + + # Load encoders with clear error messages on failure + try: + vae = get_encoder(vae_config) + except Exception as e: + self._raise_encoder_auth_error("VAE", vae_config.model_path, e) + logger.info(f"Loaded VAE from {vae_config.model_path}") + + try: + t5 = get_encoder(t5_config) + except Exception as e: + self._raise_encoder_auth_error("T5-XXL", t5_config.model_path, e) + logger.info(f"Loaded T5-XXL from {t5_config.model_path}") + + try: + clip = get_encoder(clip_config) + except Exception as e: + self._raise_encoder_auth_error("CLIP-L", clip_config.model_path, e) + logger.info(f"Loaded CLIP-L from {clip_config.model_path}") + + # Set to eval mode + vae.eval() + t5.eval() + clip.eval() + + return vae, t5, clip + + def load_data_distributed(self, **source_kwargs): + """ + Load only the data needed for this rank (distributed-aware loading). + + This method loads data more efficiently by having each rank load only + its assigned portion of the dataset, rather than loading everything + and then splitting. + + Args: + **source_kwargs: Source-specific arguments + - For directory: input_dir + - For huggingface: hf_dataset, hf_split + - For webdataset: input_path + + Returns: + List of samples assigned to this rank + """ + if self.source_type == "directory": + # Get file list (cheap operation, all ranks do this) + from pathlib import Path + + from PIL import Image + + input_path = Path(source_kwargs["input_dir"]) + images_dir = input_path / "images" + captions_dir = input_path / "captions" + + if not images_dir.exists(): + raise ValueError(f"Images directory not found: {images_dir}") + if not captions_dir.exists(): + raise ValueError(f"Captions directory not found: {captions_dir}") + + # Get all image files + image_extensions = [".jpg", ".jpeg", ".png", ".webp"] + all_image_files = [] + for ext in image_extensions: + all_image_files.extend(sorted(images_dir.glob(f"*{ext}"))) + + total_files = len(all_image_files) + logger.info(f"Found {total_files} total images") + + # Apply max_samples limit BEFORE splitting across ranks so that + # max_samples refers to the total number of samples, not per-rank. + if self.max_samples and total_files > self.max_samples: + logger.info(f"Limiting to {self.max_samples} total samples (before rank split)") + all_image_files = all_image_files[: self.max_samples] + total_files = self.max_samples + + # Split file list across ranks (before loading!) + start_idx, end_idx = split_work_for_rank(total_files, self.rank, self.world_size) + my_files = all_image_files[start_idx:end_idx] + + logger.info(f"Rank {self.rank}: Loading {len(my_files)} images (indices {start_idx}-{end_idx})") + + # Now load only this rank's files + items = [] + for img_path in my_files: + caption_path = captions_dir / f"{img_path.stem}.txt" + if not caption_path.exists(): + logger.warning(f"Caption not found for {img_path.name}, skipping") + continue + + try: + image = Image.open(img_path).convert("RGB") + with open(caption_path, "r", encoding="utf-8") as f: + caption = f.read().strip() + items.append({"image": image, "caption": caption}) + except Exception as e: + logger.warning(f"Failed to load {img_path.name}: {e}") + + return items + else: + # For non-directory sources, fall back to loading all data + logger.warning( + f"Distributed loading not yet implemented for {self.source_type}, loading all data" + ) + data_iter = self.load_data(**source_kwargs) + all_items = list(data_iter) + total_items = len(all_items) + + # Apply max_samples limit BEFORE splitting across ranks so that + # max_samples refers to the total number of samples, not per-rank. + if self.max_samples and total_items > self.max_samples: + logger.info(f"Limiting to {self.max_samples} total samples (before rank split)") + all_items = all_items[: self.max_samples] + total_items = self.max_samples + + # Split work across ranks + start_idx, end_idx = split_work_for_rank(total_items, self.rank, self.world_size) + return all_items[start_idx:end_idx] + + def _preprocess_images_batch(self, images): + """ + Preprocess batch of PIL images for VAE encoding. + + Args: + images: List of PIL Images + + Returns: + Tensor [B, C, H, W] in range [-1, 1] + """ + tensors = [] + for image in images: + # Ensure RGB mode + if image.mode != "RGB": + image = image.convert("RGB") + + # Convert to numpy array and normalize to [-1, 1] + img_array = np.array(image).astype(np.float32) / 255.0 + img_array = img_array * 2.0 - 1.0 + + # Convert to tensor (HWC -> CHW) + img_tensor = torch.from_numpy(np.transpose(img_array, (2, 0, 1))) + tensors.append(img_tensor) + + return torch.stack(tensors) + + def _encode_batch(self, images, captions): + """ + Encode a batch of images and captions. + + Position IDs are NOT generated during preprocessing - they will be + computed at runtime based on actual tensor shapes. This provides + flexibility for variable-resolution training. + + Args: + images: List of PIL Images + captions: List of caption strings + + Returns: + List of sample dicts with encoded tensors + """ + with torch.no_grad(): + # Preprocess images + images_tensor = self._preprocess_images_batch(images) + images_tensor = images_tensor.to(self.device) + + # Encode images with VAE + if self.vae_latent_mode == "resample": + latents, mean, logvar = self.vae.encode_for_resample(images_tensor) + else: + latents = self.vae.encode(images_tensor) + + # Encode text with T5 + prompt_embeds = self.t5.encode(captions) + + # Encode text with CLIP + _, pooled_prompt_embeds = self.clip.encode(captions) + + # Move to CPU and create samples (NO position IDs) + samples = [] + for i in range(len(images)): + sample = { + "latents.pth": latents[i].cpu(), + "prompt_embeds.pth": prompt_embeds[i].cpu(), + "pooled_prompt_embeds.pth": pooled_prompt_embeds[i].cpu(), + } + if self.vae_latent_mode == "resample": + sample["mean.pth"] = mean[i].cpu() + sample["logvar.pth"] = logvar[i].cpu() + samples.append(sample) + + return samples + + def run(self, **source_kwargs): + """ + Execute the encoded dataset preparation pipeline. + + Args: + **source_kwargs: Source-specific arguments (passed to load_data) + + Returns: + Dictionary with processing statistics: + - samples_processed: Number of samples successfully processed + - samples_skipped: Number of samples that failed + - shards_written: Number of output shards created + """ + import torch.distributed as dist + + # Reproducible encoding: set seeds/cuDNN flags here (not at import time). + set_reproducibility() + + start_time = time.time() + rank_prefix = f"[RANK {self.rank}]" + + tqdm.write(f"{rank_prefix} Stage 1/4: Initialization") + tqdm.write(f"{rank_prefix} Source: {self.source_type}") + tqdm.write( + f"{rank_prefix} Image preprocessing: size={self.image_size}, crop={self.center_crop}, variable_size={self.variable_size}" + ) + + # Synchronization point: ensure all ranks start together + if dist.is_initialized(): + dist.barrier() + + # Load data using distributed-aware loading (each rank loads only its portion) + tqdm.write(f"{rank_prefix} Stage 2/4: Loading data") + items_to_process = self.load_data_distributed(**source_kwargs) + tqdm.write(f"{rank_prefix} Loaded {len(items_to_process)} items") + + # Synchronization point: ensure all ranks finished loading + if dist.is_initialized(): + dist.barrier() + + # Sort items by dimensions AFTER splitting to ensure batches have uniform sizes + # This prevents tensor stacking errors while maintaining balanced workload across ranks + items_to_process = sorted( + items_to_process, key=lambda item: item["image"].size # Returns (width, height) + ) + + total_to_process = len(items_to_process) + + tqdm.write(f"{rank_prefix} Stage 3/4: Encoding samples") + + # Process in batches + batch_images = [] + batch_captions = [] + encoded_samples = [] + samples_processed = 0 + samples_skipped = 0 + shards_written = 0 + + pbar = tqdm( + total=total_to_process, + desc=f"{rank_prefix} Encoding", + unit="sample", + disable=(self.rank != 0 and self.world_size > 4), + ) + + def _flush_batch(): + """Encode and accumulate the current batch, returning count.""" + nonlocal batch_images, batch_captions, encoded_samples, samples_processed + nonlocal shards_written + batch_samples = self._encode_batch(batch_images, batch_captions) + encoded_samples.extend(batch_samples) + count = len(batch_samples) + samples_processed += count + pbar.update(count) + batch_images = [] + batch_captions = [] + + if len(encoded_samples) >= self.shard_size: + shard_offset = shards_written * self.world_size + self.rank + num_shards = save_to_webdataset( + encoded_samples, + str(self.output_dir), + self.shard_size, + shard_offset=shard_offset, + compress=self.compress, + ) + shards_written += num_shards + encoded_samples = [] + + for idx, item in enumerate(items_to_process): + try: + image = item["image"] + if self.image_size or self.center_crop or self.variable_size: + image = preprocess_image( + image, + variable_size=self.variable_size, + size=self.image_size, + center_crop=self.center_crop, + max_size=self.max_size, + ) + + current_size = image.size + + if batch_images and batch_images[0].size != current_size: + _flush_batch() + + batch_images.append(image) + batch_captions.append(item["caption"]) + + if len(batch_images) >= self.batch_size: + _flush_batch() + + except Exception as e: + tqdm.write(f"{rank_prefix} Failed to encode sample {idx}: {e}") + logger.debug(f"Sample {idx} encoding error", exc_info=True) + samples_skipped += 1 + batch_images = [] + batch_captions = [] + continue + + # Process remaining batch + if batch_images: + try: + _flush_batch() + except Exception as e: + tqdm.write(f"{rank_prefix} Failed to encode final batch: {e}") + + # Save remaining samples + if encoded_samples: + shard_offset = shards_written * self.world_size + self.rank + num_shards = save_to_webdataset( + encoded_samples, + str(self.output_dir), + self.shard_size, + shard_offset=shard_offset, + compress=self.compress, + ) + shards_written += num_shards + + pbar.close() + + elapsed_time = time.time() - start_time + tqdm.write(f"{rank_prefix} Stage 4/4: Complete") + tqdm.write( + f"{rank_prefix} Finished in {elapsed_time:.1f}s — " + f"processed: {samples_processed}, skipped: {samples_skipped}, shards: {shards_written}" + ) + + # Generate empty encodings for CFG dropout (rank 0 only). + # Uses the same T5/CLIP models already loaded with the same t5_max_length, + # guaranteeing sequence length consistency with the encoded dataset. + if self.rank == 0: + self._generate_empty_encodings() + + # Critical synchronization point: ensure all ranks finished encoding before returning + if dist.is_initialized(): + dist.barrier() + + return { + "samples_processed": samples_processed, + "samples_skipped": samples_skipped, + "shards_written": shards_written, + } + + def _generate_empty_encodings(self): + """Generate and save T5/CLIP encodings for the empty string (rank 0 only).""" + empty_dir = self.output_dir / "empty_encodings" + empty_dir.mkdir(parents=True, exist_ok=True) + + with torch.no_grad(): + t5_empty = self.t5.encode([""]) + _, clip_empty = self.clip.encode([""]) + + t5_np = t5_empty.cpu().float().numpy() + clip_np = clip_empty.cpu().float().numpy() + + np.save(str(empty_dir / "t5_empty.npy"), t5_np) + np.save(str(empty_dir / "clip_empty.npy"), clip_np) + + tqdm.write( + f"[RANK 0] Saved empty encodings to {empty_dir} " f"(t5={t5_np.shape}, clip={clip_np.shape})" + ) diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py new file mode 100644 index 000000000..5c68bc572 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/ingest.py @@ -0,0 +1,311 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Streaming Arrow-to-WebDataset pipeline for MLPerf Flux datasets. + +Downloads Apache Arrow IPC files from MLCommons R2 storage using a +concurrent prefetch pool, converts numpy-serialized bfloat16 samples +directly into WebDataset tar shards, then deletes each Arrow file to +minimize peak disk usage. + +Reduces the 6 TB disk requirement to roughly the final dataset size +(~1.2 TB for cc12m) plus a small prefetch buffer of Arrow files. +""" + +import io +import json +import logging +import queue +import tarfile +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import pyarrow.ipc + +from ..download import download_with_backoff, fetch_manifest +from .base import DatasetPipeline + +logger = logging.getLogger(__name__) + +ARROW_COLUMNS = ("t5_encodings", "clip_encodings", "mean", "logvar") +WDS_KEYS = ("t5.bytes", "clip.bytes", "mean.bytes", "logvar.bytes") + +_SENTINEL = None + + +def _arrow_to_tar( + arrow_path: Path, + tar_path: Path, + global_sample_offset: int, +) -> int: + """Convert one Arrow IPC stream file into a WebDataset tar shard. + + Reads the Arrow stream, writes each row as a set of compound-key + entries (e.g. ``00000000.t5.bytes``) into a tar archive. Uses the + ``__key__`` column from the Arrow file when available, otherwise + generates sequential keys. + + Returns the number of samples written. + """ + reader = pyarrow.ipc.open_stream(str(arrow_path)) + table = reader.read_all() + num_rows = table.num_rows + + has_key_col = "__key__" in table.schema.names + + tar_path.parent.mkdir(parents=True, exist_ok=True) + with tarfile.open(str(tar_path), "w") as tar: + for row_idx in range(num_rows): + if has_key_col: + base_name = table.column("__key__")[row_idx].as_py() + else: + base_name = f"{global_sample_offset + row_idx:08d}" + + for col_name, wds_key in zip(ARROW_COLUMNS, WDS_KEYS): + col = table.column(col_name) + raw_bytes = col[row_idx].as_py() + if raw_bytes is None: + continue + data = bytes(raw_bytes) + + entry_name = f"{base_name}.{wds_key}" + info = tarfile.TarInfo(name=entry_name) + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + + meta = json.dumps({"key": base_name}).encode("utf-8") + meta_info = tarfile.TarInfo(name=f"{base_name}.json") + meta_info.size = len(meta) + tar.addfile(meta_info, io.BytesIO(meta)) + + return num_rows + + +class StreamingIngestPipeline(DatasetPipeline): + """Streaming pipeline: download Arrow files from R2 -> WebDataset tars. + + Downloads are parallelized with a prefetch pool (``max_workers`` threads) + while conversion runs sequentially on the main thread, preserving + deterministic shard ordering. A semaphore (``prefetch_depth``) limits + how many Arrow files are kept on disk at once. + + Args: + manifest_url: URL to the .uri manifest on MLCommons R2. + input_dir: Local directory for temporary Arrow file storage. + output_dir: Output directory for WebDataset tar shards. + split_name: Subdirectory name for the split (e.g. 'train', 'val'). + max_files: Maximum number of Arrow files to process (None = all). + max_workers: Number of concurrent download threads (default 4). + prefetch_depth: Max Arrow files buffered on disk ahead of + conversion (default 6). Disk overhead ~ prefetch_depth x 200 MB. + """ + + def __init__( + self, + manifest_url: str, + input_dir: str, + output_dir: str, + split_name: str = "train", + max_files: Optional[int] = None, + max_workers: int = 4, + prefetch_depth: int = 6, + ): + self.manifest_url = manifest_url + self.input_dir = Path(input_dir) + self.output_dir = Path(output_dir) / split_name + self.split_name = split_name + self.max_files = max_files + self.max_workers = max_workers + self.prefetch_depth = prefetch_depth + + def run(self, **kwargs) -> Dict[str, int]: + """Execute the streaming conversion pipeline. + + Returns: + Dict with 'files_processed', 'samples_written', 'shards_created', + 'shards_skipped', 'files_failed'. + """ + base_url, entries = fetch_manifest(self.manifest_url) + + if self.max_files is not None: + entries = entries[: self.max_files] + + self.output_dir.mkdir(parents=True, exist_ok=True) + self.input_dir.mkdir(parents=True, exist_ok=True) + + total_files = len(entries) + + logger.info("=" * 80) + logger.info(f"Streaming Arrow->WebDataset conversion: {self.split_name}") + logger.info(f" Arrow files to process: {total_files}") + logger.info(f" Output: {self.output_dir}") + logger.info(f" Download workers: {self.max_workers}, " f"prefetch depth: {self.prefetch_depth}") + logger.info("=" * 80) + + prefetch_q: queue.Queue[Union[Tuple[int, str, Path], None]] = queue.Queue() + download_semaphore = threading.Semaphore(self.prefetch_depth) + producer_error: List[BaseException] = [] + cancel_event = threading.Event() + + failed_files: List[Dict] = [] + failed_lock = threading.Lock() + + def _download_one(file_idx: int, md5: str, fname: str) -> Tuple[int, str, Path]: + file_url = f"{base_url}/{fname}" + # fname comes from a remote manifest; never let it traverse outside + # input_dir (e.g. "../../etc/passwd"). Use the basename and verify + # the resolved path stays under input_dir before writing. + safe_name = Path(fname).name + base_dir = self.input_dir.resolve() + arrow_path = (base_dir / safe_name).resolve() + if not safe_name or arrow_path.parent != base_dir: + raise ValueError(f"Unsafe manifest filename rejected: {fname!r}") + logger.info(f" [download {file_idx + 1}/{total_files}] {fname}") + download_with_backoff(file_url, arrow_path, expected_md5=md5) + return file_idx, fname, arrow_path + + skipped_count = [0] + + def _producer() -> None: + """Submit downloads and drain results concurrently. + + A separate drain thread processes completed futures in submission + order and feeds the prefetch queue. The semaphore acquire in the + submit loop blocks when ``prefetch_depth`` files are already + in-flight or queued, bounding temporary disk usage. + """ + try: + with ThreadPoolExecutor(max_workers=self.max_workers) as pool: + futures_q: queue.Queue = queue.Queue() + + def _drain() -> None: + while True: + item = futures_q.get() + if item is _SENTINEL: + break + fut, file_idx, fname = item + if cancel_event.is_set(): + download_semaphore.release() + continue + try: + result = fut.result() + prefetch_q.put(result) + except Exception as exc: + download_semaphore.release() + logger.warning( + f" [FAILED download " f"{file_idx + 1}/{total_files}] " f"{fname}: {exc}" + ) + with failed_lock: + failed_files.append( + { + "file_idx": file_idx, + "filename": fname, + "error": str(exc), + "stage": "download", + } + ) + + drain_thread = threading.Thread(target=_drain, daemon=True) + drain_thread.start() + + for file_idx, (md5, fname) in enumerate(entries): + if cancel_event.is_set(): + break + tar_path = self.output_dir / f"shard_{file_idx:06d}.tar" + if tar_path.exists(): + skipped_count[0] += 1 + logger.info(f" [skip {file_idx + 1}/{total_files}] " f"{fname} (shard exists)") + continue + download_semaphore.acquire() + fut = pool.submit(_download_one, file_idx, md5, fname) + futures_q.put((fut, file_idx, fname)) + + futures_q.put(_SENTINEL) + drain_thread.join() + except BaseException as exc: + producer_error.append(exc) + finally: + prefetch_q.put(_SENTINEL) + + producer_thread = threading.Thread(target=_producer, daemon=True) + producer_thread.start() + + global_sample_offset = 0 + total_samples = 0 + files_processed = 0 + + try: + while True: + item = prefetch_q.get() + if item is _SENTINEL: + break + + file_idx, fname, arrow_path = item + + shard_name = f"shard_{file_idx:06d}.tar" + tar_path = self.output_dir / shard_name + + try: + num_samples = _arrow_to_tar(arrow_path, tar_path, global_sample_offset) + except Exception as exc: + logger.warning(f" [FAILED convert {file_idx + 1}/{total_files}] " f"{fname}: {exc}") + tar_path.unlink(missing_ok=True) + arrow_path.unlink(missing_ok=True) + download_semaphore.release() + with failed_lock: + failed_files.append( + { + "file_idx": file_idx, + "filename": fname, + "error": str(exc), + "stage": "conversion", + } + ) + continue + + global_sample_offset += num_samples + total_samples += num_samples + files_processed += 1 + + arrow_path.unlink(missing_ok=True) + download_semaphore.release() + + logger.info( + f"[{files_processed}/{total_files}] {fname} -> {shard_name}: " + f"{num_samples} samples (cumulative: {total_samples})" + ) + except BaseException: + cancel_event.set() + raise + finally: + producer_thread.join(timeout=10) + + if producer_error: + raise RuntimeError(f"Download failed: {producer_error[0]}") from producer_error[0] + + if failed_files: + manifest_path = self.output_dir / "failed_files.json" + manifest_path.write_text(json.dumps(failed_files, indent=2)) + logger.warning(f"{len(failed_files)} file(s) failed. " f"See {manifest_path}") + + skipped = skipped_count[0] + failed = len(failed_files) + logger.info("=" * 80) + logger.info( + f"Done {self.split_name}: {files_processed} new + {skipped} skipped" + f"{f' + {failed} failed' if failed else ''}" + f" -> {files_processed + skipped} shards, " + f"{total_samples} new samples" + ) + logger.info("=" * 80) + + return { + "files_processed": files_processed, + "samples_written": total_samples, + "shards_created": files_processed + skipped, + "shards_skipped": skipped, + "files_failed": failed, + } diff --git a/primus/backends/megatron/data/diffusion/preprocessing/pipelines/raw.py b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/raw.py new file mode 100644 index 000000000..7902c1bb1 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/pipelines/raw.py @@ -0,0 +1,234 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Raw dataset preparation pipeline for Megatron diffusion models. + +Creates raw Energon WebDatasets with images and captions for on-the-fly +encoding during training. +""" + +import logging +from pathlib import Path +from typing import Optional + +from ..utils import ( + encode_image_to_bytes, + get_distributed_info, + preprocess_image, + save_to_webdataset, +) + +logger = logging.getLogger(__name__) + + +from .base import DatasetPipeline + + +class RawDatasetPipeline(DatasetPipeline): + """ + Pipeline for creating raw Energon WebDataset with images and captions. + + This pipeline creates smaller datasets where encoding (VAE/T5/CLIP) happens + on-the-fly during training. Suitable for experimentation and when storage + is limited. + + Args: + source_type: Type of input source ('directory', 'huggingface', 'webdataset') + output_dir: Output directory for Energon WebDataset + image_size: Target image size (default: 1024) + center_crop: Whether to center crop images (default: True) + image_format: Output image format (default: 'JPEG') + image_quality: JPEG/WEBP quality 1-100 (default: 95) + shard_size: Samples per shard (default: 1000) + max_samples: Maximum samples to process (default: None for all) + compress: Whether to compress tar files with gzip (default: False) + hf_token_file: Path to HuggingFace token file (default: None) + variable_size: If True, resize to nearest multiple of 16 instead of fixed size (default: False) + max_size: Maximum dimension when variable_size is True (default: 1024) + + Example: + >>> pipeline = RawDatasetPipeline( + ... source_type='huggingface', + ... output_dir='/data/raw_pokemon', + ... image_size=1024, + ... center_crop=True, + ... ) + >>> results = pipeline.run( + ... hf_dataset='diffusers/pokemon-gpt4-captions', + ... hf_split='train' + ... ) + >>> print(f"Processed {results['samples_processed']} samples") + """ + + def __init__( + self, + source_type: str, + output_dir: str, + variable_size: bool = False, + image_size: int = 1024, + center_crop: bool = True, + max_size: int = 1024, + image_format: str = "JPEG", + image_quality: int = 95, + shard_size: int = 1000, + max_samples: Optional[int] = None, + compress: bool = False, + hf_token_file: Optional[str] = None, + ): + self.source_type = source_type + self.output_dir = Path(output_dir) + self.variable_size = variable_size + self.image_size = image_size + self.center_crop = center_crop + self.max_size = max_size + self.image_format = image_format + self.image_quality = image_quality + self.shard_size = shard_size + self.max_samples = max_samples + self.compress = compress + + # Setup HF authentication if token file provided + if hf_token_file: + from ..auth import HFAuthError, setup_hf_authentication + + try: + setup_hf_authentication(token_file=hf_token_file) + except HFAuthError as e: + raise ValueError(f"HuggingFace authentication failed: {e}") + + # Get distributed info + self.rank, self.world_size = get_distributed_info() + + # Determine output format extension + self.format_ext = { + "JPEG": "jpg", + "PNG": "png", + "WEBP": "webp", + }[self.image_format] + + logger.info(f"Initialized RawDatasetPipeline (rank {self.rank}/{self.world_size})") + logger.info(f"Output directory: {self.output_dir}") + logger.info(f"Image preprocessing: size={self.image_size}, crop={self.center_crop}") + + def process_sample(self, item): + """ + Process a single sample. + + Args: + item: Sample dict with 'image' (PIL Image) and 'caption' (str) + + Returns: + Sample dict with image bytes and caption text + """ + # Preprocess image + image = preprocess_image( + item["image"], + variable_size=self.variable_size, + size=self.image_size, + center_crop=self.center_crop, + max_size=self.max_size, + ) + + # Encode to bytes + image_bytes = encode_image_to_bytes(image, format=self.image_format, quality=self.image_quality) + + # Create sample with standard keys + return { + self.format_ext: image_bytes, + "txt": item["caption"], + } + + def run(self, **source_kwargs): + """ + Execute the raw dataset preparation pipeline. + + Args: + **source_kwargs: Source-specific arguments (passed to load_data) + + Returns: + Dictionary with processing statistics: + - samples_processed: Number of samples successfully processed + - samples_skipped: Number of samples that failed + - shards_written: Number of output shards created + """ + logger.info(f"Starting raw dataset preparation (rank {self.rank}/{self.world_size})") + logger.info(f"Source: {self.source_type}") + + # Load data as a lazy iterator (never materialized in full) + data_iter = self.load_data(**source_kwargs) + + if self.world_size > 1: + logger.info( + f"Distributed mode: rank {self.rank}/{self.world_size} " + f"processing every {self.world_size}th item (round-robin by load index)" + ) + + # Process and accumulate samples + samples = [] + samples_processed = 0 + samples_skipped = 0 + shards_written = 0 + + # Stream items and assign them to ranks round-robin by global load + # index, so the full dataset is never held in memory. max_samples caps + # the GLOBAL index (total across ranks, not per-rank), preserving the + # prior semantics; the rank split changes from contiguous ranges to + # round-robin (same overall dataset, more balanced load). The index is + # assigned at load time, before process_sample, so the rank assignment + # stays deterministic even when samples are skipped. + for global_idx, item in enumerate(data_iter): + # Truthy check (not `is not None`) mirrors the original semantics + # where max_samples=0 / None meant "no limit". + if self.max_samples and global_idx >= self.max_samples: + break + if global_idx % self.world_size != self.rank: + continue + try: + sample = self.process_sample(item) + samples.append(sample) + samples_processed += 1 + + # Log progress + if samples_processed % 100 == 0: + logger.info(f"Processed {samples_processed} samples (skipped {samples_skipped})") + + # Save shard when full + if len(samples) >= self.shard_size: + # Use distributed-aware shard naming to prevent conflicts + shard_offset = shards_written * self.world_size + self.rank + if self.world_size > 1 and shards_written == 0: + logger.info(f"Rank {self.rank}: First shard will be numbered {shard_offset:06d}.tar") + + num_shards = save_to_webdataset( + samples, + str(self.output_dir), + self.shard_size, + shard_offset=shard_offset, + compress=self.compress, + ) + shards_written += num_shards + samples = [] + + except Exception as e: + logger.warning(f"Failed to process sample {global_idx}: {e}") + samples_skipped += 1 + continue + + # Save remaining samples + if samples: + shard_offset = shards_written * self.world_size + self.rank + num_shards = save_to_webdataset( + samples, + str(self.output_dir), + self.shard_size, + shard_offset=shard_offset, + compress=self.compress, + ) + shards_written += num_shards + + return { + "samples_processed": samples_processed, + "samples_skipped": samples_skipped, + "shards_written": shards_written, + } diff --git a/primus/backends/megatron/data/diffusion/preprocessing/utils.py b/primus/backends/megatron/data/diffusion/preprocessing/utils.py new file mode 100644 index 000000000..cd57d6836 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/utils.py @@ -0,0 +1,546 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Common utilities for dataset preprocessing. + +This module provides shared functionality for creating WebDataset shards +from various input sources (HuggingFace Hub, directories, existing WebDatasets). +""" + +import io +import logging +import os +from pathlib import Path +from typing import Any, Dict, Iterator, List, Optional, Tuple, Union + +import torch +import webdataset as wds +from PIL import Image + +logger = logging.getLogger(__name__) + + +def get_distributed_info() -> Tuple[int, int]: + """ + Get distributed processing information (rank, world_size). + + Returns: + Tuple of (rank, world_size). Returns (0, 1) if not in distributed mode. + + Example: + >>> rank, world_size = get_distributed_info() + >>> if rank == 0: + ... print("This is the master process") + """ + try: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + return rank, world_size + except (ImportError, AttributeError): + pass + + # Fallback to environment variables + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + return rank, world_size + + +def split_work_for_rank(total: int, rank: int, world_size: int) -> Tuple[int, int]: + """ + Calculate start/end indices for this rank in distributed processing. + + Args: + total: Total number of items to process + rank: Current process rank (0 to world_size-1) + world_size: Total number of processes + + Returns: + Tuple of (start_idx, end_idx) for this rank + + Example: + >>> start, end = split_work_for_rank(1000, 0, 4) + >>> print(f"Process 0 handles items {start} to {end}") + Process 0 handles items 0 to 250 + """ + items_per_rank = total // world_size + start_idx = rank * items_per_rank + + if rank == world_size - 1: + # Last rank gets all remaining items (including remainder) + end_idx = total + else: + end_idx = start_idx + items_per_rank + + return start_idx, end_idx + + +def save_to_webdataset( + samples: List[Dict[str, Any]], + output_dir: str, + shard_size: int, + shard_offset: int = 0, + compress: bool = False, +) -> int: + """ + Save samples to WebDataset tar shards. + + Args: + samples: List of sample dictionaries with keys as extensions + Example: {'images': image_bytes, 'txt': caption_text} + output_dir: Directory to save shards + shard_size: Number of samples per shard + shard_offset: Starting shard number (for distributed processing) + compress: Whether to compress tar files (gzip) + + Returns: + Number of shards created + + Example: + >>> samples = [ + ... {'images': image_bytes, 'txt': 'a cat'}, + ... {'images': image_bytes2, 'txt': 'a dog'}, + ... ] + >>> num_shards = save_to_webdataset(samples, '/path/to/output', shard_size=1000) + """ + os.makedirs(output_dir, exist_ok=True) + + # Calculate number of shards needed + num_shards = (len(samples) + shard_size - 1) // shard_size + + for shard_idx in range(num_shards): + start_idx = shard_idx * shard_size + end_idx = min(start_idx + shard_size, len(samples)) + shard_samples = samples[start_idx:end_idx] + + # Create shard filename + shard_num = shard_offset + shard_idx + ext = ".tar.gz" if compress else ".tar" + shard_path = os.path.join(output_dir, f"{shard_num:06d}{ext}") + + # Write shard + with wds.TarWriter(shard_path) as sink: + for sample_idx, sample in enumerate(shard_samples): + # Generate unique key for this sample + key = f"{shard_num:06d}_{sample_idx:06d}" + + # Create WebDataset sample dict + wds_sample = {"__key__": key} + + for ext, data in sample.items(): + if ext.startswith("__"): + continue # Skip special keys + + # Handle different data types + if isinstance(data, bytes): + wds_sample[ext] = data + elif isinstance(data, str): + wds_sample[ext] = data.encode("utf-8") + elif isinstance(data, torch.Tensor): + # Save tensor to bytes + buffer = io.BytesIO() + torch.save(data, buffer) + wds_sample[ext] = buffer.getvalue() + elif isinstance(data, Image.Image): + # Save PIL image to bytes + buffer = io.BytesIO() + data.save(buffer, format="JPEG") + wds_sample[ext] = buffer.getvalue() + else: + logger.warning(f"Unknown data type for key {ext}: {type(data)}") + continue + + sink.write(wds_sample) + + logger.info(f"Created shard {shard_path} with {len(shard_samples)} samples") + + return num_shards + + +# Default keys to try when loading from HuggingFace if not specified in config +DEFAULT_IMAGE_KEYS = ["image", "jpg", "jpeg", "png", "img", "photo"] +DEFAULT_CAPTION_KEYS = ["caption", "text", "txt", "description", "prompt"] + + +def _extract_field(item: dict, field_keys: List[str], idx: int, field_type: str) -> Any: + """ + Extract field from item using list of possible keys. + Supports dot notation for JSON paths (e.g., 'json.caption'). + + Args: + item: Dataset sample dictionary + field_keys: List of keys to try (e.g., ['caption', 'json.caption']) + idx: Sample index (for logging) + field_type: 'image' or 'caption' (currently unused) + + Returns: + Field value if found, None otherwise + + Example: + >>> item = {'jpg': image_bytes, 'json': {'caption': 'a cat'}} + >>> _extract_field(item, ['caption', 'json.caption'], 0, 'caption') + 'a cat' + """ + for key in field_keys: + if "." in key: + # Handle JSON path (e.g., 'json.caption') + value = _extract_json_path(item, key, idx) + if value is not None: + return value + elif key in item: + return item[key] + return None + + +def _extract_json_path(item: dict, path: str, idx: int) -> Any: + """ + Extract value from nested JSON using dot notation. + + Args: + item: Dataset sample dictionary + path: Dot-separated path (e.g., 'json.metadata.caption') + idx: Sample index (for logging) + + Returns: + Extracted value if found, None otherwise + + Example: + >>> item = {'json': b'{"caption": "a cat"}'} + >>> _extract_json_path(item, 'json.caption', 0) + 'a cat' + """ + import json as json_module + + parts = path.split(".") + current = item + + for i, part in enumerate(parts): + if part not in current: + return None + + current = current[part] + + # If we hit a JSON string/bytes at any level, parse it + if isinstance(current, (bytes, str)): + try: + if isinstance(current, bytes): + current = json_module.loads(current.decode("utf-8")) + elif isinstance(current, str) and (current.startswith("{") or current.startswith("[")): + current = json_module.loads(current) + except Exception as e: + logger.debug(f"Sample {idx} failed to parse JSON at '{part}': {e}") + return None + + return current if isinstance(current, (str, int, float)) else None + + +def load_from_huggingface( + dataset_name: str, + split: str = "train", + streaming: bool = True, + data_files: Optional[Union[str, List[str]]] = None, + image_key: Optional[str] = None, + caption_key: Optional[str] = None, + image_keys: Optional[List[str]] = None, + caption_keys: Optional[List[str]] = None, +) -> Iterator[Dict[str, Any]]: + """ + Load dataset from HuggingFace Hub with configurable field mappings. + + Args: + dataset_name: HF dataset identifier (e.g., 'laion/laion400m') + split: Dataset split to load + streaming: Whether to stream the dataset (recommended for large datasets) + data_files: Specific files/paths to load from the dataset (e.g., 'data_1024_10K/*.tar') + image_key: Single image field name (e.g., 'jpg', 'image') + caption_key: Single caption field or JSON path (e.g., 'caption', 'json.caption') + image_keys: List of image field names to try (fallback to defaults if None) + caption_keys: List of caption fields/paths to try (fallback to defaults if None) + + Yields: + Dicts with 'image' (PIL Image) and 'caption' (str) keys + + Example: + >>> # Load with auto-detection (uses defaults) + >>> for item in load_from_huggingface('diffusers/pokemon-gpt4-captions'): + ... print(f"Caption: {item['caption']}") + + >>> # Load with specific field mappings + >>> for item in load_from_huggingface( + ... 'jackyhate/text-to-image-2M', + ... image_key='jpg', + ... caption_key='json.caption', + ... data_files='data_1024_10K/*.tar' + ... ): + ... print(f"Caption: {item['caption']}") + """ + try: + from datasets import load_dataset + except ImportError: + raise ImportError( + "The 'datasets' package is required for HuggingFace dataset loading. " + "Install with: pip install datasets" + ) + + if data_files: + logger.info(f"Loading HuggingFace dataset: {dataset_name} (split: {split}, data_files: {data_files})") + else: + logger.info(f"Loading HuggingFace dataset: {dataset_name} (split: {split})") + + dataset = load_dataset(dataset_name, split=split, streaming=streaming, data_files=data_files) + + # Determine which keys to try + if image_keys is None: + image_keys = [image_key] if image_key else DEFAULT_IMAGE_KEYS + if caption_keys is None: + caption_keys = [caption_key] if caption_key else DEFAULT_CAPTION_KEYS + + for idx, item in enumerate(dataset): + # Try to extract image using configured keys + image = _extract_field(item, image_keys, idx, "image") + + # Try to extract caption using configured keys (supports JSON paths) + caption = _extract_field(item, caption_keys, idx, "caption") + + if isinstance(caption, list): + caption = caption[0] if caption else None + + if image is None or caption is None: + logger.warning(f"Sample {idx} missing image or caption, skipping. Keys: {item.keys()}") + continue + + # Ensure image is PIL Image + if not isinstance(image, Image.Image): + try: + if isinstance(image, bytes): + image = Image.open(io.BytesIO(image)) + else: + logger.warning(f"Sample {idx} has unexpected image type: {type(image)}") + continue + except Exception as e: + logger.warning(f"Sample {idx} failed to load image: {e}") + continue + + yield { + "image": image, + "caption": caption, + } + + +def load_from_directory(input_dir: str) -> Iterator[Dict[str, Any]]: + """ + Load dataset from directory structure. + + Expected structure: + input_dir/ + images/ + 0000000.jpg + 0000001.jpg + captions/ + 0000000.txt + 0000001.txt + + Args: + input_dir: Root directory containing 'images' and 'captions' subdirectories + + Yields: + Dicts with 'image' (PIL Image) and 'caption' (str) keys + + Example: + >>> for item in load_from_directory('/path/to/dataset'): + ... print(f"Caption: {item['caption']}") + """ + input_path = Path(input_dir) + images_dir = input_path / "images" + captions_dir = input_path / "captions" + + if not images_dir.exists(): + raise ValueError(f"Images directory not found: {images_dir}") + + if not captions_dir.exists(): + raise ValueError(f"Captions directory not found: {captions_dir}") + + logger.info(f"Loading dataset from directory: {input_dir}") + + # Find all image files + image_extensions = [".jpg", ".jpeg", ".png", ".webp"] + image_files = [] + for ext in image_extensions: + image_files.extend(sorted(images_dir.glob(f"*{ext}"))) + + logger.info(f"Found {len(image_files)} image files") + + for img_path in image_files: + # Find corresponding caption file + caption_path = captions_dir / f"{img_path.stem}.txt" + + if not caption_path.exists(): + logger.warning(f"Caption not found for {img_path.name}, skipping") + continue + + try: + # Load image + image = Image.open(img_path).convert("RGB") + + # Load caption + with open(caption_path, "r", encoding="utf-8") as f: + caption = f.read().strip() + + yield { + "image": image, + "caption": caption, + } + + except Exception as e: + logger.warning(f"Failed to load {img_path.name}: {e}") + continue + + +def load_from_webdataset(input_path: str) -> Iterator[Dict[str, Any]]: + """ + Load existing WebDataset with non-standard keys. + + Converts to standard format {'image': PIL Image, 'caption': str} + + Args: + input_path: Path or glob pattern to WebDataset tar files + Examples: '/path/to/shards/*.tar', '/path/to/shard_000000.tar' + + Yields: + Dicts with 'image' (PIL Image) and 'caption' (str) keys + + Example: + >>> for item in load_from_webdataset('/path/to/dataset/*.tar'): + ... print(f"Caption: {item['caption']}") + """ + logger.info(f"Loading WebDataset from: {input_path}") + + dataset = wds.WebDataset(input_path) + + for sample in dataset: + try: + # Try to find image with various keys + image = None + for img_key in ["jpg", "png", "jpeg", "webp", "image"]: + if img_key in sample: + img_data = sample[img_key] + if isinstance(img_data, bytes): + image = Image.open(io.BytesIO(img_data)).convert("RGB") + elif isinstance(img_data, Image.Image): + image = img_data.convert("RGB") + else: + continue + break + + # Try to find caption with various keys + caption = None + for cap_key in ["txt", "caption.txt", "text", "caption"]: + if cap_key in sample: + cap_data = sample[cap_key] + if isinstance(cap_data, bytes): + caption = cap_data.decode("utf-8") + elif isinstance(cap_data, str): + caption = cap_data + else: + continue + break + + if image is None or caption is None: + logger.warning(f"Sample missing image or caption, skipping. Keys: {sample.keys()}") + continue + + yield { + "image": image, + "caption": caption, + } + + except Exception as e: + logger.warning(f"Failed to load sample: {e}") + continue + + +def preprocess_image( + image: Image.Image, + variable_size: bool = True, # if True, then image is resized to the nearest multiple of 16, otherwise it is resized to the given size + size: int = 1024, # only used if variable_size is False, then this is applied + center_crop: bool = False, # only used if variable_size is False, then this is applied + max_size: int = 1024, # maximum dimension when variable_size is True +) -> Image.Image: + """ + Preprocess image (resize, crop). + + Args: + image: PIL Image + variable_size: If True, resize to nearest multiple of 16 up to max_size + size: Target size (square) - only used if variable_size is False + center_crop: Whether to center crop before resize - only used if variable_size is False + max_size: Maximum dimension when variable_size is True (default: 1024) + + Returns: + Preprocessed PIL Image + + Example: + >>> image = Image.open('photo.jpg') + >>> processed = preprocess_image(image, variable_size=True, max_size=2048) + """ + # Center crop if requested + if center_crop and variable_size is False: + width, height = image.size + crop_size = min(width, height) + left = (width - crop_size) // 2 + top = (height - crop_size) // 2 + right = left + crop_size + bottom = top + crop_size + image = image.crop((left, top, right, bottom)) + + # Resize + if variable_size is False: + image = image.resize((size, size), Image.Resampling.LANCZOS) + else: + width, height = image.size + if max(width, height) > max_size: + scale = max_size / float(max(width, height)) + new_w = int(round(width * scale)) + new_h = int(round(height * scale)) + # High-quality downsampling + image = image.resize((new_w, new_h), resample=Image.Resampling.LANCZOS) + + width, height = image.size + new_size = (round(width / 16) * 16, round(height / 16) * 16) + image = image.resize(new_size, Image.LANCZOS) + + return image + + +def encode_image_to_bytes(image: Image.Image, format: str = "JPEG", quality: int = 95) -> bytes: + """ + Encode PIL Image to bytes. + + Args: + image: PIL Image + format: Image format ('JPEG', 'PNG', 'WEBP') + quality: JPEG/WEBP quality (1-100) + + Returns: + Image bytes + + Example: + >>> image = Image.open('photo.jpg') + >>> image_bytes = encode_image_to_bytes(image, format='JPEG', quality=95) + """ + buffer = io.BytesIO() + image.save(buffer, format=format, quality=quality) + return buffer.getvalue() + + +__all__ = [ + "get_distributed_info", + "split_work_for_rank", + "save_to_webdataset", + "load_from_huggingface", + "load_from_directory", + "load_from_webdataset", + "preprocess_image", + "encode_image_to_bytes", +] diff --git a/primus/backends/megatron/data/diffusion/preprocessing/validate.py b/primus/backends/megatron/data/diffusion/preprocessing/validate.py new file mode 100644 index 000000000..769cea725 --- /dev/null +++ b/primus/backends/megatron/data/diffusion/preprocessing/validate.py @@ -0,0 +1,302 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Energon dataset validation for Primus diffusion pipelines. + +Validates that a prepared Energon WebDataset is structurally correct and +loadable by the training code. Replaces the broken ``energon info`` CLI +tool which does not support CrudeWebdataset format. + +Standalone usage: + python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset + +Programmatic usage: + from primus.backends.megatron.data.diffusion.preprocessing.validate import ( + validate_energon_dataset, + ) + validate_energon_dataset('/path/to/dataset', encoding='preencoded') +""" + +import json +import logging +import tarfile +from pathlib import Path +from typing import Dict, List, Literal, Optional + +import torch +import yaml + +logger = logging.getLogger(__name__) + +EXPECTED_KEYS_PREENCODED = {"latents.pth", "prompt_embeds.pth", "pooled_prompt_embeds.pth"} +EXPECTED_KEYS_PREENCODED_NUMPY = {"t5.bytes", "clip.bytes", "mean.bytes", "logvar.bytes"} +EXPECTED_KEYS_RAW_IMAGE = {"jpg", "jpeg", "png", "webp"} + + +def _check_metadata(output_path: Path) -> Optional[Dict]: + """Check that .nv-meta files exist and are valid. Returns info dict or None.""" + meta_dir = output_path / ".nv-meta" + + info_path = meta_dir / ".info.json" + if not info_path.exists(): + info_path_yaml = meta_dir / ".info.yaml" + if info_path_yaml.exists(): + info_path = info_path_yaml + else: + logger.error(f"Missing metadata: neither .info.json nor .info.yaml in {meta_dir}") + return None + + try: + with open(info_path) as f: + if info_path.suffix == ".json": + info = json.load(f) + else: + info = yaml.safe_load(f) + except Exception as e: + logger.error(f"Failed to parse {info_path}: {e}") + return None + + if "shard_counts" not in info: + logger.error(f"Missing 'shard_counts' in {info_path}") + return None + + split_path = meta_dir / "split.yaml" + if not split_path.exists(): + split_path = meta_dir / "split.json" + if not split_path.exists(): + logger.error(f"Missing split config in {meta_dir}") + return None + + try: + with open(split_path) as f: + splits = yaml.safe_load(f) if split_path.suffix == ".yaml" else json.load(f) + except Exception as e: + logger.error(f"Failed to parse {split_path}: {e}") + return None + + train_parts = splits.get("split_parts", {}).get("train", []) + if not train_parts: + logger.warning("Train split is empty") + + dataset_yaml = meta_dir / "dataset.yaml" + if not dataset_yaml.exists(): + logger.error(f"Missing {dataset_yaml}") + return None + + missing_files: List[str] = [] + for shard_name in info["shard_counts"]: + if not (output_path / shard_name).exists(): + missing_files.append(shard_name) + idx_name = shard_name + ".idx" + if not (output_path / idx_name).exists(): + missing_files.append(idx_name) + + if missing_files: + logger.error(f"Missing files on disk: {missing_files[:10]}") + return None + + info["_splits"] = splits + return info + + +def _check_sample_counts(output_path: Path, info: Dict) -> bool: + """Spot-check that the first shard's entry count matches .info.json.""" + shard_counts = info["shard_counts"] + first_shard = next(iter(shard_counts)) + expected = shard_counts[first_shard] + + tar_path = output_path / first_shard + try: + with tarfile.open(str(tar_path), "r") as tar: + members = tar.getmembers() + sample_keys = set() + for m in members: + key = m.name.split(".", 1)[0] + sample_keys.add(key) + actual = len(sample_keys) + except Exception as e: + logger.error(f"Failed to read {tar_path}: {e}") + return False + + if actual != expected: + logger.error( + f"Sample count mismatch in {first_shard}: " f".info.json says {expected}, tar contains {actual}" + ) + return False + + return True + + +def _check_sample_load( + output_path: Path, + encoding: str, +) -> Optional[Dict[str, str]]: + """ + Load one sample through energon's Python API (same path as training). + Returns a dict of {key: description} for the sample, or None on failure. + """ + try: + from megatron.energon import StandardWebdatasetFactory + from megatron.energon.dataset_config import load_config + from megatron.energon.epathlib import EPath + from megatron.energon.flavors.webdataset.config import MAIN_FOLDER_NAME + from megatron.energon.loader import get_loader + from megatron.energon.worker import WorkerConfig + except ImportError as e: + logger.warning(f"Cannot import megatron.energon for sample validation: {e}") + return None + + try: + ds_path = EPath(str(output_path)) + worker_config = WorkerConfig(rank=0, world_size=1, num_workers=0) + dataset = load_config( + ds_path / MAIN_FOLDER_NAME / "dataset.yaml", + default_kwargs=dict( + path=ds_path, + split_part="train", + training=False, + worker_config=worker_config, + ), + default_type=StandardWebdatasetFactory, + ) + sample = next(iter(get_loader(dataset.build()))) + except Exception as e: + logger.error(f"Failed to load sample via energon API: {e}") + return None + + data_keys = [k for k in sample.keys() if not k.startswith("_")] + + if encoding == "preencoded": + missing = EXPECTED_KEYS_PREENCODED - set(data_keys) + if missing: + logger.error(f"Sample missing expected keys for preencoded data: {missing}") + return None + elif encoding == "preencoded_numpy": + missing = EXPECTED_KEYS_PREENCODED_NUMPY - set(data_keys) + if missing: + logger.error(f"Sample missing expected keys for preencoded_numpy data: {missing}") + return None + elif encoding == "raw": + has_image = any(k in EXPECTED_KEYS_RAW_IMAGE for k in data_keys) + has_text = "txt" in data_keys + if not has_image or not has_text: + logger.error(f"Raw sample missing image or text key. Got: {data_keys}") + return None + + descriptions = {} + for key in sorted(data_keys): + val = sample[key] + if isinstance(val, torch.Tensor): + descriptions[key] = f"shape={tuple(val.shape)}, dtype={val.dtype}" + elif isinstance(val, bytes): + descriptions[key] = f"bytes, len={len(val)}" + else: + descriptions[key] = f"{type(val).__name__}" + + return descriptions + + +def validate_energon_dataset( + output_dir: str, + encoding: Literal["preencoded", "preencoded_numpy", "raw"] = "preencoded", +) -> bool: + """ + Validate a prepared Energon WebDataset. + + Performs four checks: + 1. Metadata files (.info.json, split.yaml, dataset.yaml) are present and valid + 2. Sample count in .info.json matches actual tar contents (spot-check) + 3. One sample loads successfully through energon's Python API + 4. Prints a structured summary of the dataset + + Args: + output_dir: Path to the dataset directory + encoding: Expected encoding ('preencoded', 'preencoded_numpy', or 'raw') + + Returns: + True if all checks pass, False otherwise. + """ + output_path = Path(output_dir) + passed = True + + logger.info("Validating dataset...") + + # --- Check 1: metadata --- + info = _check_metadata(output_path) + if info is None: + logger.error(" Metadata check FAILED") + return False + logger.info(" Metadata check passed") + + # --- Check 2: sample counts --- + if not _check_sample_counts(output_path, info): + logger.error(" Sample count check FAILED") + passed = False + else: + logger.info(" Sample count check passed") + + # --- Check 3: load one sample via energon API --- + descriptions = _check_sample_load(output_path, encoding) + if descriptions is None: + logger.error(" Sample load check FAILED") + passed = False + else: + logger.info(" Sample load check passed") + + # --- Check 4: summary --- + shard_counts = info["shard_counts"] + total_samples = sum(shard_counts.values()) + num_shards = len(shard_counts) + + splits = info.get("_splits", {}).get("split_parts", {}) + split_summary = ", ".join( + ( + f"{name}={len(shards)}" + if isinstance(shards, list) and not any("{" in s for s in shards) + else f"{name}={'non-empty' if shards else 'empty'}" + ) + for name, shards in splits.items() + ) + + total_bytes = sum(f.stat().st_size for f in output_path.glob("**/*.tar")) + size_gb = total_bytes / (1024**3) + + logger.info(" " + "-" * 40) + logger.info(f" Encoding: {encoding}") + logger.info(f" Total samples: {total_samples} across {num_shards} shard(s)") + logger.info(f" Splits: {split_summary}") + if descriptions: + logger.info(" Spot check:") + for key, desc in descriptions.items(): + logger.info(f" {key}: {desc}") + logger.info(f" Dataset size: {size_gb:.2f} GB") + + if passed: + logger.info("✓ Dataset verified and ready for training") + else: + logger.warning("Dataset has issues — see errors above") + + return passed + + +if __name__ == "__main__": + import argparse + import sys + + logging.basicConfig(level=logging.INFO, format="%(message)s") + + parser = argparse.ArgumentParser( + description="Validate a Primus Energon WebDataset", + ) + parser.add_argument("dataset_path", help="Path to the dataset directory") + parser.add_argument( + "--encoding", + choices=["preencoded", "preencoded_numpy", "raw"], + default="preencoded", + help="Expected encoding type (default: preencoded)", + ) + args = parser.parse_args() + + ok = validate_energon_dataset(args.dataset_path, encoding=args.encoding) + sys.exit(0 if ok else 1) diff --git a/primus/cli/main.py b/primus/cli/main.py index b4b0b49a7..de9638be9 100644 --- a/primus/cli/main.py +++ b/primus/cli/main.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -34,6 +34,7 @@ def _prefer_checkout_primus_on_sys_path() -> None: import argparse import importlib +import logging import pkgutil import traceback from typing import Callable, Dict, Iterable, Optional, Set @@ -148,6 +149,8 @@ def main(): ... """ _ensure_project_root_on_path() + logging.basicConfig(level=logging.WARNING, format="%(message)s") + logging.getLogger("primus").setLevel(logging.INFO) parser = argparse.ArgumentParser( prog="primus", description="Primus Unified CLI for Training & Utilities", diff --git a/primus/cli/subcommands/data.py b/primus/cli/subcommands/data.py new file mode 100644 index 000000000..63f6deee7 --- /dev/null +++ b/primus/cli/subcommands/data.py @@ -0,0 +1,950 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Dataset preparation subcommands for Megatron backend diffusion training. + +Provides tools to prepare datasets in Energon WebDataset format specifically +for diffusion models using the Megatron backend. + +Commands: + primus data diffusion-raw - Prepare raw Energon WebDataset + primus data diffusion-encoded - Prepare pre-encoded Energon WebDataset + primus data diffusion-ingest - Stream pre-encoded Arrow data into WebDataset + +Supports torch.distributed for multi-GPU processing: + torchrun --nproc_per_node=8 primus data diffusion-raw ... +""" + +import argparse +import logging + +from primus.backends.megatron.data.diffusion.preprocessing.auth import ( + HFAuthError, + setup_hf_authentication, +) +from primus.backends.megatron.data.diffusion.preprocessing.pipelines.encoded import ( + EncodedDatasetPipeline, +) +from primus.backends.megatron.data.diffusion.preprocessing.pipelines.raw import ( + RawDatasetPipeline, +) +from primus.tools.utils import finalize_distributed, init_distributed + +logger = logging.getLogger(__name__) + + +def _add_common_args(parser): + """Add arguments common to all data preparation commands.""" + # Source configuration + source_group = parser.add_argument_group("Source Configuration") + source_group.add_argument( + "--source-type", + required=False, + choices=["directory", "huggingface", "webdataset"], + help="Type of input data source (required, can be set via --config)", + ) + source_group.add_argument("--input-dir", type=str, help="Input directory (for directory source)") + source_group.add_argument("--hf-dataset", type=str, help="HuggingFace dataset name (for HF source)") + source_group.add_argument( + "--hf-split", type=str, default="train", help="HuggingFace dataset split (default: train)" + ) + source_group.add_argument( + "--hf-data-files", + type=str, + default=None, + help='Specific files/paths within HF dataset (e.g., "data_1024_10K/*.tar")', + ) + source_group.add_argument("--input-path", type=str, help="WebDataset path/glob (for webdataset source)") + + # Output configuration (Energon format) + output_group = parser.add_argument_group("Output Configuration (Energon WebDataset)") + output_group.add_argument( + "--output-dir", + required=False, + help="Output directory for Energon WebDataset (required, can be set via --config)", + ) + output_group.add_argument( + "--shard-size", type=int, default=1000, help="Samples per shard (default: 1000)" + ) + output_group.add_argument( + "--max-samples", type=int, default=None, help="Maximum samples to process (default: all)" + ) + output_group.add_argument("--compress", action="store_true", help="Compress tar files with gzip") + + # Image preprocessing + image_group = parser.add_argument_group("Image Preprocessing") + image_group.add_argument( + "--variable-size", + action="store_true", + help="Resize images to the nearest multiple of 16 instead of fixed size", + ) + image_group.add_argument( + "--image-size", + type=int, + default=1024, + help="Resize images to this size (default: 1024) when variable_size is disabled", + ) + image_group.add_argument( + "--center-crop", action="store_false", help="Disable center cropping of images before resize" + ) + image_group.add_argument( + "--max-size", + type=int, + default=1024, + help="Maximum image dimension for variable size mode (default: 1024)", + ) + + # HuggingFace Authentication + auth_group = parser.add_argument_group("HuggingFace Authentication") + auth_group.add_argument( + "--hf-token-file", + type=str, + default=None, + help=( + "Path to file containing HuggingFace token. " + "File must have secure permissions (600 or 400). " + "If not provided, falls back to HF_TOKEN env variable " + "or HF CLI login (~/.cache/huggingface/token)." + ), + ) + + return parser + + +def _validate_source_args(args): + """Validate that required source arguments are provided.""" + if args.source_type == "directory" and not args.input_dir: + raise ValueError("--input-dir required for directory source") + if args.source_type == "huggingface" and not args.hf_dataset: + raise ValueError("--hf-dataset required for huggingface source") + if args.source_type == "webdataset" and not args.input_path: + raise ValueError("--input-path required for webdataset source") + + +def _prepare_raw(args): + """Prepare raw Energon WebDataset for on-the-fly encoding during training.""" + # Initialize torch.distributed if launched with torchrun + init_distributed() + + try: + _validate_source_args(args) + + # Setup HF authentication + try: + setup_hf_authentication(token_file=getattr(args, "hf_token_file", None)) + except HFAuthError as e: + logger.error(f"Authentication failed: {e}") + raise + + logger.info("=" * 80) + logger.info("Megatron Diffusion: Raw Energon WebDataset Preparation") + logger.info("=" * 80) + + pipeline = RawDatasetPipeline( + source_type=args.source_type, + output_dir=args.output_dir, + variable_size=args.variable_size, + image_size=args.image_size, + center_crop=args.center_crop, + max_size=args.max_size, + image_format=getattr(args, "image_format", "JPEG"), + image_quality=getattr(args, "image_quality", 95), + shard_size=args.shard_size, + max_samples=args.max_samples, + compress=args.compress, + ) + + source_kwargs = {} + if args.source_type == "directory": + source_kwargs["input_dir"] = args.input_dir + elif args.source_type == "huggingface": + source_kwargs["hf_dataset"] = args.hf_dataset + source_kwargs["hf_split"] = args.hf_split + source_kwargs["hf_data_files"] = getattr(args, "hf_data_files", None) + # Add data format configuration + source_kwargs["image_key"] = getattr(args, "image_key", None) + source_kwargs["caption_key"] = getattr(args, "caption_key", None) + source_kwargs["image_keys"] = getattr(args, "image_keys", None) + source_kwargs["caption_keys"] = getattr(args, "caption_keys", None) + elif args.source_type == "webdataset": + source_kwargs["input_path"] = args.input_path + + results = pipeline.run(**source_kwargs) + + logger.info("=" * 80) + logger.info("✓ Raw Energon WebDataset preparation complete!") + logger.info(f" Samples processed: {results['samples_processed']}") + logger.info(f" Samples skipped: {results['samples_skipped']}") + logger.info(f" Shards created: {results['shards_written']}") + logger.info(f" Output: {args.output_dir}") + logger.info(f" Format: Raw Energon WebDataset (images + captions)") + logger.info(f" Note: Encoding will be done on-the-fly during training") + logger.info("=" * 80) + + # Finalize dataset unless --no-finalize was passed + if not getattr(args, "no_finalize", False): + # Import distributed utilities + import torch.distributed as dist + + # Wait for all ranks to complete data preparation + if dist.is_initialized(): + logger.info("Waiting for all ranks to complete...") + dist.barrier() + + # Only rank 0 runs finalization + should_finalize = not dist.is_initialized() or dist.get_rank() == 0 + + if should_finalize: + from primus.backends.megatron.data.diffusion.preprocessing.finalize import ( + finalize_energon_dataset, + ) + + try: + finalize_energon_dataset( + output_dir=args.output_dir, + train_split=getattr(args, "train_split", 1.0), + encoding="raw", + num_workers=8, + ) + except Exception as e: + logger.error(f"Finalization failed: {e}") + raise + + # Wait again so all ranks exit together + if dist.is_initialized(): + dist.barrier() + + finally: + # Clean up distributed + finalize_distributed() + + +def _flatten_preprocessing_config(config_dict: dict) -> dict: + """ + Flatten nested YAML config to match CLI argument structure. + + Transforms hierarchical YAML structure into flat dict matching argparse namespace. + + Args: + config_dict: Nested config from YAML file + + Returns: + Flattened dict with CLI argument names as keys + + Example: + >>> config = { + ... 'source': {'type': 'huggingface', 'hf_dataset': 'pokemon'}, + ... 'model': {'batch_size': 8} + ... } + >>> flat = _flatten_preprocessing_config(config) + >>> flat['source_type'], flat['batch_size'] + ('huggingface', 8) + """ + flat = {} + + # Source configuration + if "source" in config_dict: + source = config_dict["source"] + flat["source_type"] = source.get("type") + flat["hf_dataset"] = source.get("hf_dataset") + flat["hf_split"] = source.get("hf_split", "train") + flat["hf_data_files"] = source.get("hf_data_files") + flat["input_dir"] = source.get("input_dir") + flat["input_path"] = source.get("input_path") + + # Data format configuration (field mappings for image/caption extraction) + if "data_format" in config_dict: + data_format = config_dict["data_format"] + flat["image_key"] = data_format.get("image_key") + flat["caption_key"] = data_format.get("caption_key") + flat["image_keys"] = data_format.get("image_keys") + flat["caption_keys"] = data_format.get("caption_keys") + + # Output configuration + if "output" in config_dict: + output = config_dict["output"] + flat["output_dir"] = output.get("output_dir") + flat["shard_size"] = output.get("shard_size", 1000) + flat["max_samples"] = output.get("max_samples") + flat["compress"] = output.get("compress", False) + + # Model configuration + if "model" in config_dict: + model = config_dict["model"] + flat["model_path"] = model.get("model_path", "black-forest-labs/FLUX.1-dev") + flat["vae_path"] = model.get("vae_path") + flat["t5_path"] = model.get("t5_path") + flat["clip_path"] = model.get("clip_path") + flat["precision"] = model.get("precision", "bf16") + flat["device"] = model.get("device", "cuda") + flat["batch_size"] = model.get("batch_size", 8) + flat["t5_max_length"] = model.get("t5_max_length", 512) + flat["vae_latent_mode"] = model.get("vae_latent_mode", "presampled") + + # Image preprocessing + if "image" in config_dict: + image = config_dict["image"] + flat["image_size"] = image.get("image_size", 1024) + flat["variable_size"] = image.get("variable_size", False) + flat["center_crop"] = image.get("center_crop", True) + flat["max_size"] = image.get("max_size", 1024) + + # Authentication + if "auth" in config_dict: + auth = config_dict["auth"] + flat["hf_token_file"] = auth.get("hf_token_file") + + return flat + + +def _get_encoded_parser_defaults() -> dict: + """ + Get default values from encoded parser for override detection. + + Returns dict of argument name -> default value to detect which CLI + arguments were explicitly set vs using defaults. + """ + return { + "config": None, + "source_type": None, + "hf_dataset": None, + "hf_split": "train", + "hf_data_files": None, + "input_dir": None, + "input_path": None, + "image_key": None, + "caption_key": None, + "image_keys": None, + "caption_keys": None, + "output_dir": None, + "shard_size": 1000, + "max_samples": None, + "compress": False, + "model_path": "black-forest-labs/FLUX.1-dev", + "vae_path": None, + "t5_path": None, + "clip_path": None, + "precision": "bf16", + "device": "cuda", + "batch_size": 8, + "t5_max_length": 512, + "image_size": 1024, + "variable_size": False, + "center_crop": True, + "max_size": 1024, + "hf_token_file": None, + "vae_latent_mode": "presampled", + } + + +def _load_config_with_cli_overrides(args: "argparse.Namespace") -> "argparse.Namespace": + """ + Load YAML config and merge with CLI arguments. + + Priority order (highest to lowest): + 1. Explicitly provided CLI arguments + 2. YAML config values + 3. CLI default values + + Args: + args: Parsed CLI arguments + + Returns: + Merged namespace with final configuration + + Example: + >>> # With config file specifying batch_size: 8 + >>> # And CLI arg --batch-size 16 + >>> # Result: batch_size = 16 (CLI wins) + """ + import argparse + + # If no config file, return args as-is + if not getattr(args, "config", None): + return args + + from primus.core.utils import yaml_utils + + logger.info(f"Loading config from: {args.config}") + + # Load and flatten YAML config + config_dict = yaml_utils.parse_yaml(args.config) + flat_config = _flatten_preprocessing_config(config_dict) + + # Start with config values + merged_dict = flat_config.copy() + + # Override with explicitly provided CLI arguments + parser_defaults = _get_encoded_parser_defaults() + cli_args = vars(args) + + for key, cli_value in cli_args.items(): + if key == "config": + # Keep config path for reference + merged_dict["config"] = cli_value + continue + + # If CLI value differs from default, it was explicitly set + default_value = parser_defaults.get(key) + if cli_value != default_value: + merged_dict[key] = cli_value + logger.debug(f"CLI override: {key} = {cli_value}") + + logger.info("Configuration merged (CLI args override YAML)") + return argparse.Namespace(**merged_dict) + + +def _validate_preprocessing_config(args: "argparse.Namespace") -> None: + """ + Validate preprocessing configuration after merging. + + Ensures all required parameters are present regardless of source + (CLI, YAML, or both). + + Args: + args: Merged configuration namespace + + Raises: + ValueError: If required arguments are missing or invalid + + Example: + >>> args = argparse.Namespace(source_type='huggingface', output_dir=None) + >>> _validate_preprocessing_config(args) # Raises ValueError + """ + # Check required arguments + required = ["source_type", "output_dir"] + missing = [arg for arg in required if not getattr(args, arg, None)] + + if missing: + raise ValueError( + f"Missing required arguments: {', '.join(missing)}. " + f"Provide via --config YAML or CLI arguments." + ) + + # Validate source-specific requirements + if args.source_type == "huggingface": + if not getattr(args, "hf_dataset", None): + raise ValueError( + "Missing --hf-dataset for source-type 'huggingface'. " + "Specify in config file (source.hf_dataset) or via CLI." + ) + elif args.source_type == "directory": + if not getattr(args, "input_dir", None): + raise ValueError( + "Missing --input-dir for source-type 'directory'. " + "Specify in config file (source.input_dir) or via CLI." + ) + elif args.source_type == "webdataset": + if not getattr(args, "input_path", None): + raise ValueError( + "Missing --input-path for source-type 'webdataset'. " + "Specify in config file (source.input_path) or via CLI." + ) + else: + raise ValueError( + f"Invalid source-type: {args.source_type}. " f"Must be one of: huggingface, directory, webdataset" + ) + + +def _prepare_encoded(args): + """Prepare pre-encoded Energon WebDataset with VAE/T5/CLIP for fast training.""" + # Load and merge config if provided + args = _load_config_with_cli_overrides(args) + + # Validate merged configuration + _validate_preprocessing_config(args) + + # Initialize torch.distributed if launched with torchrun + init_distributed() + + try: + _validate_source_args(args) + + # Setup HF authentication + try: + setup_hf_authentication(token_file=getattr(args, "hf_token_file", None)) + except HFAuthError as e: + logger.error(f"Authentication failed: {e}") + raise + + logger.info("=" * 80) + logger.info("Megatron Diffusion: Pre-encoded Energon WebDataset Preparation") + logger.info("=" * 80) + + pipeline = EncodedDatasetPipeline( + source_type=args.source_type, + output_dir=args.output_dir, + model_path=args.model_path, + vae_path=args.vae_path, + t5_path=args.t5_path, + clip_path=args.clip_path, + precision=args.precision, + device=args.device, + batch_size=args.batch_size, + t5_max_length=args.t5_max_length, + variable_size=args.variable_size, + image_size=args.image_size, + center_crop=args.center_crop, + max_size=args.max_size, + shard_size=args.shard_size, + max_samples=args.max_samples, + compress=args.compress, + vae_latent_mode=getattr(args, "vae_latent_mode", "presampled"), + ) + + source_kwargs = {} + if args.source_type == "directory": + source_kwargs["input_dir"] = args.input_dir + elif args.source_type == "huggingface": + source_kwargs["hf_dataset"] = args.hf_dataset + source_kwargs["hf_split"] = args.hf_split + source_kwargs["hf_data_files"] = getattr(args, "hf_data_files", None) + # Add data format configuration + source_kwargs["image_key"] = getattr(args, "image_key", None) + source_kwargs["caption_key"] = getattr(args, "caption_key", None) + source_kwargs["image_keys"] = getattr(args, "image_keys", None) + source_kwargs["caption_keys"] = getattr(args, "caption_keys", None) + + elif args.source_type == "webdataset": + source_kwargs["input_path"] = args.input_path + + results = pipeline.run(**source_kwargs) + + logger.info("=" * 80) + logger.info("✓ Pre-encoded Energon WebDataset preparation complete!") + logger.info(f" Samples processed: {results['samples_processed']}") + logger.info(f" Samples skipped: {results['samples_skipped']}") + logger.info(f" Shards created: {results['shards_written']}") + logger.info(f" Output: {args.output_dir}") + logger.info(f" Format: Pre-encoded Energon WebDataset (VAE/T5/CLIP tensors)") + logger.info(f" Note: Training will be faster (no on-the-fly encoding)") + logger.info("=" * 80) + + # Finalize dataset unless --no-finalize was passed + if not getattr(args, "no_finalize", False): + # Import distributed utilities + import torch.distributed as dist + + # Wait for all ranks to complete data preparation + if dist.is_initialized(): + logger.info("Waiting for all ranks to complete...") + dist.barrier() + + # Only rank 0 runs finalization + should_finalize = not dist.is_initialized() or dist.get_rank() == 0 + + if should_finalize: + from primus.backends.megatron.data.diffusion.preprocessing.finalize import ( + finalize_energon_dataset, + ) + + try: + finalize_energon_dataset( + output_dir=args.output_dir, + train_split=getattr(args, "train_split", 1.0), + encoding="preencoded", + num_workers=8, + ) + except Exception as e: + logger.error(f"Finalization failed: {e}") + raise + + # Wait again so all ranks exit together + if dist.is_initialized(): + dist.barrier() + + finally: + # Clean up distributed + finalize_distributed() + + +def _load_ingest_config(args): + """Load ingest YAML config and merge with CLI overrides. + + Returns a flat dict consumed by _prepare_ingest. + """ + from primus.core.utils import yaml_utils + + config_path = args.config + if not config_path: + raise ValueError("--config is required for diffusion-ingest") + + raw = yaml_utils.parse_yaml(config_path) + + config = { + "datasets": raw.get("datasets", []), + "output_dir": raw.get("output", {}).get("output_dir"), + "max_workers": raw.get("pipeline", {}).get("max_workers", 4), + "prefetch_depth": raw.get("pipeline", {}).get("prefetch_depth", 6), + "max_files": raw.get("pipeline", {}).get("max_files"), + "no_finalize": False, + } + + if raw.get("empty_encodings"): + config["empty_encodings"] = raw["empty_encodings"] + + # CLI overrides + if getattr(args, "output_dir", None): + config["output_dir"] = args.output_dir + if getattr(args, "input_dir", None): + config["input_dir"] = args.input_dir + if getattr(args, "max_workers", None) is not None: + config["max_workers"] = args.max_workers + if getattr(args, "prefetch_depth", None) is not None: + config["prefetch_depth"] = args.prefetch_depth + if getattr(args, "max_files", None) is not None: + config["max_files"] = args.max_files + if getattr(args, "no_finalize", False): + config["no_finalize"] = True + + if not config.get("output_dir"): + raise ValueError("output_dir is required. Set in config (output.output_dir) " "or via --output-dir.") + + # Default input_dir + if "input_dir" not in config: + from pathlib import Path + + config["input_dir"] = str(Path(config["output_dir"]) / "_arrow_tmp") + + return config + + +def _download_empty_encodings(config): + """Download empty_encodings files (small .npy files for CFG dropout).""" + from pathlib import Path + + from primus.backends.megatron.data.diffusion.preprocessing.download import ( + download_with_backoff, + fetch_manifest, + ) + + enc_cfg = config["empty_encodings"] + output_dir = Path(config["output_dir"]) / enc_cfg.get("output_subdir", "empty_encodings") + output_dir.mkdir(parents=True, exist_ok=True) + + base_url, entries = fetch_manifest(enc_cfg["manifest_url"]) + + for md5, fname in entries: + dest = output_dir / fname + if dest.exists(): + logger.info(f" [skip] {fname} (exists)") + continue + file_url = f"{base_url}/{fname}" + logger.info(f" [download] {fname}") + download_with_backoff(file_url, dest, expected_md5=md5) + + logger.info(f"Empty encodings downloaded to {output_dir}") + + +def _prepare_ingest(args): + """Download and convert pre-encoded Arrow data into Energon WebDataset.""" + from primus.backends.megatron.data.diffusion.preprocessing.pipelines.ingest import ( + StreamingIngestPipeline, + ) + + config = _load_ingest_config(args) + + logger.info("=" * 80) + logger.info("Megatron Diffusion: Streaming Ingest Pipeline") + logger.info("=" * 80) + + total_failed = 0 + for ds in config["datasets"]: + logger.info(f"Processing dataset: {ds['name']} (split: {ds['split_name']})") + pipeline = StreamingIngestPipeline( + manifest_url=ds["manifest_url"], + input_dir=config["input_dir"], + output_dir=config["output_dir"], + split_name=ds["split_name"], + max_files=config.get("max_files"), + max_workers=config["max_workers"], + prefetch_depth=config["prefetch_depth"], + ) + result = pipeline.run() + total_failed += result.get("files_failed", 0) + + if "empty_encodings" in config: + logger.info("Downloading empty encodings...") + _download_empty_encodings(config) + + if not config.get("no_finalize"): + from primus.backends.megatron.data.diffusion.preprocessing.finalize import ( + finalize_energon_dataset, + ) + + try: + finalize_energon_dataset( + output_dir=config["output_dir"], + encoding="preencoded_numpy", + split_parts_patterns=[("train", "train/.*"), ("val", "val/.*")], + ) + except Exception as e: + logger.error(f"Finalization failed: {e}") + raise + + logger.info("=" * 80) + logger.info("Ingest complete!") + logger.info(f" Output: {config['output_dir']}") + if total_failed: + logger.warning( + f" {total_failed} file(s) failed across all datasets. " + f"Check failed_files.json in each split directory." + ) + logger.info("=" * 80) + + +def register_subcommand(subparsers: argparse._SubParsersAction) -> argparse.ArgumentParser: + """ + Register 'primus data' subcommand for Megatron diffusion dataset preparation. + + Provides dataset preparation utilities specifically for diffusion models + using the Megatron backend, outputting Energon WebDataset format. + """ + parser = subparsers.add_parser( + "data", + help="Dataset preparation tools (Megatron diffusion, Energon format)", + description=( + "Data preparation utilities for Megatron-backend diffusion models.\n" + "Creates datasets in Energon WebDataset format.\n\n" + "Supported modes:\n" + " - diffusion-raw: Raw WebDataset (smaller, on-the-fly encoding)\n" + " - diffusion-encoded: Pre-encoded WebDataset (larger, faster training)\n" + " - diffusion-ingest: Stream pre-encoded Arrow data into WebDataset\n\n" + "Multi-GPU support via torch.distributed:\n" + " torchrun --nproc_per_node=8 primus data diffusion-raw ...\n" + " torchrun --nproc_per_node=8 primus data diffusion-encoded ..." + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + data_subparsers = parser.add_subparsers(dest="data_command", required=True, help="Data preparation mode") + + # ===== primus data diffusion-raw ===== + raw_parser = data_subparsers.add_parser( + "diffusion-raw", + help="Prepare raw Energon WebDataset (Megatron diffusion: images + captions)", + description=( + "Prepare raw Energon WebDataset for Megatron diffusion training.\n\n" + "Creates smaller datasets with raw images and captions.\n" + "Encoding (VAE/T5/CLIP) happens on-the-fly during training.\n\n" + "Output format: Energon WebDataset (tar shards)\n\n" + "Use this when:\n" + " - Storage is limited\n" + " - Experimenting with different encoders\n" + " - Rapid prototyping\n\n" + "Example:\n" + " primus data diffusion-raw \\\n" + " --source-type huggingface \\\n" + " --hf-dataset diffusers/pokemon-gpt4-captions \\\n" + " --output-dir /data/raw_pokemon \\\n" + " --image-size 1024 --center-crop\n\n" + "Multi-GPU:\n" + " torchrun --nproc_per_node=8 primus data diffusion-raw ..." + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + _add_common_args(raw_parser) + + # Raw-specific arguments + raw_specific = raw_parser.add_argument_group("Raw Dataset Options") + raw_specific.add_argument( + "--image-format", + choices=["JPEG", "PNG", "WEBP"], + default="JPEG", + help="Output image format (default: JPEG)", + ) + raw_specific.add_argument( + "--image-quality", type=int, default=95, help="JPEG/WEBP quality 1-100 (default: 95)" + ) + + # Finalization arguments + raw_finalize = raw_parser.add_argument_group("Dataset Finalization") + raw_finalize.add_argument( + "--no-finalize", + action="store_true", + dest="no_finalize", + help="Skip automatic dataset.yaml creation and energon prepare (default: finalize is ON)", + ) + raw_finalize.add_argument( + "--train-split", + type=float, + default=1.0, + help="Training split ratio for finalization (default: 1.0 = 100%% train, 0%% val/test)", + ) + + raw_parser.set_defaults(func=lambda args, unknown: _prepare_raw(args)) + + # ===== primus data diffusion-encoded ===== + encoded_parser = data_subparsers.add_parser( + "diffusion-encoded", + help="Prepare pre-encoded Energon WebDataset (Megatron diffusion: VAE/T5/CLIP)", + description=( + "Prepare pre-encoded Energon WebDataset for Megatron diffusion training.\n\n" + "Pre-encodes images with VAE and text with T5/CLIP, creating larger\n" + "datasets that enable faster training (no encoding overhead).\n\n" + "Output format: Energon WebDataset with pre-encoded tensors\n\n" + "Use this when:\n" + " - Training on production datasets\n" + " - Maximum training speed is needed\n" + " - Storage space is available\n\n" + "Example:\n" + " primus data diffusion-encoded \\\n" + " --source-type huggingface \\\n" + " --hf-dataset diffusers/pokemon-gpt4-captions \\\n" + " --output-dir /data/encoded_pokemon \\\n" + " --model-path black-forest-labs/FLUX.1-dev \\\n" + " --batch-size 8 --precision bf16\n\n" + "Config file:\n" + " primus data diffusion-encoded --config preprocessing.yaml\n\n" + "Multi-GPU:\n" + " torchrun --nproc_per_node=8 primus data diffusion-encoded ..." + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + # Config file support (optional, CLI args override config values) + encoded_parser.add_argument( + "--config", + type=str, + default=None, + help="Path to YAML config file (optional, CLI args override config values)", + ) + + _add_common_args(encoded_parser) + + # Encoded-specific arguments + model_group = encoded_parser.add_argument_group("Model Configuration") + model_group.add_argument( + "--model-path", + default="black-forest-labs/FLUX.1-dev", + help="Pretrained model path (HF or local). Default: black-forest-labs/FLUX.1-dev " + "(requires HF authentication -- see --hf-token-file)", + ) + model_group.add_argument("--vae-path", default=None, help="Custom VAE path (overrides --model-path)") + model_group.add_argument("--t5-path", default=None, help="Custom T5 path (overrides --model-path)") + model_group.add_argument("--clip-path", default=None, help="Custom CLIP path (overrides --model-path)") + model_group.add_argument( + "--precision", + choices=["bf16", "fp16", "fp32"], + default="bf16", + help="Model precision (default: bf16)", + ) + model_group.add_argument("--device", default="cuda", help="Device for encoding (default: cuda)") + model_group.add_argument("--batch-size", type=int, default=8, help="Encoding batch size (default: 8)") + model_group.add_argument( + "--t5-max-length", + type=int, + default=512, + help="T5 max sequence length (default: 512, use 256 for FLUX.1-schnell)", + ) + model_group.add_argument( + "--vae-latent-mode", + choices=["presampled", "resample"], + default="presampled", + help=( + "VAE latent storage mode (default: presampled). " + "'presampled' stores a single sampled latent per image. " + "'resample' additionally stores mean and logvar so the training " + "loop can re-draw latents via reparameterization at every step." + ), + ) + + # Finalization arguments + encoded_finalize = encoded_parser.add_argument_group("Dataset Finalization") + encoded_finalize.add_argument( + "--no-finalize", + action="store_true", + dest="no_finalize", + help="Skip automatic dataset.yaml creation and energon prepare (default: finalize is ON)", + ) + encoded_finalize.add_argument( + "--train-split", + type=float, + default=1.0, + help="Training split ratio for finalization (default: 1.0 = 100%% train, 0%% val/test)", + ) + + encoded_parser.set_defaults(func=lambda args, unknown: _prepare_encoded(args)) + + # ===== primus data diffusion-ingest ===== + ingest_parser = data_subparsers.add_parser( + "diffusion-ingest", + help="Stream pre-encoded Arrow data into Energon WebDataset", + description=( + "Stream pre-encoded Arrow data into Energon WebDataset.\n\n" + "Downloads Apache Arrow IPC files from MLCommons R2 (or similar),\n" + "converts them to WebDataset tar shards in a single pass using a\n" + "producer-consumer prefetch architecture. Minimizes disk usage by\n" + "deleting each Arrow file after conversion.\n\n" + "Requires a YAML config file specifying datasets and parameters.\n\n" + "Example:\n" + " primus data diffusion-ingest \\\n" + " --config primus/configs/data/megatron/diffusion/" + "preprocessing/mlperf_flux1.yaml\n\n" + "Override output:\n" + " primus data diffusion-ingest \\\n" + " --config .../mlperf_flux1.yaml \\\n" + " --output-dir /custom/path --max-files 5" + ), + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + ingest_parser.add_argument( + "--config", + type=str, + required=True, + help="Path to YAML config file with datasets list and pipeline settings", + ) + + ingest_output = ingest_parser.add_argument_group("Output Configuration") + ingest_output.add_argument( + "--output-dir", + type=str, + default=None, + help="Output directory (overrides config output.output_dir)", + ) + ingest_output.add_argument( + "--input-dir", + type=str, + default=None, + help="Temp directory for Arrow downloads (default: /_arrow_tmp)", + ) + + ingest_pipeline = ingest_parser.add_argument_group("Pipeline Configuration") + ingest_pipeline.add_argument( + "--max-files", + type=int, + default=None, + help="Limit Arrow files per dataset (for testing)", + ) + ingest_pipeline.add_argument( + "--max-workers", + type=int, + default=None, + help="Concurrent download threads (overrides config, default: 4)", + ) + ingest_pipeline.add_argument( + "--prefetch-depth", + type=int, + default=None, + help="Max Arrow files buffered on disk (overrides config, default: 6)", + ) + + ingest_finalize = ingest_parser.add_argument_group("Dataset Finalization") + ingest_finalize.add_argument( + "--no-finalize", + action="store_true", + dest="no_finalize", + help="Skip automatic Energon finalization (default: finalize is ON)", + ) + + ingest_parser.set_defaults(func=lambda args, unknown: _prepare_ingest(args)) + + # Set default for parent parser (required by CLI framework, but never called due to required=True on subparsers) + parser.set_defaults(func=lambda args, unknown: None) + + return parser diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/__init__.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/__init__.py new file mode 100644 index 000000000..89778402a --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_download.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_download.py new file mode 100644 index 000000000..f98126b54 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_download.py @@ -0,0 +1,251 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for download utilities (download.py). + +Covers: +- parse_md5_manifest: parsing, filtering, sorting +- download_with_backoff: retry logic, MD5 verification, timeout handling +- fetch_url_text: basic HTTP text fetch +""" + +import hashlib +import io +import tempfile +import urllib.error +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from primus.backends.megatron.data.diffusion.preprocessing.download import ( + download_with_backoff, + fetch_url_text, + parse_md5_manifest, +) + +_DOWNLOAD_MODULE = "primus.backends.megatron.data.diffusion.preprocessing.download" + + +class TestParseMd5Manifest: + def test_basic_parsing(self): + text = "abc123 data-00000.arrow\ndef456 data-00001.arrow\n" + entries = parse_md5_manifest(text) + assert entries == [("abc123", "data-00000.arrow"), ("def456", "data-00001.arrow")] + + def test_suffix_filter(self): + text = "abc123 data-00000.arrow\ndef456 readme.txt\n" + entries = parse_md5_manifest(text, suffix_filter=".arrow") + assert len(entries) == 1 + assert entries[0][1] == "data-00000.arrow" + + def test_no_filter_returns_all(self): + text = "abc123 data-00000.arrow\ndef456 readme.txt\n" + entries = parse_md5_manifest(text) + assert len(entries) == 2 + + def test_sorts_by_filename(self): + text = "bbb data-00002.arrow\naaa data-00001.arrow\nccc data-00000.arrow\n" + entries = parse_md5_manifest(text) + assert [e[1] for e in entries] == [ + "data-00000.arrow", + "data-00001.arrow", + "data-00002.arrow", + ] + + def test_skips_malformed_lines(self): + text = "abc123 data-00000.arrow\nbadline\n\n \n" + entries = parse_md5_manifest(text) + assert len(entries) == 1 + + +class TestDownloadWithBackoff: + """Tests for download_with_backoff with mocked HTTP.""" + + @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen") + def test_success_on_first_try(self, mock_urlopen): + content = b"hello world" + mock_resp = MagicMock() + mock_resp.read.side_effect = [content, b""] + mock_resp.__enter__ = MagicMock(return_value=mock_resp) + mock_resp.__exit__ = MagicMock(return_value=False) + mock_urlopen.return_value = mock_resp + + with tempfile.TemporaryDirectory() as tmp: + dest = Path(tmp) / "test.bin" + download_with_backoff("http://example.com/test.bin", dest, max_retries=0) + assert dest.exists() + + @patch(f"{_DOWNLOAD_MODULE}.time.sleep") + @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen") + def test_retries_on_429(self, mock_urlopen, mock_sleep): + content = b"success data" + mock_resp_ok = MagicMock() + mock_resp_ok.read.side_effect = [content, b""] + mock_resp_ok.__enter__ = MagicMock(return_value=mock_resp_ok) + mock_resp_ok.__exit__ = MagicMock(return_value=False) + + mock_urlopen.side_effect = [ + urllib.error.HTTPError("http://x", 429, "Too Many Requests", {}, io.BytesIO(b"")), + urllib.error.HTTPError("http://x", 429, "Too Many Requests", {}, io.BytesIO(b"")), + mock_resp_ok, + ] + + with tempfile.TemporaryDirectory() as tmp: + dest = Path(tmp) / "test.bin" + download_with_backoff("http://example.com/test.bin", dest, max_retries=3, base_delay=0.01) + assert dest.exists() + assert mock_urlopen.call_count == 3 + assert mock_sleep.call_count == 2 + + @patch(f"{_DOWNLOAD_MODULE}.time.sleep") + @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen") + def test_retries_on_503(self, mock_urlopen, mock_sleep): + mock_urlopen.side_effect = urllib.error.HTTPError( + "http://x", 503, "Service Unavailable", {}, io.BytesIO(b"") + ) + + with tempfile.TemporaryDirectory() as tmp: + dest = Path(tmp) / "test.bin" + with pytest.raises(RuntimeError, match="HTTP 503"): + download_with_backoff("http://example.com/test.bin", dest, max_retries=2, base_delay=0.01) + assert mock_urlopen.call_count == 3 # initial + 2 retries + + @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen") + def test_no_retry_on_404(self, mock_urlopen): + mock_urlopen.side_effect = urllib.error.HTTPError("http://x", 404, "Not Found", {}, io.BytesIO(b"")) + + with tempfile.TemporaryDirectory() as tmp: + dest = Path(tmp) / "test.bin" + with pytest.raises(RuntimeError, match="HTTP 404"): + download_with_backoff("http://example.com/test.bin", dest, max_retries=3, base_delay=0.01) + assert mock_urlopen.call_count == 1 + + @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen") + def test_md5_verification_pass(self, mock_urlopen): + content = b"test content" + expected_md5 = hashlib.md5(content).hexdigest() + + mock_resp = MagicMock() + mock_resp.read.side_effect = [content, b""] + mock_resp.__enter__ = MagicMock(return_value=mock_resp) + mock_resp.__exit__ = MagicMock(return_value=False) + mock_urlopen.return_value = mock_resp + + with tempfile.TemporaryDirectory() as tmp: + dest = Path(tmp) / "test.bin" + download_with_backoff("http://example.com/test.bin", dest, expected_md5=expected_md5) + assert dest.exists() + + @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen") + def test_md5_verification_fail(self, mock_urlopen): + content = b"test content" + + mock_resp = MagicMock() + mock_resp.read.side_effect = [content, b""] + mock_resp.__enter__ = MagicMock(return_value=mock_resp) + mock_resp.__exit__ = MagicMock(return_value=False) + mock_urlopen.return_value = mock_resp + + with tempfile.TemporaryDirectory() as tmp: + dest = Path(tmp) / "test.bin" + with pytest.raises(RuntimeError, match="MD5 mismatch"): + download_with_backoff( + "http://example.com/test.bin", + dest, + expected_md5="bad_md5", + max_retries=0, + ) + assert not dest.exists() + + @patch(f"{_DOWNLOAD_MODULE}.time.sleep") + @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen") + def test_retries_on_network_error(self, mock_urlopen, mock_sleep): + content = b"ok" + mock_resp_ok = MagicMock() + mock_resp_ok.read.side_effect = [content, b""] + mock_resp_ok.__enter__ = MagicMock(return_value=mock_resp_ok) + mock_resp_ok.__exit__ = MagicMock(return_value=False) + + mock_urlopen.side_effect = [ + urllib.error.URLError("Connection refused"), + mock_resp_ok, + ] + + with tempfile.TemporaryDirectory() as tmp: + dest = Path(tmp) / "test.bin" + download_with_backoff("http://example.com/test.bin", dest, max_retries=2, base_delay=0.01) + assert dest.exists() + assert mock_urlopen.call_count == 2 + + @patch(f"{_DOWNLOAD_MODULE}.time.sleep") + @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen") + def test_retries_on_md5_mismatch_then_succeeds(self, mock_urlopen, mock_sleep): + good_content = b"correct data" + bad_content = b"corrupted data" + expected_md5 = hashlib.md5(good_content).hexdigest() + + mock_resp_bad = MagicMock() + mock_resp_bad.read.side_effect = [bad_content, b""] + mock_resp_bad.__enter__ = MagicMock(return_value=mock_resp_bad) + mock_resp_bad.__exit__ = MagicMock(return_value=False) + + mock_resp_good = MagicMock() + mock_resp_good.read.side_effect = [good_content, b""] + mock_resp_good.__enter__ = MagicMock(return_value=mock_resp_good) + mock_resp_good.__exit__ = MagicMock(return_value=False) + + mock_urlopen.side_effect = [mock_resp_bad, mock_resp_good] + + with tempfile.TemporaryDirectory() as tmp: + dest = Path(tmp) / "test.bin" + download_with_backoff( + "http://example.com/test.bin", + dest, + expected_md5=expected_md5, + max_retries=3, + base_delay=0.01, + ) + assert dest.exists() + assert dest.read_bytes() == good_content + assert mock_urlopen.call_count == 2 + assert mock_sleep.call_count == 1 + + @patch(f"{_DOWNLOAD_MODULE}.time.sleep") + @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen") + def test_md5_mismatch_exhausts_retries(self, mock_urlopen, mock_sleep): + bad_content = b"always bad" + expected_md5 = "0000000000000000" + + def make_bad_resp(*args, **kwargs): + resp = MagicMock() + resp.read.side_effect = [bad_content, b""] + resp.__enter__ = MagicMock(return_value=resp) + resp.__exit__ = MagicMock(return_value=False) + return resp + + mock_urlopen.side_effect = make_bad_resp + + with tempfile.TemporaryDirectory() as tmp: + dest = Path(tmp) / "test.bin" + with pytest.raises(RuntimeError, match="MD5 mismatch"): + download_with_backoff( + "http://example.com/test.bin", + dest, + expected_md5=expected_md5, + max_retries=2, + base_delay=0.01, + ) + assert mock_urlopen.call_count == 3 # initial + 2 retries + + +class TestFetchUrlText: + @patch(f"{_DOWNLOAD_MODULE}.urllib.request.urlopen") + def test_returns_decoded_text(self, mock_urlopen): + mock_resp = MagicMock() + mock_resp.read.return_value = b"https://example.com/base" + mock_urlopen.return_value = mock_resp + + result = fetch_url_text("http://example.com/manifest.uri") + assert result == "https://example.com/base" diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_finalize.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_finalize.py new file mode 100644 index 000000000..7ab9a7c05 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_finalize.py @@ -0,0 +1,121 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for dataset finalization and encoder auth error detection. + +Tests finalize_energon_dataset from finalize.py and +_raise_encoder_auth_error from encoded.py. +""" + +import tempfile +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from primus.backends.megatron.data.diffusion.preprocessing.finalize import ( + finalize_energon_dataset, +) +from primus.backends.megatron.data.diffusion.preprocessing.pipelines.encoded import ( + EncodedDatasetPipeline, +) +from tests.utils import PrimusUT + + +def _create_dummy_tar(directory: Path, name: str = "000000.tar") -> None: + """Create an empty .tar file for testing.""" + (directory / name).write_bytes(b"") + + +class TestDatasetYamlContent(PrimusUT): + """Tests that finalize_energon_dataset writes correct dataset.yaml.""" + + @patch( + "primus.backends.megatron.data.diffusion.preprocessing.finalize._prepare_programmatic", + side_effect=ImportError("forced: exercise subprocess fallback path"), + ) + @patch("primus.backends.megatron.data.diffusion.preprocessing.validate.validate_energon_dataset") + @patch("subprocess.Popen") + def test_preencoded_yaml(self, mock_popen, mock_validate, mock_programmatic): + """dataset.yaml contains CrudeWebdataset with encoding: preencoded.""" + mock_process = MagicMock() + mock_process.communicate.return_value = ("Done\n", None) + mock_process.returncode = 0 + mock_popen.return_value = mock_process + + with tempfile.TemporaryDirectory() as tmpdir: + _create_dummy_tar(Path(tmpdir)) + finalize_energon_dataset(output_dir=tmpdir, train_split=1.0, encoding="preencoded") + + dataset_yaml = Path(tmpdir) / ".nv-meta" / "dataset.yaml" + assert dataset_yaml.exists() + content = dataset_yaml.read_text() + assert "__class__: CrudeWebdataset" in content + assert "encoding: preencoded" in content + + # Verify split_input passed to subprocess stdin + mock_process.communicate.assert_called_once_with(input="1.0, 0.0, 0.0\nn\n", timeout=300) + + @patch( + "primus.backends.megatron.data.diffusion.preprocessing.finalize._prepare_programmatic", + side_effect=ImportError("forced: exercise subprocess fallback path"), + ) + @patch("primus.backends.megatron.data.diffusion.preprocessing.validate.validate_energon_dataset") + @patch("subprocess.Popen") + def test_raw_yaml_and_split_format(self, mock_popen, mock_validate, mock_programmatic): + """dataset.yaml contains encoding: raw; split ratios are formatted correctly.""" + mock_process = MagicMock() + mock_process.communicate.return_value = ("Done\n", None) + mock_process.returncode = 0 + mock_popen.return_value = mock_process + + with tempfile.TemporaryDirectory() as tmpdir: + _create_dummy_tar(Path(tmpdir)) + finalize_energon_dataset(output_dir=tmpdir, train_split=0.8, encoding="raw") + + content = (Path(tmpdir) / ".nv-meta" / "dataset.yaml").read_text() + assert "encoding: raw" in content + + call_kwargs = mock_process.communicate.call_args + split_input = call_kwargs.kwargs.get("input") or call_kwargs[1].get("input") + assert split_input.startswith("0.8, ") + assert split_input.endswith("\nn\n") + parts = split_input.split("\n")[0].split(", ") + self.assertAlmostEqual(float(parts[0]), 0.8) + self.assertAlmostEqual(float(parts[1]), 0.1) + self.assertAlmostEqual(float(parts[2]), 0.1) + + +class TestRaiseEncoderAuthError(PrimusUT): + """Tests for EncodedDatasetPipeline._raise_encoder_auth_error.""" + + def test_401_error_gives_auth_hint(self): + """Exception containing '401' raises RuntimeError with auth hint.""" + original = Exception("HTTP 401 Unauthorized") + with self.assertRaises(RuntimeError) as ctx: + EncodedDatasetPipeline._raise_encoder_auth_error("VAE", "my-model", original) + assert "HuggingFace authentication" in str(ctx.exception) + assert "VAE" in str(ctx.exception) + + def test_token_error_gives_auth_hint(self): + """Exception containing 'token' raises RuntimeError with auth hint.""" + original = Exception("Please pass a valid token") + with self.assertRaises(RuntimeError) as ctx: + EncodedDatasetPipeline._raise_encoder_auth_error("T5-XXL", "my-model", original) + assert "HuggingFace authentication" in str(ctx.exception) + assert "T5-XXL" in str(ctx.exception) + + def test_non_auth_error_preserves_message(self): + """Non-auth exception re-raises with original message intact.""" + original = Exception("CUDA out of memory") + with self.assertRaises(RuntimeError) as ctx: + EncodedDatasetPipeline._raise_encoder_auth_error("CLIP-L", "my-model", original) + msg = str(ctx.exception) + assert "CUDA out of memory" in msg + assert "CLIP-L" in msg + assert "HuggingFace authentication" not in msg + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py new file mode 100644 index 000000000..65c1860cb --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_ingest.py @@ -0,0 +1,330 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for the streaming ingest pipeline (pipelines/ingest.py). + +Covers: +- StreamingIngestPipeline: parallel download + sequential conversion +- Resume functionality (existing shards are skipped) +- max_files parameter +- Sample offset accumulation +- Arrow file cleanup after conversion +- Skip-and-log for failed downloads and conversions +""" + +import json +import tempfile +from pathlib import Path +from unittest.mock import patch + +from primus.backends.megatron.data.diffusion.preprocessing.pipelines.ingest import ( + StreamingIngestPipeline, +) + +_INGEST_MODULE = "primus.backends.megatron.data.diffusion.preprocessing.pipelines.ingest" + + +class TestStreamingIngestPipeline: + """Tests for the StreamingIngestPipeline with mocked I/O.""" + + @patch(f"{_INGEST_MODULE}._arrow_to_tar") + @patch(f"{_INGEST_MODULE}.download_with_backoff") + @patch(f"{_INGEST_MODULE}.fetch_manifest") + def test_processes_all_files(self, mock_manifest, mock_download, mock_convert): + """Pipeline processes all entries, creates correct number of shards.""" + entries = [ + ("md5_0", "data-00000.arrow"), + ("md5_1", "data-00001.arrow"), + ("md5_2", "data-00002.arrow"), + ] + mock_manifest.return_value = ("https://base.url", entries) + mock_convert.return_value = 100 + + with tempfile.TemporaryDirectory() as tmp: + pipeline = StreamingIngestPipeline( + manifest_url="http://manifest.uri", + input_dir=f"{tmp}/arrows", + output_dir=f"{tmp}/output", + split_name="train", + max_workers=2, + prefetch_depth=2, + ) + + def fake_download(url, dest, **kwargs): + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(b"fake") + + mock_download.side_effect = fake_download + + results = pipeline.run() + + assert results["files_processed"] == 3 + assert results["samples_written"] == 300 + assert results["shards_created"] == 3 + assert mock_download.call_count == 3 + assert mock_convert.call_count == 3 + + @patch(f"{_INGEST_MODULE}._arrow_to_tar") + @patch(f"{_INGEST_MODULE}.download_with_backoff") + @patch(f"{_INGEST_MODULE}.fetch_manifest") + def test_max_files_limits_processing(self, mock_manifest, mock_download, mock_convert): + """max_files parameter limits how many files are processed.""" + entries = [(f"md5_{i}", f"data-{i:05d}.arrow") for i in range(10)] + mock_manifest.return_value = ("https://base.url", entries) + mock_convert.return_value = 50 + + with tempfile.TemporaryDirectory() as tmp: + pipeline = StreamingIngestPipeline( + manifest_url="http://manifest.uri", + input_dir=f"{tmp}/arrows", + output_dir=f"{tmp}/output", + split_name="train", + max_files=3, + max_workers=2, + prefetch_depth=2, + ) + + def fake_download(url, dest, **kwargs): + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(b"fake") + + mock_download.side_effect = fake_download + + results = pipeline.run() + + assert results["files_processed"] == 3 + assert mock_download.call_count == 3 + + @patch(f"{_INGEST_MODULE}._arrow_to_tar") + @patch(f"{_INGEST_MODULE}.download_with_backoff") + @patch(f"{_INGEST_MODULE}.fetch_manifest") + def test_download_error_skips_file(self, mock_manifest, mock_download, mock_convert): + """A failed download is skipped; the pipeline continues with remaining files.""" + entries = [("md5_0", "data-00000.arrow"), ("md5_1", "data-00001.arrow")] + mock_manifest.return_value = ("https://base.url", entries) + mock_convert.return_value = 100 + + call_count = [0] + + def selective_download(url, dest, **kwargs): + call_count[0] += 1 + if "data-00000" in url: + raise RuntimeError("Download failed: HTTP 500") + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(b"fake") + + mock_download.side_effect = selective_download + + with tempfile.TemporaryDirectory() as tmp: + pipeline = StreamingIngestPipeline( + manifest_url="http://manifest.uri", + input_dir=f"{tmp}/arrows", + output_dir=f"{tmp}/output", + split_name="train", + max_workers=1, + prefetch_depth=2, + ) + + results = pipeline.run() + + assert results["files_processed"] == 1 + assert results["files_failed"] == 1 + assert results["samples_written"] == 100 + assert mock_convert.call_count == 1 + + failed_manifest = Path(tmp) / "output" / "train" / "failed_files.json" + assert failed_manifest.exists() + failures = json.loads(failed_manifest.read_text()) + assert len(failures) == 1 + assert failures[0]["stage"] == "download" + assert failures[0]["filename"] == "data-00000.arrow" + + @patch(f"{_INGEST_MODULE}._arrow_to_tar") + @patch(f"{_INGEST_MODULE}.download_with_backoff") + @patch(f"{_INGEST_MODULE}.fetch_manifest") + def test_all_downloads_fail(self, mock_manifest, mock_download, mock_convert): + """If all downloads fail, the pipeline completes with zero processed.""" + entries = [("md5_0", "data-00000.arrow"), ("md5_1", "data-00001.arrow")] + mock_manifest.return_value = ("https://base.url", entries) + mock_download.side_effect = RuntimeError("Download failed: HTTP 500") + + with tempfile.TemporaryDirectory() as tmp: + pipeline = StreamingIngestPipeline( + manifest_url="http://manifest.uri", + input_dir=f"{tmp}/arrows", + output_dir=f"{tmp}/output", + split_name="train", + max_workers=1, + prefetch_depth=2, + ) + + results = pipeline.run() + + assert results["files_processed"] == 0 + assert results["files_failed"] == 2 + assert results["samples_written"] == 0 + assert mock_convert.call_count == 0 + + @patch(f"{_INGEST_MODULE}._arrow_to_tar") + @patch(f"{_INGEST_MODULE}.download_with_backoff") + @patch(f"{_INGEST_MODULE}.fetch_manifest") + def test_conversion_error_skips_file(self, mock_manifest, mock_download, mock_convert): + """A failed conversion is skipped; partial tar is cleaned up.""" + entries = [ + ("md5_0", "data-00000.arrow"), + ("md5_1", "data-00001.arrow"), + ] + mock_manifest.return_value = ("https://base.url", entries) + + def selective_convert(arrow_path, tar_path, offset): + if "data-00000" in str(arrow_path): + raise RuntimeError("Corrupt Arrow file") + return 100 + + mock_convert.side_effect = selective_convert + + with tempfile.TemporaryDirectory() as tmp: + pipeline = StreamingIngestPipeline( + manifest_url="http://manifest.uri", + input_dir=f"{tmp}/arrows", + output_dir=f"{tmp}/output", + split_name="train", + max_workers=1, + prefetch_depth=2, + ) + + def fake_download(url, dest, **kwargs): + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(b"fake") + + mock_download.side_effect = fake_download + + results = pipeline.run() + + assert results["files_processed"] == 1 + assert results["files_failed"] == 1 + assert results["samples_written"] == 100 + + output_train = Path(tmp) / "output" / "train" + assert not (output_train / "shard_000000.tar").exists() + assert (output_train / "failed_files.json").exists() + failures = json.loads((output_train / "failed_files.json").read_text()) + assert len(failures) == 1 + assert failures[0]["stage"] == "conversion" + + @patch(f"{_INGEST_MODULE}._arrow_to_tar") + @patch(f"{_INGEST_MODULE}.download_with_backoff") + @patch(f"{_INGEST_MODULE}.fetch_manifest") + def test_arrow_files_deleted_after_conversion(self, mock_manifest, mock_download, mock_convert): + """Arrow files are cleaned up after conversion to tar.""" + entries = [("md5_0", "data-00000.arrow")] + mock_manifest.return_value = ("https://base.url", entries) + mock_convert.return_value = 10 + + with tempfile.TemporaryDirectory() as tmp: + arrow_dir = Path(tmp) / "arrows" + pipeline = StreamingIngestPipeline( + manifest_url="http://manifest.uri", + input_dir=str(arrow_dir), + output_dir=f"{tmp}/output", + split_name="train", + max_workers=1, + prefetch_depth=2, + ) + + created_files = [] + + def fake_download(url, dest, **kwargs): + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(b"fake arrow data") + created_files.append(dest) + + mock_download.side_effect = fake_download + + pipeline.run() + + for f in created_files: + assert not f.exists(), f"Arrow file should have been deleted: {f}" + + @patch(f"{_INGEST_MODULE}._arrow_to_tar") + @patch(f"{_INGEST_MODULE}.download_with_backoff") + @patch(f"{_INGEST_MODULE}.fetch_manifest") + def test_sample_offset_accumulates(self, mock_manifest, mock_download, mock_convert): + """Global sample offset is passed correctly across shards.""" + entries = [ + ("md5_0", "data-00000.arrow"), + ("md5_1", "data-00001.arrow"), + ("md5_2", "data-00002.arrow"), + ] + mock_manifest.return_value = ("https://base.url", entries) + + sample_counts = [100, 200, 150] + mock_convert.side_effect = sample_counts + + with tempfile.TemporaryDirectory() as tmp: + pipeline = StreamingIngestPipeline( + manifest_url="http://manifest.uri", + input_dir=f"{tmp}/arrows", + output_dir=f"{tmp}/output", + split_name="train", + max_workers=1, + prefetch_depth=2, + ) + + def fake_download(url, dest, **kwargs): + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(b"fake") + + mock_download.side_effect = fake_download + + results = pipeline.run() + + offsets = [c.args[2] for c in mock_convert.call_args_list] + assert offsets == [0, 100, 300] + assert results["samples_written"] == 450 + + @patch(f"{_INGEST_MODULE}._arrow_to_tar") + @patch(f"{_INGEST_MODULE}.download_with_backoff") + @patch(f"{_INGEST_MODULE}.fetch_manifest") + def test_resume_skips_existing_shards(self, mock_manifest, mock_download, mock_convert): + """Pre-existing shard tars are skipped; only missing shards are downloaded.""" + entries = [ + ("md5_0", "data-00000.arrow"), + ("md5_1", "data-00001.arrow"), + ("md5_2", "data-00002.arrow"), + ("md5_3", "data-00003.arrow"), + ] + mock_manifest.return_value = ("https://base.url", entries) + mock_convert.return_value = 100 + + with tempfile.TemporaryDirectory() as tmp: + output_train = Path(tmp) / "output" / "train" + output_train.mkdir(parents=True) + + (output_train / "shard_000000.tar").write_bytes(b"existing") + (output_train / "shard_000002.tar").write_bytes(b"existing") + + pipeline = StreamingIngestPipeline( + manifest_url="http://manifest.uri", + input_dir=f"{tmp}/arrows", + output_dir=f"{tmp}/output", + split_name="train", + max_workers=1, + prefetch_depth=2, + ) + + def fake_download(url, dest, **kwargs): + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(b"fake") + + mock_download.side_effect = fake_download + + results = pipeline.run() + + assert results["files_processed"] == 2 + assert results["shards_skipped"] == 2 + assert results["shards_created"] == 4 # 2 new + 2 skipped + assert results["samples_written"] == 200 # only from the 2 new shards + assert mock_download.call_count == 2 + assert mock_convert.call_count == 2 diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_utils.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_utils.py new file mode 100644 index 000000000..f35d1a02b --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_utils.py @@ -0,0 +1,164 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for preprocessing utility functions. + +Tests distributed processing utilities and work splitting logic. +""" + +from unittest.mock import patch + +import pytest + +from primus.backends.megatron.data.diffusion.preprocessing.utils import ( + get_distributed_info, + split_work_for_rank, +) +from tests.utils import PrimusUT + + +class TestGetDistributedInfo(PrimusUT): + """Tests for get_distributed_info utility.""" + + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.is_available") + def test_not_distributed(self, mock_is_available, mock_is_initialized): + """Test get_distributed_info when not in distributed mode.""" + mock_is_available.return_value = True + mock_is_initialized.return_value = False + + rank, world_size = get_distributed_info() + + assert rank == 0 + assert world_size == 1 + + @patch("torch.distributed.get_world_size") + @patch("torch.distributed.get_rank") + @patch("torch.distributed.is_initialized") + @patch("torch.distributed.is_available") + def test_distributed_mode( + self, mock_is_available, mock_is_initialized, mock_get_rank, mock_get_world_size + ): + """Test get_distributed_info in distributed mode.""" + mock_is_available.return_value = True + mock_is_initialized.return_value = True + mock_get_rank.return_value = 2 + mock_get_world_size.return_value = 8 + + rank, world_size = get_distributed_info() + + assert rank == 2 + assert world_size == 8 + + +class TestSplitWorkForRank(PrimusUT): + """Tests for split_work_for_rank utility.""" + + def test_multiple_ranks_evenly_divisible(self): + """Test work splitting when items divide evenly.""" + world_size = 4 + total_items = 100 + + expected_splits = [ + (0, 25), # rank 0 + (25, 50), # rank 1 + (50, 75), # rank 2 + (75, 100), # rank 3 (last rank gets remainder) + ] + + for rank, (expected_start, expected_end) in enumerate(expected_splits): + start, end = split_work_for_rank(total_items, rank, world_size) + assert start == expected_start + assert end == expected_end + + def test_multiple_ranks_with_remainder(self): + """Test work splitting when items don't divide evenly.""" + world_size = 3 + total_items = 100 + + expected_splits = [ + (0, 33), # rank 0: 33 items + (33, 66), # rank 1: 33 items + (66, 100), # rank 2: 34 items (gets remainder) + ] + + for rank, (expected_start, expected_end) in enumerate(expected_splits): + start, end = split_work_for_rank(total_items, rank, world_size) + assert start == expected_start + assert end == expected_end + + def test_last_rank_gets_remainder(self): + """Test that last rank gets any remaining items.""" + total_items = 107 + world_size = 8 + items_per_rank = total_items // world_size # 13 + + # Last rank should get remainder + start, end = split_work_for_rank(total_items, rank=7, world_size=world_size) + + assert start == 7 * items_per_rank # 91 + assert end == total_items # 107 (gets all remaining) + + def test_covers_all_items(self): + """Test that all ranks together cover all items.""" + total_items = 137 + world_size = 7 + + all_indices = set() + for rank in range(world_size): + start, end = split_work_for_rank(total_items, rank, world_size) + rank_indices = set(range(start, end)) + # No overlap + assert len(all_indices & rank_indices) == 0 + all_indices.update(rank_indices) + + # All items covered + assert all_indices == set(range(total_items)) + + def test_fewer_items_than_ranks(self): + """Test work splitting when items < ranks.""" + total_items = 5 + world_size = 10 + + # First 5 ranks get 0 items each (5 // 10 = 0) + # Last rank gets all items + for rank in range(world_size - 1): + start, end = split_work_for_rank(total_items, rank, world_size) + # Each rank gets 0 items except last + assert start == 0 + assert end == 0 + + # Last rank gets all + start, end = split_work_for_rank(total_items, rank=world_size - 1, world_size=world_size) + assert end == total_items + + +class TestWorkSplittingEdgeCases(PrimusUT): + """Test edge cases for work splitting.""" + + def test_empty_work(self): + """Test with zero items.""" + start, end = split_work_for_rank(total=0, rank=0, world_size=4) + + assert start == 0 + assert end == 0 + + def test_one_item(self): + """Test with single item.""" + world_size = 4 + + # Ranks 0-2 get 0 items (1 // 4 = 0) + for rank in range(world_size - 1): + start, end = split_work_for_rank(1, rank, world_size) + assert start == 0 + assert end == 0 + + # Last rank gets the 1 item + start, end = split_work_for_rank(1, rank=world_size - 1, world_size=world_size) + assert start == 0 + assert end == 1 + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_validate.py b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_validate.py new file mode 100644 index 000000000..f630163c0 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/data/preprocessing/test_validate.py @@ -0,0 +1,194 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for Energon dataset validation module. + +Tests _check_metadata, _check_sample_counts, and validate_energon_dataset +using synthetic dataset directories. +""" + +import json +import tarfile +import tempfile +from pathlib import Path +from unittest.mock import patch + +import pytest +import yaml + +from primus.backends.megatron.data.diffusion.preprocessing.validate import ( + _check_metadata, + _check_sample_counts, + validate_energon_dataset, +) +from tests.utils import PrimusUT + + +def _build_valid_dataset(base: Path, num_samples: int = 3) -> None: + """Helper: build a minimal valid dataset directory on disk.""" + meta_dir = base / ".nv-meta" + meta_dir.mkdir(parents=True) + + shard_name = "000000.tar" + idx_name = shard_name + ".idx" + + # Write .info.json + info = {"shard_counts": {shard_name: num_samples}} + (meta_dir / ".info.json").write_text(json.dumps(info)) + + # Write split.yaml + split = {"split_parts": {"train": [shard_name]}} + (meta_dir / "split.yaml").write_text(yaml.dump(split)) + + # Write dataset.yaml + (meta_dir / "dataset.yaml").write_text("__module__: megatron.energon\n__class__: CrudeWebdataset\n") + + # Write shard tar with correct number of samples + tar_path = base / shard_name + with tarfile.open(str(tar_path), "w") as tar: + for i in range(num_samples): + key = f"000000_{i:06d}" + for ext in ("jpg", "txt"): + member_name = f"{key}.{ext}" + data = b"test" + info_obj = tarfile.TarInfo(name=member_name) + info_obj.size = len(data) + import io + + tar.addfile(info_obj, io.BytesIO(data)) + + # Write .idx stub + (base / idx_name).write_text("") + + +class TestCheckMetadata(PrimusUT): + """Tests for _check_metadata.""" + + def test_valid_metadata(self): + """Valid .nv-meta with all files returns info dict.""" + with tempfile.TemporaryDirectory() as tmpdir: + base = Path(tmpdir) + _build_valid_dataset(base) + info = _check_metadata(base) + + assert info is not None + assert "shard_counts" in info + assert "_splits" in info + + def test_missing_info_json(self): + """Missing .info.json returns None.""" + with tempfile.TemporaryDirectory() as tmpdir: + base = Path(tmpdir) + _build_valid_dataset(base) + (base / ".nv-meta" / ".info.json").unlink() + + assert _check_metadata(base) is None + + def test_missing_shard_counts_key(self): + """Info file without shard_counts key returns None.""" + with tempfile.TemporaryDirectory() as tmpdir: + base = Path(tmpdir) + _build_valid_dataset(base) + (base / ".nv-meta" / ".info.json").write_text(json.dumps({"other": 1})) + + assert _check_metadata(base) is None + + def test_missing_split_yaml(self): + """Missing split.yaml returns None.""" + with tempfile.TemporaryDirectory() as tmpdir: + base = Path(tmpdir) + _build_valid_dataset(base) + (base / ".nv-meta" / "split.yaml").unlink() + + assert _check_metadata(base) is None + + def test_missing_dataset_yaml(self): + """Missing dataset.yaml returns None.""" + with tempfile.TemporaryDirectory() as tmpdir: + base = Path(tmpdir) + _build_valid_dataset(base) + (base / ".nv-meta" / "dataset.yaml").unlink() + + assert _check_metadata(base) is None + + def test_missing_shard_file(self): + """Referenced shard file missing on disk returns None.""" + with tempfile.TemporaryDirectory() as tmpdir: + base = Path(tmpdir) + _build_valid_dataset(base) + (base / "000000.tar").unlink() + + assert _check_metadata(base) is None + + +class TestCheckSampleCounts(PrimusUT): + """Tests for _check_sample_counts.""" + + def test_matching_counts(self): + """Tar with correct sample count returns True.""" + with tempfile.TemporaryDirectory() as tmpdir: + base = Path(tmpdir) + _build_valid_dataset(base, num_samples=3) + + info = _check_metadata(base) + assert info is not None + assert _check_sample_counts(base, info) is True + + def test_mismatched_counts(self): + """Mismatch between .info.json and tar content returns False.""" + with tempfile.TemporaryDirectory() as tmpdir: + base = Path(tmpdir) + _build_valid_dataset(base, num_samples=3) + + # Overwrite .info.json with wrong count + wrong_info = {"shard_counts": {"000000.tar": 99}} + (base / ".nv-meta" / ".info.json").write_text(json.dumps(wrong_info)) + + info = _check_metadata(base) + assert info is not None + assert _check_sample_counts(base, info) is False + + +class TestValidateEnergonDataset(PrimusUT): + """Tests for validate_energon_dataset orchestration.""" + + @patch( + "primus.backends.megatron.data.diffusion.preprocessing.validate._check_sample_load", + return_value={"jpg": "bytes, len=100", "txt": "str"}, + ) + def test_full_pass(self, mock_load): + """Complete valid dataset returns True.""" + with tempfile.TemporaryDirectory() as tmpdir: + _build_valid_dataset(Path(tmpdir), num_samples=3) + assert validate_energon_dataset(tmpdir, encoding="raw") is True + + @patch( + "primus.backends.megatron.data.diffusion.preprocessing.validate._check_sample_load", + return_value={"jpg": "bytes, len=100", "txt": "str"}, + ) + def test_metadata_failure_returns_false(self, mock_load): + """Metadata failure returns False immediately (early exit).""" + with tempfile.TemporaryDirectory() as tmpdir: + # Empty dir has no .nv-meta at all + assert validate_energon_dataset(tmpdir, encoding="raw") is False + mock_load.assert_not_called() + + @patch( + "primus.backends.megatron.data.diffusion.preprocessing.validate._check_sample_load", + return_value={"jpg": "bytes, len=100", "txt": "str"}, + ) + def test_count_mismatch_returns_false(self, mock_load): + """Count mismatch returns False but still prints summary.""" + with tempfile.TemporaryDirectory() as tmpdir: + _build_valid_dataset(Path(tmpdir), num_samples=3) + + # Corrupt the count + wrong_info = {"shard_counts": {"000000.tar": 99}} + (Path(tmpdir) / ".nv-meta" / ".info.json").write_text(json.dumps(wrong_info)) + + assert validate_energon_dataset(tmpdir, encoding="raw") is False + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/cli/test_data_config.py b/tests/unit_tests/cli/test_data_config.py new file mode 100644 index 000000000..2db295857 --- /dev/null +++ b/tests/unit_tests/cli/test_data_config.py @@ -0,0 +1,253 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Tests for data preprocessing config parsing, validation, and authentication. + +Tests the YAML-to-CLI mapping in data.py and the auth priority chain in auth.py. +""" + +import argparse +import os +import tempfile +from unittest.mock import patch + +import pytest + +from primus.backends.megatron.data.diffusion.preprocessing.auth import ( + setup_hf_authentication, +) +from primus.cli.subcommands.data import ( + _flatten_preprocessing_config, + _get_encoded_parser_defaults, + _load_config_with_cli_overrides, + _validate_preprocessing_config, +) +from tests.utils import PrimusUT + + +class TestFlattenPreprocessingConfig(PrimusUT): + """Tests for _flatten_preprocessing_config YAML-to-flat-dict mapping.""" + + def test_all_sections_mapped(self): + """All 6 YAML sections produce correct flat keys.""" + config = { + "source": { + "type": "huggingface", + "hf_dataset": "diffusers/pokemon", + "hf_split": "train", + }, + "data_format": { + "image_key": "jpg", + "caption_key": "json.caption", + }, + "output": { + "output_dir": "/data/out", + "shard_size": 500, + }, + "model": { + "model_path": "my-model", + "batch_size": 4, + "precision": "fp16", + }, + "image": { + "image_size": 512, + "variable_size": True, + "center_crop": False, + }, + "auth": { + "hf_token_file": "/path/to/token", + }, + } + flat = _flatten_preprocessing_config(config) + + assert flat["source_type"] == "huggingface" + assert flat["hf_dataset"] == "diffusers/pokemon" + assert flat["image_key"] == "jpg" + assert flat["caption_key"] == "json.caption" + assert flat["output_dir"] == "/data/out" + assert flat["shard_size"] == 500 + assert flat["model_path"] == "my-model" + assert flat["batch_size"] == 4 + assert flat["precision"] == "fp16" + assert flat["image_size"] == 512 + assert flat["variable_size"] is True + assert flat["center_crop"] is False + assert flat["hf_token_file"] == "/path/to/token" + + def test_partial_config(self): + """Partial config with only source + output produces only those keys.""" + config = { + "source": {"type": "directory", "input_dir": "/data/images"}, + "output": {"output_dir": "/data/out"}, + } + flat = _flatten_preprocessing_config(config) + + assert flat["source_type"] == "directory" + assert flat["input_dir"] == "/data/images" + assert flat["output_dir"] == "/data/out" + assert "model_path" not in flat + assert "image_size" not in flat + + def test_model_defaults(self): + """Model section injects correct defaults when keys are absent.""" + config = {"model": {}} + flat = _flatten_preprocessing_config(config) + + assert flat["model_path"] == "black-forest-labs/FLUX.1-dev" + assert flat["batch_size"] == 8 + assert flat["precision"] == "bf16" + assert flat["device"] == "cuda" + assert flat["t5_max_length"] == 512 + + def test_image_defaults(self): + """Image section injects correct defaults when keys are absent.""" + config = {"image": {}} + flat = _flatten_preprocessing_config(config) + + assert flat["image_size"] == 1024 + assert flat["variable_size"] is False + assert flat["center_crop"] is True + assert flat["max_size"] == 1024 + + +class TestValidatePreprocessingConfig(PrimusUT): + """Tests for _validate_preprocessing_config.""" + + def _make_args(self, **kwargs): + defaults = { + "source_type": "huggingface", + "output_dir": "/data/out", + "hf_dataset": "test/dataset", + "input_dir": None, + "input_path": None, + } + defaults.update(kwargs) + return argparse.Namespace(**defaults) + + def test_missing_source_type_raises(self): + """Missing source_type raises ValueError.""" + args = self._make_args(source_type=None) + with self.assertRaises(ValueError, msg="source_type"): + _validate_preprocessing_config(args) + + def test_missing_output_dir_raises(self): + """Missing output_dir raises ValueError.""" + args = self._make_args(output_dir=None) + with self.assertRaises(ValueError, msg="output_dir"): + _validate_preprocessing_config(args) + + def test_huggingface_without_hf_dataset_raises(self): + """HuggingFace source without hf_dataset raises ValueError.""" + args = self._make_args(source_type="huggingface", hf_dataset=None) + with self.assertRaises(ValueError, msg="hf-dataset"): + _validate_preprocessing_config(args) + + def test_directory_without_input_dir_raises(self): + """Directory source without input_dir raises ValueError.""" + args = self._make_args(source_type="directory", input_dir=None) + with self.assertRaises(ValueError, msg="input-dir"): + _validate_preprocessing_config(args) + + def test_valid_config_passes(self): + """Valid config raises no errors.""" + args = self._make_args() + _validate_preprocessing_config(args) + + +class TestLoadConfigWithCliOverrides(PrimusUT): + """Tests for _load_config_with_cli_overrides merge logic.""" + + @patch("primus.core.utils.yaml_utils.parse_yaml") + def test_cli_overrides_yaml(self, mock_parse_yaml): + """Explicitly set CLI args override YAML config values.""" + mock_parse_yaml.return_value = { + "source": {"type": "huggingface", "hf_dataset": "yaml-dataset"}, + "model": {"batch_size": 4}, + } + defaults = _get_encoded_parser_defaults() + args = argparse.Namespace( + config="test.yaml", + batch_size=16, + **{k: v for k, v in defaults.items() if k not in ("config", "batch_size")}, + ) + + result = _load_config_with_cli_overrides(args) + + assert result.batch_size == 16 + assert result.hf_dataset == "yaml-dataset" + + @patch("primus.core.utils.yaml_utils.parse_yaml") + def test_yaml_used_when_cli_is_default(self, mock_parse_yaml): + """YAML values used when CLI arg equals its default.""" + mock_parse_yaml.return_value = { + "model": {"batch_size": 4}, + } + defaults = _get_encoded_parser_defaults() + args = argparse.Namespace(config="test.yaml", **{k: v for k, v in defaults.items() if k != "config"}) + + result = _load_config_with_cli_overrides(args) + + assert result.batch_size == 4 + + def test_no_config_passthrough(self): + """No config file returns args unchanged.""" + original = argparse.Namespace(config=None, batch_size=8) + result = _load_config_with_cli_overrides(original) + assert result is original + + +class TestSetupHfAuthenticationPriority(PrimusUT): + """Tests for setup_hf_authentication priority chain.""" + + def test_file_takes_priority_over_env(self): + """Token file takes priority over HF_TOKEN env var.""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".token", delete=False) as f: + f.write("hf_file_token_123") + token_path = f.name + try: + os.chmod(token_path, 0o600) + old_env = os.environ.get("HF_TOKEN") + os.environ["HF_TOKEN"] = "hf_env_token_456" + try: + token = setup_hf_authentication(token_file=token_path, use_env=True) + assert token == "hf_file_token_123" + finally: + if old_env is None: + os.environ.pop("HF_TOKEN", None) + else: + os.environ["HF_TOKEN"] = old_env + finally: + os.unlink(token_path) + + def test_env_takes_priority_over_cache(self): + """HF_TOKEN env var is used when no file is provided.""" + old_env = os.environ.get("HF_TOKEN") + os.environ["HF_TOKEN"] = "hf_env_token_789" + try: + token = setup_hf_authentication(token_file=None, use_env=True) + assert token == "hf_env_token_789" + finally: + if old_env is None: + os.environ.pop("HF_TOKEN", None) + else: + os.environ["HF_TOKEN"] = old_env + + def test_no_auth_returns_none(self): + """No auth sources returns None.""" + old_env = os.environ.pop("HF_TOKEN", None) + try: + with patch.object( + type(__import__("pathlib").Path()), + "exists", + return_value=False, + ): + token = setup_hf_authentication(token_file=None, use_env=True) + assert token is None + finally: + if old_env is not None: + os.environ["HF_TOKEN"] = old_env + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 5c68528b38d698b9ef2f7c22739242d11d9bc1d8 Mon Sep 17 00:00:00 2001 From: Andy <14128880+yeandy@users.noreply.github.com> Date: Wed, 15 Jul 2026 20:30:46 -0400 Subject: [PATCH 031/127] feat(runner): add run_preflight_direct.sh for non-container preflight (#705) SLURM-aware wrapper around `primus-cli direct -- preflight` that exports the distributed env (NNODES, NODE_RANK, MASTER_ADDR, ...) which `primus-cli direct` does not derive from SLURM itself, activates a shared Python venv via VENV_ACTIVATE, supports a wrapper-only --silent flag (preserving the report path via a saved fd), and auto-generates unique timestamped report names to avoid colliding with stale leftovers. Also add docs/preflight-direct.md with a full walkthrough (venv setup, SLURM invocation, NCCL config for Broadcom / Pensando AINIC, troubleshooting), and cross-link it from docs/README.md and docs/preflight.md. --------- Co-authored-by: Fuyuan Jing Co-authored-by: Andrew Ma Co-authored-by: amd-ama10002-2 Co-authored-by: Akash Haridas Co-authored-by: Akash Haridas <58511267+akasharidas@users.noreply.github.com> Co-authored-by: Cursor Co-authored-by: Fuyuan Jing <167437074+amd-fuyuajin@users.noreply.github.com> --- .../node-smoke-test-instruction.md | 452 +++++++ .../preflight-without-container.md | 792 ++++++++++++ docs/02-user-guide/preflight.md | 413 +++++- docs_deprecated/cli/PRIMUS-CLI-GUIDE.md | 28 +- .../node-smoke-test-instruction.md | 313 +++++ docs_deprecated/node-smoke.md | 429 +++++++ docs_deprecated/preflight-direct.md | 794 ++++++++++++ docs_deprecated/preflight.md | 514 +++++++- primus/cli/subcommands/node_smoke.py | 131 ++ primus/core/config/yaml_loader.py | 4 +- primus/core/utils/yaml_utils.py | 3 +- primus/tools/preflight/README.md | 22 +- primus/tools/preflight/global_vars.py | 27 + primus/tools/preflight/gpu/gpu_basic.py | 17 +- primus/tools/preflight/gpu/gpu_probe.py | 78 +- primus/tools/preflight/gpu/gpu_topology.py | 74 +- primus/tools/preflight/gpu/sysfs_probe.py | 385 ++++++ primus/tools/preflight/host/host_probe.py | 49 +- primus/tools/preflight/host/info.py | 15 +- primus/tools/preflight/inter_node_comm.py | 226 +++- primus/tools/preflight/inter_node_comm_p2p.py | 90 +- primus/tools/preflight/inter_node_ring_p2p.py | 41 +- primus/tools/preflight/intra_node_comm.py | 137 +- primus/tools/preflight/node_smoke/__init__.py | 58 + primus/tools/preflight/node_smoke/__main__.py | 14 + .../node_smoke/aggregator/__init__.py | 18 + .../preflight/node_smoke/aggregator/report.py | 719 +++++++++++ .../node_smoke/aggregator/summarizers.py | 449 +++++++ primus/tools/preflight/node_smoke/cli.py | 910 ++++++++++++++ .../node_smoke/collectors/__init__.py | 23 + .../preflight/node_smoke/collectors/clock.py | 36 + .../preflight/node_smoke/collectors/dmesg.py | 98 ++ .../node_smoke/collectors/fingerprint.py | 88 ++ .../node_smoke/collectors/gpu_low_level.py | 249 ++++ .../node_smoke/collectors/gpu_processes.py | 614 +++++++++ .../node_smoke/collectors/host_limits.py | 92 ++ .../preflight/node_smoke/collectors/nics.py | 452 +++++++ .../node_smoke/collectors/reused_info.py | 51 + .../node_smoke/collectors/rocm_smi.py | 409 ++++++ .../node_smoke/collectors/tooling.py | 100 ++ .../preflight/node_smoke/collectors/xgmi.py | 201 +++ .../preflight/node_smoke/logging_utils.py | 53 + .../preflight/node_smoke/orchestrator.py | 288 +++++ primus/tools/preflight/node_smoke/per_gpu.py | 327 +++++ .../tools/preflight/node_smoke/rccl_local.py | 152 +++ .../tools/preflight/node_smoke/shell_utils.py | 179 +++ .../preflight/node_smoke/tests/__init__.py | 0 .../node_smoke/tests/test_node_smoke.py | 1109 +++++++++++++++++ primus/tools/preflight/node_smoke/types.py | 36 + primus/tools/preflight/preflight_args.py | 163 ++- primus/tools/preflight/preflight_perf_test.py | 521 +++++++- primus/tools/preflight/square_gemm.py | 15 +- .../preflight/tests/test_report_naming.py | 238 ++++ primus/tools/preflight/utility.py | 52 +- runner/.primus.yaml | 5 +- runner/primus-cli-direct.sh | 197 ++- runner/primus-cli-slurm-entry.sh | 35 +- runner/primus-cli-slurm.sh | 15 +- tests/runner/test_primus_cli_direct.sh | 278 +++++ .../cli/test_preflight_subcommand.py | 4 +- .../tools/test_preflight_bisect_slurm.py | 223 ++++ tools/preflight_bisect/bisect.py | 352 ++++++ tools/preflight_bisect/fake_runner.sh | 54 + 63 files changed, 13531 insertions(+), 380 deletions(-) create mode 100644 docs/02-user-guide/node-smoke-test-instruction.md create mode 100644 docs/02-user-guide/preflight-without-container.md create mode 100644 docs_deprecated/node-smoke-test-instruction.md create mode 100644 docs_deprecated/node-smoke.md create mode 100644 docs_deprecated/preflight-direct.md create mode 100644 primus/cli/subcommands/node_smoke.py create mode 100644 primus/tools/preflight/gpu/sysfs_probe.py create mode 100644 primus/tools/preflight/node_smoke/__init__.py create mode 100644 primus/tools/preflight/node_smoke/__main__.py create mode 100644 primus/tools/preflight/node_smoke/aggregator/__init__.py create mode 100644 primus/tools/preflight/node_smoke/aggregator/report.py create mode 100644 primus/tools/preflight/node_smoke/aggregator/summarizers.py create mode 100644 primus/tools/preflight/node_smoke/cli.py create mode 100644 primus/tools/preflight/node_smoke/collectors/__init__.py create mode 100644 primus/tools/preflight/node_smoke/collectors/clock.py create mode 100644 primus/tools/preflight/node_smoke/collectors/dmesg.py create mode 100644 primus/tools/preflight/node_smoke/collectors/fingerprint.py create mode 100644 primus/tools/preflight/node_smoke/collectors/gpu_low_level.py create mode 100644 primus/tools/preflight/node_smoke/collectors/gpu_processes.py create mode 100644 primus/tools/preflight/node_smoke/collectors/host_limits.py create mode 100644 primus/tools/preflight/node_smoke/collectors/nics.py create mode 100644 primus/tools/preflight/node_smoke/collectors/reused_info.py create mode 100644 primus/tools/preflight/node_smoke/collectors/rocm_smi.py create mode 100644 primus/tools/preflight/node_smoke/collectors/tooling.py create mode 100644 primus/tools/preflight/node_smoke/collectors/xgmi.py create mode 100644 primus/tools/preflight/node_smoke/logging_utils.py create mode 100644 primus/tools/preflight/node_smoke/orchestrator.py create mode 100644 primus/tools/preflight/node_smoke/per_gpu.py create mode 100644 primus/tools/preflight/node_smoke/rccl_local.py create mode 100644 primus/tools/preflight/node_smoke/shell_utils.py create mode 100644 primus/tools/preflight/node_smoke/tests/__init__.py create mode 100644 primus/tools/preflight/node_smoke/tests/test_node_smoke.py create mode 100644 primus/tools/preflight/node_smoke/types.py create mode 100644 primus/tools/preflight/tests/test_report_naming.py create mode 100644 tests/unit_tests/tools/test_preflight_bisect_slurm.py create mode 100644 tools/preflight_bisect/bisect.py create mode 100755 tools/preflight_bisect/fake_runner.sh diff --git a/docs/02-user-guide/node-smoke-test-instruction.md b/docs/02-user-guide/node-smoke-test-instruction.md new file mode 100644 index 000000000..4aafe17cb --- /dev/null +++ b/docs/02-user-guide/node-smoke-test-instruction.md @@ -0,0 +1,452 @@ +# Node-smoke test instruction + +A lightweight, distributed-rendezvous-free preflight check that runs on every node in parallel under SLURM. It produces a **single PASS/FAIL verdict per node** plus SLURM-ready `passing_nodes.txt` / `failing_nodes.txt` you can pipe straight into `srun --nodelist=` / `--exclude=`. + +Use it to **screen a cluster fast and exclude bad nodes before launching a real training job**. A bad GPU, NIC, wedged driver, or leaked process on any node surfaces as a node FAIL — without a single global rendezvous, so a stuck node can't wedge its peers. + +- **Recommended launcher**: `runner/primus-cli slurm srun -- direct -- node_smoke ...` (auto-resolves the distributed env, applies `slurm.*` config defaults, same pattern as `train` / `benchmark`). The shorter `runner/primus-cli direct -- node_smoke ...` (bare `srun` + `direct`) is equivalent and handy for ad-hoc runs. +- **Companion tool**: [`preflight`](./preflight.md) — the heavier diagnostic with a global rendezvous and inter-node bandwidth tests. The recommended workflow is **node-smoke first, preflight second** (see [§10](#10-comparison-with-the-full-preflight)). + +--- + +## 1. What it does + +Node-smoke answers one question fast: **"which nodes are healthy enough to run anything?"** Because training jobs allocate whole nodes, a single degraded GPU (or NIC, or wedged driver) takes an entire node out of rotation. Node-smoke checks each node independently and emits a per-node verdict plus a ready-to-use exclude list, so you can prune broken nodes before committing a large job to a global rendezvous. + +It deliberately does **not** measure cross-node bandwidth — that's what [`preflight`](./preflight.md) is for. + +--- + +## 2. How it works + +- **Per-node and independent** — every node runs the checks on its own. No `MASTER_ADDR`, no `MASTER_PORT`, no global `torch.distributed` rendezvous, so a stuck node cannot wedge its peers. +- **Per-GPU isolation** — each GPU's checks run in their own Python subprocess with a hard timeout (`--per-gpu-timeout-sec`, default 15 s). A stuck `torch.cuda.set_device()` (which `signal.alarm` cannot interrupt because it sits inside a driver syscall) is `SIGKILL`'d from the parent without affecting the rest of the node's checks. +- **Local-only RCCL** — the optional Tier 2 all-reduce uses `torch.multiprocessing.spawn` over `tcp://127.0.0.1`. No cross-node communication. +- **Rank-0 aggregation** — `NODE_RANK==0` polls for the expected number of per-node JSONs (with a timeout), computes cluster-wide drift, writes the Markdown report + pass/fail lists, and returns non-zero if any node FAILs or never reports. + +--- + +## 3. Prerequisites + +| Prerequisite | How | +| ----------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| Python venv on a shared filesystem | Same venv used by `primus-cli direct -- preflight` (see [`preflight-without-container.md`](./preflight-without-container.md) §2). | +| `VENV_ACTIVATE` exported | `export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate` (optional inside the container path). | +| Inside an existing SLURM allocation | One task per node. Recommended: `runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 -- direct -- node_smoke ...`. Equivalent bare form: `srun ... --ntasks-per-node=1 runner/primus-cli direct -- node_smoke ...`. Either way the `direct -- node_smoke` path auto-selects `--single`, so each task spawns one Python process and per-GPU subprocesses are launched internally. | + +No `MASTER_ADDR`, no `MASTER_PORT`, no global rendezvous required. + +--- + +## 4. Quick start + +**Git clone the Primus repository to a shared filesystem that all nodes can read.** + +```bash +git clone --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.git +cd Primus +``` + +**Note: remember to set up the Python virtual environment and NCCL / fabric environment variables as described in [§3 Prerequisites](#3-prerequisites).** + +> ⚠ **Set the NCCL / RCCL environment first** if you plan to run with `--tier2-perf` (the local 8-GPU RCCL all-reduce). Even though the smoke test never opens a cross-node rendezvous, the Tier 2 RCCL step calls `dist.init_process_group(backend="nccl", ...)`, and RCCL **enumerates every transport at init** (XGMI / PCIe P2P + IB + sockets). A misconfigured `NCCL_IB_HCA` / `NCCL_SOCKET_IFNAME` / `NCCL_IB_GID_INDEX` can stall init or make the all-reduce silently fall back to a slow path. The launcher's `base_env.sh` auto-detects these, **but auto-detect sometimes picks the wrong values inside a container** (devices masked by the network namespace, frontend NICs picked up instead of fabric NICs, etc.), so check them and set them explicitly if auto-detection is wrong. +> +> Minimum-viable checklist before running with `--tier2-perf`: +> +> ```bash +> # Pin the RDMA / RoCE training NICs the container can actually see. +> # On a bare-metal host the auto-detect in base_env.sh usually picks +> # the right set; inside a container or on a multi-role node, list +> # them explicitly. Use the same set you would pass to a training job. +> export NCCL_IB_HCA="rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7" +> +> # Pick the RoCE v2 GID index for your fabric: +> # - Mellanox / Broadcom: typically 3 (base_env.sh default). +> # - Pensando Pollara (AINIC): 1. +> export NCCL_IB_GID_INDEX=3 +> +> # The bootstrap socket interface. Auto-detect prefers the first +> # non-loopback interface from `hostname -I`; override when that +> # picks a frontend NIC instead of the data-plane interface. +> export NCCL_SOCKET_IFNAME=eno0 +> export GLOO_SOCKET_IFNAME=eno0 +> ``` +> +> See [`preflight-without-container.md` §4 Cluster-specific NCCL configuration](./preflight-without-container.md#4-cluster-specific-nccl-configuration) for the canonical Broadcom / Pensando Pollara values (the same `NCCL_*` set is used by both tools). If you skip `--tier2-perf`, the RCCL step is not executed and none of the above applies — Tier 1 (host limits, RDMA roll-call, leaked-process detection, etc.) does not depend on RCCL. +> +> Quick verification: `runner/primus-cli direct --dry-run -- node_smoke --tier2-perf` prints the resolved `NCCL_*` block under "Environment Variables" so you can confirm the values before launching for real. + +Recommended — through the `primus-cli slurm srun` wrapper (auto-resolves `MASTER_ADDR`/`NNODES`/`NODE_RANK`, applies `slurm.*` config defaults): + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate + +# Basic Tier 1 check (~5 s/GPU, ~30 s total) +runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + -- direct -- node_smoke + +# Tier 1 + Tier 2 perf sanity (GEMM TFLOPS, HBM GB/s, local 8-GPU RCCL) +runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + -- direct -- node_smoke --tier2-perf + +# Then re-run training, excluding any node the smoke test failed: +srun --exclude=$(paste -sd, output/preflight/failing_nodes.txt) ... your-real-job +``` + +Equivalent with bare `srun` (works the same; useful when composing with custom `srun` flags): + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke + +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf +``` + +Single-node sanity check (no SLURM): + +```bash +runner/primus-cli direct -- node_smoke +``` + +> **Both forms produce the same workload.** The wrapper form is recommended because it resolves the distributed env once on the launching node and propagates it via `--env`, and applies any `slurm.`* config defaults (partition / time / etc.). The `direct` keyword between the two `--`s is **mandatory** to take the no-container path — without it the wrapper routes through the container path. See [`preflight-without-container.md` § Wrapper vs. bare-srun](./preflight-without-container.md#wrapper-vs-bare-srun) for the precedence table. + +--- + +## 5. What's checked + +A `level='fail'` finding in any check FAILs the node. Everything else is reported as info / warn. + +### Tier 1 — always runs (~5 s/GPU) + +**Per-GPU liveness** (each GPU in its own subprocess with a hard timeout): + +- `torch.cuda.set_device(i)` — proves the device is bindable (a stale / wedged GPU often fails here). +- 256 MB allocation. +- Tiny 2048² bf16 GEMM with an `isfinite()` check on the result. + +**Host / GPU / network inventory** (no rendezvous): + +- **dmesg recent-error scan** — greps the last `--dmesg-minutes` (default 15) of `dmesg` for known patterns (`xid`, `gpu reset`, `hung_task`, `mce:`, `amdgpu.*error`, ...). Matches are surfaced in the report. +- **A. Software-stack fingerprint** — kernel / OS / Python, ROCm version, amdgpu kernel-module version, PyTorch / `torch.version.hip` / RCCL versions, and per-IB-device firmware + HCA model. Used for cluster drift detection. +- **B. NIC / RDMA roll-call** — per-port state read from `/sys/class/infiniband` (works inside containers; no `ibv_devinfo` / `ibstat` dependency). Many clusters expose more RDMA ports than the training job uses, so the hard-fail rules only run against the *training-NIC* subset, selected by this precedence: + 1. `--rdma-nic-allowlist` (`NCCL_IB_HCA` syntax: comma-separated `device[:port]`, `^...` denylist, `=dev` exact-match). + 2. `NCCL_IB_HCA` env (same syntax) — so the smoke test and the training launch agree by construction. + 3. Heuristic: auto-exclude any port whose `phys_state` is `Disabled` or `Sleep` (admin-disabled). + 4. Fallback: every IB port must be ACTIVE / LinkUp. + + **Hard-fail rules** (on the included set only): port not ACTIVE / not LinkUp; active port with zero RoCE v2 GIDs (RoCE) or zero valid GIDs (IB); included-NIC count ≠ `--expected-rdma-nics N` (when set). If *every* discovered port gets excluded, the node still fails — a node with zero training NICs cannot participate in inter-node training. Excluded ports stay visible in the report for diagnostics but don't contribute to the FAIL signal. +- **C. Host limits / system tunables** — `RLIMIT_MEMLOCK` below `--ulimit-l-min-gb` (default 32 GiB) → "RDMA pin will fail under load"; `/dev/shm` below `--shm-min-gb` (default 8 GiB) → "NCCL shared-mem may fail". NUMA node count, CPU count, and `cpu0` scaling governor are collected for drift detection. +- **Foreign / leaked process detection** — foreign PIDs holding a GPU FAIL the node by default (the most common cause of "training fails to launch on a healthy-looking node"). Allowed by default: `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter`. See the [container note in §6.5](#65-allow--extend-the-foreign-process-whitelist) — running inside a container almost always needs `--allow-foreign-procs`. +- **rocm-smi self-latency** — a `rocm-smi --version` call slower than `--rocm-smi-timeout-sec` (default 5 s) is a node FAIL; a wedging amdgpu driver typically hangs `rocm-smi` for 30–60 s before the GPU itself stops responding. + +### Tier 2 — optional perf sanity (`--tier2-perf`) + +Per-GPU steady-state metrics, with iteration counts aligned to the preflight `--quick` preset so smoke and preflight numbers are directly comparable. It's a single switch — you cannot enable just one half. + +- **GEMM TFLOPS** — 8192³ bf16 `torch.matmul`; FAIL below `--gemm-tflops-min` (default 600). +- **HBM GB/s** — 512 MB device-to-device `torch.Tensor.copy_` (counts read + write); FAIL below `--hbm-gbs-min` (default 2000; a healthy MI300X is ≈ 4500–5000). +- **Local 8-GPU RCCL all-reduce GB/s** — algorithmic bandwidth `2·S·(P-1)/P / t / 1e9` at `--rccl-size-mb` (default 64 MB); FAIL below `--rccl-gbs-min` (default 100). Local only, no cross-node traffic. + +--- + +## 6. Examples (by configuration knob) + +> **Convention used below.** The examples are written with bare `srun` for brevity. Anywhere you see `srun runner/primus-cli direct -- node_smoke ...`, the equivalent wrapper form is `runner/primus-cli slurm srun -- direct -- node_smoke ...`. Pick whichever matches your habits; both target the same launcher. + +### 6.1 Hard-fail on partial NIC enumeration + +Catches "7 of 8 RDMA NICs visible" — a common cause of crashes after RoCE init. The count is compared against the *training-NIC* set (after the selector chain), so frontend / storage RoCE NICs do not inflate it. + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf --expected-rdma-nics 8 +``` + +### 6.2 Pin the training-NIC set explicitly + +When auto-detection picks the wrong ports, name the training NICs directly (otherwise `NCCL_IB_HCA` env is used; otherwise admin-disabled ports are auto-excluded): + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf \ + --rdma-nic-allowlist 'rocep158s0:1,rocep190s0:1,rocep206s0:1,rocep222s0:1,rocep28s0:1,rocep62s0:1,rocep79s0:1,rocep96s0:1' +``` + +### 6.3 Tighten Tier 2 perf thresholds + +Reject GPUs that come in below your acceptance bar. Defaults: GEMM 600 TFLOPS, HBM 2000 GB/s, local RCCL 100 GB/s. + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf \ + --gemm-tflops-min 700 --hbm-gbs-min 4500 --rccl-gbs-min 180 +``` + +### 6.4 Tighten host limits + +Fail nodes whose `RLIMIT_MEMLOCK` or `/dev/shm` is too small for production training. + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke \ + --ulimit-l-min-gb 64 --shm-min-gb 16 +``` + +### 6.5 Allow / extend the foreign-process whitelist + +By default, leaked / foreign processes holding a GPU FAIL the node. Allowed by default: `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter`. + +```bash +# Add a site-specific monitoring agent to the whitelist +srun ... runner/primus-cli direct -- node_smoke \ + --allowed-procs gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter,my-monitor + +# Don't fail at all on foreign processes (still reported in the markdown) +srun ... runner/primus-cli direct -- node_smoke --allow-foreign-procs +``` + +> ⚠ **Running node_smoke inside a container almost always trips this check.** `amd-smi process --json` reports `name="N/A"` for kernel/system PIDs like `gpuagent` whose `/proc//comm` it cannot read, and the fallback name resolution inside `node_smoke` also fails because the container's `/proc` typically does not expose host PIDs (private PID namespace without `--pid=host`, or a `hidepid=2` mount). The unresolved name doesn't match the allowlist, so the check fires and the node FAILs — even though the only "foreign" processes are well-known system daemons holding zero HBM. +> +> **In the container path, pass `--allow-foreign-procs`:** +> +> ```bash +> srun ... runner/primus-cli direct -- node_smoke --tier2-perf --allow-foreign-procs +> ``` +> +> The processes are still listed in `smoke_report.md` under "Busy GPUs / leaked processes", so a real leak is still visible; only the FAIL verdict is downgraded. +> +> **Narrower alternative** — add the literal sentinel `N/A` to the allowlist so the check still catches leaks with resolvable names (e.g. a leftover `python` rank): +> +> ```bash +> srun ... runner/primus-cli direct -- node_smoke --tier2-perf \ +> --allowed-procs gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter,N/A +> ``` +> +> Name resolution runs first, so whenever a real name *can* be resolved it overrides `N/A` and the normal allowlist applies — the `N/A` entry only matches PIDs whose name genuinely could not be recovered. +> +> **Root-cause fix** (preferred long-term): grant the container access to host PIDs so names resolve and the report shows `gpuagent` etc. instead of `N/A`. Typical fixes: launch with `--pid=host` (Docker / Podman); mount `/proc` without `hidepid=2`; or loosen `ptrace_scope` / grant `CAP_SYS_PTRACE`. + +### 6.6 Require specific tools + +Make missing CLI tools a hard FAIL (default: warn-only). + +```bash +srun ... runner/primus-cli direct -- node_smoke --require-tools amd-smi,rocm-smi,lsof +``` + +### 6.7 Skip dmesg scan (containers with no privileges) + +```bash +srun ... runner/primus-cli direct -- node_smoke --skip-dmesg +``` + +### 6.8 Custom dump path + +Keep one report per smoke run instead of overwriting the default location. + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf \ + --dump-path /shared/smoke-archive/$(date +%Y%m%d-%H%M%S) +``` + +### 6.9 Re-aggregate from existing per-node JSONs (no re-run) + +The primus-cli wrapper always runs both phases (per-node run + rank-0 aggregate). To *only* re-render the report from JSONs collected earlier, use the standalone `aggregate` subcommand — it reads the existing `/smoke/*.json` without re-running the per-node step: + +```bash +python -m primus.tools.preflight.node_smoke aggregate \ + --dump-path output/preflight --expected-nodes 6 --wait-timeout-sec 5 +``` + +### 6.10 Silent mode (for CI) + +Suppresses wrapper stdout, but the **final report path is still printed** and stderr / exit code are preserved. + +```bash +srun ... runner/primus-cli direct --silent -- node_smoke --tier2-perf +``` + +### 6.11 Combined "production-ready screen" + +A representative one-shot for a production cluster screen: + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct --silent -- node_smoke --tier2-perf \ + --expected-rdma-nics 8 \ + --gemm-tflops-min 700 --hbm-gbs-min 4500 --rccl-gbs-min 180 \ + --ulimit-l-min-gb 64 --shm-min-gb 16 \ + --require-tools amd-smi,rocm-smi,lsof \ + --dump-path /shared/smoke-archive/$(date +%Y%m%d-%H%M%S) +``` + +--- + +## 7. Outputs + +All written under `--dump-path` (default `output/preflight/`). + +| File | Purpose | +| ------------------------- | -------------------------------------------------------------------------------------------------------------------- | +| `smoke/.json` | Per-node verdict + every collected metric. One file per node. | +| `smoke_report.md` | Human-readable cluster report (status table, drift sections, perf summary, failing-node detail). | +| `passing_nodes.txt` | Newline-separated short hostnames. Pipe into `srun --nodelist=`. | +| `failing_nodes.txt` | Newline-separated short hostnames. Pipe into `srun --exclude=`. | +| `expected_nodes.txt` | Auto-populated from `scontrol show hostnames "$SLURM_JOB_NODELIST"`. Lets the report name nodes that never reported. | + +Read the cluster verdict at a glance: + +```bash +head -10 output/preflight/smoke_report.md +``` + +Feed bad nodes into a re-run: + +```bash +srun --exclude=$(paste -sd, output/preflight/failing_nodes.txt) ... your-real-job +``` + +--- + +## 8. Understanding the report + +`smoke_report.md` renders in a stable order. Each section short-circuits to a placeholder (e.g. `*All nodes match.*`, `*No NIC issues.*`) on a healthy cluster, so a clean report stays short. In order: + +1. **Status table** — one row per node: `node_rank`, hostname, PASS/FAIL, duration, top fail reason. +2. **Stack drift across cluster** — per fingerprint key, outliers vs the cluster majority (catches "1 of N nodes on a different RCCL build"). +3. **NIC firmware drift across cluster** — per-IB-device firmware drift. +4. **NIC / RDMA roll-call issues** — every offending node + port (included set only). +5. **NIC port-count summary** — cluster-majority *training-NIC* count and any node that disagrees (catches partial-NIC degradation even without `--expected-rdma-nics`). +6. **NIC excluded ports (informational)** — ports the selector chain dropped, grouped by source. Does not contribute to FAIL. +7. **Host limits issues** — per-node hard-limit violations. +8. **GPU visibility issues** — nodes where torch couldn't see the GPUs, or amd-smi sees more GPUs than torch (stale ROCm / wedged driver). +9. **GPU low-level outliers (PCIe link / HBM)** — per-GPU outliers vs the cluster majority on PCIe width/speed and HBM total. +10. **XGMI link issues** — any non-XGMI GPU pair (intra-node collectives silently fall back to PCIe). +11. **Cluster clock + time daemons** — wall-clock spread plus per-node time-daemon health. +12. **Tooling self-latency (`rocm-smi --version`)** — slow / timed-out tool calls (precursor to a wedged driver). +13. **Tooling availability** — inventory of `amd-smi` / `rocm-smi` / `lsof` per node, plus which Tier 1 checks have no working tool on each node. +14. **Busy GPUs / leaked processes** — foreign PIDs holding GPUs at smoke start. +15. **GPU pre-touch HBM usage outliers** — GPUs with non-trivial HBM in use *before* smoke touched the device. +16. **GPU compute-activity outliers** — GPUs above `--gpu-activity-warn-pct` at smoke start (warn-only). +17. **Tier 2 perf summary** (only when at least one node ran Tier 2) — per-node GEMM TFLOPS / HBM GB/s as `min / median / max`, plus local RCCL GB/s. +18. **Failing nodes — full reasons** (only when there are failing nodes) — every fail reason, expanded per node. + +--- + +## 9. Configuration reference + +### 9.1 Common knobs (cheat sheet) + +| Flag | Default | When you'd change it | +| ------------------------ | ------------------------------------------------ | ------------------------------------------------------------------------------------- | +| `--tier2-perf` | off | Always on for production screens — adds GEMM TFLOPS, HBM GB/s, local RCCL all-reduce. | +| `--gemm-tflops-min N` | 600 | Site-specific acceptance bar. | +| `--hbm-gbs-min N` | 2000 | Site-specific acceptance bar (MI300X healthy ≈ 4500–5000). | +| `--rccl-gbs-min N` | 100 | Site-specific acceptance bar. | +| `--expected-rdma-nics N` | unset | Hard-fail on partial NIC enumeration. | +| `--ulimit-l-min-gb GB` | 32 | Raise for production training profiles. | +| `--shm-min-gb GB` | 8 | Raise for large-batch / many-rank profiles. | +| `--allow-foreign-procs` | off | Co-tenant clusters, shared GPUs, or the container path (see [§6.5](#65-allow--extend-the-foreign-process-whitelist)). | +| `--allowed-procs LIST` | `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter` | Add site-specific monitoring agents. | +| `--require-tools LIST` | `""` | Fail-fast if a CLI tool is missing in PATH. | +| `--skip-dmesg` | off | Inside unprivileged containers. | +| `--dump-path DIR` | `output/preflight` | Archive each run separately. | +| `--silent` (launcher) | off | CI / scripted runs. | + +### 9.2 Full `node_smoke` (per-node) flags + +Authoritative source: `python -m primus.tools.preflight.node_smoke run --help`. + +| Flag | Default | Purpose | +|---|---|---| +| `--dump-path` | `output/preflight` | Output directory. | +| `--expected-gpus N` | auto | Override GPU count (auto-detected from `LOCAL_WORLD_SIZE` / `GPUS_PER_NODE` / `torch.cuda.device_count()`). | +| `--per-gpu-timeout-sec` | 15 | Hard timeout per per-GPU subprocess. | +| `--tier2-perf` | off | Enable Tier 2 perf sanity (per-GPU GEMM TFLOPS + HBM GB/s + node-local RCCL all-reduce). Single switch. | +| `--gemm-tflops-min` | 600 | Tier 2 GEMM threshold. | +| `--hbm-gbs-min` | 2000 | Tier 2 HBM threshold. | +| `--rccl-size-mb` | 64 | Local RCCL message size. | +| `--rccl-gbs-min` | 100 | Local RCCL bandwidth threshold. | +| `--rccl-timeout-sec` | 120 | Hard timeout for the RCCL phase. | +| `--skip-dmesg` | off | Skip dmesg scan (e.g. inside containers). | +| `--dmesg-minutes` | 15 | dmesg `--since` window. | +| `--expected-rdma-nics N` | auto (report-only) | When set, a mismatch between the included (training-NIC) count and N becomes a node FAIL. | +| `--rdma-nic-allowlist LIST` | unset | Explicit training-NIC selector in `NCCL_IB_HCA` syntax (`device[:port],...`, `^...` denylist, `=dev` exact-match). Wins over `NCCL_IB_HCA` env. When neither is set, ports whose `phys_state` is `Disabled` / `Sleep` are auto-excluded. | +| `--ulimit-l-min-gb GB` | 32 | `RLIMIT_MEMLOCK` threshold (0 disables). | +| `--shm-min-gb GB` | 8 | `/dev/shm` size threshold (0 disables). | +| `--rocm-smi-timeout-sec SEC` | 5.0 | Hard timeout for the `rocm-smi --version` self-latency canary; hitting it is a node FAIL. | +| `--hbm-busy-threshold-gib GiB` | 2.0 | FAIL if any GPU has ≥ this much HBM in use before smoke touches the device. Boundary inclusive. | +| `--allow-foreign-procs` | off | Do NOT FAIL on foreign processes holding a GPU. They are still reported. | +| `--allowed-procs LIST` | `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter` | Process names OK to find holding the GPU. Set to `""` to disable the whitelist. | +| `--gpu-activity-warn-pct PCT` | 20.0 | Warn (does NOT fail) if any GPU's `gfx_activity_pct` exceeds this at smoke start. | +| `--require-tools LIST` | `""` (warn-only) | CLI tools that MUST be in PATH (`amd-smi`, `rocm-smi`, `lsof`); anything missing becomes a hard node FAIL. | +| `--no-clean-dump-path` | off | Do NOT auto-wipe stale per-node JSONs / aggregator outputs from `--dump-path` on rank 0 at startup. | + +### 9.3 Standalone `aggregate` flags + +The primus-cli wrapper runs the aggregator automatically on rank 0 and fills `--expected-nodes` / `--expected-nodelist-file` from SLURM. These matter only when you invoke `python -m primus.tools.preflight.node_smoke aggregate` yourself (see [§6.9](#69-re-aggregate-from-existing-per-node-jsons-no-re-run)). + +| Flag | Default | Purpose | +|---|---|---| +| `--dump-path` | `output/preflight` | Same as `run`. | +| `--expected-nodes N` | none | If fewer JSONs land within `--wait-timeout-sec`, missing nodes are added as FAIL placeholders. | +| `--wait-timeout-sec` | 60 | Polling timeout. | +| `--rocm-smi-warn-sec SEC` | 1.0 | Flag (warn-only) any node where `rocm-smi --version` took longer than this. | +| `--clock-skew-warn-sec SEC` | 30.0 | Warn when wall-clock spread across nodes exceeds this many seconds (includes srun launch jitter). | +| `--hbm-busy-threshold-gib GiB` | 2.0 | Mirrors the `run` default; labels the "GPU pre-touch HBM usage outliers" section. | +| `--gpu-activity-warn-pct PCT` | 20.0 | Mirrors the `run` default; labels the "GPU compute-activity outliers" section. | +| `--expected-nodelist-file FILE` | none | One short hostname per line. Missing nodes get their real short hostname in the report and `failing_nodes.txt`. The wrapper auto-populates this from `scontrol show hostnames "$SLURM_JOB_NODELIST"` under SLURM. | + +### 9.4 Launcher-level knobs (`primus-cli direct`) + +Consumed by `primus-cli-direct.sh` **before** the `--` separator (not forwarded to the `node_smoke` Python tool): + +| Flag | Purpose | +|---|---| +| `--silent` | Redirect launcher + tool stdout to `/dev/null`. Launcher errors (`LOG_ERROR` / `LOG_WARN` on stderr) and the log file are preserved. Exit code propagated. | +| `--debug` | Verbose launcher logging. | +| `--dry-run` | Print the resolved configuration and command without executing. | +| `--env KEY=VALUE` | Inject an env var into the Python process. | + +> **Run vs. aggregate.** The primus-cli `node_smoke` subcommand always runs the per-node checks on every rank, then aggregates on rank 0 — which is what you want ~100% of the time, so there is no `--aggregate-only` wrapper flag. For the rare single-phase cases, call the standalone CLI directly: `python -m primus.tools.preflight.node_smoke run ...` (per-node only, no report) or `... aggregate ...` (report only, from existing JSONs). + +--- + +## 10. Comparison with the full `preflight` + +| Aspect | `node_smoke` | full `preflight` | +|---|---|---| +| Rendezvous | None — every node independent | Global `torch.distributed` | +| Wall clock | ~30–60 s for 6 nodes (Tier 1+2) | Minutes; scales with N for inter-node tests | +| Granularity | Per-node PASS/FAIL | Per-rank measurements (no auto-fail by default) | +| GEMM | Hard threshold per GPU | Reports per-GPU numbers, no auto-fail | +| HBM bandwidth | Yes (D2D `copy_`) | Not measured | +| Inter-node all-reduce / all-to-all | Not tested (intentionally) | Yes | +| Drift detection | Yes (versions, NIC firmware, port count) | No | +| Host limits / RDMA roll-call | Yes (hard fail) | Reported via `collect_*_info` only | +| Output format | Per-node JSON + cluster md + SLURM-ready txt | Markdown + PDF | + +Use `node_smoke` to **screen** a cluster fast and exclude bad nodes. Use the full [`preflight`](./preflight.md) when you want **deep cross-node measurements** (inter-node bandwidth matrix, ring-P2P, etc.). The recommended sequence is node-smoke first, then `preflight --quick` on the surviving nodes. + +--- + +## 11. Troubleshooting + +| Symptom | Likely cause / fix | +| ----------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `[ERROR] [direct] VENV_ACTIVATE is set but file does not exist: ...` | Fix the path (`export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate`), or `unset VENV_ACTIVATE` to fall back to system / container Python. | +| Every node FAILs with `gpu_processes: ... name='N/A'` | Container `/proc` can't resolve host PID names. See [§6.5](#65-allow--extend-the-foreign-process-whitelist): pass `--allow-foreign-procs`, or grant host-PID visibility. | +| Some nodes never produce a JSON | The aggregator names them in `failing_nodes.txt` via `expected_nodes.txt`. If `scontrol` was unavailable, they appear as ``. | +| Tier 2 perf numbers below threshold on a known-good node | Almost always insufficient CPU cores on `srun` — pass `-c ` so RCCL proxy threads have CPU. | +| Re-run on a smaller nodelist still shows the previously removed nodes as PASS | Default behavior cleans stale JSONs on rank 0. If you passed `--no-clean-dump-path`, either remove it or `rm -rf output/preflight` between runs. | + +--- + +## 12. See also + +- [`preflight.md`](./preflight.md) — the heavier `preflight` tool with a global rendezvous and inter-node bandwidth tests. +- [`preflight-without-container.md`](./preflight-without-container.md) — running `preflight` directly on the host (no container), including the shared venv + NCCL setup. +- [`primus/cli/subcommands/node_smoke.py`](../../primus/cli/subcommands/node_smoke.py) — the primus-cli subcommand wiring (two-phase dispatch: per-rank run + rank-0 aggregate). +- [`primus/tools/preflight/node_smoke/cli.py`](../../primus/tools/preflight/node_smoke/cli.py) — canonical flag definitions and per-node / aggregate phase bodies. diff --git a/docs/02-user-guide/preflight-without-container.md b/docs/02-user-guide/preflight-without-container.md new file mode 100644 index 000000000..0494fe2c7 --- /dev/null +++ b/docs/02-user-guide/preflight-without-container.md @@ -0,0 +1,792 @@ +# Run Preflight Without a Container + +> ⚠ **Run the [node-smoke test](./node-smoke-test-instruction.md) first.** `preflight` opens a global `torch.distributed` rendezvous, so a single sick node (wedged driver, leaked rank holding HBM, partial NIC enumeration, time-sync drift, etc.) can stall the whole job for up to `--dist-timeout-sec` seconds — long before any cross-node bandwidth number is produced. The node-smoke test catches those exact failure modes *without* a rendezvous in ~30–60 s and emits a SLURM-ready `failing_nodes.txt` you can pipe straight into `srun --exclude=`. Treat node-smoke as a hard prerequisite; only run `preflight` on the nodes node-smoke marked PASS. See [§0 "Which test should I run?"](#0-which-test-should-i-run) for the side-by-side comparison and the recommended 3-step workflow. + +This guide explains how to run Primus's [`preflight`](./preflight.md) cluster-diagnostic tool **directly on the host** (no Docker / Podman), via the standard Primus launcher. + +**Git clone the Primus repository to a shared filesystem that all nodes can read.** + +```bash +git clone --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.git +cd Primus +``` + +**Recommended (through the primus-cli SLURM wrapper):** + +For some clusters, you may need to explicitly request CPU and GPU resources with `srun -N -c --gpus-per-node=`. + +``` +runner/primus-cli slurm srun -N --ntasks-per-node=1 -- direct -- preflight [PREFLIGHT_ARGS...] +``` + +**Equivalent (bare srun, useful when composing with custom srun flags):** + +``` +srun -N --ntasks-per-node=1 runner/primus-cli direct -- preflight [PREFLIGHT_ARGS...] +``` + +Both forms produce the **same workload** on the same ranks. The wrapper form is recommended because it auto-resolves `MASTER_ADDR` / `MASTER_PORT` / `NNODES` / `NODE_RANK` / `GPUS_PER_NODE` once on the launching node and passes them to every rank via `--env`, applies any `slurm.`* config defaults (partition / time / etc.) from your YAML, and is the same pattern used for `train` / `benchmark` / `node_smoke`. See [§ Wrapper vs. bare-srun](#wrapper-vs-bare-srun) below for the exact precedence / caveats. + +`primus-cli direct` activates an optional Python virtualenv (`VENV_ACTIVATE`), auto-derives the distributed environment variables (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, `MASTER_PORT`, `GPUS_PER_NODE`) from `SLURM_*` when running inside a SLURM allocation, and then launches the `preflight` Python subcommand via `torchrun` (one worker per GPU). It is the recommended entry point when: + +- You're running on a SLURM cluster but cannot (or don't want to) use the container-based path. +- Your nodes share a Python virtual environment on a network-mounted filesystem. +- You want a single-node sanity check with no extra configuration. + +--- + +## 0. Which test should I run? + +Primus ships **two** complementary cluster screens. Pick the right one — and ideally run them in this order. + + +| Aspect | `node-smoke` (start here) | `preflight` (this doc) | +| ----------------- | ---------------------------------------------------------------------------------- | ------------------------------------------------------------------- | +| Purpose | "Which nodes are healthy enough to run anything?" | "What is the actual cross-node performance on the surviving nodes?" | +| Rendezvous | None — every node independent | Global `torch.distributed` rendezvous | +| Wall clock | ~30–60 s for 6 nodes (Tier 1+2) | A few minutes; scales with N for inter-node tests | +| Granularity | Per-node PASS/FAIL | Per-rank perf measurements | +| Safety | A stuck node cannot wedge its peers | A single hung NIC can stall the whole rendezvous | +| Output | Per-node JSON + cluster md + SLURM-ready `passing_nodes.txt` / `failing_nodes.txt` | Markdown + PDF perf report | +| Entry point | `primus-cli direct -- node_smoke` | `primus-cli direct -- preflight` (this doc) | +| Quick-start guide | [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) | This doc, §3+ | + + +### Recommended workflow + +> **Before running any of the commands below, complete the one-time setup:** +> +> 1. **Python virtualenv** on a shared filesystem — see [§2 Set up the Python virtual environment](#2-set-up-the-python-virtual-environment), then point the launcher at it via `export VENV_ACTIVATE=...` (details in [§2 → Tell the launcher where the venv is](#tell-the-launcher-where-the-venv-is)). +> 2. **NCCL / fabric environment variables** — usually the defaults in `base_env.sh` are fine, but multi-NIC nodes may need `NCCL_IB_HCA` / `NCCL_IB_GID_INDEX` / `NCCL_SOCKET_IFNAME` overrides. See [§4 Cluster-specific NCCL configuration](#4-cluster-specific-nccl-configuration) for known-good values per fabric (Broadcom, Pensando Pollara/AINIC). + +Through the `primus-cli slurm srun -- direct --` wrapper (recommended): + +```bash +# 1) Prune broken nodes with node-smoke (fast, no rendezvous). +runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + -- direct -- node_smoke --tier2-perf + +# 2) Re-allocate excluding the bad nodes, and run preflight --quick +# for a fast cross-node sanity check. +runner/primus-cli slurm srun -N -c 128 --gpus-per-node=8 \ + --ntasks-per-node=1 \ + --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \ + -- direct -- preflight --quick + +# 3) Optional: full preflight on the same set if --quick numbers +# look off, or if you want the full bandwidth matrix. +runner/primus-cli slurm srun -N -c 128 --gpus-per-node=8 \ + --ntasks-per-node=1 \ + --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \ + -- direct -- preflight +``` + +Equivalent with bare `srun` (works identically; useful when scripting around custom srun flags that don't compose with the wrapper): + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf + +srun -N -c 128 --gpus-per-node=8 --ntasks-per-node=1 \ + --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \ + runner/primus-cli direct -- preflight --quick + +srun -N -c 128 --gpus-per-node=8 --ntasks-per-node=1 \ + --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \ + runner/primus-cli direct -- preflight +``` + +Why this ordering matters: + +- A single broken node can stall a `torch.distributed.init_process_group()` for `--dist-timeout-sec` seconds (default 120), so feeding a known-good list to preflight is much faster. +- `node-smoke` catches things preflight cannot — leaked / foreign processes, wedged drivers, partial NIC enumeration, time-sync drift, RDMA roll-call issues — that produce *misleading* preflight failures. +- `preflight --quick` adds the cross-node bandwidth signal that `node-smoke` deliberately does not measure. + +--- + +## 1. Prerequisites + +- A working AMD ROCm installation on every node. +- Network reachability between nodes (Ethernet for bootstrap, RDMA / InfiniBand recommended for perf tests). +- A Python ≥ 3.10 virtual environment **on a shared filesystem** that all nodes can read (the same path is sourced on every node). +- The Primus repository checked out somewhere readable from every node. + +--- + +## 2. Set up the Python virtual environment + +The environment must live on a path visible from every node (e.g. NFS-mounted home, Lustre, or any shared filesystem). All nodes will `source` the same activation script. + +You can use any tool you like; `uv` is the fastest. Either of the following works. + +### What you actually need to install + +The `preflight` and `node-smoke` tools deliberately use **only a small subset** of Primus's full dependency tree. You do **not** need to install the entire `requirements.txt` — that pulls in trainer / dataset / experiment-tracking packages that neither tool ever imports. + + +| Package | Required for | Skip when | +| -------------------- | ------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ | +| `torch` (ROCm build) | Both tools — perf measurements (`torch.matmul`, `torch.distributed`, `torch.cuda.`*). | Never (mandatory). | +| `markdown2` | `preflight` PDF report only (Markdown → HTML). | You always pass `--disable-pdf`, or you only run `node-smoke` (which never produces PDFs). | +| `weasyprint` | `preflight` PDF report only (HTML → PDF). | Same as above. | +| `matplotlib` | `preflight --plot` only (per-test bandwidth bar charts). | You don't pass `--plot`. | + + +Everything else in the preflight / node-smoke code path is Python stdlib (`os`, `subprocess`, `socket`, `argparse`, `dataclasses`, `json`, `time`, ...) — no extra installs needed. + +### Option A — `uv` (recommended), minimal install + +```bash +mkdir -p ~/envs/preflight +cd ~/envs/preflight + +uv venv --python 3.12 +source .venv/bin/activate + +# 1) ROCm-built PyTorch (pin to your ROCm version; rocm7.1 shown here) +uv pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.1 --no-cache-dir + +# 2) Optional: only if you want preflight PDF reports (omit to use --disable-pdf) +uv pip install markdown2 weasyprint + +# 3) Optional: only if you want preflight --plot bar charts +uv pip install matplotlib +``` + +### Option B — `python -m venv`, minimal install + +```bash +mkdir -p ~/envs/preflight +python3.12 -m venv ~/envs/preflight/.venv +source ~/envs/preflight/.venv/bin/activate + +pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.1 --no-cache-dir +pip install markdown2 weasyprint # optional, for preflight PDFs +pip install matplotlib # optional, for preflight --plot +``` + +### Option C — full Primus runtime (only if you also want the rest of Primus) + +```bash +cd /path/to/Primus +uv pip install -r requirements.txt # or: pip install -r requirements.txt +``` + +This installs every Primus runtime dependency (trainer, dataset loaders, experiment trackers, ...). Use only if you're going to run more than just preflight / node-smoke from this environment. + +### Per-tool minimum install matrix + +If you want the absolute smallest footprint, install only what your intended invocations need: + + +| Invocation | `torch` | `markdown2` | `weasyprint` | `matplotlib` | +| ------------------------------------------------ | -------- | --------------------------------- | --------------------------------- | ------------ | +| `node-smoke` (any flags) | required | — | — | — | +| `preflight --host --gpu --network --disable-pdf` | required | — | — | — | +| `preflight --host --gpu --network` (with PDF) | required | required | required | — | +| `preflight --quick --disable-pdf` | required | — | — | — | +| `preflight --quick` (with PDF) | required | required | required | — | +| `preflight ... --plot` | required | required (unless `--disable-pdf`) | required (unless `--disable-pdf`) | required | + + +### Tell the launcher where the venv is + +`primus-cli direct` reads the `VENV_ACTIVATE` environment variable. When set, it sources the path before launching the Python process; when unset, it is a no-op (the container path, which uses the container's bundled Python, leaves this unset): + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate +``` + +`VENV_ACTIVATE` is the only optional environment variable specific to the direct flow. Everything else has a sensible default; distributed-env variables (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, ...) are auto-derived from SLURM when not pre-exported. + +--- + +## 3. Run preflight + +### Single node (no SLURM) + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate + +# Info report only (fast) +runner/primus-cli direct -- preflight --host --gpu --network + +# Info + perf report +runner/primus-cli direct -- preflight + +# Perf report only +runner/primus-cli direct -- preflight --perf-test +``` + +When SLURM is not detected the script defaults to `NNODES=1`, `NODE_RANK=0`, `MASTER_ADDR=localhost`. Any of those can be overridden by exporting them before calling the script. + +### Multi-node without SLURM (parallel SSH) + +When no scheduler is available (bare-metal, cloud VMs, lab nodes), launch +`primus-cli direct` on each node yourself via SSH. The script works +identically — you just pre-export the distributed variables that SLURM +would normally provide. + +#### Requirements + +- All nodes share the same filesystem (or at least the same Primus checkout + venv path). +- Nodes can reach each other on a **data-plane** network interface (not the management NIC). +- SSH key-based access to each node from the launching host. + +#### Required environment variables + + +| Variable | Description | +| -------------------- | ---------------------------------------------------------------------- | +| `NNODES` | Total number of nodes | +| `NODE_RANK` | This node's rank (`0` through `NNODES-1`) | +| `MASTER_ADDR` | IP of rank-0 node **on the data-plane interface** | +| `MASTER_PORT` | Rendezvous port (default `1234`; increment between concurrent runs) | +| `GPUS_PER_NODE` | GPUs per node (default `8`) | +| `NCCL_SOCKET_IFNAME` | Data-plane NIC name (e.g. `enp159s0np0`) — **critical for multi-node** | +| `GLOO_SOCKET_IFNAME` | Same as `NCCL_SOCKET_IFNAME` | +| `VENV_ACTIVATE` | Path to virtualenv `activate` script | + + +> **Warning**: `NCCL_SOCKET_IFNAME` auto-detection often picks a management interface +> (e.g. `enp28s0np0`, `eno8303`) instead of the high-bandwidth data NIC. For multi-node +> runs this causes `init_process_group` to hang or NCCL to fail silently. Always set it +> explicitly. + +#### Identifying the correct data-plane interface + +```bash +# On any node, find the interface whose IP matches the MASTER_ADDR subnet: +ip -4 addr show | grep "10.245.134" +# → enp159s0np0 inet 10.245.134.129/24 + +# Or check which interface routes to the master: +ip route get 10.245.134.129 | awk '{print $5; exit}' +``` + +### Multi-node via SLURM + +`primus-cli direct` auto-detects a SLURM allocation (via `SLURM_JOB_ID`) and derives all distributed variables from `SLURM_*` automatically. **Pre-exported values always win**, so the same launcher script also works inside the `primus-cli slurm srun ... -- direct -- ...` chain (where `slurm-entry` has already set these via `--env`): + +| Variable | Resolved as | +| --------------- | -------------------------------------------------------------------- | +| `NNODES` | `NNODES` → `SLURM_NNODES` → `SLURM_JOB_NUM_NODES` → `1` | +| `NODE_RANK` | `NODE_RANK` → `SLURM_NODEID` → `SLURM_PROCID` → `0` | +| `MASTER_ADDR` | `MASTER_ADDR` (if not empty / not `localhost`) → first hostname from `scontrol show hostnames "$SLURM_NODELIST"` | +| `MASTER_PORT` | `MASTER_PORT` → `1234` | +| `GPUS_PER_NODE` | `GPUS_PER_NODE` → `8` | + +Run it as a single task per node (the script invokes `torchrun` internally, which spawns one worker per GPU): + +> **Verify NCCL / network env first.** The script sets sensible `NCCL_`* defaults via `base_env.sh`, but auto-detection can pick the wrong device on multi-NIC nodes. Always confirm `NCCL_IB_HCA`, `NCCL_IB_GID_INDEX`, `NCCL_SOCKET_IFNAME`, and `GLOO_SOCKET_IFNAME` (set to the same value as `NCCL_SOCKET_IFNAME`) are correct for your fabric, and `export` overrides before running. See [§4](#4-cluster-specific-nccl-configuration) for cluster-specific values. + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate + +# export NCCL_IB_HCA=rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 +# export NCCL_IB_GID_INDEX=3 +# export NCCL_SOCKET_IFNAME=eno0 +# export GLOO_SOCKET_IFNAME=eno0 + +# Recommended: through the primus-cli SLURM wrapper. +runner/primus-cli slurm srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 \ + --nodelist --ntasks-per-node=1 \ + -- direct -- preflight --perf-test + +# Or, equivalently, with bare srun: +srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist \ + --ntasks-per-node=1 \ + runner/primus-cli direct -- preflight --perf-test +``` + + + +#### Wrapper vs. bare-srun + +Both forms target the **same** `primus-cli-direct.sh` launcher and produce identical workloads. The difference is only in how the SLURM context is constructed: + + +| Aspect | `primus-cli slurm srun -- direct --` (recommended) | Bare `srun ... primus-cli direct --` | +| ---------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- | +| `MASTER_ADDR` resolution | Resolved **once** on the launching node via `scontrol show hostnames "$SLURM_NODELIST" | head -n1`, then propagated to every rank via `--env MASTER_ADDR=...`. | Each rank re-derives it inside `primus-cli-direct.sh` STEP 4.7 from `SLURM_`* (same result, more `scontrol` calls). | +| `NNODES` / `NODE_RANK` / `GPUS_PER_NODE` | Set explicitly by `slurm-entry.sh` via `--env`. | Derived from `SLURM_NNODES` / `SLURM_NODEID` / `SLURM_PROCID` inside `direct.sh`. | +| `slurm.*` config defaults | Honored (partition, time, ntasks-per-node, etc. from the active YAML). | Not consulted — you pass every flag explicitly to `srun`. | +| Default wall-time | `-t 4:00:00` is auto-added if you don't pass `--time`. | None — `srun` uses the cluster default (may reject the job). | +| `direct` keyword | **Required**: `primus-cli slurm srun ... -- direct -- `. Without `direct`, the wrapper routes through the **container** path. | N/A — there's only one path. | +| `--ntasks-per-node=1` | **Not auto-added**. Pass it on the CLI (before the first `--`) or set it in the `slurm.`* config. | **Not auto-added**. Pass it as an `srun` flag. | +| Best for | Production / repeatable runs. Same pattern as `train` / `benchmark` / `node_smoke`. | Ad-hoc runs where you want to compose with arbitrary `srun` flags (`--nodelist=$(...)`, `--exclude=...` from a runtime file, etc.). | + + +For the rest of this doc the examples use bare `srun` for brevity, but every example also works with the wrapper form by substituting `srun runner/primus-cli direct --` → `runner/primus-cli slurm srun -- direct --`. + +### Key `srun` flags + + +| Flag | Why it's necessary | +| --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `-c 128` | Allocate all CPU cores per task. Without this, SLURM may default to 1 core, which starves the RCCL network proxy threads and can cause >30× slowdown on perf tests. Set this to your node's core count. | +| `--gpus-per-node=8` | Grants GPU device access (`/dev/kfd`, `/dev/dri`). Required for non-container execution. | +| `--ntasks-per-node=1` | One launcher invocation per node; `primus-cli direct` then spawns 8 workers per node via `torchrun`. | +| `-t 00:45:00` | Wall-clock limit. Full perf tests on 8N usually finish well under 10 min. | + + +> Tip — check core count: `srun -N 1 --gpus-per-node=8 bash -c 'nproc'` + +--- + +## 4. Cluster-specific NCCL configuration + +`primus-cli direct` sources `runner/helpers/envs/base_env.sh`, which sets sensible defaults for `NCCL_`* and auto-detects `NCCL_IB_HCA` / `NCCL_SOCKET_IFNAME`. Pre-exported values from your shell take precedence, so the standard pattern is: + +```bash +export VAR=value +runner/primus-cli direct -- preflight ... +``` + +### Broadcom NICs (no AINIC) + +Most clusters fall here. The defaults from `base_env.sh` are usually fine, but the two values most commonly worth overriding are: + +```bash +export NCCL_CROSS_NIC=1 # default in base_env.sh is 0 +export NCCL_PXN_DISABLE=0 # default in base_env.sh is 1 + +srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist \ + --ntasks-per-node=1 \ + runner/primus-cli direct -- preflight --perf-test +``` + +### Pensando Pollara (AINIC) RDMA + +```bash +export USING_AINIC=1 +export NCCL_IB_GID_INDEX=1 # AINIC uses index 1 (default in base_env.sh is 3) +export NCCL_PXN_DISABLE=0 + +srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist \ + --ntasks-per-node=1 \ + runner/primus-cli direct -- preflight +``` + +> `primus-cli direct` *does* accept `--env KEY=VALUE` on its own command line (placed before `--`), in addition to the conventional `export`/`srun --export=` approaches. + +--- + +## 5. Launcher flags vs. preflight flags + +Anything you place **after** the `--` separator is forwarded verbatim to the `preflight` Python tool. The launcher (`primus-cli-direct.sh`) consumes a small set of flags **before** `--`. The one most users care about is `--silent`. + +### Launcher-only flags (before `--`) + + +| Flag | Effect | +| ----------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `--silent` | Back-pocket knob: redirect the launcher's and the Python tool's `stdout` to `/dev/null`. Launcher errors (`LOG_ERROR` / `LOG_WARN`, written to `stderr`) are preserved so real failures still surface; the log file under `logs/` captures everything. Exit code is propagated unchanged. **Not recommended** for normal use — you lose live progress; prefer the log file. | +| `--debug` | Verbose launcher logging (`PRIMUS_LOG_LEVEL=DEBUG`). Forwarded to the Python tool as `--debug` too. | +| `--dry-run` | Print the resolved configuration and final `torchrun` / `python3` command without executing. | +| `--single` | Force `python3` instead of `torchrun`. `node_smoke` auto-selects this; for `preflight` you usually want the default (`torchrun`). | +| `--env KEY=VALUE` | Inject an env var into the Python process (in addition to anything `export`-ed in the shell). | +| `--log_file PATH` | Redirect the captured tee log to a specific path (default: `logs/log_.txt`). | + + +See `runner/primus-cli direct --help` for the full set. + +### Forwarded `preflight` flags (after `--`, most common) + +See [Preflight](./preflight.md) for the full list. The most common are: + +- Mode selection: `--host`, `--gpu`, `--network`, `--perf-test`, `--tests`, `--quick` +- Test tuning: `--comm-sizes-mb`, `--intra-comm-sizes-mb`, `--inter-comm-sizes-mb`, `--intra-group-sizes`, `--inter-group-sizes`, `--ring-p2p-sizes-mb` +- Reporting: `--dump-path`, `--report-file-name`, `--disable-pdf`, `--plot` +- Reliability: `--comm-cleanup-delay-sec`, `--dist-timeout-sec` + +If you do not pass `--report-file-name`, `preflight` auto-generates a unique one of the form: + +``` +preflight-${NNODES}N-YYYYMMDD-HHMMSS +``` + +This guarantees that each run lands in its own files and prevents stale leftovers from earlier runs from being mistaken for fresh output. The auto-name logic now lives in the Python tool itself, so every call site (host `srun ... primus-cli direct`, `primus-cli slurm ... -- direct`, `primus-cli slurm ... -- container`) gets the same fresh name. + +### Examples + +The examples below all assume one of the two equivalent shell-prefix conventions. Pick whichever matches your habits — every example block in this section works with either definition: + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate + +# Recommended: through the primus-cli SLURM wrapper. Auto-resolves +# MASTER_ADDR/NNODES/NODE_RANK once on the launching node and propagates +# them via --env; honors slurm.* config defaults. +SRUN="runner/primus-cli slurm srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --ntasks-per-node=1 --nodelist --" +# Then in every example below, replace `$SRUN runner/primus-cli direct --` +# with just `$SRUN direct --`. (The wrapper expects the entry-mode keyword +# `direct` as the first token after the inner `--`.) + +# Equivalent: bare srun. NNODES/NODE_RANK/MASTER_ADDR get derived inside +# primus-cli-direct.sh's STEP 4.7 directly from SLURM_*; same net effect. +SRUN="srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --ntasks-per-node=1 --nodelist " +``` + +The examples in this section use the **bare-srun** form below for brevity (since `$SRUN runner/primus-cli direct -- preflight` reads naturally as one command line). To use the wrapper form instead, substitute `$SRUN runner/primus-cli direct --` → `$SRUN direct --` after exporting `SRUN` to the wrapper variant. + +#### A. Mode selection + +```bash +# Default: info report + every perf test +$SRUN runner/primus-cli direct -- preflight + +# Info-only (fast, no torch.distributed rendezvous) +$SRUN runner/primus-cli direct -- preflight --host --gpu --network --disable-pdf + +# Perf-only, every test +$SRUN runner/primus-cli direct -- preflight --perf-test + +# Fast pre-launch sanity preset (gemm + intra-AR + inter-AR @ 64,1024 MB, +# full intra-node group, full N-node inter group, low warmup/iter) +$SRUN runner/primus-cli direct -- preflight --quick +``` + +> **Note**: Mixing perf-mode flags (`--perf-test` / `--tests` / `--quick`) with info selectors (`--host` / `--gpu` / `--network`) makes preflight drop the info selectors with a `WARN`. Run two invocations if you want both reports. + +#### B. Test selection (`--tests`) + +```bash +# Only GEMM +$SRUN runner/primus-cli direct -- preflight --tests gemm + +# Only the inter-node bandwidth tests +$SRUN runner/primus-cli direct -- preflight --tests inter-allreduce,inter-alltoall + +# Only the inter-node ring P2P +$SRUN runner/primus-cli direct -- preflight --tests inter-ring-p2p + +# Combine: GEMM + inter-AR with overridden sizes/groups +$SRUN runner/primus-cli direct -- preflight \ + --tests gemm,inter-allreduce \ + --comm-sizes-mb 64,1024 \ + --inter-group-sizes all +``` + +Valid `--tests` tokens: `gemm`, `intra-allreduce`, `intra-alltoall`, `inter-allreduce`, `inter-alltoall`, `inter-p2p`, `inter-ring-p2p`, `all`. Unknown tokens fail fast (before NCCL init). + +#### C. Message sizes + +```bash +# One CSV applied to both intra and inter +$SRUN runner/primus-cli direct -- preflight --tests intra-allreduce,inter-allreduce \ + --comm-sizes-mb 8,128 + +# Different sizes for intra vs inter (override wins over --comm-sizes-mb) +$SRUN runner/primus-cli direct -- preflight --tests intra-allreduce,inter-allreduce \ + --comm-sizes-mb 8,128 --intra-comm-sizes-mb 4,32 + +# Inter-only override (also covers inter-p2p when enabled) +$SRUN runner/primus-cli direct -- preflight --tests inter-allreduce,inter-p2p \ + --comm-sizes-mb 8,128 --inter-comm-sizes-mb 16,512 +``` + +#### D. Group sizes + +```bash +# Custom intra-node group sizes (each must divide LOCAL_WORLD_SIZE) +$SRUN runner/primus-cli direct -- preflight \ + --tests intra-allreduce \ + --intra-group-sizes 4,8 + +# Custom inter-node groups: 2-node pairs and the full N-node group +$SRUN runner/primus-cli direct -- preflight \ + --tests inter-allreduce \ + --inter-group-sizes 2,all +``` + +> Note: for `inter-alltoall` only, every requested per-group node count is internally capped at **16** (real-world MoE training rarely dispatches across more nodes; see [`preflight.md` §5.2](./preflight.md#52-group-sizes) for the rationale). The other inter-node tests use the requested sizes unchanged. So on a 128-node cluster, `--tests inter-alltoall --inter-group-sizes all` actually runs at 16-node sub-groups, while `--tests inter-allreduce --inter-group-sizes all` runs at 128 nodes as written. + +#### E. Ring P2P sizes + +```bash +$SRUN runner/primus-cli direct -- preflight \ + --tests inter-ring-p2p \ + --ring-p2p-sizes-mb 5,20,80 +``` + +#### F. Plotting + +```bash +# Generate per-test bandwidth bar charts under //*.png +$SRUN runner/primus-cli direct -- preflight \ + --tests intra-allreduce,inter-allreduce --plot +``` + +#### G. Reliability knobs + +```bash +# Bump the per-phase cleanup delay. Default 2.0 is sufficient at every +# cluster size for the comm shapes preflight exercises (inter-alltoall +# is internally capped at 16 nodes; see preflight.md §5.2). Only bump +# this on very flaky networks or unusual kernel TIME_WAIT settings. +$SRUN runner/primus-cli direct -- preflight --quick --comm-cleanup-delay-sec 5 + +# Fail fast if torch.distributed rendezvous can't complete in 30s +$SRUN runner/primus-cli direct -- preflight --perf-test --dist-timeout-sec 30 +``` + +> Operating clusters at ≥ 128 nodes? See [`preflight.md` §7](./preflight.md#7-running-on-very-large-clusters--64-nodes) for the recommended OS-level tuning (widening `ip_local_port_range`) and per-test invocation patterns. With the §5.2 inter-alltoall cap in place, a default invocation is safe at every cluster size; the §7.3 OS tuning remains best-practice for any RDMA host. + +#### H. Reporting & output layout + +```bash +# Quick info-only check on 4 nodes, no PDF +$SRUN runner/primus-cli direct -- preflight --host --gpu --network --disable-pdf \ + --report-file-name info-4N + +# Perf test only, silenced (CI-friendly), explicit name. Note that --silent +# is consumed by primus-cli-direct.sh and must appear BEFORE the `--` +# separator; everything after `--` is forwarded to the preflight Python tool. +$SRUN runner/primus-cli direct --silent -- preflight --perf-test \ + --report-file-name nightly-4N-perf + +# Archive each run under its own directory +$SRUN runner/primus-cli direct -- preflight --quick \ + --dump-path /shared/preflight-archive/$(date +%Y%m%d-%H%M%S) +``` + +#### I. Backward-compat aliases + +These still work and are equivalent to flags above. Use them only when retrofitting older scripts. + +```bash +# Same as --host --gpu --network +$SRUN runner/primus-cli direct -- preflight --check-host --check-gpu --check-network + +# Same as --inter-group-sizes all AND drops inter-p2p +$SRUN runner/primus-cli direct -- preflight --perf-test --no-split-nodes-subgroup +``` + +#### J. Combined "production-ready" pre-launch screen + +```bash +# 1) Smoke first to prune broken nodes (note: --silent goes BEFORE `--`) +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct --silent -- node_smoke --tier2-perf + +# 2) Quick perf sanity on the survivors +srun -N -c 128 --gpus-per-node=8 --ntasks-per-node=1 \ + --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \ + runner/primus-cli direct --silent -- preflight --quick \ + --comm-cleanup-delay-sec 5 --dist-timeout-sec 60 \ + --report-file-name screen-$(date +%Y%m%d-%H%M%S) +``` + +--- + +## 6. Outputs + +Reports are written to `--dump-path` (default: `output/preflight/`), with the basename from `--report-file-name` and a `_perf` suffix for performance reports: + + +| File | Produced by | Notes | +| ----------------- | ----------------------------------------------- | ------------------------- | +| `.md` | `--host --gpu --network` (or default selection) | Info report | +| `.pdf` | same, unless `--disable-pdf` | Info report PDF | +| `_perf.md` | `--perf-test` | Perf report (GEMM + comm) | +| `_perf.pdf` | same, unless `--disable-pdf` | Perf report PDF | + + +Only **rank 0** writes the report. After preflight completes, the Python tool prints the absolute path of every report file it produced to stdout. Under `--silent` these prints go to `/dev/null` along with everything else (one of the trade-offs of using `--silent`); without `--silent` the announcement is visible live. Sample output: + +``` +[Primus:Preflight] Report: /home/.../Primus/output/preflight/preflight-4N-20260428-201925.md +[Primus:Preflight] Report: /home/.../Primus/output/preflight/preflight-4N-20260428-201925_perf.md +``` + +--- + +## 7. Environment variable reference + +Variables read by `primus-cli direct` itself: + + +| Variable | Required | Default | Purpose | +| --------------- | -------- | ---------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------- | +| `VENV_ACTIVATE` | no | — | Path to the venv `bin/activate` script. Unset = no-op (use system / container Python). Set + missing file = fail-fast. | +| `NNODES` | no | `1` (or auto-derived from `SLURM_NNODES` / `SLURM_JOB_NUM_NODES`) | Number of nodes. Pre-exported always wins. | +| `NODE_RANK` | no | `0` (or auto-derived from `SLURM_NODEID` / `SLURM_PROCID`) | This node's rank. Pre-exported always wins. | +| `GPUS_PER_NODE` | no | `8` | GPUs per node | +| `MASTER_ADDR` | no | `localhost` (or first host from `scontrol show hostnames "$SLURM_NODELIST"`) | Rendezvous host. Pre-exported always wins. | +| `MASTER_PORT` | no | `1234` | Rendezvous port | + + +Variables consumed downstream by `primus-cli direct` / `base_env.sh` (set them via `export`): + + +| Variable | Default in `base_env.sh` | When to override | +| -------------------- | ------------------------ | ------------------------------------------------- | +| `NCCL_SOCKET_IFNAME` | auto-detected | Force a specific Ethernet interface for bootstrap | +| `NCCL_IB_HCA` | auto-detected | Force specific RDMA HCAs | +| `NCCL_IB_GID_INDEX` | `3` | `1` on AINIC clusters | +| `NCCL_CROSS_NIC` | `0` | `1` for multi-rail IB fabrics | +| `NCCL_PXN_DISABLE` | `1` | `0` to enable PXN multi-hop NIC sharing | +| `USING_AINIC` | unset | `1` on Pensando Pollara clusters | +| `NCCL_DEBUG` | unset | `INFO` for verbose NCCL logging | + + +--- + +## 8. Troubleshooting + +### `[ERROR] [direct] VENV_ACTIVATE is set but file does not exist: ...` + +`VENV_ACTIVATE` was set in the environment but the path it points at doesn't exist on this node. This is a fail-fast guard to prevent a silent fallback to system Python (which usually has the wrong `torch` / no ROCm). Either fix the path: + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate +``` + +… or unset it to fall back to the container / system Python: + +```bash +unset VENV_ACTIVATE +``` + +If the path looks right but the file still appears missing, confirm the venv lives on a filesystem visible from the node SLURM scheduled you onto. + +### `[Primus:Preflight] FAIL: No GPUs detected` + +The Python process inside the venv can't find ROCm. Diagnose with: + +```bash +srun --nodes=1 --nodelist= bash -c ' +echo "=== PATH ==="; echo $PATH +echo "=== LD_LIBRARY_PATH ==="; echo $LD_LIBRARY_PATH +echo "=== rocm-smi ==="; rocm-smi --showid 2>&1 +echo "=== Python torch check ===" +source ~/envs/preflight/.venv/bin/activate +python3 -c "import torch; print(\"hip:\", torch.version.hip); print(\"available:\", torch.cuda.is_available()); print(\"count:\", torch.cuda.device_count())" +' +``` + +If `LD_LIBRARY_PATH` is empty, set it explicitly: + +```bash +export LD_LIBRARY_PATH=/opt/rocm/lib:${LD_LIBRARY_PATH:-} +``` + +### Report announcement points at stale files + +This shouldn't happen with the current Python tool — the auto-generated unique report name (`preflight-${NNODES}N-`) ensures every run gets a fresh path. If you explicitly pass `--report-file-name X`, you're responsible for choosing a name that doesn't collide with prior runs. + +### Slow perf tests (~30× expected) + +Almost always a symptom of insufficient CPU cores. Pass `-c ` to `srun` so RCCL's network proxy threads have CPU to spawn on. Verify with `srun -N 1 --gpus-per-node=8 bash -c 'nproc'`. + +### Using `conda` instead of venv + +`primus-cli direct` does `source "$VENV_ACTIVATE"`, which works for venv/uv but not directly for conda. Two options: + +1. Create a venv inside the conda env and point `VENV_ACTIVATE` at that venv's activate script. +2. Write a small shim activate script (e.g. `~/envs/conda-shim.sh`) that activates conda and the desired env, then point `VENV_ACTIVATE` at it: + ```bash + # ~/envs/conda-shim.sh + source "$HOME/miniconda3/etc/profile.d/conda.sh" + conda activate + ``` + +### "Address already in use" during perf tests + +This means a node's kernel ephemeral-port pool was momentarily exhausted while preflight was building many communicators in a short window, so an outgoing `bind()` could not find a free port. It is a preflight-specific artifact of repeated communicator setup/teardown — a real training job builds its communicators once and reuses them — not a training failure mode. See [`preflight.md` §7](./preflight.md#7-running-on-very-large-clusters--64-nodes) for details. + +Preflight has two complementary defenses: + +1. The **inter-node alltoall sub-group is internally capped at 16 nodes** (see [`preflight.md` §5.2](./preflight.md#52-group-sizes)) — the only test that, at large scale, opens enough simultaneous connections to approach the per-node ephemeral-port pool. The cap eliminates this failure mode by construction. +2. A **global barrier + `--comm-cleanup-delay-sec` sleep** (default 2 s) is inserted after every comm destroy, primarily for cross-rank synchronization across the destroy → setup transition. + +If you still see `Address already in use` (e.g. on a network with an unusually narrow ephemeral-port range), the directly relevant **OS-level tuning** is widening that range — best-practice for any RDMA host: + +```bash +# Widen the ephemeral port range from ~28k to ~64k. This is the +# OS knob that directly addresses the binding constraint. +sudo sysctl -w net.ipv4.ip_local_port_range="1024 65535" +``` + +As a fallback, raise the per-phase delay: + +```bash +# Bump the per-phase delay (default 2 s) on a particularly stressed +# network. Rarely needed in practice with the §5.2 alltoall cap. +runner/primus-cli direct -- preflight --comm-cleanup-delay-sec 5 +``` + +See [`preflight.md` §7](./preflight.md#7-running-on-very-large-clusters--64-nodes) for persistence and recommended large-cluster invocation patterns (split tests into separate runs, etc.). + +If the error occurs at `init_process_group` (before tests even start), it typically means a previous job left the rendezvous port (`MASTER_PORT`, default `1234`) in `TIME_WAIT`. Either wait ~60 s or use a different port: + +```bash +export MASTER_PORT=1235 +``` + +### Capturing full output + +The launcher already writes a complete log to `logs/log_.txt` (configurable via `--log_file PATH`), even under `--silent`. If you also want a copy at the call site, redirect there: + +```bash +srun ... runner/primus-cli direct -- preflight --perf-test \ + 2>&1 | tee preflight-$(date +%Y%m%d-%H%M%S).log +``` + +--- + +## 9. Automated node bisection (finding the bad node in an NCCL hang) + +When a cluster-wide preflight run hangs or fails, use +[`tools/preflight_bisect/bisect.py`](../../tools/preflight_bisect/bisect.py) to +run `preflight --perf-test` on smaller Slurm node subsets until suspect nodes +are isolated. + +### Prerequisites + +1. Working non-container preflight setup from the sections above, with + `VENV_ACTIVATE` exported from a shared filesystem path. +2. Run from the SLURM login/head node, where both `scontrol` and `srun` are + available. +3. Run from inside a Slurm allocation, or provide a Slurm nodelist explicitly. + +### Example from inside an allocation + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate +mkdir -p output + +python tools/preflight_bisect/bisect.py \ + --nodelist "$SLURM_NODELIST" \ + --output-dir "output/bisect-$(date +%Y%m%d-%H%M%S)" \ + --trial-timeout-sec 600 \ + --slurm-time 00:15:00 \ + --preflight-env USING_AINIC=1 \ + --preflight-env NCCL_IB_GID_INDEX=1 \ + --preflight-env NCCL_CROSS_NIC=1 \ + --preflight-env NCCL_PXN_DISABLE=0 \ + 2>&1 | tee output/bisect-latest.log +``` + +Adjust the `--preflight-env` lines to match your cluster. Per-trial logs and a +final `summary.txt` are written under `--output-dir`. + +> Note: Set `--trial-timeout-sec` high enough for a healthy subset to finish. +> Too small a timeout can turn slow-but-good trials into false failures, causing +> the bisection to explore extra paths. +> +> Note: `--preflight-env KEY=VALUE` values are concatenated into a single +> `srun --export=ALL,...` argument, so values must not contain commas or +> whitespace. Keep comma-containing values as normal exported environment +> variables. + +--- + +## 10. See also + +- [Preflight](./preflight.md) — full reference for the `preflight` subcommand and its flags +- [CLI User Guide](../../docs_deprecated/cli/PRIMUS-CLI-GUIDE.md) — container-based and `primus-cli slurm` workflows +- [`runner/primus-cli-direct.sh`](../../runner/primus-cli-direct.sh) — the direct launcher itself (`primus-cli direct` dispatches here) +- [`primus/tools/preflight/`](../../primus/tools/preflight/) — preflight implementation +- [`tools/preflight_bisect/bisect.py`](../../tools/preflight_bisect/bisect.py) — bisect wrapper for narrowing down failing nodes in multi-node preflight runs diff --git a/docs/02-user-guide/preflight.md b/docs/02-user-guide/preflight.md index d88e0d8ec..eb1b325b4 100644 --- a/docs/02-user-guide/preflight.md +++ b/docs/02-user-guide/preflight.md @@ -1,136 +1,417 @@ -# Preflight diagnostics +# Preflight -`preflight` is Primus’s cluster diagnostic command. It can produce a **fast environment report** (host, GPU, and network facts) and optionally run **performance tests** (GEMM plus intra- and inter-node communication) to catch misconfiguration or outliers before large distributed training jobs. +`preflight` is Primus' cluster diagnostic tool. It produces: -`preflight` is implemented by `primus/cli/subcommands/preflight.py`, which in turn delegates to `primus.tools.preflight`. +- A **fast info report** (host / GPU / network configuration), and +- A configurable suite of **performance tests** (GEMM TFLOPS, intra-node and inter-node communication bandwidth, P2P, ring P2P). + +Use it to spot misconfiguration, hardware degradation, or perf outliers **before** committing a large distributed training run to a global rendezvous. + +- **User-facing entry**: `primus-cli ... -- preflight [args]` +- **No-container launcher**: `runner/primus-cli direct -- preflight ...` — see [`preflight-without-container.md`](./preflight-without-container.md). +- **Implementation entrypoint**: `primus/cli/subcommands/preflight.py` → `primus/tools/preflight/preflight_perf_test.py`. + +> Looking for a faster, distributed-rendezvous-free per-node screen? See [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md). The recommended workflow is **smoke first, preflight second** — see [§10 Comparison with node-smoke](#10-comparison-with-node-smoke). --- -## Overview: What preflight checks +## 1. Two run modes (and how preflight picks one) + +Preflight has two report types, controlled by a single precedence rule: + +| Mode | Triggered by | What it does | +|---|---|---| +| **Info-only** | `--host`, `--gpu`, `--network` (in any combination) | Lightweight host / GPU / network introspection. Emits a per-node report **without requiring a rendezvous**; multi-node aggregation then uses a **timeout-bounded** rendezvous (`--dist-timeout-sec`), so it never hangs indefinitely on network misconfig. | +| **Perf-only** | `--perf-test`, `--tests ...`, or `--quick` | Runs the configured perf tests under a global rendezvous. **Implied** by `--tests` and `--quick`. | +| **Default (info + perf)** | No flags at all | Runs the info report first, then every perf test. | + +### Mode precedence + +1. **Any of `--perf-test` / `--tests` / `--quick` is set → perf-only mode.** + If info selectors (`--host`/`--gpu`/`--network`) are also present, they are dropped and a `WARN` is emitted (also written as a `> Note:` at the top of the perf report). To get both reports, run two invocations. +2. **Otherwise, any of `--host`/`--gpu`/`--network` is set → info-only mode.** + Perf-only tuning knobs (e.g. `--comm-sizes-mb`) are inert in this mode and trigger a single `WARN` listing them. +3. **Otherwise (no flags) → default**: info report **first** (no rendezvous), then perf tests. -| Category | What are checked | -|----------|----------------| -| **Host** | CPU, memory, PCIe, and related system context | -| **GPU** | ROCm-visible GPU inventory and key attributes | -| **Network** | Network configuration relevant to distributed training | -| **Performance tests** | Heavier GEMM and communication tests (slower than information-only) | +The default order ensures you always get a report even if `torch.distributed` initialization later hangs. --- -## Quick start +## 2. Quick start -### Information only (fast) +### Info report only (fast) ```bash primus-cli direct -- preflight --host --gpu --network ``` -### Full preflight (information and performance tests) +### Full preflight (info + every perf test) ```bash primus-cli direct -- preflight ``` -### Performance tests only - -Skips the host, GPU, and network information report and runs GEMM + communication tests. +### Perf tests only ```bash primus-cli direct -- preflight --perf-test ``` +### Fast pre-launch sanity check + +```bash +primus-cli direct -- preflight --quick +``` + +Equivalent on SLURM via `primus-cli slurm`: + +```bash +primus-cli slurm srun -N 4 -- preflight --quick +``` + +Without a container, see [`preflight-without-container.md`](./preflight-without-container.md) for the equivalent `runner/primus-cli direct -- preflight ...` invocations. + --- -## CLI flags reference +## 3. Test selection (`--tests`) -| Flag | Purpose | -|------|---------| -| `--host` | Include host information (CPU, memory, PCIe). Alias: `--check-host`. | -| `--gpu` | Include GPU information. Alias: `--check-gpu`. | -| `--network` | Include network information. Alias: `--check-network`. | -| `--perf-test` | Run **only** performance tests (GEMM plus intra- and inter-node communication); skip the information report. | -| `--plot` | Generate plots when used with `--perf-test`. | -| `--dist-timeout-sec` | Timeout in seconds for `torch.distributed` process-group initialization (default: 120). On failure, `preflight` still attempts to write the information report and exits with a non-zero status. | -| `--dump-path` | Output directory for reports (default: `output/preflight`). | -| `--report-file-name` | Base filename for reports (default: `preflight_report`). | -| `--disable-pdf` | Disable PDF generation (PDF is enabled by default when the toolchain allows). | +`--tests` takes a comma-separated list of canonical tokens (or `all`). Implies `--perf-test`. -**Behavior notes** +| Token | What it runs | +|---|---| +| `gemm` | Single-GPU square GEMM TFLOPS sweep. | +| `intra-allreduce` | Intra-node `all_reduce` bandwidth at every selected `--intra-group-sizes` x `--intra-comm-sizes-mb`. | +| `intra-alltoall` | Intra-node `all_to_all` bandwidth, same configuration matrix. | +| `inter-allreduce` | Inter-node `all_reduce` bandwidth at every selected `--inter-group-sizes` x `--inter-comm-sizes-mb`. | +| `inter-alltoall` | Inter-node `all_to_all` bandwidth, same configuration matrix. | +| `inter-p2p` | Inter-node point-to-point send/recv between fixed **adjacent 2-node pairs** (does not use `--inter-group-sizes`). Sized by `--inter-comm-sizes-mb`, falling back to `--comm-sizes-mb`. | +| `inter-ring-p2p` | Inter-node ring-pattern P2P, sized by `--ring-p2p-sizes-mb`. | +| `all` | Every token above. Default when `--tests` is omitted. | + +Examples: + +```bash +# GEMM only +primus-cli direct -- preflight --tests gemm -- With **no** `--host`, `--gpu`, or `--network` flags and **no** `--perf-test`, `preflight` runs in the **full** workflow (information plus performance tests). -- Combine `--host`, `--gpu`, and `--network` to limit the information report to include only those sections. +# Just the inter-node bandwidth tests +primus-cli direct -- preflight --tests inter-allreduce,inter-alltoall + +# Combine with size overrides +primus-cli direct -- preflight \ + --tests gemm,inter-allreduce \ + --comm-sizes-mb 64,1024 \ + --inter-group-sizes all +``` + +Unknown tokens fail fast (before any rendezvous): + +```text +[Primus:Preflight] ERROR: invalid perf config: --tests: unknown token 'gem'. +Valid tokens: gemm, intra-allreduce, intra-alltoall, inter-allreduce, +inter-alltoall, inter-p2p, inter-ring-p2p, all +``` --- -## Usage modes +## 4. Quick preset (`--quick`) + +`--quick` is the recommended **pre-launch sanity** preset. Implies `--perf-test`. It substitutes: -### Single-node +| Knob | `--quick` value | +|---|---| +| `--tests` | `gemm,intra-allreduce,inter-allreduce` | +| `--comm-sizes-mb` | `64,1024` | +| `--intra-group-sizes` | `LOCAL_WORLD_SIZE` (full intra-node group only) | +| `--inter-group-sizes` | `all` (full N-node group only) | +| `warmup` | `5` | +| `iteration` | `20` | + +**User-supplied flags override the preset.** For example: ```bash -primus-cli direct -- preflight --host --gpu --network +# Quick preset, but with a custom size set +primus-cli direct -- preflight --quick --comm-sizes-mb 32,256 ``` -### Multi-node (Slurm mode) +A full perf run with default knobs takes minutes; `--quick` typically finishes in <60s on healthy hardware. + +--- + +## 5. Tuning the perf tests + +All perf tuning knobs default to `None` so preflight can tell whether you set them. When unset, the documented defaults below apply. + +### 5.1 Message sizes (collective + P2P) -Info report: +| Flag | Default | Applies to | +|---|---|---| +| `--comm-sizes-mb CSV` | `2,4,8,16,32,64,128,256,512,1024` | Default for both intra- and inter-node `allreduce` / `alltoall` and `inter-p2p` when no specific override is given. | +| `--intra-comm-sizes-mb CSV` | falls back to `--comm-sizes-mb` | Override for **intra-node** `allreduce` / `alltoall`. | +| `--inter-comm-sizes-mb CSV` | falls back to `--comm-sizes-mb` | Override for **inter-node** `allreduce` / `alltoall` / `inter-p2p`. | ```bash -primus-cli slurm srun -N 4 -- preflight --host --gpu --network +# Smaller, focused sweep +primus-cli direct -- preflight --comm-sizes-mb 8,128 + +# Different sizes for intra vs inter +primus-cli direct -- preflight \ + --tests intra-allreduce,inter-allreduce \ + --comm-sizes-mb 8,128 \ + --intra-comm-sizes-mb 4,32 ``` -Full `preflight`: +### 5.2 Group sizes + +| Flag | Default | Notes | +|---|---|---| +| `--intra-group-sizes CSV` | `2,4,8` | Each value must divide `LOCAL_WORLD_SIZE`. | +| `--inter-group-sizes CSV` | `2,4,all` | `all` means the full N-node group. Other values are subgroup sizes. **Only `inter-allreduce` and `inter-alltoall` consult this flag** — for `inter-alltoall`, every requested per-group node count is internally capped at **16** before deduping (see "Inter-node alltoall is capped at 16 nodes" below), while `inter-allreduce` uses the requested sizes unchanged. `inter-p2p` and `inter-ring-p2p` ignore this flag (fixed adjacent 2-node pairs and a full-cluster ring, respectively). | ```bash -primus-cli slurm srun -N 4 -- preflight +# All-GPU intra + full N-node inter only +primus-cli direct -- preflight \ + --tests intra-allreduce,inter-allreduce \ + --intra-group-sizes 8 \ + --inter-group-sizes all ``` -Performance tests only: +Validation is gated by which tests are actually selected. For example, `--tests gemm --intra-group-sizes 3` does **not** abort on a host with `LOCAL_WORLD_SIZE=8`; the intra-group constraint is only checked when an intra test is enabled. + +#### Inter-node alltoall is capped at 16 nodes + +Regardless of the cluster size or what `--inter-group-sizes` requests, the `inter-alltoall` test always runs on per-group node counts of at most **16**. Concretely, every requested value `G` is replaced with `min(G, 16)`, and the resulting list is deduped. Examples: + +| Cluster | `--inter-group-sizes` | Requested (resolved) | `inter-alltoall` actually runs | +|---|---|---|---| +| 8 N | `all` | `[8]` | `[8]` (no change) | +| 64 N | `2,4,all` | `[2, 4, 64]` | `[2, 4, 16]` | +| 128 N | `2,4,16,32,all` | `[2, 4, 16, 32, 128]` | `[2, 4, 16]` | +| 128 N | `64` | `[64]` | `[16]` | + +When the cap actually changes the list, preflight emits a single one-line WARN to stdout so the row labels in the report (e.g. `alltoall-16nodes` instead of `alltoall-128nodes`) are not surprising. + +Why the cap, and why 16: + +- **It matches real-world usage.** Production MoE training rarely dispatches tokens across more than ~8 nodes (for example, DeepSeek-V3's largest published configuration uses `EP=64` over 8 nodes with per-token dispatch capped at 4 nodes). A 16-node ceiling covers every published configuration with comfortable headroom. +- **It keeps the test from exhausting per-node network resources.** A large `inter-alltoall` sub-group opens a near-full mesh of connections per rank during communicator setup. Capping it at 16 keeps that well within a node's ephemeral-port budget and avoids spurious `Address already in use` failures at scale (see [§7](#7-running-on-very-large-clusters--64-nodes)). +- **Other inter-node tests are unaffected.** The cap applies only to `inter-alltoall`: `inter-allreduce` uses `--inter-group-sizes` unchanged, while `inter-p2p` and `inter-ring-p2p` don't consult it at all. All three also open far fewer simultaneous connections than alltoall. +- **It is intentionally not configurable.** This is a known-safe ceiling for the communication shapes preflight characterizes, not a tuning knob. + +### 5.3 Ring P2P sizes + +| Flag | Default | Applies to | +|---|---|---| +| `--ring-p2p-sizes-mb CSV` | `10,20,40,80,160` | `inter-ring-p2p` only. | ```bash -primus-cli slurm srun -N 4 -- preflight --perf-test +primus-cli direct -- preflight \ + --tests inter-ring-p2p \ + --ring-p2p-sizes-mb 5,20,80 ``` -Use the same launcher pattern you rely on for training to ensure that distributed environment variables (`WORLD_SIZE`, `RANK`, `MASTER_ADDR`, etc.) are consistent. +### 5.4 Plotting + +| Flag | Effect | +|---|---| +| `--plot` | After each perf test, write per-size bandwidth bar charts under `//` and reference them in the markdown report. | --- -## Output files and contents +## 6. Reliability knobs + +Two knobs that are inert under happy-path conditions but matter at scale or on flaky networks. + +### 6.1 `--comm-cleanup-delay-sec FLOAT` (default `2.0`) -Default output directory: `output/preflight` (override with `--dump-path`). +Delay (seconds) inserted between destroying NCCL/RCCL process groups and creating new ones. It provides cross-rank synchronization across the destroy → setup transition, so a rank doesn't try to connect to a peer whose listener hasn't finished closing. -| File(s) | Contents | -|---------|----------| -| `.md` / `.pdf` | **Information** report: host, GPU, and network sections when those checks are enabled. | -| `_perf.md` / `_perf.pdf` | **Performance** report: GEMM and communication results from the performance test path. | +- Default `2.0` is essentially free and worth keeping at every cluster size. +- Set to `0` to disable the sleep entirely (barrier only). +- Bump to e.g. `5` only on very flaky networks. -The base `` comes from `--report-file-name` (default: `preflight_report`). +```bash +# Small/medium clusters: the default is fine. Override only if you +# see port-reuse races on a very flaky network. +primus-cli slurm srun -N 8 -- preflight --quick --comm-cleanup-delay-sec 5 +``` + +See [§7](#7-running-on-very-large-clusters--64-nodes) for guidance on running at very large scale. + +### 6.2 `--dist-timeout-sec INT` (default `120`) + +Timeout (seconds) for `torch.distributed.init_process_group`. If init does not complete within this many seconds, preflight writes the info report (when applicable) plus a `Distributed Init` failure section to the markdown report, prints a clear error, and exits `2` — instead of hanging indefinitely. + +```bash +# Fail fast if rendezvous does not work +primus-cli direct -- preflight --perf-test --dist-timeout-sec 30 +``` + +--- + +## 7. Running on very large clusters (≥ 64 nodes) + +At very large scale there are a few practical considerations beyond what smaller runs encounter. A default `preflight` invocation still runs correctly at every scale we test (up to 128 nodes) without special flags — the points below are limitations to be aware of, plus recommendations that make large-cluster runs faster and easier to interpret. + +### 7.1 Limitations + +- **`inter-alltoall` is measured on at most 16 nodes per sub-group.** Regardless of cluster size or `--inter-group-sizes`, the alltoall test is capped at 16-node sub-groups (see [§5.2](#52-group-sizes)). This is intentional — it matches real-world MoE dispatch patterns and keeps the test from exhausting per-node network resources during communicator setup — but it does mean preflight will not report alltoall bandwidth for a larger topology. The cap applies only to `inter-alltoall`: `inter-allreduce` honors `--inter-group-sizes` unchanged, while `inter-p2p` and `inter-ring-p2p` don't use it at all. +- **preflight briefly builds many communicators.** Unlike a real training job — which creates its communicators once at startup and reuses them — preflight repeatedly builds and tears down large communicators in a short window. On a cluster with an unusually narrow ephemeral-port range this can occasionally surface as `Address already in use` during setup. It is a preflight-specific artifact rather than a training failure mode; the one-line OS fix is in [§7.3](#73-optional-os-tuning). + +### 7.2 Recommended: split large runs by test family + +For clusters at or beyond ~128 nodes, run one test family per invocation instead of one large run. Each invocation stays short, and it becomes easy to see which specific communication shape is degraded if a number looks off. + +```bash +# 1) GPU + intra-node fabric first (cheap, no inter-node OOB churn). +primus-cli slurm srun -N 128 -- preflight \ + --tests gemm,intra-allreduce,intra-alltoall + +# 2) Inter-node DP-style collectives, all-nodes group only. +primus-cli slurm srun -N 128 -- preflight \ + --tests inter-allreduce \ + --inter-group-sizes all +# Note: --inter-group-sizes all is honored here for inter-allreduce. +# For inter-alltoall it would be capped at 16 (see §5.2). +primus-cli slurm srun -N 128 -- preflight \ + --tests inter-alltoall \ + --inter-group-sizes all + +# 3) Inter-node PP-style ring P2P (the test that benefits most from +# isolation — it's the closest match to what real pipeline-parallel +# training actually exercises). +primus-cli slurm srun -N 128 -- preflight \ + --tests inter-ring-p2p + +# 4) Optional: pairwise inter-node P2P scan (useful for finding a +# single bad link, slower because it walks many pairs). +primus-cli slurm srun -N 128 -- preflight \ + --tests inter-p2p +``` + +Each invocation tears down its own `WORLD` on exit and touches only one `--tests` value, so you get a per-test wall clock, can re-run a single phase in isolation, and get a separate report per run via `--report-file-name`. + +### 7.3 Optional OS tuning + +`preflight` runs fine with default OS settings at every scale we test. If you do hit `Address already in use` on a cluster with a narrow ephemeral-port range, widen the range — this is good general practice for any RDMA host regardless of preflight: + +```bash +# Widen the per-node ephemeral port range (default ~28k → ~64k ports). +sudo sysctl -w net.ipv4.ip_local_port_range="1024 65535" + +# Persist across reboots: +echo 'net.ipv4.ip_local_port_range = 1024 65535' | sudo tee /etc/sysctl.d/99-large-cluster.conf +sudo sysctl --system +``` + +--- + +## 8. Reporting + +| Flag | Default | Effect | +|---|---|---| +| `--dump-path DIR` | `output/preflight` | Output directory for reports + plots. | +| `--report-file-name NAME` | auto-generated `preflight-${NNODES}N-YYYYMMDD-HHMMSS` | Base name for report files. Omit to let preflight auto-generate a unique timestamped name (prevents stale leftovers from prior runs being mistaken for fresh output). Pass an explicit value when you want a stable / well-known filename. | +| `--disable-pdf` | enabled | Skip PDF generation (Markdown only). Useful when `weasyprint`/`markdown2` aren't installed. | + +Output files: + +| File | Produced when | Notes | +|---|---|---| +| `.md` / `.pdf` | Info-only mode, or default mode | Info report. | +| `_perf.md` / `_perf.pdf` | Perf-only mode, or default mode | Perf report (GEMM + comm). | + +Only **rank 0** writes the report. + +### Perf report layout + +A `_perf.md` produced by a default run contains, in order: + +1. (Optional) `> Note:` line listing dropped info selectors. +2. `# Nodes` legend — `Node N → Hostname` table, used by every subsequent table to keep host columns compact. +3. `=======IB Bandwidth roofline (GB/s)=======` — bandwidth of the first IB device on Node 0. +4. Per enabled test, in this order: `gemm`, `intra-comm`, `inter-comm`, `inter-p2p`, `inter-ring-p2p`. Each section has a configuration line, a results table (Node / Rank / hostname / per-size GB/s), optional plots, and a per-rank wall-clock summary. +5. `[Primus:Preflight] done in s` lines on stdout for at-a-glance progress on the launching shell. --- -## Interpreting results +## 9. Backward-compat aliases -1. **Information report:** Confirm GPU count, model match expectations, and PCIe topology is sensible for your workload. Network sections should reflect the interfaces you intend for distributed training. -2. **Performance report:** Compare GEMM and collective results across nodes. Large outliers on one node often indicate driver, fabric, or process placement issues. -3. **Timeouts:** If `--dist-timeout-sec` is exceeded, inspect firewall rules, interface bindings, `MASTER_ADDR`, and `MASTER_PORT` before scaling up training. +| Flag | Equivalent | Notes | +|---|---|---| +| `--check-host`, `--check-gpu`, `--check-network` | `--host`, `--gpu`, `--network` | Same behavior. Keep working for older scripts. | +| `--no-split-nodes-subgroup` | `--inter-group-sizes all` **and** drops `inter-p2p` | Pre-`--tests`/`--inter-group-sizes` alias. Use the new flags in new scripts. | --- -## Common issues preflight helps detect +## 10. Comparison with node-smoke + +| Aspect | `node-smoke` | `preflight` | +|---|---|---| +| Rendezvous | None — every node independent | Global `torch.distributed` | +| Wall clock | ~30–60 s for 6 nodes (Tier 1+2) | Minutes; scales with N for inter-node tests | +| Granularity | Per-node PASS/FAIL | Per-rank measurements (no auto-fail by default) | +| Inter-node bandwidth matrix | Not tested (intentionally) | Yes (allreduce/alltoall/p2p/ring-p2p) | +| Drift detection | Yes (versions, NIC firmware, port count) | No | +| Host limits / RDMA roll-call | Yes (hard fail) | Reported via `collect_*_info` only | +| Output format | Per-node JSON + cluster md + SLURM-ready txt | Markdown + PDF | + +**Recommended workflow**: run `node-smoke` first to exclude broken nodes, then run `preflight` on the surviving set to get cross-node bandwidth measurements. See [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) §4 ("Quick start") for the integration commands. + +--- + +## 11. Validation & error handling + +Preflight resolves the perf config **before** any distributed rendezvous. This means typos and bad sizes/group-sizes fail in seconds, not after a 120s NCCL init: + +```text +[Primus:Preflight] ERROR: invalid perf config: --tests: unknown token 'gem'. +[Primus:Preflight] ERROR: invalid perf config: + --intra-group-sizes: [3] do not divide LOCAL_WORLD_SIZE=8 +[Primus:Preflight] ERROR: invalid perf config: --comm-sizes-mb: values must be positive (got 0) +``` + +In info-only mode, perf-only tuning knobs trigger a single warning so you notice them but they don't abort: + +```text +[Primus:Preflight] WARN: --comm-sizes-mb,--intra-group-sizes have no effect +in info-only mode (no --perf-test/--tests/--quick). +``` + +In default mode where info selectors are dropped because perf intent was set, the preserved warning is also written into the perf report header: + +```text +> Note: info selectors --host were dropped because perf mode +> (--perf-test/--tests/--quick) takes precedence. Run them in a separate +> invocation if you want both reports. +``` + +--- + +## 12. Operational tips + +- **For multi-node runs, always use `primus-cli slurm` or `primus-cli direct` under `srun`** so distributed environment variables (`NNODES` / `NODE_RANK` / `MASTER_ADDR`) are set correctly. +- **Make sure slurm requests GPU resources**. For some clusters, you may need to explicitly request GPU resources with `srun -N --gpus-per-node=`. +- **Insufficient CPU cores cause >30x perf slowdowns** — pass `srun -c ` so RCCL's network proxy threads have CPU to spawn on. Verify with `srun -N 1 --gpus-per-node=8 bash -c 'nproc'`. +- **For a quick environment snapshot**, prefer `--host --gpu --network` — you always get a local per-node report even on a broken network, and any multi-node aggregation is timeout-bounded (`--dist-timeout-sec`), so the command never hangs. +- **Between each communication test phase**, preflight performs a global barrier + `--comm-cleanup-delay-sec` sleep (default 2 s) for cross-rank sync across the destroy → setup transition. The default works at every cluster size we test up to 128 nodes. See [§7](#7-running-on-very-large-clusters--64-nodes) for large-cluster guidance. +- **For pre-launch screening of a large cluster**, the recommended sequence is: + 1. `node-smoke` to prune broken nodes (`failing_nodes.txt`). + 2. `preflight --quick` on the surviving nodes for the perf sanity numbers. + 3. `preflight` (full) on the same set if the `--quick` numbers raise a flag. + +--- -| Symptom | What to verify in reports | -|---------|---------------------------| -| Missing or wrong GPU count | GPU section: ROCm health on the node | -| Wrong network device or address | Network section: NCCL/RCCL environment | -| Slow or asymmetric inter-node comm | Performance report: compare ranks or nodes | -| Hangs at distributed process group initialization | Use `--dist-timeout-sec` to avoid, then check rendezvous and Slurm network setup | +## 13. Running preflight without a container -For deeper, single-purpose measurements, see the [Benchmark suite](./benchmarking.md). +If you cannot (or prefer not to) use a container, see [`preflight-without-container.md`](./preflight-without-container.md) for the step-by-step `runner/primus-cli direct -- preflight ...` walkthrough — Python virtual-environment setup, SLURM invocation patterns, NCCL configuration for Broadcom and Pensando (AINIC) clusters, and many configurable-knob examples. --- -## Related documentation +## 14. See also -- [Benchmark suite](./benchmarking.md) -- [Memory and performance projection](./projection.md) -- [Post-training workflows](./posttraining.md) -- [Installation and setup](../01-getting-started/installation.md) +- [`preflight-without-container.md`](./preflight-without-container.md) — quick-start guide for `primus-cli direct -- preflight` (no container). +- [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) — full guide for the per-node smoke test (screen + exclude bad nodes). +- [`runner/primus-cli-direct.sh`](../../runner/primus-cli-direct.sh) — non-container launcher (`primus-cli direct` dispatches here). +- [`primus/tools/preflight/`](../../primus/tools/preflight/) — implementation. +- [`primus/tools/preflight/preflight_args.py`](../../primus/tools/preflight/preflight_args.py) — canonical CLI definition (single source of truth for flags + defaults). diff --git a/docs_deprecated/cli/PRIMUS-CLI-GUIDE.md b/docs_deprecated/cli/PRIMUS-CLI-GUIDE.md index c884a7e89..82046e9ca 100644 --- a/docs_deprecated/cli/PRIMUS-CLI-GUIDE.md +++ b/docs_deprecated/cli/PRIMUS-CLI-GUIDE.md @@ -69,8 +69,20 @@ Primus CLI supports three execution modes, each suitable for different scenarios # Environment check (info only) ./primus-cli direct -- preflight --host --gpu --network + +# Per-node smoke test (auto-selects `--single` since node_smoke runs one +# process per node by design; rank 0 also aggregates the per-node JSONs): +./primus-cli direct -- node_smoke --tier2-perf + +# Suppress launcher + tool stdout (--silent goes BEFORE `--`; errors and +# the launcher log file are preserved; not recommended for normal use): +./primus-cli direct --silent -- preflight --quick ``` +**Optional environment variables (direct mode)**: +- `VENV_ACTIVATE` — Path to a Python virtualenv `bin/activate` script. If set, sourced before launching; if unset, no-op (the container path uses the container's bundled Python and never sets this). +- `NNODES` / `NODE_RANK` / `MASTER_ADDR` / `MASTER_PORT` / `GPUS_PER_NODE` — Pre-export to override SLURM-derived values. Inside a SLURM allocation, they are auto-derived from `SLURM_NNODES` / `SLURM_NODEID` / `SLURM_NODELIST` when not pre-exported. + **Suitable for**: - ✅ Local development and debugging - ✅ Single-node training @@ -166,13 +178,23 @@ Primus CLI supports three execution modes, each suitable for different scenarios # Run distributed GEMM benchmark ./primus-cli slurm srun -N 2 -- benchmark gemm --M 16384 --N 16384 --K 16384 -# Multi-node environment check (info only) -# this will generate a fast info report of the host, GPU, and network +# Multi-node environment check (info only). Anything after `--` that isn't +# the keyword `container` or `direct` is treated as a primus subcommand and +# routed through the default container entry chain. ./primus-cli slurm srun -N 4 -- preflight --host --gpu --network # this will generate a full preflight report of the host, GPU, and network, as well as the performance tests ./primus-cli slurm srun -N 4 -- preflight --report-file-name preflight-report-4N +# Explicit entry-mode keyword: route through primus-cli-direct.sh instead of +# the container chain. Useful when nodes share a Python venv on a shared FS +# (see docs/preflight-direct.md for the setup). +./primus-cli slurm srun -N 4 -- direct -- preflight --quick + +# Per-node smoke via the direct entry (node_smoke auto-runs in single mode; +# rank 0 aggregates after every rank finishes): +./primus-cli slurm srun -N 4 -- direct -- node_smoke --tier2-perf + # if you are using AINIC in your cluster, use the appropriate configuration file # for preflight test, set docker image to rocm/primus:v26.3 in the configuration file ./primus-cli --config runner/use_ainic.yaml slurm srun -N 2 -- preflight --report-file-name preflight-report-2N @@ -761,7 +783,7 @@ Final result: | **Entry Script** | primus-cli-direct.sh | primus-cli-container.sh | primus-cli-slurm.sh | | **Environment Prep** | Load local GPU env | Start container + mount + devices | Allocate nodes + network config | | **Execution Location** | Current host | Inside container | Slurm-allocated nodes | -| **Final Call** | Direct torchrun execution | Execute direct.sh in container | Each node executes slurm-entry.sh → direct.sh | +| **Final Call** | Direct torchrun execution (single mode auto-selected for `node_smoke`) | Execute direct.sh in container | Each node executes slurm-entry.sh → container.sh or direct.sh (via `direct` keyword) | | **Distributed Support** | Single-node multi-GPU | Single-node multi-GPU | Multi-node multi-GPU | | **Use Case** | Dev debugging | Environment isolation | Production training | diff --git a/docs_deprecated/node-smoke-test-instruction.md b/docs_deprecated/node-smoke-test-instruction.md new file mode 100644 index 000000000..639a8dd5b --- /dev/null +++ b/docs_deprecated/node-smoke-test-instruction.md @@ -0,0 +1,313 @@ +# Node-Smoke Test — Quick-Start Instructions + +A short get-started guide for the per-node preflight smoke test. For the full design / aggregator section reference / implementation history, see [node-smoke.md](./node-smoke.md). + +--- + +## 1. What it does + +A lightweight, distributed-rendezvous-free preflight check that runs on every node in parallel under SLURM. It produces a **single PASS/FAIL verdict per node** plus SLURM-ready `passing_nodes.txt` / `failing_nodes.txt` you can pipe straight into `srun --nodelist=` / `--exclude=`. + +Use it to **screen a cluster fast and exclude bad nodes before launching a real training job**. A bad GPU, NIC, wedged driver, or leaked process on any node will surface as a node FAIL — without a single global rendezvous, so a stuck node can't wedge its peers. + +--- + +## 2. Prerequisites + + +| Prerequisite | How | +| ----------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| Python venv on a shared filesystem | Same venv used by `primus-cli direct -- preflight` (see `[preflight-direct.md](./preflight-direct.md)` §2). | +| `VENV_ACTIVATE` exported | `export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate` (optional inside the container path). | +| Inside an existing SLURM allocation | One task per node. Recommended: `runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 -- direct -- node_smoke ...`. Equivalent bare form: `srun ... --ntasks-per-node=1 runner/primus-cli direct -- node_smoke ...`. Either way the `direct -- node_smoke` path auto-selects `--single`, so each task spawns one Python process and per-GPU subprocesses are launched internally. | + + +No `MASTER_ADDR`, no `MASTER_PORT`, no global rendezvous required. + +--- + +## 3. Quick start + +**Git clone the Primus repository to a shared filesystem that all nodes can read.** + +```bash +git clone --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.git +cd Primus +git checkout dev/preflight-direct-test +``` + +**Note: remember to setup the Python virtual environment and NCCL / fabric environment variables as described in [§2 Prerequisites](#2-prerequisites).** + +> ⚠ **Set the NCCL / RCCL environment first** if you plan to run with `--tier2-perf` (the local 8-GPU RCCL all-reduce). Even though the smoke test never opens a cross-node rendezvous, the Tier 2 RCCL step calls `dist.init_process_group(backend="nccl", ...)`, and RCCL **enumerates every transport at init** (XGMI / PCIe P2P + IB + sockets). A misconfigured `NCCL_IB_HCA` / `NCCL_SOCKET_IFNAME` / `NCCL_IB_GID_INDEX` can stall init or make the all-reduce silently fall back to a slow path. The launcher's `base_env.sh` auto-detects these via `get_nccl_ib_hca.sh` + `get_ip_interface.sh`, **but auto-detect sometimes picks the wrong values inside a container** (devices masked by the network namespace, frontend NICs picked up instead of fabric NICs, etc.) so you usually want to check these settings and set them explicitly if auto-detection is wrong. +> +> Minimum-viable checklist before running with `--tier2-perf`: +> +> ```bash +> # Pin the RDMA / RoCE training NICs the container can actually see. +> # On a bare-metal host the auto-detect in base_env.sh usually picks +> # the right set; inside a container or on a multi-role node, list +> # them explicitly. Use the same set you would pass to a training job. +> export NCCL_IB_HCA="rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7" +> +> # Pick the RoCE v2 GID index for your fabric: +> # - Mellanox / Broadcom: typically 3 (base_env.sh default). +> # - Pensando Pollara (AINIC): 1. +> export NCCL_IB_GID_INDEX=3 +> +> # The bootstrap socket interface. Auto-detect prefers the first +> # non-loopback interface from `hostname -I`; override when that +> # picks a frontend NIC instead of the data-plane interface. +> export NCCL_SOCKET_IFNAME=eno0 +> export GLOO_SOCKET_IFNAME=eno0 +> ``` +> +> See `[preflight-direct.md` § 4 Cluster-specific NCCL configuration](./preflight-direct.md#4-cluster-specific-nccl-configuration) for the canonical Broadcom / Pensando Pollara values (the same `NCCL_*` set is used by both tools). If you skip `--tier2-perf`, the RCCL step is not executed and none of the above applies — Tier 1 (host limits, RDMA roll-call, leaked-process detection, etc.) does not depend on RCCL. +> +> Quick verification: `runner/primus-cli direct --dry-run -- node_smoke --tier2-perf` prints the resolved `NCCL_*` block under "Environment Variables" so you can confirm the values before launching for real. + +Recommended — through the `primus-cli slurm srun` wrapper (auto-resolves `MASTER_ADDR`/`NNODES`/`NODE_RANK`, applies `slurm.*` config defaults): + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate + +# Basic Tier 1 check (~5 s/GPU, ~30 s total) +runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + -- direct -- node_smoke + +# Tier 1 + Tier 2 perf sanity (GEMM TFLOPS, HBM GB/s, local 8-GPU RCCL) +runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + -- direct -- node_smoke --tier2-perf + +# Then re-run training, excluding any node the smoke test failed: +srun --exclude=$(paste -sd, output/preflight/failing_nodes.txt) ... your-real-job +``` + +Equivalent with bare `srun` (works the same; useful when composing with custom `srun` flags): + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke + +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf +``` + +Single-node sanity check (no SLURM): + +```bash +runner/primus-cli direct -- node_smoke +``` + +> **Both forms produce the same workload.** The wrapper form is recommended because it resolves the distributed env once on the launching node and propagates it via `--env`, and applies any `slurm.`* config defaults (partition / time / etc.). See `[preflight-direct.md` § Wrapper vs. bare-srun](./preflight-direct.md#wrapper-vs-bare-srun) for the precedence table. + +--- + +## 4. More examples (by configuration knob) + +> **Convention used below.** The examples in this section are written with bare `srun` for brevity. Anywhere you see `srun runner/primus-cli direct -- node_smoke ...`, the equivalent wrapper form is `runner/primus-cli slurm srun -- direct -- node_smoke ...`. Pick whichever matches your habits; both target the same launcher. + +### 4.1 Hard-fail on partial NIC enumeration + +Catches "7 of 8 RDMA NICs visible" — common cause of crashes after RoCE init. + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf --expected-rdma-nics 8 +``` + +### 4.2 Tighten Tier 2 perf thresholds + +Reject GPUs that come in below your acceptance bar. Defaults: GEMM 600 TFLOPS, HBM 2000 GB/s, local RCCL 100 GB/s. + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf \ + --gemm-tflops-min 700 --hbm-gbs-min 4500 --rccl-gbs-min 180 +``` + +### 4.3 Tighten host limits + +Fail nodes whose `RLIMIT_MEMLOCK` or `/dev/shm` is too small for production training. + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke \ + --ulimit-l-min-gb 64 --shm-min-gb 16 +``` + +### 4.4 Custom dump path + +Keep one report per smoke run instead of overwriting the default location. + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf \ + --dump-path /shared/smoke-archive/$(date +%Y%m%d-%H%M%S) +``` + +### 4.5 Allow / extend the foreign-process whitelist + +By default, leaked / foreign processes holding a GPU FAIL the node (most common cause of "training fails to launch on a healthy-looking node"). Allowed by default: `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter`. + +```bash +# Add a site-specific monitoring agent to the whitelist +srun ... runner/primus-cli direct -- node_smoke \ + --allowed-procs gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter,my-monitor + +# Don't fail at all on foreign processes (still reported in the markdown) +srun ... runner/primus-cli direct -- node_smoke --allow-foreign-procs +``` + +#### Containers: `name='N/A'` false positives → use `--allow-foreign-procs` + +> ⚠ **Running node_smoke inside a container almost always trips this check.** `amd-smi process --json` reports `name="N/A"` for kernel/system PIDs like `gpuagent` whose `/proc//comm` it cannot read, and the fallback `_resolve_proc_name(pid)` inside `node_smoke` then also fails because the container's `/proc` typically does not expose host PIDs (private PID namespace without `--pid=host`, or a `hidepid=2` mount). The unresolved name doesn't match the allowlist (`gpuagent,rocm-smi-daemon,...`), so the check fires and the node FAILs — even though the only "foreign" processes are well-known system daemons holding zero HBM. +> +> **In the container path, pass `--allow-foreign-procs`:** +> +> ```bash +> srun ... runner/primus-cli direct -- node_smoke --tier2-perf --allow-foreign-procs +> ``` +> +> The processes are still listed in `smoke_report.md` under "Busy GPUs / leaked processes" so a real leak is still visible; only the FAIL verdict is downgraded. +> +> **Narrower alternative** if you want the check to still catch leaks with resolvable names (e.g. a leftover `python` rank), add the literal sentinel `N/A` to the allowlist: +> +> ```bash +> srun ... runner/primus-cli direct -- node_smoke --tier2-perf \ +> --allowed-procs gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter,N/A +> ``` +> +> The annotator runs `_resolve_proc_name` first, so whenever a real name *can* be resolved (on the host, or after fixing `/proc` visibility) it overrides "N/A" and the normal allowlist applies. The `N/A` entry only matches PIDs whose name genuinely could not be recovered — strictly narrower than `--allow-foreign-procs`. +> +> **Root-cause fix** (preferred long-term): grant the container access to host PIDs so `_resolve_proc_name` works and the report shows real names (`gpuagent`, etc.) instead of `N/A`. Typical fixes: +> +> - Launch with `--pid=host` (Docker / Podman) so host PIDs are directly addressable. +> - Mount `/proc` without `hidepid=2`. +> - Loosen `ptrace_scope` or grant `CAP_SYS_PTRACE`. +> +> Once any of those is in place, `_resolve_proc_name` finds the names, the default allowlist matches them, and you no longer need `--allow-foreign-procs`. + +### 4.6 Require specific tools + +Make missing CLI tools a hard FAIL (default: warn-only). + +```bash +srun ... runner/primus-cli direct -- node_smoke --require-tools amd-smi,rocm-smi,lsof +``` + +### 4.7 Skip dmesg scan (containers with no privileges) + +```bash +srun ... runner/primus-cli direct -- node_smoke --skip-dmesg +``` + +### 4.8 Re-aggregate from existing per-node JSONs (no re-run) + +Useful when you only want to refresh the markdown report, or when you've collected JSONs separately. + +```bash +# From any node, no allocation needed if you're just reading local files. +# The primus-cli wrapper always runs both phases, so use the standalone +# aggregate subcommand for "aggregate only" -- it reads the existing +# /smoke/*.json without re-running the per-node smoke step. +python -m primus.tools.preflight.node_smoke aggregate \ + --dump-path output/preflight --expected-nodes 6 --wait-timeout-sec 5 +``` + +### 4.9 Silent mode (for CI) + +Suppresses wrapper stdout, but the **final report path is still printed** and stderr / exit code are preserved. + +```bash +srun ... runner/primus-cli direct --silent -- node_smoke --tier2-perf +``` + +### 4.10 Combined "production-ready screen" + +A representative one-shot for a production cluster screen: + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct --silent -- node_smoke --tier2-perf \ + --expected-rdma-nics 8 \ + --gemm-tflops-min 700 --hbm-gbs-min 4500 --rccl-gbs-min 180 \ + --ulimit-l-min-gb 64 --shm-min-gb 16 \ + --require-tools amd-smi,rocm-smi,lsof \ + --dump-path /shared/smoke-archive/$(date +%Y%m%d-%H%M%S) +``` + +--- + +## 5. Outputs + +All written under `--dump-path` (default `output/preflight/`). + + +| File | Purpose | +| ------------------------- | -------------------------------------------------------------------------------------------------------------------- | +| `smoke/.json` | Per-node verdict + every collected metric. One file per node. | +| `smoke_report.md` | Human-readable cluster report (status table, drift sections, perf summary, failing-node detail). | +| `passing_nodes.txt` | Newline-separated short hostnames. Pipe into `srun --nodelist=`. | +| `failing_nodes.txt` | Newline-separated short hostnames. Pipe into `srun --exclude=`. | +| `expected_nodes.txt` | Auto-populated from `scontrol show hostnames "$SLURM_JOB_NODELIST"`. Lets the report name nodes that never reported. | + + +Read the cluster verdict at a glance: + +```bash +head -10 output/preflight/smoke_report.md +``` + +Feed bad nodes into a re-run: + +```bash +srun --exclude=$(paste -sd, output/preflight/failing_nodes.txt) ... your-real-job +``` + +--- + +## 6. Common knobs (cheat sheet) + + +| Flag | Default | When you'd change it | +| ---------------------------- | ------------------------------------------------ | ------------------------------------------------------------------------------------- | +| `--tier2-perf` | off | Always on for production screens — adds GEMM TFLOPS, HBM GB/s, local RCCL all-reduce. | +| `--gemm-tflops-min N` | 600 | Site-specific acceptance bar. | +| `--hbm-gbs-min N` | 2000 | Site-specific acceptance bar (MI300X healthy ≈ 4500–5000). | +| `--rccl-gbs-min N` | 100 | Site-specific acceptance bar. | +| `--expected-rdma-nics N` | unset | Hard-fail on partial NIC enumeration. | +| `--ulimit-l-min-gb GB` | 32 | Raise for production training profiles. | +| `--shm-min-gb GB` | 8 | Raise for large-batch / many-rank profiles. | +| `--allow-foreign-procs` | off | Co-tenant clusters or shared GPUs. | +| `--allowed-procs LIST` | `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter` | Add site-specific monitoring agents. | +| `--require-tools LIST` | `""` | Fail-fast if a CLI tool is missing in PATH. | +| `--skip-dmesg` | off | Inside unprivileged containers. | +| `--dump-path DIR` | `output/preflight` | Archive each run separately. | +| `--silent` (wrapper) | off | CI / scripted runs. | +| `--aggregate-only` (wrapper) | off | Re-render report without re-running per-node checks. | + + +For the full flag list and the aggregator subcommand, see `python -m primus.tools.preflight.node_smoke run --help` and `... aggregate --help`, or `[node-smoke.md](./node-smoke.md)` §"Configuration knobs". + +--- + +## 7. Troubleshooting + + +| Symptom | Likely cause / fix | +| ----------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `[ERROR] [direct] VENV_ACTIVATE is set but file does not exist: ...` | Fix the path (`export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate`), or `unset VENV_ACTIVATE` to fall back to system / container Python. | +| Every node FAILs with `gpu_processes: ... name='N/A'` | Should no longer happen after the `/proc//comm` fallback fix. If it does, check that `/proc//comm` is readable on the node (`hidepid` mount?). Workaround: `--allow-foreign-procs`. | +| Some nodes never produce a JSON | Aggregator names them in `failing_nodes.txt` via `expected_nodes.txt`. If `scontrol` was unavailable, they'll appear as ``. | +| Tier 2 perf numbers below threshold on a known-good node | Almost always insufficient CPU cores on `srun` — pass `-c ` so RCCL proxy threads have CPU. | +| Re-run on a smaller nodelist still shows the previously removed nodes as PASS | Default behavior cleans stale JSONs on rank 0. If you passed `--no-clean-dump-path`, either remove it or `rm -rf output/preflight` between runs. | + + +--- + +## 8. See also + +- `[node-smoke.md](./node-smoke.md)` — full design, aggregator sections, configuration reference, implementation history. +- `[preflight-direct.md](./preflight-direct.md)` — the heavier `preflight` tool with global rendezvous and inter-node bandwidth tests. +- `[primus/cli/subcommands/node_smoke.py](../primus/cli/subcommands/node_smoke.py)` — the primus-cli subcommand wiring (two-phase dispatch: rank-N run + rank-0 aggregate). +- `[primus/tools/preflight/node_smoke/cli.py](../primus/tools/preflight/node_smoke/cli.py)` — canonical flag definitions and per-node / aggregate phase bodies. diff --git a/docs_deprecated/node-smoke.md b/docs_deprecated/node-smoke.md new file mode 100644 index 000000000..98471ea80 --- /dev/null +++ b/docs_deprecated/node-smoke.md @@ -0,0 +1,429 @@ +# Node-Local Smoke Test + +> **Just want to run it?** See [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) for the short quick-start guide. This document is the full reference (architecture, every report section, every flag, design history). + +A lightweight, distributed-rendezvous-free preflight check that runs on every node in parallel under SLURM. Designed to **quickly identify broken nodes before a large training job commits to a global rendezvous**. Because training jobs allocate whole nodes, a under-performing GPU (or NIC, or wedged driver) takes the entire node out of rotation -- so the smoke test produces a single PASS/FAIL verdict per node and SLURM-ready `passing_nodes.txt` / `failing_nodes.txt` you can pipe straight into `srun --nodelist=` / `--exclude=`. + +- **Implementation**: `primus/tools/preflight/node_smoke/` (Python sub-package; entry point `python -m primus.tools.preflight.node_smoke`). +- **Recommended launcher**: `runner/primus-cli slurm srun -- direct -- node_smoke ...` (auto-resolves `MASTER_ADDR`/`NNODES`/`NODE_RANK` via `--env`, applies `slurm.*` config defaults, same pattern as `train` / `benchmark`). The shorter `runner/primus-cli direct -- node_smoke ...` form (bare `srun` + direct) is equivalent and useful for ad-hoc runs. +- **Companion**: see `docs/preflight.md` for the full preflight tool (with global rendezvous and richer perf tests). + +## Quick start + +Recommended — through the `primus-cli slurm srun` wrapper: + +```bash +# Inside an existing SLURM allocation (the normal case): +runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + -- direct -- node_smoke + +# With perf sanity (GEMM TFLOPS, HBM GB/s, local 8-GPU RCCL all-reduce): +runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + -- direct -- node_smoke --tier2-perf + +# Hard-fail on partial NIC enumeration (e.g. 7 of 8 RDMA NICs). +# The count is compared against the *training-NIC* set after the +# selector chain runs, so frontend / storage RoCE NICs do not count. +runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + -- direct -- node_smoke --tier2-perf --expected-rdma-nics 8 + +# Explicitly pin the training-NIC selector (otherwise NCCL_IB_HCA env +# is used; otherwise admin-disabled phys_state ports are auto-excluded): +runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + -- direct -- node_smoke --tier2-perf \ + --rdma-nic-allowlist 'rocep158s0:1,rocep190s0:1,rocep206s0:1,rocep222s0:1,rocep28s0:1,rocep62s0:1,rocep79s0:1,rocep96s0:1' +``` + +Equivalent with bare `srun` (works the same; useful when composing with custom `srun` flags): + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke + +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf +``` + +Single-node local check (no SLURM, both forms collapse to the same call): + +```bash +runner/primus-cli direct -- node_smoke +``` + +> **Note on the `direct` keyword**: with the `primus-cli slurm srun` wrapper, the entry-mode keyword `direct` between the two `--`s is mandatory to take the direct (no-container) path. Without it the wrapper routes through the **container** path. See [`preflight-direct.md` § Wrapper vs. bare-srun](./preflight-direct.md#wrapper-vs-bare-srun) for the full precedence and caveats. + +When `VENV_ACTIVATE` is set, `primus-cli direct` sources it before launching `node_smoke` (same convention as `primus-cli direct -- preflight`). When unset (e.g. inside the container path), it is a no-op. + +## Outputs + +After a run, `/` (default `output/preflight/`) contains: + +| File | Purpose | +|---|---| +| `smoke/.json` | Per-node verdict + every collected metric. One file per node. | +| `smoke_report.md` | Human-readable cluster report (status table, drift sections, perf summary, failing-node detail). | +| `passing_nodes.txt` | Newline-separated short hostnames. Pipe into `srun --nodelist=`. | +| `failing_nodes.txt` | Newline-separated short hostnames. Pipe into `srun --exclude=`. | +| `expected_nodes.txt` | (auto-populated from `scontrol show hostnames "$SLURM_JOB_NODELIST"`) Used by the aggregator to name nodes that never reported. | + +```bash +# Re-run training, excluding the bad nodes from the previous smoke: +srun --exclude=$(paste -sd, output/preflight/failing_nodes.txt) ... your-real-job +``` + +## Architecture + +- **Per-node Python entry** (`node_smoke.py run`) — runs independently on every node. No `MASTER_ADDR`, no global `torch.distributed` rendezvous; a stuck node cannot wedge its peers. +- **Per-GPU isolation** — each GPU's checks run in their own Python subprocess with a hard timeout. A stuck `torch.cuda.set_device()` (which can't be aborted by `signal.alarm` because it sits inside a non-interruptible driver syscall) is `SIGKILL`'d from the parent without affecting the rest of the node's checks. +- **Local-only RCCL** — Tier 2 all-reduce uses `torch.multiprocessing.spawn` over `tcp://127.0.0.1`. No cross-node communication. +- **Aggregator on `NODE_RANK==0`** — polls `/smoke/` for the expected number of JSONs (with a timeout), computes drift across the cluster, writes the markdown report and pass/fail txt files. Returns non-zero if any node FAILs or never reports. + +## Module layout — where each check lives + +The implementation is a Python sub-package under `primus/tools/preflight/node_smoke/`, split so each Tier 1 sub-section, the per-GPU subprocess body, the orchestrator, and the aggregator each live in their own file. The single public entry point is `main` (re-exported from `__init__`); `python -m primus.tools.preflight.node_smoke ...` resolves to `__main__.py` which calls it. + +``` +primus/tools/preflight/node_smoke/ +├── __init__.py # re-export `main` +├── __main__.py # `python -m primus.tools.preflight.node_smoke` +├── cli.py # `_build_parser`, `_cmd_run`, `_cmd_aggregate`, +│ # `_cmd_per_gpu`, `main` +├── types.py # `GPUResult`, `NodeResult` dataclasses +├── logging_utils.py # `_ts`, `_log`, `_warn`, hostname normalisation +├── shell_utils.py # `_which`, `_read_text`, `_resolve_gpu_bdf`, +│ # `_systemctl_is_active`, `_parse_size_with_unit`, +│ # `_findings_to_dicts` +├── per_gpu.py # `_per_gpu_body` (Tier 1 + optional Tier 2 perf, +│ # GEMM/HBM bandwidth measurement) +├── rccl_local.py # node-local RCCL all-reduce (Tier 2) +├── orchestrator.py # `_spawn_per_gpu`, `_node_status_from`, +│ # `_clean_dump_path` +├── collectors/ # one module per Tier 1 sub-section +│ ├── dmesg.py # recent dmesg error scan +│ ├── fingerprint.py # Tier 1 A — software-stack fingerprint +│ ├── nics.py # Tier 1 B — NIC / RDMA roll-call +│ ├── host_limits.py # Tier 1 C — ulimit / shm / NUMA / governor +│ ├── gpu_low_level.py # Tier 1 D-1 — amd-smi metric (ECC, throttle, +│ │ # clocks, power) +│ ├── xgmi.py # Tier 1 D-2 — XGMI link matrix +│ ├── clock.py # Tier 1 E — wall time + time-daemon health +│ ├── rocm_smi.py # Tier 1 F + cross-tool fallbacks for D-1/2/G +│ ├── gpu_processes.py # Tier 1 G — foreign / leaked PID detection +│ ├── tooling.py # tooling availability inventory +│ └── reused_info.py # reused gpu/host/network info collectors +└── aggregator/ + ├── summarizers.py # `_*_rows` / `_*_summary` data shapers + └── report.py # `write_smoke_report` + one `_write_
` + # helper per Markdown `##` section +``` + +Dependency graph (acyclic; arrows mean "imports"): + +```mermaid +flowchart TD + cli[cli.py] --> orchestrator[orchestrator.py] + cli --> per_gpu[per_gpu.py] + cli --> rccl_local[rccl_local.py] + cli --> aggReport["aggregator/report.py"] + cli --> collectorsAll["collectors/*"] + cli --> logging[logging_utils.py] + cli --> types[types.py] + + aggReport --> aggSum["aggregator/summarizers.py"] + aggReport --> logging + aggSum --> tooling["collectors/tooling.py"] + + orchestrator --> types + + per_gpu --> shell[shell_utils.py] + + collectorsAll --> shell + collectorsAll --> rocmsmi["collectors/rocm_smi.py"] + rocmsmi --> shell + logging -.->|stdlib only| std[(socket / sys / time)] + shell -.->|stdlib only| std + types -.->|stdlib only| std +``` + +* `collectors/` are leaf modules (depend on `shell_utils` + sometimes `collectors/rocm_smi`); they never import `cli` / `orchestrator` / `per_gpu`. +* `per_gpu` is the body of the `_per_gpu` subprocess and is the ONLY module loaded inside that subprocess via `python -m primus.tools.preflight.node_smoke _per_gpu N` — so its dependency surface is intentionally narrow (only `shell_utils`). +* `aggregator/` only depends on its own `summarizers` plus `logging_utils` (and `collectors/tooling` for the static `_TRACKED_TOOLS` constant). + +## What's checked + +### Tier 1 — mandatory (~5 s / GPU, always runs) + +**Per-GPU subprocess (with hard timeout):** +- `torch.cuda.set_device(i)` — proves the device is bindable (a stale GPU often fails here) +- 256 MB allocation +- Tiny GEMM (2048² bf16) with `isfinite()` check on the result + +**Reused from existing preflight collectors** (no rendezvous needed): +- `collect_gpu_info` — `level='fail'` Findings cause node FAIL +- `collect_host_info` — same +- `collect_network_info(expect_distributed=False)` — same + +**dmesg recent-error scan** — greps the last `--dmesg-minutes` (default 15) of `dmesg` for known patterns (`xid`, `gpu reset`, `hung_task`, `mce:`, `amdgpu.*error`, ...). Matches are surfaced in the report. + +**A. Software-stack fingerprint** (`tier1.fingerprint`): +- Kernel, OS, Python +- ROCm version (`/opt/rocm/.info/version`) +- amdgpu kernel-module version (`/sys/module/amdgpu/version`) +- PyTorch version, `torch.version.hip`, RCCL version (via `torch.cuda.nccl.version()`), librccl path +- Per-IB-device firmware (`/sys/class/infiniband//fw_ver`) and HCA model + +**B. NIC / RDMA roll-call** (`tier1.nics`): +- Per port (read entirely from `/sys/class/infiniband` — no `ibv_devinfo`/`ibstat` dependency, works inside containers): `state`, `phys_state`, `rate`, netdev + MTU, total non-zero GIDs, RoCE v2 GID count +- **Training-NIC selector** — many clusters expose more RDMA-capable ports than the training job uses (frontend / management / storage NICs). The hard-fail rules only run against the *included* subset. Precedence: + 1. `--rdma-nic-allowlist 'rocep158s0:1,rocep190s0:1,...'` (full `NCCL_IB_HCA` syntax: comma-separated `device[:port]`, `^...` for denylist, `=dev` for exact-match, no `:port` to match any port on the device). + 2. `NCCL_IB_HCA` env (same syntax) — mirrors what NCCL/RCCL itself will use, so the smoke test and the training launch agree by construction. + 3. Heuristic: auto-exclude any port whose `phys_state` is `Disabled` or `Sleep` (admin-disabled at firmware/driver level — no SFP, BIOS port-disable, netdev admin-down). Real failure modes on a port that *is* meant to be used produce a different `phys_state` (`Polling`, `LinkErrorRecovery`, or `LinkUp` with `state!=ACTIVE`), so this heuristic does not mask cable / driver problems. + 4. Fallback: every IB port must be ACTIVE / LinkUp. +- Excluded ports stay visible in `tier1.nics.ports` and are summarised in `tier1.nics.excluded_ports` + `info_issues` for diagnostics. They do NOT contribute to the node FAIL signal. +- **Hard fail rules** (only on the included set): port not `ACTIVE` / not `LinkUp`, active port with zero RoCE v2 GIDs (RoCE) or zero valid GIDs (IB), included-NIC count ≠ `--expected-rdma-nics N` (when set). +- **Empty-set guard**: if every discovered port gets excluded, the node still hard-fails — a node with zero training NICs cannot participate in inter-node training. + +**C. Host limits / system tunables** (`tier1.host_limits`): +- `RLIMIT_MEMLOCK`, `RLIMIT_NOFILE`, `RLIMIT_NPROC` +- `/dev/shm` size + free +- NUMA node count, CPU count, `cpu0` scaling governor +- **Hard fail rules**: `RLIMIT_MEMLOCK` finite and below `--ulimit-l-min-gb` (default 32 GiB) → "RDMA pin will fail under load"; `/dev/shm` size below `--shm-min-gb` (default 8 GiB) → "NCCL shared-mem may fail" + +### Tier 2 — optional perf sanity (`--tier2-perf`) + +Per-GPU steady-state metrics, with iteration counts aligned to the preflight `--quick` preset (`warmup=5, iters=20` for GEMM/RCCL; `warmup=10, iters=20` for HBM) so smoke and preflight numbers are directly comparable. + +- **GEMM TFLOPS** — 8192³ bf16 `torch.matmul`, threshold `--gemm-tflops-min` (default 600). +- **HBM GB/s** — 512 MB device-to-device `torch.Tensor.copy_` (counts read + write), threshold `--hbm-gbs-min` (default 2000). +- **Local 8-GPU RCCL all-reduce GB/s** — algorithmic bandwidth `2·S·(P-1)/P / t / 1e9` at 64 MB, threshold `--rccl-gbs-min` (default 100). + +## Aggregator report sections + +Every section short-circuits to a placeholder (`*All nodes match.*` / `*No NIC issues.*` / `*No host-limit issues.*`) on a healthy cluster, so the report stays short. Each section header is part of the operator-facing contract — order and wording are stable across releases (some Slack bots / CI scripts grep for them). + +In order: + +1. **Status table** — one row per node with `node_rank`, hostname, PASS/FAIL, duration, top fail reason. +2. **Stack drift across cluster** — for every scalar fingerprint key, outliers vs the cluster majority. +3. **NIC firmware drift across cluster** — per-IB-device firmware drift. +4. **NIC / RDMA roll-call issues** — every offending node + port (included set only). +5. **NIC port-count summary** — cluster-majority *training-NIC* count and any node that disagrees (catches partial-NIC degradation without `--expected-rdma-nics`). The count is taken from the included set, so nodes that legitimately have extra frontend / storage RoCE NICs don't show up as anomalies. +6. **NIC excluded ports (informational)** — ports the selector chain dropped from the training-NIC set, grouped by source (`--rdma-nic-allowlist` / `NCCL_IB_HCA` / heuristic). Informational only; does not contribute to FAIL. +7. **Host limits issues** — per-node hard-limit violations. +8. **GPU visibility issues** — nodes where torch couldn't see the GPUs or amd-smi sees more GPUs than torch (stale ROCm / wedged amdgpu driver). Independent of every other collector. +9. **GPU low-level outliers (PCIe link / HBM)** — per-GPU outliers vs the cluster majority on PCIe width/speed and HBM total. +10. **XGMI link issues** — any non-XGMI GPU pair (intra-node collectives silently fall back to PCIe). +11. **Cluster clock + time daemons** — wall-clock spread plus per-node time-daemon health. +12. **Tooling self-latency (`rocm-smi --version`)** — slow / timed-out tool calls (precursor to a wedged amdgpu driver). +13. **Tooling availability** — always-on inventory of `amd-smi` / `rocm-smi` / `lsof` per node, plus which Tier 1 checks have NO working tool on each node. +14. **Busy GPUs / leaked processes** — foreign PIDs holding GPUs at smoke start (most common cause of training failing to launch on an otherwise-healthy node). +15. **GPU pre-touch HBM usage outliers** — GPUs with non-trivial HBM in use BEFORE smoke touched the device. +16. **GPU compute-activity outliers** — GPUs with `gfx_activity_pct >= --gpu-activity-warn-pct` at smoke start (warn-only). +17. **Tier 2 perf summary** (conditional, only when at least one node ran Tier 2) — per-node GEMM TFLOPS / HBM GB/s as `min / median / max`, plus local RCCL GB/s. +18. **Failing nodes — full reasons** (conditional, only when there are failing nodes) — every fail reason, expanded per node. + +Each section that does pure data shaping is wrapped in its own `try / except`, so a future schema bug in one section can't truncate the rest of the report. The two intentional EXCEPTIONS are **Tier 2 perf summary** and **Failing nodes — full reasons** — both deliberately propagate exceptions so a regression in either bubbles up rather than silently rendering a half-empty section. + +## Configuration knobs + +The authoritative source of flags + defaults is `python -m primus.tools.preflight.node_smoke run --help` (and `... aggregate --help`). The tables below mirror the parser as of the package-split refactor. + +### `run` subcommand + +| Flag | Default | Purpose | +|---|---|---| +| `--dump-path` | `output/preflight` | Output directory. | +| `--expected-gpus N` | auto | Override GPU count (auto-detected from `LOCAL_WORLD_SIZE` / `GPUS_PER_NODE` / `torch.cuda.device_count()`). | +| `--per-gpu-timeout-sec` | 15 | Hard timeout per per-GPU subprocess. | +| `--tier2-perf` | off | Enable Tier 2 perf sanity (per-GPU GEMM TFLOPS + HBM GB/s + node-local RCCL all-reduce). Single switch — you cannot enable just one half. | +| `--gemm-tflops-min` | 600 | Tier 2 GEMM threshold. | +| `--hbm-gbs-min` | 2000 | Tier 2 HBM threshold. | +| `--rccl-size-mb` | 64 | Local RCCL message size. | +| `--rccl-gbs-min` | 100 | Local RCCL bandwidth threshold. | +| `--rccl-timeout-sec` | 120 | Hard timeout for the RCCL phase. | +| `--skip-dmesg` | off | Skip dmesg scan (e.g. inside containers). | +| `--dmesg-minutes` | 15 | dmesg `--since` window. | +| `--expected-rdma-nics N` | auto-report-only | When set, a mismatch between the **included (training-NIC) count** and N becomes a node FAIL. Compares against the post-selector count, not the raw number of devices under `/sys/class/infiniband`. | +| `--rdma-nic-allowlist LIST` | unset | Explicit training-NIC selector in `NCCL_IB_HCA` syntax (`device[:port],...`, `^...` denylist, `=dev` exact-match). Wins over `NCCL_IB_HCA` env. When neither this flag nor the env is set, the collector auto-excludes ports whose `phys_state` is `Disabled` or `Sleep`. | +| `--ulimit-l-min-gb GB` | 32 | RLIMIT_MEMLOCK threshold (0 disables). | +| `--shm-min-gb GB` | 8 | `/dev/shm` size threshold (0 disables). | +| `--rocm-smi-timeout-sec SEC` | 5.0 | Hard timeout for the `rocm-smi --version` self-latency canary; hitting it is a node FAIL (driver likely wedging). | +| `--hbm-busy-threshold-gib GiB` | 2.0 | FAIL the node if any GPU has at least this many GiB of HBM in use BEFORE smoke touches the device (i.e. someone else is holding it). Boundary is inclusive. | +| `--allow-foreign-procs` | off | Do NOT FAIL the node when foreign processes are found holding a GPU. They will still be reported. | +| `--allowed-procs LIST` | `gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter` | Comma-separated process names that are OK to find holding the GPU. Set to `""` to disable the whitelist. | +| `--gpu-activity-warn-pct PCT` | 20.0 | Warn (does NOT fail) if amd-smi reports any GPU's `gfx_activity_pct` above this when smoke starts. | +| `--require-tools LIST` | `""` (warn-only) | Comma-separated CLI tool names that MUST be in PATH (`amd-smi`, `rocm-smi`, `lsof`); anything missing becomes a hard node FAIL. | +| `--no-clean-dump-path` | off | Do NOT auto-wipe stale per-node JSONs / aggregator outputs from `--dump-path` on rank 0 at startup. Default behavior is to clean so re-runs on a different (smaller) nodelist don't inherit ghost PASS verdicts from removed nodes. | + +### `aggregate` subcommand + +| Flag | Default | Purpose | +|---|---|---| +| `--dump-path` | `output/preflight` | Same as `run`. | +| `--expected-nodes N` | none | If fewer JSONs land within `--wait-timeout-sec`, missing nodes are added as FAIL placeholders. | +| `--wait-timeout-sec` | 60 | Polling timeout. | +| `--rocm-smi-warn-sec SEC` | 1.0 | Flag (warn-only) any node where `rocm-smi --version` took longer than this. | +| `--clock-skew-warn-sec SEC` | 30.0 | Warn when wall-clock spread across nodes exceeds this many seconds. Includes srun launch jitter so the default is loose. | +| `--hbm-busy-threshold-gib GiB` | 2.0 | Mirrors the `run`-side default; used to label the **GPU pre-touch HBM usage outliers** section. | +| `--gpu-activity-warn-pct PCT` | 20.0 | Mirrors the `run`-side default; used to label the **GPU compute-activity outliers** section. | +| `--expected-nodelist-file FILE` | none | One short hostname per line. Missing nodes get their **real short hostname** in the report and `failing_nodes.txt` (instead of `` placeholders). The primus-cli wrapper auto-populates this from `scontrol show hostnames "$SLURM_JOB_NODELIST"` under SLURM. | + +### Launcher-level knobs (`primus-cli direct`) + +These are consumed by `primus-cli-direct.sh` **before** the `--` separator (not forwarded to the `node_smoke` Python tool): + +| Flag | Purpose | +|---|---| +| `--silent` | Back-pocket knob: redirect launcher + tool stdout to `/dev/null`. Launcher errors (`LOG_ERROR` / `LOG_WARN` on stderr) and the log file are preserved. Exit code propagated. | +| `--debug` | Verbose launcher logging. | +| `--dry-run`| Show the resolved command without executing. | +| `--env KEY=VALUE` | Inject an env var into the Python process. | + +### Rare advanced control (run-only / aggregate-only / no-aggregate) + +The primus-cli `node_smoke` subcommand always runs `_cmd_run` on every rank followed by rank-0 `_cmd_aggregate`. That is what users want ~100% of the time. For the rare cases where you need just one phase (e.g. re-aggregate yesterday's JSONs without re-running per-node, or smoke a single node without producing a cluster report), reach for the standalone CLI directly: + +```bash +# Per-node only, no aggregator (useful when scheduling phases separately): +python -m primus.tools.preflight.node_smoke run --tier2-perf + +# Aggregate only (read existing /smoke/*.json, produce cluster report): +python -m primus.tools.preflight.node_smoke aggregate \ + --dump-path output/preflight --expected-nodes 6 --wait-timeout-sec 5 +``` + +## Comparison with the full `preflight` + +| Aspect | `node_smoke` | full `preflight` | +|---|---|---| +| Rendezvous | None — every node independent | Global `torch.distributed` | +| Wall clock | ~50–60 s for 6 nodes (Tier 1+2) | Minutes; scales with N for inter-node tests | +| GEMM threshold | Hard threshold per GPU | Reports per-GPU numbers, no auto-fail | +| HBM bandwidth | Yes (D2D `copy_`) | Not measured | +| Inter-node all-reduce/all-to-all | Not tested (intentionally) | Yes | +| Drift detection | Yes (versions, NIC firmware, port count) | No | +| Host limits / RDMA roll-call | Yes (hard fail) | Reported via `collect_*_info` only | +| Output format | Per-node JSON + cluster md + SLURM-ready txt | Markdown + PDF | + +Use `node_smoke` to **screen** a cluster fast and exclude bad nodes. Use the full `preflight` when you want **deep cross-node measurements** (inter-node bandwidth matrix, ring-P2P, etc.). + +--- + +## Implementation history + +Captured here so future contributors understand *why* the design looks the way it does. + +### 1. Configurable preflight (predecessor work) + +Before `node_smoke` existed, the goal was simply to make the full `preflight` perf phase configurable: which tests to run, which message sizes, which subgroup sizes. Outcome (committed before `node_smoke`): + +- `--tests gemm,intra-allreduce,inter-allreduce,...` to select tests +- `--comm-sizes-mb 2,8,64,1024`, `--intra-comm-sizes-mb`, `--inter-comm-sizes-mb`, `--ring-p2p-sizes-mb` +- `--intra-group-sizes 2,4,8`, `--inter-group-sizes 2,4,all` +- `--quick` preset (small warmup/iters, single message size) +- New flag-precedence rules: `--tests` / `--quick` imply `--perf-test`; mixing perf and info selectors warns and drops info; tuning knobs without perf intent are inert with a quieter warn +- Report improvements: Node/Rank columns compressed into ranges (e.g. `0-7`), a Node→Hostname legend at the top, "Leader hostname" column showing only the first host of each group + +This work is in `primus/tools/preflight/preflight_perf_test.py` and the comm modules. It set up the global accessors (`set_warmup` / `set_iteration` / `get_*`) that `node_smoke` later mirrored to keep iteration counts comparable. + +### 2. Why a separate node-local smoke test + +A user pointed out that for large jobs, what really matters is "which node has a problem", not which GPU. They proposed a per-node distributed-environment with simple checks (`set_device`, quick bandwidth/speed tests) that returns a success/fail flag per node, *before* the real training job opens its global rendezvous. This was the motivation for `node_smoke.py` — much faster, no dependency on a healthy cluster, and a stuck node can't take down its peers. + +### 3. Tiering decision + +Two tiers, picked interactively: + +- **Tier 1**: mandatory, fast (~5 s/GPU) — `set_device`, alloc, tiny GEMM, plus reused info collectors. The bar is "the GPU enumerated and runs ops". +- **Tier 2**: optional perf sanity (`--tier2 / --tier2-rccl`) — GEMM TFLOPS, HBM bandwidth, local 8-GPU RCCL. The bar is "the GPU is at expected steady-state performance". + +HBM bandwidth was clarified to mean device-to-device `torch.Tensor.copy_` of a 512 MB buffer, counting read + write, which gives ~70–80 % of the MI300X HBM3 roofline (~5300 GB/s) on healthy hardware. + +### 4. First scaled run + measurement-quality bug + +The first 6-node run produced: + +- GEMM 8192³ bf16: smoke median **724 TFLOPS**, full preflight median **765 TFLOPS** (~6 % gap) +- Local AR 64 MB / 8 GPU: smoke median **199 GB/s**, full preflight median **230 GB/s** (~12 % gap) + +Formula audit confirmed `2·S·(P-1)/P / t` and `2·N³/t` are identical to `intra_node_comm.py` and `square_gemm.py`. The systematic offset traced to **iteration counts being too low**. The original RCCL loop was `1 warmup + 1 timed iter` — basically a kernel-launch latency test, not a bandwidth test. Fixed by aligning to the preflight `--quick` preset: + +- GEMM: `warmup 3→5, iters 10→20` +- HBM: `warmup 5→10, iters 10→20` +- RCCL: `warmup 1→5, iters 1→20` + +Per-node runtime cost rounded to <1 s additional. Aggregator gained a "Tier 2 perf summary" section so per-node GEMM/HBM/RCCL outliers are visible without grepping JSONs. + +### 5. A + B + C (drift, NIC roll-call, host limits) + +Discussion of what the smoke test was *missing* led to six categories. A, B, C landed: + +- **A. Stack drift detection** — per-node `fingerprint` (kernel, ROCm, amdgpu, RCCL, torch, NIC firmware, HCA model) + aggregator-side cluster-majority-vs-outlier comparison. Catches "1 of N nodes on a different RCCL build" — a frequent cause of "job dies at minute 3". +- **B. NIC / RDMA roll-call** — sysfs-only inventory of `/sys/class/infiniband` (no `ibv_devinfo` dependency), per-port hard-fail rules for state ≠ ACTIVE, missing RoCE v2 GIDs, count mismatch (when `--expected-rdma-nics` is set). +- **C. Host limits** — `RLIMIT_MEMLOCK` (32 GiB default threshold), `/dev/shm` (8 GiB default threshold), plus collected-only NUMA / governor / kernel for drift detection. + +Verified with a synthetic two-node drift test (one real + one edited copy with mismatched RCCL, mismatched amdgpu, mismatched `rdma3` firmware, `rdma2:1` DOWN, and `memlock=64 MiB`): every section lit up correctly, `fail_reasons` were prefixed with `nic:` / `host_limits:` for traceability, exit code propagated. + +### 6. Aggregator crash on heterogeneous fingerprints + +An 18-node run on a different cluster crashed the aggregator with `TypeError: unhashable type: 'dict'`. Root cause: `_stack_drift_rows()` added a key to its scalar-comparison set whenever **any** node reported it as `None` (or a scalar), then iterated **all** nodes' values for that key into `Counter(...)`. On the failing cluster `nic_fw` was `None` on one node and a dict on others — the dict wasn't hashable. Fix: + +1. Only collect a key when at least one node reports it as a real scalar (drop the "None counts as scalar" path). +2. Defense-in-depth `isinstance(v, (str, int, float))` check inside the per-host loop. +3. Each report section wrapped in its own `try / except` so one section's bug can't truncate the rest of the report. +4. New "NIC port-count summary" section that always renders and lists nodes whose port count differs from the cluster majority (so partial-NIC degradation like 7-of-8 is visible without `--expected-rdma-nics`). + +### 7. Package split (refactor of the 4.5k-line monolith) + +`node_smoke.py` had grown to ~4500 lines with all collectors, the orchestrator, the per-GPU subprocess body, and the ~700-line aggregator markdown writer in a single file. The refactor turned it into a Python sub-package (`primus/tools/preflight/node_smoke/`) with one module per Tier 1 sub-section (`collectors/`), the per-GPU subprocess body, the orchestrator, and the aggregator's data shapers (`aggregator/summarizers.py`) and Markdown writer (`aggregator/report.py`, with one `_write_
` helper per `##` heading). The single public entry point — `main` — is re-exported from `__init__.py`, so the existing `python -m primus.tools.preflight.node_smoke ...` invocation (used by the primus-cli wrapper and by `_spawn_per_gpu` for per-GPU subprocesses) keeps working unchanged. Behavior parity was checked by diffing the per-node JSON and `smoke_report.md` against a baseline (with a small allowlist for run-variant fields like PIDs, hardware cycle counters, and `available_gb`/`free_gb`/`cached_gb`); CLI help text, JSON schema, report section order, and exit-code semantics for `run` / `_per_gpu` / `aggregate` are byte-identical to pre-refactor. + +### 8. Short hostnames + naming nodes that never reported + +`failing_nodes.txt` held FQDNs (`socket.gethostname()` returned the FQDN on the failing cluster) — not pipeable into `srun --exclude=`. Nodes that never produced a JSON only showed up as `` placeholders, so operators couldn't act on them. + +Fix: + +1. Normalize `host = socket.gethostname().split(".", 1)[0]` in `_cmd_run` for both the JSON filename and the `host` field; logs use the short name too. +2. Aggregator defensively short-normalizes every loaded JSON, so legacy FQDN files produce SLURM-ready txt outputs without re-running the smoke step. +3. New `aggregate --expected-nodelist-file FILE` flag — missing nodes appended with their real short hostname (and a self-describing `expected hostname '' from --expected-nodelist-file` reason), written to `failing_nodes.txt` directly. +4. Wrapper resolves `SLURM_JOB_NODELIST` via `scontrol show hostnames` into `/expected_nodes.txt` and forwards it to the aggregator. Best-effort: silent fallback to count-only behaviour when `scontrol` is unavailable. + +This also makes "the node that SLURM marked as `task X: unknown`" visible in the report under its real hostname. + +--- + +## Future work + +These were proposed but not yet built. In rough priority order: + +### D. GPU low-level health (beyond "alloc + small GEMM works") + +Reveals hardware that *enumerates* but is degraded. Most map to one `rocm-smi` query or one sysfs read in the existing per-GPU subprocess. + +- GPU count == expected and `lspci -d 1002:` agrees +- PCIe link width/speed per GPU (`/sys/bus/pci/devices//current_link_{speed,width}`) — catches "GPU at Gen3 x8 because the slot needs reseating" +- XGMI link matrix between every GPU pair (reuse `primus/tools/preflight/gpu/gpu_topology.py`) +- HBM size per GPU matches expected +- ECC counters: uncorrectable as hard fail, correctable as info-only with a cluster-median baseline +- GPU clock state: flag any GPU stuck at idle GFX clock (stale-state symptom) +- Throttle reasons from `rocm-smi --showperflevel` (`power_throttle` / `thermal_throttle`) +- Power cap drift across the cluster + +Aggregator gets a "GPU-level drift" section that pinpoints `host:gpu` outliers, not just node-level. + +### E. Time / cluster sync + +- Wall-clock skew vs `node_rank=0`: each node writes its `time.time()` into its JSON; aggregator computes `max - min` and warns at > 1 s, fails at > 5 s +- Time-daemon health (`systemctl is-active chronyd / ntpd / systemd-timesyncd`) + +### F. Storage / runtime liveness (site-specific) + +- Shared-FS latency probe: 1 KB write + `stat` to a unique path, aggregator flags nodes far above the cluster median (Lustre/NFS hiccups) +- Shared-FS quota / free space +- DNS resolution sanity for peer hostnames +- `rocm-smi --version` self-latency (5 s timeout) — catches drivers that have started to wedge but haven't crashed yet (we've seen 30–60 s `rocm-smi` calls precede a full GPU hang by minutes) +- Container / image hash drift — if the launcher exports `CONTAINER_IMAGE_TAG`, fold it into the existing fingerprint + +### Bigger architectural item (lower priority) + +- If `NODE_RANK==0` itself fails to start, no aggregator runs anywhere. Possible mitigations: separate aggregator step submitted after the smoke step, or polling watchdog on the submit host. Out of scope for now — handled in practice by always passing `--time=` to `srun` so SLURM force-terminates a stuck job and you can re-run the aggregator alone with `--aggregate-only` + `--expected-nodelist-file`. diff --git a/docs_deprecated/preflight-direct.md b/docs_deprecated/preflight-direct.md new file mode 100644 index 000000000..416ce4457 --- /dev/null +++ b/docs_deprecated/preflight-direct.md @@ -0,0 +1,794 @@ +# Run Preflight Without a Container + +> ⚠ **Run the [node-smoke test](./node-smoke-test-instruction.md) first.** `preflight` opens a global `torch.distributed` rendezvous, so a single sick node (wedged driver, leaked rank holding HBM, partial NIC enumeration, time-sync drift, etc.) can stall the whole job for up to `--dist-timeout-sec` seconds — long before any cross-node bandwidth number is produced. The node-smoke test catches those exact failure modes *without* a rendezvous in ~30–60 s and emits a SLURM-ready `failing_nodes.txt` you can pipe straight into `srun --exclude=`. Treat node-smoke as a hard prerequisite; only run `preflight` on the nodes node-smoke marked PASS. See [§0 "Which test should I run?"](#0-which-test-should-i-run) for the side-by-side comparison and the recommended 3-step workflow. + +This guide explains how to run Primus's `[preflight](./preflight.md)` cluster-diagnostic tool **directly on the host** (no Docker / Podman), via the standard Primus launcher. + +**Git clone the Primus repository to a shared filesystem that all nodes can read.** + +```bash +git clone --recurse-submodules https://github.com/AMD-AIG-AIMA/Primus.git +cd Primus +git checkout dev/preflight-direct-test +``` + +**Recommended (through the primus-cli SLURM wrapper):** + +``` +runner/primus-cli slurm srun -N --ntasks-per-node=1 -- direct -- preflight [PREFLIGHT_ARGS...] +``` + +**Equivalent (bare srun, useful when composing with custom srun flags):** + +``` +srun -N --ntasks-per-node=1 runner/primus-cli direct -- preflight [PREFLIGHT_ARGS...] +``` + +Both forms produce the **same workload** on the same ranks. The wrapper form is recommended because it auto-resolves `MASTER_ADDR` / `MASTER_PORT` / `NNODES` / `NODE_RANK` / `GPUS_PER_NODE` once on the launching node and passes them to every rank via `--env`, applies any `slurm.`* config defaults (partition / time / etc.) from your YAML, and is the same pattern used for `train` / `benchmark` / `node_smoke`. See [§ Wrapper vs. bare-srun](#wrapper-vs-bare-srun) below for the exact precedence / caveats. + +`primus-cli direct` activates an optional Python virtualenv (`VENV_ACTIVATE`), auto-derives the distributed environment variables (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, `MASTER_PORT`, `GPUS_PER_NODE`) from `SLURM_*` when running inside a SLURM allocation, and then launches the `preflight` Python subcommand via `torchrun` (one worker per GPU). It is the recommended entry point when: + +- You're running on a SLURM cluster but cannot (or don't want to) use the container-based path. +- Your nodes share a Python virtual environment on a network-mounted filesystem. +- You want a single-node sanity check with no extra configuration. + +--- + +## 0. Which test should I run? + +Primus ships **two** complementary cluster screens. Pick the right one — and ideally run them in this order. + + +| Aspect | `node-smoke` (start here) | `preflight` (this doc) | +| ----------------- | ---------------------------------------------------------------------------------- | ------------------------------------------------------------------- | +| Purpose | "Which nodes are healthy enough to run anything?" | "What is the actual cross-node performance on the surviving nodes?" | +| Rendezvous | None — every node independent | Global `torch.distributed` rendezvous | +| Wall clock | ~30–60 s for 6 nodes (Tier 1+2) | A few minutes; scales with N for inter-node tests | +| Granularity | Per-node PASS/FAIL | Per-rank perf measurements | +| Safety | A stuck node cannot wedge its peers | A single hung NIC can stall the whole rendezvous | +| Output | Per-node JSON + cluster md + SLURM-ready `passing_nodes.txt` / `failing_nodes.txt` | Markdown + PDF perf report | +| Entry point | `primus-cli direct -- node_smoke` | `primus-cli direct -- preflight` (this doc) | +| Quick-start guide | `[node-smoke-test-instruction.md](./node-smoke-test-instruction.md)` | This doc, §3+ | + + +### Recommended workflow + +> **Before running any of the commands below, complete the one-time setup:** +> +> 1. **Python virtualenv** on a shared filesystem — see [§2 Set up the Python virtual environment](#2-set-up-the-python-virtual-environment), then point the launcher at it via `export VENV_ACTIVATE=...` (details in [§2 → Tell the launcher where the venv is](#tell-the-launcher-where-the-venv-is)). +> 2. **NCCL / fabric environment variables** — usually the defaults in `base_env.sh` are fine, but multi-NIC nodes may need `NCCL_IB_HCA` / `NCCL_IB_GID_INDEX` / `NCCL_SOCKET_IFNAME` overrides. See [§4 Cluster-specific NCCL configuration](#4-cluster-specific-nccl-configuration) for known-good values per fabric (Broadcom, Pensando Pollara/AINIC). + +Through the `primus-cli slurm srun -- direct --` wrapper (recommended): + +```bash +# 1) Prune broken nodes with node-smoke (fast, no rendezvous). +runner/primus-cli slurm srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + -- direct -- node_smoke --tier2-perf + +# 2) Re-allocate excluding the bad nodes, and run preflight --quick +# for a fast cross-node sanity check. +runner/primus-cli slurm srun -N -c 128 --gpus-per-node=8 \ + --ntasks-per-node=1 \ + --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \ + -- direct -- preflight --quick + +# 3) Optional: full preflight on the same set if --quick numbers +# look off, or if you want the full bandwidth matrix. +runner/primus-cli slurm srun -N -c 128 --gpus-per-node=8 \ + --ntasks-per-node=1 \ + --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \ + -- direct -- preflight +``` + +Equivalent with bare `srun` (works identically; useful when scripting around custom srun flags that don't compose with the wrapper): + +```bash +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct -- node_smoke --tier2-perf + +srun -N -c 128 --gpus-per-node=8 --ntasks-per-node=1 \ + --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \ + runner/primus-cli direct -- preflight --quick + +srun -N -c 128 --gpus-per-node=8 --ntasks-per-node=1 \ + --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \ + runner/primus-cli direct -- preflight +``` + +Why this ordering matters: + +- A single broken node can stall a `torch.distributed.init_process_group()` for `--dist-timeout-sec` seconds (default 120), so feeding a known-good list to preflight is much faster. +- `node-smoke` catches things preflight cannot — leaked / foreign processes, wedged drivers, partial NIC enumeration, time-sync drift, RDMA roll-call issues — that produce *misleading* preflight failures. +- `preflight --quick` adds the cross-node bandwidth signal that `node-smoke` deliberately does not measure. + +--- + +## 1. Prerequisites + +- A working AMD ROCm installation on every node. +- Network reachability between nodes (Ethernet for bootstrap, RDMA / InfiniBand recommended for perf tests). +- A Python ≥ 3.10 virtual environment **on a shared filesystem** that all nodes can read (the same path is sourced on every node). +- The Primus repository checked out somewhere readable from every node. + +--- + +## 2. Set up the Python virtual environment + +The environment must live on a path visible from every node (e.g. NFS-mounted home, Lustre, or any shared filesystem). All nodes will `source` the same activation script. + +You can use any tool you like; `uv` is the fastest. Either of the following works. + +### What you actually need to install + +The `preflight` and `node-smoke` tools deliberately use **only a small subset** of Primus's full dependency tree. You do **not** need to install the entire `requirements.txt` — that pulls in trainer / dataset / experiment-tracking packages that neither tool ever imports. + + +| Package | Required for | Skip when | +| -------------------- | ------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ | +| `torch` (ROCm build) | Both tools — perf measurements (`torch.matmul`, `torch.distributed`, `torch.cuda.`*). | Never (mandatory). | +| `markdown2` | `preflight` PDF report only (Markdown → HTML). | You always pass `--disable-pdf`, or you only run `node-smoke` (which never produces PDFs). | +| `weasyprint` | `preflight` PDF report only (HTML → PDF). | Same as above. | +| `matplotlib` | `preflight --plot` only (per-test bandwidth bar charts). | You don't pass `--plot`. | + + +Everything else in the preflight / node-smoke code path is Python stdlib (`os`, `subprocess`, `socket`, `argparse`, `dataclasses`, `json`, `time`, ...) — no extra installs needed. + +### Option A — `uv` (recommended), minimal install + +```bash +mkdir -p ~/envs/preflight +cd ~/envs/preflight + +uv venv --python 3.12 +source .venv/bin/activate + +# 1) ROCm-built PyTorch (pin to your ROCm version; rocm7.1 shown here) +uv pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.1 --no-cache-dir + +# 2) Optional: only if you want preflight PDF reports (omit to use --disable-pdf) +uv pip install markdown2 weasyprint + +# 3) Optional: only if you want preflight --plot bar charts +uv pip install matplotlib +``` + +### Option B — `python -m venv`, minimal install + +```bash +mkdir -p ~/envs/preflight +python3.12 -m venv ~/envs/preflight/.venv +source ~/envs/preflight/.venv/bin/activate + +pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.1 --no-cache-dir +pip install markdown2 weasyprint # optional, for preflight PDFs +pip install matplotlib # optional, for preflight --plot +``` + +### Option C — full Primus runtime (only if you also want the rest of Primus) + +```bash +cd /path/to/Primus +uv pip install -r requirements.txt # or: pip install -r requirements.txt +``` + +This installs every Primus runtime dependency (trainer, dataset loaders, experiment trackers, ...). Use only if you're going to run more than just preflight / node-smoke from this environment. + +### Per-tool minimum install matrix + +If you want the absolute smallest footprint, install only what your intended invocations need: + + +| Invocation | `torch` | `markdown2` | `weasyprint` | `matplotlib` | +| ------------------------------------------------ | -------- | --------------------------------- | --------------------------------- | ------------ | +| `node-smoke` (any flags) | required | — | — | — | +| `preflight --host --gpu --network --disable-pdf` | required | — | — | — | +| `preflight --host --gpu --network` (with PDF) | required | required | required | — | +| `preflight --quick --disable-pdf` | required | — | — | — | +| `preflight --quick` (with PDF) | required | required | required | — | +| `preflight ... --plot` | required | required (unless `--disable-pdf`) | required (unless `--disable-pdf`) | required | + + +### Tell the launcher where the venv is + +`primus-cli direct` reads the `**VENV_ACTIVATE**` environment variable. When set, it sources the path before launching the Python process; when unset, it is a no-op (the container path, which uses the container's bundled Python, leaves this unset): + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate +``` + +`VENV_ACTIVATE` is the only optional environment variable specific to the direct flow. Everything else has a sensible default; distributed-env variables (`NNODES`, `NODE_RANK`, `MASTER_ADDR`, ...) are auto-derived from SLURM when not pre-exported. + +--- + +## 3. Run preflight + +### Single node (no SLURM) + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate + +# Info report only (fast) +runner/primus-cli direct -- preflight --host --gpu --network + +# Info + perf report +runner/primus-cli direct -- preflight + +# Perf report only +runner/primus-cli direct -- preflight --perf-test +``` + +When SLURM is not detected the script defaults to `NNODES=1`, `NODE_RANK=0`, `MASTER_ADDR=localhost`. Any of those can be overridden by exporting them before calling the script. + +### Multi-node without SLURM (parallel SSH) + +When no scheduler is available (bare-metal, cloud VMs, lab nodes), launch +`primus-cli direct` on each node yourself via SSH. The script works +identically — you just pre-export the distributed variables that SLURM +would normally provide. + +#### Requirements + +- All nodes share the same filesystem (or at least the same Primus checkout + venv path). +- Nodes can reach each other on a **data-plane** network interface (not the management NIC). +- SSH key-based access to each node from the launching host. + +#### Required environment variables + + +| Variable | Description | +| -------------------- | ---------------------------------------------------------------------- | +| `NNODES` | Total number of nodes | +| `NODE_RANK` | This node's rank (`0` through `NNODES-1`) | +| `MASTER_ADDR` | IP of rank-0 node **on the data-plane interface** | +| `MASTER_PORT` | Rendezvous port (default `1234`; increment between concurrent runs) | +| `GPUS_PER_NODE` | GPUs per node (default `8`) | +| `NCCL_SOCKET_IFNAME` | Data-plane NIC name (e.g. `enp159s0np0`) — **critical for multi-node** | +| `GLOO_SOCKET_IFNAME` | Same as `NCCL_SOCKET_IFNAME` | +| `VENV_ACTIVATE` | Path to virtualenv `activate` script | + + +> **Warning**: `NCCL_SOCKET_IFNAME` auto-detection often picks a management interface +> (e.g. `enp28s0np0`, `eno8303`) instead of the high-bandwidth data NIC. For multi-node +> runs this causes `init_process_group` to hang or NCCL to fail silently. Always set it +> explicitly. + +#### Identifying the correct data-plane interface + +```bash +# On any node, find the interface whose IP matches the MASTER_ADDR subnet: +ip -4 addr show | grep "10.245.134" +# → enp159s0np0 inet 10.245.134.129/24 + +# Or check which interface routes to the master: +ip route get 10.245.134.129 | awk '{print $5; exit}' +### Multi-node via SLURM + +`primus-cli direct` auto-detects a SLURM allocation (via `SLURM_JOB_ID`) and derives all distributed variables from `SLURM_*` automatically. **Pre-exported values always win**, so the same launcher script also works inside the `primus-cli slurm srun ... -- direct -- ...` chain (where `slurm-entry` has already set these via `--env`): + +| Variable | Resolved as | +| --------------- | -------------------------------------------------------------------- | +| `NNODES` | `NNODES` → `SLURM_NNODES` → `SLURM_JOB_NUM_NODES` → `1` | +| `NODE_RANK` | `NODE_RANK` → `SLURM_NODEID` → `SLURM_PROCID` → `0` | +| `MASTER_ADDR` | `MASTER_ADDR` (if not empty / not `localhost`) → first hostname from `scontrol show hostnames "$SLURM_NODELIST"` | +| `MASTER_PORT` | `MASTER_PORT` → `1234` | +| `GPUS_PER_NODE` | `GPUS_PER_NODE` → `8` | + +Run it as a single task per node (the script invokes `torchrun` internally, which spawns one worker per GPU): + +> **Verify NCCL / network env first.** The script sets sensible `NCCL_`* defaults via `base_env.sh`, but auto-detection can pick the wrong device on multi-NIC nodes. Always confirm `NCCL_IB_HCA`, `NCCL_IB_GID_INDEX`, `NCCL_SOCKET_IFNAME`, and `GLOO_SOCKET_IFNAME` (set to the same value as `NCCL_SOCKET_IFNAME`) are correct for your fabric, and `export` overrides before running. See [§4](#4-cluster-specific-nccl-configuration) for cluster-specific values. + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate + +# export NCCL_IB_HCA=rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 +# export NCCL_IB_GID_INDEX=3 +# export NCCL_SOCKET_IFNAME=eno0 +# export GLOO_SOCKET_IFNAME=eno0 + +# Recommended: through the primus-cli SLURM wrapper. +runner/primus-cli slurm srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 \ + --nodelist --ntasks-per-node=1 \ + -- direct -- preflight --perf-test + +# Or, equivalently, with bare srun: +srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist \ + --ntasks-per-node=1 \ + runner/primus-cli direct -- preflight --perf-test +``` + + + +#### Wrapper vs. bare-srun + +Both forms target the **same** `primus-cli-direct.sh` launcher and produce identical workloads. The difference is only in how the SLURM context is constructed: + + +| Aspect | `primus-cli slurm srun -- direct --` (recommended) | Bare `srun ... primus-cli direct --` | +| ---------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- | +| `MASTER_ADDR` resolution | Resolved **once** on the launching node via `scontrol show hostnames "$SLURM_NODELIST" | head -n1`, then propagated to every rank via `--env MASTER_ADDR=...`. | Each rank re-derives it inside `primus-cli-direct.sh` STEP 4.7 from `SLURM_`* (same result, more `scontrol` calls). | +| `NNODES` / `NODE_RANK` / `GPUS_PER_NODE` | Set explicitly by `slurm-entry.sh` via `--env`. | Derived from `SLURM_NNODES` / `SLURM_NODEID` / `SLURM_PROCID` inside `direct.sh`. | +| `slurm.*` config defaults | Honored (partition, time, ntasks-per-node, etc. from the active YAML). | Not consulted — you pass every flag explicitly to `srun`. | +| Default wall-time | `-t 4:00:00` is auto-added if you don't pass `--time`. | None — `srun` uses the cluster default (may reject the job). | +| `direct` keyword | **Required**: `primus-cli slurm srun ... -- direct -- `. Without `direct`, the wrapper routes through the **container** path. | N/A — there's only one path. | +| `--ntasks-per-node=1` | **Not auto-added**. Pass it on the CLI (before the first `--`) or set it in the `slurm.`* config. | **Not auto-added**. Pass it as an `srun` flag. | +| Best for | Production / repeatable runs. Same pattern as `train` / `benchmark` / `node_smoke`. | Ad-hoc runs where you want to compose with arbitrary `srun` flags (`--nodelist=$(...)`, `--exclude=...` from a runtime file, etc.). | + + +For the rest of this doc the examples use bare `srun` for brevity, but every example also works with the wrapper form by substituting `srun runner/primus-cli direct --` → `runner/primus-cli slurm srun -- direct --`. + +### Key `srun` flags + + +| Flag | Why it's necessary | +| --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `-c 128` | Allocate all CPU cores per task. Without this, SLURM may default to 1 core, which starves the RCCL network proxy threads and can cause >30× slowdown on perf tests. Set this to your node's core count. | +| `--gpus-per-node=8` | Grants GPU device access (`/dev/kfd`, `/dev/dri`). Required for non-container execution. | +| `--ntasks-per-node=1` | One launcher invocation per node; `primus-cli direct` then spawns 8 workers per node via `torchrun`. | +| `-t 00:45:00` | Wall-clock limit. Full perf tests on 8N usually finish well under 10 min. | + + +> Tip — check core count: `srun -N 1 --gpus-per-node=8 bash -c 'nproc'` + +--- + +## 4. Cluster-specific NCCL configuration + +`primus-cli direct` sources `runner/helpers/envs/base_env.sh`, which sets sensible defaults for `NCCL_`* and auto-detects `NCCL_IB_HCA` / `NCCL_SOCKET_IFNAME`. Pre-exported values from your shell take precedence, so the standard pattern is: + +```bash +export VAR=value +runner/primus-cli direct -- preflight ... +``` + +### Broadcom NICs (no AINIC) + +Most clusters fall here. The defaults from `base_env.sh` are usually fine, but the two values most commonly worth overriding are: + +```bash +export NCCL_CROSS_NIC=1 # default in base_env.sh is 0 +export NCCL_PXN_DISABLE=0 # default in base_env.sh is 1 + +srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist \ + --ntasks-per-node=1 \ + runner/primus-cli direct -- preflight --perf-test +``` + +### Pensando Pollara (AINIC) RDMA + +```bash +export USING_AINIC=1 +export NCCL_IB_GID_INDEX=1 # AINIC uses index 1 (default in base_env.sh is 3) +export NCCL_PXN_DISABLE=0 + +srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --nodelist \ + --ntasks-per-node=1 \ + runner/primus-cli direct -- preflight +``` + +> `primus-cli direct` *does* accept `--env KEY=VALUE` on its own command line (placed before `--`), in addition to the conventional `export`/`srun --export=` approaches. + +--- + +## 5. Launcher flags vs. preflight flags + +Anything you place **after** the `--` separator is forwarded verbatim to the `preflight` Python tool. The launcher (`primus-cli-direct.sh`) consumes a small set of flags **before** `--`. The one most users care about is `--silent`. + +### Launcher-only flags (before `--`) + + +| Flag | Effect | +| ----------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `--silent` | Back-pocket knob: redirect the launcher's and the Python tool's `stdout` to `/dev/null`. Launcher errors (`LOG_ERROR` / `LOG_WARN`, written to `stderr`) are preserved so real failures still surface; the log file under `logs/` captures everything. Exit code is propagated unchanged. **Not recommended** for normal use — you lose live progress; prefer the log file. | +| `--debug` | Verbose launcher logging (`PRIMUS_LOG_LEVEL=DEBUG`). Forwarded to the Python tool as `--debug` too. | +| `--dry-run` | Print the resolved configuration and final `torchrun` / `python3` command without executing. | +| `--single` | Force `python3` instead of `torchrun`. `node_smoke` auto-selects this; for `preflight` you usually want the default (`torchrun`). | +| `--env KEY=VALUE` | Inject an env var into the Python process (in addition to anything `export`-ed in the shell). | +| `--log_file PATH` | Redirect the captured tee log to a specific path (default: `logs/log_.txt`). | + + +See `runner/primus-cli direct --help` for the full set. + +### Forwarded `preflight` flags (after `--`, most common) + +See [Preflight](./preflight.md) for the full list. The most common are: + +- Mode selection: `--host`, `--gpu`, `--network`, `--perf-test`, `--tests`, `--quick` +- Test tuning: `--comm-sizes-mb`, `--intra-comm-sizes-mb`, `--inter-comm-sizes-mb`, `--intra-group-sizes`, `--inter-group-sizes`, `--ring-p2p-sizes-mb` +- Reporting: `--dump-path`, `--report-file-name`, `--disable-pdf`, `--plot` +- Reliability: `--comm-cleanup-delay-sec`, `--dist-timeout-sec` + +If you do not pass `--report-file-name`, `preflight` auto-generates a unique one of the form: + +``` +preflight-${NNODES}N-YYYYMMDD-HHMMSS +``` + +This guarantees that each run lands in its own files and prevents stale leftovers from earlier runs from being mistaken for fresh output. The auto-name logic now lives in the Python tool itself, so every call site (host `srun ... primus-cli direct`, `primus-cli slurm ... -- direct`, `primus-cli slurm ... -- container`) gets the same fresh name. + +### Examples + +The examples below all assume one of the two equivalent shell-prefix conventions. Pick whichever matches your habits — every example block in this section works with either definition: + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate + +# Recommended: through the primus-cli SLURM wrapper. Auto-resolves +# MASTER_ADDR/NNODES/NODE_RANK once on the launching node and propagates +# them via --env; honors slurm.* config defaults. +SRUN="runner/primus-cli slurm srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --ntasks-per-node=1 --nodelist --" +# Then in every example below, replace `$SRUN runner/primus-cli direct --` +# with just `$SRUN direct --`. (The wrapper expects the entry-mode keyword +# `direct` as the first token after the inner `--`.) + +# Equivalent: bare srun. NNODES/NODE_RANK/MASTER_ADDR get derived inside +# primus-cli-direct.sh's STEP 4.7 directly from SLURM_*; same net effect. +SRUN="srun -t 00:45:00 -N 4 -c 128 --gpus-per-node=8 --ntasks-per-node=1 --nodelist " +``` + +The examples in this section use the **bare-srun** form below for brevity (since `$SRUN runner/primus-cli direct -- preflight` reads naturally as one command line). To use the wrapper form instead, substitute `$SRUN runner/primus-cli direct --` → `$SRUN direct --` after exporting `SRUN` to the wrapper variant. + +#### A. Mode selection + +```bash +# Default: info report + every perf test +$SRUN runner/primus-cli direct -- preflight + +# Info-only (fast, no torch.distributed rendezvous) +$SRUN runner/primus-cli direct -- preflight --host --gpu --network --disable-pdf + +# Perf-only, every test +$SRUN runner/primus-cli direct -- preflight --perf-test + +# Fast pre-launch sanity preset (gemm + intra-AR + inter-AR @ 64,1024 MB, +# full intra-node group, full N-node inter group, low warmup/iter) +$SRUN runner/primus-cli direct -- preflight --quick +``` + +> **Note**: Mixing perf-mode flags (`--perf-test` / `--tests` / `--quick`) with info selectors (`--host` / `--gpu` / `--network`) makes preflight drop the info selectors with a `WARN`. Run two invocations if you want both reports. + +#### B. Test selection (`--tests`) + +```bash +# Only GEMM +$SRUN runner/primus-cli direct -- preflight --tests gemm + +# Only the inter-node bandwidth tests +$SRUN runner/primus-cli direct -- preflight --tests inter-allreduce,inter-alltoall + +# Only the inter-node ring P2P +$SRUN runner/primus-cli direct -- preflight --tests inter-ring-p2p + +# Combine: GEMM + inter-AR with overridden sizes/groups +$SRUN runner/primus-cli direct -- preflight \ + --tests gemm,inter-allreduce \ + --comm-sizes-mb 64,1024 \ + --inter-group-sizes all +``` + +Valid `--tests` tokens: `gemm`, `intra-allreduce`, `intra-alltoall`, `inter-allreduce`, `inter-alltoall`, `inter-p2p`, `inter-ring-p2p`, `all`. Unknown tokens fail fast (before NCCL init). + +#### C. Message sizes + +```bash +# One CSV applied to both intra and inter +$SRUN runner/primus-cli direct -- preflight --tests intra-allreduce,inter-allreduce \ + --comm-sizes-mb 8,128 + +# Different sizes for intra vs inter (override wins over --comm-sizes-mb) +$SRUN runner/primus-cli direct -- preflight --tests intra-allreduce,inter-allreduce \ + --comm-sizes-mb 8,128 --intra-comm-sizes-mb 4,32 + +# Inter-only override (also covers inter-p2p when enabled) +$SRUN runner/primus-cli direct -- preflight --tests inter-allreduce,inter-p2p \ + --comm-sizes-mb 8,128 --inter-comm-sizes-mb 16,512 +``` + +#### D. Group sizes + +```bash +# Custom intra-node group sizes (each must divide LOCAL_WORLD_SIZE) +$SRUN runner/primus-cli direct -- preflight \ + --tests intra-allreduce \ + --intra-group-sizes 4,8 + +# Custom inter-node groups: 2-node pairs and the full N-node group +$SRUN runner/primus-cli direct -- preflight \ + --tests inter-allreduce \ + --inter-group-sizes 2,all +``` + +> Note: for `inter-alltoall` only, every requested per-group node count is internally clamped to **16** (real-world MoE training rarely dispatches across more nodes; see `[preflight.md` §5.2](./preflight.md#52-group-sizes) for the rationale). The other inter-node tests use the requested sizes unchanged. So on a 128-node cluster, `--tests inter-alltoall --inter-group-sizes all` actually runs at 16-node sub-groups, while `--tests inter-allreduce --inter-group-sizes all` runs at 128 nodes as written. + +#### E. Ring P2P sizes + +```bash +$SRUN runner/primus-cli direct -- preflight \ + --tests inter-ring-p2p \ + --ring-p2p-sizes-mb 5,20,80 +``` + +#### F. Plotting + +```bash +# Generate per-test bandwidth bar charts under //*.png +$SRUN runner/primus-cli direct -- preflight \ + --tests intra-allreduce,inter-allreduce --plot +``` + +#### G. Reliability knobs + +```bash +# Bump the per-phase cleanup delay. Default 2.0 is sufficient at every +# cluster size for the comm shapes preflight exercises (inter-alltoall +# is internally capped at 16 nodes; see preflight.md §5.2). Only bump +# this on very flaky networks or unusual kernel TIME_WAIT settings. +$SRUN runner/primus-cli direct -- preflight --quick --comm-cleanup-delay-sec 5 + +# Fail fast if torch.distributed rendezvous can't complete in 30s +$SRUN runner/primus-cli direct -- preflight --perf-test --dist-timeout-sec 30 +``` + +> Operating clusters at ≥ 128 nodes? See `[preflight.md` §7](./preflight.md#7-running-on-very-large-clusters--64-nodes) for the recommended OS-level tunings (`tcp_tw_reuse`, wider `ip_local_port_range`) and per-test invocation patterns. With the §5.2 inter-alltoall cap in place, a default invocation is safe at every cluster size; the §7.2 sysctls remain best-practice for any RDMA workload. + +#### H. Reporting & output layout + +```bash +# Quick info-only check on 4 nodes, no PDF +$SRUN runner/primus-cli direct -- preflight --host --gpu --network --disable-pdf \ + --report-file-name info-4N + +# Perf test only, silenced (CI-friendly), explicit name. Note that --silent +# is consumed by primus-cli-direct.sh and must appear BEFORE the `--` +# separator; everything after `--` is forwarded to the preflight Python tool. +$SRUN runner/primus-cli direct --silent -- preflight --perf-test \ + --report-file-name nightly-4N-perf + +# Archive each run under its own directory +$SRUN runner/primus-cli direct -- preflight --quick \ + --dump-path /shared/preflight-archive/$(date +%Y%m%d-%H%M%S) +``` + +#### I. Backward-compat aliases + +These still work and are equivalent to flags above. Use them only when retrofitting older scripts. + +```bash +# Same as --host --gpu --network +$SRUN runner/primus-cli direct -- preflight --check-host --check-gpu --check-network + +# Same as --inter-group-sizes all AND drops inter-p2p +$SRUN runner/primus-cli direct -- preflight --perf-test --no-split-nodes-subgroup +``` + +#### J. Combined "production-ready" pre-launch screen + +```bash +# 1) Smoke first to prune broken nodes (note: --silent goes BEFORE `--`) +srun -N "$SLURM_NNODES" --ntasks-per-node=1 \ + runner/primus-cli direct --silent -- node_smoke --tier2-perf + +# 2) Quick perf sanity on the survivors +srun -N -c 128 --gpus-per-node=8 --ntasks-per-node=1 \ + --exclude=$(paste -sd, output/preflight/failing_nodes.txt) \ + runner/primus-cli direct --silent -- preflight --quick \ + --comm-cleanup-delay-sec 5 --dist-timeout-sec 60 \ + --report-file-name screen-$(date +%Y%m%d-%H%M%S) +``` + +--- + +## 6. Outputs + +Reports are written to `--dump-path` (default: `output/preflight/`), with the basename from `--report-file-name` and a `_perf` suffix for performance reports: + + +| File | Produced by | Notes | +| ----------------- | ----------------------------------------------- | ------------------------- | +| `.md` | `--host --gpu --network` (or default selection) | Info report | +| `.pdf` | same, unless `--disable-pdf` | Info report PDF | +| `_perf.md` | `--perf-test` | Perf report (GEMM + comm) | +| `_perf.pdf` | same, unless `--disable-pdf` | Perf report PDF | + + +Only **rank 0** writes the report. After preflight completes, the Python tool prints the absolute path of every report file it produced to stdout. Under `--silent` these prints go to `/dev/null` along with everything else (one of the trade-offs of using `--silent`); without `--silent` the announcement is visible live. Sample output: + +``` +[Primus:Preflight] Report: /home/.../Primus/output/preflight/preflight-4N-20260428-201925.md +[Primus:Preflight] Report: /home/.../Primus/output/preflight/preflight-4N-20260428-201925_perf.md +``` + +--- + +## 7. Environment variable reference + +Variables read by `primus-cli direct` itself: + + +| Variable | Required | Default | Purpose | +| --------------- | -------- | ---------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------- | +| `VENV_ACTIVATE` | no | — | Path to the venv `bin/activate` script. Unset = no-op (use system / container Python). Set + missing file = fail-fast. | +| `NNODES` | no | `1` (or auto-derived from `SLURM_NNODES` / `SLURM_JOB_NUM_NODES`) | Number of nodes. Pre-exported always wins. | +| `NODE_RANK` | no | `0` (or auto-derived from `SLURM_NODEID` / `SLURM_PROCID`) | This node's rank. Pre-exported always wins. | +| `GPUS_PER_NODE` | no | `8` | GPUs per node | +| `MASTER_ADDR` | no | `localhost` (or first host from `scontrol show hostnames "$SLURM_NODELIST"`) | Rendezvous host. Pre-exported always wins. | +| `MASTER_PORT` | no | `1234` | Rendezvous port | + + +Variables consumed downstream by `primus-cli direct` / `base_env.sh` (set them via `export`): + + +| Variable | Default in `base_env.sh` | When to override | +| -------------------- | ------------------------ | ------------------------------------------------- | +| `NCCL_SOCKET_IFNAME` | auto-detected | Force a specific Ethernet interface for bootstrap | +| `NCCL_IB_HCA` | auto-detected | Force specific RDMA HCAs | +| `NCCL_IB_GID_INDEX` | `3` | `1` on AINIC clusters | +| `NCCL_CROSS_NIC` | `0` | `1` for multi-rail IB fabrics | +| `NCCL_PXN_DISABLE` | `1` | `0` to enable PXN multi-hop NIC sharing | +| `USING_AINIC` | unset | `1` on Pensando Pollara clusters | +| `NCCL_DEBUG` | unset | `INFO` for verbose NCCL logging | + + +--- + +## 8. Troubleshooting + +### `[ERROR] [direct] VENV_ACTIVATE is set but file does not exist: ...` + +`VENV_ACTIVATE` was set in the environment but the path it points at doesn't exist on this node. This is a fail-fast guard to prevent a silent fallback to system Python (which usually has the wrong `torch` / no ROCm). Either fix the path: + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate +``` + +… or unset it to fall back to the container / system Python: + +```bash +unset VENV_ACTIVATE +``` + +If the path looks right but the file still appears missing, confirm the venv lives on a filesystem visible from the node SLURM scheduled you onto. + +### `[Primus:Preflight] FAIL: No GPUs detected` + +The Python process inside the venv can't find ROCm. Diagnose with: + +```bash +srun --nodes=1 --nodelist= bash -c ' +echo "=== PATH ==="; echo $PATH +echo "=== LD_LIBRARY_PATH ==="; echo $LD_LIBRARY_PATH +echo "=== rocm-smi ==="; rocm-smi --showid 2>&1 +echo "=== Python torch check ===" +source ~/envs/preflight/.venv/bin/activate +python3 -c "import torch; print(\"hip:\", torch.version.hip); print(\"available:\", torch.cuda.is_available()); print(\"count:\", torch.cuda.device_count())" +' +``` + +If `LD_LIBRARY_PATH` is empty, set it explicitly: + +```bash +export LD_LIBRARY_PATH=/opt/rocm/lib:${LD_LIBRARY_PATH:-} +``` + +### Report announcement points at stale files + +This shouldn't happen with the current Python tool — the auto-generated unique report name (`preflight-${NNODES}N-`) ensures every run gets a fresh path. If you explicitly pass `--report-file-name X`, you're responsible for choosing a name that doesn't collide with prior runs. + +### Slow perf tests (~30× expected) + +Almost always a symptom of insufficient CPU cores. Pass `-c ` to `srun` so RCCL's network proxy threads have CPU to spawn on. Verify with `srun -N 1 --gpus-per-node=8 bash -c 'nproc'`. + +### Using `conda` instead of venv + +`primus-cli direct` does `source "$VENV_ACTIVATE"`, which works for venv/uv but not directly for conda. Two options: + +1. Create a venv inside the conda env and point `VENV_ACTIVATE` at that venv's activate script. +2. Write a small shim activate script (e.g. `~/envs/conda-shim.sh`) that activates conda and the desired env, then point `VENV_ACTIVATE` at it: + ```bash + # ~/envs/conda-shim.sh + source "$HOME/miniconda3/etc/profile.d/conda.sh" + conda activate + ``` + +### "Address already in use" during perf tests + +This error occurs when peak simultaneous ESTAB sockets per node during an `ncclCommInit` exhausts the kernel ephemeral-port pool, so the next outgoing `bind()` walks the entire range without finding an allocatable port. (Despite the name and the `TIME_WAIT` framing in the kernel docs, accumulated `TIME_WAIT` count does *not* gate this for NCCL inter-node OOB — see `[preflight.md` §7.1](./preflight.md#71-why-address-already-in-use-used-to-surface-at-scale) for the mechanism and the empirical evidence.) + +Preflight has two complementary defenses: + +1. The **inter-node alltoall sub-group is internally capped at 16 nodes** (see `[preflight.md` §5.2](./preflight.md#52-group-sizes)) — the only test that, at large scale, can push peak ESTAB anywhere near the per-node ephemeral pool. The cap eliminates this failure mode by construction. +2. A **global barrier + `--comm-cleanup-delay-sec` sleep** (default 2 s) is inserted after every comm destroy, primarily for cross-rank synchronization across the destroy → setup transition. + +If you still see `Address already in use` (e.g. on a network with an unusually narrow ephemeral-port range), the directly relevant **OS-level tuning** is widening that range — best-practice for any RDMA host: + +```bash +# Widen the ephemeral port range from ~28k to ~64k. This is the only +# OS knob that directly addresses the binding constraint. +sudo sysctl -w net.ipv4.ip_local_port_range="1024 65535" + +# General hygiene for hosts running mixed RDMA + repeated outgoing +# TCP workloads (NCCL inter-node OOB by itself doesn't benefit from +# this -- see preflight.md §7.2 for why). +sudo sysctl -w net.ipv4.tcp_tw_reuse=1 +``` + +As a fallback, raise the per-phase delay: + +```bash +# Bump the per-phase delay (default 2 s) on a particularly stressed +# network. Rarely needed in practice with the §5.2 alltoall cap. +runner/primus-cli direct -- preflight --comm-cleanup-delay-sec 5 +``` + +See `[preflight.md` §7](./preflight.md#7-running-on-very-large-clusters--64-nodes) for the full explanation, persistence, and recommended large-cluster invocation patterns (split tests into separate runs, etc.). + +If the error occurs at `init_process_group` (before tests even start), it typically means a previous job left port 29500 in `TIME_WAIT`. Either wait ~60 s or use a different port: + +```bash +export MASTER_PORT=29501 +``` + +### Capturing full output + +The launcher already writes a complete log to `logs/log_.txt` (configurable via `--log_file PATH`), even under `--silent`. If you also want a copy at the call site, redirect there: + +```bash +srun ... runner/primus-cli direct -- preflight --perf-test \ + 2>&1 | tee preflight-$(date +%Y%m%d-%H%M%S).log +``` + +--- + +## 9. Automated node bisection (finding the bad node in an NCCL hang) + +When a cluster-wide preflight run hangs or fails, use +`[tools/preflight_bisect/bisect.py](../tools/preflight_bisect/bisect.py)` to +run `preflight --perf-test` on smaller Slurm node subsets until suspect nodes +are isolated. + +### Prerequisites + +1. Working non-container preflight setup from the sections above, with + `VENV_ACTIVATE` exported from a shared filesystem path. +2. Run from the SLURM login/head node, where both `scontrol` and `srun` are + available. +3. Run from inside a Slurm allocation, or provide a Slurm nodelist explicitly. + +### Example from inside an allocation + +```bash +export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate +mkdir -p output + +python tools/preflight_bisect/bisect.py \ + --nodelist "$SLURM_NODELIST" \ + --output-dir "output/bisect-$(date +%Y%m%d-%H%M%S)" \ + --trial-timeout-sec 600 \ + --slurm-time 00:15:00 \ + --preflight-env USING_AINIC=1 \ + --preflight-env NCCL_IB_GID_INDEX=1 \ + --preflight-env NCCL_CROSS_NIC=1 \ + --preflight-env NCCL_PXN_DISABLE=0 \ + 2>&1 | tee output/bisect-latest.log +``` + +Adjust the `--preflight-env` lines to match your cluster. Per-trial logs and a +final `summary.txt` are written under `--output-dir`. + +> Note: Set `--trial-timeout-sec` high enough for a healthy subset to finish. +> Too small a timeout can turn slow-but-good trials into false failures, causing +> the bisection to explore extra paths. +> +> Note: `--preflight-env KEY=VALUE` values are concatenated into a single +> `srun --export=ALL,...` argument, so values must not contain commas or +> whitespace. Keep comma-containing values as normal exported environment +> variables. + +--- + +## 10. See also + +- [Preflight](./preflight.md) — full reference for the `preflight` subcommand and its flags +- [CLI User Guide](./cli/PRIMUS-CLI-GUIDE.md) — container-based and `primus-cli slurm` workflows +- `[runner/primus-cli-direct.sh](../runner/primus-cli-direct.sh)` — the direct launcher itself (`primus-cli direct` dispatches here) +- `[primus/tools/preflight/](../primus/tools/preflight/)` — preflight implementation +- `[tools/preflight_bisect/bisect.py](../tools/preflight_bisect/bisect.py)` — bisect wrapper for narrowing down failing nodes in multi-node preflight runs diff --git a/docs_deprecated/preflight.md b/docs_deprecated/preflight.md index b4290e04e..7f3e90c3f 100644 --- a/docs_deprecated/preflight.md +++ b/docs_deprecated/preflight.md @@ -1,19 +1,51 @@ # Preflight -`preflight` is Primus’ cluster diagnostic tool. It can generate a **fast info report** (host/GPU/network) and can also run **performance tests** (GEMM + intra/inter-node comm) to help spot misconfiguration or outliers before large distributed runs. +`preflight` is Primus' cluster diagnostic tool. It produces: -- **User-facing entry**: `primus-cli … -- preflight [args]` -- **Implementation entrypoint**: `primus/cli/subcommands/preflight.py` +- A **fast info report** (host / GPU / network configuration), and +- A configurable suite of **performance tests** (GEMM TFLOPS, intra-node and inter-node communication bandwidth, P2P, ring P2P). -## Quick start +Use it to spot misconfiguration, hardware degradation, or perf outliers **before** committing a large distributed training run to a global rendezvous. -### Info report only (fast) +- **User-facing entry**: `primus-cli ... -- preflight [args]` +- **No-container launcher**: `runner/primus-cli direct -- preflight ...` — see [`preflight-direct.md`](./preflight-direct.md). +- **Implementation entrypoint**: `primus/cli/subcommands/preflight.py` → `primus/tools/preflight/preflight_perf_test.py`. + +> Looking for a faster, distributed-rendezvous-free per-node screen? See [`node-smoke.md`](./node-smoke.md) (and the [quick-start guide](./node-smoke-test-instruction.md)). The recommended workflow is **smoke first, preflight second** — see [§10 Comparison with node-smoke](#10-comparison-with-node-smoke). + +--- + +## 1. Two run modes (and how preflight picks one) + +Preflight has two report types, controlled by a single precedence rule: + +| Mode | Triggered by | What it does | +|---|---|---| +| **Info-only** | `--host`, `--gpu`, `--network` (in any combination) | Lightweight host / GPU / network introspection. **No `torch.distributed` rendezvous.** Cannot hang on network misconfig. | +| **Perf-only** | `--perf-test`, `--tests ...`, or `--quick` | Runs the configured perf tests under a global rendezvous. **Implied** by `--tests` and `--quick`. | +| **Default (info + perf)** | No flags at all | Runs the info report first, then every perf test. | + +### Mode precedence + +1. **Any of `--perf-test` / `--tests` / `--quick` is set → perf-only mode.** + If info selectors (`--host`/`--gpu`/`--network`) are also present, they are dropped and a `WARN` is emitted (also written as a `> Note:` at the top of the perf report). To get both reports, run two invocations. +2. **Otherwise, any of `--host`/`--gpu`/`--network` is set → info-only mode.** + Perf-only tuning knobs (e.g. `--comm-sizes-mb`) are inert in this mode and trigger a single `WARN` listing them. +3. **Otherwise (no flags) → default**: info report **first** (no rendezvous), then perf tests. + +The default order ensures you always get a report even if `torch.distributed` initialization later hangs. + +--- + +## 2. Quick start + +### Info report only (fast, no rendezvous) ```bash primus-cli direct -- preflight --host --gpu --network ``` -### Full preflight (info + perf tests) +### Full preflight (info + every perf test) ```bash primus-cli direct -- preflight @@ -25,54 +57,466 @@ primus-cli direct -- preflight primus-cli direct -- preflight --perf-test ``` -## Common usage (Slurm) +### Fast pre-launch sanity check + +```bash +primus-cli direct -- preflight --quick +``` + +Equivalent on SLURM via `primus-cli slurm`: + +```bash +primus-cli slurm srun -N 4 -- preflight --quick +``` + +Without a container, see [`preflight-direct.md`](./preflight-direct.md) for the equivalent `runner/primus-cli direct -- preflight ...` invocations. + +--- + +## 3. Test selection (`--tests`) + +`--tests` takes a comma-separated list of canonical tokens (or `all`). Implies `--perf-test`. + +| Token | What it runs | +|---|---| +| `gemm` | Single-GPU square GEMM TFLOPS sweep. | +| `intra-allreduce` | Intra-node `all_reduce` bandwidth at every selected `--intra-group-sizes` x `--intra-comm-sizes-mb`. | +| `intra-alltoall` | Intra-node `all_to_all` bandwidth, same configuration matrix. | +| `inter-allreduce` | Inter-node `all_reduce` bandwidth at every selected `--inter-group-sizes` x `--inter-comm-sizes-mb`. | +| `inter-alltoall` | Inter-node `all_to_all` bandwidth, same configuration matrix. | +| `inter-p2p` | Inter-node 2-rank P2P send/recv. Requires `--inter-group-sizes` to actually contain pair-able sizes. | +| `inter-ring-p2p` | Inter-node ring-pattern P2P, sized by `--ring-p2p-sizes-mb`. | +| `all` | Every token above. Default when `--tests` is omitted. | + +Examples: + +```bash +# GEMM only +primus-cli direct -- preflight --tests gemm + +# Just the inter-node bandwidth tests +primus-cli direct -- preflight --tests inter-allreduce,inter-alltoall + +# Combine with size overrides +primus-cli direct -- preflight \ + --tests gemm,inter-allreduce \ + --comm-sizes-mb 64,1024 \ + --inter-group-sizes all +``` + +Unknown tokens fail fast (before any rendezvous): + +```text +[Primus:Preflight] ERROR: invalid perf config: --tests: unknown token 'gem'. +Valid tokens: gemm, intra-allreduce, intra-alltoall, inter-allreduce, +inter-alltoall, inter-p2p, inter-ring-p2p, all +``` + +--- + +## 4. Quick preset (`--quick`) + +`--quick` is the recommended **pre-launch sanity** preset. Implies `--perf-test`. It substitutes: + +| Knob | `--quick` value | +|---|---| +| `--tests` | `gemm,intra-allreduce,inter-allreduce` | +| `--comm-sizes-mb` | `64,1024` | +| `--intra-group-sizes` | `LOCAL_WORLD_SIZE` (full intra-node group only) | +| `--inter-group-sizes` | `all` (full N-node group only) | +| `warmup` | `5` | +| `iteration` | `20` | + +**User-supplied flags override the preset.** For example: + +```bash +# Quick preset, but with a custom size set +primus-cli direct -- preflight --quick --comm-sizes-mb 32,256 +``` + +A full perf run with default knobs takes minutes; `--quick` typically finishes in <60s on healthy hardware. + +--- + +## 5. Tuning the perf tests + +All perf tuning knobs default to `None` so preflight can tell whether you set them. When unset, the documented defaults below apply. + +### 5.1 Message sizes (collective + P2P) + +| Flag | Default | Applies to | +|---|---|---| +| `--comm-sizes-mb CSV` | `2,4,8,16,32,64,128,256,512,1024` | Default for both intra- and inter-node `allreduce` / `alltoall` and `inter-p2p` when no specific override is given. | +| `--intra-comm-sizes-mb CSV` | falls back to `--comm-sizes-mb` | Override for **intra-node** `allreduce` / `alltoall`. | +| `--inter-comm-sizes-mb CSV` | falls back to `--comm-sizes-mb` | Override for **inter-node** `allreduce` / `alltoall` / `inter-p2p`. | + +```bash +# Smaller, focused sweep +primus-cli direct -- preflight --comm-sizes-mb 8,128 + +# Different sizes for intra vs inter +primus-cli direct -- preflight \ + --tests intra-allreduce,inter-allreduce \ + --comm-sizes-mb 8,128 \ + --intra-comm-sizes-mb 4,32 +``` + +### 5.2 Group sizes + +| Flag | Default | Notes | +|---|---|---| +| `--intra-group-sizes CSV` | `2,4,8` | Each value must divide `LOCAL_WORLD_SIZE`. | +| `--inter-group-sizes CSV` | `2,4,all` | `all` means the full N-node group. Other values are subgroup sizes. **For `inter-alltoall` only**, every requested per-group node count is internally clamped to **16** before deduping (see "Inter-node alltoall is capped at 16 nodes" below). The other inter-node tests (`inter-allreduce`, `inter-p2p`, `inter-ring-p2p`) use the requested sizes unchanged. | + +```bash +# All-GPU intra + full N-node inter only +primus-cli direct -- preflight \ + --tests intra-allreduce,inter-allreduce \ + --intra-group-sizes 8 \ + --inter-group-sizes all +``` + +Validation is gated by which tests are actually selected. For example, `--tests gemm --intra-group-sizes 3` does **not** abort on a host with `LOCAL_WORLD_SIZE=8`; the intra-group constraint is only checked when an intra test is enabled. + +#### Inter-node alltoall is capped at 16 nodes + +Regardless of the cluster size or what `--inter-group-sizes` requests, the `inter-alltoall` test always runs on per-group node counts of at most **16**. Concretely, every requested value `G` is replaced with `min(G, 16)`, and the resulting list is deduped. Examples: + +| Cluster | `--inter-group-sizes` | Requested (resolved) | `inter-alltoall` actually runs | +|---|---|---|---| +| 8 N | `all` | `[8]` | `[8]` (no change) | +| 64 N | `2,4,all` | `[2, 4, 64]` | `[2, 4, 16]` | +| 128 N | `2,4,16,32,all` | `[2, 4, 16, 32, 128]` | `[2, 4, 16]` | +| 128 N | `64` | `[64]` | `[16]` | + +When the cap actually changes the list, preflight emits a single one-line WARN to stdout so the row labels in the report (e.g. `alltoall-16nodes` instead of `alltoall-128nodes`) are not surprising. + +Why a hard cap and why 16: + +- **Real-world MoE training rarely dispatches across more than ~8 nodes.** DeepSeek-V3's largest published configuration uses `EP=64` over 8 nodes with per-token dispatch capped at 4 nodes; NVIDIA Megatron-Core's published EP recipes follow the same shape. 16 covers every published configuration with comfortable headroom while staying well clear of the per-node ephemeral-port pressure described in §7. +- **The cap eliminates the dominant source of `Address already in use` at large scale.** During each `ncclCommInit`, an inter-node alltoall sub-group of `K` nodes opens a near-full mesh of IB OOB sockets per local rank — empirically, peak simultaneous ESTAB sockets per node grow ~linearly with `K` at **~145 sockets per added node** (linear fit `peak_ESTAB ≈ 145·K + 683` over measurements at 24/32/48/56 N; see §7.1.2). Once peak ESTAB approaches the size of the kernel's ephemeral-port pool (default `28 232` ports), the next outgoing `bind()` walks the entire range without finding an allocatable port and fails. Capping the sub-group at 16 holds peak ESTAB at **~3.7 k** — measured directly on a 56 N cluster with the cap active — comfortably under any sensible pool, so the failure cannot occur regardless of cluster size or OS tuning. +- **Other inter-node tests are unaffected.** `inter-allreduce` (ring/tree, ~`log K` peers per rank) and `inter-ring-p2p` (ring, 2 peers per rank) and `inter-p2p` (pairwise) all open far fewer simultaneous OOB sockets than alltoall and continue to honor `--inter-group-sizes` exactly as written. As a concrete reference point: at 56 nodes, a default-configured `inter-allreduce` peaks at ~1.8 k ESTAB; an `inter-alltoall` over the same 56 nodes peaks at ~8.8 k. +- **Intentionally not configurable.** This is a known-safe ceiling for the comm shapes preflight is supposed to characterize, not a tuning knob; raising it would re-introduce the very failure mode preflight is meant to *detect* in the cluster, not *cause*. + +### 5.3 Ring P2P sizes + +| Flag | Default | Applies to | +|---|---|---| +| `--ring-p2p-sizes-mb CSV` | `10,20,40,80,160` | `inter-ring-p2p` only. | + +```bash +primus-cli direct -- preflight \ + --tests inter-ring-p2p \ + --ring-p2p-sizes-mb 5,20,80 +``` + +### 5.4 Plotting + +| Flag | Effect | +|---|---| +| `--plot` | After each perf test, write per-size bandwidth bar charts under `//` and reference them in the markdown report. | + +--- + +## 6. Reliability knobs + +Two knobs that are inert under happy-path conditions but matter at scale or on flaky networks. + +### 6.1 `--comm-cleanup-delay-sec FLOAT` (default `2.0`) + +Delay (seconds) inserted between destroying NCCL/RCCL process groups and creating new ones. Provides cross-rank synchronization across the destroy → setup transition (so a rank doesn't try to connect to a peer whose listener hasn't finished closing) and gives the kernel a moment to unlink closed-socket bookkeeping. See §7.1 for why this knob is *not* primarily defending against `TIME_WAIT` pressure (which doesn't apply to NCCL's connection pattern) — the actual `Address already in use` defense is the §5.2 inter-alltoall cap. + +- Default `2.0` is essentially free and worth keeping as cheap insurance at every cluster size. +- Set to `0` to disable the sleep entirely (barrier only). +- Bump to e.g. `5` only on very flaky networks or pathological kernel scheduling; widening `ip_local_port_range` (§7.2) is a more direct fix when the binding constraint is genuinely hit. + +```bash +# Small/medium clusters: defaults are fine. Override only if you see +# port-reuse races on very flaky networks or unusual kernel TIME_WAIT +# settings. +primus-cli slurm srun -N 8 -- preflight --quick --comm-cleanup-delay-sec 5 +``` + +See §7 ("Running on very large clusters") for the rationale behind why the default works at scale and for the OS-level best-practices that apply to any RDMA workload. + +### 6.2 `--dist-timeout-sec INT` (default `120`) + +Timeout (seconds) for `torch.distributed.init_process_group`. If init does not complete within this many seconds, preflight writes the info report (when applicable) plus a `Distributed Init` failure section to the markdown report, prints a clear error, and exits `2` — instead of hanging indefinitely. -Info report only (fast): +> Note: §6 used to also document `--comm-cleanup-large-threshold-nodes`, which forced a 60 s drain whenever a destroyed subgroup met or exceeded a size threshold (default 64 nodes). That flag was removed because the underlying failure mode it tried to paper over — peak simultaneous ESTAB exhausting the per-node ephemeral-port pool during a large inter-node alltoall `ncclCommInit` — is now prevented at the source by the §5.2 alltoall cap. The `--comm-cleanup-delay-sec` knob remains, but its primary role is now cross-rank synchronization across the destroy → setup transition rather than draining `TIME_WAIT`; the default 2 s is essentially free and worth keeping as cheap insurance. ```bash -primus-cli slurm srun -N 4 -- preflight --host --gpu --network +# Fail fast if rendezvous does not work +primus-cli direct -- preflight --perf-test --dist-timeout-sec 30 +``` + +--- + +## 7. Running on very large clusters (≥ 64 nodes) + +Beyond ~64 nodes, two practical concerns dominate that smaller runs never see. Read this section once if you operate clusters in this range; it explains the failure mode, the OS knobs that fix it at the system level, and the recommended preflight invocation patterns. + +### 7.1 Why "Address already in use" used to surface at scale + +The failure is a **per-node ephemeral port exhaustion** during a single inter-node alltoall `ncclCommInit`, not a chronic accumulation of `TIME_WAIT` sockets across phases. Understanding this distinction is what motivates both the §5.2 cap and the §7.2 OS tuning recommendations. + +#### 7.1.1 What actually consumes the per-node ephemeral pool + +Linux's outgoing-connection allocator (`__inet_hash_connect()`) does **not** reject ports just because some other socket is in any state on them. It rejects a port only when the new connection's full 4-tuple `(saddr, sport, daddr, dport)` collides with an existing socket's 4-tuple: + +- **Live ESTAB sockets** sit on a specific 4-tuple and prevent the kernel from reusing that exact 4-tuple. Each new outgoing connection that lands on a port already holding an ESTAB socket has to walk to the next port. As ESTAB density approaches one-socket-per-port across the entire ephemeral range, the walk takes longer and longer until eventually no port is allocatable — that's the EADDRINUSE. +- **TIME_WAIT sockets** also live on specific 4-tuples but are governed by `tcp_tw_reuse`. Critically, they only block a new connection when the new connection's *desired* 4-tuple collides with the historical one — i.e. when the new connection is to the *same* `(daddr, dport)` from the *same* `(saddr, sport)`. + +For NCCL inter-node OOB traffic, the second case essentially never happens: each `ncclCommInit` connects to **fresh peer OOB listening ports** (the peer chooses an ephemeral listener per setup), so successive comm setups always have different `dport`. TIME_WAIT entries from a previous destroy sit on `(local, P_old, peer, dport_OLD)`; the next setup wants `(local, ?, peer, dport_NEW)`. Even when the new connection happens to land on `sport == P_old`, the dports differ → no 4-tuple collision → the TIME_WAIT entry is invisible to the allocator regardless of `tcp_tw_reuse`. + +The practical consequence: under NCCL's connection pattern, **the binding constraint reduces to peak simultaneous ESTAB ≤ size of the ephemeral pool**. + +#### 7.1.2 Why inter-node alltoall is the test that exhausts it + +Per-node peak ESTAB scales very differently across the inter-node test families: + +| Test (56 N, default knobs) | Peer topology per rank | Peak ESTAB per node | +|---|---|---| +| `inter-allreduce` | ring / tree, ~`log K` peers | ~1.8 k | +| `inter-alltoall` | full mesh, `K-1` peers | ~8.8 k | + +The empirical scaling of `inter-alltoall` peak ESTAB in the default perf-test sweep. Four measurements at 24/32/48/56 N fit a near-perfect line: + +``` +peak_ESTAB(per node) ≈ 145.09 · K + 683 ``` -Full preflight (info + perf tests): +| K (nodes) | Measured peak ESTAB | Fit (145·K + 683) | Source | +|---|---|---|---| +| 16 (capped run) | **3 687** | 3 003 | direct measurement — fit slightly under-predicts at the low-N extrapolation | +| 24 | 4 165 | 4 165 | calibration point | +| 32 | 5 326 | 5 326 | calibration point | +| 48 | 7 647 | 7 647 | calibration point | +| 56 | 8 808 | 8 808 | calibration point | +| 128 (uncapped, extrapolated) | — | **~19 240** | linear extrapolation | + +Two practical reads from this: + +- **At 56 N the workload already sits at ~31 % of the default 28 232-port pool**, with ~19 k ports of headroom. That's why every default-pool 56 N run in our experiment succeeded. +- **An uncapped 128-N inter-alltoall would peak around ~19 k ESTAB** — still inside the default pool but with ~9 k ports of headroom, and *over the cliff* on any cluster that has narrowed `ip_local_port_range`, that runs additional outgoing TCP work concurrently, or that has the source-port allocator's random-walk hit a bad starting offset. + +#### 7.1.3 Empirical confirmation: the binding constraint really is `peak_ESTAB ≤ pool_size` + +Holding the workload constant (56 N, only inter-alltoall) and varying *only* the per-node ephemeral pool size: + +| Pool | Peak ESTAB | Headroom | Result | +|---|---|---|---| +| 28 231 (default) | 8 805 | +19 426 | OK | +| 9 000 | 8 806 | +194 | OK (barely) | +| 8 000 | (~8 800 expected) | −800 | **FAIL** | + +The transition is sharp and right at the predicted boundary. Note that peak `TIME_WAIT` count in the same runs ranged from ~16 k to ~24 k — **far above** the pool size in the 9 k and 8 k rows. If TIME_WAIT count drove the failure, both narrow-pool rows should fail. They don't; only the row where `peak_ESTAB > pool_size` does. + +The mechanism is further corroborated by a direct cap experiment on the same 56 N cluster, default pool: re-running the inter-node alltoall test with the per-group node count capped at 16 (matching the §5.2 cap) yields peak ESTAB **3 687** — well under both the default 28 k pool and the 9 k / 8 k stressed pools — and the run completes cleanly. The cap acts directly on the binding constraint by holding peak ESTAB low enough that the pool cannot be exhausted, regardless of cluster size. + +#### 7.1.4 The cap closes the failure mode at the source + +The §5.2 inter-node alltoall cap (16 nodes max) holds peak ESTAB at **~3.7 k** regardless of cluster size — measured directly in §7.1.3, comfortably under the default 28 k pool (~13 % utilization) and still safe at half-default pool widths. The OS-level tunings in §7.2 remain useful general hygiene for any RDMA / multi-NIC workload, but a default preflight invocation no longer needs them to avoid `Address already in use`. + +This is **not** a real training failure mode in any case — production training jobs create their TP / DP / PP / EP communicators *once* at startup and reuse them, and real-world MoE training rarely dispatches across more than ~8 nodes. The preflight tool is the only thing that builds many large communicators in a short window, which is why the issue was preflight-specific to begin with. + +### 7.2 OS-level tuning (best-practice for any large-cluster node) + +Given §7.1's mechanism (binding constraint = peak ESTAB ≤ pool size), the OS knobs split cleanly into "directly relevant" and "general hygiene": ```bash -primus-cli slurm srun -N 4 -- preflight +# 1) DIRECTLY RELEVANT: widen the per-node ephemeral pool. +# The default range is 32768-60999 (~28 k ports). Widening it +# to 1024-65535 (~64 k ports) more than doubles the headroom +# for peak simultaneous ESTAB — and that is the only thing that +# can produce EADDRINUSE under the NCCL connection pattern. +sudo sysctl -w net.ipv4.ip_local_port_range="1024 65535" + +# 2) GENERAL HYGIENE: allow TIME_WAIT reuse for outgoing connections. +# For the NCCL inter-node OOB pattern this is largely a no-op +# (each ncclCommInit picks fresh peer destination ports, so +# historical TIME_WAIT 4-tuples never collide with what the +# next setup wants). It is still recommended for any host that +# runs additional outgoing TCP workloads where the SAME +# (daddr, dport) is hit repeatedly from the same source IP -- +# the textbook scenario the kernel doc is written around. +sudo sysctl -w net.ipv4.tcp_tw_reuse=1 + +# Persist across reboots: +cat <<'EOF' | sudo tee /etc/sysctl.d/99-large-cluster.conf +net.ipv4.ip_local_port_range = 1024 65535 +net.ipv4.tcp_tw_reuse = 1 +EOF +sudo sysctl --system ``` -Perf tests only: +If you only have time to set one of these, pick `ip_local_port_range`. With the §5.2 inter-alltoall cap holding peak ESTAB at ~3.7 k even at 1024 N, neither knob is *required* for preflight — but the wider port range is the right insurance for any host that might also run other outgoing TCP traffic concurrently. + +Note on `tcp_tw_reuse=2`: as of Linux 4.12, value `2` enables TIME_WAIT reuse **only for loopback** (`127.0.0.0/8`, `::1`). For inter-node IB OOB connections this is equivalent to `tcp_tw_reuse=0`. The fact that several of our 56 N runs succeeded under `tcp_tw_reuse=2` with `peak_WAIT > 16 k` is direct evidence that `TIME_WAIT` *count* doesn't gate inter-node bind() — only `peak_ESTAB > pool_size` does (see §7.1.3). + +### 7.3 In-tool defenses + +Preflight has two complementary defenses: + +1. **The §5.2 inter-node alltoall cap (16 nodes max).** *This is the actual fix.* Regardless of cluster size or `--inter-group-sizes`, the alltoall test never builds a sub-group large enough to push peak ESTAB anywhere near the per-node ephemeral-port pool. This eliminates the historical EADDRINUSE failure mode by construction (see §7.1). +2. **The §6.1 per-phase cleanup delay (`--comm-cleanup-delay-sec`, default `2.0`).** A global barrier + sleep inserted after every comm destroy. Its primary job is **cross-rank synchronization** across the destroy → setup transition (so a rank doesn't try to connect to a peer whose listener hasn't finished closing) and giving the kernel a moment to unlink closed-socket bookkeeping. It is *not* protecting against `TIME_WAIT` 4-tuple collisions — those don't occur in NCCL's connection pattern (see §7.1.1). The default 2 s is essentially free and worth keeping as cheap insurance. + +Together, a default preflight invocation is safe at every cluster size we test up to 1024 nodes. The only situation where you would consider raising `--comm-cleanup-delay-sec` is on a network with unusually narrow ephemeral-port ranges or pathological kernel scheduling — and even there, widening `ip_local_port_range` (§7.2) is the cleaner fix because it directly addresses the binding constraint. + +### 7.4 Recommended invocation patterns at very large scale + +For clusters at or beyond ~128 nodes, the most reliable and most informative way to use preflight is still to **split the run into one test family per invocation** rather than one big run — not for `Address already in use` reasons (the §7.3 defenses handle that), but because it keeps wall-clock per invocation small and makes it trivial to identify which specific comm shape is degraded if a metric looks off. ```bash -primus-cli slurm srun -N 4 -- preflight --perf-test +# 1) GPU + intra-node fabric first (cheap, no inter-node OOB churn). +primus-cli slurm srun -N 128 -- preflight \ + --tests gemm,intra-allreduce,intra-alltoall + +# 2) Inter-node DP-style collectives, all-nodes group only. +primus-cli slurm srun -N 128 -- preflight \ + --tests inter-allreduce \ + --inter-group-sizes all +# Note: --inter-group-sizes all is honored here for inter-allreduce. +# For inter-alltoall it would be clamped to 16 (see §5.2). +primus-cli slurm srun -N 128 -- preflight \ + --tests inter-alltoall \ + --inter-group-sizes all + +# 3) Inter-node PP-style ring P2P (the test that benefits most from +# isolation — it's the closest match to what real pipeline-parallel +# training actually exercises). +primus-cli slurm srun -N 128 -- preflight \ + --tests inter-ring-p2p + +# 4) Optional: pairwise inter-node P2P scan (useful for finding a +# single bad link, slower because it walks many pairs). +primus-cli slurm srun -N 128 -- preflight \ + --tests inter-p2p +``` + +Each invocation: + +- Tears down its own `WORLD` at exit, so the next invocation starts with a fresh per-node port pool. +- Touches only one `--tests` value, so you get a per-test wall clock and can re-run a single phase if its numbers look off without paying for the others. +- Carries its own report under `--report-file-name` (or the wrapper-generated default), which makes archiving and comparison across runs straightforward. + +### 7.5 Decision flow + +| Cluster size | Recommended approach | +|---|---| +| ≤ 32 nodes | Single command, default knobs. Nothing special. | +| 33-127 nodes | Single command, default knobs. | +| ≥ 128 nodes | Single command works with default knobs (the §5.2 alltoall cap holds peak ESTAB at ~3.7 k, ≈ 7.6× under the default 28 k ephemeral pool). Splitting per `--tests` token as in §7.4 is recommended for diagnostic clarity rather than for safety. Widening `ip_local_port_range` (§7.2) is best-practice for any RDMA workload but no longer required for preflight specifically. | +| ≥ 256 nodes | Always split as in §7.4 — keeps every invocation snappy and makes regressions much easier to localize. | + +--- + +## 8. Reporting + +| Flag | Default | Effect | +|---|---|---| +| `--dump-path DIR` | `output/preflight` | Output directory for reports + plots. | +| `--report-file-name NAME` | auto-generated `preflight-${NNODES}N-YYYYMMDD-HHMMSS` | Base name for report files. Omit to let preflight auto-generate a unique timestamped name (prevents stale leftovers from prior runs being mistaken for fresh output). Pass an explicit value when you want a stable / well-known filename. | +| `--disable-pdf` | enabled | Skip PDF generation (Markdown only). Useful when `weasyprint`/`markdown2` aren't installed. | + +Output files: + +| File | Produced when | Notes | +|---|---|---| +| `.md` / `.pdf` | Info-only mode, or default mode | Info report. | +| `_perf.md` / `_perf.pdf` | Perf-only mode, or default mode | Perf report (GEMM + comm). | + +Only **rank 0** writes the report. + +### Perf report layout + +A `_perf.md` produced by a default run contains, in order: + +1. (Optional) `> Note:` line listing dropped info selectors. +2. `# Nodes` legend — `Node N → Hostname` table, used by every subsequent table to keep host columns compact. +3. `=======IB Bandwidth roofline (GB/s)=======` — bandwidth of the first IB device on Node 0. +4. Per enabled test, in this order: `gemm`, `intra-comm`, `inter-comm`, `inter-p2p`, `inter-ring-p2p`. Each section has a configuration line, a results table (Node / Rank / hostname / per-size GB/s), optional plots, and a per-rank wall-clock summary. +5. `[Primus:Preflight] done in s` lines on stdout for at-a-glance progress on the launching shell. + +--- + +## 9. Backward-compat aliases + +| Flag | Equivalent | Notes | +|---|---|---| +| `--check-host`, `--check-gpu`, `--check-network` | `--host`, `--gpu`, `--network` | Same behavior. Keep working for older scripts. | +| `--no-split-nodes-subgroup` | `--inter-group-sizes all` **and** drops `inter-p2p` | Pre-`--tests`/`--inter-group-sizes` alias. Use the new flags in new scripts. | + +--- + +## 10. Comparison with node-smoke + +| Aspect | `node-smoke` | `preflight` | +|---|---|---| +| Rendezvous | None — every node independent | Global `torch.distributed` | +| Wall clock | ~30–60 s for 6 nodes (Tier 1+2) | Minutes; scales with N for inter-node tests | +| Granularity | Per-node PASS/FAIL | Per-rank measurements (no auto-fail by default) | +| Inter-node bandwidth matrix | Not tested (intentionally) | Yes (allreduce/alltoall/p2p/ring-p2p) | +| Drift detection | Yes (versions, NIC firmware, port count) | No | +| Host limits / RDMA roll-call | Yes (hard fail) | Reported via `collect_*_info` only | +| Output format | Per-node JSON + cluster md + SLURM-ready txt | Markdown + PDF | + +**Recommended workflow**: run `node-smoke` first to exclude broken nodes, then run `preflight` on the surviving set to get cross-node bandwidth measurements. See [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) §3 ("Quick start") for the integration commands. + +--- + +## 11. Validation & error handling + +Preflight resolves the perf config **before** any distributed rendezvous. This means typos and bad sizes/group-sizes fail in seconds, not after a 120s NCCL init: + +```text +[Primus:Preflight] ERROR: invalid perf config: --tests: unknown token 'gem'. +[Primus:Preflight] ERROR: invalid perf config: + --intra-group-sizes: [3] do not divide LOCAL_WORLD_SIZE=8 +[Primus:Preflight] ERROR: invalid perf config: --comm-sizes-mb: '0' must be positive ``` -## CLI flags +In info-only mode, perf-only tuning knobs trigger a single warning so you notice them but they don't abort: + +```text +[Primus:Preflight] WARN: --comm-sizes-mb,--intra-group-sizes have no effect +in info-only mode (no --perf-test/--tests/--quick). +``` + +In default mode where info selectors are dropped because perf intent was set, the preserved warning is also written into the perf report header: + +```text +> Note: info selectors --host were dropped because perf mode +> (--perf-test/--tests/--quick) takes precedence. Run them in a separate +> invocation if you want both reports. +``` -Selection: -- `--host`: host info (CPU, memory, PCIe) -- `--gpu`: GPU info -- `--network`: network info -- `--perf-test`: run perf tests only (GEMM + comm). This is slower. +--- -Reporting: -- `--dump-path`: output directory (default: `output/preflight`) -- `--report-file-name`: base report name (default: `preflight_report`) -- `--disable-pdf`: disable PDF generation +## 12. Operational tips -Perf-test extras: -- `--plot`: generate plots (only used with `--perf-test`) +- **For multi-node runs, always use `primus-cli slurm` or `primus-cli direct` under `srun`** so distributed environment variables (`NNODES` / `NODE_RANK` / `MASTER_ADDR`) are set correctly. +- **Insufficient CPU cores cause >30x perf slowdowns** — pass `srun -c ` so RCCL's network proxy threads have CPU to spawn on. Verify with `srun -N 1 --gpus-per-node=8 bash -c 'nproc'`. +- **For a quick environment snapshot**, prefer `--host --gpu --network` (no rendezvous, finishes in seconds even on broken networks). +- **Between each communication test phase**, preflight performs a global barrier + `--comm-cleanup-delay-sec` sleep (default 2 s) for cross-rank sync across the destroy → setup transition. The default works at every cluster size we test up to 1024 nodes because the inter-node alltoall sub-group is internally capped at 16 (see §5.2), which holds peak simultaneous ESTAB sockets per node well under the kernel's ephemeral-port pool — the only constraint that actually produces `Address already in use` under NCCL's connection pattern (see §7.1). Widening `ip_local_port_range` (§7.2) is best-practice for any RDMA workload. +- **For pre-launch screening of a large cluster**, the recommended sequence is: + 1. `node-smoke` to prune broken nodes (`failing_nodes.txt`). + 2. `preflight --quick` on the surviving nodes for the perf sanity numbers. + 3. `preflight` (full) on the same set if the `--quick` numbers raise a flag. -Backward compatibility: -- `--check-host/--check-gpu/--check-network` are supported as aliases for `--host/--gpu/--network`. +--- -## Outputs +## 13. Running preflight without a container -By default, outputs are written under `output/preflight`. +If you cannot (or prefer not to) use a container, see [`preflight-direct.md`](./preflight-direct.md) for the step-by-step `runner/primus-cli direct -- preflight ...` walkthrough — Python virtual-environment setup, SLURM invocation patterns, NCCL configuration for Broadcom and Pensando (AINIC) clusters, and many configurable-knob examples. -Typical report files: -- `preflight_report.md` / `preflight_report.pdf`: **info report** (host/GPU/network) -- `preflight_report_perf.md` / `preflight_report_perf.pdf`: **perf report** (GEMM + comm tests) +--- -## Notes +## 14. See also -- For multi-node runs, use `primus-cli slurm …` (or your preferred launcher) so distributed environment variables are set correctly. -- If you only want a quick environment snapshot, prefer `--host --gpu --network`. +- [`preflight-direct.md`](./preflight-direct.md) — quick-start guide for `primus-cli direct -- preflight` (no container). +- [`node-smoke.md`](./node-smoke.md) — full reference for the per-node smoke test. +- [`node-smoke-test-instruction.md`](./node-smoke-test-instruction.md) — short quick-start for the smoke test. +- [`runner/primus-cli-direct.sh`](../runner/primus-cli-direct.sh) — non-container launcher (`primus-cli direct` dispatches here). +- [`primus/tools/preflight/`](../primus/tools/preflight/) — implementation. +- [`primus/tools/preflight/preflight_args.py`](../primus/tools/preflight/preflight_args.py) — canonical CLI definition (single source of truth for flags + defaults). diff --git a/primus/cli/subcommands/node_smoke.py b/primus/cli/subcommands/node_smoke.py new file mode 100644 index 000000000..5e1c695a4 --- /dev/null +++ b/primus/cli/subcommands/node_smoke.py @@ -0,0 +1,131 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +node_smoke CLI subcommand. + +Surfaces ``primus.tools.preflight.node_smoke`` as a first-class primus-cli +subcommand so users can run it via the same dispatch chain as ``preflight``, +``train``, ``benchmark``, etc.: + + primus-cli direct -- node_smoke --tier2-perf + primus-cli slurm srun -N 4 -- direct -- node_smoke --tier2-perf + primus-cli slurm srun -N 4 -- container -- node_smoke --tier2-perf + +Design choices (consolidate-preflight-direct-wrappers plan, section 4): + +1. **No inner ``run`` keyword.** The standalone CLI has + ``run / aggregate / _per_gpu`` subparsers, but only ``run`` is user-facing. + Hoisting its flags onto the top-level ``node_smoke`` parser means the user + types ``primus-cli direct -- node_smoke --tier2-perf`` instead of + ``... -- node_smoke run --tier2-perf``. + +2. **Always aggregate on rank 0.** Every wrapper invocation runs ``_cmd_run`` + on every rank, then ``_cmd_aggregate`` on rank 0 only. Aggregation takes + a few seconds (reads per-node JSONs, writes one report) and is what users + want ~100% of the time. The rare exceptions ("run only, don't aggregate" + / "aggregate only") stay reachable through the unchanged standalone CLI: + + python -m primus.tools.preflight.node_smoke run ... + python -m primus.tools.preflight.node_smoke aggregate ... + +3. **``allow_abbrev=False``.** Mirrors the standalone ``run`` subparser at + ``primus/tools/preflight/node_smoke/cli.py:542``. Without this, an old + script that still passes ``--tier2`` (legacy flag name) would silently + match ``--tier2-perf`` as a prefix and run the wrong test set. + +4. **No ``--silent`` flag.** Silencing is handled exclusively by the bash + launcher (``primus-cli-direct.sh`` ``--silent`` before ``--``). Passing + ``--silent`` here will be rejected by argparse, which is the desired + behavior -- one knob in one place. +""" + +from __future__ import annotations + +import argparse +import os +from typing import Any, List + + +def run(args: Any, extra_args: List[str]) -> None: + """Two-phase dispatch: per-node ``_cmd_run`` on every rank, then rank-0 + ``_cmd_aggregate`` with SLURM-resolved aggregator args. + + Exit-code rule (matches the deleted wrapper): + - Non-rank-0: propagate ``_cmd_run`` exit code. + - Rank 0: aggregator exit code wins (it knows about MISSING nodes that + ``_cmd_run`` can't see). If aggregator returns 0 but run failed, we + still surface the run failure so a sick rank-0 box can't paint itself + green via a successful aggregate. + """ + from primus.tools.preflight.node_smoke.cli import ( + _cmd_aggregate, + _cmd_run, + _resolve_aggregate_args_from_slurm, + ) + + if extra_args: + # node_smoke uses allow_abbrev=False at the argparse level; reaching + # this path means the user passed something the parser didn't claim. + # Surface it loudly so a typo'd flag doesn't get silently dropped. + print( + f"[Primus:NodeSmoke] Unknown arguments: {extra_args}. " + f"Run `primus-cli node_smoke --help` for valid options.", + ) + raise SystemExit(2) + + rc_run = int(_cmd_run(args)) + rank = int(os.environ.get("NODE_RANK", os.environ.get("SLURM_NODEID", "0"))) + if rank != 0: + raise SystemExit(rc_run) + + agg_ns = _resolve_aggregate_args_from_slurm(args) + rc_agg = int(_cmd_aggregate(agg_ns)) + + # Aggregator's exit code wins on rank 0 (it can detect MISSING nodes), + # but a failed per-node run on rank 0 itself must not be hidden by a + # successful aggregate. Surface the higher of the two. + raise SystemExit(max(rc_agg, rc_run)) + + +def register_subcommand(subparsers: argparse._SubParsersAction) -> argparse.ArgumentParser: + """Register ``node_smoke`` with the primus-cli main parser. + + Hoists the ``run`` subparser's flags directly onto the top-level + ``node_smoke`` parser (no inner ``run`` keyword) plus the aggregator + tuning knobs the rank-0 aggregate step needs. + """ + from primus.tools.preflight.node_smoke.cli import ( + _add_aggregate_flags, + _add_run_flags, + ) + + parser = subparsers.add_parser( + "node_smoke", + help="Run per-node preflight smoke test on every node and aggregate on rank 0.", + description=( + "Node-local preflight smoke test. Each node runs independently " + "(no global rendezvous, no torch.distributed). On rank 0 the per-node " + "verdicts are aggregated into smoke_report.md + passing_nodes.txt + " + "failing_nodes.txt -- the latter two are directly consumable by " + "`srun --nodelist=` / `srun --exclude=`." + ), + # See module docstring point 3. + allow_abbrev=False, + ) + + # Attach the canonical run-side flag surface. Anything new added to + # `_add_run_flags` automatically flows here. + _add_run_flags(parser) + + # Attach the aggregator-only flags. `--dump-path` / `--hbm-busy-threshold-gib` + # / `--gpu-activity-warn-pct` are already on the parser from `_add_run_flags` + # with identical defaults, so we skip them here to avoid argparse + # conflicting-option errors. + _add_aggregate_flags(parser, include_dump_path=False) + + parser.set_defaults(func=run) + return parser diff --git a/primus/core/config/yaml_loader.py b/primus/core/config/yaml_loader.py index 89cb8b58a..4ec3a5722 100644 --- a/primus/core/config/yaml_loader.py +++ b/primus/core/config/yaml_loader.py @@ -7,8 +7,6 @@ import os import re -import yaml - from primus.core.config.merge_utils import deep_merge ENV_PATTERN = re.compile(r"\${([^:{}]+)(?::([^}]*))?}") @@ -30,6 +28,8 @@ def parse_yaml(path: str) -> dict: # 1. Load YAML # ================================================================ def _load_yaml(path: str): + import yaml + with open(path, "r") as f: return yaml.load(f, Loader=yaml.SafeLoader) diff --git a/primus/core/utils/yaml_utils.py b/primus/core/utils/yaml_utils.py index 1f39b1f2f..27cca8b13 100755 --- a/primus/core/utils/yaml_utils.py +++ b/primus/core/utils/yaml_utils.py @@ -8,8 +8,6 @@ from types import SimpleNamespace from typing import Any, Mapping -import yaml - from primus.core.config.merge_utils import deep_merge from primus.core.config.yaml_loader import parse_yaml as _parse_yaml_core @@ -140,6 +138,7 @@ def dump_namespace_to_yaml(ns: SimpleNamespace, file_path: str): >>> ns = SimpleNamespace(a=1, b=SimpleNamespace(c=2)) >>> dump_namespace_to_yaml(ns, "config.yaml") """ + import yaml def ns_to_dict(obj): if isinstance(obj, SimpleNamespace): diff --git a/primus/tools/preflight/README.md b/primus/tools/preflight/README.md index d7ca011f0..ff7ae7a8a 100644 --- a/primus/tools/preflight/README.md +++ b/primus/tools/preflight/README.md @@ -9,26 +9,32 @@ Torch / single node: ```bash primus-cli preflight \ - --dump-path output/preflight \ - --report-file-name preflight_report + --dump-path output/preflight ``` Slurm (multi-node example): ```bash NUM_NODES=8 srun -N ${NUM_NODES} --ntasks-per-node=1 --cpus-per-task=256 \ - primus-cli preflight --dump-path output/preflight --report-file-name preflight_report + primus-cli preflight --dump-path output/preflight ``` +If you omit `--report-file-name`, preflight auto-generates a unique +timestamped basename of the form `preflight-${NNODES}N-YYYYMMDD-HHMMSS` +so each run writes to a fresh path and never overwrites prior output. +Pass `--report-file-name NAME` only when you want a stable, well-known +filename. + ## 📂 Output Directory After running **Preflight**, all test results and reports are generated under the `output/preflight` directory. -The final reports are: +The final reports (basename shown here is the auto-generated default; it +reflects whatever `--report-file-name` resolves to) are: -- `preflight_report.md` – a Markdown version of the test report -- `preflight_report.pdf` – a PDF version of the same report +- `.md` – a Markdown version of the test report +- `.pdf` – a PDF version of the same report These reports summarize GPU performance, intra-node and inter-node communication results, and help identify potential issues within the cluster. @@ -40,8 +46,8 @@ These reports summarize GPU performance, intra-node and inter-node communication output/preflight ├── inter_node_comm ├── intra_node_comm -├── preflight_report.md -├── preflight_report.pdf +├── preflight-8N-20260715-142530.md +├── preflight-8N-20260715-142530.pdf ├── square_gemm_tflops └── ... ``` diff --git a/primus/tools/preflight/global_vars.py b/primus/tools/preflight/global_vars.py index a70f5a3d8..a99b74cec 100644 --- a/primus/tools/preflight/global_vars.py +++ b/primus/tools/preflight/global_vars.py @@ -13,9 +13,36 @@ MASTER_ADDR = os.environ.get("MASTER_ADDR", "127.0.0.1") MASTER_PORT = os.environ.get("MASTER_PORT", "29500") +# Default warmup / iteration counts. These remain importable as module-level +# constants for backward compatibility, but new code should call +# `get_warmup()` / `get_iteration()` so that `--quick` (or any future override) +# can lower the counts at runtime. WARMUP = 10 ITERATION = 50 +# Internal mutable counters used by the accessors below. +_WARMUP = WARMUP +_ITERATION = ITERATION + + +def set_warmup(value: int) -> None: + global _WARMUP + _WARMUP = int(value) + + +def set_iteration(value: int) -> None: + global _ITERATION + _ITERATION = int(value) + + +def get_warmup() -> int: + return _WARMUP + + +def get_iteration() -> int: + return _ITERATION + + _HOST_NAMES = [None] diff --git a/primus/tools/preflight/gpu/gpu_basic.py b/primus/tools/preflight/gpu/gpu_basic.py index 99f07b32d..54a79e448 100644 --- a/primus/tools/preflight/gpu/gpu_basic.py +++ b/primus/tools/preflight/gpu/gpu_basic.py @@ -178,14 +178,19 @@ def run_gpu_basic_checks( else: findings.append(Finding("info", "GPU occupancy", {"note": "amd-smi JSON not available; skipped"})) - # ROCm runtime availability: best-effort presence via tooling. + # ROCm runtime availability: sysfs is always available; amd-smi/rocm-smi + # are supplementary and only probed on LOCAL_RANK 0. if probe.backend == "rocm": - if "amd-smi" in probe.tooling or "rocm-smi" in probe.tooling: - findings.append( - Finding("info", "ROCm runtime/tooling detected", {"tooling": list(probe.tooling.keys())}) - ) + detected = [k for k in ("sysfs", "amd-smi", "rocm-smi") if k in probe.tooling] + if detected: + findings.append(Finding("info", "ROCm runtime/tooling detected", {"tooling": detected})) else: - findings.append(Finding("warn", "ROCm tooling not found (amd-smi/rocm-smi)", {})) + from primus.tools.preflight.global_vars import LOCAL_RANK + + if LOCAL_RANK != 0: + findings.append(Finding("info", "ROCm tooling skipped (non-zero LOCAL_RANK)", {})) + else: + findings.append(Finding("warn", "ROCm tooling not found (sysfs/amd-smi/rocm-smi)", {})) ok = not any(f.level == "fail" for f in findings) return {"ok": ok, "probe": probe, "findings": findings} diff --git a/primus/tools/preflight/gpu/gpu_probe.py b/primus/tools/preflight/gpu/gpu_probe.py index 5238c65ac..45bfbfb3f 100644 --- a/primus/tools/preflight/gpu/gpu_probe.py +++ b/primus/tools/preflight/gpu/gpu_probe.py @@ -7,10 +7,14 @@ from __future__ import annotations import json +import logging from typing import Any, Dict, List, Optional +from .sysfs_probe import sysfs_probe from .utils import ProbeResult, run_cmd, which +logger = logging.getLogger(__name__) + def _normalize_gfx_arch(raw: str) -> str: """ @@ -120,27 +124,19 @@ def _probe_with_torch() -> Dict[str, Any]: def _probe_amd_smi() -> Optional[Dict[str, Any]]: + """Best-effort amd-smi JSON probe for process occupancy detection only.""" if which("amd-smi") is None: return None - - # Prefer JSON if available (newer amd-smi); fall back to `list` output. - rc, out, err = run_cmd(["amd-smi", "list", "--json"], timeout_s=10) - if rc == 0 and out: - try: + try: + rc, out, err = run_cmd(["amd-smi", "list", "--json"], timeout_s=10) + if rc == 0 and out: return {"rc": rc, "json": json.loads(out), "err": err} - except Exception: - # Fall through to non-json. - pass - - rc, out, err = run_cmd(["amd-smi", "list"], timeout_s=10) - return {"rc": rc, "out": out, "err": err} + except Exception as e: + logger.debug("amd-smi probe failed: %s", e) + return None -def _probe_rocm_smi() -> Optional[Dict[str, Any]]: - if which("rocm-smi") is None: - return None - rc, out, err = run_cmd(["rocm-smi", "-a"], timeout_s=10) - return {"rc": rc, "out": out, "err": err} +_PROBE_CACHE: Optional[ProbeResult] = None def probe_gpus() -> ProbeResult: @@ -149,21 +145,54 @@ def probe_gpus() -> ProbeResult: Returns a normalized structure; individual fields may be missing depending on the environment/tooling availability. + + Results are cached since the probe is called multiple times per rank + (from gpu_basic, gpu_topology, and gpu_perf) and the output is static + during a single preflight run. + + Primary GPU enumeration uses sysfs (KFD topology), which is safe to call + from any rank (no subprocesses, no /dev/shm mutex). amd-smi is kept as + an optional probe on LOCAL_RANK == 0 for process-occupancy data only; + its failure never crashes the run. """ + global _PROBE_CACHE + if _PROBE_CACHE is not None: + return _PROBE_CACHE + + from primus.tools.preflight.global_vars import LOCAL_RANK + torch_info = _probe_with_torch() tooling: Dict[str, Any] = {"torch": torch_info} - amd = _probe_amd_smi() - if amd is not None: - tooling["amd-smi"] = amd - rocmsmi = _probe_rocm_smi() - if rocmsmi is not None: - tooling["rocm-smi"] = rocmsmi + # sysfs probe: safe from every rank, no subprocess calls. + sysfs_result = sysfs_probe() + if sysfs_result.ok: + tooling["sysfs"] = { + "gpu_count": sysfs_result.gpu_count, + "gpus": [ + { + "index": i, + "pci_bdf": g.pci_bdf_str, + "unique_id": hex(g.unique_id) if g.unique_id else None, + "numa_node": g.numa_node, + } + for i, g in enumerate(sysfs_result.gpus) + ], + "link_count": len(sysfs_result.links), + } + else: + logger.debug("sysfs probe unavailable: %s", sysfs_result.error) + + # amd-smi JSON: subprocess-based, LOCAL_RANK 0 only, best-effort. + # Used solely for process-occupancy detection in gpu_basic.py. + if LOCAL_RANK == 0: + amd = _probe_amd_smi() + if amd is not None: + tooling["amd-smi"] = amd backend = torch_info.get("backend") if torch_info.get("ok") else "unknown" devices = torch_info.get("devices", []) if torch_info.get("ok") else [] - # Version metadata (best-effort) tooling["amdgpu_version"] = _probe_amdgpu_version() tooling["rocm_version"] = _probe_rocm_version() @@ -172,4 +201,5 @@ def probe_gpus() -> ProbeResult: and bool(torch_info.get("cuda_is_available")) and int(torch_info.get("device_count", 0)) > 0 ) - return ProbeResult(ok=ok, backend=str(backend), devices=list(devices), tooling=tooling) + _PROBE_CACHE = ProbeResult(ok=ok, backend=str(backend), devices=list(devices), tooling=tooling) + return _PROBE_CACHE diff --git a/primus/tools/preflight/gpu/gpu_topology.py b/primus/tools/preflight/gpu/gpu_topology.py index e86198764..b047fc9bb 100644 --- a/primus/tools/preflight/gpu/gpu_topology.py +++ b/primus/tools/preflight/gpu/gpu_topology.py @@ -6,13 +6,17 @@ from __future__ import annotations +import logging import os from collections import Counter from typing import Any, Dict, List, Optional from .gpu_probe import probe_gpus +from .sysfs_probe import sysfs_gpu_bdfs, sysfs_topology_summary from .utils import Finding, ProbeResult, run_cmd, which +logger = logging.getLogger(__name__) + def _device_consistency(devices: List[Dict[str, Any]]) -> List[Finding]: findings: List[Finding] = [] @@ -29,9 +33,22 @@ def _device_consistency(devices: List[Dict[str, Any]]) -> List[Finding]: def _numa_mapping_best_effort() -> Optional[Dict[str, Any]]: """ - Best-effort GPU<->NUMA mapping. - On many systems this requires PCI bus IDs; we attempt to use amd-smi if present. + Best-effort GPU<->NUMA mapping via sysfs (primary) or amd-smi (fallback). + + Sysfs reads are safe from any rank (no subprocesses, no mutex). + amd-smi fallback only attempted on LOCAL_RANK == 0. """ + bdfs = sysfs_gpu_bdfs() + if bdfs: + return {"rc": 0, "gpus": bdfs, "source": "sysfs"} + + logger.debug("sysfs NUMA mapping unavailable; trying amd-smi fallback") + + from primus.tools.preflight.global_vars import LOCAL_RANK + + if LOCAL_RANK != 0: + return None + if which("amd-smi") is None: return None @@ -39,17 +56,15 @@ def _numa_mapping_best_effort() -> Optional[Dict[str, Any]]: if rc != 0 or not out: return {"rc": rc, "err": err, "note": "amd-smi list --csv failed"} - # amd-smi --csv format may vary; we do a very small heuristic: - # take 2nd column as PCI BDF if it looks like xxxx:xx:xx.x lines = [l for l in out.splitlines() if l.strip()] - bdfs: List[str] = [] + bdf_list: List[str] = [] for ln in lines[1:]: cols = [c.strip() for c in ln.split(",")] if len(cols) >= 2 and ":" in cols[1] and "." in cols[1]: - bdfs.append(cols[1]) + bdf_list.append(cols[1]) mapping: List[Dict[str, Any]] = [] - for i, bdf in enumerate(bdfs): + for i, bdf in enumerate(bdf_list): node_path = f"/sys/bus/pci/devices/{bdf}/numa_node" numa = None if os.path.exists(node_path): @@ -59,14 +74,29 @@ def _numa_mapping_best_effort() -> Optional[Dict[str, Any]]: numa = None mapping.append({"gpu": i, "pci_bdf": bdf, "numa_node": numa}) - return {"rc": rc, "gpus": mapping} + return {"rc": rc, "gpus": mapping, "source": "amd-smi"} def _xgmi_presence_best_effort() -> Optional[Dict[str, Any]]: - # Best-effort: use amd-smi topology if available; otherwise skip. + """ + Best-effort topology (XGMI vs PCIe) via sysfs (primary) or amd-smi (fallback). + + Sysfs reads are safe from any rank. amd-smi fallback only on LOCAL_RANK == 0. + """ + topo = sysfs_topology_summary() + if topo is not None and topo.get("rc") == 0: + return topo + + logger.debug("sysfs topology unavailable; trying amd-smi fallback") + + from primus.tools.preflight.global_vars import LOCAL_RANK + + if LOCAL_RANK != 0: + return None + if which("amd-smi") is None: return None - rc, out, err = run_cmd(["amd-smi", "topo"], timeout_s=10) + rc, out, err = run_cmd(["amd-smi", "topology"], timeout_s=10) if rc != 0: return {"rc": rc, "err": err} return {"rc": rc, "out": out} @@ -105,9 +135,14 @@ def run_gpu_standard_checks(*, force_topology: bool = False) -> Dict[str, Any]: findings.extend(_device_consistency(probe.devices)) # NUMA mapping / imbalance detection (best-effort). + # Prefers sysfs (safe for all ranks); falls back to amd-smi on LOCAL_RANK 0. + from primus.tools.preflight.global_vars import LOCAL_RANK + numa = _numa_mapping_best_effort() - if numa is None: - findings.append(Finding("warn", "NUMA mapping unavailable (amd-smi not found); skipped", {})) + if numa is None and LOCAL_RANK != 0: + findings.append(Finding("info", "NUMA mapping skipped (non-zero LOCAL_RANK, no sysfs)", {})) + elif numa is None: + findings.append(Finding("warn", "NUMA mapping unavailable (sysfs/amd-smi not found); skipped", {})) else: nodes = [x.get("numa_node") for x in numa.get("gpus", []) if x.get("numa_node") is not None] imbalance = False @@ -123,13 +158,16 @@ def run_gpu_standard_checks(*, force_topology: bool = False) -> Dict[str, Any]: # Topology (XGMI vs PCIe) best-effort. if force_topology or probe.backend == "rocm": topo = _xgmi_presence_best_effort() - if topo is None: - findings.append(Finding("warn", "Topology check skipped (amd-smi not found)", {})) + if topo is None and LOCAL_RANK != 0: + findings.append(Finding("info", "Topology check skipped (non-zero LOCAL_RANK, no sysfs)", {})) + elif topo is None: + findings.append(Finding("warn", "Topology check skipped (sysfs/amd-smi not found)", {})) else: - findings.append(Finding("info", "GPU topology (amd-smi topo)", topo)) - # Detect obvious PCIe fallback hint (heuristic on output). - out = str(topo.get("out", "")) - if out and ("PCIE" in out.upper() or "PCIe" in out): + source = topo.get("source", "amd-smi") + findings.append(Finding("info", f"GPU topology ({source})", topo)) + out = str(topo.get("out", "") or topo.get("matrix", "")) + has_xgmi = topo.get("has_xgmi") + if has_xgmi is False or (out and "PCIE" in out.upper() and "XGMI" not in out.upper()): findings.append(Finding("warn", "Topology indicates PCIe paths; XGMI may be absent", {})) # RCCL/NCCL env sanity (WARN-only). diff --git a/primus/tools/preflight/gpu/sysfs_probe.py b/primus/tools/preflight/gpu/sysfs_probe.py new file mode 100644 index 000000000..3c033a61f --- /dev/null +++ b/primus/tools/preflight/gpu/sysfs_probe.py @@ -0,0 +1,385 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Direct sysfs GPU probing — reads KFD topology and DRM sysfs to enumerate AMD +GPUs, PCI BDFs, NUMA mapping, and inter-GPU link topology (XGMI vs PCIe) +without invoking amd-smi or rocm-smi subprocesses. + +Approach adapted from RCCL's alt_rsmi.cc +(https://github.com/ROCm/rocm-systems/blob/develop/projects/rccl/src/misc/alt_rsmi.cc) +which reads the same KFD sysfs paths to avoid rocm_smi_lib's /dev/shm mutex +contention (rocm_smi_lib#88). +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +KFD_NODES_ROOT = "/sys/class/kfd/kfd/topology/nodes" +AMD_VENDOR_ID = 0x1002 + +LINK_TYPE_XGMI = 11 +LINK_TYPE_PCIE = 2 + + +@dataclass +class GpuNode: + node_id: int = 0 + gpu_id: int = 0 + unique_id: int = 0 + location_id: int = 0 + domain: int = 0 + bdf: int = 0 + bus: int = 0 + device: int = 0 + function: int = 0 + partition_id: int = 0 + pci_bdf_str: str = "" + numa_node: Optional[int] = None + properties: Dict[str, int] = field(default_factory=dict) + + +@dataclass +class LinkInfo: + src_idx: int = 0 + dst_idx: int = 0 + link_type: str = "unknown" + weight: int = 0 + min_bandwidth: int = 0 + max_bandwidth: int = 0 + hops: int = 0 + + +@dataclass +class SysfsProbeResult: + ok: bool = False + gpu_count: int = 0 + gpus: List[GpuNode] = field(default_factory=list) + links: List[LinkInfo] = field(default_factory=list) + error: Optional[str] = None + + +def _read_sysfs_file(path: str) -> Optional[str]: + try: + with open(path, "r", encoding="utf-8") as f: + return f.read().strip() + except Exception: + return None + + +def _read_sysfs_int(path: str) -> Optional[int]: + val = _read_sysfs_file(path) + if val is None: + return None + try: + return int(val) + except ValueError: + return None + + +def _read_kfd_properties(node_id: int) -> Dict[str, int]: + """Parse /sys/class/kfd/kfd/topology/nodes/{node_id}/properties → dict.""" + props: Dict[str, int] = {} + path = f"{KFD_NODES_ROOT}/{node_id}/properties" + content = _read_sysfs_file(path) + if content is None: + return props + for line in content.splitlines(): + parts = line.split() + if len(parts) >= 2: + try: + props[parts[0]] = int(parts[1]) + except ValueError: + pass + return props + + +def _read_kfd_gpu_id(node_id: int) -> Optional[int]: + """Read /sys/class/kfd/kfd/topology/nodes/{node_id}/gpu_id.""" + return _read_sysfs_int(f"{KFD_NODES_ROOT}/{node_id}/gpu_id") + + +def _read_link_properties(node_id: int, link_id: int) -> Dict[str, int]: + """Parse /sys/class/kfd/kfd/topology/nodes/{node_id}/io_links/{link_id}/properties.""" + path = f"{KFD_NODES_ROOT}/{node_id}/io_links/{link_id}/properties" + props: Dict[str, int] = {} + content = _read_sysfs_file(path) + if content is None: + return props + for line in content.splitlines(): + parts = line.split() + if len(parts) >= 2: + try: + props[parts[0]] = int(parts[1]) + except ValueError: + pass + return props + + +def _count_io_links(node_id: int) -> int: + """Count subdirectories in io_links/ for a given node.""" + links_dir = f"{KFD_NODES_ROOT}/{node_id}/io_links" + if not os.path.isdir(links_dir): + return 0 + count = 0 + try: + for entry in os.listdir(links_dir): + if os.path.isdir(os.path.join(links_dir, entry)) and entry not in (".", ".."): + count += 1 + except OSError: + pass + return count + + +def _bdf_to_string(domain: int, bus: int, device: int, function: int) -> str: + """Format PCI BDF as string like '0000:c1:00.0'.""" + return f"{domain:04x}:{bus:02x}:{device:02x}.{function:x}" + + +_GPU_NODES_CACHE: Optional[List["GpuNode"]] = None + + +def _enumerate_gpu_nodes() -> List[GpuNode]: + """ + Scan KFD topology nodes, filter AMD GPUs (vendor_id == 0x1002), + compute PCI BDF from location_id + domain, sort by BDF. + + Results are cached at module level since GPU topology is static + during a single preflight run and this is called from multiple paths. + + Logic mirrors ARSMI_init() from alt_rsmi.cc. + """ + global _GPU_NODES_CACHE + if _GPU_NODES_CACHE is not None: + return _GPU_NODES_CACHE + + if not os.path.isdir(KFD_NODES_ROOT): + return [] + + raw_nodes: List[Tuple[int, GpuNode]] = [] + + try: + entries = os.listdir(KFD_NODES_ROOT) + except OSError: + return [] + + for entry in entries: + if not entry.isdigit(): + continue + + node_id = int(entry) + gpu_id = _read_kfd_gpu_id(node_id) + if gpu_id is None or gpu_id == 0: + continue + + props = _read_kfd_properties(node_id) + vendor_id = props.get("vendor_id", 0) + if vendor_id != AMD_VENDOR_ID: + continue + + unique_id = props.get("unique_id", 0) + location_id = props.get("location_id", 0) + domain = props.get("domain", 0) & 0xFFFFFFFF + + bdf_raw = (domain << 32) | location_id + bus = (location_id >> 8) & 0xFF + device = (location_id >> 3) & 0x1F + function = location_id & 0x7 + partition_id = (location_id >> 28) & 0xF + + pci_bdf_str = _bdf_to_string(domain, bus, device, function) + + numa = _read_sysfs_int(f"/sys/bus/pci/devices/{pci_bdf_str}/numa_node") + + node = GpuNode( + node_id=node_id, + gpu_id=gpu_id, + unique_id=unique_id, + location_id=location_id, + domain=domain, + bdf=bdf_raw, + bus=bus, + device=device, + function=function, + partition_id=partition_id, + pci_bdf_str=pci_bdf_str, + numa_node=numa, + properties=props, + ) + raw_nodes.append((bdf_raw, node)) + + raw_nodes.sort(key=lambda x: x[0]) + _GPU_NODES_CACHE = [n for _, n in raw_nodes] + return _GPU_NODES_CACHE + + +def _build_link_matrix(gpu_nodes: List[GpuNode]) -> List[LinkInfo]: + """ + Read io_links for each GPU node to determine XGMI vs PCIe connectivity. + Returns a flat list of LinkInfo for GPU-to-GPU links only. + + Logic mirrors the link matrix construction in ARSMI_init(). + """ + if not gpu_nodes: + return [] + + node_id_to_idx = {g.node_id: i for i, g in enumerate(gpu_nodes)} + links: List[LinkInfo] = [] + + for src_idx, src_node in enumerate(gpu_nodes): + n_links = _count_io_links(src_node.node_id) + for link_id in range(n_links): + props = _read_link_properties(src_node.node_id, link_id) + if not props: + continue + + dst_node_id = props.get("node_to") + if dst_node_id is None: + continue + + dst_idx = node_id_to_idx.get(dst_node_id) + if dst_idx is None: + continue + + link_type_raw = props.get("type", 0) + weight = props.get("weight", 0) + min_bw = props.get("min_bandwidth", 0) + max_bw = props.get("max_bandwidth", 0) + + if link_type_raw == LINK_TYPE_XGMI: + link_type = "XGMI" + hops = 1 + elif link_type_raw == LINK_TYPE_PCIE: + link_type = "PCIe" + hops = 2 + else: + link_type = "unknown" + hops = 0 + + links.append( + LinkInfo( + src_idx=src_idx, + dst_idx=dst_idx, + link_type=link_type, + weight=weight, + min_bandwidth=min_bw, + max_bandwidth=max_bw, + hops=hops, + ) + ) + + return links + + +def sysfs_probe() -> SysfsProbeResult: + """ + Main entry point: enumerate AMD GPUs and topology via sysfs. + + No subprocess calls, no /dev/shm mutex, safe to call from any rank. + """ + try: + gpu_nodes = _enumerate_gpu_nodes() + if not gpu_nodes: + return SysfsProbeResult(ok=False, error="No AMD GPUs found via KFD sysfs") + + links = _build_link_matrix(gpu_nodes) + return SysfsProbeResult( + ok=True, + gpu_count=len(gpu_nodes), + gpus=gpu_nodes, + links=links, + ) + except Exception as e: + return SysfsProbeResult(ok=False, error=f"sysfs probe failed: {e}") + + +# ── Convenience helpers for integration with existing preflight code ── + + +def sysfs_gpu_count() -> int: + """GPU count via KFD sysfs. Returns 0 on failure.""" + try: + nodes = _enumerate_gpu_nodes() + return len(nodes) + except Exception: + return 0 + + +def sysfs_gpu_bdfs() -> List[Dict[str, Any]]: + """ + Return per-GPU BDF + NUMA mapping, compatible with the format + _numa_mapping_best_effort() currently returns. + """ + try: + nodes = _enumerate_gpu_nodes() + return [{"gpu": i, "pci_bdf": n.pci_bdf_str, "numa_node": n.numa_node} for i, n in enumerate(nodes)] + except Exception: + return [] + + +def sysfs_has_xgmi() -> Optional[bool]: + """ + Check if any GPU-to-GPU link is XGMI. Returns None if probe fails. + """ + try: + nodes = _enumerate_gpu_nodes() + if not nodes: + return None + links = _build_link_matrix(nodes) + return any(lk.link_type == "XGMI" for lk in links) + except Exception: + return None + + +def sysfs_topology_summary() -> Optional[Dict[str, Any]]: + """ + Build a topology summary similar to what `amd-smi topo` provides, + formatted for preflight reporting. + """ + try: + nodes = _enumerate_gpu_nodes() + if not nodes: + return None + + links = _build_link_matrix(nodes) + n = len(nodes) + + matrix: List[List[str]] = [["" for _ in range(n)] for _ in range(n)] + for i in range(n): + matrix[i][i] = "self" + + for lk in links: + matrix[lk.src_idx][lk.dst_idx] = lk.link_type + + header = [f"GPU{i}" for i in range(n)] + lines = [" " + " ".join(f"{h:>6}" for h in header)] + for i in range(n): + row = f"GPU{i:<3} " + " ".join(f"{matrix[i][j]:>6}" for j in range(n)) + lines.append(row) + + has_xgmi = any(lk.link_type == "XGMI" for lk in links) + + return { + "rc": 0, + "source": "sysfs", + "gpu_count": n, + "has_xgmi": has_xgmi, + "matrix": "\n".join(lines), + "links": [ + { + "src": lk.src_idx, + "dst": lk.dst_idx, + "type": lk.link_type, + "weight": lk.weight, + } + for lk in links + ], + } + except Exception as e: + return {"rc": 1, "error": str(e)} diff --git a/primus/tools/preflight/host/host_probe.py b/primus/tools/preflight/host/host_probe.py index fde5eb0e5..238a72db0 100644 --- a/primus/tools/preflight/host/host_probe.py +++ b/primus/tools/preflight/host/host_probe.py @@ -488,31 +488,25 @@ def _get_pcie_link_info_sysfs() -> Dict[str, Any]: return info -def get_gpu_count_rocm() -> int: - """Get GPU count using rocm-smi (more reliable in containers).""" +def get_gpu_count_sysfs() -> int: + """GPU count via KFD sysfs — no subprocesses, no /dev/shm mutex.""" try: - result = subprocess.run( - ["rocm-smi", "--showid"], - capture_output=True, - text=True, - timeout=10, - ) - if result.returncode == 0: - # Count lines that look like GPU entries (contain "GPU[") - count = sum(1 for line in result.stdout.splitlines() if "GPU[" in line) - if count > 0: - return count - except FileNotFoundError: - pass + from primus.tools.preflight.gpu.sysfs_probe import sysfs_gpu_count + + count = sysfs_gpu_count() + if count > 0: + return count except Exception: pass + return 0 + - # Fallback: try HIP_VISIBLE_DEVICES or count /dev/dri/renderD* devices +def get_gpu_count_rocm_fallback() -> int: + """GPU count without invoking rocm-smi (HIP_VISIBLE_DEVICES / /dev/dri).""" hip_devices = os.environ.get("HIP_VISIBLE_DEVICES", "") if hip_devices: return len([x for x in hip_devices.split(",") if x.strip()]) - # Count render devices try: import glob @@ -525,6 +519,27 @@ def get_gpu_count_rocm() -> int: return 0 +def get_gpu_count_rocm() -> int: + """Get GPU count using rocm-smi (more reliable in containers).""" + try: + result = subprocess.run( + ["rocm-smi", "--showid"], + capture_output=True, + text=True, + timeout=10, + ) + if result.returncode == 0: + count = sum(1 for line in result.stdout.splitlines() if "GPU[" in line) + if count > 0: + return count + except FileNotFoundError: + pass + except Exception: + pass + + return get_gpu_count_rocm_fallback() + + def get_pcie_topology() -> Optional[str]: """Get PCIe topology using lstopo or lspci -t.""" try: diff --git a/primus/tools/preflight/host/info.py b/primus/tools/preflight/host/info.py index ed99dd60d..fd5c1aa80 100644 --- a/primus/tools/preflight/host/info.py +++ b/primus/tools/preflight/host/info.py @@ -16,6 +16,8 @@ from .host_probe import ( get_cpu_info, get_gpu_count_rocm, + get_gpu_count_rocm_fallback, + get_gpu_count_sysfs, get_hostname, get_kernel_version, get_memory_info, @@ -95,9 +97,16 @@ def collect_host_info() -> List[Finding]: ib_count = sum(1 for d in pcie_devices if d["type"] == "Infiniband") eth_count = sum(1 for d in pcie_devices if d["type"] == "Ethernet") - # Use rocm-smi as fallback for GPU count (more reliable in containers) - gpu_count_rocm = get_gpu_count_rocm() - gpu_count = max(gpu_count_pcie, gpu_count_rocm) + # Primary: sysfs (KFD topology) — safe from any rank, no subprocess. + # Fallback: rocm-smi on LOCAL_RANK 0 only (subprocess, /dev/shm mutex). + gpu_count_sysfs = get_gpu_count_sysfs() + if gpu_count_sysfs > 0: + gpu_count_extra = gpu_count_sysfs + else: + from primus.tools.preflight.global_vars import LOCAL_RANK + + gpu_count_extra = get_gpu_count_rocm() if LOCAL_RANK == 0 else get_gpu_count_rocm_fallback() + gpu_count = max(gpu_count_pcie, gpu_count_extra) findings.append( Finding( diff --git a/primus/tools/preflight/inter_node_comm.py b/primus/tools/preflight/inter_node_comm.py index 2ac5213fe..8f1ca06a5 100644 --- a/primus/tools/preflight/inter_node_comm.py +++ b/primus/tools/preflight/inter_node_comm.py @@ -5,32 +5,98 @@ ############################################################################### import time +from typing import Iterable, List, Optional, Sequence, Union -import matplotlib.pyplot as plt import torch import torch.distributed as dist from primus.tools.preflight.global_vars import ( - ITERATION, LOCAL_RANK, LOCAL_WORLD_SIZE, RANK, - WARMUP, WORLD_SIZE, get_hostnames, + get_iteration, + get_warmup, ) from primus.tools.preflight.utility import ( + barrier_after_comm_destroy, create_dir, extract_first_middle_last, extract_number, + format_int_range, log, ) +# Hard cap on the per-group node count for inter-node alltoall. Real-world +# MoE training rarely dispatches across more than ~8 nodes (e.g. DeepSeek-V3's +# largest published EP=64 spans 8 nodes with per-token dispatch capped at +# 4 nodes), so 16 covers every published configuration with comfortable +# headroom. Capping here also eliminates the dominant source of EADDRINUSE at +# >=128 nodes -- uncapped all-N inter-node alltoall destroys generate enough +# IB OOB TIME_WAIT per node to exhaust the default ephemeral-port pool. +# Other inter-node tests (allreduce / p2p / ring-p2p) are unaffected. +# Intentionally not exposed as a CLI flag: this is a known-safe ceiling for +# the comm shapes preflight is supposed to characterize, not a tuning knob. +_INTER_ALLTOALL_MAX_NODES = 16 -def run_inter_node_comm(args): + +def _resolve_inter_group_sizes( + group_sizes: Optional[Sequence[Union[int, str]]], + num_nodes: int, +) -> List[int]: + """Translate user-supplied inter-node group sizes (with 'all') to ints. + + - 'all' is mapped to num_nodes. + - values > num_nodes are dropped. + - duplicates are removed and the result is sorted ascending. + """ + if group_sizes is None or len(group_sizes) == 0: + candidates = [2, 4, num_nodes] + else: + candidates = [] + for g in group_sizes: + if isinstance(g, str) and g.strip().lower() == "all": + candidates.append(num_nodes) + else: + candidates.append(int(g)) + candidates = [c for c in candidates if c >= 2 and c <= num_nodes] + return sorted(set(candidates)) + + +def run_inter_node_comm( + args, + enabled_comms: Optional[Iterable[str]] = None, + sizes_mb: Optional[Sequence[int]] = None, + group_sizes: Optional[Sequence[Union[int, str]]] = None, +): + """Inter-node allreduce / alltoall benchmark. + + Args: + args: parsed namespace. + enabled_comms: subset of {"allreduce", "alltoall"} to run. Defaults to both. + sizes_mb: message sizes in MB. + group_sizes: list of node group sizes; values > num_nodes are dropped. + 'all' is accepted as a synonym for num_nodes. Defaults to [2, 4, num_nodes]. + + Note: + The alltoall path additionally clamps every requested per-group node + count to ``_INTER_ALLTOALL_MAX_NODES`` (16) before deduping, regardless + of the cluster size or the user's ``--inter-group-sizes`` choice. The + allreduce path uses the requested sizes unchanged. + """ device = torch.device(f"cuda:{LOCAL_RANK}") - sizes = [2**i * 1024 * 1024 for i in range(1, 11)] - # sizes = [2**i * 1024 * 1024 for i in range(1, 5)] + + if sizes_mb is None or len(sizes_mb) == 0: + sizes_mb = [2**i for i in range(1, 11)] + sizes = [int(mb) * 1024 * 1024 for mb in sizes_mb] + + enabled_set = set(enabled_comms) if enabled_comms else {"allreduce", "alltoall"} + enabled_set &= {"allreduce", "alltoall"} + if not enabled_set: + log("Skip inter-node comm benchmark (no enabled comms)") + return + assert WORLD_SIZE % LOCAL_WORLD_SIZE == 0 num_nodes = WORLD_SIZE // LOCAL_WORLD_SIZE @@ -38,13 +104,30 @@ def run_inter_node_comm(args): log(f"Skip inter node comm benchmark, {num_nodes=}") return - # N-node allreduce & alltoall (adjacent pairs) - # 2-node allreduce, pair nodes: [0, 1], [2, 3], ... - # 4-node allreduce, pair nodes: [0, 1, 2, 3], [4, 5, 6, 7]... - cases = { - "allreduce": list(set([2, 4] + [num_nodes])), - "alltoall": list(set([2, 4] + [num_nodes])), - } + node_counts = _resolve_inter_group_sizes(group_sizes, num_nodes) + if not node_counts: + log("Skip inter-node comm benchmark, no valid group sizes") + return + + cases = {} + for comm in ("allreduce", "alltoall"): + if comm not in enabled_set: + continue + if comm == "alltoall": + capped = sorted({min(c, _INTER_ALLTOALL_MAX_NODES) for c in node_counts}) + if capped != list(node_counts): + log( + f" inter-alltoall: per-group node count capped at " + f"{_INTER_ALLTOALL_MAX_NODES} (requested {list(node_counts)} -> " + f"running {capped}). Real-world MoE alltoall rarely exceeds " + f"~8 nodes; see docs/preflight.md \u00a75.2 for the rationale." + ) + cases[comm] = capped + else: + cases[comm] = list(node_counts) + + warmup = get_warmup() + iteration = get_iteration() if RANK == 0: with open(args.markdown_file, "a", encoding="utf-8") as f: @@ -62,19 +145,39 @@ def run_inter_node_comm(args): latency_results = {} bandwidth_results = {} - num_procs = adjacent_nodes * LOCAL_WORLD_SIZE - num_adjacent_groups = num_nodes // adjacent_nodes + num_full_groups = num_nodes // adjacent_nodes + remainder_nodes = num_nodes % adjacent_nodes adjacent_group = None - for i_group in range(num_adjacent_groups): - group_ranks = [ - i_group * adjacent_nodes * LOCAL_WORLD_SIZE + r - for r in range(adjacent_nodes * LOCAL_WORLD_SIZE) - ] + # Track per-group member ranks for compact reporting. + all_group_ranks: List[List[int]] = [] + group_node_counts: List[int] = [] + + for i_group in range(num_full_groups): + group_start = i_group * adjacent_nodes * LOCAL_WORLD_SIZE + group_ranks = [group_start + r for r in range(adjacent_nodes * LOCAL_WORLD_SIZE)] tmp_group = dist.new_group(ranks=group_ranks) if RANK in group_ranks: assert adjacent_group is None adjacent_group = tmp_group - if RANK < num_adjacent_groups * adjacent_nodes * LOCAL_WORLD_SIZE: + all_group_ranks.append(group_ranks) + group_node_counts.append(adjacent_nodes) + + if remainder_nodes >= 2: + group_start = num_full_groups * adjacent_nodes * LOCAL_WORLD_SIZE + group_ranks = [group_start + r for r in range(remainder_nodes * LOCAL_WORLD_SIZE)] + tmp_group = dist.new_group(ranks=group_ranks) + if RANK in group_ranks: + assert adjacent_group is None + adjacent_group = tmp_group + all_group_ranks.append(group_ranks) + group_node_counts.append(remainder_nodes) + + num_procs = dist.get_world_size(adjacent_group) if adjacent_group is not None else 0 + + total_grouped_ranks = num_full_groups * adjacent_nodes * LOCAL_WORLD_SIZE + if remainder_nodes >= 2: + total_grouped_ranks += remainder_nodes * LOCAL_WORLD_SIZE + if RANK < total_grouped_ranks: assert adjacent_group is not None for size in sizes: @@ -83,7 +186,7 @@ def run_inter_node_comm(args): tensor = torch.rand(size // 2, dtype=torch.bfloat16, device=device) dist.barrier(group=adjacent_group, device_ids=[torch.cuda.current_device()]) - for _ in range(WARMUP): + for _ in range(warmup): if "allreduce" == comm: dist.all_reduce(tensor, group=adjacent_group) elif "alltoall" == comm: @@ -92,7 +195,7 @@ def run_inter_node_comm(args): assert False torch.cuda.synchronize() start = time.time() - for _ in range(ITERATION): + for _ in range(iteration): if "allreduce" == comm: dist.all_reduce(tensor, group=adjacent_group) elif "alltoall" == comm: @@ -100,7 +203,7 @@ def run_inter_node_comm(args): else: assert False torch.cuda.synchronize() - elapsed = (time.time() - start) / ITERATION + elapsed = (time.time() - start) / iteration scale = 2 if comm == "allreduce" else 1 comm_size = scale * size * (num_procs - 1) / num_procs gb_per_sec = comm_size / elapsed / 1e9 @@ -110,6 +213,7 @@ def run_inter_node_comm(args): dist.barrier(device_ids=[torch.cuda.current_device()]) if adjacent_group is not None: dist.destroy_process_group(adjacent_group) + barrier_after_comm_destroy(args.comm_cleanup_delay_sec) all_latency_results = [None for _ in range(WORLD_SIZE)] all_bandwidth_results = [None for _ in range(WORLD_SIZE)] @@ -120,68 +224,82 @@ def run_inter_node_comm(args): keys = sorted( list({k for r in all_bandwidth_results for k in (r or {}).keys()}), key=extract_number ) - max_len = max(len(s) for s in get_hostnames()) + 2 + hostnames = get_hostnames() + + # Show only the leader node's hostname; the Node range plus the + # legend at the top of the report cover the rest. + def _row_for(group_ranks: List[int], results): + leader = group_ranks[0] + host_str = hostnames[leader] + node_str = format_int_range([r // LOCAL_WORLD_SIZE for r in group_ranks]) + rank_str = format_int_range(group_ranks) + return host_str, node_str, rank_str, results[leader] + + formatted_keys = [f"{key:<6}" for key in keys] + host_col_label = "Leader hostname" + host_col_w = max(20, len(host_col_label) + 2) + header_line = ( + f"{host_col_label:<{host_col_w}} {'Node':<10} {'Rank':<10} " f"{' '.join(formatted_keys)}" + ) with open(args.markdown_file, "a", encoding="utf-8") as f: f.write(f"=======InterNodeComm - {case_name} (us)=======\n") log(f"=======InterNodeComm - {case_name} (us)=======") + log(header_line) - f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n") + f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n") f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n") - - formatted_keys = [f"{key:<6}" for key in keys] - log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}") - for rank, r in enumerate(all_latency_results): - hostname = get_hostnames()[rank] - if rank % num_procs != 0: - continue - node_id = rank // LOCAL_WORLD_SIZE - + for group_ranks in all_group_ranks: + host_str, node_str, rank_str, r = _row_for(group_ranks, all_latency_results) formatted_values = [f"{r.get(key, 0):<6.2f}" for key in keys] - log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}") - f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n") + log( + f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} " + f"{' '.join(formatted_values)}" + ) + f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n") f.write(f"\n") f.write(f"=======InterNodeComm - {case_name} (GB/s)=======\n") log(f"=======InterNodeComm - {case_name} (GB/s)=======") + log(header_line) - f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n") + f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n") f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n") - formatted_keys = [f"{key:<6}" for key in keys] - log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}") - for rank, r in enumerate(all_bandwidth_results): - hostname = get_hostnames()[rank] - if rank % num_procs != 0: - continue - node_id = rank // LOCAL_WORLD_SIZE - + for group_ranks in all_group_ranks: + host_str, node_str, rank_str, r = _row_for(group_ranks, all_bandwidth_results) formatted_values = [f"{r.get(key, 0):<6.2f}" for key in keys] - log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}") - f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n") + log( + f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} " + f"{' '.join(formatted_values)}" + ) + f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n") f.write(f"\n") if not args.plot: continue - log(f"=======Plot IntraNode {case_name} Bandwidth=======") + import matplotlib.pyplot as plt + + log(f"=======Plot InterNode {case_name} Bandwidth=======") with open(args.markdown_file, "a", encoding="utf-8") as f: f.write(f"=======Plot InterNode {case_name} Bandwidth=======\n") plot_case = f"inter_node_comm/{comm}" dump_path = f"{args.dump_path}/{plot_case}" create_dir(dump_path) print_keys = extract_first_middle_last(keys) - first_rank_bandwidth_results = [ - all_bandwidth_results[i] for i in range(len(all_bandwidth_results)) if i % num_procs == 0 - ] + leader_ranks = [g[0] for g in all_group_ranks] + first_rank_bandwidth_results = [all_bandwidth_results[i] for i in leader_ranks] num_print_ranks = len(first_rank_bandwidth_results) for size_key in print_keys: values = [r[size_key] for r in first_rank_bandwidth_results] plt.figure(figsize=(10, 4)) bars = plt.bar(range(num_print_ranks), values) - plt.xlabel(f"RankPair ({num_procs} ranks)") + plt.xlabel(f"Group (starting rank)") plt.ylabel("Bandwidth") plt.title(f"Inter Node {case_name} Bandwidth for {size_key}") - xtick_labels = [f"{i*num_procs}" for i in range(num_print_ranks)] + xtick_labels = [ + f"{leader_ranks[i]} ({group_node_counts[i]}N)" for i in range(num_print_ranks) + ] plt.xticks(range(num_print_ranks), xtick_labels) plt.grid(True, axis="y") diff --git a/primus/tools/preflight/inter_node_comm_p2p.py b/primus/tools/preflight/inter_node_comm_p2p.py index 9b38feee4..8f2574e03 100644 --- a/primus/tools/preflight/inter_node_comm_p2p.py +++ b/primus/tools/preflight/inter_node_comm_p2p.py @@ -5,32 +5,37 @@ ############################################################################### import time +from typing import Optional, Sequence -import matplotlib.pyplot as plt import torch import torch.distributed as dist from primus.tools.preflight.global_vars import ( - ITERATION, LOCAL_RANK, LOCAL_WORLD_SIZE, RANK, - WARMUP, WORLD_SIZE, get_hostnames, + get_iteration, + get_warmup, ) from primus.tools.preflight.utility import ( + barrier_after_comm_destroy, create_dir, extract_first_middle_last, extract_number, + format_int_range, log, ) -def run_inter_node_comm_p2p(args): +def run_inter_node_comm_p2p(args, sizes_mb: Optional[Sequence[int]] = None): device = torch.device(f"cuda:{LOCAL_RANK}") - sizes = [2**i * 1024 * 1024 for i in range(1, 11)] - # sizes = [2**i * 1024 * 1024 for i in range(1, 5)] + if sizes_mb is None or len(sizes_mb) == 0: + sizes_mb = [2**i for i in range(1, 11)] + sizes = [int(mb) * 1024 * 1024 for mb in sizes_mb] + warmup = get_warmup() + iteration = get_iteration() assert WORLD_SIZE % LOCAL_WORLD_SIZE == 0 num_nodes = WORLD_SIZE // LOCAL_WORLD_SIZE @@ -49,10 +54,14 @@ def run_inter_node_comm_p2p(args): bandwidth_results = {} num_adjacent_groups = num_nodes // adjacent_nodes + num_paired_ranks = num_adjacent_groups * adjacent_nodes * LOCAL_WORLD_SIZE p2p_group = None is_src_rank = ((RANK // LOCAL_WORLD_SIZE) % 2) == 0 - peer_rank = RANK + LOCAL_WORLD_SIZE if is_src_rank else RANK - LOCAL_WORLD_SIZE - assert peer_rank >= 0 and peer_rank < WORLD_SIZE + if RANK < num_paired_ranks: + peer_rank = RANK + LOCAL_WORLD_SIZE if is_src_rank else RANK - LOCAL_WORLD_SIZE + assert peer_rank >= 0 and peer_rank < WORLD_SIZE + else: + peer_rank = -1 for i_group in range(num_adjacent_groups): for i_r in range(LOCAL_WORLD_SIZE): group_ranks = [ @@ -76,20 +85,20 @@ def run_inter_node_comm_p2p(args): tensor = torch.rand(size // 2, dtype=torch.bfloat16, device=device) dist.barrier(group=p2p_group, device_ids=[torch.cuda.current_device()]) - for _ in range(WARMUP): + for _ in range(warmup): if is_src_rank: dist.send(tensor, dst=peer_rank, group=p2p_group) else: dist.recv(tensor, src=peer_rank, group=p2p_group) torch.cuda.synchronize() start = time.time() - for _ in range(ITERATION): + for _ in range(iteration): if is_src_rank: dist.send(tensor, dst=peer_rank, group=p2p_group) else: dist.recv(tensor, src=peer_rank, group=p2p_group) torch.cuda.synchronize() - elapsed = (time.time() - start) / ITERATION + elapsed = (time.time() - start) / iteration comm_size = size gb_per_sec = comm_size / elapsed / 1e9 latency_results[f"{size//1024//1024}MB"] = elapsed * 1e6 @@ -98,6 +107,7 @@ def run_inter_node_comm_p2p(args): dist.barrier(device_ids=[torch.cuda.current_device()]) if p2p_group is not None: dist.destroy_process_group(p2p_group) + barrier_after_comm_destroy(args.comm_cleanup_delay_sec) all_latency_results = [None for _ in range(WORLD_SIZE)] all_bandwidth_results = [None for _ in range(WORLD_SIZE)] @@ -106,7 +116,7 @@ def run_inter_node_comm_p2p(args): if RANK == 0: keys = sorted(list({k for r in all_bandwidth_results for k in (r or {}).keys()}), key=extract_number) - max_len = max(len(s) for s in get_hostnames()) + 2 + hostnames = get_hostnames() # result of src ranks will be print src_ranks = [] @@ -114,6 +124,8 @@ def run_inter_node_comm_p2p(args): src_rank_latency_results = [] src_rank_bandwidth_results = [] for rank, r in enumerate(all_bandwidth_results): + if rank >= num_paired_ranks: + continue is_src_rank = ((rank // LOCAL_WORLD_SIZE) % 2) == 0 peer_rank = rank + LOCAL_WORLD_SIZE if is_src_rank else rank - LOCAL_WORLD_SIZE assert peer_rank >= 0 and peer_rank < WORLD_SIZE @@ -124,45 +136,61 @@ def run_inter_node_comm_p2p(args): src_rank_latency_results.append(all_latency_results[rank]) src_rank_bandwidth_results.append(r) + # Show only the leader (src) host. Both nodes appear in the Node column, + # and the Node->Hostname legend at the top of the report covers the rest. + def _row_for(src: int, peer: int): + host_str = hostnames[src] + node_str = format_int_range([src // LOCAL_WORLD_SIZE, peer // LOCAL_WORLD_SIZE]) + rank_str = format_int_range([src, peer]) + return host_str, node_str, rank_str + + formatted_keys = [f"{key:<6}" for key in keys] + host_col_label = "Leader hostname" + host_col_w = max(20, len(host_col_label) + 2) + header_line = ( + f"{host_col_label:<{host_col_w}} {'Node':<10} {'Rank':<10} " f"{' '.join(formatted_keys)}" + ) + with open(args.markdown_file, "a", encoding="utf-8") as f: f.write(f"=======InterNodeComm - {case_name} (us)=======\n") log(f"=======InterNodeComm - {case_name} (us)=======") + log(header_line) - f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n") + f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n") f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n") - - formatted_keys = [f"{key:<6}" for key in keys] - log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}") for i_r in range(len(src_ranks)): - rank = src_ranks[i_r] - hostname = get_hostnames()[rank] - node_id = rank // LOCAL_WORLD_SIZE - + src = src_ranks[i_r] + peer = peer_ranks[i_r] + host_str, node_str, rank_str = _row_for(src, peer) formatted_values = [f"{src_rank_latency_results[i_r].get(key, 0):<6.2f}" for key in keys] - log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}") - f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n") + log( + f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} " f"{' '.join(formatted_values)}" + ) + f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n") f.write(f"\n") f.write(f"=======InterNodeComm - {case_name} (GB/s)=======\n") log(f"=======InterNodeComm - {case_name} (GB/s)=======") + log(header_line) - f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n") + f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n") f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n") - formatted_keys = [f"{key:<6}" for key in keys] - log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}") for i_r in range(len(src_ranks)): - rank = src_ranks[i_r] - hostname = get_hostnames()[rank] - node_id = rank // LOCAL_WORLD_SIZE - + src = src_ranks[i_r] + peer = peer_ranks[i_r] + host_str, node_str, rank_str = _row_for(src, peer) formatted_values = [f"{src_rank_bandwidth_results[i_r].get(key, 0):<6.2f}" for key in keys] - log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}") - f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n") + log( + f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} " f"{' '.join(formatted_values)}" + ) + f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n") f.write(f"\n") if not args.plot: return + import matplotlib.pyplot as plt + log(f"=======Plot InterNode {case_name} Bandwidth=======") with open(args.markdown_file, "a", encoding="utf-8") as f: f.write(f"=======Plot InterNode {case_name} Bandwidth=======\n") diff --git a/primus/tools/preflight/inter_node_ring_p2p.py b/primus/tools/preflight/inter_node_ring_p2p.py index 931685ae9..d5b0ab306 100644 --- a/primus/tools/preflight/inter_node_ring_p2p.py +++ b/primus/tools/preflight/inter_node_ring_p2p.py @@ -1,19 +1,20 @@ +from typing import Optional, Sequence + import torch import torch.distributed as dist from torch.profiler import ProfilerActivity from primus.tools.preflight.global_vars import ( - ITERATION, LOCAL_WORLD_SIZE, RANK, - WARMUP, WORLD_SIZE, + get_iteration, + get_warmup, ) -from primus.tools.preflight.utility import log +from primus.tools.preflight.utility import barrier_after_comm_destroy, log # profile parameters _ENABLE_PROFILE = False -_PROFILE_STEPS = min(ITERATION, 10) _NODE_RANK = RANK // LOCAL_WORLD_SIZE _GLOBAL_PIPELINE_GROUP = None @@ -76,10 +77,14 @@ def run_ring_p2p(num_bytes: int): start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) - for it in range(WARMUP): + warmup = get_warmup() + iteration = get_iteration() + profile_steps = min(iteration, 10) + + for it in range(warmup): reqs = send_recv_once(x, y) - if WARMUP > 0: + if warmup > 0: # TODO(limou) # check, does this create a cuda event # making default stream waiting for NCCL stream ? @@ -87,20 +92,20 @@ def run_ring_p2p(num_bytes: int): req.wait() if _ENABLE_PROFILE and RANK == 0: - assert ITERATION >= _PROFILE_STEPS + assert iteration >= profile_steps prof = torch.profiler.profile( activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True, schedule=torch.profiler.schedule( - wait=ITERATION - _PROFILE_STEPS, + wait=iteration - profile_steps, warmup=0, - active=_PROFILE_STEPS, + active=profile_steps, ), ) prof.start() start_event.record() - for it in range(ITERATION): + for it in range(iteration): reqs = send_recv_once(x, y) if _ENABLE_PROFILE and RANK == 0: prof.step() @@ -110,7 +115,7 @@ def run_ring_p2p(num_bytes: int): end_event.record() end_event.synchronize() - avg_time_elapsed = start_event.elapsed_time(end_event) / ITERATION + avg_time_elapsed = start_event.elapsed_time(end_event) / iteration if _ENABLE_PROFILE and RANK == 0: prof.stop() prof.export_chrome_trace(f"inter-node_ring_p2p_trace_{num_bytes}.json") @@ -177,7 +182,7 @@ def write_markdown(args, sizes_in_mb, time_statistics): """ -def run_inter_node_ring_p2p(args): +def run_inter_node_ring_p2p(args, sizes_mb: Optional[Sequence[int]] = None): assert WORLD_SIZE % LOCAL_WORLD_SIZE == 0 num_nodes = WORLD_SIZE // LOCAL_WORLD_SIZE @@ -190,16 +195,22 @@ def run_inter_node_ring_p2p(args): log("Skip inter node ring p2p benchmark") return - SIZES_IN_MB_TO_BENCH = [10, 20, 40, 80, 160] + if sizes_mb is None or len(sizes_mb) == 0: + sizes_mb = [10, 20, 40, 80, 160] + sizes_in_mb_to_bench = [int(mb) for mb in sizes_mb] time_statistics = [] - for size_in_mb in SIZES_IN_MB_TO_BENCH: + for size_in_mb in sizes_in_mb_to_bench: avg_time_elapsed = run_ring_p2p(size_in_mb * (2**20)) all_latency_results = [-1.0 for _ in range(WORLD_SIZE)] dist.gather_object(avg_time_elapsed, all_latency_results if RANK == 0 else None, dst=0) time_statistics.append(all_latency_results[:LOCAL_WORLD_SIZE]) - write_markdown(args, SIZES_IN_MB_TO_BENCH, time_statistics) + write_markdown(args, sizes_in_mb_to_bench, time_statistics) + + if _GLOBAL_PIPELINE_GROUP is not None: + dist.destroy_process_group(_GLOBAL_PIPELINE_GROUP) + barrier_after_comm_destroy(args.comm_cleanup_delay_sec) # TODO (limou) # support plot diff --git a/primus/tools/preflight/intra_node_comm.py b/primus/tools/preflight/intra_node_comm.py index 982f94b46..87eb97c4b 100644 --- a/primus/tools/preflight/intra_node_comm.py +++ b/primus/tools/preflight/intra_node_comm.py @@ -5,36 +5,69 @@ ############################################################################### import time +from typing import Iterable, List, Optional, Sequence -import matplotlib.pyplot as plt import torch import torch.distributed as dist from primus.tools.preflight.global_vars import ( - ITERATION, LOCAL_RANK, LOCAL_WORLD_SIZE, RANK, - WARMUP, WORLD_SIZE, get_hostnames, + get_iteration, + get_warmup, ) from primus.tools.preflight.utility import ( + barrier_after_comm_destroy, create_dir, extract_first_middle_last, extract_number, + format_int_range, log, ) -def run_intra_node_comm(args): +def run_intra_node_comm( + args, + enabled_comms: Optional[Iterable[str]] = None, + sizes_mb: Optional[Sequence[int]] = None, + group_sizes: Optional[Sequence[int]] = None, +): + """Intra-node allreduce / alltoall benchmark. + + Args: + args: parsed namespace (must have markdown_file, dump_path, plot, ib_bw). + enabled_comms: subset of {"allreduce", "alltoall"} to run. Defaults to both. + sizes_mb: message sizes in MB. Defaults to powers of two from 2..1024. + group_sizes: GPU group sizes; each must divide LOCAL_WORLD_SIZE. + Defaults to [2, 4, 8]. + """ device = torch.device(f"cuda:{LOCAL_RANK}") - sizes = [2**i * 1024 * 1024 for i in range(1, 11)] - # sizes = [2**i * 1024 * 1024 for i in range(1, 5)] - cases = { - "allreduce": [2, 4, 8], - "alltoall": [2, 4, 8], - } + + if sizes_mb is None or len(sizes_mb) == 0: + sizes_mb = [2**i for i in range(1, 11)] + sizes = [int(mb) * 1024 * 1024 for mb in sizes_mb] + + enabled_set = set(enabled_comms) if enabled_comms else {"allreduce", "alltoall"} + enabled_set &= {"allreduce", "alltoall"} + if not enabled_set: + log("Skip intra-node comm benchmark (no enabled comms)") + return + + if group_sizes is None or len(group_sizes) == 0: + group_sizes = [2, 4, 8] + # Filter out invalid group sizes (must divide LOCAL_WORLD_SIZE) and de-dupe. + group_sizes = sorted({int(g) for g in group_sizes if int(g) > 0 and LOCAL_WORLD_SIZE % int(g) == 0}) + if not group_sizes: + log(f"Skip intra-node comm benchmark, no valid group sizes for LOCAL_WORLD_SIZE={LOCAL_WORLD_SIZE}") + return + + cases = {comm: list(group_sizes) for comm in ("allreduce", "alltoall") if comm in enabled_set} + + warmup = get_warmup() + iteration = get_iteration() if RANK == 0: with open(args.markdown_file, "a", encoding="utf-8") as f: @@ -54,6 +87,9 @@ def run_intra_node_comm(args): num_nodes = WORLD_SIZE // LOCAL_WORLD_SIZE num_groups_per_node = LOCAL_WORLD_SIZE // num_procs group = None + # Track which ranks belong to each group so we can render compact rows. + all_group_ranks: List[List[int]] = [] + my_group_index = -1 for i_node in range(num_nodes): for i_group in range(num_groups_per_node): group_ranks = [ @@ -63,12 +99,15 @@ def run_intra_node_comm(args): if RANK in group_ranks: assert group is None group = tmp_group + my_group_index = len(all_group_ranks) + all_group_ranks.append(group_ranks) assert group is not None + assert my_group_index >= 0 for size in sizes: tensor = torch.rand(size // 2, dtype=torch.bfloat16, device=device) dist.barrier(group=group, device_ids=[torch.cuda.current_device()]) - for _ in range(WARMUP): + for _ in range(warmup): if "allreduce" == comm: dist.all_reduce(tensor, group=group) elif "alltoall" == comm: @@ -77,7 +116,7 @@ def run_intra_node_comm(args): assert False torch.cuda.synchronize() start = time.time() - for _ in range(ITERATION): + for _ in range(iteration): if "allreduce" == comm: dist.all_reduce(tensor, group=group) elif "alltoall" == comm: @@ -85,7 +124,7 @@ def run_intra_node_comm(args): else: assert False torch.cuda.synchronize() - elapsed = (time.time() - start) / ITERATION + elapsed = (time.time() - start) / iteration scale = 2 if comm == "allreduce" else 1 comm_size = scale * size * (num_procs - 1) / num_procs gb_per_sec = comm_size / elapsed / 1e9 @@ -96,6 +135,7 @@ def run_intra_node_comm(args): # destroy this parallel group dist.destroy_process_group(group) + barrier_after_comm_destroy(args.comm_cleanup_delay_sec) all_latency_results = [None for _ in range(WORLD_SIZE)] all_bandwidth_results = [None for _ in range(WORLD_SIZE)] @@ -106,48 +146,65 @@ def run_intra_node_comm(args): keys = sorted( list({k for r in all_bandwidth_results for k in (r or {}).keys()}), key=extract_number ) - max_len = max(len(s) for s in get_hostnames()) + 2 + hostnames = get_hostnames() + + # Each row corresponds to one group (group_ranks). Use the first + # rank's results since all members observe the same collective. + # The Hostname column shows only the leader's host (compact); + # use the Node range plus the legend at the top of the report + # to look up the rest. + def _row_for(group_ranks: List[int], results): + leader = group_ranks[0] + host_str = hostnames[leader] + node_str = format_int_range([r // LOCAL_WORLD_SIZE for r in group_ranks]) + rank_str = format_int_range(group_ranks) + return host_str, node_str, rank_str, results[leader] + + formatted_keys = [f"{key:<6}" for key in keys] + host_col_label = "Leader hostname" + host_col_w = max(20, len(host_col_label) + 2) + header_line = ( + f"{host_col_label:<{host_col_w}} {'Node':<10} {'Rank':<10} " f"{' '.join(formatted_keys)}" + ) with open(args.markdown_file, "a", encoding="utf-8") as f: f.write(f"=======IntraNodeComm - {case_name} (us)=======\n") log(f"=======IntraNodeComm - {case_name} (us)=======") + log(header_line) - f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n") + f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n") f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n") - formatted_keys = [f"{key:<6}" for key in keys] - log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}") - for rank, r in enumerate(all_latency_results): - hostname = get_hostnames()[rank] - if rank % num_procs != 0: - continue - node_id = rank // LOCAL_WORLD_SIZE - + for group_ranks in all_group_ranks: + host_str, node_str, rank_str, r = _row_for(group_ranks, all_latency_results) formatted_values = [f"{r.get(key, 0):<6.2f}" for key in keys] - log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}") - f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n") + log( + f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} " + f"{' '.join(formatted_values)}" + ) + f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n") f.write(f"\n") f.write(f"=======IntraNodeComm - {case_name} (GB/s)=======\n") log(f"=======IntraNodeComm - {case_name} (GB/s)=======") + log(header_line) - f.write(f"| Hostname | Node | Rank | {' | '.join(keys)}|\n") + f.write(f"| {host_col_label} | Node | Rank | {' | '.join(keys)}|\n") f.write(f"|----------|----------|----------{'|----------' * len(keys)}|\n") - formatted_keys = [f"{key:<6}" for key in keys] - log(f"{'Hostname':<{max_len}} {'Node':<5} {'Rank':<5} {' '.join(formatted_keys)}") - for rank, r in enumerate(all_bandwidth_results): - hostname = get_hostnames()[rank] - if rank % num_procs != 0: - continue - node_id = rank // LOCAL_WORLD_SIZE - + for group_ranks in all_group_ranks: + host_str, node_str, rank_str, r = _row_for(group_ranks, all_bandwidth_results) formatted_values = [f"{r.get(key, 0):<6.2f}" for key in keys] - log(f"{hostname:<{max_len}} {node_id:<5} {rank:<5} {' '.join(formatted_values)}") - f.write(f"| {hostname} | {node_id} | {rank} | {' | '.join(formatted_values)}|\n") + log( + f"{host_str:<{host_col_w}} {node_str:<10} {rank_str:<10} " + f"{' '.join(formatted_values)}" + ) + f.write(f"| {host_str} | {node_str} | {rank_str} | {' | '.join(formatted_values)}|\n") f.write(f"\n") if not args.plot: continue + import matplotlib.pyplot as plt + log(f"=======Plot IntraNode {case_name} Bandwidth=======") with open(args.markdown_file, "a", encoding="utf-8") as f: f.write(f"=======Plot IntraNode {case_name} Bandwidth=======\n") @@ -155,9 +212,9 @@ def run_intra_node_comm(args): dump_path = f"{args.dump_path}/{plot_case}" create_dir(dump_path) print_keys = extract_first_middle_last(keys) - first_rank_bandwidth_results = [ - all_bandwidth_results[i] for i in range(len(all_bandwidth_results)) if i % num_procs == 0 - ] + # Use leader rank of each group for plotting. + leader_ranks = [g[0] for g in all_group_ranks] + first_rank_bandwidth_results = [all_bandwidth_results[i] for i in leader_ranks] num_print_ranks = len(first_rank_bandwidth_results) for size_key in print_keys: values = [r[size_key] for r in first_rank_bandwidth_results] @@ -166,7 +223,7 @@ def run_intra_node_comm(args): plt.xlabel(f"RankPair ({num_procs} ranks)") plt.ylabel("Bandwidth") plt.title(f"Intra Node {case_name} bandwidth for {size_key}") - xtick_labels = [f"{i*num_procs}" for i in range(num_print_ranks)] + xtick_labels = [f"{leader_ranks[i]}" for i in range(num_print_ranks)] plt.xticks(range(num_print_ranks), xtick_labels) plt.grid(True, axis="y") diff --git a/primus/tools/preflight/node_smoke/__init__.py b/primus/tools/preflight/node_smoke/__init__.py new file mode 100644 index 000000000..b4fce7177 --- /dev/null +++ b/primus/tools/preflight/node_smoke/__init__.py @@ -0,0 +1,58 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Node-local preflight smoke test. + +Goal: quickly identify which nodes have problems before launching a real +training job. Each node tests itself in parallel with no global rendezvous, +writes a per-node JSON verdict, and an aggregator (rank 0) reads all JSONs +and emits PASS/FAIL lists usable by SLURM ``--exclude=`` / ``--nodelist=``. + +Subcommands +----------- + +* ``run`` -- per-node entry. Runs Tier 1 (per-GPU sanity + reused + host/gpu/network info collectors + a small dmesg scan) and, when + ``--tier2-perf`` is set, Tier 2 perf sanity (GEMM TFLOPS, HBM bandwidth, + and node-local RCCL all-reduce). Writes ``/smoke/.json``. + Always exits 0 when the JSON was written (the per-node verdict lives + in the JSON's ``status`` field, not in the exit code) -- otherwise an + intentionally-detected unhealthy node would make srun pollute its + output with one ``error: ... task N: Exited with exit code 1`` per + failing node, which is misleading: the smoke test is succeeding at + identifying bad nodes, not failing. + +* ``aggregate`` -- read all per-node JSONs and emit + ``/smoke_report.md``, ``/passing_nodes.txt``, and + ``/failing_nodes.txt``. Exits non-zero if any node FAILs or is + missing -- this is the single CI-friendly cluster-health exit signal. + +* ``_per_gpu`` -- internal subcommand spawned by ``run`` to test a single + GPU in an isolated subprocess with a hard timeout. Not for direct use. + +Why per-GPU subprocesses? +------------------------- + +A stuck ``torch.cuda.set_device(i)`` cannot be aborted reliably with +``signal.alarm`` because the call may be inside a non-interruptible driver +syscall. By running each per-GPU test in its own subprocess we can SIGKILL +it on timeout without affecting the rest of the node's checks. + +Package layout +-------------- + +This module is a sub-package -- the implementation is split across many +small files mirroring the Tier 1 A-G section structure. The single +public entry point exported here is :func:`main`, which is also the +target of ``python -m primus.tools.preflight.node_smoke``. +""" + +from __future__ import annotations + +from .cli import main + +__all__ = ["main"] diff --git a/primus/tools/preflight/node_smoke/__main__.py b/primus/tools/preflight/node_smoke/__main__.py new file mode 100644 index 000000000..86cef62ab --- /dev/null +++ b/primus/tools/preflight/node_smoke/__main__.py @@ -0,0 +1,14 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Entry point for ``python -m primus.tools.preflight.node_smoke``.""" + +from __future__ import annotations + +from .cli import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/primus/tools/preflight/node_smoke/aggregator/__init__.py b/primus/tools/preflight/node_smoke/aggregator/__init__.py new file mode 100644 index 000000000..74a9229d5 --- /dev/null +++ b/primus/tools/preflight/node_smoke/aggregator/__init__.py @@ -0,0 +1,18 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Aggregator -- combines per-node JSONs into a cluster report. + +* :mod:`.summarizers` -- pure data-shaping helpers used by the report + writer (one ``_*_rows`` / ``_*_summary`` function per Markdown + section). They never raise; missing data degrades to empty rows. +* :mod:`.report` -- the markdown writer. The single + :func:`.report.write_smoke_report` entry point composes the report + by calling small per-section ``_write_
`` helpers, in the + exact order the original monolithic block produced. +""" + +from __future__ import annotations diff --git a/primus/tools/preflight/node_smoke/aggregator/report.py b/primus/tools/preflight/node_smoke/aggregator/report.py new file mode 100644 index 000000000..58a9ca95d --- /dev/null +++ b/primus/tools/preflight/node_smoke/aggregator/report.py @@ -0,0 +1,719 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Cluster smoke-report markdown writer. + +The single entry point :func:`write_smoke_report` composes the report by +calling small per-section ``_write_
`` helpers in the exact +order the original monolithic block produced. Behaviour is preserved +verbatim, including: + +* every section header +* the per-section ``try/except`` wrapping pattern +* the two intentional EXCEPTIONS to that pattern -- ``Tier 2 perf + summary`` and ``Failing nodes -- full reasons`` -- which deliberately + do NOT have their own ``try/except`` in the original (so a regression + in either bubbles up rather than being swallowed). Per refactor + guardrails these stay un-wrapped. +""" + +from __future__ import annotations + +from typing import IO, Any, Dict, List + +from ..logging_utils import _warn +from .summarizers import ( + _busy_gpu_rows, + _clock_summary, + _gpu_activity_rows, + _gpu_low_level_outlier_rows, + _host_limits_issue_rows, + _nic_excluded_rows, + _nic_fw_drift_rows, + _nic_issue_rows, + _pretouch_hbm_rows, + _stack_drift_rows, + _tooling_inventory_rows, + _tooling_latency_rows, + _xgmi_issue_rows, +) + + +def _write_header( + f: IO[str], + nodes: List[Dict[str, Any]], + passing: List[Dict[str, Any]], + failing: List[Dict[str, Any]], + expected: Any, +) -> None: + """Write the report title, summary counts, and per-node status table.""" + f.write("# Node-Local Smoke Test Report\n\n") + f.write(f"- **Expected nodes**: `{expected if expected is not None else 'unknown'}`\n") + f.write(f"- **Reported nodes**: `{len(nodes)}`\n") + f.write(f"- **PASS**: `{len(passing)}` **FAIL**: `{len(failing)}`\n\n") + f.write("| Node | Hostname | Status | Duration | Top fail reason |\n") + f.write("|------|----------|--------|----------|-----------------|\n") + for n in nodes: + reasons = n.get("fail_reasons") or [] + top = (reasons[0] if reasons else "").replace("|", "/") + if len(top) > 120: + top = top[:117] + "..." + f.write( + f"| {n.get('node_rank', '?')} | {n.get('host', '?')} | " + f"{n.get('status', '?')} | {n.get('duration_sec', 0)}s | {top} |\n" + ) + + +def _write_stack_drift(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + # ----- A. Stack drift across cluster ----- + # Empty section when every node reports the same value for every + # scalar fingerprint key. We always print the section header so the + # operator can see at a glance that the check ran. + # Each helper is wrapped so a single bug in one section can never + # truncate the whole report; the failure is recorded inline so the + # operator still sees something for that section. + f.write("\n## Stack drift across cluster\n\n") + try: + drift = _stack_drift_rows(nodes) + if not drift: + f.write("*All nodes match.*\n") + else: + f.write("| Key | Majority (count/total) | Outlier nodes |\n") + f.write("|------|-------------------------|----------------|\n") + for row in drift: + outliers = "; ".join(f"`{h}` = `{v}`" for h, v in row["outliers"]) + f.write( + f"| `{row['key']}` | `{row['majority']}` " + f"({row['count']}/{row['total']}) | {outliers} |\n" + ) + except Exception as e: + f.write(f"*Stack-drift section failed to render: {e}*\n") + _warn(f"stack-drift render failed: {e}") + + +def _write_nic_fw_drift(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + # ----- A.2 NIC firmware drift across cluster ----- + f.write("\n## NIC firmware drift across cluster\n\n") + try: + nic_drift = _nic_fw_drift_rows(nodes) + if not nic_drift: + f.write("*All NIC firmwares match (or no NICs reported).*\n") + else: + f.write("| NIC | Majority FW (count/total) | Outlier nodes |\n") + f.write("|-----|---------------------------|----------------|\n") + for row in nic_drift: + outliers = "; ".join(f"`{h}` = `{v}`" for h, v in row["outliers"]) + f.write( + f"| `{row['device']}` | `{row['majority']}` " + f"({row['count']}/{row['total']}) | {outliers} |\n" + ) + except Exception as e: + f.write(f"*NIC firmware drift section failed to render: {e}*\n") + _warn(f"nic-fw-drift render failed: {e}") + + +def _write_nic_issues(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + # ----- B. NIC / RDMA roll-call issues ----- + f.write("\n## NIC / RDMA roll-call issues\n\n") + try: + nic_issues = _nic_issue_rows(nodes) + if not nic_issues: + f.write("*No NIC issues.*\n") + else: + f.write("| Node | Hostname | Issue |\n") + f.write("|------|----------|-------|\n") + for row in nic_issues: + msg = str(row["issue"]).replace("|", "/") + if len(msg) > 160: + msg = msg[:157] + "..." + f.write(f"| {row['node_rank']} | {row['host']} | {msg} |\n") + except Exception as e: + f.write(f"*NIC issues section failed to render: {e}*\n") + _warn(f"nic-issues render failed: {e}") + + +def _write_nic_port_count(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + # ----- B.2 NIC port-count summary (helps spot "node X has fewer + # training NICs than the cluster") -- always rendered, even when no + # per-port issue tripped. We flag any node whose **included** port + # count differs from the cluster majority so operators can act on + # partial-degradation cases like 7/8 ports without having to set + # --expected-rdma-nics. The count is taken from `included_ports` + # (i.e. after the training-NIC selector ran) so nodes that legitimately + # have extra frontend / storage RoCE NICs don't show up as anomalies. + f.write("\n## NIC port-count summary\n\n") + try: + from collections import Counter + + counts = [] + for n in nodes: + nic = (n.get("tier1") or {}).get("nics") or {} + # Prefer the post-selector count; fall back to total ports + # for legacy JSONs written before --rdma-nic-allowlist + # existed (those don't have an `included_ports` key). + included = nic.get("included_ports") + if included is None: + count = len(nic.get("ports") or []) + else: + count = len(included) + counts.append( + ( + n.get("node_rank", "?"), + n.get("host", "?"), + count, + ) + ) + if not counts: + f.write("*No NIC data reported.*\n") + else: + cnt = Counter(c for *_, c in counts) + majority_count, _ = cnt.most_common(1)[0] + anomalies = [(nr, h, c) for nr, h, c in counts if c != majority_count] + f.write( + f"Cluster-majority training-NIC count: **{majority_count}** " + f"(seen on {cnt[majority_count]}/{len(counts)} nodes).\n\n" + ) + if not anomalies: + f.write("*Every node reports the majority count.*\n") + else: + f.write("| Node | Hostname | Training NICs found |\n") + f.write("|------|----------|---------------------|\n") + for nr, h, c in anomalies: + f.write(f"| {nr} | {h} | {c} |\n") + except Exception as e: + f.write(f"*NIC port-count summary failed to render: {e}*\n") + _warn(f"nic-port-count render failed: {e}") + + +def _write_nic_excluded(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + # ----- B.3 Excluded NIC ports -- informational only. Surfaces ports + # that the selector chain dropped from the training-NIC set so the + # operator can verify the heuristic / NCCL_IB_HCA / --rdma-nic-allowlist + # did what they expected (e.g. "are my front-end NICs still + # admin-down?"). Excluded ports do NOT contribute to the node FAIL + # signal; they live in `tier1.nics.excluded_ports` + `info_issues`. + f.write("\n## NIC excluded ports (informational)\n\n") + try: + rows = _nic_excluded_rows(nodes) + if not rows: + f.write( + "*No NIC ports were excluded -- every discovered RDMA " + "port is in the training-NIC set on every node.*\n" + ) + else: + # Brief per-source summary so the operator immediately knows + # whether to expect output here (e.g. "yes, NCCL_IB_HCA is + # filtering out 4 ports/node like I configured"). + from collections import Counter + + src_counts = Counter(r["source"] for r in rows) + src_bits = [] + for src, c in sorted(src_counts.items()): + pretty = { + "cli": "--rdma-nic-allowlist", + "env": "NCCL_IB_HCA env", + "heuristic": "phys_state heuristic", + }.get(src, src) + src_bits.append(f"{c} via `{pretty}`") + f.write( + "Ports excluded from the training-NIC set " + "(" + ", ".join(src_bits) + "). " + "Hard-fail rules (state ACTIVE, phys_state LinkUp, " + "RoCE v2 GIDs) did NOT run on these ports.\n\n" + ) + f.write("| Node | Hostname | Source | Port + reason |\n") + f.write("|------|----------|--------|---------------|\n") + for row in rows: + msg = str(row["issue"]).replace("|", "/") + if len(msg) > 160: + msg = msg[:157] + "..." + f.write(f"| {row['node_rank']} | {row['host']} | " f"`{row['source']}` | {msg} |\n") + except Exception as e: + f.write(f"*NIC excluded-ports section failed to render: {e}*\n") + _warn(f"nic-excluded render failed: {e}") + + +def _write_host_limits(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + # ----- C. Host limits issues ----- + f.write("\n## Host limits issues\n\n") + try: + limits_issues = _host_limits_issue_rows(nodes) + if not limits_issues: + f.write("*No host-limit issues.*\n") + else: + f.write("| Node | Hostname | Issue |\n") + f.write("|------|----------|-------|\n") + for row in limits_issues: + msg = str(row["issue"]).replace("|", "/") + if len(msg) > 200: + msg = msg[:197] + "..." + f.write(f"| {row['node_rank']} | {row['host']} | {msg} |\n") + except Exception as e: + f.write(f"*Host limits section failed to render: {e}*\n") + _warn(f"host-limits render failed: {e}") + + +def _write_gpu_visibility(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + # ----- GPU visibility issues (no GPUs / amd-smi vs torch mismatch) ----- + # Independent guard -- doesn't rely on the reused gpu_info collector + # emitting a level=fail finding, which has been known to silently + # downgrade to warn when collect_gpu_info() raises. + f.write("\n## GPU visibility issues\n\n") + try: + vis_rows: List[Dict[str, Any]] = [] + for n in nodes: + vis = (n.get("tier1") or {}).get("gpu_visibility") or {} + for issue in vis.get("fail_reasons", []) or []: + vis_rows.append( + { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + "torch": vis.get("torch_visible"), + "amd_smi": vis.get("amd_smi_visible"), + "expected": vis.get("expected_gpus"), + "issue": issue, + } + ) + if not vis_rows: + f.write( + "*Every node resolved expected_gpus >= 1 and torch + " "amd-smi agree on the GPU count.*\n" + ) + else: + f.write( + "Nodes where the GPU is invisible to torch, or where " + "amd-smi sees more GPUs than torch (stale ROCm / wedged " + "amdgpu driver). These are hard fails independent of " + "every other collector.\n\n" + ) + f.write("| Node | Hostname | expected | torch | amd-smi | Issue |\n") + f.write("|------|----------|----------|-------|---------|-------|\n") + for row in vis_rows: + msg = str(row["issue"]).replace("|", "/") + if len(msg) > 200: + msg = msg[:197] + "..." + f.write( + f"| {row['node_rank']} | {row['host']} | " + f"{row['expected']} | {row['torch']} | " + f"{row['amd_smi']} | {msg} |\n" + ) + except Exception as e: + f.write(f"*GPU visibility section failed to render: {e}*\n") + _warn(f"gpu-visibility render failed: {e}") + + +def _write_gpu_low_level(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + # ----- D-1: GPU low-level outliers (PCIe link, HBM total) ----- + f.write("\n## GPU low-level outliers (PCIe link / HBM)\n\n") + try: + gpu_outliers = _gpu_low_level_outlier_rows(nodes) + if not gpu_outliers: + f.write("*All GPUs match the cluster majority on PCIe link " "and HBM total.*\n") + else: + f.write( + "Per-GPU values that differ from the cluster majority. A " + "GPU sitting at half PCIe width / half HBM is almost " + "always a hardware fault on that single device.\n\n" + ) + f.write("| Metric | Cluster majority (count/total) | " "Outliers (`host:gpu` = value) |\n") + f.write("|--------|---------------------------------|" "-------------------------------|\n") + for row in gpu_outliers: + out_str = "; ".join(f"`{h}:{g}` = `{v}`" for h, g, v in row["outliers"]) + f.write( + f"| {row['label']} | `{row['majority']}` " + f"({row['count']}/{row['total']}) | {out_str} |\n" + ) + except Exception as e: + f.write(f"*GPU low-level section failed to render: {e}*\n") + _warn(f"gpu-low-level render failed: {e}") + + +def _write_xgmi(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + # ----- D-2: XGMI link issues ----- + f.write("\n## XGMI link issues\n\n") + try: + xgmi_issues = _xgmi_issue_rows(nodes) + if not xgmi_issues: + f.write("*All GPU pairs report XGMI on every node " "(or amd-smi topology was unavailable).*\n") + else: + f.write( + "Any non-XGMI GPU pair is a hard fail -- intra-node " + "collectives silently fall back to PCIe and lose 5-10x " + "of the bandwidth NCCL/RCCL expects.\n\n" + ) + f.write("| Node | Hostname | Issue |\n") + f.write("|------|----------|-------|\n") + for row in xgmi_issues: + msg = str(row["summary"]).replace("|", "/") + if len(msg) > 200: + msg = msg[:197] + "..." + f.write(f"| {row['node_rank']} | {row['host']} | {msg} |\n") + except Exception as e: + f.write(f"*XGMI section failed to render: {e}*\n") + _warn(f"xgmi render failed: {e}") + + +def _write_clock( + f: IO[str], + nodes: List[Dict[str, Any]], + skew_warn_sec: float, +) -> None: + # ----- E: cluster wall-clock spread + time-daemon roll-call ----- + f.write("\n## Cluster clock + time daemons\n\n") + try: + clk = _clock_summary(nodes, skew_warn_sec=skew_warn_sec) + spread = clk["spread_sec"] + if spread is None: + f.write("*Not enough nodes reported a wall-clock timestamp.*\n") + else: + marker = " (**warn** -- exceeds " f"{clk['spread_warn_sec']}s)" if clk["spread_warn"] else "" + f.write( + f"- Wall-clock spread across {clk['n_nodes_with_time']} " f"nodes: **{spread}s**{marker}.\n" + ) + f.write(f"- Earliest: `{clk['earliest_host']}`, " f"latest: `{clk['latest_host']}`.\n") + f.write( + "- (Spread is an upper bound on real clock skew -- it " "also includes srun launch jitter.)\n" + ) + if clk["no_daemon_hosts"]: + f.write( + "\n**Nodes with no active time-sync daemon " "(chronyd / ntpd / systemd-timesyncd):**\n\n" + ) + f.write("| Node | Hostname |\n") + f.write("|------|----------|\n") + for nr, h in clk["no_daemon_hosts"]: + f.write(f"| {nr} | {h} |\n") + else: + f.write("\n*Every node has at least one active time-sync " "daemon.*\n") + except Exception as e: + f.write(f"*Clock section failed to render: {e}*\n") + _warn(f"clock render failed: {e}") + + +def _write_tooling_latency( + f: IO[str], + nodes: List[Dict[str, Any]], + rocm_smi_warn_sec: float, +) -> None: + # ----- F-partial: rocm-smi self-latency ----- + f.write("\n## Tooling self-latency (`rocm-smi --version`)\n\n") + try: + tool_rows = _tooling_latency_rows( + nodes, + warn_sec=float(rocm_smi_warn_sec), + ) + if not tool_rows: + f.write("*No nodes exceeded the warn threshold " f"({rocm_smi_warn_sec}s) and no timeouts.*\n") + else: + f.write( + "Slow `rocm-smi --version` calls historically precede a " + "wedged amdgpu driver. Hitting the hard timeout is a " + "node FAIL; slow-but-completed calls are warn-only.\n\n" + ) + f.write("| Node | Hostname | Latency (s) | Flag |\n") + f.write("|------|----------|-------------|------|\n") + for r in tool_rows: + lat = r.get("latency_sec") + lat_s = f"{lat:.2f}" if isinstance(lat, (int, float)) else "?" + f.write(f"| {r['node_rank']} | {r['host']} | " f"{lat_s} | {r['flag']} |\n") + except Exception as e: + f.write(f"*Tooling section failed to render: {e}*\n") + _warn(f"tooling render failed: {e}") + + +def _write_tooling_availability( + f: IO[str], + nodes: List[Dict[str, Any]], +) -> None: + # ----- Tooling availability (always-on; loud counterweight to + # the silent skips that happen when amd-smi / rocm-smi / lsof + # are missing from PATH) ----- + f.write("\n## Tooling availability\n\n") + try: + inv = _tooling_inventory_rows(nodes) + tracked = inv["tracked"] + mc = inv["missing_counts"] + if not inv["any_missing"]: + f.write( + "*Every tracked tool (`" + "`, `".join(tracked) + "`) was present in PATH on every node.*\n" + ) + else: + summary_bits = [] + for t in tracked: + if mc[t]: + summary_bits.append(f"`{t}` missing on **{mc[t]}** node(s)") + # Compute cluster-wide uncovered-checks summary so the + # operator immediately knows whether the missing tools + # actually leave a coverage hole or whether the rocm-smi + # / lsof fallbacks are picking up the slack. + uncovered_counts: Dict[str, int] = {} + for n in nodes: + uc = ((n.get("tier1") or {}).get("tooling_inventory") or {}).get("uncovered") or [] + for c in uc: + uncovered_counts[c] = uncovered_counts.get(c, 0) + 1 + if uncovered_counts: + uc_bits = [f"`{c}` on **{n}** node(s)" for c, n in sorted(uncovered_counts.items())] + coverage_line = ( + "Checks with NO working tool (truly silent-skipped): " + "; ".join(uc_bits) + "." + ) + else: + coverage_line = ( + "Every check is still covered via the rocm-smi " + "or lsof fallback on every node -- no checks are " + "silently skipped." + ) + f.write( + "Several Tier 1 checks (ECC, XGMI, foreign-process, " + "GPU activity, wedged-driver) prefer `amd-smi` but " + "fall back to `rocm-smi` (and `lsof` for foreign-" + "process) when amd-smi is missing. " + + "; ".join(summary_bits) + + ". " + + coverage_line + + " Add `--require-tools amd-smi,rocm-smi` to `run` to " + "promote a missing tool to a node FAIL anyway.\n\n" + ) + # Per-node table -- only the nodes that ARE missing something, + # so a healthy cluster doesn't get a giant N-row table. + f.write("| Node | Hostname | " + " | ".join(tracked) + " |\n") + f.write("|------|----------| " + " | ".join("---" for _ in tracked) + " |\n") + for r in inv["rows"]: + if all(r.get(t) for t in tracked): + continue + cells = [] + for t in tracked: + cells.append("OK" if r.get(t) else "**MISSING**") + f.write(f"| {r['node_rank']} | {r['host']} | " + " | ".join(cells) + " |\n") + except Exception as e: + f.write(f"*Tooling availability section failed to render: {e}*\n") + _warn(f"tooling-availability render failed: {e}") + + +def _write_busy_gpus(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + # ----- G: Busy GPUs / leaked processes ----- + f.write("\n## Busy GPUs / leaked processes\n\n") + try: + busy_rows = _busy_gpu_rows(nodes) + if not busy_rows: + f.write( + "*No foreign processes detected on any GPU " + "(or `amd-smi process` was unavailable on every node).*\n" + ) + else: + f.write( + "Foreign PIDs found holding GPUs at smoke start. The most " + "common cause is leaked Python ranks from a previous " + "training job (look for `python` / `torchrun` / `train.py`). " + "Clean up with `pkill -9 -f train.py` (or similar) on the " + "listed nodes BEFORE launching the next job.\n\n" + ) + f.write("| Node | Hostname | GPU | PID | Process | HBM held (GiB) |\n") + f.write("|------|----------|-----|-----|---------|----------------|\n") + for r in busy_rows: + name = str(r.get("name", "")).replace("|", "/")[:40] + hbm = r.get("hbm_gib") + hbm_s = f"{hbm}" if hbm is not None else "?" + f.write( + f"| {r['node_rank']} | {r['host']} | " f"{r['gpu']} | {r['pid']} | `{name}` | {hbm_s} |\n" + ) + except Exception as e: + f.write(f"*Busy-GPU section failed to render: {e}*\n") + _warn(f"busy-gpu render failed: {e}") + + +def _write_pretouch_hbm( + f: IO[str], + nodes: List[Dict[str, Any]], + hbm_busy_threshold_gib: float, +) -> None: + # ----- G: Pre-touch HBM-used outliers ----- + f.write("\n## GPU pre-touch HBM usage outliers\n\n") + try: + threshold = float(hbm_busy_threshold_gib) + pt_rows = _pretouch_hbm_rows(nodes, threshold_gib=threshold) + if not pt_rows: + f.write( + f"*No GPU reached the pre-touch HBM threshold " + f"({threshold} GiB or more) -- every GPU started clean.*\n" + ) + else: + f.write( + f"GPUs with **at least {threshold} GiB** of HBM already " + f"in use BEFORE smoke touched the device. This number is " + "not polluted by our own caching allocator (it's measured " + "before any allocation), so it directly reflects foreign " + "or leaked occupancy.\n\n" + ) + f.write("| Node | Hostname | GPU | HBM used pre-touch (GiB) |\n") + f.write("|------|----------|-----|---------------------------|\n") + for r in pt_rows: + f.write(f"| {r['node_rank']} | {r['host']} | " f"{r['gpu']} | {r['used_gib']} |\n") + except Exception as e: + f.write(f"*Pre-touch HBM section failed to render: {e}*\n") + _warn(f"pretouch-hbm render failed: {e}") + + +def _write_gpu_activity( + f: IO[str], + nodes: List[Dict[str, Any]], + gpu_activity_warn_pct: float, +) -> None: + # ----- G: GPU compute activity outliers ----- + f.write("\n## GPU compute-activity outliers\n\n") + try: + warn_pct = float(gpu_activity_warn_pct) + act_rows = _gpu_activity_rows(nodes, warn_pct=warn_pct) + if not act_rows: + f.write( + f"*No GPU exceeded `gfx_activity_pct >= {warn_pct}%` at " + "smoke start (or amd-smi did not report activity).*\n" + ) + else: + f.write( + f"GPUs reporting **>= {warn_pct}%** compute activity at " + "smoke start. Short bursts are normal; sustained " + "non-trivial activity across multiple GPUs strongly " + "suggests a leaked rank still running compute. Warn-only " + "(does not by itself fail the node).\n\n" + ) + f.write("| Node | Hostname | GPU | Activity % |\n") + f.write("|------|----------|-----|------------|\n") + for r in act_rows: + f.write(f"| {r['node_rank']} | {r['host']} | " f"{r['gpu']} | {r['activity_pct']} |\n") + except Exception as e: + f.write(f"*Activity section failed to render: {e}*\n") + _warn(f"activity render failed: {e}") + + +def _write_tier2_perf_summary(f: IO[str], nodes: List[Dict[str, Any]]) -> None: + """Write the Tier 2 perf summary section. + + INTENTIONALLY UN-WRAPPED in try/except (matching the original + monolithic block): a regression in this loop should bubble up rather + than be silently swallowed. Only emitted when at least one node + actually ran Tier 2. + """ + # Tier 2 perf summary -- only emitted when at least one node ran Tier 2. + # Surfaces per-node GEMM TFLOPS / HBM GB/s (min/median/max across the + # node's GPUs) plus the local RCCL all-reduce GB/s, so outliers across + # the cluster are visible without opening every per-node JSON. + perf_rows: List[str] = [] + any_tier2 = False + for n in nodes: + t2 = n.get("tier2") or {} + per_gpu = (n.get("tier1") or {}).get("per_gpu") or [] + gemm = [ + p.get("details", {}).get("gemm_tflops") + for p in per_gpu + if isinstance(p.get("details", {}).get("gemm_tflops"), (int, float)) + ] + hbm = [ + p.get("details", {}).get("hbm_gbs") + for p in per_gpu + if isinstance(p.get("details", {}).get("hbm_gbs"), (int, float)) + ] + rccl_gbs = (t2.get("rccl") or {}).get("gbs") + if not gemm and not hbm and rccl_gbs is None: + perf_rows.append(f"| {n.get('node_rank', '?')} | {n.get('host', '?')} | | | |") + continue + any_tier2 = True + + def _fmt_stats(xs): + if not xs: + return "" + xs_sorted = sorted(xs) + med = xs_sorted[len(xs_sorted) // 2] + return f"{min(xs):.1f} / {med:.1f} / {max(xs):.1f}" + + perf_rows.append( + f"| {n.get('node_rank', '?')} | {n.get('host', '?')} | " + f"{_fmt_stats(gemm)} | {_fmt_stats(hbm)} | " + f"{rccl_gbs if rccl_gbs is not None else ''} |" + ) + + if any_tier2: + f.write("\n## Tier 2 perf summary\n\n") + f.write( + "Per-node GEMM TFLOPS (8192^3 bf16) and HBM GB/s shown as " + "`min / median / max` across the node's GPUs. RCCL GB/s is the " + "node-local 8-GPU all-reduce algorithmic bandwidth at 64 MB.\n\n" + ) + f.write( + "| Node | Hostname | GEMM TFLOPS (min/med/max) | " "HBM GB/s (min/med/max) | Local RCCL GB/s |\n" + ) + f.write( + "|------|----------|----------------------------|" + "------------------------|------------------|\n" + ) + for r in perf_rows: + f.write(r + "\n") + + +def _write_failing_reasons( + f: IO[str], + failing: List[Dict[str, Any]], +) -> None: + """Write the per-node fail-reason dump for failing nodes. + + INTENTIONALLY UN-WRAPPED in try/except (matching the original + monolithic block): if iterating fail_reasons raises, surface it + instead of swallowing it. + """ + if failing: + f.write("\n## Failing nodes -- full reasons\n\n") + for n in failing: + f.write(f"### {n.get('host', '?')}\n\n") + for r in n.get("fail_reasons") or []: + f.write(f"- {r}\n") + f.write("\n") + + +def write_smoke_report( + report_path: str, + *, + nodes: List[Dict[str, Any]], + passing: List[Dict[str, Any]], + failing: List[Dict[str, Any]], + expected: Any, + clock_skew_warn_sec: float, + rocm_smi_warn_sec: float, + hbm_busy_threshold_gib: float, + gpu_activity_warn_pct: float, +) -> None: + """Write the cluster smoke report to ``report_path``. + + Section ORDER and HEADINGS are part of the operator-facing contract + -- many CI/CD pipelines and slack bots scrape ``smoke_report.md`` for + specific ``##`` headings. Do NOT reorder or rename without a deliberate + behavior change. The order matches the original monolithic writer + (header, A, A.2, B, B.2, B.3, C, GPU visibility, D-1, D-2, E, + F-partial, Tooling availability, G x3, Tier 2 perf summary + [conditional], Failing nodes [conditional]). + """ + with open(report_path, "w", encoding="utf-8") as f: + _write_header(f, nodes, passing, failing, expected) + _write_stack_drift(f, nodes) + _write_nic_fw_drift(f, nodes) + _write_nic_issues(f, nodes) + _write_nic_port_count(f, nodes) + _write_nic_excluded(f, nodes) + _write_host_limits(f, nodes) + _write_gpu_visibility(f, nodes) + _write_gpu_low_level(f, nodes) + _write_xgmi(f, nodes) + _write_clock(f, nodes, skew_warn_sec=clock_skew_warn_sec) + _write_tooling_latency(f, nodes, rocm_smi_warn_sec=rocm_smi_warn_sec) + _write_tooling_availability(f, nodes) + _write_busy_gpus(f, nodes) + _write_pretouch_hbm( + f, + nodes, + hbm_busy_threshold_gib=hbm_busy_threshold_gib, + ) + _write_gpu_activity( + f, + nodes, + gpu_activity_warn_pct=gpu_activity_warn_pct, + ) + _write_tier2_perf_summary(f, nodes) + _write_failing_reasons(f, failing) diff --git a/primus/tools/preflight/node_smoke/aggregator/summarizers.py b/primus/tools/preflight/node_smoke/aggregator/summarizers.py new file mode 100644 index 000000000..fcfb8b92d --- /dev/null +++ b/primus/tools/preflight/node_smoke/aggregator/summarizers.py @@ -0,0 +1,449 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Pure data-shaping helpers used by the cluster-level report writer. + +Every function takes the loaded per-node JSON list (``nodes``) and +returns plain rows or summary dicts. None of them format markdown -- the +report writer in :mod:`.report` is solely responsible for layout. +Each helper is best-effort: missing data degrades to an empty list rather +than raising. +""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from ..collectors.tooling import _TRACKED_TOOLS + +# --------------------------------------------------------------------------- +# Aggregator helpers -- A. stack/NIC drift, B. NIC issues, C. host limits +# --------------------------------------------------------------------------- + + +def _stack_drift_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """For every *scalar* fingerprint key, find the cluster-majority value and + list the nodes that disagree. + + Returns one row per key that has at least one outlier. Keys missing from + every node, or where every node reported the same value, are omitted so + a healthy cluster produces an empty list. + """ + from collections import Counter + + # Only collect keys that at least ONE node reported as a scalar. We + # ignore None here so a key that happens to be None on one node and a + # dict on another (e.g. nic_fw on a node without an IB stack) doesn't + # leak into the scalar-drift loop and crash Counter() with an unhashable + # value. + keys: set = set() + for n in nodes: + fp = ((n.get("tier1") or {}).get("fingerprint") or {}) or {} + for k, v in fp.items(): + if isinstance(v, (str, int, float)): + keys.add(k) + + rows: List[Dict[str, Any]] = [] + for k in sorted(keys): + per_host: List[tuple] = [] + for n in nodes: + fp = ((n.get("tier1") or {}).get("fingerprint") or {}) or {} + v = fp.get(k) + # Defense in depth: skip non-scalar values per-host too, in case + # different nodes disagree on the type for the same key. + if not isinstance(v, (str, int, float)): + continue + per_host.append((n.get("host", "?"), v)) + if not per_host: + continue + c = Counter(v for _, v in per_host) + majority, count = c.most_common(1)[0] + outliers = [(h, v) for h, v in per_host if v != majority] + if not outliers: + continue + rows.append( + { + "key": k, + "majority": majority, + "count": count, + "total": len(per_host), + "outliers": outliers, + } + ) + return rows + + +def _nic_fw_drift_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Per-IB-device firmware drift across the cluster (e.g. rdma0 mismatch).""" + from collections import Counter + + all_devs: set = set() + for n in nodes: + fp = ((n.get("tier1") or {}).get("fingerprint") or {}) or {} + all_devs.update((fp.get("nic_fw") or {}).keys()) + + rows: List[Dict[str, Any]] = [] + for dev in sorted(all_devs): + per_host: List[tuple] = [] + for n in nodes: + fp = ((n.get("tier1") or {}).get("fingerprint") or {}) or {} + v = (fp.get("nic_fw") or {}).get(dev) + if v is None: + continue + per_host.append((n.get("host", "?"), v)) + if not per_host: + continue + c = Counter(v for _, v in per_host) + majority, count = c.most_common(1)[0] + outliers = [(h, v) for h, v in per_host if v != majority] + if not outliers: + continue + rows.append( + { + "device": dev, + "majority": majority, + "count": count, + "total": len(per_host), + "outliers": outliers, + } + ) + return rows + + +def _nic_issue_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Per-node NIC roll-call issues (port DOWN / no GIDs / count mismatch).""" + rows: List[Dict[str, Any]] = [] + for n in nodes: + nic = (n.get("tier1") or {}).get("nics") or {} + for issue in nic.get("issues", []) or []: + rows.append( + { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + "issue": issue, + } + ) + return rows + + +def _nic_excluded_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Per-node NIC ports that were excluded from the training-NIC set. + + Excluded ports are informational: they did not contribute to the + node FAIL signal, but operators sometimes want to see them + (e.g. "did the heuristic do what I expected?", "is my front-end + NIC still disabled or did it come back?"). Each row carries the + selector source so the operator can tell whether the exclusion + came from --rdma-nic-allowlist, NCCL_IB_HCA, or the + phys_state=Disabled/Sleep heuristic. + """ + rows: List[Dict[str, Any]] = [] + for n in nodes: + nic = (n.get("tier1") or {}).get("nics") or {} + sel = nic.get("selector") or {} + source = sel.get("source") or "?" + for issue in nic.get("info_issues", []) or []: + rows.append( + { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + "source": source, + "issue": issue, + } + ) + return rows + + +def _host_limits_issue_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Per-node host-limit hard violations (ulimit -l / /dev/shm too low).""" + rows: List[Dict[str, Any]] = [] + for n in nodes: + hl = (n.get("tier1") or {}).get("host_limits") or {} + for issue in hl.get("fail_reasons", []) or []: + rows.append( + { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + "issue": issue, + } + ) + return rows + + +# --------------------------------------------------------------------------- +# Aggregator helpers -- D-1 / D-2 / E / F +# --------------------------------------------------------------------------- + + +def _gpu_low_level_outlier_rows( + nodes: List[Dict[str, Any]], +) -> List[Dict[str, Any]]: + """Find per-GPU outliers in PCIe link + HBM total across the cluster. + + For each scalar metric the cluster has a strong majority value (e.g. + 16 lanes, 32 GT/s, 191 GiB HBM). A single GPU below the majority on + any of these is almost always a hardware issue -- a cold-soldered + socket, a degraded PCIe link, or HBM that the firmware refused to + bring online. We surface every such (host, gpu, metric, value) + tuple, with the cluster majority for context. + + Power cap and ECC counters from amd-smi are intentionally NOT included + here; they have their own narrower checks (ECC = hard fail in + ``_node_status_from``; power cap = informational only because cluster + operators sometimes set per-rack caps deliberately). + """ + from collections import Counter + + fields = ( + ("pcie_link_width", "PCIe width (lanes)"), + ("pcie_link_speed_gts", "PCIe speed (GT/s)"), + ("hbm_total_gib", "HBM total (GiB)"), + ) + rows: List[Dict[str, Any]] = [] + for key, label in fields: + per_gpu: List[tuple] = [] # (host, gpu_idx, value) + for n in nodes: + for p in (n.get("tier1") or {}).get("per_gpu") or []: + low = (p.get("details") or {}).get("low_level") or {} + v = low.get(key) + if isinstance(v, (int, float)): + per_gpu.append((n.get("host", "?"), p.get("gpu", "?"), v)) + if not per_gpu: + continue + c = Counter(v for _, _, v in per_gpu) + majority, count = c.most_common(1)[0] + outliers = [(h, g, v) for h, g, v in per_gpu if v != majority] + if not outliers: + continue + rows.append( + { + "key": key, + "label": label, + "majority": majority, + "count": count, + "total": len(per_gpu), + "outliers": outliers, + } + ) + return rows + + +def _xgmi_issue_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Per-node XGMI link issues (any non-XGMI GPU pair).""" + rows: List[Dict[str, Any]] = [] + for n in nodes: + xg = (n.get("tier1") or {}).get("xgmi") or {} + if not xg.get("ok"): + err = xg.get("error") + if err: + rows.append( + { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + "summary": f"could not collect topology: {err}", + } + ) + continue + bad = xg.get("non_xgmi_pairs") or [] + if not bad: + continue + # Show up to 6 sample pairs to keep the table readable; the full + # matrix lives in the per-node JSON. + sample = ", ".join(f"({i},{j})={t}" for i, j, t in bad[:6]) + suffix = "" if len(bad) <= 6 else f" (+{len(bad) - 6} more)" + rows.append( + { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + "summary": f"{len(bad)} non-XGMI pair(s): {sample}{suffix}", + } + ) + return rows + + +def _clock_summary( + nodes: List[Dict[str, Any]], + skew_warn_sec: float, +) -> Dict[str, Any]: + """Compute wall-clock spread + per-node time-daemon health.""" + times: List[tuple] = [] # (host, wall_time_unix) + no_daemon_hosts: List[tuple] = [] # (node_rank, host) + for n in nodes: + clk = (n.get("tier1") or {}).get("clock") or {} + wt = clk.get("wall_time_unix") + if isinstance(wt, (int, float)): + times.append((n.get("host", "?"), float(wt))) + if clk and not clk.get("any_active", True): + no_daemon_hosts.append((n.get("node_rank", "?"), n.get("host", "?"))) + + spread_sec = None + earliest_h = latest_h = None + if len(times) >= 2: + earliest_h, earliest = min(times, key=lambda x: x[1]) + latest_h, latest = max(times, key=lambda x: x[1]) + spread_sec = round(latest - earliest, 3) + return { + "n_nodes_with_time": len(times), + "spread_sec": spread_sec, + "spread_warn_sec": skew_warn_sec, + "spread_warn": (spread_sec is not None and spread_sec > skew_warn_sec), + "earliest_host": earliest_h, + "latest_host": latest_h, + "no_daemon_hosts": no_daemon_hosts, + } + + +def _tooling_latency_rows( + nodes: List[Dict[str, Any]], + warn_sec: float, +) -> List[Dict[str, Any]]: + """Per-node `rocm-smi --version` self-latency outliers (timed-out + slow).""" + rows: List[Dict[str, Any]] = [] + for n in nodes: + t = (n.get("tier1") or {}).get("tooling") or {} + lat = t.get("latency_sec") + timed_out = bool(t.get("timed_out")) + if t.get("error") and lat is None: + # Tool missing -- not interesting for a slow-tool report. + continue + flag = "" + if timed_out: + flag = "TIMEOUT" + elif isinstance(lat, (int, float)) and lat > warn_sec: + flag = f">{warn_sec}s" + if not flag: + continue + rows.append( + { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + "latency_sec": lat, + "flag": flag, + "timeout_sec": t.get("timeout_sec"), + } + ) + return rows + + +def _tooling_inventory_rows(nodes: List[Dict[str, Any]]) -> Dict[str, Any]: + """Per-node tooling-presence rows + a count of nodes missing each tool. + + Returns:: + + { + "rows": [ + {"node_rank": 0, "host": "tus1-p3-g25", + "amd-smi": True, "rocm-smi": True, "lsof": True}, + ... + ], + "missing_counts": {"amd-smi": 0, "rocm-smi": 1, "lsof": 0}, + "any_missing": True, + "tracked": ["amd-smi", "rocm-smi", "lsof"], + } + + Always-on: even when every tool is present everywhere, we still emit + a (small, reassuring) summary so the operator can see at a glance + that the toolchain was healthy on every node. + """ + tracked = list(_TRACKED_TOOLS) + rows: List[Dict[str, Any]] = [] + missing_counts: Dict[str, int] = {t: 0 for t in tracked} + for n in nodes: + inv = (n.get("tier1") or {}).get("tooling_inventory") or {} + tools = inv.get("tools") or {} + row: Dict[str, Any] = { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + } + for t in tracked: + present = bool((tools.get(t) or {}).get("present")) + row[t] = present + if not present: + missing_counts[t] += 1 + rows.append(row) + return { + "rows": rows, + "missing_counts": missing_counts, + "any_missing": any(c > 0 for c in missing_counts.values()), + "tracked": tracked, + } + + +def _busy_gpu_rows(nodes: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Per-node foreign GPU process listing (for the "Busy GPUs" section).""" + rows: List[Dict[str, Any]] = [] + for n in nodes: + gp = (n.get("tier1") or {}).get("gpu_processes") or {} + if not gp.get("ok"): + continue + per_gpu = gp.get("per_gpu") or [] + for g in per_gpu: + for p in g.get("processes") or []: + if not p.get("is_foreign"): + continue + hbm_b = p.get("hbm_bytes") + rows.append( + { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + "gpu": g.get("gpu", "?"), + "pid": p.get("pid"), + "name": p.get("name") or "", + "hbm_gib": ( + round(hbm_b / (1 << 30), 2) if isinstance(hbm_b, int) and hbm_b > 0 else None + ), + } + ) + return rows + + +def _pretouch_hbm_rows( + nodes: List[Dict[str, Any]], + threshold_gib: float, +) -> List[Dict[str, Any]]: + """Per-GPU pre-touch HBM-used outliers (above threshold).""" + rows: List[Dict[str, Any]] = [] + for n in nodes: + per_gpu = (n.get("tier1") or {}).get("per_gpu") or [] + for p in per_gpu: + d = p.get("details") or {} + used_gib = d.get("hbm_pre_touch_used_gib") + if not isinstance(used_gib, (int, float)): + continue + if used_gib >= threshold_gib: + rows.append( + { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + "gpu": p.get("gpu", "?"), + "used_gib": round(float(used_gib), 2), + } + ) + return rows + + +def _gpu_activity_rows( + nodes: List[Dict[str, Any]], + warn_pct: float, +) -> List[Dict[str, Any]]: + """Per-GPU compute-activity outliers (above warn threshold).""" + rows: List[Dict[str, Any]] = [] + for n in nodes: + amd = (n.get("tier1") or {}).get("gpu_low_level") or {} + for rec in amd.get("per_gpu") or []: + pct = rec.get("gfx_activity_pct") + if not isinstance(pct, (int, float)): + continue + if float(pct) >= warn_pct: + rows.append( + { + "node_rank": n.get("node_rank", "?"), + "host": n.get("host", "?"), + "gpu": rec.get("gpu", "?"), + "activity_pct": round(float(pct), 1), + } + ) + return rows diff --git a/primus/tools/preflight/node_smoke/cli.py b/primus/tools/preflight/node_smoke/cli.py new file mode 100644 index 000000000..2462dc55d --- /dev/null +++ b/primus/tools/preflight/node_smoke/cli.py @@ -0,0 +1,910 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""CLI wiring + the three subcommand entry points (`run`, `aggregate`, +`_per_gpu`). + +The argparse layout is part of the operator-facing contract: flag names, +defaults, help strings, and abbreviation behaviour are preserved exactly. +``allow_abbrev=False`` on the ``run`` subparser keeps an old ``--tier2`` +flag from a stale script silently matching the new ``--tier2-perf``. +""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from dataclasses import asdict +from typing import Any, Dict, List, Optional + +from .aggregator.report import write_smoke_report +from .collectors.clock import _collect_clock_state +from .collectors.dmesg import _collect_dmesg_errors +from .collectors.fingerprint import _collect_node_fingerprint +from .collectors.gpu_low_level import _collect_amd_smi_metrics +from .collectors.gpu_processes import _collect_gpu_processes +from .collectors.host_limits import _collect_host_limits +from .collectors.nics import _collect_nic_status +from .collectors.reused_info import _collect_reused_info +from .collectors.rocm_smi import _collect_rocm_smi_self_latency +from .collectors.tooling import _collect_tooling_inventory +from .collectors.xgmi import _collect_xgmi_topology +from .logging_utils import _log, _short_name, _this_host_short, _warn +from .orchestrator import _clean_dump_path, _node_status_from, _spawn_per_gpu +from .per_gpu import _per_gpu_body +from .rccl_local import _run_local_rccl +from .types import GPUResult, NodeResult + + +def _cmd_per_gpu(ns: argparse.Namespace) -> int: + """Internal subcommand: run all per-GPU checks for a single GPU index.""" + result = _per_gpu_body( + gpu=int(ns.gpu), + tier2_perf=bool(ns.tier2_perf), + gemm_tflops_min=float(ns.gemm_tflops_min), + hbm_gbs_min=float(ns.hbm_gbs_min), + hbm_busy_threshold_bytes=int(float(ns.hbm_busy_threshold_gib) * (1 << 30)), + ) + # Single JSON line on stdout; nothing else. + print(json.dumps(result), flush=True) + return 0 if result.get("status") == "PASS" else 1 + + +def _cmd_run(ns: argparse.Namespace) -> int: + """Per-node entry: orchestrate Tier 1 + optional Tier 2 + write JSON.""" + # Always store the short hostname so consumers (passing_nodes.txt / + # failing_nodes.txt and SLURM tools that read them) get a name they + # can use directly. + host = _this_host_short() + node_rank = int(os.environ.get("NODE_RANK", os.environ.get("SLURM_NODEID", "0"))) + + # Rank-0 only: wipe stale JSONs / aggregator outputs from a previous + # run so re-runs on a different (smaller) nodelist cannot inherit + # ghost PASS verdicts from removed nodes. This MUST run before any + # rank's _spawn_per_gpu loop completes (and thus before any rank + # writes its JSON) -- which is why it lives at the top of _cmd_run + # on rank 0 only, not in the wrapper. + if node_rank == 0 and not ns.no_clean_dump_path: + removed = _clean_dump_path(ns.dump_path) + if removed: + _log( + f"cleaned {len(removed)} stale file(s) from {ns.dump_path} " + f"(per-node JSONs + aggregator outputs)" + ) + + expected_gpus = ns.expected_gpus + if expected_gpus is None: + expected_gpus = int(os.environ.get("LOCAL_WORLD_SIZE", os.environ.get("GPUS_PER_NODE", "0")) or 0) + if expected_gpus <= 0: + try: + import torch # type: ignore + + expected_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0 + except Exception: + expected_gpus = 0 + expected_gpus = max(0, int(expected_gpus)) + + # GPU visibility guard. We capture each independent source (the + # --expected-gpus flag, env vars, torch, and -- below -- amd-smi) so + # the JSON tells the operator *why* we resolved to N. The hard-fail + # rules live here, decoupled from any other collector, because we have + # seen `_collect_reused_info()` downgrade the "No GPUs detected" fail + # to a warn when collect_gpu_info() raises -- which would otherwise + # let a CPU-only or stale-GPU node PASS smoke silently. + torch_visible = 0 + torch_is_available = False + try: + import torch # type: ignore + + torch_is_available = bool(torch.cuda.is_available()) + torch_visible = int(torch.cuda.device_count()) if torch_is_available else 0 + except Exception: + pass + gpu_visibility: Dict[str, Any] = { + "expected_gpus": expected_gpus, + "explicit_expected_gpus": ns.expected_gpus, + "torch_visible": torch_visible, + "torch_is_available": torch_is_available, + "env_local_world_size": int(os.environ.get("LOCAL_WORLD_SIZE", "0") or 0), + "env_gpus_per_node": int(os.environ.get("GPUS_PER_NODE", "0") or 0), + "amd_smi_visible": None, # filled in after _collect_amd_smi_metrics + "fail_reasons": [], + } + if expected_gpus < 1: + msg = ( + f"expected_gpus={expected_gpus}: no per-GPU sanity tests will " + f"run (torch_is_available={torch_is_available}, " + f"torch_visible={torch_visible}, " + f"LOCAL_WORLD_SIZE={gpu_visibility['env_local_world_size']}, " + f"GPUS_PER_NODE={gpu_visibility['env_gpus_per_node']})" + ) + gpu_visibility["fail_reasons"].append(msg) + _warn(msg) + + _log( + f"start node-smoke: node_rank={node_rank} expected_gpus={expected_gpus} " + f"tier2_perf={ns.tier2_perf}" + ) + if ns.tier2_perf and expected_gpus < 2: + # Tier 2 also includes a node-local RCCL all-reduce, which needs at + # least 2 GPUs. Surface the skip up front instead of silently doing + # only GEMM/HBM and giving the operator false coverage confidence. + _warn( + f"--tier2-perf requested but expected_gpus={expected_gpus} < 2; " + "the node-local RCCL all-reduce phase will be skipped. " + "Per-GPU GEMM and HBM checks will still run." + ) + + # G: enumerate processes currently holding each GPU BEFORE we spawn + # any per-GPU subprocess. Anything we see here is, by definition, not + # us -- it's a leaked rank from a previous job, a foreign tenant, or + # an in-band monitoring agent. The aggregator + _node_status_from + # turn this into a hard FAIL unless the operator opted out via + # --allow-foreign-procs (or whitelisted the agent name). + allowed_proc_names = [s for s in (getattr(ns, "allowed_procs", "") or "").split(",") if s.strip()] + tier1_extra_pre: Dict[str, Any] = {} + tier1_extra_pre["gpu_processes"] = _collect_gpu_processes( + self_pid=os.getpid(), + allowed_proc_names=allowed_proc_names, + ) + gp = tier1_extra_pre["gpu_processes"] + if gp.get("ok"): + _log( + f"gpu_processes ({gp.get('tool')}): " + f"{gp.get('foreign_count', 0)} foreign PID(s) across " + f"{len(gp.get('per_gpu') or [])} GPU bucket(s)" + ) + if gp.get("foreign_count", 0) > 0: + for g in gp.get("per_gpu") or []: + for p in g.get("processes") or []: + if p.get("is_foreign"): + hbm = p.get("hbm_bytes") + hbm_s = ( + f"{round(hbm / (1 << 30), 2)} GiB" if isinstance(hbm, int) and hbm > 0 else "?" + ) + _warn( + f"foreign process on gpu{g.get('gpu')}: " + f"pid={p.get('pid')} name={p.get('name')!r} " + f"hbm={hbm_s}" + ) + else: + _warn( + f"gpu_processes: enumeration unavailable " + f"({gp.get('error') or gp.get('json_error') or gp.get('text_error') or '?'})" + ) + + t0 = time.time() + per_gpu: List[GPUResult] = [] + for i in range(expected_gpus): + r = _spawn_per_gpu( + i, + timeout_sec=ns.per_gpu_timeout_sec, + tier2_perf=bool(ns.tier2_perf), + gemm_tflops_min=ns.gemm_tflops_min, + hbm_gbs_min=ns.hbm_gbs_min, + hbm_busy_threshold_gib=float(ns.hbm_busy_threshold_gib), + ) + per_gpu.append(r) + _log( + f"gpu{i}: {r.status} ({r.duration_sec:.1f}s)" + + (f" -- {r.reason}" if r.reason else "") + + (f" -- {r.details}" if r.details else "") + ) + + # Tier 1 reused info collectors + tier1_extra: Dict[str, Any] = {} + tier1_extra.update(tier1_extra_pre) # carry forward gpu_processes + tier1_extra.update(_collect_reused_info()) + if not ns.skip_dmesg: + tier1_extra["dmesg"] = _collect_dmesg_errors(window_minutes=ns.dmesg_minutes) + else: + tier1_extra["dmesg"] = {"ok": True, "matches": [], "error": "skipped"} + + # A/B/C: software-stack fingerprint, NIC roll-call, host limits. + # All three are pure data-collection (millisecond-scale sysfs reads); the + # heavy cluster-level drift detection happens at aggregation time. + tier1_extra["fingerprint"] = _collect_node_fingerprint() + tier1_extra["nics"] = _collect_nic_status( + expected_count=ns.expected_rdma_nics, + allowlist=getattr(ns, "rdma_nic_allowlist", None), + ) + tier1_extra["host_limits"] = _collect_host_limits( + ulimit_l_min_gb=ns.ulimit_l_min_gb, + shm_min_gb=ns.shm_min_gb, + ) + + # D-1 heavy: per-GPU ECC / throttle / clocks / power via amd-smi (one + # node-level call, results indexed by gpu). + # D-2: XGMI link matrix via amd-smi topology (one node-level call). + # E: wall-time + time-daemon active states. + # F-partial: rocm-smi --version self-latency with a hard timeout to + # catch drivers that are starting to wedge. + # Tooling inventory FIRST so we can warn loudly before running the + # collectors that depend on each tool. Several downstream collectors + # (gpu_low_level, xgmi, gpu_processes, tooling) silently no-op when + # their tool is missing, which can let a broken node pass smoke + # unnoticed -- this section is the operator-visible counterweight. + tier1_extra["tooling_inventory"] = _collect_tooling_inventory() + inv = tier1_extra["tooling_inventory"] + inv_missing = inv["missing"] + inv_uncovered = inv["uncovered"] + if inv_missing: + # First line: which tools are missing. Loud regardless of whether + # a fallback covers everything. + _warn(f"tooling: {len(inv_missing)} tracked tool(s) NOT in PATH: " f"{', '.join(inv_missing)}.") + if inv_uncovered: + # Second line: which checks have NO working tool at all. This + # is the actually-dangerous case (a check that will silently + # no-op no matter which tool we try). + _warn( + f"tooling: {len(inv_uncovered)} check(s) have NO working " + f"tool and will be silently skipped: " + + ", ".join(inv_uncovered) + + ". Use --require-tools to promote missing tools to a " + "node FAIL." + ) + else: + # Reassuring line: every check is covered via fallback. + _warn( + "tooling: every check is still covered via fallback " + "(rocm-smi / lsof) -- no checks will be silently skipped. " + "Use --require-tools to promote missing tools to a node " + "FAIL anyway if your environment requires them." + ) + + tier1_extra["gpu_low_level"] = _collect_amd_smi_metrics() + tier1_extra["xgmi"] = _collect_xgmi_topology() + tier1_extra["clock"] = _collect_clock_state() + tier1_extra["tooling"] = _collect_rocm_smi_self_latency(timeout_sec=float(ns.rocm_smi_timeout_sec)) + + # Visibility cross-check: if amd-smi successfully enumerated GPUs but + # torch couldn't see them, that's a high-signal sign of a stale ROCm + # install / wedged amdgpu driver -- exactly the case where a "smoke + # test" is supposed to pull the node out of rotation. We only treat + # the JSON path as authoritative for counting (the text fallback + # cannot be reliably parsed for a count). + amd_low = tier1_extra["gpu_low_level"] + if amd_low.get("ok") and amd_low.get("tool") == "amd-smi metric --json": + per = amd_low.get("per_gpu") or [] + n_amd = len(per) if isinstance(per, list) else 0 + gpu_visibility["amd_smi_visible"] = n_amd + if n_amd > 0 and torch_visible < n_amd: + mismatch = ( + f"gpu_visibility_mismatch: amd-smi sees {n_amd} GPU(s) " + f"but torch.cuda.device_count()={torch_visible} " + f"(torch_is_available={torch_is_available}); ROCm install " + f"or amdgpu driver may be broken on this node" + ) + gpu_visibility["fail_reasons"].append(mismatch) + _warn(mismatch) + tier1_extra["gpu_visibility"] = gpu_visibility + xg = tier1_extra["xgmi"] + if xg.get("ok"): + bad = xg.get("non_xgmi_pairs") or [] + _log(f"xgmi: {xg.get('n_gpus', 0)}x{xg.get('n_gpus', 0)} matrix, " f"{len(bad)} non-XGMI pair(s)") + elif xg.get("error"): + _warn(f"xgmi: {xg.get('error')}") + tool = tier1_extra["tooling"] + if tool.get("ok"): + _log(f"rocm-smi --version: {tool.get('latency_sec')}s") + elif tool.get("timed_out"): + _warn(f"rocm-smi --version timed out after {tool.get('timeout_sec')}s " "-- driver may be wedging") + elif tool.get("error"): + _warn(f"tooling: {tool.get('error')}") + nic_summary = tier1_extra["nics"] + _log( + f"nics: {len(nic_summary.get('ports', []))} port(s) found, " + f"{len(nic_summary.get('issues', []))} issue(s)" + ) + if tier1_extra["host_limits"].get("fail_reasons"): + for r in tier1_extra["host_limits"]["fail_reasons"]: + _warn(f"host_limits: {r}") + + # Tier 2 local RCCL all-reduce. Gated on a single flag now (--tier2-perf) + # so users cannot accidentally end up with only the per-GPU half running. + tier2_extra: Dict[str, Any] = {} + if ns.tier2_perf and expected_gpus > 1: + _log(f"tier2 local RCCL all-reduce: {expected_gpus} ranks, {ns.rccl_size_mb}MB") + rccl = _run_local_rccl( + local_world_size=expected_gpus, + size_mb=ns.rccl_size_mb, + timeout_sec=ns.rccl_timeout_sec, + ) + if rccl.get("status") == "PASS" and rccl.get("gbs") is not None: + if float(rccl["gbs"]) < ns.rccl_gbs_min: + rccl = { + "status": "FAIL", + "gbs": rccl["gbs"], + "error": (f"local RCCL {rccl['gbs']} GB/s < threshold {ns.rccl_gbs_min}"), + } + tier2_extra["rccl"] = rccl + _log(f"tier2 RCCL: {rccl}") + + required_tools = [s.strip() for s in (getattr(ns, "require_tools", "") or "").split(",") if s.strip()] + fail_reasons = _node_status_from( + per_gpu, + tier1_extra, + tier2_extra, + allow_foreign_procs=bool(ns.allow_foreign_procs), + required_tools=required_tools, + ) + status = "PASS" if not fail_reasons else "FAIL" + + node_result = NodeResult( + host=host, + node_rank=node_rank, + status=status, + duration_sec=round(time.time() - t0, 3), + fail_reasons=fail_reasons, + tier1={ + "per_gpu": [asdict(r) for r in per_gpu], + **tier1_extra, + }, + tier2=tier2_extra, + ) + + smoke_dir = os.path.join(ns.dump_path, "smoke") + os.makedirs(smoke_dir, exist_ok=True) + out_path = os.path.join(smoke_dir, f"{host}.json") + with open(out_path, "w", encoding="utf-8") as f: + json.dump(asdict(node_result), f, indent=2, default=str) + _log(f"wrote {out_path} status={status} duration={node_result.duration_sec}s") + if fail_reasons: + for r in fail_reasons[:5]: + _warn(r) + + # Per-node `run` exits 0 whenever the smoke test ran to completion -- + # the node verdict (PASS/FAIL) is in the JSON, in failing_nodes.txt, + # and in the aggregator's exit code. Conflating "this node is broken" + # with "this tool crashed" makes srun output look like the smoke test + # itself is failing, when in fact it's correctly DOING ITS JOB of + # identifying broken nodes. The aggregator (rank 0) is the single + # source of truth for the CI-friendly cluster-health exit signal. + # Tool failures (couldn't import, couldn't write JSON, etc.) still + # propagate as non-zero via Python's default exception handling. + return 0 + + +def _cmd_aggregate(ns: argparse.Namespace) -> int: + """Read all per-node JSONs from ``/smoke/`` and emit summary outputs.""" + smoke_dir = os.path.join(ns.dump_path, "smoke") + os.makedirs(smoke_dir, exist_ok=True) + + expected = int(ns.expected_nodes) if ns.expected_nodes is not None else None + deadline = time.time() + max(0, int(ns.wait_timeout_sec)) + found_paths: List[str] = [] + while True: + found_paths = sorted(os.path.join(smoke_dir, p) for p in os.listdir(smoke_dir) if p.endswith(".json")) + if expected is None or len(found_paths) >= expected: + break + if time.time() >= deadline: + break + time.sleep(1) + + nodes: List[Dict[str, Any]] = [] + for p in found_paths: + try: + with open(p, "r", encoding="utf-8") as f: + nodes.append(json.load(f)) + except Exception as e: + nodes.append( + { + "host": os.path.basename(p).rsplit(".json", 1)[0], + "status": "FAIL", + "fail_reasons": [f"failed to parse {p}: {e}"], + "duration_sec": 0, + "node_rank": -1, + } + ) + + # Normalize every loaded ``host`` to its short form so legacy JSON files + # that hold an FQDN (older runs of node_smoke) still produce SLURM-ready + # passing/failing lists. + for n in nodes: + n["host"] = _short_name(str(n.get("host", ""))) + + # Optional: an explicit expected hostname list (one per line). When + # provided, we name missing nodes by their real short hostname instead + # of synthetic ```` placeholders, so the failing nodes list + # is directly usable with ``srun --exclude=``. + expected_hosts_short: List[str] = [] + nodelist_file = getattr(ns, "expected_nodelist_file", None) + if nodelist_file: + try: + with open(nodelist_file, "r", encoding="utf-8") as f: + expected_hosts_short = [_short_name(line.strip()) for line in f if line.strip()] + _log(f"loaded {len(expected_hosts_short)} expected hostnames from " f"{nodelist_file}") + except Exception as e: + _warn(f"failed to read --expected-nodelist-file {nodelist_file}: {e}") + + seen_hosts_short = {n.get("host", "") for n in nodes} + + if expected_hosts_short: + # An explicit list always wins over --expected-nodes for both the + # count and (more importantly) the identity of missing nodes. + if expected is None or expected != len(expected_hosts_short): + expected = len(expected_hosts_short) + missing_hosts = sorted(set(expected_hosts_short) - seen_hosts_short) + for h in missing_hosts: + nodes.append( + { + "host": h, + "status": "FAIL", + "fail_reasons": [ + f"no JSON received within {ns.wait_timeout_sec}s " + f"(expected hostname '{h}' from --expected-nodelist-file)" + ], + "duration_sec": 0, + "node_rank": -1, + } + ) + elif expected is not None and len(seen_hosts_short) < expected: + # Fallback: we know the count but not the identities -> emit + # synthetic placeholders. These intentionally do NOT land in + # passing/failing txt files (see _is_real_host below). + for i in range(expected - len(seen_hosts_short)): + nodes.append( + { + "host": f"", + "status": "FAIL", + "fail_reasons": [ + f"no JSON received within {ns.wait_timeout_sec}s " + f"(expected_nodes={expected}, " + f"found={len(seen_hosts_short)})" + ], + "duration_sec": 0, + "node_rank": -1, + } + ) + + # Sort by node_rank if present, otherwise by hostname. + def _key(n: Dict[str, Any]): + nr = n.get("node_rank", 0) + return ( + int(nr) if isinstance(nr, (int, str)) and str(nr).lstrip("-").isdigit() else 1 << 30, + str(n.get("host", "")), + ) + + nodes.sort(key=_key) + + passing = [n for n in nodes if n.get("status") == "PASS"] + failing = [n for n in nodes if n.get("status") != "PASS"] + + report_path = os.path.join(ns.dump_path, "smoke_report.md") + pass_path = os.path.join(ns.dump_path, "passing_nodes.txt") + fail_path = os.path.join(ns.dump_path, "failing_nodes.txt") + + write_smoke_report( + report_path, + nodes=nodes, + passing=passing, + failing=failing, + expected=expected, + clock_skew_warn_sec=float(ns.clock_skew_warn_sec), + rocm_smi_warn_sec=float(ns.rocm_smi_warn_sec), + hbm_busy_threshold_gib=float(getattr(ns, "hbm_busy_threshold_gib", 2.0)), + gpu_activity_warn_pct=float(getattr(ns, "gpu_activity_warn_pct", 20.0)), + ) + + # Only write REAL hostnames to the txt files so they can be piped directly + # into `srun --nodelist=` / `srun --exclude=`. Synthetic "" + # placeholders for nodes that never reported are surfaced in the markdown + # report instead. + def _is_real_host(h: str) -> bool: + return bool(h) and not (h.startswith("")) + + with open(pass_path, "w", encoding="utf-8") as f: + for n in passing: + h = str(n.get("host", "")) + if _is_real_host(h): + f.write(h + "\n") + with open(fail_path, "w", encoding="utf-8") as f: + for n in failing: + h = str(n.get("host", "")) + if _is_real_host(h): + f.write(h + "\n") + + _log( + f"aggregate: {len(passing)}/{len(nodes)} PASS " + f"report={report_path} passing={pass_path} failing={fail_path}" + ) + + # R5: announce absolute report paths on stdout so the operator can + # copy/paste them. Always fires (this function only runs on rank 0). + # Under bash-side --silent these prints go to /dev/null along with + # everything else, which is acceptable per plan. + _announce_aggregate_paths(ns.dump_path, report_path, pass_path, fail_path) + + return 0 if not failing and (expected is None or len(nodes) == expected) else 1 + + +def _announce_aggregate_paths(dump_path: str, report_path: str, pass_path: str, fail_path: str) -> None: + """Print absolute paths of aggregator outputs (R5). + + Replaces the bash wrapper's ``log_always Report: ...`` lines. Always + rank-0 only by construction (only ``_cmd_aggregate`` calls this, and the + primus-cli wrapper only runs aggregate on rank 0). + """ + for label, p in (("Report", report_path), ("Passing", pass_path), ("Failing", fail_path)): + try: + if os.path.isfile(p): + print(f"[Primus:NodeSmoke] {label}: {os.path.abspath(p)}", flush=True) + except OSError: + continue + + +# --------------------------------------------------------------------------- +# Argparse wiring +# +# `_add_run_flags` and `_add_aggregate_flags` are the canonical definition of +# the user-facing flag surface. Both the standalone ``run`` / ``aggregate`` +# subparsers and the primus-cli ``node_smoke`` top-level parser +# (primus/cli/subcommands/node_smoke.py) attach via these helpers, so a flag +# added once here automatically reaches every entry point. +# --------------------------------------------------------------------------- + + +def _add_run_flags(parser: argparse.ArgumentParser) -> None: + """Attach every flag that `run` accepts to ``parser``. + + Used by both the standalone ``run`` subparser and the primus-cli + top-level ``node_smoke`` parser. Adding a flag here picks it up + everywhere. + """ + parser.add_argument( + "--dump-path", default="output/preflight", help="Directory under which smoke/.json is written." + ) + parser.add_argument( + "--expected-gpus", + type=int, + default=None, + help="Expected GPU count on this node (default: LOCAL_WORLD_SIZE/GPUS_PER_NODE/torch.cuda.device_count()).", + ) + parser.add_argument( + "--per-gpu-timeout-sec", type=int, default=15, help="Hard timeout for each per-GPU subprocess." + ) + parser.add_argument( + "--tier2-perf", + action="store_true", + help="Enable Tier 2 perf sanity: per-GPU GEMM TFLOPS, " + "HBM bandwidth, AND node-local RCCL all-reduce. " + "All three are fast (< 30 s/node total).", + ) + parser.add_argument( + "--gemm-tflops-min", + type=float, + default=600.0, + help="FAIL if Tier 2 GEMM TFLOPS is below this. Default: 600 (MI300X-class).", + ) + parser.add_argument( + "--hbm-gbs-min", + type=float, + default=2000.0, + help="FAIL if Tier 2 HBM GB/s is below this. Default: 2000.", + ) + parser.add_argument( + "--rccl-size-mb", type=int, default=64, help="Tensor size for local RCCL all-reduce (MB)." + ) + parser.add_argument( + "--rccl-gbs-min", + type=float, + default=100.0, + help="FAIL if local RCCL GB/s is below this. Default: 100.", + ) + parser.add_argument( + "--rccl-timeout-sec", type=int, default=120, help="Hard timeout for the local RCCL all-reduce phase." + ) + parser.add_argument( + "--skip-dmesg", action="store_true", help="Skip the dmesg recent-error scan (e.g. inside containers)." + ) + parser.add_argument("--dmesg-minutes", type=int, default=15, help="Window for dmesg --since (minutes).") + # NIC / RDMA roll-call (B). expected_count=None means "report only"; + # set this to e.g. 8 to make a missing or down NIC port a node FAIL. + # The count is compared against the **included** (training-NIC) set + # resolved by the selector chain below, NOT the raw number of devices + # under /sys/class/infiniband (which on multi-role nodes can include + # frontend / management / storage RoCE NICs). + parser.add_argument( + "--expected-rdma-nics", + type=int, + default=None, + help="Expected RDMA training-NIC port count (compared against " + "the included set, not /sys/class/infiniband total). If " + "set, a count mismatch becomes a node FAIL.", + ) + # Training-NIC selector. Many clusters expose more RDMA-capable + # ports than the training job will actually use (frontend / + # management / storage NICs). The hard-fail rules (state must be + # ACTIVE, phys_state must be LinkUp, RoCE v2 GID present, ...) only + # run against the included subset. Excluded ports stay visible in + # the JSON (under tier1.nics.excluded_ports + info_issues) for + # diagnostics. Precedence: this flag > NCCL_IB_HCA env > heuristic + # (auto-exclude phys_state in {Disabled, Sleep}). + parser.add_argument( + "--rdma-nic-allowlist", + type=str, + default=None, + help="Comma-separated device[:port] list of training NICs " + "(NCCL_IB_HCA syntax: `^...` for denylist, `=dev` for " + "exact-match, no `:port` = match any port on device). " + "Wins over the NCCL_IB_HCA env. When neither is set, the " + "collector auto-excludes ports whose phys_state is " + "Disabled or Sleep (admin-disabled ports) and treats every " + "other port as a training NIC.", + ) + # Host-limits hard thresholds (C). Set to 0 to disable a check. + parser.add_argument( + "--ulimit-l-min-gb", + type=float, + default=32.0, + help="FAIL the node if RLIMIT_MEMLOCK is finite and below " + "this many GiB (RDMA pin will fail). 0 disables.", + ) + parser.add_argument( + "--shm-min-gb", + type=float, + default=8.0, + help="FAIL the node if /dev/shm is below this many GiB " "(NCCL shared-mem may fail). 0 disables.", + ) + # F-partial: rocm-smi self-latency. Hitting this timeout is treated as a + # hard fail because a wedging amdgpu driver typically makes rocm-smi + # hang for 30-60 s before the GPU itself stops responding. + parser.add_argument( + "--rocm-smi-timeout-sec", + type=float, + default=5.0, + help="Hard timeout for `rocm-smi --version`. Hitting it is " "a node FAIL (driver likely wedging).", + ) + # G: foreign / leaked process detection. Hard-fail by default; the + # operator can opt out for partitions that legitimately co-tenant the + # GPU, or whitelist known in-band agents by name. + parser.add_argument( + "--hbm-busy-threshold-gib", + type=float, + default=2.0, + help="FAIL the node if any GPU has at least this much " + "HBM in use BEFORE we touch the device (i.e. someone " + "else is holding it). Boundary is inclusive. Default: 2.0 GiB.", + ) + parser.add_argument( + "--allow-foreign-procs", + action="store_true", + help="Do NOT FAIL the node when foreign processes are " + "found holding a GPU. They will still be reported.", + ) + parser.add_argument( + "--allowed-procs", + type=str, + default="gpuagent,rocm-smi-daemon,amd-smi,dcgm-exporter", + help="Comma-separated process names that are OK to find " + "holding the GPU. The default whitelists the AMD " + "system agents (`gpuagent`, `rocm-smi-daemon`, " + "`amd-smi`) and a common observability sidecar " + "(`dcgm-exporter`) so they don't fail every node. " + "Set to an empty string to disable the whitelist.", + ) + parser.add_argument( + "--gpu-activity-warn-pct", + type=float, + default=20.0, + help="Aggregator warns (does NOT fail) if amd-smi reports " + "any GPU's gfx_activity_pct above this when smoke " + "starts. Default: 20.", + ) + # Tooling availability. Default empty = warn-only (the WARN already + # fires before the per-GPU subprocesses run, and the aggregator + # always renders a "Tooling availability" section). Strict + # environments can pass --require-tools amd-smi,rocm-smi to promote + # a missing tool to a hard node FAIL. + parser.add_argument( + "--require-tools", + type=str, + default="", + help="Comma-separated CLI tool names that MUST be " + "present in PATH for the node to PASS. Anything " + "missing becomes a hard node FAIL. Tracked tools: " + "amd-smi, rocm-smi, lsof. Default: warn-only.", + ) + parser.add_argument( + "--no-clean-dump-path", + action="store_true", + help="Do NOT auto-wipe stale per-node JSONs and aggregator " + "outputs from --dump-path on rank 0 at startup. " + "Default behavior is to clean so re-runs on a " + "different nodelist don't inherit ghost PASS " + "verdicts from removed nodes.", + ) + + +def _add_aggregate_flags(parser: argparse.ArgumentParser, include_dump_path: bool = True) -> None: + """Attach every flag that `aggregate` accepts to ``parser``. + + ``include_dump_path=False`` is used by the primus-cli wrapper because + ``--dump-path`` is already attached by ``_add_run_flags`` -- both share + that flag with identical defaults / meaning, so we attach it once. + """ + if include_dump_path: + parser.add_argument("--dump-path", default="output/preflight", help="Same as `run --dump-path`.") + parser.add_argument( + "--expected-nodes", + type=int, + default=None, + help="Number of nodes expected to report. Missing nodes are FAIL.", + ) + parser.add_argument( + "--wait-timeout-sec", + type=int, + default=60, + help="How long to wait for all expected JSONs to land before aggregating anyway.", + ) + parser.add_argument( + "--rocm-smi-warn-sec", + type=float, + default=1.0, + help="Flag (warn-only) any node where `rocm-smi --version` " "took longer than this many seconds.", + ) + parser.add_argument( + "--clock-skew-warn-sec", + type=float, + default=30.0, + help="Warn (info-only) when wall-clock spread across nodes " + "exceeds this many seconds. Includes srun launch " + "jitter so the default is loose.", + ) + # Mirror the run-side thresholds so the report can label its sections + # using the same numbers each node's `run` was configured with. When + # `_add_aggregate_flags` is called on the same parser as + # `_add_run_flags` (the primus-cli wrapper), these are skipped to avoid + # argparse `conflicting option` errors -- the run-side defaults already + # cover the same names with identical defaults. + if include_dump_path: + parser.add_argument( + "--hbm-busy-threshold-gib", + type=float, + default=2.0, + help="Pre-touch HBM-used threshold (GiB) used by the " + "'GPU pre-touch HBM usage outliers' section.", + ) + parser.add_argument( + "--gpu-activity-warn-pct", + type=float, + default=20.0, + help="GPU activity %% threshold used by the " "'GPU compute-activity outliers' section.", + ) + parser.add_argument( + "--expected-nodelist-file", + type=str, + default=None, + help="Optional file with one expected (short) hostname per line. " + "When provided, missing nodes are reported with their real " + "hostname instead of synthetic placeholders, and " + "are written to failing_nodes.txt directly. The primus-cli " + "wrapper auto-populates this from `scontrol show hostnames` " + "when running under SLURM.", + ) + + +def _resolve_aggregate_args_from_slurm(args: argparse.Namespace) -> argparse.Namespace: + """Build an aggregator-flavored Namespace from ``args``. + + Used by the primus-cli ``node_smoke`` wrapper to populate + ``expected_nodes`` and ``expected_nodelist_file`` from the SLURM + environment when the user did not pass them explicitly. All other + aggregator fields pass through unchanged. The standalone CLI's + ``aggregate`` subcommand is unaffected (it constructs its own + Namespace via argparse). + """ + dump_path = getattr(args, "dump_path", "output/preflight") or "output/preflight" + + # Expected node count: CLI > SLURM_NNODES > SLURM_JOB_NUM_NODES > NNODES. + expected_nodes = getattr(args, "expected_nodes", None) + if expected_nodes is None: + for var in ("SLURM_NNODES", "SLURM_JOB_NUM_NODES", "NNODES"): + v = os.environ.get(var, "") + if v.isdigit() and int(v) > 0: + expected_nodes = int(v) + break + + # Expected nodelist file: CLI > SLURM_JOB_NODELIST (via `scontrol show + # hostnames`). Mirrors the deleted run_node_smoke_direct.sh behavior so + # the aggregator can name MISSING nodes by their real short hostname + # instead of synthetic placeholders. + expected_nodelist_file = getattr(args, "expected_nodelist_file", None) + if not expected_nodelist_file: + slurm_nodelist = os.environ.get("SLURM_JOB_NODELIST", "") + if slurm_nodelist and _which("scontrol"): + try: + os.makedirs(dump_path, exist_ok=True) + candidate = os.path.join(dump_path, "expected_nodes.txt") + import subprocess + + proc = subprocess.run( + ["scontrol", "show", "hostnames", slurm_nodelist], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + check=False, + ) + if proc.returncode == 0 and proc.stdout.strip(): + with open(candidate, "wb") as f: + f.write(proc.stdout) + expected_nodelist_file = candidate + n_hosts = sum(1 for line in proc.stdout.splitlines() if line.strip()) + _log(f"resolved expected nodelist ({n_hosts} nodes) -> {candidate}") + else: + _warn( + "scontrol show hostnames returned empty / non-zero; " + "aggregator will use placeholders" + ) + except OSError as e: + _warn(f"failed to resolve expected nodelist via scontrol: {e}") + + return argparse.Namespace( + dump_path=dump_path, + expected_nodes=expected_nodes, + wait_timeout_sec=int(getattr(args, "wait_timeout_sec", 60)), + rocm_smi_warn_sec=float(getattr(args, "rocm_smi_warn_sec", 1.0)), + clock_skew_warn_sec=float(getattr(args, "clock_skew_warn_sec", 30.0)), + hbm_busy_threshold_gib=float(getattr(args, "hbm_busy_threshold_gib", 2.0)), + gpu_activity_warn_pct=float(getattr(args, "gpu_activity_warn_pct", 20.0)), + expected_nodelist_file=expected_nodelist_file, + ) + + +def _which(cmd: str) -> Optional[str]: + """Lightweight ``shutil.which`` import (avoid module-level import cost).""" + import shutil + + return shutil.which(cmd) + + +def _build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + prog="python -m primus.tools.preflight.node_smoke", + description=( + "Node-local preflight smoke test. Each node runs independently " + "(no global rendezvous) and writes a per-node JSON verdict." + ), + ) + sub = p.add_subparsers(dest="cmd", required=True) + + # ---- run ---- + # allow_abbrev=False so abbreviated forms (e.g. an old --tier2 left in + # a script) do NOT silently match the new --tier2-perf as a prefix. + # We want them to error out loudly so behavior never changes silently + # underneath an unsuspecting caller. + pr = sub.add_parser( + "run", + help="Run per-node smoke test on this node.", + allow_abbrev=False, + ) + _add_run_flags(pr) + pr.set_defaults(func=_cmd_run) + + # ---- aggregate ---- + pa = sub.add_parser("aggregate", help="Aggregate per-node JSONs into report + passing/failing lists.") + _add_aggregate_flags(pa) + pa.set_defaults(func=_cmd_aggregate) + + # ---- _per_gpu (internal) ---- + pg = sub.add_parser( + "_per_gpu", help="(internal) Run smoke checks for a single GPU index. Spawned by `run`." + ) + pg.add_argument("gpu", type=int) + pg.add_argument("--tier2-perf", action="store_true") + pg.add_argument("--gemm-tflops-min", type=float, default=600.0) + pg.add_argument("--hbm-gbs-min", type=float, default=2000.0) + pg.add_argument("--hbm-busy-threshold-gib", type=float, default=2.0) + pg.set_defaults(func=_cmd_per_gpu) + + return p + + +def main(argv: Optional[List[str]] = None) -> int: + parser = _build_parser() + ns = parser.parse_args(argv) + return int(ns.func(ns)) diff --git a/primus/tools/preflight/node_smoke/collectors/__init__.py b/primus/tools/preflight/node_smoke/collectors/__init__.py new file mode 100644 index 000000000..611ef9105 --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/__init__.py @@ -0,0 +1,23 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 1 collectors -- one module per Tier 1 sub-section. + +* ``dmesg`` : recent dmesg error scan +* ``fingerprint`` : Tier 1 A -- software-stack fingerprint +* ``nics`` : Tier 1 B -- NIC / RDMA roll-call +* ``host_limits`` : Tier 1 C -- host limits (ulimit, /dev/shm, NUMA, governor) +* ``gpu_low_level``: Tier 1 D-1 -- per-GPU ECC / clocks / power via amd-smi +* ``xgmi`` : Tier 1 D-2 -- XGMI topology matrix +* ``clock`` : Tier 1 E -- wall time + time-daemon active states +* ``rocm_smi`` : Tier 1 F + fallbacks -- rocm-smi self-latency and + amd-smi fallback parsers +* ``gpu_processes``: Tier 1 G -- foreign / leaked process detection +* ``tooling`` : tooling-availability inventory +* ``reused_info`` : reused gpu/host/network info collectors +""" + +from __future__ import annotations diff --git a/primus/tools/preflight/node_smoke/collectors/clock.py b/primus/tools/preflight/node_smoke/collectors/clock.py new file mode 100644 index 000000000..86b13b0c9 --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/clock.py @@ -0,0 +1,36 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 1 -- E: clock state (wall time + time-daemon active states).""" + +from __future__ import annotations + +import time +from typing import Any, Dict + +from ..shell_utils import _systemctl_is_active + + +def _collect_clock_state() -> Dict[str, Any]: + """Capture this node's wall time and time-daemon health. + + Wall time is captured early so the aggregator can compute a + cluster-wide spread. Note this includes srun launch jitter, so the + spread is an *upper bound* on the real clock skew. The aggregator + uses loose thresholds (warn at 30 s, no hard fail) for the spread, + and reserves the hard fail for "no time-sync daemon active". + """ + out: Dict[str, Any] = { + "wall_time_unix": time.time(), + "monotonic": time.monotonic(), + "daemons": {}, + } + for unit in ("chronyd", "ntp", "ntpd", "systemd-timesyncd"): + out["daemons"][unit] = _systemctl_is_active(unit) + active = [u for u, s in out["daemons"].items() if s == "active"] + out["any_active"] = bool(active) + out["active_units"] = active + return out diff --git a/primus/tools/preflight/node_smoke/collectors/dmesg.py b/primus/tools/preflight/node_smoke/collectors/dmesg.py new file mode 100644 index 000000000..6ae3628fb --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/dmesg.py @@ -0,0 +1,98 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Recent dmesg error scan. + +Greps the last ``window_minutes`` of dmesg for known-bad patterns. Best +effort: missing dmesg / failed read becomes ``ok=False`` rather than +raising. +""" + +from __future__ import annotations + +import subprocess +from typing import Any, Dict, List + +# Regex (NOT substring) patterns matched case-insensitively against each +# dmesg line. They MUST be valid Python regex -- if you only want a literal +# substring (e.g. ``mce: ``), it's still a valid regex with no specials. +# Why regex: real amdgpu failure lines look like +# ``amdgpu 0000:05:00.0: amdgpu_device_resume failed: -19`` +# ``amdgpu: [drm] *ERROR* ring sdma0 timeout`` +# so we need ``amdgpu.*(error|fail|timeout)`` -- a substring match against +# the literal pattern ``amdgpu.*error`` would essentially never fire. +_DMESG_PATTERNS = ( + r"\bxid\b", + r"hardware error", + r"gpu reset", + r"hung_task", + r"hung task", + r"page allocation failure", + r"soft lockup", + r"amdgpu.*(error|fail|timeout)", + r"\*error\*", # [drm] *ERROR* + r"mce: ", +) + + +def _collect_dmesg_errors(window_minutes: int = 15) -> Dict[str, Any]: + """Best-effort grep of recent dmesg lines for known-bad patterns. + + Returns a dict with ``ok`` (bool), ``matches`` (list of matched lines, capped), + and ``error`` (str) when dmesg cannot be read. + """ + out: Dict[str, Any] = {"ok": True, "matches": [], "error": None} + try: + # ``--since`` requires recent util-linux; fall back to the last 2000 lines. + try: + cp = subprocess.run( + ["dmesg", "--since", f"-{window_minutes}min"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=10, + check=False, + ) + if cp.returncode != 0: + raise RuntimeError(cp.stderr.strip() or f"rc={cp.returncode}") + text = cp.stdout + except Exception: + cp = subprocess.run( + ["dmesg"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=10, + check=False, + ) + if cp.returncode != 0: + out["ok"] = False + out["error"] = (cp.stderr or "").strip() or f"rc={cp.returncode}" + return out + text = "\n".join(cp.stdout.splitlines()[-2000:]) + + import re + + matches: List[str] = [] + # Pre-compile each pattern with re.IGNORECASE; a malformed regex is + # logged into ``out['pattern_errors']`` but never aborts the scan. + compiled: List[Any] = [] + for p in _DMESG_PATTERNS: + try: + compiled.append(re.compile(p, re.IGNORECASE)) + except re.error as e: + out.setdefault("pattern_errors", []).append(f"{p!r}: {e}") + for line in text.splitlines(): + if any(pat.search(line) for pat in compiled): + matches.append(line) + if len(matches) >= 50: + break + out["matches"] = matches + return out + except Exception as e: + out["ok"] = False + out["error"] = str(e) + return out diff --git a/primus/tools/preflight/node_smoke/collectors/fingerprint.py b/primus/tools/preflight/node_smoke/collectors/fingerprint.py new file mode 100644 index 000000000..225a9b94a --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/fingerprint.py @@ -0,0 +1,88 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 1 -- A. Software-stack fingerprint (drift detection happens at aggregate).""" + +from __future__ import annotations + +import os +import sys +from typing import Any, Dict + +from ..shell_utils import _parse_os_release_pretty, _read_text + + +def _collect_node_fingerprint() -> Dict[str, Any]: + """Collect a deterministic, hashable fingerprint of the software stack + on this node so the aggregator can detect drift across the cluster. + + Every value is best-effort: missing tools / files become ``None`` rather + than raising. The aggregator skips ``None`` values when computing the + cluster majority for a given key. + """ + fp: Dict[str, Any] = {} + + # Kernel + OS + try: + fp["kernel"] = os.uname().release + except Exception: + fp["kernel"] = None + fp["os_release"] = _parse_os_release_pretty() + fp["python"] = sys.version.split()[0] + + # ROCm / HIP / amdgpu + fp["rocm"] = _read_text("/opt/rocm/.info/version") or None + fp["amdgpu_driver"] = _read_text("/sys/module/amdgpu/version") or None + + # PyTorch + (R)CCL + try: + import torch # type: ignore + + fp["torch"] = getattr(torch, "__version__", None) + fp["torch_hip"] = getattr(getattr(torch, "version", None), "hip", None) + try: + v = torch.cuda.nccl.version() # type: ignore[attr-defined] + if isinstance(v, tuple): + fp["rccl"] = ".".join(str(x) for x in v) + else: + fp["rccl"] = str(v) + except Exception: + fp["rccl"] = None + + # Locate librccl.so under torch's lib dir for a stable per-node path. + try: + torch_lib = os.path.join(os.path.dirname(torch.__file__), "lib") + for n in sorted(os.listdir(torch_lib)): + if n.startswith("librccl.so"): + fp["rccl_path"] = os.path.join(torch_lib, n) + break + except Exception: + pass + except Exception: + fp["torch"] = None + fp["torch_hip"] = None + fp["rccl"] = None + + # Per-IB-device firmware + HCA model fingerprints. Both are critical for + # detecting "1 of N nodes flashed differently" silent regressions. + nic_fw: Dict[str, str] = {} + nic_hca: Dict[str, str] = {} + ib_root = "/sys/class/infiniband" + if os.path.isdir(ib_root): + try: + for dev in sorted(os.listdir(ib_root)): + fw = _read_text(os.path.join(ib_root, dev, "fw_ver")) + if fw: + nic_fw[dev] = fw + hca = _read_text(os.path.join(ib_root, dev, "hca_type")) + if hca: + nic_hca[dev] = hca + except Exception: + pass + fp["nic_fw"] = nic_fw or None + fp["nic_hca"] = nic_hca or None + + return fp diff --git a/primus/tools/preflight/node_smoke/collectors/gpu_low_level.py b/primus/tools/preflight/node_smoke/collectors/gpu_low_level.py new file mode 100644 index 000000000..27b99dbef --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/gpu_low_level.py @@ -0,0 +1,249 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 1 -- D-1 heavy: per-GPU low-level via amd-smi (ECC, throttle, clocks, +power cap). Runs ONCE per node (not per per-GPU subprocess) so the smoke +step doesn't pay an amd-smi startup tax 8x. Best-effort: missing amd-smi +or unparseable output degrades to {"ok": False, ...} without raising. +""" + +from __future__ import annotations + +import json +import subprocess +from typing import Any, Dict, List + +from ..shell_utils import _which +from .rocm_smi import _rocm_smi_ras_info_text, _rocm_smi_use_json + + +def _collect_amd_smi_metrics() -> Dict[str, Any]: + """Best-effort capture of per-GPU low-level metrics via ``amd-smi``. + + We try ``amd-smi metric --json`` first (newer builds emit valid JSON); + if that fails we fall back to text output and surface the raw text under + ``raw`` so an operator can still grep it. The on-disk shape is: + + { + "ok": bool, + "tool": "amd-smi metric --json" | "amd-smi metric" | None, + "per_gpu": [ {gpu, gfx_clock_mhz, hbm_used_bytes, + power_avg_w, power_cap_w, temp_edge_c, + ecc_uncorrectable_total, ecc_correctable_total, + throttle_status_raw, ...}, ... ], + "error": "..." (only when ok is False) + } + + Hard-fail semantics live in ``_node_status_from``: any non-zero + uncorrectable ECC count becomes a node FAIL. Throttle status is + captured under ``throttle_status_raw`` for operator inspection but + is NOT failed-on -- the amd-smi throttle schema varies too much + across releases to make a robust default rule. + """ + out: Dict[str, Any] = {"ok": False, "tool": None, "per_gpu": []} + + if _which("amd-smi") is not None: + # Try JSON first. + try: + cp = subprocess.run( + ["amd-smi", "metric", "--json"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=15, + check=False, + ) + if cp.returncode == 0 and cp.stdout.strip(): + try: + doc = json.loads(cp.stdout) + out["ok"] = True + out["tool"] = "amd-smi metric --json" + out["per_gpu"] = _flatten_amd_smi_metric_json(doc) + except Exception as e: + out["json_parse_error"] = str(e) + else: + out["json_rc"] = cp.returncode + out["json_stderr"] = (cp.stderr or "").strip()[:200] + except subprocess.TimeoutExpired: + out["json_error"] = "amd-smi metric --json timed out" + except Exception as e: + out["json_error"] = str(e) + + # If JSON didn't work, capture the raw text output so the operator + # can still grep it. We don't try to parse the human-readable text + # -- per-GPU outliers will still show up via the sysfs/torch-side + # details we capture in _per_gpu_body, and the rocm-smi fallback + # below will fill in ECC + activity in the per_gpu records. + if not out["ok"]: + try: + cp = subprocess.run( + ["amd-smi", "metric"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=15, + check=False, + ) + if cp.returncode == 0: + out["ok"] = True + out["tool"] = "amd-smi metric" + out["raw"] = cp.stdout[:8000] # cap to keep JSON small + else: + out["error"] = (cp.stderr or "").strip()[:200] or f"rc={cp.returncode}" + except subprocess.TimeoutExpired: + out["error"] = "amd-smi metric timed out" + except Exception as e: + out["error"] = str(e) + else: + out["error"] = "amd-smi not found in PATH" + + # rocm-smi fallback / fill-in. We always run it -- even when amd-smi + # produced JSON, the schema may have missed ECC or activity fields, + # and rocm-smi gives us a second independent source. Per-GPU records + # are merged by GPU index so amd-smi-provided power/clocks/temp stay + # alongside rocm-smi-provided ECC/activity. + if _which("rocm-smi") is not None: + # Build an index of existing records so we can merge in place. + by_gpu: Dict[int, Dict[str, Any]] = {} + for rec in out.get("per_gpu") or []: + g = rec.get("gpu") + if isinstance(g, int): + by_gpu[g] = rec + + # ECC fill-in (only if not already populated by amd-smi) + ras = _rocm_smi_ras_info_text() + ecc_filled = 0 + if ras.get("ok"): + for rec in ras.get("per_gpu") or []: + g = rec.get("gpu") + if not isinstance(g, int): + continue + tgt = by_gpu.setdefault(g, {"gpu": g}) + if tgt.get("ecc_uncorrectable_total") is None: + tgt["ecc_uncorrectable_total"] = rec.get("ecc_uncorrectable_total") + tgt["ecc_correctable_total"] = rec.get("ecc_correctable_total") + tgt["ecc_source"] = "rocm-smi --showrasinfo" + ecc_filled += 1 + if ecc_filled: + out.setdefault("fallback_tools", []).append( + f"rocm-smi --showrasinfo (ECC for {ecc_filled} GPU(s))" + ) + + # Activity fill-in (only if not already populated by amd-smi) + use = _rocm_smi_use_json() + act_filled = 0 + if use.get("ok"): + for rec in use.get("per_gpu") or []: + g = rec.get("gpu") + if not isinstance(g, int): + continue + tgt = by_gpu.setdefault(g, {"gpu": g}) + if tgt.get("gfx_activity_pct") is None: + tgt["gfx_activity_pct"] = rec.get("gfx_activity_pct") + tgt["activity_source"] = "rocm-smi --showuse" + act_filled += 1 + if act_filled: + out.setdefault("fallback_tools", []).append( + f"rocm-smi --showuse (activity for {act_filled} GPU(s))" + ) + + # Rebuild per_gpu in stable index order if rocm-smi added entries + if ecc_filled or act_filled: + out["per_gpu"] = [by_gpu[g] for g in sorted(by_gpu.keys())] + # If amd-smi produced nothing at all, this is now a successful + # collection, just sourced entirely from rocm-smi. + if not out["ok"] and out["per_gpu"]: + out["ok"] = True + out["tool"] = "rocm-smi (ECC/activity fallback)" + + return out + + +def _flatten_amd_smi_metric_json(doc: Any) -> List[Dict[str, Any]]: + """Pull the fields we care about out of `amd-smi metric --json` output. + + The exact schema varies between amd-smi releases. We touch only the + most-stable nesting -- a top-level list of per-GPU dicts, each with + sub-blocks like ``power``, ``clock``, ``temperature``, ``ecc``, + ``throttle_status`` -- and tolerate missing fields silently. + """ + out: List[Dict[str, Any]] = [] + items = doc if isinstance(doc, list) else (doc.get("gpus", []) if isinstance(doc, dict) else []) + for i, g in enumerate(items): + if not isinstance(g, dict): + continue + rec: Dict[str, Any] = {"gpu": i} + # gpu id may be in g["gpu"] or g["device_id"] depending on schema + if isinstance(g.get("gpu"), int): + rec["gpu"] = g["gpu"] + # power + power = g.get("power") or {} + if isinstance(power, dict): + for k_src, k_dst in ( + ("average_socket_power", "power_avg_w"), + ("current_socket_power", "power_avg_w"), + ("socket_power", "power_avg_w"), + ("power_cap", "power_cap_w"), + ("power_limit", "power_cap_w"), + ): + v = power.get(k_src) + if isinstance(v, (int, float)) and rec.get(k_dst) is None: + rec[k_dst] = v + # clocks (gfx clock most useful) + clk = g.get("clock") or g.get("clocks") or {} + if isinstance(clk, dict): + gfx = clk.get("gfx") or clk.get("gfx_0") or clk.get("gfx_clock") or {} + if isinstance(gfx, dict): + for k in ("clk", "current", "value", "frequency"): + if isinstance(gfx.get(k), (int, float)): + rec["gfx_clock_mhz"] = gfx[k] + break + elif isinstance(gfx, (int, float)): + rec["gfx_clock_mhz"] = gfx + # temperature + temp = g.get("temperature") or {} + if isinstance(temp, dict): + for k in ("edge", "current", "value"): + v = temp.get(k) + if isinstance(v, (int, float)): + rec["temp_edge_c"] = v + break + # ECC + ecc = g.get("ecc") or g.get("ecc_count") or {} + if isinstance(ecc, dict): + ue = ecc.get("uncorrectable") or ecc.get("uncorrectable_total") or ecc.get("ue") or 0 + ce = ecc.get("correctable") or ecc.get("correctable_total") or ecc.get("ce") or 0 + try: + rec["ecc_uncorrectable_total"] = int(ue) + rec["ecc_correctable_total"] = int(ce) + except Exception: + pass + # Throttle + thr = g.get("throttle_status") or g.get("throttle") or {} + if isinstance(thr, dict): + rec["throttle_status_raw"] = thr + elif isinstance(thr, (str, list)): + rec["throttle_status_raw"] = thr + # GPU compute activity %. Used by the aggregator to surface GPUs + # that are running someone else's compute right now (warn-only; + # short bursts are normal but a sustained pegged-100% across + # multiple GPUs is a strong signal of a leaked rank from a + # previous job). + usage = g.get("usage") or g.get("activity") or g.get("utilization") or {} + if isinstance(usage, dict): + for k in ("gfx_activity", "gfx", "gfx_busy_percent", "gpu_busy_percent"): + v = usage.get(k) + if isinstance(v, (int, float)): + rec["gfx_activity_pct"] = float(v) + break + if rec.get("gfx_activity_pct") is None: + v = usage.get("activity") or usage.get("value") + if isinstance(v, (int, float)): + rec["gfx_activity_pct"] = float(v) + elif isinstance(usage, (int, float)): + rec["gfx_activity_pct"] = float(usage) + out.append(rec) + return out diff --git a/primus/tools/preflight/node_smoke/collectors/gpu_processes.py b/primus/tools/preflight/node_smoke/collectors/gpu_processes.py new file mode 100644 index 000000000..9e24ba19b --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/gpu_processes.py @@ -0,0 +1,614 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 1 -- G: foreign / leaked process detection on each GPU. + +The single most common reason a "healthy" cluster fails to launch a large +training job is that a previous job's Python ranks are still attached to +the GPUs (held HBM, half-torn-down NCCL communicators, or just stuck in +__del__). Symptoms in the new training job: torch.cuda.OutOfMemoryError +at model init with a misleading "free=Y" message, NCCL/RCCL bootstrap +hang, or random ranks failing the first all-reduce due to compute +contention. node-smoke catches these BEFORE the operator launches the +real job by enumerating PIDs that hold each GPU and FAILing the node +unless the operator explicitly opted in via --allow-foreign-procs. +""" + +from __future__ import annotations + +import json +import os +import subprocess +from typing import Any, Dict, List, Optional, Tuple + +from ..shell_utils import _parse_size_with_unit, _which +from .rocm_smi import _rocm_smi_processes + +# Process-name placeholders that some versions of `amd-smi process` and +# `rocm-smi --showpids` emit when they cannot read the real name from +# /proc//comm (typically for kernel/system-owned PIDs like +# `gpuagent`). We treat these as "missing" and fall back to /proc. +_MISSING_NAME_TOKENS = frozenset({"", "n/a", "na", "none", "null", "-", "unknown", "?"}) + + +def _is_missing_name(name: str) -> bool: + return (name or "").strip().lower() in _MISSING_NAME_TOKENS + + +def _resolve_proc_name(pid: int) -> str: + """Best-effort recover the process name for ``pid`` from /proc. + + Tries ``/proc//comm`` first (15-char kernel name, world-readable + on default Linux), then falls back to the ``Name:`` line in + ``/proc//status``. Returns ``""`` when both reads fail (PID gone, + hidepid mount, ptrace_scope, etc.) so the caller can decide whether + to keep the original placeholder. + """ + try: + with open(f"/proc/{int(pid)}/comm", "r", encoding="utf-8") as f: + name = f.read().strip() + if name: + return name + except OSError: + pass + try: + with open(f"/proc/{int(pid)}/status", "r", encoding="utf-8") as f: + for line in f: + if line.startswith("Name:"): + return line.split(":", 1)[1].strip() + except OSError: + pass + return "" + + +def _resolve_self_pid_view(self_pid: int) -> Dict[str, Any]: + """Resolve our own PID as the host kernel sees it, even from inside a + private PID namespace (e.g. a Docker / Kubernetes container). + + ``amd-smi process``, ``rocm-smi --showpids`` and ``lsof /dev/kfd`` + all report PIDs in the **root (host) PID namespace** because KFD is + a host-kernel resource that knows nothing about user namespaces. + ``os.getpid()``, by contrast, returns the PID *as our own namespace + sees it*. In a private PID namespace these are different numbers, + and the naive ``reported_pid == os.getpid()`` test would always + return False -- causing our own training rank to be flagged + ``is_foreign=True`` and (with the default policy) failing the node. + + The kernel exposes the full mapping in ``/proc/self/status``: + + NSpid: 2005679 42 + + where the first entry is the deepest-namespace PID (= the host PID + on bare metal) and the LAST entry is the most-deeply-nested PID + (= what ``os.getpid()`` returns inside a private namespace). On a + bare-metal host or a non-namespaced container, the line has a + single field equal to ``os.getpid()``. + + Returns:: + + { + "host_pid": int, # what amd-smi/rocm-smi report + "container_pid": int, # == self_pid (passed in) + "pid_namespaced": bool, # True if the two differ + "ns_chain": [int, ...], # full NSpid chain (for the report) + } + + Best-effort: if ``/proc/self/status`` cannot be read or has no + ``NSpid`` line, we assume bare-metal and fall back to ``self_pid``. + """ + out: Dict[str, Any] = { + "host_pid": int(self_pid), + "container_pid": int(self_pid), + "pid_namespaced": False, + "ns_chain": [int(self_pid)], + } + try: + with open("/proc/self/status", "r", encoding="utf-8") as f: + for line in f: + if not line.startswith("NSpid:"): + continue + parts = line.split()[1:] # drop "NSpid:" + chain: List[int] = [] + for tok in parts: + try: + chain.append(int(tok)) + except ValueError: + continue + if not chain: + break + out["ns_chain"] = chain + # Convention: NSpid lists outermost (host) namespace first + # and the current namespace last; matches /proc man page. + out["host_pid"] = chain[0] + out["container_pid"] = chain[-1] + out["pid_namespaced"] = chain[0] != chain[-1] + break + except Exception: + pass + return out + + +def _collect_gpu_processes( + self_pid: int, + allowed_proc_names: Optional[List[str]] = None, +) -> Dict[str, Any]: + """Enumerate processes currently holding each GPU on this node. + + Tries, in order: + 1. ``amd-smi process --json`` (preferred -- structured) + 2. ``amd-smi process`` (text fallback -- best-effort parse) + 3. ``lsof /dev/kfd /dev/dri/renderD*`` (last resort: just openers) + + Output shape:: + + { + "ok": bool, + "tool": "amd-smi process --json" | "amd-smi process" | "lsof" | None, + "self_pid": int, # container/local view (legacy) + "self_pgid": int, # pgid in our own namespace + "self_host_pid": int, # PID amd-smi/rocm-smi report for us + "pid_namespaced": bool, # True inside a private PID ns + "ns_pid_chain": [int, ...], # full NSpid chain (root..ours) + "allowed_proc_names": [str, ...], # passthrough for aggregator + "per_gpu": [ + {"gpu": 0, "processes": [ + {"pid": int, "name": str, "hbm_bytes": int|None, + "is_self": bool, "is_allowed": bool, "is_foreign": bool}, + ... + ]}, + ... + ], + "foreign_count": int, # PIDs not us and not allowed + "error": str # only when ok is False + } + + Filtering: a PID is treated as "ours" (and thus not foreign) if it + matches our **host-namespace PID** (because amd-smi / rocm-smi / + lsof always report root-ns PIDs) or, when we are NOT inside a + private PID namespace, our pgid (which catches per-GPU subprocesses + we may have spawned). Inside a private PID namespace the pgid match + is intentionally skipped: we cannot ``os.getpgid()`` PIDs we cannot + see, and any spawned subprocess will appear as a sibling host PID + that the operator should treat the same way as a leaked rank. + + ``allowed_proc_names`` is a case-insensitive name allow-list for + known node-resident agents (``rocm-smi-daemon``, ``amd-smi``, + ``dcgm-exporter``, ``gpuagent``, ...). + """ + allowed = {n.strip().lower() for n in (allowed_proc_names or []) if n.strip()} + + # Resolve host-side PID. Inside a private PID namespace this is the + # number amd-smi / rocm-smi / lsof will report for us; on bare metal + # or a shared-PID-ns container (the SLURM + pyxis/enroot default) it + # equals self_pid. + pid_view = _resolve_self_pid_view(int(self_pid)) + host_self_pid = int(pid_view["host_pid"]) + pid_namespaced = bool(pid_view["pid_namespaced"]) + + # pgid is only meaningful within OUR own PID namespace -- attempting + # os.getpgid() on a host-side PID we cannot see would raise ESRCH. + try: + self_pgid = os.getpgid(int(self_pid)) + except OSError: + self_pgid = int(self_pid) + + out: Dict[str, Any] = { + "ok": False, + "tool": None, + "self_pid": int(self_pid), + "self_pgid": int(self_pgid), + "self_host_pid": host_self_pid, + "pid_namespaced": pid_namespaced, + "ns_pid_chain": list(pid_view.get("ns_chain") or [int(self_pid)]), + "allowed_proc_names": sorted(allowed), + "per_gpu": [], + "foreign_count": 0, + } + + def _annotate(pid: int, name: str, hbm_bytes: Optional[int]) -> Dict[str, Any]: + # Direct PID match against our HOST-side PID (works in every + # mode: bare metal, shared PID ns, private PID ns). + is_self = int(pid) == host_self_pid + # pgid match is only safe outside a private PID namespace; inside + # one we cannot see host PIDs at all and os.getpgid() would + # ESRCH for every reported pid. + if not is_self and not pid_namespaced: + try: + pgid = os.getpgid(int(pid)) + except OSError: + pgid = -1 + if pgid == self_pgid: + is_self = True + + # Some amd-smi / rocm-smi builds report `name="N/A"` (or empty) + # for kernel/system PIDs like `gpuagent` because they cannot read + # /proc//comm themselves. We can: do it inline so the + # allowlist actually matches and the report shows the real name. + raw_name = name or "" + name_resolved_from_proc = False + if _is_missing_name(raw_name): + recovered = _resolve_proc_name(int(pid)) + if recovered: + name = recovered + name_resolved_from_proc = True + + is_allowed = (name or "").strip().lower() in allowed + return { + "pid": int(pid), + "name": name or "", + "name_raw": raw_name, + "name_resolved_from_proc": bool(name_resolved_from_proc), + "hbm_bytes": int(hbm_bytes) if isinstance(hbm_bytes, (int, float)) else None, + "is_self": bool(is_self), + "is_allowed": bool(is_allowed), + "is_foreign": bool(not is_self and not is_allowed), + } + + if _which("amd-smi") is not None: + # 1) amd-smi process --json + try: + cp = subprocess.run( + ["amd-smi", "process", "--json"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=15, + check=False, + ) + if cp.returncode == 0 and cp.stdout.strip(): + try: + doc = json.loads(cp.stdout) + parsed, schema_drift = _flatten_amd_smi_process_json(doc, _annotate) + if parsed and not schema_drift: + # Schema fully recognized (Shape A/A'/B): trust the + # result even when each per-GPU bucket is empty. + out["per_gpu"] = parsed + out["tool"] = "amd-smi process --json" + out["ok"] = True + elif schema_drift: + # Top-level shape matched (Shape A bucket layout or + # Shape B per-process dicts) but at least one record + # had no usable pid / process_id -- the smoking gun + # for a future amd-smi rename. Fall through to the + # text / rocm-smi / lsof chain rather than silently + # reporting the node clean. + out["json_parse_error"] = ( + "amd-smi process --json: per-process schema not " + "recognized (pid / process_id missing on at least " + "one record); falling back to text / rocm-smi / lsof" + ) + else: + # Valid JSON but no recognizable per-GPU entries -- a + # future amd-smi schema we don't speak yet. Fall + # through to text / rocm-smi / lsof rather than + # silently masking the node clean. + out["json_parse_error"] = ( + "amd-smi process --json returned no recognizable " + "per-GPU dicts; falling back to text / rocm-smi / lsof" + ) + except Exception as e: + out["json_parse_error"] = str(e) + else: + out["json_rc"] = cp.returncode + out["json_stderr"] = (cp.stderr or "").strip()[:200] + except subprocess.TimeoutExpired: + out["json_error"] = "amd-smi process --json timed out" + except Exception as e: + out["json_error"] = str(e) + + # 2) amd-smi process (text) + if not out["ok"]: + try: + cp = subprocess.run( + ["amd-smi", "process"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=15, + check=False, + ) + if cp.returncode == 0: + parsed = _parse_amd_smi_process_text(cp.stdout, _annotate) + out["per_gpu"] = parsed + out["tool"] = "amd-smi process" + out["ok"] = True + out["raw"] = cp.stdout[:4000] + else: + out["text_rc"] = cp.returncode + out["text_stderr"] = (cp.stderr or "").strip()[:200] + except subprocess.TimeoutExpired: + out["text_error"] = "amd-smi process timed out" + except Exception as e: + out["text_error"] = str(e) + + # 3) rocm-smi --showpids --json. Doesn't give us per-GPU mapping + # (--showpidgpus emits "WARNING: No JSON data to report"), so all + # PIDs go into the gpu=-1 bucket -- same convention as the lsof + # fallback. This is still a HUGE win over lsof: rocm-smi reports + # the actual KFD process name which lets the operator decide if + # it's a leaked rank vs a known agent. + if not out["ok"]: + rocm = _rocm_smi_processes(_annotate) + if rocm.get("ok"): + out["per_gpu"] = rocm.get("per_gpu") or [] + out["tool"] = rocm.get("tool") + out["ok"] = True + else: + out["rocm_smi_error"] = rocm.get("error") + + # 4) lsof on /dev/kfd + /dev/dri/renderD* (we cannot map back to specific + # GPUs reliably this way, but at least we surface foreign openers). + if not out["ok"]: + try: + import glob as _glob + + paths = ["/dev/kfd"] + sorted(_glob.glob("/dev/dri/renderD*")) + existing = [p for p in paths if os.path.exists(p)] + if existing and _which("lsof") is not None: + cp = subprocess.run( + ["lsof", "-Fpcn", "--", *existing], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=10, + check=False, + ) + if cp.returncode in (0, 1): # lsof exits 1 when nothing open + procs = _parse_lsof_pcn(cp.stdout, _annotate) + # Without per-GPU mapping, surface as a single bucket + # under gpu=-1 so the aggregator still flags foreigners. + out["per_gpu"] = [{"gpu": -1, "processes": procs}] if procs else [] + out["tool"] = "lsof" + out["ok"] = True + else: + out["lsof_rc"] = cp.returncode + except Exception as e: + out["lsof_error"] = str(e) + + if not out["ok"] and "error" not in out: + out["error"] = "no working enumeration tool " "(amd-smi process / rocm-smi --showpids / lsof)" + + out["foreign_count"] = sum( + 1 for g in out["per_gpu"] for p in (g.get("processes") or []) if p.get("is_foreign") + ) + return out + + +def _flatten_amd_smi_process_json( + doc: Any, + annotate: Any, +) -> Tuple[List[Dict[str, Any]], bool]: + """Coerce the amd-smi process JSON into a stable per-GPU shape. + + The schema varies across releases; we tolerate missing fields silently. + Three top-level shapes are seen in the wild: + + A) Per-GPU dicts each carrying ``process_list``, where each list + item wraps the actual process under ``process_info`` and reports + memory as ``{"value": N, "unit": "B"}`` -- this is the modern + ``amd-smi`` (>= 6.x) layout:: + + [{"gpu": 0, "process_list": [ + {"process_info": { + "pid": 12345, "name": "python", + "memory_usage": {"vram_mem": {"value": 256, "unit": "MB"}} + }} + ]}] + + A') Same as A but each ``process_list`` item is the process dict + directly (older amd-smi releases), with memory as a flat int or + formatted string:: + + [{"gpu": 0, "process_list": [ + {"pid": 12345, "name": "python", + "memory_usage": {"vram_mem": "256 MB"}} + ]}] + + B) Top-level list of per-process dicts each carrying explicit + ``gpu`` / ``gpus`` (uncommon, but seen on a couple of branches):: + + [{"pid": ..., "name": ..., "gpu": 0, "memory_usage": {...}}] + + Returns ``(per_gpu, schema_drift_suspected)``. The drift flag is raised + when at least one record looked like a process dict at the structural + level (i.e. we entered ``_push``) but the per-process pid/process_id + field was missing or had a non-numeric value -- the smoking gun for a + future amd-smi schema rename. The caller uses this to fall through to + the text/rocm-smi/lsof fallbacks instead of trusting an empty result. + A genuinely clean node (empty ``process_list``) never enters ``_push`` + so the flag stays False and the fast path is preserved. + """ + out_by_gpu: Dict[int, List[Dict[str, Any]]] = {} + schema_drift_suspected = False + + def _push(gpu: int, pid: Any, name: Any, hbm: Any) -> None: + nonlocal schema_drift_suspected + if not isinstance(pid, (int, float)): + schema_drift_suspected = True + return + ann = annotate(int(pid), str(name or ""), hbm) + out_by_gpu.setdefault(int(gpu), []).append(ann) + + def _value_unit_to_bytes(v: Any) -> Optional[int]: + """Resolve a size value that may be int, formatted string, or + ``{"value": N, "unit": "B|KB|MB|GB|..."}`` (modern amd-smi shape) + into bytes.""" + if isinstance(v, (int, float)): + return int(v) + if isinstance(v, str): + return _parse_size_with_unit(v) + if isinstance(v, dict): + val = v.get("value") + unit = v.get("unit") + if isinstance(val, (int, float)): + if isinstance(unit, str) and unit.strip(): + return _parse_size_with_unit(f"{val} {unit.strip()}") + return int(val) + if isinstance(val, str): + if isinstance(unit, str) and unit.strip(): + return _parse_size_with_unit(f"{val} {unit.strip()}") + return _parse_size_with_unit(val) + return None + + def _hbm_of(d: Dict[str, Any]) -> Optional[int]: + # Preferred path: memory_usage.vram_mem (covers shapes A, A', B). + mu = d.get("memory_usage") if isinstance(d.get("memory_usage"), dict) else None + if mu is not None: + for k in ("vram_mem", "vram_memory", "vram"): + if k in mu: + n = _value_unit_to_bytes(mu.get(k)) + if n is not None: + return n + # Fallback: mem_usage at the same level (older shapes -- usually + # mirrors vram_mem on AMD GPUs since GTT/CPU are negligible). + if "mem_usage" in d: + n = _value_unit_to_bytes(d.get("mem_usage")) + if n is not None: + return n + # Last resort: a flat "vram" key directly under d. + if "vram" in d: + n = _value_unit_to_bytes(d.get("vram")) + if n is not None: + return n + return None + + def _unwrap_proc(p: Dict[str, Any]) -> Dict[str, Any]: + """Modern amd-smi wraps each process under ``process_info``; + unwrap so the rest of the parser can read pid/name/memory at + the top level uniformly.""" + if isinstance(p.get("process_info"), dict): + return p["process_info"] + return p + + items = ( + doc + if isinstance(doc, list) + else (doc.get("processes") or doc.get("gpus") or [] if isinstance(doc, dict) else []) + ) + if not isinstance(items, list): + items = [] + + for item in items: + if not isinstance(item, dict): + continue + # Shape A / A': per-GPU dict with process_list. Use explicit + # presence checks rather than `... or ...` -- an empty list is + # the *clean GPU* signal we want to honor, and `[] or x` would + # short-circuit past it to ``processes``, then to None, and miss + # the Shape A bucket pre-registration entirely. + plist = item.get("process_list") + if not isinstance(plist, list): + plist = item.get("processes") + if isinstance(plist, list): + gpu_idx = item.get("gpu") + if not isinstance(gpu_idx, int): + gpu_idx = -1 + # Register the GPU bucket as soon as we recognize a Shape A + # entry, even if its process_list is empty. This makes an + # empty return value mean "schema didn't match" rather than + # the ambiguous "schema matched but no processes" -- so the + # caller can safely fall through to the text / rocm-smi / + # lsof fallbacks on a future amd-smi schema. + out_by_gpu.setdefault(int(gpu_idx), []) + for p in plist: + if not isinstance(p, dict): + continue + proc = _unwrap_proc(p) + _push( + gpu_idx, + proc.get("pid") if proc.get("pid") is not None else proc.get("process_id"), + proc.get("name") or proc.get("process_name"), + _hbm_of(proc), + ) + continue + # Shape B: per-process dict with explicit gpu + proc = _unwrap_proc(item) + if "pid" in proc or "process_id" in proc: + g = proc.get("gpu") + gpus = ( + proc.get("gpus") + if isinstance(proc.get("gpus"), list) + else ([g] if isinstance(g, int) else [-1]) + ) + for gpu_idx in gpus: + if not isinstance(gpu_idx, int): + gpu_idx = -1 + _push( + gpu_idx, + proc.get("pid") if proc.get("pid") is not None else proc.get("process_id"), + proc.get("name") or proc.get("process_name"), + _hbm_of(proc), + ) + + return ( + [{"gpu": k, "processes": v} for k, v in sorted(out_by_gpu.items())], + schema_drift_suspected, + ) + + +def _parse_amd_smi_process_text(text: str, annotate: Any) -> List[Dict[str, Any]]: + """Best-effort parser for ``amd-smi process`` plain-text output. + + Format varies, but typical layout is one record per process with lines + like ``GPU: 0`` / ``PID: 12345`` / ``NAME: python`` / ``VRAM_MEM: 256 MB``. + We tokenise key-value pairs case-insensitively and group records by + blank-line separators. + """ + out_by_gpu: Dict[int, List[Dict[str, Any]]] = {} + cur: Dict[str, Any] = {} + blocks: List[Dict[str, Any]] = [] + for raw in text.splitlines(): + line = raw.strip() + if not line: + if cur: + blocks.append(cur) + cur = {} + continue + if ":" in line: + k, _, v = line.partition(":") + cur[k.strip().lower()] = v.strip() + if cur: + blocks.append(cur) + for b in blocks: + gpu_s = b.get("gpu") or b.get("gpu_id") or b.get("device") or "-1" + try: + gpu = int(gpu_s.split()[0]) + except Exception: + gpu = -1 + try: + pid = int(b.get("pid", "").split()[0]) + except Exception: + continue + name = b.get("name") or b.get("process") or b.get("process_name") or "" + hbm_raw = b.get("vram_mem") or b.get("vram") or b.get("memory_usage") or b.get("mem_usage") or "" + hbm = _parse_size_with_unit(hbm_raw) if hbm_raw else None + out_by_gpu.setdefault(gpu, []).append(annotate(pid, name, hbm)) + return [{"gpu": k, "processes": v} for k, v in sorted(out_by_gpu.items())] + + +def _parse_lsof_pcn(text: str, annotate: Any) -> List[Dict[str, Any]]: + """Parse ``lsof -Fpcn`` output (one field per line, prefixed by f-code). + + We only need ``p`` and ``c``. Returns one annotated + record per unique PID; HBM bytes unknown (lsof can't measure that). + """ + out: Dict[int, Dict[str, Any]] = {} + cur_pid: Optional[int] = None + cur_name = "" + for line in text.splitlines(): + if not line: + continue + tag, val = line[0], line[1:] + if tag == "p": + try: + cur_pid = int(val) + except ValueError: + cur_pid = None + cur_name = "" + elif tag == "c" and cur_pid is not None: + cur_name = val + out[cur_pid] = annotate(cur_pid, cur_name, None) + return list(out.values()) diff --git a/primus/tools/preflight/node_smoke/collectors/host_limits.py b/primus/tools/preflight/node_smoke/collectors/host_limits.py new file mode 100644 index 000000000..d0e9a0e5c --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/host_limits.py @@ -0,0 +1,92 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 1 -- C. Host limits (ulimit -l, /dev/shm, NUMA, CPU governor).""" + +from __future__ import annotations + +import os +from typing import Any, Dict, List + +from ..shell_utils import _read_text + + +def _collect_host_limits(*, ulimit_l_min_gb: float, shm_min_gb: float) -> Dict[str, Any]: + """Capture training-relevant kernel/process limits and tunables and + return hard-failure reasons for the ones that block training under load. + + Hard fail today (cause node FAIL): + + * ``ulimit -l`` (RLIMIT_MEMLOCK) is not unlimited and below + ``ulimit_l_min_gb`` -- RDMA pin failures look like NCCL hangs. + * ``/dev/shm`` total size below ``shm_min_gb`` -- NCCL shared-memory + transport falls back or fails. + + Soft (collected for drift detection only): + + * NUMA node count, CPU count, CPU governor, kernel/OS version. The + aggregator flags drift across the cluster but does not FAIL nodes + individually for these. + """ + out: Dict[str, Any] = {} + + # Resource limits. + try: + import resource # type: ignore + + soft_l, _ = resource.getrlimit(resource.RLIMIT_MEMLOCK) + out["memlock_soft_bytes"] = -1 if soft_l == resource.RLIM_INFINITY else int(soft_l) + soft_n, _ = resource.getrlimit(resource.RLIMIT_NOFILE) + out["nofile_soft"] = int(soft_n) + soft_p, _ = resource.getrlimit(resource.RLIMIT_NPROC) + out["nproc_soft"] = -1 if soft_p == resource.RLIM_INFINITY else int(soft_p) + except Exception as e: + out["resource_error"] = str(e) + + # /dev/shm size + free. + try: + st = os.statvfs("/dev/shm") + out["shm_size_bytes"] = int(st.f_blocks) * int(st.f_frsize) + out["shm_avail_bytes"] = int(st.f_bavail) * int(st.f_frsize) + except Exception as e: + out["shm_error"] = str(e) + + # NUMA topology. + try: + nodes = [ + n for n in os.listdir("/sys/devices/system/node") if n.startswith("node") and n[4:].isdigit() + ] + out["numa_nodes"] = len(nodes) + except Exception: + out["numa_nodes"] = None + + # CPU count + governor. + try: + out["cpu_count"] = os.cpu_count() + except Exception: + out["cpu_count"] = None + out["cpu_governor"] = _read_text("/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor") or None + + # Hard checks. + fail_reasons: List[str] = [] + memlock = out.get("memlock_soft_bytes") + if memlock is not None and memlock != -1 and ulimit_l_min_gb > 0: + if memlock < ulimit_l_min_gb * (1 << 30): + fail_reasons.append( + f"ulimit -l (memlock) = {memlock // (1 << 20)} MiB; " + f"required: unlimited or >= {ulimit_l_min_gb} GiB. " + "RDMA pin will fail under load." + ) + shm = out.get("shm_size_bytes") + if shm is not None and shm_min_gb > 0: + if shm < shm_min_gb * (1 << 30): + fail_reasons.append( + f"/dev/shm size = {shm / (1 << 30):.2f} GiB; " + f"required: >= {shm_min_gb} GiB. NCCL shared-mem may fail." + ) + out["fail_reasons"] = fail_reasons + + return out diff --git a/primus/tools/preflight/node_smoke/collectors/nics.py b/primus/tools/preflight/node_smoke/collectors/nics.py new file mode 100644 index 000000000..8266b18e6 --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/nics.py @@ -0,0 +1,452 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 1 -- B. NIC / RDMA roll-call (per-port state + GIDs from sysfs). + +Selector chain +============== + +Many clusters expose more RDMA-capable ports in ``/sys/class/infiniband/`` +than the training job will actually use. Common reasons: + +* a separate front-end / management RoCE NIC that is physically present + but admin-disabled (no SFP, BIOS port-disable, netdev down) -- typical + ``state=DOWN phys_state=Disabled`` ports; +* a storage / control-plane RoCE NIC that is fully up but is reserved + for sockets (it shows up in ``NCCL_SOCKET_IFNAME`` rather than + ``NCCL_IB_HCA``). + +Failing the node on those ports is wrong: the operator already told +NCCL/RCCL not to use them, and the training job will run fine. So the +collector resolves a *training-NIC selector* and only enforces the hard +rules (state must be ACTIVE, phys_state must be LinkUp, RoCE v2 GID +present, ...) on the included subset. Excluded ports stay in +``ports`` for diagnostics and are summarised in ``info_issues``. + +Selector precedence (highest first): + +1. ``allowlist`` argument -- typically wired to ``--rdma-nic-allowlist`` + on the CLI. NCCL syntax (see :func:`_parse_nic_selector` below). +2. ``NCCL_IB_HCA`` env -- mirrors what NCCL/RCCL itself will use, so the + smoke test and the training launch agree by construction. +3. ``phys_state in {Disabled, Sleep}`` heuristic -- ports in those + administratively-down phys_states are intentionally not in use + (admin-disabled at firmware/driver level, no SFP, BIOS-disabled, ...) + and are auto-excluded. Crucially, real failure modes on a port that + *is* intended to be used produce a different phys_state + (``Polling``, ``LinkErrorRecovery``, ``LinkUp`` with ``state!=ACTIVE``, + ...), so this heuristic does not mask cable / driver problems. +4. Fallback: every IB port must be ACTIVE / LinkUp (the strict rule). + +Empty-set guard: if the resolved included set is empty (every port was +excluded by some combination of the above) but the host has at least +one IB port at all, that's still a node FAIL -- a node with zero +training NICs cannot participate in inter-node training. +""" + +from __future__ import annotations + +import os +from typing import Any, Dict, List, Optional, Set, Tuple + +from ..shell_utils import _read_text + +# Phys-states we treat as "administratively off, not a failure". +# See module docstring for why this set is conservative. +_ADMIN_DOWN_PHYS_STATES: Set[str] = {"disabled", "sleep"} + + +def _parse_nic_selector(raw: str) -> Dict[str, Any]: + """Parse an ``NCCL_IB_HCA``-style selector string. + + Accepts the full NCCL syntax: + + * Comma-separated entries of the form ``device[:port]``. ``port`` + is an integer; if omitted, the entry matches any port on the + device. + * Optional ``^`` prefix on the *whole* string makes it a denylist + (every device matches except the listed ones). + * Optional ``=`` prefix on an individual entry forces an exact + device-name match (e.g. ``=mlx5`` matches device ``mlx5`` only, + not ``mlx5_0``). Without ``=``, an entry is a *prefix* match, + matching ``mlx5`` against ``mlx5_0`` / ``mlx5_1`` / ... + + Returns a dict with: + + * ``mode``: ``"allowlist"`` or ``"denylist"``. + * ``entries``: list of ``(device_pattern, exact_match, port_or_None)``. + * ``raw``: the original input, after stripping the global ``^``. + + Returns ``{"mode": "passthrough", ...}`` for empty / whitespace-only + input, signalling the caller to fall through to the next selector. + """ + s = (raw or "").strip() + if not s: + return {"mode": "passthrough", "entries": [], "raw": raw or ""} + + mode = "allowlist" + if s.startswith("^"): + mode = "denylist" + s = s[1:].strip() + + entries: List[Tuple[str, bool, Optional[int]]] = [] + for token in s.split(","): + t = token.strip() + if not t: + continue + exact = False + if t.startswith("="): + exact = True + t = t[1:].strip() + if not t: + continue + # Split off an optional :port suffix. We split from the right so + # device names containing ':' (none in the wild today, but we + # don't want to be the thing that breaks if they appear) are + # preserved. + port: Optional[int] = None + if ":" in t: + dev_part, _, port_str = t.rpartition(":") + try: + port = int(port_str) + t = dev_part + except ValueError: + # Not actually a port suffix -- treat the whole thing + # as the device name. + port = None + entries.append((t, exact, port)) + + if not entries: + return {"mode": "passthrough", "entries": [], "raw": raw or ""} + return {"mode": mode, "entries": entries, "raw": raw or ""} + + +def _selector_matches( + sel: Dict[str, Any], + device: str, + port: int, +) -> bool: + """Return True iff ``(device, port)`` matches an allow/denylist selector + parsed by :func:`_parse_nic_selector`. + + ``passthrough`` selectors trivially match everything (the caller + should not normally pass them here; instead it should treat them as + "no selector at all" and fall through to the next layer). + """ + mode = sel.get("mode", "passthrough") + if mode == "passthrough": + return True + entries = sel.get("entries") or [] + listed = False + for dev_pat, exact, want_port in entries: + # Device name match: exact iff the entry started with '='; else + # prefix match (NCCL semantics). + if exact: + dev_ok = device == dev_pat + else: + dev_ok = device.startswith(dev_pat) + if not dev_ok: + continue + # Port filter: an entry with no :port matches any port on that + # device. + if want_port is not None and want_port != port: + continue + listed = True + break + if mode == "allowlist": + return listed + # Denylist: ports NOT in the list are accepted. + return not listed + + +def _resolve_selector( + allowlist_arg: Optional[str], + env: Optional[Dict[str, str]] = None, +) -> Dict[str, Any]: + """Resolve the training-NIC selector for this node. + + Precedence: + + 1. ``allowlist_arg`` (typically the value of ``--rdma-nic-allowlist``). + 2. ``NCCL_IB_HCA`` env var. + 3. Heuristic fallback (``"heuristic"``) -- the caller handles this + by auto-excluding ports whose ``phys_state`` is in + :data:`_ADMIN_DOWN_PHYS_STATES`. + + Returns a dict with ``source`` (``"cli"`` / ``"env"`` / ``"heuristic"``) + plus the parsed-selector fields when applicable. + """ + if env is None: + env = dict(os.environ) + # 1. CLI / explicit arg. + if allowlist_arg: + parsed = _parse_nic_selector(allowlist_arg) + if parsed.get("mode") != "passthrough": + parsed["source"] = "cli" + return parsed + # 2. NCCL_IB_HCA env. + env_raw = env.get("NCCL_IB_HCA") or "" + if env_raw.strip(): + parsed = _parse_nic_selector(env_raw) + if parsed.get("mode") != "passthrough": + parsed["source"] = "env" + return parsed + # 3. Heuristic fallback. + return { + "source": "heuristic", + "mode": "heuristic", + "entries": [], + "raw": "", + "admin_down_phys_states": sorted(_ADMIN_DOWN_PHYS_STATES), + } + + +def _collect_nic_status( + expected_count: Optional[int], + *, + allowlist: Optional[str] = None, +) -> Dict[str, Any]: + """Inventory every RDMA port on this node and flag the ones that would + silently break inter-node training. + + Reads everything from ``/sys/class/infiniband`` (kernel ``ib_core`` + ABI) so the check is **vendor- and fabric-agnostic** -- works on + Mellanox/NVIDIA (``mlx5_ib``), Broadcom (``bnxt_re``), Intel + (``irdma``), Marvell (``qedr``), AWS EFA (``efa``), Huawei + (``hns_roce``), etc., and on either RoCE-over-Ethernet or true + InfiniBand fabrics. We do not depend on ``ibv_devinfo`` / + ``ibstat`` / vendor SDKs being present in the container. + + Per port we capture: + + * ``link_layer`` (``Ethernet`` for RoCE, ``InfiniBand`` for IB) -- + determines which GID rule applies below; + * link state (``state``: ``ACTIVE``/``DOWN``/``INIT``) and physical + state (``phys_state``: ``LinkUp``/``Polling``/...); + * link rate (Gb/s); + * netdev + MTU (so the aggregator can detect MTU drift, which silently + tanks RoCE all-reduce throughput); + * GID counts -- total non-zero GIDs and the subset configured as + ``RoCE v2`` (an empty RoCE v2 set is a frequent cause of training + jobs hanging at the first inter-node collective on RoCE clusters). + + Per-port hard issues are only emitted for ports in the resolved + *training-NIC selector* (see module docstring). Excluded ports stay + in ``ports`` for diagnostics and are summarised in ``info_issues``. + + Issues are pushed into ``out["issues"]`` (each a short string). Hard + issues (port not Active / missing GIDs / wrong NIC count) are + treated as node FAIL by ``_node_status_from``. The GID check is + fabric-aware: + + * RoCE/Ethernet port must have at least one ``RoCE v2`` GID + configured (RoCEv1 is essentially unused for AI training). + * InfiniBand port must have at least one valid (non-zero) GID -- + it is normally auto-populated by the SM; an empty GID table on + an ACTIVE IB port indicates a subnet-manager problem. + * Unknown ``link_layer`` (very old kernels) falls back to "must + have any valid GID" so we don't false-FAIL exotic configurations. + """ + selector = _resolve_selector(allowlist) + out: Dict[str, Any] = { + "expected_count": expected_count, + "selector": selector, + "ports": [], + "included_ports": [], + "excluded_ports": [], + "issues": [], + "info_issues": [], + } + base = "/sys/class/infiniband" + if not os.path.isdir(base): + # Container may not expose the IB stack; report and let the operator + # decide. We only mark this as a hard issue when the user explicitly + # asked for a positive expected_count. + msg = f"{base} missing -- no RDMA stack visible" + if expected_count and expected_count > 0: + out["issues"].append(msg) + else: + out["info"] = msg + return out + + try: + devs = sorted(os.listdir(base)) + except Exception as e: + out["issues"].append(f"failed to list {base}: {e}") + return out + + for dev in devs: + port_dir = os.path.join(base, dev, "ports") + if not os.path.isdir(port_dir): + continue + try: + ports = sorted(os.listdir(port_dir)) + except Exception: + continue + for port_str in ports: + try: + port = int(port_str) + except ValueError: + continue + p = os.path.join(port_dir, port_str) + + # Sysfs values look like "4: ACTIVE" / "5: LinkUp" / "400 Gb/sec (4X NDR)" + state_raw = _read_text(os.path.join(p, "state")) + phys_raw = _read_text(os.path.join(p, "phys_state")) + rate_raw = _read_text(os.path.join(p, "rate")) + state = state_raw.split(":", 1)[-1].strip() if state_raw else "" + phys = phys_raw.split(":", 1)[-1].strip() if phys_raw else "" + rate_gbps: Optional[int] = None + try: + rate_gbps = int(rate_raw.split()[0]) + except Exception: + pass + + # Fabric type: "Ethernet" -> RoCE, "InfiniBand" -> IB. + # Determines which GID rule to apply below. Provided by + # ib_core for every RDMA driver since Linux 3.x; missing + # only on extremely old kernels. + link_layer = (_read_text(os.path.join(p, "link_layer")) or "").strip() or None + + # GID inventory. A GID is "all-zero" until configured. + gid_count = 0 + rocev2_count = 0 + gids_dir = os.path.join(p, "gids") + types_dir = os.path.join(p, "gid_attrs", "types") + valid_gid_indices: List[int] = [] + if os.path.isdir(gids_dir): + try: + for gn in sorted(os.listdir(gids_dir), key=lambda s: int(s) if s.isdigit() else 0): + if not gn.isdigit(): + continue + g = _read_text(os.path.join(gids_dir, gn)) + if g and g != "0000:0000:0000:0000:0000:0000:0000:0000": + gid_count += 1 + valid_gid_indices.append(int(gn)) + except Exception: + pass + if os.path.isdir(types_dir): + for idx in valid_gid_indices: + t = _read_text(os.path.join(types_dir, str(idx))) + if "RoCE v2" in t or "RoCEv2" in t: + rocev2_count += 1 + + # Linked netdev + MTU. + ifname: Optional[str] = None + mtu: Optional[int] = None + net_dir = os.path.join(base, dev, "device", "net") + if os.path.isdir(net_dir): + try: + nets = sorted(os.listdir(net_dir)) + if nets: + ifname = nets[0] + mtu_raw = _read_text(f"/sys/class/net/{ifname}/mtu") + try: + mtu = int(mtu_raw) + except Exception: + mtu = None + except Exception: + pass + + port_rec = { + "device": dev, + "port": port, + "link_layer": link_layer, + "state": state or None, + "phys_state": phys or None, + "rate_gbps": rate_gbps, + "ifname": ifname, + "mtu": mtu, + "gid_count": gid_count, + "rocev2_gid_count": rocev2_count, + } + out["ports"].append(port_rec) + + # ------------------------------------------------------------ + # Selector: is this port part of the training-NIC set? + # ------------------------------------------------------------ + label = f"{dev}:{port}" + include = True + exclude_reason = "" + if selector.get("source") == "heuristic": + # Auto-exclude admin-down phys_states; keep everything + # else and let the strict per-port rules run. + if (phys or "").lower() in _ADMIN_DOWN_PHYS_STATES: + include = False + exclude_reason = f"phys_state={phys} (admin-disabled, not used for training)" + else: + # Explicit allow/denylist (CLI or env). + if not _selector_matches(selector, dev, port): + include = False + src = selector.get("source") or "selector" + src_label = ( + "NCCL_IB_HCA" if src == "env" else "--rdma-nic-allowlist" if src == "cli" else src + ) + extra = "" + if (state or "").upper() != "ACTIVE" or (phys or "").lower() != "linkup": + extra = f" (state={state} phys_state={phys})" + exclude_reason = f"excluded by {src_label}{extra}" + + if not include: + out["excluded_ports"].append(label) + out["info_issues"].append(f"{label} {exclude_reason}") + continue + out["included_ports"].append(label) + + # ------------------------------------------------------------ + # Per-port hard issues -> node FAIL. Only for included ports. + # ------------------------------------------------------------ + if state and state.upper() != "ACTIVE": + out["issues"].append(f"{dev}:{port} state={state} (expected ACTIVE)") + if phys and phys.upper() != "LINKUP": + out["issues"].append(f"{dev}:{port} phys_state={phys} (expected LinkUp)") + # Fabric-aware GID requirement on ACTIVE ports. + if state.upper() == "ACTIVE": + ll = (link_layer or "").lower() + if ll == "ethernet": + # RoCE: needs at least one RoCE v2 GID; RoCEv1 is + # essentially unused for AI training and we treat + # its absence as a hard fail. + if rocev2_count == 0: + out["issues"].append( + f"{dev}:{port} no RoCE v2 GIDs configured " f"(RoCE/Ethernet fabric)" + ) + elif ll == "infiniband": + # True IB: subnet manager normally populates GIDs. + # Empty GID table on an ACTIVE IB port = SM problem. + if gid_count == 0: + out["issues"].append( + f"{dev}:{port} no GIDs populated " f"(InfiniBand fabric -- check subnet manager)" + ) + else: + # Unknown / missing link_layer (very old kernel): + # require at least one valid GID rather than a + # specific type, to avoid false FAIL. + if gid_count == 0: + out["issues"].append( + f"{dev}:{port} no valid GIDs configured " f"(link_layer unknown)" + ) + + # The expected-count check operates on the *included* set: operators + # almost always set --expected-rdma-nics to mean "training NIC count" + # (e.g. 8 = 1 NIC per GPU), not "everything visible under + # /sys/class/infiniband/" -- which on multi-role nodes includes + # frontend / management / storage RoCE NICs. + if expected_count is not None and len(out["included_ports"]) != expected_count: + out["issues"].append(f"RDMA NIC port count {len(out['included_ports'])} != expected {expected_count}") + + # Empty-set guard: if we saw IB ports on this host but every single + # one got excluded, the node cannot participate in inter-node + # training. This catches "the whole RoCE card disabled itself" and + # similar disasters that would otherwise PASS silently because every + # port met the exclusion criteria. + if out["ports"] and not out["included_ports"]: + out["issues"].append( + "no included RDMA NIC ports remain after selector " + f"({selector.get('source')}); node cannot participate in inter-node training" + ) + + return out diff --git a/primus/tools/preflight/node_smoke/collectors/reused_info.py b/primus/tools/preflight/node_smoke/collectors/reused_info.py new file mode 100644 index 000000000..ef0f3236e --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/reused_info.py @@ -0,0 +1,51 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Reused gpu/host/network info collectors from the rest of the preflight tree. + +These collectors already work without a global PG and produce ``Finding`` +objects (level='fail' counts as a node failure). We import each one +lazily so a missing dependency in one section doesn't cascade. +""" + +from __future__ import annotations + +from typing import Any, Dict + +from ..shell_utils import _findings_to_dicts + + +def _collect_reused_info() -> Dict[str, Any]: + """Run the existing host/gpu/network info collectors. They already work + without a global PG and produce ``Finding`` objects (level='fail' counts + as a node failure).""" + section: Dict[str, Any] = {"gpu_info": [], "host_info": [], "network_info": []} + try: + from primus.tools.preflight.gpu.info import collect_gpu_info + + section["gpu_info"] = _findings_to_dicts(collect_gpu_info()) + except Exception as e: + section["gpu_info"] = [ + {"level": "warn", "message": "collect_gpu_info raised", "details": {"error": str(e)}} + ] + try: + from primus.tools.preflight.host.info import collect_host_info + + section["host_info"] = _findings_to_dicts(collect_host_info()) + except Exception as e: + section["host_info"] = [ + {"level": "warn", "message": "collect_host_info raised", "details": {"error": str(e)}} + ] + try: + from primus.tools.preflight.network.info import collect_network_info + + # expect_distributed=False so we don't WARN about a missing world PG. + section["network_info"] = _findings_to_dicts(collect_network_info(expect_distributed=False)) + except Exception as e: + section["network_info"] = [ + {"level": "warn", "message": "collect_network_info raised", "details": {"error": str(e)}} + ] + return section diff --git a/primus/tools/preflight/node_smoke/collectors/rocm_smi.py b/primus/tools/preflight/node_smoke/collectors/rocm_smi.py new file mode 100644 index 000000000..a5b301dab --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/rocm_smi.py @@ -0,0 +1,409 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 1 -- F: rocm-smi self-latency + cross-tool fallbacks for amd-smi checks. + +rocm-smi is "usually available" even on nodes that lack amd-smi (older +ROCm installs, stripped-down containers). Every amd-smi check we +silently no-op when amd-smi is missing has a rocm-smi equivalent, so +we can keep ECC / XGMI / foreign-process / activity coverage even +without amd-smi. Each helper produces the same record shape the +upstream amd-smi parser already emits, so _node_status_from and the +aggregator helpers don't need any per-tool conditionals. + +Output schemas (cross-tool stable): + + ECC -> per-GPU {gpu, ecc_correctable_total, ecc_uncorrectable_total} + XGMI -> {ok, tool, n_gpus, link_types: [[...]], non_xgmi_pairs} + procs -> per-GPU {gpu, processes: [annotated PID dicts]} + activity -> per-GPU {gpu, gfx_activity_pct} + +The self-latency canary (``_collect_rocm_smi_self_latency``) is its own +Tier 1 F check: a wedged amdgpu driver makes ``rocm-smi`` calls take +30-60 s before failing outright -- usually minutes before the GPU +itself stops responding. Hitting the timeout is treated as a hard fail +in ``_node_status_from``. +""" + +from __future__ import annotations + +import json +import subprocess +import time +from typing import Any, Dict, List, Optional + +from ..shell_utils import _which + + +def _collect_rocm_smi_self_latency(*, timeout_sec: float) -> Dict[str, Any]: + """Time a single ``rocm-smi --version`` call against ``timeout_sec``. + + A wedged amdgpu driver makes ``rocm-smi`` calls take 30-60 s before + failing outright -- usually minutes before the GPU itself stops + responding. Catching this in preflight gives operators a chance to + drain the node before a real training job starts hanging on it. + """ + out: Dict[str, Any] = { + "ok": False, + "tool": None, + "latency_sec": None, + "timeout_sec": float(timeout_sec), + } + binpath = _which("rocm-smi") + if binpath is None: + out["error"] = "rocm-smi not found in PATH" + return out + out["tool"] = binpath + t0 = time.monotonic() + try: + cp = subprocess.run( + [binpath, "--version"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=timeout_sec, + check=False, + ) + out["latency_sec"] = round(time.monotonic() - t0, 3) + out["rc"] = cp.returncode + out["ok"] = cp.returncode == 0 + if cp.returncode != 0: + out["error"] = (cp.stderr or cp.stdout or "").strip()[:200] + except subprocess.TimeoutExpired: + out["latency_sec"] = round(time.monotonic() - t0, 3) + out["timed_out"] = True + out["error"] = f"rocm-smi --version did not finish in {timeout_sec}s -- driver may be wedging" + except Exception as e: + out["error"] = str(e) + return out + + +def _rocm_smi_ras_info_text(timeout_sec: float = 15.0) -> Dict[str, Any]: + """ECC counts via ``rocm-smi --showrasinfo`` (TEXT only). + + The ``--json`` form returns "WARNING: No JSON data to report" so we + parse the text. Format is one block per GPU:: + + GPU[0]: RAS INFO + Block Status Correctable Error Uncorrectable Error + UMC ENABLED 0 0 + SDMA ENABLED 0 0 + GFX ENABLED 0 0 + ... + + Some blocks (ATHUB, PCIE_BIF, HDP, ...) report only Status. We sum + every numeric Correctable/Uncorrectable cell per GPU so a hardware + error in any block surfaces in ``ecc_uncorrectable_total``. + """ + out: Dict[str, Any] = {"ok": False, "tool": None, "per_gpu": []} + if _which("rocm-smi") is None: + out["error"] = "rocm-smi not found in PATH" + return out + try: + cp = subprocess.run( + ["rocm-smi", "--showrasinfo"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=timeout_sec, + check=False, + ) + if cp.returncode != 0: + out["error"] = (cp.stderr or cp.stdout or "").strip()[:200] or f"rc={cp.returncode}" + return out + out["per_gpu"] = _parse_rocm_smi_ras_info_text(cp.stdout) + out["tool"] = "rocm-smi --showrasinfo" + out["ok"] = True + except subprocess.TimeoutExpired: + out["error"] = "rocm-smi --showrasinfo timed out" + except Exception as e: + out["error"] = str(e) + return out + + +def _parse_rocm_smi_ras_info_text(text: str) -> List[Dict[str, Any]]: + """Parse the per-GPU ``GPU[N]: RAS INFO`` blocks from rocm-smi text.""" + import re + + out: List[Dict[str, Any]] = [] + cur_gpu: Optional[int] = None + cur_corr = 0 + cur_uncorr = 0 + in_block = False + gpu_hdr = re.compile(r"^GPU\[(\d+)\]:\s*RAS INFO", re.IGNORECASE) + for raw in text.splitlines(): + line = raw.strip() + m = gpu_hdr.match(line) + if m: + # Flush previous GPU + if cur_gpu is not None: + out.append( + { + "gpu": cur_gpu, + "ecc_correctable_total": cur_corr, + "ecc_uncorrectable_total": cur_uncorr, + } + ) + cur_gpu = int(m.group(1)) + cur_corr = 0 + cur_uncorr = 0 + in_block = True + continue + if not in_block or cur_gpu is None: + continue + # End-of-block separator + if line.startswith("__") or line.startswith("=="): + continue + # Per-block row: "BLOCK STATUS CORR UNCORR" -- last two are ints + # when present. Header row has the words "Correctable Error" so + # we filter out non-numeric rows naturally. + toks = line.split() + if len(toks) < 4: + continue + try: + corr = int(toks[-2]) + uncorr = int(toks[-1]) + except ValueError: + continue + cur_corr += corr + cur_uncorr += uncorr + if cur_gpu is not None: + out.append( + { + "gpu": cur_gpu, + "ecc_correctable_total": cur_corr, + "ecc_uncorrectable_total": cur_uncorr, + } + ) + return out + + +def _rocm_smi_topotype_json(timeout_sec: float = 15.0) -> Dict[str, Any]: + """XGMI link-type matrix via ``rocm-smi --showtopotype --json``. + + Output shape is keyed by pair-string:: + + {"system": { + "(Topology) Link type between DRM devices 0 and 1": "XGMI", + "(Topology) Link type between DRM devices 0 and 2": "XGMI", + ... # upper triangle only + }} + + We parse the indices out of each key with a regex, build a symmetric + NxN matrix, and emit it in the same shape ``_collect_xgmi_topology`` + already produces (so downstream consumers don't care which tool + populated it). ``non_xgmi_pairs`` contains the (i, j, link_type) + triples for any off-diagonal pair whose link_type is not XGMI. + """ + import re + + out: Dict[str, Any] = {"ok": False, "tool": None, "n_gpus": 0, "link_types": [], "non_xgmi_pairs": []} + if _which("rocm-smi") is None: + out["error"] = "rocm-smi not found in PATH" + return out + try: + cp = subprocess.run( + ["rocm-smi", "--showtopotype", "--json"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=timeout_sec, + check=False, + ) + if cp.returncode != 0 or not cp.stdout.strip(): + out["error"] = (cp.stderr or "").strip()[:200] or f"rc={cp.returncode}" + return out + try: + doc = json.loads(cp.stdout) + except Exception as e: + out["error"] = f"json parse failed: {e}" + return out + sys_block = doc.get("system") if isinstance(doc, dict) else None + if not isinstance(sys_block, dict): + out["error"] = "no `system` key in rocm-smi --showtopotype output" + return out + pat = re.compile(r"DRM\s+devices?\s+(\d+)\s+and\s+(\d+)", re.IGNORECASE) + pairs: Dict[tuple, str] = {} + max_idx = -1 + for k, v in sys_block.items(): + m = pat.search(str(k)) + if not m: + continue + i = int(m.group(1)) + j = int(m.group(2)) + pairs[(i, j)] = str(v) + max_idx = max(max_idx, i, j) + if max_idx < 0: + out["error"] = "no parseable DRM-device pair keys" + return out + n = max_idx + 1 + mat = [["?" for _ in range(n)] for _ in range(n)] + for (i, j), t in pairs.items(): + mat[i][j] = t + mat[j][i] = t + for i in range(n): + mat[i][i] = "SELF" + non_xgmi: List[Any] = [] + for i in range(n): + for j in range(i + 1, n): + t = mat[i][j] + if t and t.upper() != "XGMI": + non_xgmi.append((i, j, t)) + out["n_gpus"] = n + out["link_types"] = mat + out["non_xgmi_pairs"] = non_xgmi + out["tool"] = "rocm-smi --showtopotype --json" + out["ok"] = True + except subprocess.TimeoutExpired: + out["error"] = "rocm-smi --showtopotype timed out" + except Exception as e: + out["error"] = str(e) + return out + + +def _rocm_smi_processes( + annotate: Any, + timeout_sec: float = 15.0, +) -> Dict[str, Any]: + """Foreign-process enumeration via ``rocm-smi --showpids --json``. + + Output shape (verified against rocm-smi on a busy MI300X):: + + {"system": { + "PID2683309": "python3.11, 1, 24556904448, 0, 0", + "PID29324": "gpuagent, 0, 0, 0, 0" + }} + + Comma-separated fields are: ``name, num_gpus_used, vram_bytes, + sdma_bytes, cu_occupancy``. We extract ``name`` (field 0) and + ``vram_bytes`` (field 2) -- enough to surface leaked training PIDs + holding gigabytes of HBM. ``--showpidgpus --json`` returns + "No JSON data to report", so per-GPU mapping is not available -- + all PIDs go into the gpu=-1 bucket (same convention as lsof). + """ + out: Dict[str, Any] = {"ok": False, "tool": None, "per_gpu": []} + if _which("rocm-smi") is None: + out["error"] = "rocm-smi not found in PATH" + return out + try: + cp = subprocess.run( + ["rocm-smi", "--showpids", "--json"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=timeout_sec, + check=False, + ) + if cp.returncode != 0 or not cp.stdout.strip(): + out["error"] = (cp.stderr or "").strip()[:200] or f"rc={cp.returncode}" + return out + try: + doc = json.loads(cp.stdout) + except Exception as e: + out["error"] = f"json parse failed: {e}" + return out + sys_block = doc.get("system") if isinstance(doc, dict) else None + if not isinstance(sys_block, dict): + # Empty or unexpected -- treat as "no PIDs" + out["tool"] = "rocm-smi --showpids --json" + out["ok"] = True + return out + procs: List[Dict[str, Any]] = [] + for k, v in sys_block.items(): + try: + pid = int(str(k).lstrip("PID").lstrip("pid")) + except ValueError: + continue + name = "" + hbm: Optional[int] = None + if isinstance(v, str): + parts = [s.strip() for s in v.split(",")] + if parts: + name = parts[0] + # field 2 = VRAM bytes (rocm-smi --showpids documented format) + if len(parts) > 2: + try: + hbm = int(parts[2]) + except ValueError: + hbm = None + elif isinstance(v, dict): + name = str(v.get("name") or v.get("process_name") or "") + vram = v.get("vram") or v.get("vram_bytes") + if isinstance(vram, (int, float)): + hbm = int(vram) + elif isinstance(v, list) and v: + name = str(v[0]) + procs.append(annotate(pid, name, hbm)) + if procs: + out["per_gpu"] = [{"gpu": -1, "processes": procs}] + out["tool"] = "rocm-smi --showpids --json" + out["ok"] = True + except subprocess.TimeoutExpired: + out["error"] = "rocm-smi --showpids timed out" + except Exception as e: + out["error"] = str(e) + return out + + +def _rocm_smi_use_json(timeout_sec: float = 15.0) -> Dict[str, Any]: + """GPU compute-activity % via ``rocm-smi --showuse --json``. + + Output shape:: + + {"card0": {"GPU use (%)": "0", "GFX Activity": "465307149"}, ...} + + "GPU use (%)" is the percentage we want for ``gfx_activity_pct`` + (matching the amd-smi field name). "GFX Activity" is a cumulative + cycle counter, not a percentage -- ignored. + """ + out: Dict[str, Any] = {"ok": False, "tool": None, "per_gpu": []} + if _which("rocm-smi") is None: + out["error"] = "rocm-smi not found in PATH" + return out + try: + cp = subprocess.run( + ["rocm-smi", "--showuse", "--json"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=timeout_sec, + check=False, + ) + if cp.returncode != 0 or not cp.stdout.strip(): + out["error"] = (cp.stderr or "").strip()[:200] or f"rc={cp.returncode}" + return out + try: + doc = json.loads(cp.stdout) + except Exception as e: + out["error"] = f"json parse failed: {e}" + return out + if not isinstance(doc, dict): + out["error"] = "unexpected top-level shape" + return out + per_gpu: List[Dict[str, Any]] = [] + for k, v in doc.items(): + if not isinstance(v, dict): + continue + if not str(k).startswith("card"): + continue + try: + gpu = int(str(k)[len("card") :]) + except ValueError: + continue + pct_raw = v.get("GPU use (%)") or v.get("GPU use") + try: + pct = float(str(pct_raw).strip()) + except (TypeError, ValueError): + continue + per_gpu.append({"gpu": gpu, "gfx_activity_pct": pct}) + out["per_gpu"] = per_gpu + out["tool"] = "rocm-smi --showuse --json" + out["ok"] = True + except subprocess.TimeoutExpired: + out["error"] = "rocm-smi --showuse timed out" + except Exception as e: + out["error"] = str(e) + return out diff --git a/primus/tools/preflight/node_smoke/collectors/tooling.py b/primus/tools/preflight/node_smoke/collectors/tooling.py new file mode 100644 index 000000000..811ee3f50 --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/tooling.py @@ -0,0 +1,100 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 1 -- tooling availability inventory. + +Several Tier 1 checks (ECC via amd-smi metric, XGMI via amd-smi topology, +foreign-process enumeration via amd-smi process / lsof, wedged-driver +canary via rocm-smi --version) are best-effort: each collector returns +ok=False and the FAIL rules in _node_status_from then iterate over empty +data, silently no-op'ing. That's fine on a node where the tools are +legitimately absent (containers stripped down for size), but DANGEROUS +in a production cluster -- a node with ECC errors, broken XGMI, leaked +ranks, or a stale ROCm install can PASS smoke just because amd-smi is +missing. + +This collector runs ONCE per node, captures which tools were resolvable +in PATH, and feeds three downstream consumers: + + 1. A loud `_warn` at run-time so missing tools are visible in the + srun log right next to "start node-smoke". + 2. An always-on "Tooling availability" section in the aggregator + report -- a per-node table that's loud even when nothing else is. + 3. The optional `--require-tools` flag, which promotes a missing + required tool to a hard FAIL via _node_status_from. +""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from ..shell_utils import _which + +_TRACKED_TOOLS = ("amd-smi", "rocm-smi", "lsof") + + +def _collect_tooling_inventory() -> Dict[str, Any]: + """Resolve each tracked tool in PATH and compute per-check coverage. + + Output shape:: + + { + "ok": True, # always True; collector itself can't fail + "tools": { + "amd-smi": {"present": True, "path": "/usr/bin/amd-smi"}, + "rocm-smi": {"present": True, "path": "/usr/bin/rocm-smi"}, + "lsof": {"present": True, "path": "/usr/bin/lsof"}, + }, + "missing": ["rocm-smi"], # convenience list of absent tools + "coverage": { + "ECC": True, + "XGMI": True, + "foreign-process": True, + "GPU activity warn": True, + "amd-smi/torch GPU-count cross-check": True, + "wedged-driver canary": True, + }, + "uncovered": [], # checks with no working tool + } + + Coverage rules (each check is "covered" if ANY of the listed tools + is present): + + * ECC -> amd-smi OR rocm-smi (--showrasinfo) + * XGMI link matrix -> amd-smi OR rocm-smi (--showtopotype) + * foreign-process enumeration -> amd-smi OR rocm-smi OR lsof + * GPU activity warn (gfx_activity_pct)-> amd-smi OR rocm-smi (--showuse) + * amd-smi/torch GPU-count cross-check -> amd-smi only (rocm-smi enumerates + by `cardN` not by torch's index) + * wedged-driver canary (rocm-smi --version) -> rocm-smi only + """ + tools: Dict[str, Dict[str, Any]] = {} + missing: List[str] = [] + for name in _TRACKED_TOOLS: + path = _which(name) + tools[name] = {"present": path is not None, "path": path} + if path is None: + missing.append(name) + + has_amd = tools["amd-smi"]["present"] + has_rocm = tools["rocm-smi"]["present"] + has_lsof = tools["lsof"]["present"] + coverage = { + "ECC": has_amd or has_rocm, + "XGMI": has_amd or has_rocm, + "foreign-process": has_amd or has_rocm or has_lsof, + "GPU activity warn": has_amd or has_rocm, + "amd-smi/torch GPU-count cross-check": has_amd, + "wedged-driver canary": has_rocm, + } + uncovered = [name for name, ok in coverage.items() if not ok] + return { + "ok": True, + "tools": tools, + "missing": missing, + "coverage": coverage, + "uncovered": uncovered, + } diff --git a/primus/tools/preflight/node_smoke/collectors/xgmi.py b/primus/tools/preflight/node_smoke/collectors/xgmi.py new file mode 100644 index 000000000..0d6d55ea8 --- /dev/null +++ b/primus/tools/preflight/node_smoke/collectors/xgmi.py @@ -0,0 +1,201 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 1 -- D-2: XGMI topology matrix via ``amd-smi topology`` (text parser), +with rocm-smi --showtopotype as a cross-tool fallback.""" + +from __future__ import annotations + +import subprocess +from typing import Any, Dict, List + +from ..shell_utils import _which +from .rocm_smi import _rocm_smi_topotype_json + + +def _collect_xgmi_topology() -> Dict[str, Any]: + """Try amd-smi topology first; fall back to rocm-smi --showtopotype. + + Both tools produce slightly different per-row labels (amd-smi uses + PCIe BDFs, rocm-smi uses DRM device indices), but the link_types + matrix and non_xgmi_pairs computation is identical, so downstream + consumers in _node_status_from / the aggregator don't need to know + which tool produced the data. + """ + out = _collect_xgmi_topology_amd_smi() + if out.get("ok"): + return out + rocm = _rocm_smi_topotype_json() + if rocm.get("ok"): + return { + "ok": True, + "tool": rocm.get("tool") or "rocm-smi --showtopotype --json", + "bdfs": [], # rocm-smi reports DRM indices, not PCIe BDFs + "matrix": rocm.get("link_types") or [], + "n_gpus": rocm.get("n_gpus") or 0, + "non_xgmi_pairs": rocm.get("non_xgmi_pairs") or [], + "amd_smi_error": out.get("error"), + } + # Both failed -- preserve amd-smi's error for the operator and + # surface rocm-smi's separately so they can debug both paths. + out["rocm_smi_error"] = rocm.get("error") + return out + + +def _collect_xgmi_topology_amd_smi() -> Dict[str, Any]: + """Parse ``amd-smi topology`` and return a square link-type matrix. + + ``amd-smi topology`` emits several BDF-labelled sub-tables (ACCESS, + WEIGHT, HOPS, LINK TYPE, NUMA BW, ...). We pick the ``LINK TYPE TABLE`` + sub-section, which contains values like ``SELF`` (diagonal) and + ``XGMI`` / ``PCIE`` / ``PIX`` / ``SOC`` etc. Off-diagonal cells that + aren't ``XGMI`` are recorded as ``non_xgmi_pairs`` and treated as a + hard fail by ``_node_status_from`` -- the moment a single GPU pair + falls back to PCIe inside a node, intra-node collectives lose 5-10x + of the bandwidth NCCL/RCCL expects. + + The on-disk shape: + + { + "ok": bool, + "tool": "amd-smi topology" | None, + "bdfs": ["0000:05:00.0", ...], + "matrix": [["SELF","XGMI",...], ["XGMI","SELF",...], ...], + "n_gpus": int, + "non_xgmi_pairs": [(i, j, link_type), ...], + "error": "..." + } + """ + out: Dict[str, Any] = { + "ok": False, + "tool": None, + "bdfs": [], + "matrix": [], + "non_xgmi_pairs": [], + } + if _which("amd-smi") is None: + out["error"] = "amd-smi not found in PATH" + return out + try: + cp = subprocess.run( + ["amd-smi", "topology"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=15, + check=False, + ) + except subprocess.TimeoutExpired: + out["error"] = "amd-smi topology timed out" + return out + except Exception as e: + out["error"] = str(e) + return out + if cp.returncode != 0: + out["error"] = (cp.stderr or "").strip()[:200] or f"rc={cp.returncode}" + return out + + text = cp.stdout + + # Parse: find the `LINK TYPE TABLE:` section, then the BDF header row, + # then the per-BDF data rows. Stop at the next section header (any all- + # caps label ending in `TABLE:`) or end of text. + import re + + bdf_re = re.compile(r"\b([0-9a-fA-F]{4}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.\d)\b") + section_header_re = re.compile(r"^\s*[A-Z][A-Z0-9 -]+TABLE:\s*$") + + lines = text.splitlines() + try: + idx = next(i for i, l in enumerate(lines) if l.strip().upper() == "LINK TYPE TABLE:") + except StopIteration: + out["error"] = "no `LINK TYPE TABLE:` section in `amd-smi topology` output" + out["raw"] = text[:4000] + return out + + # Header row is the next non-empty line after the label, and contains + # no leading BDF -- only column BDFs. + header_bdfs: List[str] = [] + data_start = None + for j in range(idx + 1, len(lines)): + l = lines[j].rstrip() + if not l.strip(): + continue + if section_header_re.match(l): + break + toks = bdf_re.findall(l) + if not toks: + continue + # The header line has only column BDFs (no leading row label), and + # the first non-whitespace char position lines up with the columns. + # Heuristic: header has BDFs but no other tokens that look like + # link-type values (XGMI/PCIE/SELF/...). Data rows always have + # exactly one leading BDF followed by N value tokens. + non_bdf_toks = [t for t in l.split() if not bdf_re.fullmatch(t)] + if not non_bdf_toks: + header_bdfs = toks + data_start = j + 1 + break + + if not header_bdfs or data_start is None: + out["error"] = "could not find header row inside LINK TYPE TABLE" + out["raw"] = text[:4000] + return out + + n = len(header_bdfs) + bdf_to_idx = {b: i for i, b in enumerate(header_bdfs)} + matrix: List[List[str]] = [[""] * n for _ in range(n)] + seen_rows = 0 + for j in range(data_start, len(lines)): + l = lines[j].rstrip() + if not l.strip(): + continue + if section_header_re.match(l): + break + toks = l.split() + # First token must be a BDF, the remaining N tokens are the row. + if not bdf_re.fullmatch(toks[0]): + continue + row_bdf = toks[0] + cells = toks[1:] + if row_bdf not in bdf_to_idx: + continue + row_idx = bdf_to_idx[row_bdf] + for k, cell in enumerate(cells[:n]): + matrix[row_idx][k] = cell + seen_rows += 1 + + if seen_rows == 0: + out["error"] = "no BDF-labelled rows found inside LINK TYPE TABLE" + out["raw"] = text[:4000] + return out + + healthy_diag = {"SELF", "X", "-", "0"} + healthy_link = {"XGMI"} + non_xgmi: List[Any] = [] + for i, row in enumerate(matrix): + for j_idx, cell in enumerate(row): + cu = cell.strip().upper() + if i == j_idx: + # Diagonal: must be SELF (or empty if the row was missing). + if cu and cu not in healthy_diag and cu not in healthy_link: + non_xgmi.append((i, j_idx, cell)) + continue + if not cu: + # Missing cell -> can't certify XGMI -> flag. + non_xgmi.append((i, j_idx, "")) + continue + if cu in healthy_link or cu in healthy_diag: + continue + non_xgmi.append((i, j_idx, cell)) + + out["ok"] = True + out["tool"] = "amd-smi topology" + out["bdfs"] = header_bdfs + out["matrix"] = matrix + out["n_gpus"] = n + out["non_xgmi_pairs"] = non_xgmi + return out diff --git a/primus/tools/preflight/node_smoke/logging_utils.py b/primus/tools/preflight/node_smoke/logging_utils.py new file mode 100644 index 000000000..7c853e165 --- /dev/null +++ b/primus/tools/preflight/node_smoke/logging_utils.py @@ -0,0 +1,53 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Hostname normalization + log/warn helpers used everywhere in node-smoke. + +The log prefix ``[HH:MM:SS][node-smoke][]`` is part of the +operator-facing contract -- existing log-scraping tooling assumes it, +so it is preserved verbatim across the refactor. +""" + +from __future__ import annotations + +import socket +import sys +import time + + +def _ts() -> str: + return time.strftime("%H:%M:%S") + + +def _short_name(h: str) -> str: + """Return the leading short-hostname segment. + + SLURM tools (`scontrol show hostnames`, `srun --nodelist=`, + `srun --exclude=`) all operate on short hostnames, so we normalize + everywhere so the produced ``passing_nodes.txt`` / ``failing_nodes.txt`` + can be piped straight into them. ``socket.gethostname()`` returns the + FQDN on some clusters, hence this helper. + """ + if not h: + return h + return h.split(".", 1)[0] + + +def _this_host_short() -> str: + """This node's short hostname (first segment of socket.gethostname()).""" + return _short_name(socket.gethostname()) + + +def _log(msg: str) -> None: + print(f"[{_ts()}][node-smoke][{_this_host_short()}] {msg}", flush=True) + + +def _warn(msg: str) -> None: + print( + f"[{_ts()}][node-smoke][{_this_host_short()}] WARN: {msg}", + file=sys.stderr, + flush=True, + ) diff --git a/primus/tools/preflight/node_smoke/orchestrator.py b/primus/tools/preflight/node_smoke/orchestrator.py new file mode 100644 index 000000000..8c4123af6 --- /dev/null +++ b/primus/tools/preflight/node_smoke/orchestrator.py @@ -0,0 +1,288 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Per-node orchestration helpers. + +* :func:`_spawn_per_gpu` -- launch the ``_per_gpu`` subcommand for one + GPU index with a hard timeout. +* :func:`_node_status_from` -- compute the list of fail reasons for the + whole node from collected per-GPU + tier1 + + tier2 state. +* :func:`_clean_dump_path` -- wipe stale artifacts from a previous run on + rank 0 before any rank writes its current + JSON. +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import time +from typing import Any, Dict, List, Optional + +from .types import GPUResult + + +def _spawn_per_gpu( + gpu: int, + *, + timeout_sec: int, + tier2_perf: bool, + gemm_tflops_min: float, + hbm_gbs_min: float, + hbm_busy_threshold_gib: float, +) -> GPUResult: + """Spawn ``python -m primus.tools.preflight.node_smoke _per_gpu ...`` + with a hard timeout so a stuck driver call cannot wedge the parent.""" + cmd = [ + sys.executable, + "-m", + "primus.tools.preflight.node_smoke", + "_per_gpu", + str(gpu), + "--gemm-tflops-min", + str(gemm_tflops_min), + "--hbm-gbs-min", + str(hbm_gbs_min), + "--hbm-busy-threshold-gib", + str(hbm_busy_threshold_gib), + ] + if tier2_perf: + cmd.append("--tier2-perf") + + t0 = time.time() + try: + cp = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=timeout_sec, + check=False, + ) + except subprocess.TimeoutExpired: + return GPUResult( + gpu=gpu, + status="TIMEOUT", + reason=f"per-gpu subprocess hit hard timeout {timeout_sec}s", + duration_sec=round(time.time() - t0, 3), + ) + except Exception as e: + return GPUResult( + gpu=gpu, + status="FAIL", + reason=f"failed to spawn per-gpu subprocess: {e}", + duration_sec=round(time.time() - t0, 3), + ) + + # The subprocess prints exactly one JSON line on stdout for the result. + raw = (cp.stdout or "").strip().splitlines() + if not raw: + return GPUResult( + gpu=gpu, + status="FAIL", + reason=( + f"per-gpu subprocess produced no JSON (rc={cp.returncode}, " + f"stderr={cp.stderr.strip()[:200]})" + ), + duration_sec=round(time.time() - t0, 3), + ) + try: + data = json.loads(raw[-1]) + except Exception as e: + return GPUResult( + gpu=gpu, + status="FAIL", + reason=f"per-gpu JSON parse failed: {e}; raw={raw[-1][:200]}", + duration_sec=round(time.time() - t0, 3), + ) + + return GPUResult( + gpu=int(data.get("gpu", gpu)), + status=str(data.get("status", "FAIL")), + reason=str(data.get("reason", "")), + duration_sec=float(data.get("duration_sec", time.time() - t0)), + details=dict(data.get("details", {})), + ) + + +def _node_status_from( + per_gpu: List[GPUResult], + tier1_extra: Dict[str, Any], + tier2_extra: Dict[str, Any], + *, + allow_foreign_procs: bool = False, + required_tools: Optional[List[str]] = None, +) -> List[str]: + """Compute a list of ``fail_reasons`` for the node from collected results. + + Empty list -> node PASS. Any non-empty result -> node FAIL. + + ``allow_foreign_procs`` downgrades the foreign-process FAIL to a + silent inclusion in the JSON (still surfaced by the aggregator's + "Busy GPUs" section, just not a hard fail). + + ``required_tools`` is the operator-supplied list of CLI tools that + MUST be present on this node (e.g. ``["amd-smi", "rocm-smi"]``). + Anything in this list that is not in the tooling-inventory becomes + a hard FAIL. Empty / None means "warn only" (the default). + """ + reasons: List[str] = [] + + # Self-contained GPU visibility guard. Decoupled from any other + # collector so a wrapped/downgraded "No GPUs detected" finding can + # never silently turn a CPU-only or stale-GPU node into a PASS. + vis = tier1_extra.get("gpu_visibility") or {} + for r in vis.get("fail_reasons", []) or []: + reasons.append(f"gpu_visibility: {r}") + + for r in per_gpu: + if r.status != "PASS": + reasons.append(f"gpu{r.gpu}: {r.status}: {r.reason}") + + for section_name in ("gpu_info", "host_info", "network_info"): + for f in tier1_extra.get(section_name, []): + if f.get("level") == "fail": + reasons.append(f"{section_name}: {f.get('message', '')}") + + dmesg = tier1_extra.get("dmesg") or {} + if dmesg.get("matches"): + first = dmesg["matches"][0] + reasons.append(f"dmesg ({len(dmesg['matches'])} match(es), e.g.): {first[:200]}") + + # B. NIC / RDMA roll-call -- every issue here is a hard fail because each + # one (port DOWN, missing RoCE v2 GID, wrong NIC count) silently breaks + # inter-node training the moment the first global collective runs. + for issue in (tier1_extra.get("nics") or {}).get("issues", []) or []: + reasons.append(f"nic: {issue}") + + # C. Host limits -- only the entries the collector flagged as hard + # (ulimit -l below threshold, /dev/shm too small) become node FAIL. + for issue in (tier1_extra.get("host_limits") or {}).get("fail_reasons", []) or []: + reasons.append(f"host_limits: {issue}") + + rccl = tier2_extra.get("rccl") or {} + if rccl and rccl.get("status") not in (None, "PASS"): + reasons.append(f"rccl: {rccl.get('status')}: {rccl.get('error', '')}") + + # D-1 heavy: any per-GPU uncorrectable ECC count is a hard fail. The + # amd-smi schema isn't stable across releases so we trust only the + # values our flattener was able to coerce to int. Throttle reasons stay + # informational (the schema is too vendor-specific to fail on). + amd = tier1_extra.get("gpu_low_level") or {} + for rec in amd.get("per_gpu", []) or []: + ue = rec.get("ecc_uncorrectable_total") + if isinstance(ue, int) and ue > 0: + reasons.append(f"gpu{rec.get('gpu', '?')}: ECC uncorrectable count = {ue}") + + # D-2: any non-XGMI GPU pair is a hard fail -- intra-node collectives + # silently fall back to PCIe and lose 5-10x bandwidth. + xg = tier1_extra.get("xgmi") or {} + bad = xg.get("non_xgmi_pairs") or [] + if bad: + sample = ", ".join(f"({i},{j})={t}" for i, j, t in bad[:3]) + reasons.append(f"xgmi: {len(bad)} non-XGMI GPU pair(s) detected, e.g. {sample}") + + # F-partial: rocm-smi --version that timed out -> driver is wedging. + # Slow-but-completed calls are surfaced by the aggregator only. + tool = tier1_extra.get("tooling") or {} + if tool.get("timed_out"): + reasons.append( + f"tooling: rocm-smi --version did not return within " + f"{tool.get('timeout_sec', '?')}s -- driver may be wedging" + ) + + # Tooling availability. Missing CLI tools (amd-smi, rocm-smi, lsof) + # cause silent skips of several Tier 1 checks. Operators in strict + # environments can pass --require-tools to convert "missing" into a + # node FAIL so the node is pulled from rotation until the toolchain + # is fixed. Default (empty list) is warn-only (the WARN already fires + # in _cmd_run before the per-GPU subprocesses run). + if required_tools: + inv = (tier1_extra.get("tooling_inventory") or {}).get("tools") or {} + missing_required = [t for t in required_tools if not (inv.get(t) or {}).get("present")] + if missing_required: + reasons.append( + f"tooling_inventory: required tool(s) NOT in PATH: " + f"{', '.join(missing_required)} -- silent-skips several " + f"Tier 1 checks (ECC, XGMI, foreign-process, wedged-driver). " + f"Pass --require-tools '' to disable this fail." + ) + + # G: foreign processes holding the GPU. Hard-fail by default because + # this is the single most common cause of training failing to launch + # on an otherwise-healthy node (leaked Python ranks from a previous + # job, a profiler that never detached, a foreign tenant on a shared + # partition). The operator can downgrade with --allow-foreign-procs + # if their workflow legitimately co-tenants the GPU. + gp = tier1_extra.get("gpu_processes") or {} + if not allow_foreign_procs and gp.get("foreign_count", 0) > 0: + examples: List[str] = [] + for g in gp.get("per_gpu") or []: + for p in g.get("processes") or []: + if not p.get("is_foreign"): + continue + hbm = p.get("hbm_bytes") + hbm_s = f" hbm={round(hbm / (1 << 30), 2)}GiB" if isinstance(hbm, int) and hbm > 0 else "" + examples.append(f"gpu{g.get('gpu')}: pid={p.get('pid')} " f"name={p.get('name')!r}{hbm_s}") + if len(examples) >= 3: + break + if len(examples) >= 3: + break + reasons.append( + f"gpu_processes: {gp.get('foreign_count', 0)} foreign process(es) " + f"holding GPU(s) (e.g. " + "; ".join(examples) + ") -- likely " + f"leaked rank(s) from a previous job. Clean up with " + f"`pkill -9 -f train.py` (or similar) or pass --allow-foreign-procs." + ) + + return reasons + + +def _clean_dump_path(dump_path: str) -> List[str]: + """Wipe stale per-node JSONs and aggregator outputs from a previous run. + + Without this, a re-run on a different (smaller) nodelist would leave + JSONs from removed nodes in ``/smoke/`` and the aggregator would + happily count them as PASS, contaminating the report. We clean only + artifacts that ``run`` and ``aggregate`` produce (per-node JSONs and + the four top-level outputs); anything else under ``--dump-path`` is + left untouched. + + Race safety: this is called only on rank 0 in ``_cmd_run``, BEFORE any + rank can have written its current-run JSON (each rank's per-GPU + subprocess loop + collector phase takes seconds; rank 0's cleanup + finishes in milliseconds). Other ranks never delete anything. + + Returns the list of files actually removed (for logging). + """ + removed: List[str] = [] + smoke_dir = os.path.join(dump_path, "smoke") + if os.path.isdir(smoke_dir): + for name in os.listdir(smoke_dir): + if name.endswith(".json"): + p = os.path.join(smoke_dir, name) + try: + os.remove(p) + removed.append(p) + except OSError: + pass + for name in ( + "smoke_report.md", + "passing_nodes.txt", + "failing_nodes.txt", + "expected_nodes.txt", + ): + p = os.path.join(dump_path, name) + if os.path.isfile(p): + try: + os.remove(p) + removed.append(p) + except OSError: + pass + return removed diff --git a/primus/tools/preflight/node_smoke/per_gpu.py b/primus/tools/preflight/node_smoke/per_gpu.py new file mode 100644 index 000000000..64a29173e --- /dev/null +++ b/primus/tools/preflight/node_smoke/per_gpu.py @@ -0,0 +1,327 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Per-GPU subprocess body (Tier 1 + optional Tier 2 perf). + +This module is the body of the ``_per_gpu`` subcommand: it runs every +GPU-local check on a single device index and returns a dict result. +The orchestrator spawns one subprocess per GPU with a hard timeout so a +stuck driver call cannot wedge the parent. + +Kept intact (no internal split) -- every code path returns a complete +verdict dict, and inlining each stage is currently easier to follow +than splitting the function up. +""" + +from __future__ import annotations + +import time +from typing import Any, Dict + +from .shell_utils import _read_text, _resolve_gpu_bdf + + +def _per_gpu_body( + gpu: int, + *, + tier2_perf: bool, + gemm_tflops_min: float, + hbm_gbs_min: float, + hbm_busy_threshold_bytes: int = 2 * (1 << 30), +) -> Dict[str, Any]: + """Run all per-GPU tests for a single GPU and return a dict result. + + Tier 1 (always): set_device, **pre-touch HBM-busy check** (FAIL if more + than ``hbm_busy_threshold_bytes`` already in use before our test + allocates anything), allocate 256 MB, tiny GEMM 2048x2048 bf16 with + finite-value check. + + Tier 2 (when ``tier2_perf`` is True): GEMM 8192x8192 bf16 TFLOPS + measurement against ``gemm_tflops_min``, and HBM device-to-device + copy bandwidth against ``hbm_gbs_min``. Each metric below threshold + yields FAIL. + + Pre-touch HBM check: ``torch.cuda.mem_get_info`` is called BEFORE we + allocate anything on this GPU, so the "used" reading reflects only + foreign / leaked allocations. The post-test reading is also captured + (under ``low_level.hbm_free_bytes``) for completeness, but the FAIL + rule uses only the pre-touch number to avoid being polluted by our + own caching-allocator footprint. + """ + t0 = time.time() + details: Dict[str, Any] = {} + + try: + import torch # type: ignore + except Exception as e: + return { + "gpu": gpu, + "status": "FAIL", + "reason": f"torch import failed: {e}", + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + + if not torch.cuda.is_available(): + return { + "gpu": gpu, + "status": "FAIL", + "reason": "torch.cuda.is_available() is False", + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + if gpu >= torch.cuda.device_count(): + return { + "gpu": gpu, + "status": "FAIL", + "reason": (f"gpu index {gpu} >= visible device_count {torch.cuda.device_count()}"), + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + + # --- set_device --- + try: + torch.cuda.set_device(gpu) + except Exception as e: + return { + "gpu": gpu, + "status": "FAIL", + "reason": f"set_device({gpu}) raised: {e}", + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + + # --- pre-touch HBM-busy check (BEFORE we allocate anything) --- + # Captured here, NOT in the low_level block at the end, because by + # then PyTorch's caching allocator has already taken pages we won't + # truly release on empty_cache(). The pre-touch reading is the only + # honest answer to "is someone else holding this GPU?". + try: + free_b, total_b = torch.cuda.mem_get_info(gpu) + used_b = max(0, int(total_b) - int(free_b)) + details["hbm_pre_touch_total_bytes"] = int(total_b) + details["hbm_pre_touch_free_bytes"] = int(free_b) + details["hbm_pre_touch_used_bytes"] = used_b + details["hbm_pre_touch_used_gib"] = round(used_b / (1 << 30), 3) + if used_b >= hbm_busy_threshold_bytes: + return { + "gpu": gpu, + "status": "FAIL", + "reason": ( + f"pre-touch HBM busy: {round(used_b / (1 << 30), 2)} GiB " + f"already in use (threshold " + f"{round(hbm_busy_threshold_bytes / (1 << 30), 2)} GiB) " + f"-- likely leaked process from a previous job; " + f"see node-level gpu_processes section to identify the PID" + ), + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + except Exception as e: + details["hbm_pre_touch_error"] = f"mem_get_info failed: {e}" + + # --- 256 MB tensor alloc + simple write + sync --- + try: + nbytes_alloc = 256 * 1024 * 1024 + n_elem = nbytes_alloc // 2 # bf16 + x = torch.empty(n_elem, dtype=torch.bfloat16, device=f"cuda:{gpu}") + x.fill_(1.0) + torch.cuda.synchronize() + del x + torch.cuda.empty_cache() + except Exception as e: + return { + "gpu": gpu, + "status": "FAIL", + "reason": f"256MB bf16 alloc/fill/sync failed: {e}", + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + + # --- tiny GEMM 2048x2048 bf16, finite-value check --- + try: + m = n = k = 2048 + a = torch.randn((m, k), dtype=torch.bfloat16, device=f"cuda:{gpu}") + b = torch.randn((k, n), dtype=torch.bfloat16, device=f"cuda:{gpu}") + c = torch.matmul(a, b) + torch.cuda.synchronize() + if not torch.isfinite(c).all().item(): + return { + "gpu": gpu, + "status": "FAIL", + "reason": "tiny GEMM produced non-finite values (possible HW corruption)", + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + del a, b, c + torch.cuda.empty_cache() + except Exception as e: + return { + "gpu": gpu, + "status": "FAIL", + "reason": f"tiny GEMM 2048x2048 failed: {e}", + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + + # --- D-1 light: PCIe link + HBM (sysfs + torch only, fast & no shell-out) --- + # Captured into details.low_level so the aggregator can flag drift across + # the cluster (e.g. a single GPU running at Gen3 x8 because the slot + # needs reseating, or a GPU that exposes only half of its HBM). + # Each sub-capture is independent: a missing/unparseable PCIe BDF must + # not cost us the HBM size, and vice-versa. + low: Dict[str, Any] = {} + props = None + try: + props = torch.cuda.get_device_properties(gpu) + except Exception as e: + low["error"] = f"get_device_properties failed: {e}" + + if props is not None: + # PCIe link details (sysfs) + try: + bdf = _resolve_gpu_bdf(props) + if bdf: + low["pci_bdf"] = bdf + sysdir = f"/sys/bus/pci/devices/{bdf}" + speed = _read_text(f"{sysdir}/current_link_speed") + width = _read_text(f"{sysdir}/current_link_width") + low["pcie_link_speed_raw"] = speed or None + low["pcie_link_width"] = int(width) if width.isdigit() else None + # speed is e.g. "32.0 GT/s PCIe" -> 32.0 + try: + low["pcie_link_speed_gts"] = float(speed.split()[0]) if speed else None + except Exception: + low["pcie_link_speed_gts"] = None + else: + low["pcie_error"] = ( + f"could not resolve PCIe BDF (pci_bus_id=" f"{getattr(props, 'pci_bus_id', None)!r})" + ) + except Exception as e: + low["pcie_error"] = f"PCIe sysfs capture failed: {e}" + + # HBM total/free (torch). Independent of BDF resolution. + try: + free_b, total_b = torch.cuda.mem_get_info(gpu) + low["hbm_total_bytes"] = int(total_b) + low["hbm_free_bytes"] = int(free_b) + low["hbm_total_gib"] = round(total_b / (1 << 30), 2) + except Exception as e: + low["hbm_error"] = f"mem_get_info failed: {e}" + tm = int(getattr(props, "total_memory", 0) or 0) + if tm: + low["hbm_total_bytes"] = tm + low["hbm_total_gib"] = round(tm / (1 << 30), 2) + if low: + details["low_level"] = low + + # --- Tier 2 perf sanity (optional) --- + if tier2_perf: + # GEMM TFLOPS. Warmup/iter counts mirror the preflight `--quick` preset + # (`square_gemm.py` with WARMUP=5, ITERATION=20) so smoke and preflight + # report comparable steady-state numbers. + try: + tflops = _measure_gemm_tflops(gpu, size=8192, warmup=5, iters=20) + details["gemm_tflops"] = round(tflops, 2) + if tflops < gemm_tflops_min: + return { + "gpu": gpu, + "status": "FAIL", + "reason": (f"GEMM TFLOPS {tflops:.0f} < threshold {gemm_tflops_min:.0f}"), + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + except Exception as e: + return { + "gpu": gpu, + "status": "FAIL", + "reason": f"GEMM TFLOPS measurement failed: {e}", + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + + # HBM device-to-device copy bandwidth. HBM is fast enough that we + # need a healthy number of timed iterations for stable timing. + try: + gbs = _measure_hbm_gbs(gpu, size_bytes=512 * 1024 * 1024, warmup=10, iters=20) + details["hbm_gbs"] = round(gbs, 1) + if gbs < hbm_gbs_min: + return { + "gpu": gpu, + "status": "FAIL", + "reason": f"HBM GB/s {gbs:.0f} < threshold {hbm_gbs_min:.0f}", + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + except Exception as e: + return { + "gpu": gpu, + "status": "FAIL", + "reason": f"HBM bandwidth measurement failed: {e}", + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + + return { + "gpu": gpu, + "status": "PASS", + "reason": "", + "duration_sec": round(time.time() - t0, 3), + "details": details, + } + + +def _measure_gemm_tflops(gpu: int, *, size: int, warmup: int, iters: int) -> float: + """Measure GEMM TFLOPS for square ``size x size`` bf16 matmul on ``cuda:gpu``.""" + import torch # type: ignore + + torch.cuda.set_device(gpu) + a = torch.randn((size, size), dtype=torch.bfloat16, device=f"cuda:{gpu}") + b = torch.randn((size, size), dtype=torch.bfloat16, device=f"cuda:{gpu}") + for _ in range(warmup): + _ = torch.matmul(a, b) + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + _ = torch.matmul(a, b) + end.record() + end.synchronize() + + elapsed_s = start.elapsed_time(end) / 1000.0 / iters + flops = 2.0 * size * size * size + return (flops / elapsed_s) / 1e12 + + +def _measure_hbm_gbs(gpu: int, *, size_bytes: int, warmup: int, iters: int) -> float: + """Measure local HBM bandwidth via device-to-device ``copy_``. + + Each iteration: 1 read of ``src`` + 1 write to ``dst`` = ``2 * size_bytes`` + of HBM traffic. Uses ``torch.cuda.Event`` for accurate GPU-side timing. + """ + import torch # type: ignore + + torch.cuda.set_device(gpu) + n = size_bytes // 2 # bf16 = 2 bytes/element + src = torch.empty(n, dtype=torch.bfloat16, device=f"cuda:{gpu}") + dst = torch.empty_like(src) + src.fill_(1.0) + for _ in range(warmup): + dst.copy_(src) + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + dst.copy_(src) + end.record() + end.synchronize() + + elapsed_s = start.elapsed_time(end) / 1000.0 / iters + return (2.0 * size_bytes / elapsed_s) / 1e9 diff --git a/primus/tools/preflight/node_smoke/rccl_local.py b/primus/tools/preflight/node_smoke/rccl_local.py new file mode 100644 index 000000000..6e34dca19 --- /dev/null +++ b/primus/tools/preflight/node_smoke/rccl_local.py @@ -0,0 +1,152 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tier 2 -- node-local RCCL all-reduce (optional). + +Tier 2 also includes a node-local RCCL all-reduce as a steady-state +intra-node bandwidth check. We use ``torch.multiprocessing.spawn`` to +launch one worker per local GPU on a process group bound to +``tcp://127.0.0.1:``. No cross-node communication. +""" + +from __future__ import annotations + +import json +import os +import socket +import time +from typing import Any, Dict, List + + +def _rccl_worker( + local_rank: int, + world_size: int, + port: int, + size_mb: int, + out_path: str, +) -> None: + """Subprocess body for ``torch.multiprocessing.spawn``. + + Runs ``warmup`` warmup + ``iters`` timed all-reduces of ``size_mb`` MB on a + local-only NCCL/RCCL process group bound to ``tcp://127.0.0.1:port``. + Local rank 0 writes the resulting GB/s to ``out_path``. + + Iteration counts are intentionally aligned with the preflight `--quick` + preset (`intra_node_comm.py` with WARMUP=5, ITERATION=20) so smoke and + preflight report comparable steady-state bandwidth. + """ + import torch # type: ignore + import torch.distributed as dist # type: ignore + + warmup = 5 + iters = 20 + + try: + torch.cuda.set_device(local_rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + world_size=world_size, + rank=local_rank, + ) + nbytes = size_mb * 1024 * 1024 + n_elem = nbytes // 2 # bf16 + t = torch.ones(n_elem, dtype=torch.bfloat16, device=f"cuda:{local_rank}") + + for _ in range(warmup): + dist.all_reduce(t) + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + dist.all_reduce(t) + end.record() + end.synchronize() + + elapsed_s = start.elapsed_time(end) / 1000.0 / iters + # NCCL all-reduce effective bandwidth: 2*S*(P-1)/P bytes per rank. + comm_bytes = 2.0 * nbytes * (world_size - 1) / world_size + gbs = comm_bytes / elapsed_s / 1e9 + + if local_rank == 0: + with open(out_path, "w", encoding="utf-8") as f: + json.dump({"status": "PASS", "gbs": round(gbs, 1)}, f) + dist.barrier() + dist.destroy_process_group() + except Exception as e: + if local_rank == 0: + try: + with open(out_path, "w", encoding="utf-8") as f: + json.dump({"status": "FAIL", "error": str(e)}, f) + except Exception: + pass + + +def _run_local_rccl(*, local_world_size: int, size_mb: int, timeout_sec: int) -> Dict[str, Any]: + """Spawn local-only RCCL workers to measure intra-node all-reduce bandwidth. + + Returns ``{"status": "PASS"|"FAIL"|"TIMEOUT", ...}``. + """ + import tempfile + + if local_world_size <= 1: + return {"status": "PASS", "gbs": None, "skipped": "local_world_size<=1"} + + try: + import torch # type: ignore + import torch.multiprocessing as mp # type: ignore + except Exception as e: + return {"status": "FAIL", "error": f"torch import failed: {e}"} + + # Pick a free local TCP port. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + + with tempfile.NamedTemporaryFile(prefix="node_smoke_rccl_", suffix=".json", delete=False) as tf: + out_path = tf.name + + ctx = mp.get_context("spawn") + procs: List[Any] = [] + try: + for r in range(local_world_size): + p = ctx.Process( + target=_rccl_worker, + args=(r, local_world_size, port, size_mb, out_path), + ) + p.start() + procs.append(p) + + deadline = time.time() + timeout_sec + for p in procs: + remaining = max(0.0, deadline - time.time()) + p.join(timeout=remaining) + if p.is_alive(): + # One worker stuck -> kill all and report TIMEOUT. + for q in procs: + if q.is_alive(): + q.terminate() + for q in procs: + q.join(timeout=5) + if q.is_alive(): + q.kill() + return { + "status": "TIMEOUT", + "error": f"local RCCL all-reduce did not finish in {timeout_sec}s", + } + + if not os.path.exists(out_path): + return {"status": "FAIL", "error": "no result file produced"} + with open(out_path, "r", encoding="utf-8") as f: + data = json.load(f) + return data + finally: + try: + os.unlink(out_path) + except Exception: + pass diff --git a/primus/tools/preflight/node_smoke/shell_utils.py b/primus/tools/preflight/node_smoke/shell_utils.py new file mode 100644 index 000000000..952044125 --- /dev/null +++ b/primus/tools/preflight/node_smoke/shell_utils.py @@ -0,0 +1,179 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Stateless utility helpers shared across collectors and orchestration. + +Every function here is pure (no I/O effects beyond reading sysfs / running +a tiny subprocess) and never raises -- callers expect best-effort +defaults so a missing tool / file degrades gracefully. +""" + +from __future__ import annotations + +import os +import subprocess +from typing import Any, Dict, List, Optional + + +def _which(prog: str) -> Optional[str]: + """Tiny shutil.which() replacement that doesn't pull in shutil at import.""" + for d in (os.environ.get("PATH") or "").split(os.pathsep): + p = os.path.join(d, prog) + if os.path.isfile(p) and os.access(p, os.X_OK): + return p + return None + + +def _read_text(path: str, default: str = "") -> str: + """Best-effort read of a small sysfs/proc text file. Never raises.""" + try: + with open(path, "r", encoding="utf-8", errors="replace") as f: + return f.read().strip() + except Exception: + return default + + +def _parse_os_release_pretty() -> Optional[str]: + """Return PRETTY_NAME from /etc/os-release, or None.""" + try: + with open("/etc/os-release", "r", encoding="utf-8") as f: + for line in f: + if line.startswith("PRETTY_NAME="): + v = line.split("=", 1)[1].strip().strip('"').strip("'") + return v + except Exception: + pass + return None + + +def _resolve_gpu_bdf(props: Any) -> Optional[str]: + """Return the PCIe BDF (e.g. ``"0000:75:00.0"``) for a torch device. + + ``torch.cuda.get_device_properties(i).pci_bus_id`` is annoyingly + polymorphic across PyTorch + ROCm versions: sometimes a string in the + canonical ``domain:bus:device.function`` form (lowercase or uppercase), + sometimes an int (just the bus byte). We coerce all variants into the + canonical lowercase form and verify the sysfs directory actually + exists before returning -- so the caller can read PCIe link info + without an extra existence check. + + Returns None if the BDF cannot be resolved (caller should still + capture HBM via ``mem_get_info`` and skip PCIe sysfs reads). + """ + raw = getattr(props, "pci_bus_id", None) + if raw is None: + return None + # 1) String form. Could be "0000:05:00.0", "05:00.0", or uppercase. + if isinstance(raw, str): + s = raw.strip().lower() + if not s: + return None + candidates = [s, f"0000:{s}" if not s.startswith("0000:") else s] + for c in candidates: + if os.path.isdir(f"/sys/bus/pci/devices/{c}"): + return c + return None + # 2) Int form (just the bus byte). Standard layout for AMD GPUs is + # 0000::00.0; verify with sysfs and fall back to a glob if the + # device.function differs from 00.0 on this host. + if isinstance(raw, int): + bus_hex = f"{raw:02x}" + primary = f"0000:{bus_hex}:00.0" + if os.path.isdir(f"/sys/bus/pci/devices/{primary}"): + return primary + import glob + + matches = sorted(glob.glob(f"/sys/bus/pci/devices/0000:{bus_hex}:*")) + if matches: + return os.path.basename(matches[0]) + return None + return None + + +def _systemctl_is_active(unit: str) -> Optional[str]: + """Return ``systemctl is-active `` ('active'/'inactive'/'failed'/...) + or None if systemctl is missing / errors. Always best-effort.""" + if _which("systemctl") is None: + return None + try: + cp = subprocess.run( + ["systemctl", "is-active", unit], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=3, + check=False, + ) + # systemctl returns non-zero for inactive/failed -- that's fine, + # we just want the textual state. + return (cp.stdout or "").strip() or "unknown" + except Exception: + return None + + +def _parse_size_with_unit(s: str) -> Optional[int]: + """Parse a size string into bytes. + + Accepts (case-insensitive):: + + "12345" -> 12345 + "256 MB" -> 268435456 + "256MB" -> 268435456 (no space) + "12.5 GiB" -> 13421772800 + "12.5GiB" -> 13421772800 + " -1 " -> -1 (sentinel for "unlimited") + + Returns ``None`` for empty input, an unrecognised unit (e.g. + ``"500 MHz"``), or any input the regex cannot fully match (e.g. + ``"12 GB extra"``). This is deliberate: a silent fallthrough + to ``num * 1`` would let frequency / count strings masquerade + as byte counts and corrupt downstream comparisons. + """ + if not s: + return None + import re + + m = re.match( + r"\s*([+-]?\d+(?:\.\d+)?)\s*([a-zA-Z]+)?\.?\s*$", + s, + ) + if not m: + return None + units = { + "b": 1, + "k": 1 << 10, + "kb": 1 << 10, + "kib": 1 << 10, + "m": 1 << 20, + "mb": 1 << 20, + "mib": 1 << 20, + "g": 1 << 30, + "gb": 1 << 30, + "gib": 1 << 30, + "t": 1 << 40, + "tb": 1 << 40, + "tib": 1 << 40, + } + unit = (m.group(2) or "b").lower() + if unit not in units: + return None + try: + return int(float(m.group(1)) * units[unit]) + except ValueError: + return None + + +def _findings_to_dicts(findings: List[Any]) -> List[Dict[str, Any]]: + """Normalize Finding dataclasses (different modules each define their own) + into plain dicts.""" + return [ + { + "level": getattr(f, "level", "info"), + "message": getattr(f, "message", str(f)), + "details": getattr(f, "details", {}), + } + for f in findings + ] diff --git a/primus/tools/preflight/node_smoke/tests/__init__.py b/primus/tools/preflight/node_smoke/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/primus/tools/preflight/node_smoke/tests/test_node_smoke.py b/primus/tools/preflight/node_smoke/tests/test_node_smoke.py new file mode 100644 index 000000000..7e2d2d052 --- /dev/null +++ b/primus/tools/preflight/node_smoke/tests/test_node_smoke.py @@ -0,0 +1,1109 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Regression tests for the node-smoke package. + +Each test is a deliberate guard against a behaviour that has either +broken in production (history items in `docs/node-smoke.md`) or is part +of the operator-facing contract (CLI flags, JSON schema, report section +order). They are pure-Python: no GPU, no subprocess fanout, no real +amd-smi/rocm-smi/lsof dependency. + +Run directly: + + pytest primus/tools/preflight/node_smoke/tests/test_node_smoke.py -v +""" + +from __future__ import annotations + +import subprocess +import sys + +import pytest + +from primus.tools.preflight.node_smoke.aggregator.report import write_smoke_report +from primus.tools.preflight.node_smoke.aggregator.summarizers import ( + _busy_gpu_rows, + _clock_summary, + _pretouch_hbm_rows, + _stack_drift_rows, +) +from primus.tools.preflight.node_smoke.collectors.gpu_processes import ( + _flatten_amd_smi_process_json, + _parse_lsof_pcn, +) +from primus.tools.preflight.node_smoke.collectors.nics import ( + _parse_nic_selector, + _resolve_selector, + _selector_matches, +) +from primus.tools.preflight.node_smoke.collectors.rocm_smi import ( + _parse_rocm_smi_ras_info_text, +) +from primus.tools.preflight.node_smoke.logging_utils import _short_name +from primus.tools.preflight.node_smoke.orchestrator import _node_status_from +from primus.tools.preflight.node_smoke.shell_utils import _parse_size_with_unit + +# --------------------------------------------------------------------------- +# A. Pure-helper unit tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "s,expected", + [ + # Plain ints + decimals (no unit -> bytes) + ("12345", 12345), + ("", None), + # SI / IEC mix, both with and without space + ("256 MB", 256 * (1 << 20)), + ("256MB", 256 * (1 << 20)), + ("12.5 GiB", int(12.5 * (1 << 30))), + ("12.5GiB", int(12.5 * (1 << 30))), + # The "unlimited" sentinel + (" -1 ", -1), + # Things that MUST NOT silently become byte-counts. `500 MHz` was + # the historical bug -- frequency masquerading as bytes. + ("500 MHz", None), + ("12 GB extra", None), + ("not-a-number", None), + ], +) +def test_parse_size_with_unit(s, expected): + """A.1 -- regex + unit table; unrecognized units MUST return None.""" + assert _parse_size_with_unit(s) == expected + + +def test_short_name_strips_fqdn(): + """A.2 -- FQDN normalisation (history item: SLURM-ready txt outputs).""" + assert _short_name("tus1-p3-g25.cluster.example.com") == "tus1-p3-g25" + assert _short_name("tus1-p3-g25") == "tus1-p3-g25" + assert _short_name("") == "" + + +def test_parse_rocm_smi_ras_sums_per_gpu_and_skips_status_only_rows(): + """A.3 -- ECC text parser: one GPU block per `GPU[N]: RAS INFO` header, + summing the last two int columns of each block row, ignoring rows that + have only Status (e.g. ATHUB UNAVAILABLE).""" + sample = ( + "GPU[0]: RAS INFO\n" + " Block Status Correctable Error Uncorrectable Error\n" + " UMC ENABLED 3 1\n" + " SDMA ENABLED 0 0\n" + " ATHUB UNAVAILABLE\n" + "GPU[1]: RAS INFO\n" + " UMC ENABLED 0 0\n" + ) + out = _parse_rocm_smi_ras_info_text(sample) + assert out == [ + {"gpu": 0, "ecc_correctable_total": 3, "ecc_uncorrectable_total": 1}, + {"gpu": 1, "ecc_correctable_total": 0, "ecc_uncorrectable_total": 0}, + ] + + +def test_parse_lsof_pcn_handles_multi_open_pid(): + """A.4 -- lsof -Fpcn field-prefix parser: same PID across multiple open + files collapses to a single annotated record per PID.""" + text = "p123\ncpython\nf3\nf4\np123\ncpython\np456\nctrain.py\n" + rows = _parse_lsof_pcn( + text, + lambda pid, name, hbm: {"pid": pid, "name": name, "hbm_bytes": hbm}, + ) + assert sorted(r["pid"] for r in rows) == [123, 456] + assert {r["pid"]: r["name"] for r in rows} == { + 123: "python", + 456: "train.py", + } + + +# --------------------------------------------------------------------------- +# B. Aggregator-summarizer regression guards +# --------------------------------------------------------------------------- + + +def test_stack_drift_does_not_crash_on_heterogeneous_dict_vs_none(): + """B.1 -- guards against the production crash captured in + docs/node-smoke.md history item 6: a fingerprint key that is None on + one node and a dict on another USED to crash Counter() with + `TypeError: unhashable type: 'dict'`. After the fix the key must + simply be excluded from scalar drift (since it is not a scalar on + any node).""" + nodes = [ + { + "host": "a", + "tier1": {"fingerprint": {"kernel": "5.15", "rocm": "6.2", "nic_fw": None}}, + }, + { + "host": "b", + "tier1": {"fingerprint": {"kernel": "5.15", "rocm": "6.2", "nic_fw": {"rdma0": "20.0"}}}, + }, + ] + # Must return without raising. + rows = _stack_drift_rows(nodes) + # And nic_fw must NOT appear in scalar drift (it's not a scalar + # on any node). + assert all(r["key"] != "nic_fw" for r in rows) + + +def test_clock_summary_spread_and_warn(): + """B.2 -- spread = max - min, warn flag fires above threshold, + nodes with no active time daemon are listed.""" + nodes = [ + {"host": "a", "tier1": {"clock": {"wall_time_unix": 1000.0, "any_active": True}}}, + {"host": "b", "tier1": {"clock": {"wall_time_unix": 1042.5, "any_active": True}}}, + {"host": "c", "tier1": {"clock": {"wall_time_unix": 1010.0, "any_active": False}}}, + ] + s = _clock_summary(nodes, skew_warn_sec=30.0) + assert s["spread_sec"] == 42.5 + assert s["spread_warn"] is True # 42.5 > 30 + # node_rank defaults to "?" when missing from the input dict. + assert ("?", "c") in s["no_daemon_hosts"] + assert s["earliest_host"] == "a" + assert s["latest_host"] == "b" + + +def test_busy_gpu_rows_filters_to_is_foreign_only(): + """B.3 -- only processes flagged is_foreign make it into the + Busy-GPU table; HBM is converted from bytes to GiB.""" + nodes = [ + { + "host": "a", + "node_rank": 0, + "tier1": { + "gpu_processes": { + "ok": True, + "per_gpu": [ + { + "gpu": 0, + "processes": [ + { + "pid": 1, + "name": "self", + "hbm_bytes": 1 << 30, + "is_self": True, + "is_allowed": False, + "is_foreign": False, + }, + { + "pid": 2, + "name": "agent", + "hbm_bytes": 0, + "is_self": False, + "is_allowed": True, + "is_foreign": False, + }, + { + "pid": 3, + "name": "leak", + "hbm_bytes": 4 * (1 << 30), + "is_self": False, + "is_allowed": False, + "is_foreign": True, + }, + ], + } + ], + } + }, + } + ] + rows = _busy_gpu_rows(nodes) + assert [r["pid"] for r in rows] == [3] + assert rows[0]["hbm_gib"] == 4.0 + assert rows[0]["name"] == "leak" + + +# --------------------------------------------------------------------------- +# C. Orchestrator decision-logic tests +# --------------------------------------------------------------------------- + + +def test_node_status_empty_means_pass(): + """C.1 -- no per-GPU results, no Tier 1 issues, no Tier 2 issues + -> empty reasons list -> node PASS.""" + assert _node_status_from([], {}, {}) == [] + + +def test_node_status_nonzero_ecc_fails(): + """C.2 -- any non-zero uncorrectable ECC count is an unconditional + node FAIL. The amd-smi schema is unstable so we trust only ints.""" + tier1 = {"gpu_low_level": {"per_gpu": [{"gpu": 3, "ecc_uncorrectable_total": 7}]}} + reasons = _node_status_from([], tier1, {}) + assert any("gpu3" in r and "uncorrectable" in r for r in reasons) + + +def test_node_status_allow_foreign_procs_downgrades(): + """C.3 -- foreign processes hard-fail the node by default, but + --allow-foreign-procs downgrades to silent inclusion in the JSON.""" + tier1 = { + "gpu_processes": { + "ok": True, + "foreign_count": 1, + "per_gpu": [ + { + "gpu": 0, + "processes": [ + { + "pid": 99, + "name": "leak", + "hbm_bytes": 1 << 30, + "is_foreign": True, + } + ], + } + ], + } + } + # Default (allow_foreign_procs=False) -> FAIL. + assert _node_status_from([], tier1, {}, allow_foreign_procs=False) + # Operator opted in -> PASS. + assert _node_status_from([], tier1, {}, allow_foreign_procs=True) == [] + + +def test_node_status_require_tools_missing_amd_smi_fails(): + """C.4 -- --require-tools promotes a missing CLI tool to a hard + node FAIL; satisfied requirements remain silent.""" + tier1 = { + "tooling_inventory": { + "tools": { + "amd-smi": {"present": False, "path": None}, + "rocm-smi": {"present": True, "path": "/usr/bin/rocm-smi"}, + "lsof": {"present": True, "path": "/usr/bin/lsof"}, + } + } + } + # amd-smi required but absent -> FAIL. + reasons = _node_status_from([], tier1, {}, required_tools=["amd-smi", "rocm-smi"]) + assert any("amd-smi" in r for r in reasons) + # Only rocm-smi required and present -> no reason added. + assert _node_status_from([], tier1, {}, required_tools=["rocm-smi"]) == [] + + +# --------------------------------------------------------------------------- +# D. CLI / report parity tests +# --------------------------------------------------------------------------- + + +# Section ORDER + headings are part of the operator-facing contract -- +# slack bots and CI scripts grep for these. Update this list ONLY when +# you intentionally change the contract. +EXPECTED_SECTIONS = [ + "## Stack drift across cluster", + "## NIC firmware drift across cluster", + "## NIC / RDMA roll-call issues", + "## NIC port-count summary", + "## NIC excluded ports (informational)", + "## Host limits issues", + "## GPU visibility issues", + "## GPU low-level outliers (PCIe link / HBM)", + "## XGMI link issues", + "## Cluster clock + time daemons", + "## Tooling self-latency (`rocm-smi --version`)", + "## Tooling availability", + "## Busy GPUs / leaked processes", + "## GPU pre-touch HBM usage outliers", + "## GPU compute-activity outliers", + # Tier 2 perf summary intentionally omitted -- only renders when at + # least one node ran Tier 2; this fixture has none. + "## Failing nodes -- full reasons", # only when failing is non-empty +] + + +def test_report_section_order_stable(tmp_path): + """D.1 -- report section order + headings are stable. This catches + anyone who reorders the _write_
calls in + aggregator.report.write_smoke_report.""" + nodes = [ + { + "host": "a", + "node_rank": 0, + "status": "FAIL", + "duration_sec": 1.0, + "fail_reasons": ["xgmi: 1 non-XGMI pair"], + "tier1": {}, + "tier2": {}, + } + ] + out = tmp_path / "report.md" + write_smoke_report( + str(out), + nodes=nodes, + passing=[], + failing=nodes, + expected=1, + clock_skew_warn_sec=30.0, + rocm_smi_warn_sec=1.0, + hbm_busy_threshold_gib=2.0, + gpu_activity_warn_pct=20.0, + ) + seen = [line for line in out.read_text().splitlines() if line.startswith("## ")] + assert seen == EXPECTED_SECTIONS + + +def test_module_help_exits_zero(): + """D.2 -- `python -m primus.tools.preflight.node_smoke --help` must + exit 0. Doubles as a smoke test that the package imports cleanly + under `-m` (i.e. __main__.py + __init__.py + cli.py are all + consistent).""" + cp = subprocess.run( + [sys.executable, "-m", "primus.tools.preflight.node_smoke", "--help"], + capture_output=True, + text=True, + timeout=15, + ) + assert cp.returncode == 0, cp.stderr + assert "node-local preflight smoke test" in cp.stdout.lower() + + +# --------------------------------------------------------------------------- +# E. amd-smi schema-drift detection + HBM threshold boundary guards +# --------------------------------------------------------------------------- +# +# These tests lock in two operator-facing contracts: +# +# 1. _flatten_amd_smi_process_json must return (parsed, drift) where +# `drift` is True only when at least one record looked like a process +# dict structurally but had no usable pid/process_id. This lets the +# caller fall through to text/rocm-smi/lsof on a future amd-smi +# schema rename instead of silently reporting the node clean. +# +# 2. The aggregator's pre-touch HBM threshold uses `>=` (inclusive +# boundary). A GPU sitting at exactly the threshold MUST appear in +# the rendered `smoke_report.md` outliers section; the operator- +# facing help / docs / report wording is aligned to "at least". +# (We test the aggregator side because that's the artifact +# operators read; the runner-side `>=` in per_gpu.py is the +# one-line counterpart, kept honest by argparse help + code review.) + + +def _annotate_passthrough(pid, name, hbm): + """Minimal annotate() shim for the schema-drift tests.""" + return {"pid": pid, "name": name, "hbm_bytes": hbm} + + +def test_flatten_amd_smi_top_level_unknown_returns_empty_no_drift(): + """E.1 -- a doc that is neither Shape A/A'/B (no `process_list`, no + pid-bearing items) returns ([], False). False on `drift` is what + distinguishes "schema we do not speak yet at the top level" from + "schema we do speak but per-process fields drifted"; the caller + falls through in either case but the message differs.""" + parsed, drift = _flatten_amd_smi_process_json({"unexpected_key": 1}, _annotate_passthrough) + assert parsed == [] + assert drift is False + parsed, drift = _flatten_amd_smi_process_json([{"some_other_field": 42}], _annotate_passthrough) + assert parsed == [] + assert drift is False + + +def test_flatten_amd_smi_shape_a_empty_process_list_no_drift(): + """E.2 -- Shape A with empty process_list is the clean-node case. + Buckets are pre-registered (so the caller can tell schema-matched + from schema-not-matched), processes lists are empty, drift is False + -- the caller trusts the result and skips the fallback chain.""" + doc = [ + {"gpu": 0, "process_list": []}, + {"gpu": 1, "process_list": []}, + ] + parsed, drift = _flatten_amd_smi_process_json(doc, _annotate_passthrough) + assert parsed == [ + {"gpu": 0, "processes": []}, + {"gpu": 1, "processes": []}, + ] + assert drift is False + + +def test_flatten_amd_smi_shape_a_renamed_pid_field_raises_drift(): + """E.3 -- the future-schema-drift smoking gun: a Shape A entry whose + process_list items are dicts (so the structural shape matches) but + use a renamed key (`proc_pid`) instead of `pid` / `process_id`. + The bucket is registered, processes is empty, but drift=True so the + caller falls through to amd-smi text / rocm-smi / lsof rather than + silently reporting the GPU clean.""" + doc = [ + { + "gpu": 0, + "process_list": [ + {"proc_pid": 1234, "name": "python"}, + ], + } + ] + parsed, drift = _flatten_amd_smi_process_json(doc, _annotate_passthrough) + assert parsed == [{"gpu": 0, "processes": []}] + assert drift is True + + +def test_pretouch_hbm_rows_includes_gpu_at_exact_threshold(): + """E.4 -- the aggregator's `_pretouch_hbm_rows` uses `>= threshold`, + which is what populates the `GPU pre-touch HBM usage outliers` + section in `smoke_report.md`. Lock in the inclusive boundary at the + exact value so the wording fix in `aggregator/report.py` and the + runtime stay aligned: a GPU sitting at threshold MUST be listed, + a GPU below MUST NOT, a GPU above MUST.""" + nodes = [ + { + "host": "host-a", + "node_rank": 0, + "tier1": { + "per_gpu": [ + {"gpu": 0, "details": {"hbm_pre_touch_used_gib": 2.0}}, # == + {"gpu": 1, "details": {"hbm_pre_touch_used_gib": 1.99}}, # < + {"gpu": 2, "details": {"hbm_pre_touch_used_gib": 2.01}}, # > + ], + }, + }, + ] + rows = _pretouch_hbm_rows(nodes, threshold_gib=2.0) + assert sorted(r["gpu"] for r in rows) == [0, 2] + + +# --------------------------------------------------------------------------- +# F. NIC training-NIC selector (NCCL_IB_HCA-style allowlist + heuristic +# fallback + precedence chain) +# +# These guard the operator-facing contract introduced when multi-role +# clusters (front-end / storage RoCE NICs co-resident with the training +# NICs) started showing up. The fail mode they protect against: 56 of 57 +# real production nodes were reported as FAIL because two unplugged +# front-end ports showed `state=DOWN phys_state=Disabled`, even though +# `NCCL_IB_HCA` explicitly listed only the 8 healthy back-end NICs. +# --------------------------------------------------------------------------- + + +def _install_synthetic_ib_tree(monkeypatch, base): + """Redirect `_collect_nic_status`'s sysfs reads at a tmp tree. + + The collector hard-codes ``/sys/class/infiniband`` so we monkey-patch + its ``os.path.isdir`` / ``os.listdir`` / ``_read_text`` symbols to + rewrite that prefix to ``base``. We DELIBERATELY rewrite at the + string level (not via a chroot-style abstraction) because the + rewrite has to also fire for the recursive sub-paths the collector + constructs (``//ports//state`` etc.) -- the rewrite + must therefore intercept every path-shaped call inside the + collector, not just the entry-point listdir. + """ + import os + + real_listdir = os.listdir + real_isdir = os.path.isdir + + def _isdir(p: str) -> bool: + if p == "/sys/class/infiniband": + return True + return real_isdir(p.replace("/sys/class/infiniband", str(base))) + + def _listdir(p: str): + return real_listdir(p.replace("/sys/class/infiniband", str(base))) + + def _read_text(path: str) -> str: + try: + with open(path.replace("/sys/class/infiniband", str(base)), "r") as fh: + return fh.read().strip() + except Exception: + return "" + + monkeypatch.setattr("primus.tools.preflight.node_smoke.collectors.nics.os.path.isdir", _isdir) + monkeypatch.setattr("primus.tools.preflight.node_smoke.collectors.nics.os.listdir", _listdir) + monkeypatch.setattr("primus.tools.preflight.node_smoke.collectors.nics._read_text", _read_text) + + +def test_parse_nic_selector_allowlist_with_ports(): + """F.1 -- baseline NCCL_IB_HCA syntax: comma-separated `device:port`.""" + sel = _parse_nic_selector("rocep158s0:1,rocep190s0:1") + assert sel["mode"] == "allowlist" + assert sel["entries"] == [ + ("rocep158s0", False, 1), + ("rocep190s0", False, 1), + ] + + +def test_parse_nic_selector_denylist_prefix(): + """F.2 -- a leading `^` flips the whole selector to denylist.""" + sel = _parse_nic_selector("^roceo12399,roceo12409") + assert sel["mode"] == "denylist" + # The ^ is stripped from the entries; only the global mode changes. + assert [e[0] for e in sel["entries"]] == ["roceo12399", "roceo12409"] + + +def test_parse_nic_selector_exact_prefix(): + """F.3 -- `=name` forces exact (not prefix) device-name matching.""" + sel = _parse_nic_selector("=mlx5,mlx5_other") + assert sel["entries"][0] == ("mlx5", True, None) # exact-match flag set + assert sel["entries"][1] == ("mlx5_other", False, None) + + +def test_parse_nic_selector_no_port_matches_any_port(): + """F.4 -- entries without `:port` accept any port on the device.""" + sel = _parse_nic_selector("mlx5_0") + assert sel["entries"] == [("mlx5_0", False, None)] + assert _selector_matches(sel, "mlx5_0", 1) is True + assert _selector_matches(sel, "mlx5_0", 2) is True + + +def test_parse_nic_selector_empty_is_passthrough(): + """F.5 -- empty / whitespace-only input means 'no selector', so the + caller falls through to the next precedence layer (env / heuristic) + rather than treating it as 'allowlist nothing' (which would silently + exclude every port).""" + assert _parse_nic_selector("")["mode"] == "passthrough" + assert _parse_nic_selector(" ")["mode"] == "passthrough" + assert _parse_nic_selector("^")["mode"] == "passthrough" # ^ but no entries + assert _parse_nic_selector(",, ,")["mode"] == "passthrough" + + +def test_selector_matches_prefix_matches_mlx5_to_mlx5_0(): + """F.6 -- NCCL prefix semantics: `mlx5` allowlist entry matches every + device that starts with `mlx5` (the historical NCCL behavior). This + is why operators must use `=mlx5` if they want exact-only matching.""" + sel = _parse_nic_selector("mlx5") + assert _selector_matches(sel, "mlx5", 1) is True + assert _selector_matches(sel, "mlx5_0", 1) is True + assert _selector_matches(sel, "mlx5_bond_0", 1) is True + assert _selector_matches(sel, "rocep158s0", 1) is False + + +def test_selector_matches_denylist_inverts(): + """F.7 -- denylist semantics: ports NOT in the list pass; ports in + the list are rejected. This is the more ergonomic form for clusters + with many training NICs and a small handful of front-end ports.""" + sel = _parse_nic_selector("^roceo12399,roceo12409") + # Listed -> excluded. + assert _selector_matches(sel, "roceo12399", 1) is False + assert _selector_matches(sel, "roceo12409", 1) is False + # Not listed -> included. + assert _selector_matches(sel, "rocep158s0", 1) is True + + +def test_selector_matches_port_filter_is_per_port(): + """F.8 -- `:port` suffix narrows the match to that specific port on + the device; the same device on a different port is NOT matched.""" + sel = _parse_nic_selector("mlx5_0:2") + assert _selector_matches(sel, "mlx5_0", 2) is True + assert _selector_matches(sel, "mlx5_0", 1) is False + + +def test_resolve_selector_cli_beats_env(): + """F.9 -- precedence: when both the CLI flag and NCCL_IB_HCA env are + set, the CLI wins. Operators must be able to override their shell + env without touching it.""" + sel = _resolve_selector( + allowlist_arg="mlx5_0:1", + env={"NCCL_IB_HCA": "rocep158s0:1"}, + ) + assert sel["source"] == "cli" + assert sel["entries"][0][0] == "mlx5_0" + + +def test_resolve_selector_env_used_when_cli_absent(): + """F.10 -- precedence: env is consulted only when the CLI flag is + unset. This is the most common case in production (operator sets + NCCL_IB_HCA once in their job script; smoke picks it up + transparently).""" + sel = _resolve_selector( + allowlist_arg=None, + env={"NCCL_IB_HCA": "rocep158s0:1"}, + ) + assert sel["source"] == "env" + + +def test_resolve_selector_heuristic_when_both_absent(): + """F.11 -- precedence fallback: when neither CLI nor env is set, + return a `heuristic` marker. The caller (`_collect_nic_status`) + handles this by auto-excluding ports whose `phys_state` is in + {Disabled, Sleep}.""" + sel = _resolve_selector(allowlist_arg=None, env={}) + assert sel["source"] == "heuristic" + # Document the heuristic's exact admin-down set so a future widening + # (e.g. to also exclude `Polling`) has to update this test + # deliberately. `Polling` MUST NOT be in here -- it means + # "actively looking for a link partner" which is a real failure on + # a port intended to be used. + assert sel["admin_down_phys_states"] == ["disabled", "sleep"] + + +def test_resolve_selector_empty_string_falls_through_to_env(): + """F.12 -- an empty `--rdma-nic-allowlist ''` must not silently + block out everything; it should behave the same as not passing the + flag at all and fall through to the env. Defensive against + shell-quoting accidents (`--rdma-nic-allowlist "$VAR"` when + `$VAR` is unset).""" + sel = _resolve_selector( + allowlist_arg="", + env={"NCCL_IB_HCA": "rocep158s0:1"}, + ) + assert sel["source"] == "env" + + +def test_resolve_selector_blank_env_falls_through_to_heuristic(): + """F.13 -- same defensive behavior for an empty NCCL_IB_HCA env + (some clusters export it unset / empty by accident). MUST fall + through to the heuristic rather than allowlist-nothing.""" + sel = _resolve_selector(allowlist_arg=None, env={"NCCL_IB_HCA": ""}) + assert sel["source"] == "heuristic" + + +def test_collect_nic_status_env_allowlist_excludes_disabled_frontend_ports(tmp_path, monkeypatch): + """F.14 -- end-to-end against a synthetic /sys/class/infiniband + mirroring the production failure mode: 12 IB devices (8 training + + 2 storage ACTIVE + 2 frontend Disabled). With NCCL_IB_HCA listing + only the 8 training NICs, the collector must: + * include exactly the 8 listed devices, + * exclude the other 4 with a source=`env` info_issue each, + * emit ZERO hard issues, + * tag the Disabled ports with their state in the info_issue + (so an operator who reads excluded_ports can still see they + were unhealthy, just not relevant). + """ + from primus.tools.preflight.node_smoke.collectors.nics import _collect_nic_status + + # Synthetic sysfs tree. + base = tmp_path / "ib" + devices = [ + # (name, state, phys_state, rate_str) + ("rocep158s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + ("rocep190s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + ("rocep206s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + ("rocep222s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + ("rocep28s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + ("rocep62s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + ("rocep79s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + ("rocep96s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + # Storage / control-plane: ACTIVE but NOT in NCCL_IB_HCA. + ("rocep159s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + ("rocep29s0", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + # Front-end: Disabled, NOT in NCCL_IB_HCA -- the user's bug. + ("roceo12399", "1: DOWN", "3: Disabled", ""), + ("roceo12409", "1: DOWN", "3: Disabled", ""), + ] + + for name, state, phys, rate in devices: + port_dir = base / name / "ports" / "1" + port_dir.mkdir(parents=True) + (port_dir / "state").write_text(state) + (port_dir / "phys_state").write_text(phys) + (port_dir / "rate").write_text(rate) + (port_dir / "link_layer").write_text("Ethernet") + # One valid RoCE v2 GID on every ACTIVE port so the GID rule + # doesn't accidentally fail the test (which would mask the + # selector behavior we're trying to verify). + gids = port_dir / "gids" + gids.mkdir() + (gids / "0").write_text("fe80:0000:0000:0000:0000:0000:0000:0001") + types = port_dir / "gid_attrs" / "types" + types.mkdir(parents=True) + (types / "0").write_text("IB/RoCE v2") + + _install_synthetic_ib_tree(monkeypatch, base) + + monkeypatch.setenv( + "NCCL_IB_HCA", + "rocep158s0:1,rocep190s0:1,rocep206s0:1,rocep222s0:1," + "rocep28s0:1,rocep62s0:1,rocep79s0:1,rocep96s0:1", + ) + + out = _collect_nic_status(expected_count=None) + + # 12 ports total, 8 included, 4 excluded -- the headline assertion. + assert len(out["ports"]) == 12 + assert len(out["included_ports"]) == 8 + assert len(out["excluded_ports"]) == 4 + # The 2 Disabled frontend ports and the 2 not-in-HCA storage ports. + assert set(out["excluded_ports"]) == { + "rocep159s0:1", + "rocep29s0:1", + "roceo12399:1", + "roceo12409:1", + } + # Selector metadata records the source so the operator can verify + # which precedence layer fired. + assert out["selector"]["source"] == "env" + # Zero hard issues -> node would PASS. + assert out["issues"] == [] + # Every excluded port produced an info_issue mentioning the env + # source. The two Disabled ports MUST also carry their state info + # so operators investigating "are my frontend NICs still down?" can + # see it without opening every per-port record. + info = "\n".join(out["info_issues"]) + assert "NCCL_IB_HCA" in info + assert "phys_state=Disabled" in info # frontend ports + # ACTIVE-but-not-in-HCA ports do NOT need the state suffix (they're + # healthy, just reserved for sockets/storage). + + +def test_collect_nic_status_heuristic_only_excludes_disabled_phys_state(tmp_path, monkeypatch): + """F.15 -- with no env and no CLI selector, the heuristic must + auto-exclude `phys_state=Disabled` ports but MUST keep `phys_state= + Polling` (cable unplugged on an intended-up port) and `phys_state= + LinkUp with state=INIT` (driver/SM didn't finish bringup) in the + included set so they hard-fail. The whole point of choosing Disabled + as the exclusion signal is that it's the only phys_state that + unambiguously means "admin-down, not used".""" + from primus.tools.preflight.node_smoke.collectors.nics import _collect_nic_status + + base = tmp_path / "ib" + devices = [ + ("training_ok", "4: ACTIVE", "5: LinkUp", "400 Gb/sec"), + ("frontend_disabled", "1: DOWN", "3: Disabled", ""), + ("training_cable_pulled", "1: DOWN", "2: Polling", ""), + ("training_unconfigured", "2: INIT", "5: LinkUp", "400 Gb/sec"), + ] + for name, state, phys, rate in devices: + port_dir = base / name / "ports" / "1" + port_dir.mkdir(parents=True) + (port_dir / "state").write_text(state) + (port_dir / "phys_state").write_text(phys) + (port_dir / "rate").write_text(rate) + (port_dir / "link_layer").write_text("Ethernet") + gids = port_dir / "gids" + gids.mkdir() + (gids / "0").write_text("fe80:0000:0000:0000:0000:0000:0000:0001") + types = port_dir / "gid_attrs" / "types" + types.mkdir(parents=True) + (types / "0").write_text("IB/RoCE v2") + + _install_synthetic_ib_tree(monkeypatch, base) + monkeypatch.delenv("NCCL_IB_HCA", raising=False) + + out = _collect_nic_status(expected_count=None) + + assert out["selector"]["source"] == "heuristic" + # Only the Disabled port is excluded by the heuristic. + assert out["excluded_ports"] == ["frontend_disabled:1"] + # The 3 remaining are included; the 2 broken-but-included ones produce + # hard issues. + assert set(out["included_ports"]) == { + "training_ok:1", + "training_cable_pulled:1", + "training_unconfigured:1", + } + issues_text = " ".join(out["issues"]) + assert "training_cable_pulled" in issues_text # state=DOWN + assert "training_unconfigured" in issues_text # state=INIT + assert "training_ok" not in issues_text + + +def test_collect_nic_status_empty_set_guard_when_everything_excluded(tmp_path, monkeypatch): + """F.16 -- defense in depth: if the selector ends up excluding every + discovered port (e.g. operator typo'd --rdma-nic-allowlist, or the + whole RoCE card got admin-disabled), the node MUST still hard-fail. + A node with zero training NICs cannot participate in inter-node + training and silently passing it would be the worst possible + regression introduced by the selector feature.""" + from primus.tools.preflight.node_smoke.collectors.nics import _collect_nic_status + + base = tmp_path / "ib" + port_dir = base / "training_ok" / "ports" / "1" + port_dir.mkdir(parents=True) + (port_dir / "state").write_text("4: ACTIVE") + (port_dir / "phys_state").write_text("5: LinkUp") + (port_dir / "rate").write_text("400 Gb/sec") + (port_dir / "link_layer").write_text("Ethernet") + + _install_synthetic_ib_tree(monkeypatch, base) + # CLI allowlist that matches NOTHING (typo or operator mistake). + out = _collect_nic_status(expected_count=None, allowlist="this_device_does_not_exist:1") + + assert out["included_ports"] == [] + assert out["excluded_ports"] == ["training_ok:1"] + # The empty-set guard fires -- node FAIL. + assert any("no included RDMA NIC ports" in issue for issue in out["issues"]), out["issues"] + + +def test_collect_nic_status_expected_count_compares_included(tmp_path, monkeypatch): + """F.17 -- the behavior change for --expected-rdma-nics: it must + compare against the *included* count, not the total /sys/class/ + infiniband count. Otherwise `--expected-rdma-nics 8` would fail on + the very clusters this feature exists to help (12 devices, 8 + training).""" + from primus.tools.preflight.node_smoke.collectors.nics import _collect_nic_status + + base = tmp_path / "ib" + # 8 training + 2 frontend Disabled = 10 total ports, 8 training NICs. + for i in range(8): + d = base / f"trainnic{i}" / "ports" / "1" + d.mkdir(parents=True) + (d / "state").write_text("4: ACTIVE") + (d / "phys_state").write_text("5: LinkUp") + (d / "rate").write_text("400 Gb/sec") + (d / "link_layer").write_text("Ethernet") + gids = d / "gids" + gids.mkdir() + (gids / "0").write_text("fe80:0000:0000:0000:0000:0000:0000:0001") + types = d / "gid_attrs" / "types" + types.mkdir(parents=True) + (types / "0").write_text("IB/RoCE v2") + for i in range(2): + d = base / f"frontnic{i}" / "ports" / "1" + d.mkdir(parents=True) + (d / "state").write_text("1: DOWN") + (d / "phys_state").write_text("3: Disabled") + (d / "rate").write_text("") + (d / "link_layer").write_text("Ethernet") + + _install_synthetic_ib_tree(monkeypatch, base) + monkeypatch.delenv("NCCL_IB_HCA", raising=False) + + # With --expected-rdma-nics 8 on the heuristic path: the 2 Disabled + # ports get auto-excluded, leaving 8 included = matches expected. + out = _collect_nic_status(expected_count=8) + assert out["issues"] == [], out["issues"] + # And the inverse: --expected-rdma-nics 10 (treating the total + # `/sys/class/infiniband` count as the expected) MUST fail, because + # the comparison happens against the included set. + out = _collect_nic_status(expected_count=10) + assert any("!= expected 10" in i for i in out["issues"]) + + +# --------------------------------------------------------------------------- +# Primus-cli subcommand wrapper (consolidate-preflight-direct-wrappers). +# These tests pin the contract for the new `primus-cli node_smoke` +# subcommand layer: hoisted flags, abbreviation rejection, no Python-side +# --silent, two-phase dispatch (run on all ranks, aggregate on rank 0), +# and SLURM-aware aggregator-arg resolution. +# --------------------------------------------------------------------------- + + +def _build_primus_cli_node_smoke_parser(): + """Helper: build a parser that mirrors how primus.cli.main wires the + `node_smoke` subcommand, without going through the full main() + dispatch. Centralized here so individual tests don't redo the wiring. + """ + import argparse as _argparse + + from primus.cli.subcommands.node_smoke import register_subcommand + + top = _argparse.ArgumentParser(prog="primus") + sub = top.add_subparsers(dest="command", required=True) + register_subcommand(sub) + return top + + +def test_primus_cli_node_smoke_known_flags_parse_cleanly(): + """`primus-cli node_smoke --tier2-perf --hbm-busy-threshold-gib 3.5 ...` + must parse without unknown-args; the namespace must merge the run + surface with the aggregator surface in a single object.""" + parser = _build_primus_cli_node_smoke_parser() + ns, unknown = parser.parse_known_args( + ["node_smoke", "--tier2-perf", "--hbm-busy-threshold-gib", "3.5", "--expected-nodes", "4"] + ) + assert unknown == [] + assert ns.command == "node_smoke" + # Run-side knob + assert ns.tier2_perf is True + assert ns.hbm_busy_threshold_gib == 3.5 + # Aggregator-side knob hoisted to the same namespace + assert ns.expected_nodes == 4 + # Run-side default preserved (we only attach --dump-path once via _add_run_flags) + assert ns.dump_path == "output/preflight" + + +def test_primus_cli_node_smoke_rejects_silent(): + """`--silent` is a bash-launcher knob only; the python subcommand + must surface it as an unknown argument so a misplaced --silent never + silently disappears into argparse.""" + parser = _build_primus_cli_node_smoke_parser() + _ns, unknown = parser.parse_known_args(["node_smoke", "--silent"]) + assert "--silent" in unknown, ( + "primus-cli node_smoke must NOT accept --silent (silencing lives " + "exclusively in primus-cli-direct.sh's bash layer; this test " + "guards against dual handling regressing)" + ) + + +def test_primus_cli_node_smoke_rejects_abbreviated_tier2(): + """`--tier2` is an abbreviation of `--tier2-perf` -- argparse's + `allow_abbrev=False` must make this a hard error so a stale script + using the old flag name doesn't silently match the new one.""" + import argparse as _argparse + + parser = _build_primus_cli_node_smoke_parser() + # parse_known_args with `allow_abbrev=False` returns --tier2 in + # `unknown`; the main() dispatcher then rejects unknown args with + # SystemExit(2). The subparser-level guarantee we care about here is + # that --tier2 NEVER ends up bound to tier2_perf=True via prefix + # matching. + ns, unknown = parser.parse_known_args(["node_smoke", "--tier2"]) + assert "--tier2" in unknown + assert ns.tier2_perf is False, "--tier2 must NOT silently match --tier2-perf via prefix abbreviation" + + # Defense-in-depth: parse_args (strict mode) raises SystemExit. + with pytest.raises(SystemExit): + parser.parse_args(["node_smoke", "--tier2"]) + _ = _argparse # silence unused + + +def test_resolve_aggregate_args_from_slurm_uses_slurm_nnodes(monkeypatch): + """When `--expected-nodes` is not supplied, the helper must fall back + to `SLURM_NNODES` (then `SLURM_JOB_NUM_NODES`, then `NNODES`) so the + aggregator correctly identifies missing nodes.""" + import argparse as _argparse + + from primus.tools.preflight.node_smoke.cli import _resolve_aggregate_args_from_slurm + + monkeypatch.setenv("SLURM_NNODES", "6") + monkeypatch.delenv("SLURM_JOB_NUM_NODES", raising=False) + monkeypatch.delenv("NNODES", raising=False) + monkeypatch.delenv("SLURM_JOB_NODELIST", raising=False) + + args = _argparse.Namespace( + dump_path="output/preflight", + expected_nodes=None, + wait_timeout_sec=60, + rocm_smi_warn_sec=1.0, + clock_skew_warn_sec=30.0, + hbm_busy_threshold_gib=2.0, + gpu_activity_warn_pct=20.0, + expected_nodelist_file=None, + ) + agg = _resolve_aggregate_args_from_slurm(args) + assert agg.expected_nodes == 6 + assert agg.dump_path == "output/preflight" + # No SLURM_JOB_NODELIST + no pre-supplied file -> stays None (the + # aggregator falls back to `` placeholders, by design). + assert agg.expected_nodelist_file is None + + +def test_resolve_aggregate_args_from_slurm_cli_wins(monkeypatch): + """User-supplied `--expected-nodes` takes precedence over SLURM_*.""" + import argparse as _argparse + + from primus.tools.preflight.node_smoke.cli import _resolve_aggregate_args_from_slurm + + monkeypatch.setenv("SLURM_NNODES", "6") + args = _argparse.Namespace( + dump_path="output/preflight", + expected_nodes=4, # user-set wins + wait_timeout_sec=60, + rocm_smi_warn_sec=1.0, + clock_skew_warn_sec=30.0, + hbm_busy_threshold_gib=2.0, + gpu_activity_warn_pct=20.0, + expected_nodelist_file=None, + ) + agg = _resolve_aggregate_args_from_slurm(args) + assert agg.expected_nodes == 4 + + +def test_primus_cli_node_smoke_two_phase_dispatch_rank0(monkeypatch): + """On rank 0 the subcommand must call `_cmd_run` AND `_cmd_aggregate`, + propagate the max(rc_run, rc_agg) exit code, and pass the SLURM- + resolved aggregator args (not the raw `args`) to the aggregator.""" + import argparse as _argparse + + from primus.cli.subcommands import node_smoke as smoke_cmd + + monkeypatch.setenv("NODE_RANK", "0") + monkeypatch.delenv("SLURM_NODEID", raising=False) + monkeypatch.setenv("SLURM_NNODES", "5") + monkeypatch.delenv("SLURM_JOB_NODELIST", raising=False) + + calls = [] + + def fake_run(ns): + calls.append(("run", ns)) + return 0 + + def fake_agg(ns): + calls.append(("aggregate", ns)) + return 1 # simulate one MISSING node + + monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_run", fake_run, raising=True) + monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_aggregate", fake_agg, raising=True) + + args = _argparse.Namespace( + dump_path="/tmp/test-smoke", + expected_nodes=None, + wait_timeout_sec=60, + rocm_smi_warn_sec=1.0, + clock_skew_warn_sec=30.0, + hbm_busy_threshold_gib=2.0, + gpu_activity_warn_pct=20.0, + expected_nodelist_file=None, + ) + with pytest.raises(SystemExit) as exc: + smoke_cmd.run(args, extra_args=[]) + assert exc.value.code == 1, "aggregator's FAIL must surface as the exit code" + + # Run AND aggregate both fired, in that order. + assert [c[0] for c in calls] == ["run", "aggregate"] + # Aggregate received the SLURM-resolved namespace (expected_nodes + # auto-filled from SLURM_NNODES=5), not the user-supplied args. + agg_ns = calls[1][1] + assert agg_ns.expected_nodes == 5 + # And the aggregator namespace is a *different* object from `args`. + assert agg_ns is not args + + +def test_primus_cli_node_smoke_two_phase_dispatch_non_rank0(monkeypatch): + """Non-rank-0 ranks must run `_cmd_run` only and propagate its exit + code; calling `_cmd_aggregate` on every rank would race on the + shared output directory.""" + import argparse as _argparse + + from primus.cli.subcommands import node_smoke as smoke_cmd + + monkeypatch.setenv("NODE_RANK", "3") + monkeypatch.delenv("SLURM_NODEID", raising=False) + + calls = [] + + def fake_run(ns): + calls.append("run") + return 0 + + def fake_agg(ns): # MUST NOT be called on a non-rank-0 rank. + calls.append("aggregate") + raise AssertionError("aggregator must not run on non-rank-0") + + monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_run", fake_run, raising=True) + monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_aggregate", fake_agg, raising=True) + + args = _argparse.Namespace( + dump_path="/tmp/test-smoke", + expected_nodes=None, + wait_timeout_sec=60, + rocm_smi_warn_sec=1.0, + clock_skew_warn_sec=30.0, + hbm_busy_threshold_gib=2.0, + gpu_activity_warn_pct=20.0, + expected_nodelist_file=None, + ) + with pytest.raises(SystemExit) as exc: + smoke_cmd.run(args, extra_args=[]) + assert exc.value.code == 0 + assert calls == ["run"] + + +def test_primus_cli_node_smoke_rank0_run_failure_surfaces(monkeypatch): + """If `_cmd_run` fails on rank 0 but the aggregator (which only knows + about missing-from-`/smoke/*.json`) reports success, the + subcommand must still surface the rank-0 run failure -- otherwise a + sick rank-0 node could paint itself green via a successful aggregate.""" + import argparse as _argparse + + from primus.cli.subcommands import node_smoke as smoke_cmd + + monkeypatch.setenv("NODE_RANK", "0") + monkeypatch.delenv("SLURM_NODEID", raising=False) + monkeypatch.delenv("SLURM_NNODES", raising=False) + monkeypatch.delenv("NNODES", raising=False) + monkeypatch.delenv("SLURM_JOB_NODELIST", raising=False) + + monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_run", lambda _ns: 1, raising=True) + monkeypatch.setattr("primus.tools.preflight.node_smoke.cli._cmd_aggregate", lambda _ns: 0, raising=True) + + args = _argparse.Namespace( + dump_path="/tmp/test-smoke", + expected_nodes=None, + wait_timeout_sec=60, + rocm_smi_warn_sec=1.0, + clock_skew_warn_sec=30.0, + hbm_busy_threshold_gib=2.0, + gpu_activity_warn_pct=20.0, + expected_nodelist_file=None, + ) + with pytest.raises(SystemExit) as exc: + smoke_cmd.run(args, extra_args=[]) + assert exc.value.code == 1, "rank-0 _cmd_run=1 must not be masked by _cmd_aggregate=0" diff --git a/primus/tools/preflight/node_smoke/types.py b/primus/tools/preflight/node_smoke/types.py new file mode 100644 index 000000000..97f1f4ae5 --- /dev/null +++ b/primus/tools/preflight/node_smoke/types.py @@ -0,0 +1,36 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Shared dataclasses for the node-smoke pipeline.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Dict, List + + +@dataclass +class GPUResult: + """Result of all checks for a single GPU on this node.""" + + gpu: int + status: str # "PASS" | "FAIL" | "TIMEOUT" + reason: str = "" + duration_sec: float = 0.0 + details: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class NodeResult: + """Whole-node verdict written to ``/smoke/.json``.""" + + host: str + node_rank: int + status: str # "PASS" | "FAIL" + duration_sec: float + fail_reasons: List[str] + tier1: Dict[str, Any] + tier2: Dict[str, Any] diff --git a/primus/tools/preflight/preflight_args.py b/primus/tools/preflight/preflight_args.py index a816b2b47..1710e14b4 100644 --- a/primus/tools/preflight/preflight_args.py +++ b/primus/tools/preflight/preflight_args.py @@ -12,18 +12,49 @@ import argparse +# Canonical perf-test tokens accepted by --tests. +PERF_TEST_TOKENS = ( + "gemm", + "intra-allreduce", + "intra-alltoall", + "inter-allreduce", + "inter-alltoall", + "inter-p2p", + "inter-ring-p2p", +) + +# Names of the "intent-bearing" perf flags. Setting any of these implies perf +# mode (no need to also pass --perf-test). When mixed with info selectors +# (--host/--gpu/--network), perf wins and the info selectors are dropped with +# a warning. +PERF_INTENT_FLAGS = ("--perf-test", "--tests", "--quick") + def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: """ Register arguments for `primus-cli preflight`. + Mode precedence: + 1. Any of --perf-test / --tests / --quick is set -> perf mode wins. + If info selectors (--host/--gpu/--network) are also set, they are + dropped with a warning. Perf-only tuning knobs (--comm-sizes-mb, + --intra-group-sizes, etc.) take effect. + 2. Otherwise, any of --host/--gpu/--network is set -> info-only mode. + Perf-only tuning knobs, if set, are inert and a WARN is emitted. + 3. Otherwise (no flags) -> default: run info AND all perf tests. + Usage: - primus-cli preflight # Show all info (Host + GPU + Network) - primus-cli preflight --host # Host only - primus-cli preflight --gpu # GPU only - primus-cli preflight --network # Network only - primus-cli preflight --gpu --network # GPU + Network - primus-cli preflight --perf-test # Run perf tests ONLY (skip info) + primus-cli preflight # Default: info + all perf + primus-cli preflight --host # Host info only + primus-cli preflight --gpu # GPU info only + primus-cli preflight --network # Network info only + primus-cli preflight --gpu --network # GPU + Network info + primus-cli preflight --perf-test # Perf only, all tests + primus-cli preflight --quick # Perf only, fast preset + primus-cli preflight --tests gemm # Perf only, GEMM only + primus-cli preflight --tests gemm,inter-allreduce \\ + --comm-sizes-mb 64,1024 \\ + --inter-group-sizes all """ # Check selection flags # Keep --check-* as compatibility aliases. @@ -54,11 +85,98 @@ def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPa "--perf-test", action="store_true", help="Run perf tests ONLY (GEMM, intra/inter node communication). " - "This is slower and skips the host/gpu/network info report.", + "Skips the host/gpu/network info report. Implied by --tests/--quick.", + ) + + # Performance test specific options. + parser.add_argument( + "--plot", + action="store_true", + help="Generate plots (perf mode only).", + ) + + # Test selection (CSV). Tokens: gemm,intra-allreduce,intra-alltoall, + # inter-allreduce,inter-alltoall,inter-p2p,inter-ring-p2p, or 'all'. + parser.add_argument( + "--tests", + type=str, + default=None, + help="Comma-separated list of perf tests to run. Tokens: " + "gemm, intra-allreduce, intra-alltoall, inter-allreduce, inter-alltoall, " + "inter-p2p, inter-ring-p2p, all. Implies --perf-test. " + "When unset, runs every test.", + ) + + # Message size config (CSV in MB) for comm tests. + parser.add_argument( + "--comm-sizes-mb", + type=str, + default=None, + help="Default message sizes (CSV in MB) used for intra-/inter-node " + "allreduce, alltoall, and inter-node p2p tests when no specific override is given. " + "Default: 2,4,8,16,32,64,128,256,512,1024.", + ) + parser.add_argument( + "--intra-comm-sizes-mb", + type=str, + default=None, + help="Override message sizes (CSV in MB) for intra-node allreduce/alltoall. " + "Falls back to --comm-sizes-mb when unset.", + ) + parser.add_argument( + "--inter-comm-sizes-mb", + type=str, + default=None, + help="Override message sizes (CSV in MB) for inter-node allreduce/alltoall/p2p. " + "Falls back to --comm-sizes-mb when unset.", ) - # Performance test specific options (only used with --perf-test) - parser.add_argument("--plot", action="store_true", help="Generate plots (only with --perf-test)") + # Group size config. + parser.add_argument( + "--intra-group-sizes", + type=str, + default=None, + help="Comma-separated list of intra-node GPU group sizes to test " + "(each must divide LOCAL_WORLD_SIZE). Default: 2,4,8.", + ) + parser.add_argument( + "--inter-group-sizes", + type=str, + default=None, + help="Comma-separated list of inter-node group sizes to test. Use 'all' for " + "the full N-node group. Default: 2,4,all. Note: for inter-node alltoall " + "each value is clamped to 16 (real-world MoE training rarely dispatches " + "across more nodes); other inter-node tests use the requested sizes " + "unchanged. See docs/preflight.md \u00a75.2 for details.", + ) + + # Inter-node ring p2p sizes. + parser.add_argument( + "--ring-p2p-sizes-mb", + type=str, + default=None, + help="Message sizes (CSV in MB) for the inter-node ring P2P test. " "Default: 10,20,40,80,160.", + ) + + # Quick preset. + parser.add_argument( + "--quick", + action="store_true", + help="Fast pre-launch preset. Implies --perf-test. Selects gemm + " + "intra-allreduce + inter-allreduce, uses sizes 64,1024 MB, full " + "intra-node group only, full N-node inter-node group only, and lowers " + "warmup/iterations. User-supplied flags override.", + ) + + # Back-compat alias: kept so existing scripts keep working. + # Internally it maps to --inter-group-sizes all and disables inter-p2p. + parser.add_argument( + "--no-split-nodes-subgroup", + dest="split_nodes_subgroup", + action="store_false", + help="[Deprecated] Skip inter-node comm tests on node subgroups (2-node, 4-node). " + "Equivalent to --inter-group-sizes all and dropping inter-p2p.", + ) # Distributed init timeout (prevents hangs when network/rendezvous is misconfigured) parser.add_argument( @@ -69,6 +187,24 @@ def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPa "If init times out, preflight will write the info report and exit with failure.", ) + # Communicator cleanup delay (cross-rank sync across destroy -> setup; + # the actual "Address already in use" defense is the inter-alltoall cap + # in --inter-group-sizes / inter_node_comm.py, not this delay). + parser.add_argument( + "--comm-cleanup-delay-sec", + type=float, + default=2.0, + help="Delay (seconds) inserted between destroying NCCL/RCCL process " + "groups and creating new ones. Provides cross-rank synchronization " + "across the destroy/setup transition and gives the kernel a moment " + "to unlink closed-socket bookkeeping. Set to 0 to disable the sleep " + "(barrier only). The actual 'Address already in use' defense at " + "scale is the inter-node alltoall sub-group cap of 16 (see " + "--inter-group-sizes); widening net.ipv4.ip_local_port_range " + "(docs/preflight.md \u00a77.2) is the directly relevant OS knob if " + "you ever need more headroom.", + ) + # Report output options parser.add_argument( "--dump-path", @@ -79,8 +215,13 @@ def add_preflight_parser(parser: argparse.ArgumentParser) -> argparse.ArgumentPa parser.add_argument( "--report-file-name", type=str, - default="preflight_report", - help="Base name for report files (default: preflight_report).", + default=None, + help=( + "Base name for report files. When omitted, an auto-generated " + "timestamped name of the form 'preflight-{NNODES}N-{YYYYMMDD-HHMMSS}' " + "is used so each run writes to a fresh path and never overwrites " + "or collides with stale leftovers." + ), ) parser.add_argument( "--disable-pdf", diff --git a/primus/tools/preflight/preflight_perf_test.py b/primus/tools/preflight/preflight_perf_test.py index e3eb2607a..636dee728 100644 --- a/primus/tools/preflight/preflight_perf_test.py +++ b/primus/tools/preflight/preflight_perf_test.py @@ -8,11 +8,20 @@ import os import socket import sys +import time from dataclasses import asdict, dataclass from datetime import datetime -from typing import Any, Dict, List - -from primus.tools.preflight.global_vars import LOCAL_RANK, RANK, set_hostnames +from typing import Any, Dict, List, Optional, Set + +from primus.tools.preflight.global_vars import ( + LOCAL_RANK, + LOCAL_WORLD_SIZE, + RANK, + WORLD_SIZE, + set_hostnames, + set_iteration, + set_warmup, +) from primus.tools.preflight.gpu.info import collect_gpu_info, write_gpu_report from primus.tools.preflight.host.info import collect_host_info, write_host_report from primus.tools.preflight.inter_node_comm import run_inter_node_comm @@ -23,6 +32,7 @@ collect_network_info, write_network_report, ) +from primus.tools.preflight.preflight_args import PERF_TEST_TOKENS from primus.tools.preflight.square_gemm import run_square_gemm from primus.tools.preflight.utility import ( gather_hostnames, @@ -34,6 +44,115 @@ from primus.tools.utils import gather_records, get_rank_world +def _parse_csv_int_list(value: Optional[str], name: str) -> List[int]: + """Parse a CSV of positive ints. Empty / None -> [].""" + if value is None: + return [] + s = str(value).strip() + if not s: + return [] + out: List[int] = [] + for tok in s.split(","): + tok = tok.strip() + if not tok: + continue + try: + v = int(tok) + except ValueError: + raise ValueError(f"{name}: '{tok}' is not an integer") + if v <= 0: + raise ValueError(f"{name}: values must be positive (got {v})") + out.append(v) + return out + + +def _parse_csv_inter_group_list(value: Optional[str], name: str) -> List[Any]: + """Parse a CSV that may include the special token 'all'.""" + if value is None: + return [] + s = str(value).strip() + if not s: + return [] + out: List[Any] = [] + for tok in s.split(","): + tok = tok.strip() + if not tok: + continue + if tok.lower() == "all": + out.append("all") + continue + try: + v = int(tok) + except ValueError: + raise ValueError(f"{name}: '{tok}' is not an integer or 'all'") + if v <= 0: + raise ValueError(f"{name}: values must be positive (got {v})") + out.append(v) + return out + + +def _write_node_hostname_legend(markdown_file: str) -> None: + """Write a Node -> Hostname legend at the top of the perf markdown report + and mirror it to the console on rank 0. + + Uses the per-rank hostnames already gathered via `set_hostnames()`, picking + `LOCAL_RANK == 0` of each node as the canonical hostname for that node. + """ + from primus.tools.preflight.global_vars import get_hostnames + + hostnames = get_hostnames() + if not hostnames: + return + num_nodes = WORLD_SIZE // LOCAL_WORLD_SIZE + + log("=======Nodes=======") + log(f"{'Node':<6} Hostname") + with open(markdown_file, "a", encoding="utf-8") as f: + f.write("# Nodes\n\n") + f.write("| Node | Hostname |\n") + f.write("|------|----------|\n") + for n in range(num_nodes): + rank = n * LOCAL_WORLD_SIZE + host = hostnames[rank] if rank < len(hostnames) else "" + f.write(f"| {n} | {host} |\n") + log(f"{n:<6} {host}") + f.write("\n") + log("") + + +def _parse_perf_tests(value: Optional[str]) -> Set[str]: + """Parse --tests CSV into a set of canonical tokens. + + None / '' / 'all' -> every token. A non-empty value that yields zero valid + tokens (e.g. ',,,') raises ValueError so users notice the typo instead of + silently running no perf tests. + """ + if value is None: + return set(PERF_TEST_TOKENS) + raw = str(value) + s = raw.strip().lower() + if not s or s == "all": + return set(PERF_TEST_TOKENS) + selected: Set[str] = set() + for tok in s.split(","): + tok = tok.strip() + if not tok: + continue + if tok == "all": + return set(PERF_TEST_TOKENS) + if tok not in PERF_TEST_TOKENS: + raise ValueError( + f"--tests: unknown token '{tok}'. " f"Valid tokens: {', '.join(PERF_TEST_TOKENS)}, all" + ) + selected.add(tok) + if not selected: + raise ValueError( + f"--tests: no valid tokens parsed from {raw!r}. " + f"Valid tokens: {', '.join(PERF_TEST_TOKENS)}, all" + ) + return selected + + @dataclass class Finding: level: str # "info" | "warn" | "fail" @@ -49,6 +168,90 @@ def _status_from_counts(fail_count: int, warn_count: int) -> str: return "OK" +def _ensure_report_file_name(args: Any) -> str: + """Ensure ``args.report_file_name`` is set; auto-generate a timestamped + default when the user did not pass ``--report-file-name``. + + The auto-generated form is ``preflight-{NNODES}N-{YYYYMMDD-HHMMSS}``. This + guarantees each run writes to a fresh, never-before-used path so the + report-path announcement (R5) cannot point at a stale leftover from a + previous run. + + Returns the resolved name. Safe to call multiple times: subsequent calls + are a no-op once ``args.report_file_name`` is populated. + """ + name = getattr(args, "report_file_name", None) + if name: + return name + + # Prefer NNODES from the env (set by the runner wrapper before exec) over + # deriving from torch's world size, because in info-only mode the + # distributed process group is not initialized and world reports 1. + nnodes_env = os.environ.get("NNODES") + if nnodes_env and nnodes_env.isdigit() and int(nnodes_env) > 0: + nnodes = int(nnodes_env) + else: + try: + _rank, world = get_rank_world() + except Exception: + world = 1 + try: + local_world = int(os.environ.get("LOCAL_WORLD_SIZE") or os.environ.get("GPUS_PER_NODE") or 8) + except (TypeError, ValueError): + local_world = 8 + nnodes = max(1, (world or 1) // max(1, local_world)) + + name = f"preflight-{nnodes}N-{datetime.now():%Y%m%d-%H%M%S}" + try: + args.report_file_name = name + except (AttributeError, TypeError): + # If args is a frozen / read-only container, the caller will see the + # returned name; downstream code in this module always uses the + # attribute, so the only impact is that the auto-name regenerates on + # the next call. Acceptable for the edge case. + pass + return name + + +def _announce_report_paths(args: Any) -> None: + """Print absolute paths of report files on rank 0 (R5). + + Scans ``args.dump_path`` for files matching + ``{,_perf}.{md,pdf}`` and prints absolute paths to + stdout. Under bash-side ``--silent`` (primus-cli-direct.sh), fd 1 is + ``/dev/null`` by inheritance, so these prints are silenced along with + every other stdout write -- acceptable per plan. + + Always rank-0 only. Best-effort: failures here must never mask the + preflight exit code. + """ + try: + rank, _world = get_rank_world() + except Exception: + rank = 0 + if rank != 0: + return + dump_path = getattr(args, "dump_path", "output/preflight") or "output/preflight" + name = getattr(args, "report_file_name", None) + if not name: + return + found_any = False + for suffix in ("", "_perf"): + for ext in ("md", "pdf"): + p = os.path.join(dump_path, f"{name}{suffix}.{ext}") + try: + if os.path.isfile(p): + print(f"[Primus:Preflight] Report: {os.path.abspath(p)}", flush=True) + found_any = True + except OSError: + continue + if not found_any: + print( + f"[Primus:Preflight] WARN: no report files found at " f"{dump_path}/{name}{{,_perf}}.{{md,pdf}}", + flush=True, + ) + + def run_preflight_info(args: Any, expect_distributed: bool = True) -> int: """ Run lightweight preflight info collection (host/gpu/network), aggregate across ranks, @@ -85,7 +288,10 @@ def run_preflight_info(args: Any, expect_distributed: bool = True) -> int: check_host = check_gpu = check_network = True dump_path = getattr(args, "dump_path", "output/preflight") - report_file_name = getattr(args, "report_file_name", "preflight_report") + # Defensive: if a caller reached this helper without going through + # run_preflight() (e.g. a test fixture), normalize the report name here so + # we never write a file literally called "None.md". + report_file_name = _ensure_report_file_name(args) save_pdf = bool(getattr(args, "save_pdf", True)) findings: List[Finding] = [] @@ -176,15 +382,58 @@ def run_preflight_info(args: Any, expect_distributed: bool = True) -> int: return rc +def _list_set_perf_tuning_knobs(args) -> List[str]: + """Return CLI flag names of perf tuning knobs the user explicitly set. + + Used to warn that these knobs are inert in info-only mode (i.e. when none + of --perf-test/--tests/--quick are set but at least one of + --host/--gpu/--network is). + """ + set_flags: List[str] = [] + if getattr(args, "comm_sizes_mb", None) is not None: + set_flags.append("--comm-sizes-mb") + if getattr(args, "intra_comm_sizes_mb", None) is not None: + set_flags.append("--intra-comm-sizes-mb") + if getattr(args, "inter_comm_sizes_mb", None) is not None: + set_flags.append("--inter-comm-sizes-mb") + if getattr(args, "intra_group_sizes", None) is not None: + set_flags.append("--intra-group-sizes") + if getattr(args, "inter_group_sizes", None) is not None: + set_flags.append("--inter-group-sizes") + if getattr(args, "ring_p2p_sizes_mb", None) is not None: + set_flags.append("--ring-p2p-sizes-mb") + if getattr(args, "plot", False): + set_flags.append("--plot") + if not getattr(args, "split_nodes_subgroup", True): + set_flags.append("--no-split-nodes-subgroup") + return set_flags + + def run_preflight(args): """ Preflight entry point with dispatch logic. - - If any of --host/--gpu/--network is set → show only selected info sections - - If no selection flags are set (plain `preflight`) → run ALL: info + perf tests - - If --perf-test is set → run perf tests ONLY (skip info) + Mode precedence (single rule): + + 1. Any of --perf-test / --tests / --quick is set -> perf-only mode. + If info selectors (--host/--gpu/--network) are also present, they + are dropped with a WARN. + 2. Otherwise, any of --host/--gpu/--network is set -> info-only mode. + Perf tuning knobs (e.g. --comm-sizes-mb), if set, are inert and a + WARN is emitted. + 3. Otherwise (no flags) -> default: info AND all perf tests. """ - perf_test = getattr(args, "perf_test", False) + # R4: canonical normalization point for the report file name. Done here + # before any downstream code reads args.report_file_name, so every code + # path (info-only, perf-only, info+perf, dist-init failure) sees the same + # value -- and so the dist-init failure paths below can interpolate it + # safely without a None-guard. + _ensure_report_file_name(args) + + perf_test = bool(getattr(args, "perf_test", False)) + tests_set = getattr(args, "tests", None) is not None + quick_set = bool(getattr(args, "quick", False)) + has_perf_intent = perf_test or tests_set or quick_set def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Exception) -> None: try: @@ -209,11 +458,50 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti # If any selection flags are set, only run info collection/report. # IMPORTANT: do NOT initialize torch.distributed for info-only mode; preflight must not hang # when networking/rendezvous is misconfigured. - check_host = getattr(args, "check_host", False) - check_gpu = getattr(args, "check_gpu", False) - check_network = getattr(args, "check_network", False) + check_host = bool(getattr(args, "check_host", False)) + check_gpu = bool(getattr(args, "check_gpu", False)) + check_network = bool(getattr(args, "check_network", False)) any_selection = bool(check_host or check_gpu or check_network) + # Precedence: perf-mode wins over info selectors. + info_dropped_warning: Optional[str] = None + if has_perf_intent: + if any_selection: + dropped = [ + flag + for present, flag in ( + (check_host, "--host"), + (check_gpu, "--gpu"), + (check_network, "--network"), + ) + if present + ] + info_dropped_warning = ( + f"info selectors {','.join(dropped)} were dropped because perf " + f"mode (--perf-test/--tests/--quick) takes precedence. " + f"Run them in a separate invocation if you want both reports." + ) + print(f"[Primus:Preflight] WARN: {info_dropped_warning}", file=sys.stderr) + check_host = check_gpu = check_network = False + any_selection = False + # Reflect the override on `args` so info-section helpers downstream + # (e.g. run_preflight_info getattr) see consistent state. + args.check_host = False + args.check_gpu = False + args.check_network = False + # Auto-imply --perf-test so the rest of the dispatch treats this as perf-only. + perf_test = True + args.perf_test = True + elif any_selection: + # Info-only mode: tuning knobs are inert. Emit a single WARN listing them. + inert = _list_set_perf_tuning_knobs(args) + if inert: + print( + f"[Primus:Preflight] WARN: {','.join(inert)} have no effect in " + f"info-only mode (no --perf-test/--tests/--quick).", + file=sys.stderr, + ) + # 1) Info-only mode: run without distributed init. if not perf_test and any_selection: # First, emit a local-only report immediately (so user gets output even if PG init hangs). @@ -231,20 +519,23 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti if world > 1: init_distributed(timeout=timedelta(seconds=dist_timeout_sec)) try: - return run_preflight_info(args) + rc = run_preflight_info(args) + _announce_report_paths(args) + return rc finally: finalize_distributed() except Exception as e: if rank == 0: dump_path = getattr(args, "dump_path", "output/preflight") - report_file_name = getattr(args, "report_file_name", "preflight_report") os.makedirs(dump_path, exist_ok=True) - markdown_file = f"{dump_path}/{report_file_name}.md" + markdown_file = f"{dump_path}/{args.report_file_name}.md" _append_dist_init_failure(markdown_file, dist_timeout_sec, e) print(f"[Primus:Preflight] ERROR: distributed init failed: {e}", file=sys.stderr) + _announce_report_paths(args) return 2 # world==1 fallback + _announce_report_paths(args) return local_rc # 2) Plain `preflight` (no flags): run info FIRST (no dist init) so we always get a report. @@ -252,6 +543,20 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti if not perf_test and not any_selection: info_rc = run_preflight_info(args, expect_distributed=False) + # 2.5) Resolve perf config NOW, before any distributed rendezvous, so that + # invalid CLI input (typos, bad sizes/group-sizes for selected tests) fails + # fast instead of after a 120s NCCL init. This is also where we apply the + # `--quick` warmup/iteration overrides on every rank. + try: + perf_cfg = _resolve_perf_config(args) + except ValueError as e: + print(f"[Primus:Preflight] ERROR: invalid perf config: {e}", file=sys.stderr) + return 2 + if perf_cfg["warmup"] is not None: + set_warmup(perf_cfg["warmup"]) + if perf_cfg["iteration"] is not None: + set_iteration(perf_cfg["iteration"]) + # 3) Perf tests (perf-only OR plain preflight after info): now attempt distributed init # with a timeout so we fail fast instead of hanging. from datetime import timedelta @@ -265,16 +570,16 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti # We already wrote the info report in plain preflight mode; append a clear failure note. rank, _world = get_rank_world() dump_path = getattr(args, "dump_path", "output/preflight") - report_file_name = getattr(args, "report_file_name", "preflight_report") if rank == 0: try: os.makedirs(dump_path, exist_ok=True) - markdown_file = f"{dump_path}/{report_file_name}.md" + markdown_file = f"{dump_path}/{args.report_file_name}.md" _append_dist_init_failure(markdown_file, dist_timeout_sec, e) except Exception as ee: print(f"[Primus:Preflight] [rank0] WARN: failed to write report: {ee}", file=sys.stderr) print(f"[Primus:Preflight] ERROR: distributed init failed: {e}", file=sys.stderr) + _announce_report_paths(args) return 2 try: @@ -311,12 +616,70 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti args.pdf_file = f"{args.dump_path}/{args.report_file_name}{perf_suffix}.pdf" remove_file(args.markdown_file) - # run tests - run_square_gemm(args) - run_intra_node_comm(args) - run_inter_node_comm(args) - run_inter_node_comm_p2p(args) - run_inter_node_ring_p2p(args) + # Write a Node -> Hostname legend at the top of the perf report so + # subsequent tables can use compact (and possibly truncated) hostname + # representations without losing the host<->node mapping. + if RANK == 0: + if info_dropped_warning: + with open(args.markdown_file, "a", encoding="utf-8") as f: + f.write(f"> Note: {info_dropped_warning}\n\n") + _write_node_hostname_legend(args.markdown_file) + + if RANK == 0: + log( + f"[Primus:Preflight] perf tests selected: " + f"{','.join(sorted(perf_cfg['enabled_tests'])) or '(none)'}" + ) + + # ------------------------------------------------------------------ + # Dispatch tests by token, with per-test wall-clock logging. + # ------------------------------------------------------------------ + enabled = perf_cfg["enabled_tests"] + intra_comms = {c for c in ("allreduce", "alltoall") if f"intra-{c}" in enabled} + inter_comms = {c for c in ("allreduce", "alltoall") if f"inter-{c}" in enabled} + + def _timed(name: str, fn): + t0 = time.time() + fn() + if RANK == 0: + log(f"[Primus:Preflight] {name} done in {time.time() - t0:.1f}s") + + if "gemm" in enabled: + _timed("gemm", lambda: run_square_gemm(args)) + + if intra_comms: + _timed( + f"intra-comm ({','.join(sorted(intra_comms))})", + lambda: run_intra_node_comm( + args, + enabled_comms=intra_comms, + sizes_mb=perf_cfg["intra_sizes_mb"], + group_sizes=perf_cfg["intra_group_sizes"], + ), + ) + + if inter_comms: + _timed( + f"inter-comm ({','.join(sorted(inter_comms))})", + lambda: run_inter_node_comm( + args, + enabled_comms=inter_comms, + sizes_mb=perf_cfg["inter_sizes_mb"], + group_sizes=perf_cfg["inter_group_sizes"], + ), + ) + + if "inter-p2p" in enabled: + _timed( + "inter-p2p", + lambda: run_inter_node_comm_p2p(args, sizes_mb=perf_cfg["inter_sizes_mb"]), + ) + + if "inter-ring-p2p" in enabled: + _timed( + "inter-ring-p2p", + lambda: run_inter_node_ring_p2p(args, sizes_mb=perf_cfg["ring_sizes_mb"]), + ) if RANK == 0 and args.save_pdf: md_to_pdf(args.markdown_file, args.pdf_file) @@ -327,9 +690,123 @@ def _append_dist_init_failure(markdown_file: str, timeout_sec: int, err: Excepti return info_rc return 0 finally: + _announce_report_paths(args) finalize_distributed() +# Built-in defaults for tuning knobs. These mirror the documented defaults in +# the parser; we keep argparse's `default=None` so that "user did not pass it" +# is detectable, and substitute these defaults when resolving the config. +_DEFAULT_COMM_SIZES_MB = "2,4,8,16,32,64,128,256,512,1024" +_DEFAULT_INTRA_GROUP_SIZES = "2,4,8" +_DEFAULT_INTER_GROUP_SIZES = "2,4,all" +_DEFAULT_RING_P2P_SIZES_MB = "10,20,40,80,160" + + +def _resolve_perf_config(args) -> Dict[str, Any]: + """Resolve perf-test selection, sizes, group sizes, and quick-preset overrides. + + User-supplied values take precedence over `--quick` defaults. Validates and + raises `ValueError` on bad input. + + This function is side-effect free: it returns `warmup` / `iteration` for + the caller to apply (e.g. via `set_warmup` / `set_iteration`) instead of + mutating module globals itself. + + Validation of intra-/inter-/ring-specific knobs is gated by which tests + are actually selected, so e.g. `--tests gemm --intra-group-sizes 3` does + NOT abort on a host with LOCAL_WORLD_SIZE=8. + """ + # All perf-related CLI args have argparse default=None, so a non-None value + # unambiguously means "the user passed this flag". + tests_str = getattr(args, "tests", None) + comm_sizes_str = getattr(args, "comm_sizes_mb", None) + intra_comm_sizes_str = getattr(args, "intra_comm_sizes_mb", None) + inter_comm_sizes_str = getattr(args, "inter_comm_sizes_mb", None) + intra_group_sizes_str = getattr(args, "intra_group_sizes", None) + inter_group_sizes_str = getattr(args, "inter_group_sizes", None) + ring_sizes_str = getattr(args, "ring_p2p_sizes_mb", None) + quick = bool(getattr(args, "quick", False)) + split_nodes_subgroup = bool(getattr(args, "split_nodes_subgroup", True)) + + # Apply --quick defaults only where the user did not override (i.e. value is None). + warmup: Optional[int] = None + iteration: Optional[int] = None + if quick: + if tests_str is None: + tests_str = "gemm,intra-allreduce,inter-allreduce" + if comm_sizes_str is None: + comm_sizes_str = "64,1024" + if intra_group_sizes_str is None: + intra_group_sizes_str = str(LOCAL_WORLD_SIZE) + if inter_group_sizes_str is None: + inter_group_sizes_str = "all" + # Lower warmup / iterations for a fast go/no-go signal. Caller applies. + warmup = 5 + iteration = 20 + + # Substitute built-in defaults for any values still unset. + if comm_sizes_str is None: + comm_sizes_str = _DEFAULT_COMM_SIZES_MB + if intra_group_sizes_str is None: + intra_group_sizes_str = _DEFAULT_INTRA_GROUP_SIZES + if inter_group_sizes_str is None: + inter_group_sizes_str = _DEFAULT_INTER_GROUP_SIZES + if ring_sizes_str is None: + ring_sizes_str = _DEFAULT_RING_P2P_SIZES_MB + + enabled_tests = _parse_perf_tests(tests_str) + + # Honor the legacy --no-split-nodes-subgroup alias: drop subgroup sizes + # AND skip the inter-p2p test (matching the prior behavior of the script). + if not split_nodes_subgroup: + inter_group_sizes_str = "all" + enabled_tests.discard("inter-p2p") + + needs_intra = any(t in enabled_tests for t in ("intra-allreduce", "intra-alltoall")) + needs_inter_coll = any(t in enabled_tests for t in ("inter-allreduce", "inter-alltoall")) + needs_inter_p2p = "inter-p2p" in enabled_tests + needs_inter = needs_inter_coll or needs_inter_p2p + needs_ring = "inter-ring-p2p" in enabled_tests + + default_sizes = _parse_csv_int_list(comm_sizes_str, "--comm-sizes-mb") + intra_sizes_override = _parse_csv_int_list(intra_comm_sizes_str, "--intra-comm-sizes-mb") + inter_sizes_override = _parse_csv_int_list(inter_comm_sizes_str, "--inter-comm-sizes-mb") + intra_sizes_mb = intra_sizes_override or default_sizes + inter_sizes_mb = inter_sizes_override or default_sizes + if needs_intra and not intra_sizes_mb: + raise ValueError("--comm-sizes-mb / --intra-comm-sizes-mb yielded no sizes") + if needs_inter and not inter_sizes_mb: + raise ValueError("--comm-sizes-mb / --inter-comm-sizes-mb yielded no sizes") + + intra_group_sizes = _parse_csv_int_list(intra_group_sizes_str, "--intra-group-sizes") + if needs_intra: + if not intra_group_sizes: + raise ValueError("--intra-group-sizes yielded no values") + bad = [g for g in intra_group_sizes if LOCAL_WORLD_SIZE % g != 0] + if bad: + raise ValueError(f"--intra-group-sizes: {bad} do not divide LOCAL_WORLD_SIZE={LOCAL_WORLD_SIZE}") + + inter_group_sizes = _parse_csv_inter_group_list(inter_group_sizes_str, "--inter-group-sizes") + if needs_inter and not inter_group_sizes: + raise ValueError("--inter-group-sizes yielded no values") + + ring_sizes_mb = _parse_csv_int_list(ring_sizes_str, "--ring-p2p-sizes-mb") + if needs_ring and not ring_sizes_mb: + raise ValueError("--ring-p2p-sizes-mb yielded no sizes") + + return { + "enabled_tests": enabled_tests, + "intra_sizes_mb": intra_sizes_mb, + "inter_sizes_mb": inter_sizes_mb, + "intra_group_sizes": intra_group_sizes, + "inter_group_sizes": inter_group_sizes, + "ring_sizes_mb": ring_sizes_mb, + "warmup": warmup, + "iteration": iteration, + } + + def main(): parser = argparse.ArgumentParser() from primus.tools.preflight.preflight_args import add_preflight_parser diff --git a/primus/tools/preflight/square_gemm.py b/primus/tools/preflight/square_gemm.py index daa592062..c947ba121 100644 --- a/primus/tools/preflight/square_gemm.py +++ b/primus/tools/preflight/square_gemm.py @@ -6,18 +6,17 @@ import time -import matplotlib.pyplot as plt import torch import torch.distributed as dist from primus.tools.preflight.global_vars import ( - ITERATION, LOCAL_RANK, LOCAL_WORLD_SIZE, RANK, - WARMUP, WORLD_SIZE, get_hostnames, + get_iteration, + get_warmup, ) from primus.tools.preflight.utility import create_dir, log @@ -26,19 +25,21 @@ def run_square_gemm(args): sizes = [1024, 2048, 4096, 8192, 10240] latency_results = {} flops_results = {} + warmup = get_warmup() + iteration = get_iteration() for size in sizes: a = torch.randn((size, size), device=f"cuda:{LOCAL_RANK}", dtype=torch.bfloat16) b = torch.randn((size, size), device=f"cuda:{LOCAL_RANK}", dtype=torch.bfloat16) torch.cuda.synchronize() - for _ in range(WARMUP): + for _ in range(warmup): torch.matmul(a, b) torch.cuda.synchronize() start = time.time() - for _ in range(ITERATION): + for _ in range(iteration): torch.matmul(a, b) torch.cuda.synchronize() end = time.time() - t = (end - start) / ITERATION + t = (end - start) / iteration tflops = 2 * size * size * size / (t * 1e12) latency_results[f"{size}x{size}x{size}"] = t flops_results[f"{size}x{size}x{size}"] = tflops @@ -87,6 +88,8 @@ def run_square_gemm(args): if not args.plot: return + import matplotlib.pyplot as plt + log("=======Plot Square GEMM TFLOPS=======") with open(args.markdown_file, "a", encoding="utf-8") as f: f.write(f"=======Plot Square GEMM TFLOPS=======\n") diff --git a/primus/tools/preflight/tests/test_report_naming.py b/primus/tools/preflight/tests/test_report_naming.py new file mode 100644 index 000000000..13f78ccb2 --- /dev/null +++ b/primus/tools/preflight/tests/test_report_naming.py @@ -0,0 +1,238 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tests for ``preflight_perf_test._ensure_report_file_name`` / +``_announce_report_paths`` (consolidate-preflight-direct-wrappers, R5). + +These tests pin the contract of the Python-side replacements for the +old bash ``run_preflight_direct.sh`` behavior: + +* Every preflight run that doesn't get an explicit ``--report-file-name`` + must write to a unique, timestamped path. Stale reports from a previous + run must NEVER be mistaken for the current run's output. +* Once reports exist, rank 0 must print their absolute paths so an + operator running under ``primus-cli direct`` can find them without + guessing the dump-path layout. +* Best-effort: the announcement must never raise, even when nothing was + written (the per-run output may have failed earlier). +""" + +from __future__ import annotations + +import argparse +import os +import re +from datetime import datetime, timedelta +from typing import Optional + +from primus.tools.preflight.preflight_args import add_preflight_parser +from primus.tools.preflight.preflight_perf_test import ( + _announce_report_paths, + _ensure_report_file_name, +) + + +def _build_preflight_parser() -> argparse.ArgumentParser: + """Construct the same parser that `primus-cli preflight` registers, + isolated from the broader primus.cli wiring so tests can exercise + just the argparse contract. + """ + p = argparse.ArgumentParser(prog="preflight-test") + add_preflight_parser(p) + return p + + +# --------------------------------------------------------------------------- +# `--report-file-name` default contract +# --------------------------------------------------------------------------- + + +def test_preflight_args_report_file_name_default_is_none(): + """The argparse default for ``--report-file-name`` must be ``None`` so + ``_ensure_report_file_name`` can detect "user did not pass it" and + auto-generate a unique name. A non-None default would silently + suppress the auto-naming logic and re-introduce the "every run + overwrites preflight_report.md" footgun the plan removed.""" + parser = _build_preflight_parser() + ns = parser.parse_args([]) + assert ns.report_file_name is None, ( + "argparse default must stay None -- _ensure_report_file_name " + "uses falsy check to decide whether to auto-generate" + ) + + +# --------------------------------------------------------------------------- +# `_ensure_report_file_name` +# --------------------------------------------------------------------------- + + +_AUTO_NAME_RE = re.compile(r"^preflight-(\d+)N-(\d{8})-(\d{6})$") + + +def test_ensure_report_file_name_user_value_wins(monkeypatch): + """When the user passed ``--report-file-name=foo`` it must survive + untouched -- the auto-name path must never overwrite an explicit + value.""" + monkeypatch.setenv("NNODES", "32") # would otherwise affect auto-name + ns = argparse.Namespace(report_file_name="my-custom-name") + out = _ensure_report_file_name(ns) + assert out == "my-custom-name" + assert ns.report_file_name == "my-custom-name" + + +def test_ensure_report_file_name_auto_uses_nnodes_env(monkeypatch): + """When ``NNODES`` is set, the auto-name must use it directly. This + is the critical case for info-only mode where no process group is + initialized and ``get_rank_world()`` would return world=1.""" + monkeypatch.setenv("NNODES", "42") + ns = argparse.Namespace(report_file_name=None) + out = _ensure_report_file_name(ns) + m = _AUTO_NAME_RE.match(out) + assert m is not None, f"auto-name {out!r} doesn't match expected pattern" + assert m.group(1) == "42" + assert ns.report_file_name == out + + +def test_ensure_report_file_name_idempotent(monkeypatch): + """Calling ``_ensure_report_file_name`` twice in the same run must + return the same name -- otherwise different code paths (run / + info / announce) would each compute their own timestamped name and + write to different files.""" + monkeypatch.setenv("NNODES", "8") + ns = argparse.Namespace(report_file_name=None) + first = _ensure_report_file_name(ns) + # Sleep is unnecessary -- if the function recomputed it would use + # `datetime.now()` again, which is enough to diverge at second + # granularity over a slow run. We test the stronger invariant: the + # second call returns the stored value verbatim, without recomputing. + monkeypatch.setattr( + "primus.tools.preflight.preflight_perf_test.datetime", + _ExplodingDatetime, # any access to datetime would raise + ) + second = _ensure_report_file_name(ns) + assert second == first + + +class _ExplodingDatetime: + """datetime stub that raises on any use -- proves a method doesn't + reach the auto-name regeneration path.""" + + @classmethod + def now(cls, *_a, **_kw): + raise AssertionError( + "_ensure_report_file_name regenerated the timestamp on " "the second call (it must be idempotent)" + ) + + +def test_ensure_report_file_name_unique_per_run(monkeypatch): + """Two independent invocations (different `args` objects with + `report_file_name=None`) at different timestamps must produce + distinct names -- this is the guarantee that protects against + stale-report aliasing under back-to-back runs.""" + monkeypatch.setenv("NNODES", "4") + + class _FrozenDatetime: + _now = datetime(2026, 5, 7, 12, 0, 0) + + @classmethod + def now(cls, *_a, **_kw): + return cls._now + + monkeypatch.setattr("primus.tools.preflight.preflight_perf_test.datetime", _FrozenDatetime) + ns_a = argparse.Namespace(report_file_name=None) + name_a = _ensure_report_file_name(ns_a) + # Bump the clock by one second to simulate the next run. + _FrozenDatetime._now = _FrozenDatetime._now + timedelta(seconds=1) + ns_b = argparse.Namespace(report_file_name=None) + name_b = _ensure_report_file_name(ns_b) + + assert name_a != name_b + assert name_a.endswith("12-00-00".replace("-", "")) + assert name_b.endswith("12-00-01".replace("-", "")) + + +def test_ensure_report_file_name_handles_invalid_nnodes(monkeypatch): + """A malformed ``NNODES`` (non-digit) must NOT crash the auto-name + helper; it must fall back to deriving from torch's world size, with a + safe minimum of 1.""" + monkeypatch.setenv("NNODES", "") # explicitly empty + monkeypatch.delenv("LOCAL_WORLD_SIZE", raising=False) + monkeypatch.delenv("GPUS_PER_NODE", raising=False) + ns = argparse.Namespace(report_file_name=None) + out = _ensure_report_file_name(ns) + m = _AUTO_NAME_RE.match(out) + assert m is not None, f"auto-name {out!r} doesn't match expected pattern" + # When world isn't initialized we end up with nnodes=max(1, 1//8)=1. + assert int(m.group(1)) >= 1 + + +# --------------------------------------------------------------------------- +# `_announce_report_paths` +# --------------------------------------------------------------------------- + + +def _make_args(dump_path: str, name: Optional[str]) -> argparse.Namespace: + return argparse.Namespace(dump_path=dump_path, report_file_name=name) + + +def test_announce_report_paths_finds_existing_md(tmp_path, capsys, monkeypatch): + """Smoke: ``_announce_report_paths`` must print absolute paths for + every report variant that exists on disk.""" + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + (tmp_path / "preflight-2N-20260507-120000.md").write_text("# report") + (tmp_path / "preflight-2N-20260507-120000_perf.md").write_text("# perf") + + args = _make_args(str(tmp_path), "preflight-2N-20260507-120000") + _announce_report_paths(args) + + out = capsys.readouterr().out + assert "preflight-2N-20260507-120000.md" in out + assert "preflight-2N-20260507-120000_perf.md" in out + # Paths must be absolute -- operators copy them into srun nodelists + # / scp commands; a relative path here is a footgun under SLURM + # where the cwd doesn't match across nodes. + for line in out.splitlines(): + if line.startswith("[Primus:Preflight] Report:"): + path = line.split("Report:", 1)[1].strip() + assert os.path.isabs(path), f"non-absolute report path: {path!r}" + + +def test_announce_report_paths_warns_when_no_files(tmp_path, capsys, monkeypatch): + """When the helper is called but no report files exist (e.g. preflight + crashed before writing), it must emit a WARN line so the operator + isn't left wondering where the report went -- but it must NOT raise.""" + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + args = _make_args(str(tmp_path), "ghost-name") + _announce_report_paths(args) # must not raise + out = capsys.readouterr().out + assert "WARN" in out + assert "ghost-name" in out + + +def test_announce_report_paths_skips_non_rank0(tmp_path, capsys, monkeypatch): + """Non-rank-0 ranks must stay silent -- otherwise every node would + print the same paths and clutter the launcher output.""" + monkeypatch.setenv("RANK", "3") + monkeypatch.setenv("WORLD_SIZE", "8") + (tmp_path / "preflight-1N-20260507-120000.md").write_text("# report") + args = _make_args(str(tmp_path), "preflight-1N-20260507-120000") + _announce_report_paths(args) + out = capsys.readouterr().out + assert out == "", f"non-rank-0 announce produced output: {out!r}" + + +def test_announce_report_paths_no_name_is_noop(tmp_path, capsys, monkeypatch): + """Defensive: if ``args.report_file_name`` is somehow still None when + we reach the announcement (e.g. an early bail-out before + ``_ensure_report_file_name`` was called), the helper must be a + silent no-op rather than crash on the format-string path build.""" + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + args = _make_args(str(tmp_path), None) + _announce_report_paths(args) # must not raise + assert capsys.readouterr().out == "" diff --git a/primus/tools/preflight/utility.py b/primus/tools/preflight/utility.py index 9c71d6335..e27ec1a73 100644 --- a/primus/tools/preflight/utility.py +++ b/primus/tools/preflight/utility.py @@ -6,15 +6,17 @@ import os import socket +import time from pathlib import Path +from typing import Iterable, List -import markdown2 import torch import torch.distributed as dist -from weasyprint import HTML from primus.tools.preflight.global_vars import RANK, WORLD_SIZE +DEFAULT_COMM_CLEANUP_DELAY_SEC = 2.0 + def log(msg): if RANK == 0: @@ -48,6 +50,23 @@ def gather_hostnames(): return None +def barrier_after_comm_destroy(delay_sec: float = DEFAULT_COMM_CLEANUP_DELAY_SEC): + """Global barrier + sleep after destroying NCCL/RCCL process groups. + + When communicators are destroyed and recreated quickly (e.g. in test loops), + the underlying sockets may still be in TIME_WAIT or the remote side may not + have fully cleaned up. This causes "Address already in use" errors on bind. + + This function ensures all ranks have completed destruction (barrier) and then + waits a short period for the OS to release ports before the next group creation. + """ + log(f" barrier_after_comm_destroy: barrier synced, sleeping {delay_sec:.1f}s") + dist.barrier(device_ids=[torch.cuda.current_device()]) + if delay_sec > 0: + time.sleep(delay_sec) + log(f" barrier_after_comm_destroy: done (delay={delay_sec:.1f}s)") + + def remove_file(file_path): if RANK == 0: if os.path.exists(file_path): @@ -56,6 +75,32 @@ def remove_file(file_path): dist.barrier(device_ids=[torch.cuda.current_device()]) +def format_int_range(values: Iterable[int]) -> str: + """Compactly format a list of ints as a comma-separated set of contiguous ranges. + + Examples: + [0,1,2,3] -> "0-3" + [0,1,2,3,8,9] -> "0-3,8-9" + [0,2,4,6] -> "0,2,4,6" + [5] -> "5" + [] -> "" + """ + vs = sorted(set(int(v) for v in values)) + if not vs: + return "" + + parts: List[str] = [] + start = prev = vs[0] + for v in vs[1:]: + if v == prev + 1: + prev = v + continue + parts.append(f"{start}-{prev}" if start != prev else f"{start}") + start = prev = v + parts.append(f"{start}-{prev}" if start != prev else f"{start}") + return ",".join(parts) + + def extract_first_middle_last(lst): if not lst: return [] @@ -70,6 +115,9 @@ def extract_first_middle_last(lst): def md_to_pdf(md_path, pdf_path): + import markdown2 + from weasyprint import HTML + with open(md_path, "r", encoding="utf-8") as f: markdown_text = f.read() diff --git a/runner/.primus.yaml b/runner/.primus.yaml index f8b244498..afe6948a6 100644 --- a/runner/.primus.yaml +++ b/runner/.primus.yaml @@ -121,7 +121,10 @@ direct: master_addr: "localhost" # Direct mode specific options - run_mode: "torchrun" + # run_mode: deliberately omitted here so that primus-cli-direct.sh can + # auto-select 'single' for subcommands that require it (e.g. node_smoke). + # Set run_mode explicitly in your own config to force a specific launcher. + # Precedence: CLI --single > user config direct.run_mode > auto-detect > torchrun. script: "primus/cli/main.py" numa: "auto" log_file: "" diff --git a/runner/primus-cli-direct.sh b/runner/primus-cli-direct.sh index d87de0d01..c90a0d86d 100755 --- a/runner/primus-cli-direct.sh +++ b/runner/primus-cli-direct.sh @@ -33,6 +33,11 @@ Options: --log_file PATH Save log to a specific file (default: logs/log_TIMESTAMP.txt) --numa Force enable NUMA binding for processes --no-numa Force disable NUMA binding for processes + --silent Suppress all launcher and python tool stdout (back-pocket option). + Must be placed BEFORE the '--' separator. Launcher errors + (LOG_ERROR / LOG_WARN, written to stderr) and the log file + are preserved. Exit code is propagated. NOT recommended for + normal use -- you lose live progress; prefer the log file. Distributed Environment Variables: NNODES Number of nodes participating in distributed run [default: 1] @@ -41,6 +46,17 @@ Distributed Environment Variables: MASTER_ADDR Hostname or IP of master node [default: localhost] MASTER_PORT Port of master node [default: 1234] + When running inside a SLURM allocation (SLURM_JOB_ID set), the above are + auto-derived from SLURM_NNODES / SLURM_NODEID / SLURM_NODELIST when not + pre-exported. Pre-exported values always win, so the standard + slurm-entry -> container -> direct chain is unaffected. + +Optional Environment Variables: + VENV_ACTIVATE Path to a Python virtualenv activate script. If set, sourced + before the primus run. Unset = no-op (use system python or + the container's bundled python). Set + missing file = + fail-fast (avoid silent torch-version mismatches). + You can set these variables in either of the following ways: # (1) Export variables before launch (recommended for scripts or single-node runs) export NNODES=2 GPUS_PER_NODE=8 NODE_RANK=0 MASTER_ADDR=host1 @@ -50,6 +66,10 @@ You can set these variables in either of the following ways: primus-cli direct --env NNODES=2 --env GPUS_PER_NODE=8 --env NODE_RANK=1 --env MASTER_ADDR=host1 -- \\ benchmark gemm -M 4096 -N 4096 -K 4096 + # (3) Let SLURM provide them (inside an existing allocation) + export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate + srun -N 4 --ntasks-per-node=1 ./runner/primus-cli direct -- preflight --quick + Examples: # Pretrain with a config file (single node) primus-cli direct -- train pretrain --config examples/megatron/exp_pretrain.yaml @@ -77,12 +97,19 @@ Examples: # Force enable NUMA binding for better performance primus-cli direct --numa -- benchmark gemm -M 8192 -N 8192 -K 8192 + # Run silently (back-pocket option; launcher errors and log file preserved) + primus-cli direct --silent -- preflight --quick + Notes: - If --single is specified, Primus skips torchrun and uses python3 directly. + - run_mode auto-detection: when the primus subcommand is 'node_smoke', run_mode + defaults to 'single' (node_smoke runs one process per node by design). + Explicit --single or config direct.run_mode still wins. - If --script is not specified, defaults to primus/cli/main.py. - Always separate Primus arguments from launcher options using '--'. - Environment variables can be mixed: 'export' takes precedence unless overridden by '--env'. - Multi-node jobs require MASTER_ADDR set to the master node's hostname/IP. + Inside a SLURM allocation it is auto-resolved from SLURM_NODELIST. - Patch scripts are executed in order before running the main script (useful for env setup, hot fixes, etc.). - NUMA binding is auto-disabled by default; use --numa to enable for better memory locality. - GPU-specific optimizations: The script automatically sources primus-env.sh, which detects your @@ -97,6 +124,43 @@ if [[ "${1:-}" == "--help" || "${1:-}" == "-h" ]]; then exit 0 fi +############################################################################### +# Runner-level --silent (back-pocket option for clean stdout) +# +# Contract: +# - launcher LOG_INFO / LOG_INFO_RANK0 / LOG_DEBUG_* (stdout) -> /dev/null +# - launcher LOG_WARN / LOG_ERROR (stderr, written via >&2) -> terminal +# - python tool's stdout + stderr -> /dev/null +# (the launched CMD ends with `2>&1 | tee `, so the python +# child's fd2 is merged into the pipe and tee writes to the log file +# + this script's fd1 = /dev/null) +# - log file -> captures all +# - exit code -> propagated +# +# --silent is only honored BEFORE the `--` separator. Anywhere after `--` it +# is forwarded to the python tool whose argparse will reject it (intentional; +# python tools have no --silent flag and never will -- silencing is a bash-side +# concern). +# +# Not recommended for normal use: read the log file or omit --silent if you +# want to see live tool output and progress. +############################################################################### +SILENT=0 +for _arg in "$@"; do + if [[ "$_arg" == "--" ]]; then + break + fi + if [[ "$_arg" == "--silent" ]]; then + SILENT=1 + break + fi +done +if [[ "$SILENT" == "1" ]]; then + # fd1 -> /dev/null. fd2 is left attached to the terminal, so LOG_ERROR / + # LOG_WARN (which write via >&2) still reach the operator. + exec >/dev/null +fi + # Resolve runner directory # Use RUNNER_DIR instead of SCRIPT_DIR to avoid conflicts with sourced scripts RUNNER_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" @@ -180,6 +244,12 @@ while [[ $# -gt 0 ]]; do PRE_PARSE_ARGS+=("$1") shift ;; + --silent) + # Already consumed by the pre-scan above (which has applied + # `exec >/dev/null`). Swallow it here so it never reaches the + # python parser, which has no --silent flag. + shift + ;; --) # Explicit separator: remaining args are for primus Python module shift # skip the '--' @@ -241,8 +311,11 @@ if [[ "$DRY_RUN_MODE" == "0" ]]; then fi fi -# Set default values for parameters not in config -direct_config[run_mode]="${direct_config[run_mode]:-torchrun}" +# Set default values for parameters not in config. +# NOTE: run_mode default is deliberately NOT applied here. The auto-detect step +# after STEP 4 needs to distinguish "user did not specify run_mode" (so we can +# auto-select `single` for node_smoke) from "default got applied in STEP 3". +# Defaulting to `torchrun` therefore happens after STEP 4 instead. direct_config[script]="${direct_config[script]:-primus/cli/main.py}" direct_config[numa]="${direct_config[numa]:-auto}" direct_config[log_file]="${direct_config[log_file]:-}" @@ -315,6 +388,36 @@ set -- "${primus_args[@]}" # Done here so it flows through hooks/patches and matches container/slurm pattern. [[ "$DEBUG_MODE" == "1" ]] && set -- --debug "$@" +############################################################################### +# STEP 4.4: Auto-select run_mode based on primus subcommand +# +# Some primus subcommands MUST run as a single process per srun task (no +# torchrun fan-out, no inter-node rendezvous). node_smoke is the canonical +# example: it runs one process per node by design and the per-GPU phase is +# launched as subprocesses internally. Listed here so users don't need to +# remember --single. Explicit --single / config direct.run_mode still wins +# (those paths set direct_config[run_mode] before this block runs). +############################################################################### +SINGLE_MODE_SUBCOMMANDS=(node_smoke) +if [[ -z "${direct_config[run_mode]:-}" ]]; then + _detected_subcmd="" + for _arg in "${primus_args[@]}"; do + case "$_arg" in + --*|-*) continue ;; + *) _detected_subcmd="$_arg"; break ;; + esac + done + _default_run_mode="torchrun" + for _sc in "${SINGLE_MODE_SUBCOMMANDS[@]}"; do + if [[ "$_detected_subcmd" == "$_sc" ]]; then + _default_run_mode="single" + LOG_INFO_RANK0 "[direct] Auto-selected run_mode=single for subcommand '$_detected_subcmd'" + break + fi + done + direct_config[run_mode]="$_default_run_mode" +fi + ############################################################################### # STEP 4.5: Process non-cumulative parameters (use last value only) ############################################################################### @@ -337,6 +440,83 @@ if [[ -z "${direct_config[log_file]:-}" ]]; then fi mkdir -p "$(dirname "${direct_config[log_file]:-}")" +############################################################################### +# STEP 4.7: Activate virtualenv (R1) and derive distributed env from SLURM (R2) +# +# Previously these lived in runner/run_preflight_direct.sh and +# runner/run_node_smoke_direct.sh (now deleted). Hoisting them here makes every +# `primus-cli direct -- ...` call site (host srun, slurm-entry -> direct, +# slurm-entry -> container -> direct) inherit identical behavior. +############################################################################### + +# R1 -- Python virtualenv. VENV_ACTIVATE unset = no-op (this is the right +# default for the container path: primus-cli-container.sh does not auto-forward +# VENV_ACTIVATE through its env passthrough whitelist, so inside the container +# we use the container's bundled python). Set + missing = fail-fast: better a +# loud error than a silent torch-version mismatch. +if [[ -n "${VENV_ACTIVATE:-}" ]]; then + if [[ ! -f "$VENV_ACTIVATE" ]]; then + LOG_ERROR "[direct] VENV_ACTIVATE is set but file does not exist: $VENV_ACTIVATE" + exit 1 + fi + # shellcheck disable=SC1090 + source "$VENV_ACTIVATE" + LOG_INFO_RANK0 "[direct] Activated virtualenv: $VENV_ACTIVATE" +fi + +# R2 -- distributed env. Pre-set values always win (the existing +# slurm-entry -> container -> direct chain already passes them via --env, so +# this block is a no-op there). When called directly under srun on the host, +# this block derives them from SLURM_*. +export GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +export MASTER_PORT="${MASTER_PORT:-1234}" + +if [[ -n "${SLURM_JOB_ID:-}" ]]; then + # Pre-exported values always win (per plan). The standard + # slurm-entry -> container -> direct chain already exports NNODES / + # NODE_RANK with the SLURM-derived values before direct.sh runs, so this + # block is a no-op there. When called directly under srun on bare metal + # without pre-set values, derive from SLURM_*. + export NNODES="${NNODES:-${SLURM_NNODES:-${SLURM_JOB_NUM_NODES:-1}}}" + export NODE_RANK="${NODE_RANK:-${SLURM_NODEID:-${SLURM_PROCID:-0}}}" + if [[ -z "${MASTER_ADDR:-}" || "${MASTER_ADDR}" == "localhost" ]]; then + if command -v scontrol >/dev/null 2>&1 && [[ -n "${SLURM_NODELIST:-}" ]]; then + if ! MASTER_ADDR="$(scontrol show hostnames "$SLURM_NODELIST" | head -n1)" \ + || [[ -z "$MASTER_ADDR" ]]; then + LOG_ERROR "[direct] Failed to resolve MASTER_ADDR from SLURM_NODELIST=${SLURM_NODELIST:-}" + exit 1 + fi + export MASTER_ADDR + else + # In a SLURM context but we cannot resolve a real address from + # scontrol -- either scontrol is not installed (CI / dev VM) or + # SLURM_NODELIST was not propagated (rare; can happen in stubbed + # tests or single-node allocations). Fall back to localhost so + # the downstream sanity check at line 503 has a valid value; + # NODE_RANK=0 + MASTER_ADDR=localhost is correct for single-node + # SLURM runs and for dry-run smoke tests. Real multi-node + # bare-srun on a SLURM head node always has both, so this + # branch never fires in production. + export MASTER_ADDR="localhost" + fi + fi + LOG_INFO_RANK0 "[direct] SLURM detected: JOB_ID=$SLURM_JOB_ID NNODES=$NNODES NODE_RANK=$NODE_RANK MASTER_ADDR=${MASTER_ADDR:-}" +else + export NNODES="${NNODES:-1}" + export NODE_RANK="${NODE_RANK:-0}" + export MASTER_ADDR="${MASTER_ADDR:-localhost}" +fi + +# Sanity-check the resolved distributed env (lifted from the deleted wrappers). +[[ "$NNODES" =~ ^[1-9][0-9]*$ ]] || { LOG_ERROR "[direct] NNODES must be a positive integer (got '$NNODES')"; exit 1; } +[[ "$NODE_RANK" =~ ^[0-9]+$ ]] || { LOG_ERROR "[direct] NODE_RANK must be a non-negative integer (got '$NODE_RANK')"; exit 1; } +[[ "$MASTER_PORT" =~ ^[0-9]+$ ]] || { LOG_ERROR "[direct] MASTER_PORT must be a non-negative integer (got '$MASTER_PORT')"; exit 1; } +[[ -n "$MASTER_ADDR" ]] || { LOG_ERROR "[direct] MASTER_ADDR is empty"; exit 1; } +(( NODE_RANK < NNODES )) || { LOG_ERROR "[direct] NODE_RANK ($NODE_RANK) must be < NNODES ($NNODES)"; exit 1; } +if [[ "$MASTER_ADDR" == "localhost" && "${NNODES:-1}" -gt 1 ]]; then + LOG_WARN "[direct] MASTER_ADDR=localhost with NNODES=$NNODES — multi-node will likely fail" +fi + ############################################################################### # STEP 5: Source GPU environment and helper modules ############################################################################### @@ -461,7 +641,14 @@ fi # STEP 9: Build launch command ############################################################################### -# Allow RUN_MODE to be overridden by environment variable +# Final run-mode resolution. RUN_MODE-from-env wins over direct_config[run_mode] +# because framework prepare-hooks (e.g. runner/helpers/hooks/train/pretrain/ +# maxtext/prepare.py) emit `env.RUN_MODE=single` from STEP 6 -- the hook layer +# knows things the launcher can't (e.g. "this framework is JAX, not torch, so +# torchrun would be wrong"). $RUN_MODE is the authoritative value for the rest +# of the script -- the display block and the torchrun-only Distributed Settings +# gate at STEP 10 both read $RUN_MODE, NOT direct_config[run_mode], so a hook +# that flips the mode is faithfully reflected in the printed configuration. RUN_MODE="${RUN_MODE:-${direct_config[run_mode]:-torchrun}}" # Resolve the launch target. Normally this is the script path @@ -533,7 +720,7 @@ else print_section "Primus Direct Launch Configuration" fi -PRINT_INFO_RANK0 " Run Mode : ${direct_config[run_mode]:-}" +PRINT_INFO_RANK0 " Run Mode : ${RUN_MODE}" PRINT_INFO_RANK0 " Script Path : ${LAUNCH_TARGET[*]:-${direct_config[script]:-}}" PRINT_INFO_RANK0 " Config File : ${CONFIG_FILE:-}" PRINT_INFO_RANK0 " Log File : ${direct_config[log_file]:-}" @@ -554,7 +741,7 @@ if [[ ${#primus_env_kv[@]} -gt 0 ]]; then PRINT_INFO_RANK0 "" fi -if [[ "${direct_config[run_mode]:-}" == "torchrun" ]]; then +if [[ "${RUN_MODE}" == "torchrun" ]]; then PRINT_INFO_RANK0 " Distributed Settings:" PRINT_INFO_RANK0 " NNODES : ${NNODES:-1}" PRINT_INFO_RANK0 " NODE_RANK : ${NODE_RANK:-0}" diff --git a/runner/primus-cli-slurm-entry.sh b/runner/primus-cli-slurm-entry.sh index f45f57585..dd8662fab 100755 --- a/runner/primus-cli-slurm-entry.sh +++ b/runner/primus-cli-slurm-entry.sh @@ -118,8 +118,19 @@ if [[ -z "${SLURM_NODELIST:-}" ]]; then exit 2 fi -# Get all node hostnames (sorted, as needed) -readarray -t NODE_ARRAY < <(scontrol show hostnames "$SLURM_NODELIST") +# Get all node hostnames (sorted, as needed). Prefer scontrol, which correctly +# expands compressed nodelists (e.g. "node[01-04]"). When scontrol is +# unavailable -- CI containers / dev VMs without the Slurm client tools -- fall +# back to parsing SLURM_NODELIST directly. This mirrors the scontrol-optional +# handling in primus-cli-direct.sh so every launcher behaves consistently +# off-cluster (range expansion is skipped in the fallback, which is fine for the +# single-host / comma-list forms used in CI and single-node runs). +if command -v scontrol >/dev/null 2>&1; then + readarray -t NODE_ARRAY < <(scontrol show hostnames "$SLURM_NODELIST") +else + LOG_WARN "[slurm-entry] scontrol not found; parsing SLURM_NODELIST without range expansion" + readarray -t NODE_ARRAY < <(tr ',' '\n' <<< "$SLURM_NODELIST") +fi SLURM_MASTER_ADDR="${NODE_ARRAY[0]:-}" if [[ -z "$SLURM_MASTER_ADDR" ]]; then LOG_ERROR "[slurm-entry] Failed to resolve the first host from SLURM_NODELIST=$SLURM_NODELIST" @@ -156,8 +167,20 @@ validate_distributed_params || LOG_WARN "[slurm-entry] Failed to validate distri # ------------- Dispatch based on mode --------------- -# Parse mode (default: container) +# Strip the entry-mode separator if present. +[[ "${1:-}" == "--" ]] && shift + +# Parse entry mode (default: container). Supported keywords are the ones +# advertised by primus-cli-slurm.sh: container | direct. Anything else is +# treated as a primus subcommand and routed through the default container +# chain (preserves the documented terse forms like `-- preflight`). +ENTRY_MODE="container" +if [[ "${1:-}" == "container" || "${1:-}" == "direct" ]]; then + ENTRY_MODE="$1" + shift +fi [[ "${1:-}" == "--" ]] && shift +LOG_INFO_RANK0 "[slurm-entry] Entry mode: $ENTRY_MODE" # Build arguments based on mode SCRIPT_ARGS=() @@ -175,8 +198,10 @@ SCRIPT_ARGS+=( --env "GPUS_PER_NODE=$GPUS_PER_NODE" ) -# Build script path (container mode only) -script_path="$RUNNER_DIR/primus-cli-container.sh" +# Dispatch to the chosen entry script. Both primus-cli-container.sh and +# primus-cli-direct.sh already accept --env / --config / --debug, so no +# downstream changes are needed. +script_path="$RUNNER_DIR/primus-cli-${ENTRY_MODE}.sh" require_file "$script_path" "[slurm-entry] Script not found: $script_path" # Build full command diff --git a/runner/primus-cli-slurm.sh b/runner/primus-cli-slurm.sh index b98d5c1e2..4c8194ffb 100755 --- a/runner/primus-cli-slurm.sh +++ b/runner/primus-cli-slurm.sh @@ -17,8 +17,11 @@ Usage: Description: Launch distributed Primus jobs via Slurm. - Everything before the first '--' is passed to Slurm (srun/sbatch and flags). - - specifies Primus execution mode: container | direct | preflight (see below). + - specifies Primus execution mode: container | direct (default: container). - The second '--' (if any) separates Primus entry args from Primus CLI arguments. + - If is anything other than 'container' or 'direct' (e.g. a primus + subcommand like 'preflight'), it is routed through the default container + chain unchanged. Options: --config FILE Load configuration from specified file @@ -32,8 +35,14 @@ Examples: # Launch with sbatch, log to file, run benchmark primus-cli slurm sbatch --output=run.log -N 2 -- container -- benchmark gemm -M 4096 -N 4096 -K 4096 - # Run preflight environment check across 4 nodes - primus-cli slurm srun -N 4 -- preflight + # Run preflight environment check across 4 nodes (terse form: defaults to container entry) + primus-cli slurm srun -N 4 -- preflight --quick + + # Run preflight via the direct entry (no container; uses host python or VENV_ACTIVATE) + primus-cli slurm srun -N 4 -- direct -- preflight --quick + + # Run node smoke test via direct entry (auto-selects single mode for node_smoke) + primus-cli slurm srun -N 4 -- direct -- node_smoke --tier2-perf # Dry-run to see what would be executed primus-cli slurm --dry-run srun -N 4 -- container -- train diff --git a/tests/runner/test_primus_cli_direct.sh b/tests/runner/test_primus_cli_direct.sh index a7c6582cb..86ea1a542 100755 --- a/tests/runner/test_primus_cli_direct.sh +++ b/tests/runner/test_primus_cli_direct.sh @@ -338,6 +338,277 @@ test_help_output() { assert_contains "$output" "--numa" "NUMA option documented" assert_contains "$output" "--env" "Env option documented" assert_contains "$output" "--patch" "Patch option documented" + assert_contains "$output" "--silent" "Silent flag documented" + assert_contains "$output" "VENV_ACTIVATE" "VENV_ACTIVATE documented" + assert_contains "$output" "SLURM" "SLURM auto-derivation documented" +} + +# ============================================================================ +# Test 11: VENV_ACTIVATE handling (R1 -- consolidate-preflight-direct-wrappers) +# ============================================================================ +test_venv_activate() { + local_print_section "Test 11: VENV_ACTIVATE (R1)" + + # Sub-test 11a: VENV_ACTIVATE unset = no-op (dry-run succeeds without + # any "VENV_ACTIVATE" error, and the script never tries to source a file). + local out_unset + out_unset=$(unset VENV_ACTIVATE; timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true) + assert_contains "$out_unset" "End of Dry Run" "Unset VENV_ACTIVATE is a no-op (dry-run completes)" + assert_not_contains "$out_unset" "VENV_ACTIVATE is set but" "No spurious 'missing file' error when unset" + assert_not_contains "$out_unset" "Activated virtualenv:" "No 'Activated virtualenv' message when unset" + + # Sub-test 11b: VENV_ACTIVATE set + valid file = sourced. + local tmpvenv + tmpvenv="$(mktemp)" + echo 'export PRIMUS_TEST_VENV_SOURCED=yes' > "$tmpvenv" + local out_valid + out_valid=$(VENV_ACTIVATE="$tmpvenv" timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true) + assert_contains "$out_valid" "Activated virtualenv:" "Valid VENV_ACTIVATE is sourced" + rm -f "$tmpvenv" + + # Sub-test 11c: VENV_ACTIVATE set + missing file = fail-fast (LOG_ERROR + # on stderr, non-zero exit). + local ec_missing=0 + VENV_ACTIVATE=/does/not/exist/anywhere timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm >/dev/null 2>/tmp/test_venv_err_$$ || ec_missing=$? + if [[ "$ec_missing" -ne 0 ]] && grep -q "VENV_ACTIVATE is set but file does not exist" /tmp/test_venv_err_$$; then + assert_pass "Missing VENV_ACTIVATE file fails fast with clear error" + else + assert_fail "Missing VENV_ACTIVATE file should fail fast" \ + "exit=$ec_missing stderr=$(cat /tmp/test_venv_err_$$)" + fi + rm -f /tmp/test_venv_err_$$ +} + +# ============================================================================ +# Test 12: SLURM env derivation (R2) +# ============================================================================ +test_slurm_env_derivation() { + local_print_section "Test 12: SLURM env derivation (R2)" + + # Sub-test 12a: SLURM_JOB_ID + SLURM_NNODES + SLURM_NODEID -> NNODES / + # NODE_RANK derived. + local out_slurm + out_slurm=$(SLURM_JOB_ID=999 SLURM_NNODES=4 SLURM_NODEID=0 SLURM_NODELIST=tus1-p3-g25 \ + timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true) + assert_contains "$out_slurm" "SLURM detected" "SLURM detection log fires" + assert_contains "$out_slurm" "NNODES=4" "NNODES derived from SLURM_NNODES" + assert_contains "$out_slurm" "--nnodes 4" "torchrun gets --nnodes 4" + + # Sub-test 12b: pre-exported NNODES wins over SLURM_NNODES. + local out_preset + out_preset=$(SLURM_JOB_ID=999 SLURM_NNODES=4 NNODES=7 NODE_RANK=0 SLURM_NODELIST=tus1-p3-g25 \ + timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true) + assert_contains "$out_preset" "NNODES=7" "Pre-exported NNODES=7 wins over SLURM_NNODES=4" + assert_contains "$out_preset" "--nnodes 7" "torchrun honors pre-exported NNODES=7" + + # Sub-test 12c: sanity check rejects NODE_RANK >= NNODES. + local ec_sanity=0 + NNODES=2 NODE_RANK=5 timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm >/dev/null 2>/tmp/test_slurm_err_$$ || ec_sanity=$? + if [[ "$ec_sanity" -ne 0 ]] && grep -q "NODE_RANK (5) must be < NNODES (2)" /tmp/test_slurm_err_$$; then + assert_pass "Sanity check rejects NODE_RANK >= NNODES" + else + assert_fail "Sanity check should reject NODE_RANK >= NNODES" \ + "exit=$ec_sanity stderr=$(cat /tmp/test_slurm_err_$$)" + fi + rm -f /tmp/test_slurm_err_$$ + + # Sub-test 12d: SLURM context without SLURM_NODELIST must still complete. + # This exercises the same code path as "no scontrol on the host" (the + # `command -v scontrol && [[ -n "$SLURM_NODELIST" ]]` short-circuits to + # false either way), and used to crash with "MASTER_ADDR: unbound + # variable" under set -u. The launcher must fall back to MASTER_ADDR=localhost + # and the dry-run must reach the torchrun command line. + local out_no_nodelist + out_no_nodelist=$(env -u SLURM_NODELIST SLURM_JOB_ID=999 SLURM_NNODES=4 SLURM_NODEID=0 \ + timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true) + assert_not_contains "$out_no_nodelist" "unbound variable" \ + "SLURM context with no SLURM_NODELIST does not crash on unbound MASTER_ADDR" + assert_contains "$out_no_nodelist" "MASTER_ADDR=localhost" \ + "MASTER_ADDR falls back to localhost when scontrol/NODELIST unavailable" + assert_contains "$out_no_nodelist" "--nnodes 4" \ + "torchrun cmd still gets --nnodes 4 even without NODELIST" +} + +# ============================================================================ +# Test 13: Auto-single run_mode for node_smoke +# ============================================================================ +test_auto_single_for_node_smoke() { + local_print_section "Test 13: Auto-single run_mode for node_smoke" + + # Sub-test 13a: preflight stays in torchrun mode by default. + local out_preflight + out_preflight=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- preflight --quick 2>&1 || true) + assert_contains "$out_preflight" "Run Mode : torchrun" "preflight defaults to torchrun" + assert_contains "$out_preflight" "torchrun --nproc_per_node" "preflight uses torchrun command" + + # Sub-test 13b: node_smoke auto-selects single mode. + local out_smoke + out_smoke=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- node_smoke --tier2-perf 2>&1 || true) + assert_contains "$out_smoke" "Auto-selected run_mode=single for subcommand 'node_smoke'" \ + "node_smoke auto-detect fires" + assert_contains "$out_smoke" "Run Mode : single" "node_smoke ends up in single mode" + assert_contains "$out_smoke" "python3" "node_smoke uses python3 launcher" + assert_not_contains "$out_smoke" "torchrun --nproc_per_node" "node_smoke does NOT use torchrun" + + # Sub-test 13c: explicit --single on a non-node_smoke subcommand still + # works (regression guard). + local out_explicit + out_explicit=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --single --dry-run -- benchmark gemm 2>&1 || true) + assert_contains "$out_explicit" "Run Mode : single" "Explicit --single still works for benchmark" +} + +# ============================================================================ +# Test 14: --silent contract (bash-level) +# ============================================================================ +test_silent_flag() { + local_print_section "Test 14: --silent flag contract" + + # Sub-test 14a: --silent produces empty stdout. + local out_silent="/tmp/test_silent_out_$$" + local err_silent="/tmp/test_silent_err_$$" + timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --silent --dry-run -- benchmark gemm >"$out_silent" 2>"$err_silent" || true + if [[ ! -s "$out_silent" ]]; then + assert_pass "--silent silences stdout (out file is empty)" + else + assert_fail "--silent should silence stdout" "size=$(wc -c <"$out_silent") head=$(head -3 "$out_silent")" + fi + rm -f "$out_silent" "$err_silent" + + # Sub-test 14b: --silent does NOT silence launcher errors (stderr survives). + local err_err_silent="/tmp/test_silent_err2_$$" + VENV_ACTIVATE=/does/not/exist timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --silent --dry-run -- benchmark gemm >/dev/null 2>"$err_err_silent" || true + if grep -q "VENV_ACTIVATE is set but file does not exist" "$err_err_silent"; then + assert_pass "--silent preserves launcher LOG_ERROR on stderr" + else + assert_fail "--silent should preserve launcher errors on stderr" \ + "stderr=$(cat "$err_err_silent")" + fi + rm -f "$err_err_silent" + + # Sub-test 14c: --silent is consumed by the launcher and NOT forwarded + # to the python tool. The forwarded args list (logged before silencing + # applied? actually after -- it shouldn't appear at all). Easiest check: + # the dry-run "Would Execute" line must NOT contain --silent. + local out_no_forward + out_no_forward=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- benchmark gemm 2>&1 || true) + # Sanity: this baseline run has no --silent anywhere. + assert_not_contains "$out_no_forward" "--silent" "Baseline dry-run has no --silent leakage" +} + +# ============================================================================ +# Test 15: slurm-entry direct mode dispatch +# ============================================================================ +test_slurm_entry_direct_dispatch() { + local_print_section "Test 15: slurm-entry direct/container dispatch" + + # NOTE: primus-cli-slurm-entry.sh resolves the node list via scontrol when + # available and falls back to parsing SLURM_NODELIST directly otherwise, so + # this test works both on-cluster (scontrol present) and in CI containers + # (scontrol absent). We set a plain SLURM_NODELIST so both paths agree. + # Sub-test 15a: -- direct -- preflight ... routes through primus-cli-direct.sh. + local out_direct + out_direct=$(SLURM_NODELIST=tus1-p3-g25 SLURM_JOB_ID=1 SLURM_NNODES=1 SLURM_NODEID=0 \ + timeout 30 bash "$RUNNER_DIR/primus-cli-slurm-entry.sh" --dry-run -- direct -- preflight --quick 2>&1 || true) + assert_contains "$out_direct" "Entry mode: direct" "slurm-entry parses 'direct' keyword" + assert_contains "$out_direct" "primus-cli-direct.sh" "slurm-entry dispatches to primus-cli-direct.sh" + + # Sub-test 15b: -- container -- ... routes through primus-cli-container.sh (existing path). + local out_container + out_container=$(SLURM_NODELIST=tus1-p3-g25 SLURM_JOB_ID=1 SLURM_NNODES=1 SLURM_NODEID=0 \ + timeout 30 bash "$RUNNER_DIR/primus-cli-slurm-entry.sh" --dry-run -- container -- train pretrain 2>&1 || true) + assert_contains "$out_container" "Entry mode: container" "slurm-entry parses 'container' keyword" + assert_contains "$out_container" "primus-cli-container.sh" "slurm-entry dispatches to primus-cli-container.sh" + + # Sub-test 15c: terse form `-- preflight` (no keyword) defaults to container. + local out_terse + out_terse=$(SLURM_NODELIST=tus1-p3-g25 SLURM_JOB_ID=1 SLURM_NNODES=1 SLURM_NODEID=0 \ + timeout 30 bash "$RUNNER_DIR/primus-cli-slurm-entry.sh" --dry-run -- preflight --quick 2>&1 || true) + assert_contains "$out_terse" "Entry mode: container" "Terse form defaults to container" + assert_contains "$out_terse" "primus-cli-container.sh" "Terse form dispatches to primus-cli-container.sh" +} + +# ============================================================================ +# Test 16: RUN_MODE env override (framework prepare-hook contract) +# +# Framework hooks (e.g. runner/helpers/hooks/train/pretrain/maxtext/prepare.py) +# emit `env.RUN_MODE=single` so JAX/MaxText runs as `python3 ...` instead of +# `torchrun ...`. This used to launch correctly but display incorrectly: +# STEP 10 showed `Run Mode: torchrun` and printed the Distributed Settings +# block, because both reads went against the pre-hook direct_config[run_mode] +# instead of the final $RUN_MODE. +# +# This test pins the post-fix contract: when RUN_MODE is exported in the env +# before primus-cli-direct.sh runs (a sufficient stand-in for "a hook +# exported it"), every visible knob -- the displayed Run Mode, the +# Distributed Settings gate, and the actual launch command -- reflects the +# env-override, not the pre-hook default. +# ============================================================================ +test_run_mode_env_override() { + local_print_section "Test 16: RUN_MODE env override (framework hook contract)" + + # Baseline: `train pretrain` with NO hook override -- the launcher's + # auto-default (torchrun) wins and the display shows torchrun + the + # Distributed Settings block. This anchors the "before" state we are + # protecting users from. + local out_baseline + out_baseline=$(timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" --dry-run -- train pretrain 2>&1 || true) + assert_contains "$out_baseline" "Run Mode : torchrun" \ + "Baseline train pretrain shows torchrun (no env override)" + assert_contains "$out_baseline" "Distributed Settings:" \ + "Baseline train pretrain prints Distributed Settings" + assert_contains "$out_baseline" "torchrun --nproc_per_node" \ + "Baseline train pretrain uses torchrun in Full Command" + + # Sub-test 16a: pre-exporting RUN_MODE=single (the MaxText hook's + # effect, modeled in-process so the test doesn't depend on the + # framework hook actually firing) must flip ALL three views: + # - displayed "Run Mode" line + # - "Distributed Settings" gate + # - launch command (`python3 ...`, not `torchrun ...`) + local out_override + out_override=$(RUN_MODE=single timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" \ + --dry-run -- train pretrain 2>&1 || true) + assert_contains "$out_override" "Run Mode : single" \ + "RUN_MODE=single env override surfaces in displayed Run Mode" + assert_not_contains "$out_override" "Distributed Settings:" \ + "RUN_MODE=single env override suppresses Distributed Settings block" + assert_not_contains "$out_override" "torchrun --nproc_per_node" \ + "RUN_MODE=single env override drops torchrun from Full Command" + assert_contains "$out_override" "python3" \ + "RUN_MODE=single env override uses python3 launcher in Full Command" + + # Sub-test 16b: env-override beats an EXPLICIT --single on the CLI. + # This matches the launcher's documented precedence at STEP 9: + # RUN_MODE="${RUN_MODE:-${direct_config[run_mode]:-torchrun}}" + # The env layer wins because the hook layer (which is where RUN_MODE + # actually originates in real runs) knows things the user / config + # can't and gets the final word. We use --single here purely as a + # convenient stand-in for "user/config set direct_config[run_mode]" + # -- there is no symmetric --torchrun CLI flag. + local out_vs_cli + out_vs_cli=$(RUN_MODE=torchrun timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" \ + --single --dry-run -- benchmark gemm 2>&1 || true) + assert_contains "$out_vs_cli" "Run Mode : torchrun" \ + "RUN_MODE=torchrun env override wins over --single CLI flag" + assert_contains "$out_vs_cli" "torchrun --nproc_per_node" \ + "RUN_MODE=torchrun env override produces torchrun launcher despite --single" + + # Sub-test 16c: exporting RUN_MODE must override even the auto-single + # detection for node_smoke. Auto-detect sets direct_config[run_mode] + # BEFORE the env-override check at line 610, so env still wins. This + # is the corner case behind the bug we fixed: the OLD code printed + # `Run Mode: single` here (from direct_config[run_mode]) even though + # the launch command was actually torchrun -- a display divergence + # in the exact opposite direction from the MaxText case. + local out_reverse + out_reverse=$(RUN_MODE=torchrun timeout 30 bash "$RUNNER_DIR/primus-cli-direct.sh" \ + --dry-run -- node_smoke --tier2-perf 2>&1 || true) + assert_contains "$out_reverse" "Run Mode : torchrun" \ + "RUN_MODE=torchrun env override surfaces in display for node_smoke" + assert_contains "$out_reverse" "Distributed Settings:" \ + "RUN_MODE=torchrun env override re-enables Distributed Settings for node_smoke" + assert_contains "$out_reverse" "torchrun --nproc_per_node" \ + "RUN_MODE=torchrun env override produces torchrun launch for node_smoke" } # ============================================================================ @@ -358,6 +629,13 @@ main() { test_debug_mode test_config_priority test_help_output + # New tests from the consolidate-preflight-direct-wrappers refactor: + test_venv_activate + test_slurm_env_derivation + test_auto_single_for_node_smoke + test_silent_flag + test_slurm_entry_direct_dispatch + test_run_mode_env_override # Print summary echo "" diff --git a/tests/unit_tests/cli/test_preflight_subcommand.py b/tests/unit_tests/cli/test_preflight_subcommand.py index 6ffb4eeea..47a997b5d 100644 --- a/tests/unit_tests/cli/test_preflight_subcommand.py +++ b/tests/unit_tests/cli/test_preflight_subcommand.py @@ -39,7 +39,9 @@ def test_defaults(): assert args.perf_test is False assert args.dist_timeout_sec == 120 assert args.dump_path == "output/preflight" - assert args.report_file_name == "preflight_report" + # Default is None so the tool auto-generates a unique timestamped name + # (preflight-{NNODES}N-{YYYYMMDD-HHMMSS}) at run time. + assert args.report_file_name is None assert args.save_pdf is True diff --git a/tests/unit_tests/tools/test_preflight_bisect_slurm.py b/tests/unit_tests/tools/test_preflight_bisect_slurm.py new file mode 100644 index 000000000..a4267e250 --- /dev/null +++ b/tests/unit_tests/tools/test_preflight_bisect_slurm.py @@ -0,0 +1,223 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Integration tests for tools/preflight_bisect/bisect.py. These require a live +# Slurm allocation and skip automatically when the required environment is +# missing. +# +# Required env vars +# ----------------- +# Both tests require a Slurm nodelist. If BISECT_NODELIST is unset, the tests +# use SLURM_NODELIST from the current allocation. +# +# Running from inside a Slurm allocation +# -------------------------- +# cd ~/Primus +# export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate +# python3 -m pytest tests/unit_tests/tools/test_preflight_bisect_slurm.py -v +# More complex example +# -------------------------- +# export NCCL_SOCKET_IFNAME=tw-eth0 +# export GLOO_SOCKET_IFNAME=tw-eth0 +# export NCCL_IB_HCA="rdma0:1,rdma1:1,rdma2:1,rdma3:1,rdma4:1,rdma5:1,rdma6:1,rdma7:1" +# export BISECT_PREFLIGHT_ENV="USING_AINIC=1 NCCL_IB_GID_INDEX=3 NCCL_CROSS_NIC=0 NCCL_PXN_DISABLE=0" +# +# python3 -m pytest \ +# tests/unit_tests/tools/test_preflight_bisect_slurm.py \ +# -v \ +# --basetemp=/tmp/preflight-bisect-pytest +# +# Inspect pytest artifacts: +# ls -R /tmp/preflight-bisect-pytest +# cat /tmp/preflight-bisect-pytest/test_bisect_all_nodes_pass*/trial-*.log +# cat /tmp/preflight-bisect-pytest/test_bisect_all_nodes_pass*/summary.txt +# +# Defaults +# -------- +# - BISECT_NODELIST defaults to SLURM_NODELIST. +# - BISECT_PARTITION defaults to SLURM_JOB_PARTITION; otherwise --partition is omitted. +# - BISECT_BAD_NODE defaults to the last host in the resolved nodelist. +# +# Optional tuning: +# BISECT_TRIAL_TIMEOUT, BISECT_SLURM_TIME, BISECT_PREFLIGHT_ENV +# Keep comma-containing values such as NCCL_IB_HCA as normal exported env vars; +# BISECT_PREFLIGHT_ENV is passed through srun --export=ALL,... +# +############################################################################### + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[3] +BISECT_PY = REPO_ROOT / "tools" / "preflight_bisect" / "bisect.py" +FAKE_RUNNER = REPO_ROOT / "tools" / "preflight_bisect" / "fake_runner.sh" +REAL_RUNNER = REPO_ROOT / "runner" / "primus-cli" + + +def _require_env(*names: str) -> dict[str, str]: + """Return a dict of the requested env var values, skipping the test if any are missing.""" + missing = [n for n in names if not os.environ.get(n)] + if missing: + pytest.skip(f"Required env var(s) not set: {', '.join(missing)}") + return {n: os.environ[n] for n in names} + + +def _resolve_nodelist() -> str: + nodelist = os.environ.get("BISECT_NODELIST") or os.environ.get("SLURM_NODELIST") + if not nodelist: + pytest.skip("Set BISECT_NODELIST or run from inside a Slurm allocation with SLURM_NODELIST set") + return nodelist + + +def _resolve_partition() -> str: + return os.environ.get("BISECT_PARTITION") or os.environ.get("SLURM_JOB_PARTITION") or "" + + +def _resolve_hosts(nodelist: str) -> list[str]: + try: + result = subprocess.run( + ["scontrol", "show", "hostnames", nodelist], + capture_output=True, + check=True, + text=True, + ) + except FileNotFoundError: + pytest.skip("scontrol not found; run this test on a Slurm login/head node") + except subprocess.CalledProcessError as exc: + pytest.skip(f"scontrol failed to resolve nodelist {nodelist!r}: {exc.stderr.strip()}") + + hosts = [line.strip() for line in result.stdout.splitlines() if line.strip()] + if not hosts: + pytest.skip(f"scontrol produced no hostnames for nodelist {nodelist!r}") + return hosts + + +def _resolve_bad_node(nodelist: str) -> str: + return os.environ.get("BISECT_BAD_NODE") or _resolve_hosts(nodelist)[-1] + + +def _resolve_bisect_env() -> dict[str, str]: + env = {"BISECT_NODELIST": _resolve_nodelist()} + partition = _resolve_partition() + if partition: + env["BISECT_PARTITION"] = partition + return env + + +def _run_bisect(extra_args: list[str], env: dict[str, str], tmp_path: Path, timeout: int) -> str: + """Invoke bisect.py as a subprocess and return the contents of summary.txt.""" + preflight_env_args: list[str] = [] + for kv in os.environ.get("BISECT_PREFLIGHT_ENV", "").split(): + preflight_env_args += ["--preflight-env", kv] + + cmd = [ + sys.executable, + str(BISECT_PY), + "--nodelist", + env["BISECT_NODELIST"], + "--output-dir", + str(tmp_path), + ] + if env.get("BISECT_PARTITION"): + cmd.extend(["--partition", env["BISECT_PARTITION"]]) + cmd.extend([*preflight_env_args, *extra_args]) + + result = subprocess.run( + cmd, + env={**os.environ, **env}, + capture_output=True, + text=True, + timeout=timeout, + ) + summary_path = tmp_path / "summary.txt" + summary_text = summary_path.read_text(encoding="utf-8") if summary_path.exists() else "" + + if result.returncode != 0: + pytest.fail( + f"bisect.py exited {result.returncode} and wrote no summary.\n" + f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + + return summary_text + + +def test_bisect_all_nodes_pass(tmp_path): + """Run bisect.py with the real preflight runner against a known-healthy nodeset. + + Requires VENV_ACTIVATE and either BISECT_NODELIST or SLURM_NODELIST. + All nodes are expected to pass the preflight perf-test, so bisect should + report SUSPECT_NODES: (none). + """ + env = {**_resolve_bisect_env(), **_require_env("VENV_ACTIVATE")} + trial_timeout = int(os.environ.get("BISECT_TRIAL_TIMEOUT", "600")) + slurm_time = os.environ.get("BISECT_SLURM_TIME", "00:15:00") + + # subprocess timeout: give a comfortable margin above the per-trial timeout + # to account for bisect recursion and Slurm scheduling overhead. + # Each level of bisection can run up to 2 concurrent trials, and a nodeset + # of N nodes has at most log2(N)+1 levels, so multiply generously. + subprocess_timeout = trial_timeout * 8 + + summary_text = _run_bisect( + extra_args=[ + "--trial-timeout-sec", + str(trial_timeout), + "--slurm-time", + slurm_time, + "--runner", + str(REAL_RUNNER), + ], + env=env, + tmp_path=tmp_path, + timeout=subprocess_timeout, + ) + + assert ( + "SUSPECT_NODES: (none)" in summary_text + ), f"Expected no suspect nodes for a healthy nodeset, but got:\n{summary_text}" + + +def test_bisect_identifies_bad_node(tmp_path): + """Run bisect.py with fake_runner.sh seeding one bad node. + + Requires either BISECT_NODELIST or SLURM_NODELIST. + BISECT_BAD_NODE can override the default bad node, which is the last host in + the resolved nodelist. + bisect.py is expected to identify exactly that node as the sole suspect. + """ + env = _resolve_bisect_env() + bad_node = _resolve_bad_node(env["BISECT_NODELIST"]) + env["BISECT_BAD_NODE"] = bad_node + trial_timeout = int(os.environ.get("BISECT_TRIAL_TIMEOUT", "30")) + slurm_time = os.environ.get("BISECT_SLURM_TIME", "00:02:00") + + subprocess_timeout = trial_timeout * 8 + + summary_text = _run_bisect( + extra_args=[ + "--trial-timeout-sec", + str(trial_timeout), + "--slurm-time", + slurm_time, + "--runner", + str(FAKE_RUNNER), + "--preflight-env", + f"BAD_NODE={bad_node}", + "--max-concurrent-trials", + "2", + ], + env=env, + tmp_path=tmp_path, + timeout=subprocess_timeout, + ) + + assert ( + f"SUSPECT_NODES: {bad_node}" in summary_text + ), f"Expected '{bad_node}' to be identified as the sole suspect, but got:\n{summary_text}" diff --git a/tools/preflight_bisect/bisect.py b/tools/preflight_bisect/bisect.py new file mode 100644 index 000000000..3c0c9f2c9 --- /dev/null +++ b/tools/preflight_bisect/bisect.py @@ -0,0 +1,352 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# Minimal Slurm nodelist bisection for Primus preflight --perf-test. +# Run from repo root (or any cwd; script resolves repo root for runner/ path). +# +# Usage (typical): +# export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate +# cd /path/to/Primus +# python tools/preflight_bisect/bisect.py --nodelist "node[01-32]" -p gpus ... +# +# Caveats (see also --help): +# - Scale-only hangs: subsets may all PASS while full N fails; suspects may be empty. +# - Multiple bad nodes: union of singleton suspects from failing subtrees. +# - Tune --trial-timeout-sec to ~2-3x healthy full-N runtime; too short causes false HANG. +# - --scancel-user-on-hang kills ALL your Slurm jobs; do not use if you have other work. +############################################################################### +from __future__ import annotations + +import argparse +import os +import signal +import subprocess +import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + + +def _repo_root() -> Path: + # tools/preflight_bisect/bisect.py -> repo root is parent.parent.parent + return Path(__file__).resolve().parent.parent.parent + + +def expand_nodelist(nodelist: str) -> list[str]: + out = subprocess.check_output( + ["scontrol", "show", "hostnames", nodelist], + text=True, + stderr=subprocess.PIPE, + ) + hosts = [ln.strip() for ln in out.splitlines() if ln.strip()] + if not hosts: + raise SystemExit(f"scontrol produced no hostnames for nodelist={nodelist!r}") + return hosts + + +def _format_node_range(nodes: list[str]) -> str: + if not nodes: + return "" + if len(nodes) == 1: + return nodes[0] + return f"{nodes[0]}..{nodes[-1]}" + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed < 1: + raise argparse.ArgumentTypeError("must be >= 1") + return parsed + + +@dataclass +class BisectState: + max_concurrent_trials: int + idx: int = 0 + trials: list[dict[str, Any]] = field(default_factory=list) + _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False) + _trial_slots: threading.BoundedSemaphore = field(init=False, repr=False) + + def __post_init__(self) -> None: + self._trial_slots = threading.BoundedSemaphore(self.max_concurrent_trials) + + def next_trial_idx(self) -> int: + with self._lock: + idx = self.idx + self.idx += 1 + return idx + + def record_trial(self, idx: int, nodes: list[str], status: str) -> None: + record = {"idx": idx, "n": len(nodes), "status": status, "nodes": list(nodes)} + with self._lock: + self.trials.append(record) + + def ordered_trials(self) -> list[dict[str, Any]]: + with self._lock: + return sorted(self.trials, key=lambda trial: trial["idx"]) + + def acquire_trial_slot(self) -> None: + self._trial_slots.acquire() + + def release_trial_slot(self) -> None: + self._trial_slots.release() + + +def run_trial( + nodes: list[str], + trial_idx: int, + state: BisectState, + args: argparse.Namespace, + runner: Path, + out_dir: Path, +) -> str: + """Run one preflight perf trial. Returns 'pass', 'fail', or 'hang'.""" + subset = ",".join(nodes) + log_path = out_dir / f"trial-{trial_idx:03d}.log" + cmd: list[str] = [ + "srun", + f"-N{len(nodes)}", + f"--nodelist={subset}", + "-n", + str(len(nodes)), + "--ntasks-per-node=1", + "-c", + str(args.cpus_per_task), + f"--gres=gpu:{args.gpus_per_node}", + f"-t{args.slurm_time}", + ] + if args.partition: + cmd.extend(["-p", args.partition]) + # Per-trial env overrides are propagated via srun --export so every rank on + # every node sees them (the consolidated primus-cli launcher does accept + # --env KEY=VALUE, but only rank 0 would see those; --export covers all + # ranks). ALL keeps the caller's environment (notably VENV_ACTIVATE) + # intact; the trailing K=V pairs override / add on top. SLURM tokenizes + # --export on commas, so values must not contain ',' or whitespace + # (NCCL flags never do). + export_val = "ALL" + if args.preflight_env: + export_val += "," + ",".join(args.preflight_env) + cmd.append(f"--export={export_val}") + cmd.append(str(runner)) + cmd.extend( + [ + "direct", + "--", + "preflight", + "--perf-test", + "--report-file-name", + f"trial-{trial_idx:03d}", + ] + ) + + state.acquire_trial_slot() + try: + header = ( + f"CMD: {' '.join(cmd)}\n" + f"NODES ({len(nodes)}): {subset}\n" + f"START: {datetime.now(timezone.utc).isoformat()}\n\n" + ) + print( + f"[trial {trial_idx:03d}] N={len(nodes)} {_format_node_range(nodes)} -> {log_path.name}", + flush=True, + ) + log_path.parent.mkdir(parents=True, exist_ok=True) + with log_path.open("wb", buffering=0) as logf: + logf.write(header.encode()) + + proc = subprocess.Popen( + cmd, + cwd=str(_repo_root()), + stdout=logf, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + try: + rc = proc.wait(timeout=args.trial_timeout_sec) + return "pass" if rc == 0 else "fail" + except subprocess.TimeoutExpired: + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + try: + proc.wait(timeout=60) + except subprocess.TimeoutExpired: + pass # process is stuck in D-state; SIGKILL is pending, move on + if args.scancel_user_on_hang: + subprocess.run( + ["scancel", "--signal=KILL", "--user", os.environ.get("USER", "")], + stderr=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + ) + return "hang" + finally: + state.release_trial_slot() + + +def bisect( + nodes: list[str], + state: BisectState, + args: argparse.Namespace, + runner: Path, + out_dir: Path, +) -> list[str]: + idx = state.next_trial_idx() + status = run_trial(nodes, idx, state, args, runner, out_dir) + state.record_trial(idx, nodes, status) + + if status == "pass": + return [] + if len(nodes) == 1: + return list(nodes) + + mid = len(nodes) // 2 + if args.max_concurrent_trials == 1: + left = bisect(nodes[:mid], state, args, runner, out_dir) + right = bisect(nodes[mid:], state, args, runner, out_dir) + return left + right + + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="preflight-bisect") as executor: + right_future = executor.submit(bisect, nodes[mid:], state, args, runner, out_dir) + left = bisect(nodes[:mid], state, args, runner, out_dir) + right = right_future.result() + return left + right + + +def write_summary(out_dir: Path, nodes: list[str], suspects: list[str], trials: list[dict[str, Any]]) -> None: + path = out_dir / "summary.txt" + lines = [ + f"{datetime.now(timezone.utc).isoformat()} bisect nodes={len(nodes)}", + ] + for t in sorted(trials, key=lambda trial: trial["idx"]): + nlist = t["nodes"] + r = _format_node_range(nlist) + lines.append(f"[{t['idx']:03d}] N={t['n']:2d} {t['status'].upper():4s} nodes={r}") + if suspects: + lines.append("SUSPECT_NODES: " + " ".join(sorted(set(suspects)))) + else: + lines.append("SUSPECT_NODES: (none)") + summary = "\n".join(lines) + "\n" + path.write_text(summary, encoding="utf-8") + print(summary, end="") + + +def main() -> int: + p = argparse.ArgumentParser( + description="Recursively bisect a Slurm nodelist using Primus preflight --perf-test.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Environment: + Export VENV_ACTIVATE to your venv activate script before running (required by + runner/primus-cli direct -> primus-cli-direct.sh). See docs/preflight-direct.md. + +Caveats: + - If the hang only reproduces at full scale, all subsets may PASS -> SUSPECT_NODES empty. + - Multiple faulty nodes yield a union of suspects from failing singleton trials. + - By default, failing sibling subsets launch in parallel (up to 2 concurrent trials). + - --trial-timeout-sec too low marks healthy runs as HANG. + - --scancel-user-on-hang cancels ALL jobs for $USER; only use it with --max-concurrent-trials=1. +""", + ) + p.add_argument("--nodelist", required=True, help='Slurm nodelist expression, e.g. "node[01-32]"') + p.add_argument("-p", "--partition", default="", help="Slurm partition (-p), optional") + p.add_argument( + "--output-dir", + type=Path, + default=Path("bisect-out"), + help="Directory for trial-*.log and summary.txt (default: ./bisect-out)", + ) + p.add_argument( + "--trial-timeout-sec", + type=int, + default=900, + help="Wall-clock timeout per trial in seconds (default: 900)", + ) + p.add_argument( + "--slurm-time", + default="00:45:00", + help="srun -t limit per trial (default: 00:45:00)", + ) + p.add_argument( + "--max-concurrent-trials", + type=_positive_int, + default=2, + help="Maximum concurrent subset trials (default: 2). Set to 1 to force sequential execution.", + ) + p.add_argument("--cpus-per-task", type=int, default=128, help="srun -c (default: 128)") + p.add_argument( + "--gpus-per-node", + type=int, + default=8, + help="GPUs per node; emitted as srun --gres=gpu:N (default: 8)", + ) + p.add_argument( + "--preflight-env", + action="append", + default=[], + metavar="KEY=VALUE", + help=( + "Repeatable; propagated into each trial via " + "'srun --export=ALL,KEY=VALUE,...'. Values must not contain ',' or whitespace." + ), + ) + p.add_argument( + "--runner", + type=Path, + default=None, + help=( + "Override path to the launcher invoked per trial " + "(default: repo runner/primus-cli). The bisector always appends " + "'direct -- preflight --perf-test --report-file-name ...' after " + "the runner path; custom runners that ignore positional args " + "(e.g. fake_runner.sh) work transparently." + ), + ) + p.add_argument( + "--scancel-user-on-hang", + action="store_true", + help="On timeout, also run: scancel --signal=KILL --user $USER (DANGEROUS)", + ) + args = p.parse_args() + if args.scancel_user_on_hang and args.max_concurrent_trials > 1: + p.error("--scancel-user-on-hang is only supported with --max-concurrent-trials=1") + + if not os.environ.get("VENV_ACTIVATE"): + print( + "ERROR: VENV_ACTIVATE is not set. Export the path to your venv's " + "bin/activate before running, e.g.\n" + " export VENV_ACTIVATE=~/envs/preflight/.venv/bin/activate", + file=sys.stderr, + ) + return 2 + + runner = args.runner or (_repo_root() / "runner" / "primus-cli") + if not runner.is_file(): + print(f"ERROR: runner not found: {runner}", file=sys.stderr) + return 2 + + try: + hosts = expand_nodelist(args.nodelist) + except subprocess.CalledProcessError as e: + print(f"ERROR: scontrol failed: {e}", file=sys.stderr) + return 2 + except FileNotFoundError: + print("ERROR: scontrol not found; run this script on a Slurm login/head node.", file=sys.stderr) + return 2 + + out_dir = args.output_dir.resolve() + state = BisectState(max_concurrent_trials=args.max_concurrent_trials) + suspects = bisect(hosts, state, args, runner, out_dir) + write_summary(out_dir, hosts, suspects, state.ordered_trials()) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/preflight_bisect/fake_runner.sh b/tools/preflight_bisect/fake_runner.sh new file mode 100755 index 000000000..3068b3904 --- /dev/null +++ b/tools/preflight_bisect/fake_runner.sh @@ -0,0 +1,54 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# fake_runner.sh — drop-in runner replacement for bisect.py testing. +# +# Simulates a bad node without running a real preflight, so that +# bisect.py's identification logic can be validated on a live cluster. +# +# Environment variables (pass via bisect.py --preflight-env): +# BAD_NODE hostname of the node that should appear faulty (required) +# HANG_SECONDS if >0, the bad node sleeps this long before exiting 1, +# which lets you exercise the --trial-timeout-sec hang path +# (default: 0 → instant failure) +# GOOD_SLEEP seconds a healthy node sleeps to simulate a real run +# (default: 2) +# +# Usage with bisect.py: +# python tools/preflight_bisect/bisect.py \ +# --nodelist "chi2867,chi2879" \ +# --partition mi355x \ +# --output-dir "output/bisect-$(date +%Y%m%d-%H%M%S)" \ +# --trial-timeout-sec 600 \ +# --slurm-time 00:15:00 \ +# --runner tools/preflight_bisect/fake_runner.sh \ +# --preflight-env BAD_NODE=chi2879 +############################################################################### +set -euo pipefail + +BAD_NODE="${BAD_NODE:-}" +HANG_SECONDS="${HANG_SECONDS:-0}" +GOOD_SLEEP="${GOOD_SLEEP:-2}" + +THIS_HOST="$(hostname -s)" + +if [[ -z "$BAD_NODE" ]]; then + echo "[fake_runner] ERROR: BAD_NODE is not set. Pass --preflight-env BAD_NODE= to bisect.py." >&2 + exit 2 +fi + +if [[ "$THIS_HOST" == "$BAD_NODE" ]]; then + echo "[fake_runner] $THIS_HOST == BAD_NODE ($BAD_NODE): simulating failure" + if (( HANG_SECONDS > 0 )); then + echo "[fake_runner] sleeping ${HANG_SECONDS}s to simulate a hang" + sleep "$HANG_SECONDS" + fi + exit 1 +fi + +echo "[fake_runner] $THIS_HOST != BAD_NODE ($BAD_NODE): simulating healthy run (sleep ${GOOD_SLEEP}s)" +sleep "$GOOD_SLEEP" +exit 0 From 05613e34a881fe3cb8781b16da6912f93e9cc445 Mon Sep 17 00:00:00 2001 From: Andy <14128880+yeandy@users.noreply.github.com> Date: Wed, 15 Jul 2026 20:31:30 -0400 Subject: [PATCH 032/127] feat(maxtext): support MaxText v26.4 with v26.3 backward compatibility (#869) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Updates the Primus MaxText launcher to support **MaxText v26.4** while remaining backward-compatible with **v26.3 and earlier**. v26.4 relocated MaxText into a `maxtext.*` package (training loop at `maxtext.trainers.pre_train.train`, config at `maxtext.configs.pyconfig`, logging at `maxtext.utils.max_logging`), whereas older releases exposed `MaxText.*`. In addition, v26.4's `pyconfig` now hard-fails on unknown config fields instead of warning. This PR resolves the correct import layout at runtime so a single Primus checkout works against both v26.4 and v26.3 images. ## Changes - **`maxtext_pretrain_trainer.py`**: add `_resolve_maxtext_train()` that imports `initialize`/`run` from `maxtext.trainers.pre_train.train`, falling back to legacy `MaxText.train`. The resolved module name is used as `argv[0]` for `initialize`. - **Strip Primus-internal params before config export** (`file_sink_level`, `stderr_sink_level`, `sink_level`, `trainable`, `model`), since v26.4's `pyconfig` rejects unknown fields. - **`patches/train_patches.py`**: add `_resolve_train_and_pyconfig()`; skip the override-forwarding patch gracefully if neither layout is importable. - **`patches/logger_patches.py`**: add `_resolve_max_logging()` (`maxtext.utils.max_logging` → `MaxText.max_logging`). - **Submodule**: bump `third_party/maxtext` to `release/v26.4` (`80f431d0`) to match the shipped v26.4 image. ## Testing Validated end-to-end on **2 nodes (MI355X)**, both images reaching steady-state training with matching loss curves: | Image | Result | Steady state | |---|---|---| | `rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0` | steps 0–4, exit 0 | ~1043 TFLOP/s/device | | `rocm/jax-training:maxtext-v26.3` (backward compat) | steps 0–4, exit 0 | ~1066 TFLOP/s/device | Loss descended identically (12.26 → ~12.10) on both, confirming the new `maxtext.*` and legacy `MaxText.*` code paths both work. ## Notes for reviewers - The Python changes are non-breaking on their own (work with either image). The **submodule bump makes v26.4 the repo default** --- .gitmodules | 2 +- primus/_thirdparty.lock | 2 +- .../maxtext/maxtext_pretrain_trainer.py | 43 +++++++++++++++++-- .../maxtext/patches/logger_patches.py | 35 +++++++++++---- .../backends/maxtext/patches/train_patches.py | 31 +++++++++++-- .../configs/modules/maxtext/pre_trainer.yaml | 4 ++ third_party/maxtext | 2 +- 7 files changed, 101 insertions(+), 18 deletions(-) diff --git a/.gitmodules b/.gitmodules index 1d234e782..819157d3a 100644 --- a/.gitmodules +++ b/.gitmodules @@ -8,7 +8,7 @@ [submodule "third_party/maxtext"] path = third_party/maxtext url = https://github.com/ROCm/maxtext.git - branch = release/v26.3 + branch = release/v26.4 [submodule "third_party/Emerging-Optimizers"] path = third_party/Emerging-Optimizers url = https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git diff --git a/primus/_thirdparty.lock b/primus/_thirdparty.lock index 465d3a116..250b6dfb7 100644 --- a/primus/_thirdparty.lock +++ b/primus/_thirdparty.lock @@ -16,7 +16,7 @@ "name": "maxtext", "path": "third_party/maxtext", "url": "https://github.com/ROCm/maxtext.git", - "commit": "851c0935a4e9c9e4cd0771603e3bdf75a6578ad1" + "commit": "80f431d0ebbf6465257e2a5d931046dfb0b465b0" }, { "name": "Emerging-Optimizers", diff --git a/primus/backends/maxtext/maxtext_pretrain_trainer.py b/primus/backends/maxtext/maxtext_pretrain_trainer.py index 17d207ad1..f4b79d8a5 100644 --- a/primus/backends/maxtext/maxtext_pretrain_trainer.py +++ b/primus/backends/maxtext/maxtext_pretrain_trainer.py @@ -35,6 +35,39 @@ warning_rank_0, ) +# Primus-internal params that are not part of MaxText's config schema. MaxText +# v26.4's pyconfig raises on unknown fields (v26.3 merely warns), so these must +# be stripped before the config is handed to ``pyconfig.initialize``. +_PRIMUS_ONLY_PARAMS = ( + "file_sink_level", + "stderr_sink_level", + "sink_level", + "trainable", + "model", +) + + +def _resolve_maxtext_train(): + """Resolve MaxText's train entrypoints across MaxText versions. + + MaxText v26.4+ ships as the ``maxtext`` package with the training loop at + ``maxtext.trainers.pre_train.train``. MaxText v26.3 and earlier expose it + as ``MaxText.train``. Prefer the newer layout and fall back to the legacy + one so a single Primus checkout works against both images. + + Returns: + Tuple of ``(initialize, run, module_name)`` where ``module_name`` is the + importable module string to use as ``argv[0]`` for ``initialize``. + """ + try: + from maxtext.trainers.pre_train.train import initialize, run + + return initialize, run, "maxtext.trainers.pre_train.train" + except ImportError: + from MaxText.train import initialize, run + + return initialize, run, "MaxText.train" + class MaxTextPretrainTrainer(BaseTrainer): """ @@ -75,7 +108,7 @@ def init(self): """ log_rank_0("MaxTextPretrainTrainer.init() - initializing MaxText training") - from MaxText.train import initialize + initialize, _, module_name = _resolve_maxtext_train() from primus.backends.maxtext.argument_builder import ( export_params_to_yaml, @@ -86,9 +119,13 @@ def init(self): params_dict = namespace_to_dict(self.backend_args) params_dict.pop("override_model", None) + # Strip Primus-internal params that MaxText's config schema rejects. + for key in _PRIMUS_ONLY_PARAMS: + params_dict.pop(key, None) + yaml_path = export_params_to_yaml(params_dict) try: - argv = ["MaxText.train", yaml_path] + argv = [module_name, yaml_path] self.train_config, self.recorder, self.diagnostic_config = initialize(argv, **override_model_args) finally: try: @@ -173,7 +210,7 @@ def train(self): log_rank_0("Executing MaxText pretrain...") - from MaxText.train import run + _, run, _ = _resolve_maxtext_train() run(self.train_config, self.recorder, self.diagnostic_config) diff --git a/primus/backends/maxtext/patches/logger_patches.py b/primus/backends/maxtext/patches/logger_patches.py index 4c5cd2bf2..69a33ca02 100644 --- a/primus/backends/maxtext/patches/logger_patches.py +++ b/primus/backends/maxtext/patches/logger_patches.py @@ -18,6 +18,25 @@ from primus.core.utils.module_utils import error_rank_0, log_rank_0, warning_rank_0 +def _resolve_max_logging(): + """Resolve MaxText's ``max_logging`` module across MaxText versions. + + v26.4+ ships it at ``maxtext.utils.max_logging``; v26.3 and earlier expose + it as ``MaxText.max_logging``. Returns ``None`` if neither is importable. + """ + try: + from maxtext.utils import max_logging as maxtext_logging + + return maxtext_logging + except ImportError: + try: + import MaxText.max_logging as maxtext_logging + + return maxtext_logging + except ImportError: + return None + + @register_patch( patch_id="maxtext.logger", backend="maxtext", @@ -31,16 +50,14 @@ def patch_maxtext_logger(ctx: PatchContext) -> None: """ log_rank_0("[Patch:maxtext.logger] Patching MaxText logger...") - try: - import MaxText.max_logging as maxtext_logging - - if hasattr(maxtext_logging, "log"): - maxtext_logging.log = primus_logger.info - warning_rank_0("[Patch:maxtext.logger] MaxText logger patched successfully.") - else: - error_rank_0("[Patch:maxtext.logger] MaxText logging module does not have a 'log' function.") - except ImportError: + maxtext_logging = _resolve_max_logging() + if maxtext_logging is None: error_rank_0("[Patch:maxtext.logger] Failed to import MaxText's logging module.") + elif hasattr(maxtext_logging, "log"): + maxtext_logging.log = primus_logger.info + warning_rank_0("[Patch:maxtext.logger] MaxText logger patched successfully.") + else: + error_rank_0("[Patch:maxtext.logger] MaxText logging module does not have a 'log' function.") # Configure JAX logger level based on module config level_map = {"DEBUG": 10, "INFO": 20, "WARNING": 30, "ERROR": 40} diff --git a/primus/backends/maxtext/patches/train_patches.py b/primus/backends/maxtext/patches/train_patches.py index 39aa017fe..0e99db535 100644 --- a/primus/backends/maxtext/patches/train_patches.py +++ b/primus/backends/maxtext/patches/train_patches.py @@ -19,6 +19,25 @@ from primus.core.utils.module_utils import log_rank_0, warning_rank_0 +def _resolve_train_and_pyconfig(): + """Resolve MaxText's train module and pyconfig across MaxText versions. + + MaxText v26.4+ exposes the training loop at + ``maxtext.trainers.pre_train.train`` and config at ``maxtext.configs.pyconfig``. + MaxText v26.3 and earlier expose ``MaxText.train`` and ``MaxText.pyconfig``. + """ + try: + from maxtext.configs import pyconfig + from maxtext.trainers.pre_train import train as orig_train + + return orig_train, pyconfig + except ImportError: + import MaxText.train as orig_train + from MaxText import pyconfig + + return orig_train, pyconfig + + @register_patch( patch_id="maxtext.train", backend="maxtext", @@ -28,13 +47,19 @@ ) def patch_train(ctx: PatchContext) -> None: """ - Monkey-patch ``MaxText.train.initialize`` so that callers can pass + Monkey-patch MaxText's ``train.initialize`` so that callers can pass ``**kwargs`` which are transparently forwarded to ``pyconfig.initialize``. """ log_rank_0("[Patch:maxtext.train] Patching MaxText train module...") - import MaxText.train as orig_train - from MaxText import pyconfig + try: + orig_train, pyconfig = _resolve_train_and_pyconfig() + except ImportError as e: + warning_rank_0( + f"[Patch:maxtext.train] Could not locate MaxText train/pyconfig module; " + f"skipping override-forwarding patch: {e}" + ) + return _upstream_initialize = orig_train.initialize diff --git a/primus/configs/modules/maxtext/pre_trainer.yaml b/primus/configs/modules/maxtext/pre_trainer.yaml index b16135df7..973f58787 100644 --- a/primus/configs/modules/maxtext/pre_trainer.yaml +++ b/primus/configs/modules/maxtext/pre_trainer.yaml @@ -16,6 +16,10 @@ hf_path: "allenai/c4" # for using https://huggingface.co/datasets/allenai/c4 hf_data_dir: "en" hf_train_files: "" packing: true +# Required by MaxText v26.4's strict config validation: when hardware=gpu, +# packing=true, and attention=cudnn_flash_te, max_segments_per_seq must be > 0. +# 32 matches upstream MaxText GPU model configs and is harmless on v26.3. +max_segments_per_seq: 32 shardy: false diff --git a/third_party/maxtext b/third_party/maxtext index 851c0935a..80f431d0e 160000 --- a/third_party/maxtext +++ b/third_party/maxtext @@ -1 +1 @@ -Subproject commit 851c0935a4e9c9e4cd0771603e3bdf75a6578ad1 +Subproject commit 80f431d0ebbf6465257e2a5d931046dfb0b465b0 From 6bad69566bc64829181e63f209ef8259301b1f3f Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Thu, 16 Jul 2026 09:01:06 +0800 Subject: [PATCH 033/127] fix(megatron): ROCm-safe attention_backend + Mamba/SFT E2E, dead-code & CI cleanup (#870) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - **fix(megatron): coerce `attention_backend` enum args + make it ROCm-safe.** Primus builds Megatron's arg namespace straight from YAML/CLI, bypassing argparse type coercion, so enum args like `attention_backend` arrived as plain `str` and lost every downstream `== AttnBackend.member` check — `--attention_backend fused` was a silent no-op. `MegatronArgBuilder` now coerces enum args through Megatron's own argparse converter. A ROCm-gated `before_train` patch also lets the selected backend *win* over the image's baked `NVTE_*_ATTN` (ROCm images bake `NVTE_FLASH_ATTN=0`, which makes stock megatron's `auto`/explicit-backend assertion crash): `auto` respects the baked `FLASH=0`, explicit backends force their combination. - **test(megatron): add Mamba/Hybrid and native SFT+LoRA E2E cases.** `test_mamba_370M` / `test_zebra_llama_1B_hybrid` cover the MambaStack/HybridStack `auto` attention path; `test_qwen2_sft_lora` is the first posttrain (MegatronSFTTrainer) + `peft/*` E2E, using a tiny offline config/fixture. The conversion hook (`01_convert_checkpoints.py`) neutralizes `NVTE_*_ATTN` around `AutoBridge.import_ckpt()` since it runs against Bridge's bundled megatron that never sees the patch. - **chore(megatron): remove dead code** — unreachable Muon optimizer patch + vendored impl, and superseded `te_gemm_patch_wgrad` / `te_group_gemm_patch_wgrad`. - **ci(dependabot): monthly, major-only** GitHub Actions updates to cut PR churn (security updates unaffected). --------- Co-authored-by: Xiaoming-AMD --- .github/dependabot.yml | 12 +- .github/workflows/ci.yaml | 40 +- .../configs/MI300X/mamba_130M_pretrain.yaml | 54 ++ .../configs/MI355X/mamba_130M_pretrain.yaml | 54 ++ primus/backends/megatron/argument_builder.py | 68 +- .../core/extensions/te_gemm_patch_wgrad.py | 736 ------------------ .../extensions/te_group_gemm_patch_wgrad.py | 372 --------- .../core/optimizer/layer_wise_optimizer.py | 307 -------- .../backends/megatron/core/optimizer/moun.py | 353 --------- .../core/optimizer/moun_optimizer_config.py | 287 ------- .../megatron/megatron_base_trainer.py | 10 +- .../patches/attention_backend_patches.py | 87 +++ .../patches/muon_optimizer_patches.py | 109 --- .../models/megatron_bridge/mamba_130M.yaml | 2 + .../modules/megatron/trainer_base.yaml | 6 +- .../megatron/01_convert_checkpoints.py | 36 +- skills/backend-patch-explorer/SKILL.md | 2 +- tests/trainer/fixtures/sft_lora_smoke.jsonl | 8 + tests/trainer/test_megatron_trainer.py | 155 ++++ .../test_megatron_trainer_sft_lora.yaml | 103 +++ .../test_attention_backend_patches.py | 85 ++ .../test_megatron_argument_builder.py | 111 +++ .../megatron/test_megatron_base_trainer.py | 21 + .../megatron/test_muon_optimizer_patches.py | 214 ----- tests/unit_tests/ci/test_select_tests.py | 71 +- tools/ci/coverage_summary.py | 150 ++-- tools/ci/junit_summary.py | 34 +- tools/ci/runtime_summary.py | 116 ++- tools/ci/select_tests.py | 99 +-- 29 files changed, 1079 insertions(+), 2623 deletions(-) create mode 100644 examples/megatron_bridge/configs/MI300X/mamba_130M_pretrain.yaml create mode 100644 examples/megatron_bridge/configs/MI355X/mamba_130M_pretrain.yaml delete mode 100644 primus/backends/megatron/core/extensions/te_gemm_patch_wgrad.py delete mode 100644 primus/backends/megatron/core/extensions/te_group_gemm_patch_wgrad.py delete mode 100644 primus/backends/megatron/core/optimizer/layer_wise_optimizer.py delete mode 100644 primus/backends/megatron/core/optimizer/moun.py delete mode 100644 primus/backends/megatron/core/optimizer/moun_optimizer_config.py create mode 100644 primus/backends/megatron/patches/attention_backend_patches.py delete mode 100644 primus/backends/megatron/patches/muon_optimizer_patches.py create mode 100644 primus/configs/models/megatron_bridge/mamba_130M.yaml create mode 100644 tests/trainer/fixtures/sft_lora_smoke.jsonl create mode 100644 tests/trainer/test_megatron_trainer_sft_lora.yaml create mode 100644 tests/unit_tests/backends/megatron/test_attention_backend_patches.py delete mode 100644 tests/unit_tests/backends/megatron/test_muon_optimizer_patches.py diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 1bfc167d9..87b66dc14 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -1,16 +1,24 @@ version: 2 # Version updates for GitHub Actions only; CVEs are handled by Dependabot -# security updates (enabled repo-wide). +# security updates (enabled repo-wide), which are unaffected by the schedule +# and ignore rules below. updates: - package-ecosystem: "github-actions" directory: "/" schedule: - interval: "weekly" + interval: "monthly" open-pull-requests-limit: 5 groups: github-actions: patterns: - "*" + # Only open PRs for major version bumps; skip the churn of minor/patch + # version updates (security fixes still come through regardless). + ignore: + - dependency-name: "*" + update-types: + - "version-update:semver-minor" + - "version-update:semver-patch" labels: - "dependencies" - "github-actions" diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index c19a82543..bda09abab 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -222,6 +222,9 @@ jobs: echo "> build-docker success" run-unittest-torch: + permissions: + contents: read + actions: read # let "Write runtime summary" list this job's own steps for auto-discovered timing env: PRIMUS_WORKDIR: /mnt/apps_proxy/tas/0_public/primus_ci/actions-runner-torch # PRIMUS_WORKDIR: /wekafs/primus-data/primus_safe_ci/torch @@ -246,7 +249,6 @@ jobs: - run: echo "Begin AITER + Primus-Turbo Install." - name: Install AITER run: | - : > "$RUNNER_TEMP/runtime.tsv" # reset the CI runtime log for this job echo "✅ [Uninstall old aiter] started at: $(date)" pip3 uninstall aiter amd-aiter -y || true rm -rf /tmp/aiter || true @@ -263,7 +265,6 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [Build aiter] ended at: $(date)" echo "⏱️ [Build aiter] Total elapsed time: ${elapsed} seconds" - echo -e "Build aiter\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" - name: Install Primus-Turbo run: | rm -rf /tmp/Primus-Turbo || true @@ -279,7 +280,6 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [Pip install requirements] ended at: $(date)" echo "⏱️ [Pip install requirements] Total elapsed time: ${elapsed} seconds" - echo -e "primus-turbo: pip install requirements\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" start_time=$(date +%s) echo "✅ [build primus-turbo] started at: $(date)" pip3 install --no-build-isolation -e . -v @@ -287,7 +287,6 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [build primus-turbo] ended at: $(date)" echo "⏱️ [build primus-turbo] Total elapsed time: ${elapsed} seconds" - echo -e "primus-turbo: build/install\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" - run: echo "🎉 Begin Primus Unit Test." - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 with: @@ -344,21 +343,7 @@ jobs: # Note HSA_NO_SCRATCH_RECLAIM=1 must be set to avoid RCCL perf hit (TAS-8N Node), rocm ver:70125424 export HSA_NO_SCRATCH_RECLAIM=1 mkdir -p "${GITHUB_WORKSPACE}/test-reports" - # Component-aware selection on PRs; full suite on push/release/dispatch. - # Fail-safe: any failure to compute the diff falls back to the full suite. - TARGETS="./tests/unit_tests/" - if [[ "${{ github.event_name }}" == "pull_request" ]]; then - base_sha="${{ github.event.pull_request.base.sha }}" - git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true - changed="$(git diff --name-only "${base_sha}" HEAD 2>/dev/null || true)" - if [[ -n "${changed}" ]]; then - sel="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py)" - [[ -n "${sel}" ]] && TARGETS="${sel}" - fi - fi - echo "Selected unit-test targets: ${TARGETS}" - # shellcheck disable=SC2086 # intentional word-splitting of multiple paths - pytest --maxfail=1 -s ${TARGETS} \ + pytest --maxfail=1 -s ./tests/unit_tests/ \ --cov=primus --cov-report=term-missing:skip-covered \ --junitxml="${GITHUB_WORKSPACE}/test-reports/core-unit.xml" \ --deselect=tests/unit_tests/megatron/cco/test_tp_overlap.py::TPOverlapTestCase::test_fp8_te_linear \ @@ -398,7 +383,7 @@ jobs: git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true changed="$(git diff --name-only "${base}" HEAD 2>/dev/null || true)" if [[ -n "${changed}" ]]; then - e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py --e2e)" + e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py)" echo "Selected torch E2E scope: ${e2e:-}" if [[ "${e2e}" != "all" ]]; then echo "${e2e}" | grep -qw megatron || M=0 @@ -467,8 +452,10 @@ jobs: python tools/ci/junit_summary.py --title torch test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" || true - name: Write runtime summary if: always() + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: | - python tools/ci/runtime_summary.py --title torch "$RUNNER_TEMP/runtime.tsv" >> "$GITHUB_STEP_SUMMARY" || true + python tools/ci/runtime_summary.py --title torch >> "$GITHUB_STEP_SUMMARY" || true - name: Build torch coverage json (unit + E2E) if: always() continue-on-error: true @@ -502,6 +489,9 @@ jobs: SP=$(python -c "import site; print(site.getsitepackages()[0])" 2>/dev/null) && rm -f "$SP/primus_e2e_coverage.pth" || true run-unittest-jax: + permissions: + contents: read + actions: read # let "Write runtime summary" list this job's own steps for auto-discovered timing env: # PRIMUS_WORKDIR: /wekafs/primus-data/primus_safe_ci/jax PRIMUS_WORKDIR: /mnt/apps_proxy/tas/0_public/primus_docker_jax_ci/actions-runner @@ -534,7 +524,6 @@ jobs: echo "Primus-Turbo dir: /tmp/Primus-Turbo" git config --global --add safe.directory /tmp/Primus-Turbo cd /tmp/Primus-Turbo - : > "$RUNNER_TEMP/runtime.tsv" # reset the CI runtime log for this job start_time=$(date +%s) echo "✅ [Pip install requirements] started at: $(date)" mkdir -p ${PRIMUS_WORKDIR}/primus-cache @@ -543,7 +532,6 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [Pip install requirements] ended at: $(date)" echo "⏱️ [Pip install requirements] Total elapsed time: ${elapsed} seconds" - echo -e "pip install/upgrade\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" start_time=$(date +%s) echo "✅ [build primus-turbo] started at: $(date)" end_time=$(date +%s) @@ -607,7 +595,7 @@ jobs: git fetch --no-tags --depth=200 origin "${{ github.base_ref }}" 2>/dev/null || true changed="$(git diff --name-only "${base}" HEAD 2>/dev/null || true)" if [[ -n "${changed}" ]]; then - e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py --e2e)" + e2e="$(printf '%s\n' "${changed}" | python tools/ci/select_tests.py)" echo "Selected jax E2E scope: ${e2e:-}" if [[ "${e2e}" != "all" ]]; then echo "${e2e}" | grep -qw maxtext || X=0 @@ -644,8 +632,10 @@ jobs: python tools/ci/junit_summary.py --title jax test-reports/*.xml >> "$GITHUB_STEP_SUMMARY" || true - name: Write runtime summary if: always() + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: | - python tools/ci/runtime_summary.py --title jax "$RUNNER_TEMP/runtime.tsv" >> "$GITHUB_STEP_SUMMARY" || true + python tools/ci/runtime_summary.py --title jax >> "$GITHUB_STEP_SUMMARY" || true - name: Build jax coverage json (MaxText E2E) if: always() continue-on-error: true diff --git a/examples/megatron_bridge/configs/MI300X/mamba_130M_pretrain.yaml b/examples/megatron_bridge/configs/MI300X/mamba_130M_pretrain.yaml new file mode 100644 index 000000000..5ba988890 --- /dev/null +++ b/examples/megatron_bridge/configs/MI300X/mamba_130M_pretrain.yaml @@ -0,0 +1,54 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:mamba_130M_pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron_bridge + config: pretrain_trainer.yaml + + # Model to run + model: mamba_130M.yaml + + overrides: + stderr_sink_level: DEBUG + + # Recipe flavor: upstream Megatron-Bridge Mamba2 130M pretrain config + flavor: mamba2_130m_pretrain_config + + # --- Parameters accepted by _mamba2_common() --- + + # Training configuration + train_iters: 50 + global_batch_size: 32 + micro_batch_size: 4 + seq_length: ${PRIMUS_SEQ_LENGTH:2048} + + # Nested overrides (applied by _apply_nested_overrides) + log_interval: 1 + eval_interval: 500 + eval_iters: 0 + skip_save: true + + # Optimizer + lr: 3.0e-4 + min_lr: 3.0e-5 + lr_warmup_iters: 2 + lr_decay_iters: null + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: false + use_megatron_fsdp: false + enable_primus_turbo: false + + # Data + mock: true + # Real data path should point to a Megatron indexed dataset prefix: + # .bin + .idx + data_paths: ${PRIMUS_TOKENIZED_DATA_PATH:null} + train_data_path: null + valid_data_path: null + test_data_path: null diff --git a/examples/megatron_bridge/configs/MI355X/mamba_130M_pretrain.yaml b/examples/megatron_bridge/configs/MI355X/mamba_130M_pretrain.yaml new file mode 100644 index 000000000..5ba988890 --- /dev/null +++ b/examples/megatron_bridge/configs/MI355X/mamba_130M_pretrain.yaml @@ -0,0 +1,54 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:mamba_130M_pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + pre_trainer: + framework: megatron_bridge + config: pretrain_trainer.yaml + + # Model to run + model: mamba_130M.yaml + + overrides: + stderr_sink_level: DEBUG + + # Recipe flavor: upstream Megatron-Bridge Mamba2 130M pretrain config + flavor: mamba2_130m_pretrain_config + + # --- Parameters accepted by _mamba2_common() --- + + # Training configuration + train_iters: 50 + global_batch_size: 32 + micro_batch_size: 4 + seq_length: ${PRIMUS_SEQ_LENGTH:2048} + + # Nested overrides (applied by _apply_nested_overrides) + log_interval: 1 + eval_interval: 500 + eval_iters: 0 + skip_save: true + + # Optimizer + lr: 3.0e-4 + min_lr: 3.0e-5 + lr_warmup_iters: 2 + lr_decay_iters: null + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: false + use_megatron_fsdp: false + enable_primus_turbo: false + + # Data + mock: true + # Real data path should point to a Megatron indexed dataset prefix: + # .bin + .idx + data_paths: ${PRIMUS_TOKENIZED_DATA_PATH:null} + train_data_path: null + valid_data_path: null + test_data_path: null diff --git a/primus/backends/megatron/argument_builder.py b/primus/backends/megatron/argument_builder.py index 89b31e850..696e3c610 100644 --- a/primus/backends/megatron/argument_builder.py +++ b/primus/backends/megatron/argument_builder.py @@ -7,11 +7,13 @@ from __future__ import annotations import argparse +import enum from functools import lru_cache from types import SimpleNamespace -from typing import Any, Dict, Mapping, Union +from typing import Any, Callable, Dict, Mapping, Union from primus.core.utils.env import get_torchrun_env +from primus.core.utils.module_utils import warning_rank_0 # ------------------------------------------------------------ @@ -52,6 +54,48 @@ def _load_megatron_defaults() -> Dict[str, Any]: return vars(args).copy() # Convert Namespace → dict and cache +# ------------------------------------------------------------ +# Load Megatron's enum argparse type converters (cached) +# ------------------------------------------------------------ +@lru_cache(maxsize=1) +def _load_megatron_enum_types() -> Dict[str, Callable[[str], Any]]: + """Map each *enum-typed* Megatron arg ``dest`` to its argparse converter. + + Primus feeds config/CLI values straight into the namespace instead of + through Megatron's argparse, so an enum arg like ``attention_backend`` + arrives as a plain ``str`` and silently loses against every downstream + ``== SomeEnum.member`` comparison. We only touch enum args (the class of + bug this addresses); int/float/str args already arrive well-typed from + YAML. Reuse the parser's own converters so nothing is hand-maintained. + Returns ``{}`` if the parser cannot be built. + """ + try: + parser = _build_megatron_parser() + except Exception: # noqa: BLE001 - megatron may be unavailable (e.g. unit tests) + return {} + types: Dict[str, Callable[[str], Any]] = {} + for action in parser._actions: + convert = getattr(action, "type", None) + choices = getattr(action, "choices", None) + if callable(convert) and choices and all(isinstance(c, enum.Enum) for c in choices): + types[action.dest] = convert + return types + + +def _coerce_value(convert: Callable[[str], Any], value: Any) -> Any: + """Apply an argparse ``type`` converter to a raw override value. + + Mirrors what argparse does for a CLI string. Non-string values (already the + right type, e.g. from YAML) pass through untouched; list values are + converted element-wise to support ``nargs='+'`` args. + """ + if isinstance(value, str): + return convert(value) + if isinstance(value, list): + return [convert(v) if isinstance(v, str) else v for v in value] + return value + + # ------------------------------------------------------------ # MegatronArgBuilder: merge Primus → Megatron # ------------------------------------------------------------ @@ -93,9 +137,10 @@ def update(self, values: Union[Mapping[str, Any], SimpleNamespace]) -> "Megatron - None values are allowed and will override Megatron defaults. - Non-Megatron parameters are silently ignored. """ - # Get Megatron's supported parameters + # Get Megatron's supported parameters and enum argparse converters megatron_defaults = _load_megatron_defaults() megatron_keys = set(megatron_defaults.keys()) + enum_types = _load_megatron_enum_types() # Normalize input to a (key, value) iterable if isinstance(values, SimpleNamespace): @@ -105,11 +150,26 @@ def update(self, values: Union[Mapping[str, Any], SimpleNamespace]) -> "Megatron for key, value in items: # Only accept parameters that Megatron recognizes (including None overrides) - if key in megatron_keys: - self.overrides[key] = value + if key not in megatron_keys: + continue + + self.overrides[key] = self._coerce_enum(key, value, enum_types) return self + @staticmethod + def _coerce_enum(key: str, value: Any, enum_types: Mapping[str, Callable[[str], Any]]) -> Any: + """Coerce an enum override to its enum type (see _load_megatron_enum_types).""" + convert = enum_types.get(key) + if convert is None or value is None: + return value + + try: + return _coerce_value(convert, value) + except Exception as exc: # noqa: BLE001 - keep raw value, don't abort the build + warning_rank_0(f"[MegatronArgBuilder] could not coerce '{key}'={value!r}: {exc}") + return value + # ------------------------------------------------------------------ # Produce the final Megatron Namespace # ------------------------------------------------------------------ diff --git a/primus/backends/megatron/core/extensions/te_gemm_patch_wgrad.py b/primus/backends/megatron/core/extensions/te_gemm_patch_wgrad.py deleted file mode 100644 index a529aa7d2..000000000 --- a/primus/backends/megatron/core/extensions/te_gemm_patch_wgrad.py +++ /dev/null @@ -1,736 +0,0 @@ -# This file was modified for portability to AMDGPU -# Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Linear API""" -import functools -from functools import reduce -from operator import mul as multiply_op -from typing import Optional, Tuple, Union - -import torch -import transformer_engine_torch as tex -from transformer_engine.pytorch.constants import dist_group_type -from transformer_engine.pytorch.cpp_extensions import general_gemm -from transformer_engine.pytorch.cpu_offload import set_offloading_param -from transformer_engine.pytorch.distributed import ( - _fsdp_gather_tensors, - _fsdp_scatter_tensors, - allreduce, - gather_along_first_dim, - get_distributed_world_size, - in_fp8_activation_recompute_phase, - is_fp8_activation_recompute_enabled, - reduce_scatter_along_first_dim, -) -from transformer_engine.pytorch.fp8 import FP8GlobalStateManager -from transformer_engine.pytorch.graph import is_graph_capturing -from transformer_engine.pytorch.module._common import _fix_gathered_fp8_transpose -from transformer_engine.pytorch.module.base import ( - _2X_ACC_DGRAD, - _2X_ACC_FPROP, - _2X_ACC_WGRAD, - TransformerEngineBaseModule, - get_ub, - get_workspace, -) -from transformer_engine.pytorch.rocm_utils import ( - clear_fp8_weight_transpose_cache, - create_fp8_weight_transpose_cache, -) - -try: # TE >= 2.12: base classes moved to tensor.storage and renamed *Storage - from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import ( - MXFP8TensorStorage as MXFP8TensorBase, - ) -except ModuleNotFoundError: # TE <= 2.8 - from transformer_engine.pytorch.tensor._internal.mxfp8_tensor_base import ( - MXFP8TensorBase, - ) - -from transformer_engine.pytorch.utils import ( - assert_dim_for_fp8_exec, - cast_if_needed, - clear_tensor_data, - non_tn_fp8_gemm_supported, - nvtx_range_pop, - nvtx_range_push, - requires_grad, -) - -from primus.backends.megatron.core.pipeline_parallel.wgrad_adapter import ( - insert_wgrad_func_into_cache, -) - -try: - from transformer_engine.pytorch.tensor.quantized_tensor import ( - QuantizedTensor, - Quantizer, - prepare_for_saving, - restore_from_saved, - ) -except ModuleNotFoundError: - from transformer_engine.pytorch.quantized_tensor import ( - QuantizedTensor, - Quantizer, - prepare_for_saving, - restore_from_saved, - ) - - -class _LinearWithWGradSplit(torch.autograd.Function): - """Linear semi-top level module - Calls custom cuda extensions. - """ - - @staticmethod - def forward( - ctx, - weight: torch.Tensor, - inp: torch.Tensor, - bias: Optional[torch.Tensor], - is_first_microbatch: Union[bool, None], - fp8: bool, - fp8_calibration: bool, - input_quantizer: Optional[Quantizer], - weight_quantizer: Optional[Quantizer], - output_quantizer: Optional[Quantizer], - grad_output_quantizer: Optional[Quantizer], - grad_input_quantizer: Optional[Quantizer], - fuse_wgrad_accumulation: bool, - cpu_offloading: bool, - tp_group: Union[dist_group_type, None], - tp_size: int, - sequence_parallel: bool, - tensor_parallel: bool, - activation_dtype: torch.dtype, - parallel_mode: Union[str, None], - is_grad_enabled: bool, - ub_overlap_rs_fprop: bool, - ub_overlap_ag_dgrad: bool, - ub_overlap_ag_fprop: bool, - ub_overlap_rs_dgrad: bool, - ub_bulk_dgrad: bool, - ub_bulk_wgrad: bool, - ub_name: str, - fp8_output: bool, # pylint: disable=unused-argument - fsdp_group: Union[dist_group_type, None], - module: torch.nn.Module, - skip_fp8_weight_update: bool, - keep_fp8_weight_transpose_cache: bool, - ) -> torch.Tensor: - # pylint: disable=missing-function-docstring - assert not ub_bulk_wgrad, "not support for ZeroBubble" - assert bias is None, "not support bias yet" - - # NVTX label for profiling - nvtx_label = "transformer_engine._Linear.forward" - if ub_name is not None: - nvtx_label = f"{nvtx_label}.{ub_name}" - - # Make sure input dimensions are compatible - out_features, in_features = weight.shape - inp_shape = inp.shape - assert inp_shape[-1] == in_features, "GEMM not possible" - - tp_world_size = get_distributed_world_size(tp_group) - backward_needs_input = is_grad_enabled and weight.requires_grad - - # Prepare input tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - nvtx_range_push(f"{nvtx_label}.input_cast_comm") - inputmat = inp.view(-1, in_features) - inputmat_total = None - with_input_all_gather_nccl = ( - parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop - ) - own_quantized_input = False - if fp8: - assert_dim_for_fp8_exec(inputmat, weight) - if any([ub_overlap_ag_fprop, ub_overlap_rs_fprop]) and not ( - FP8GlobalStateManager.get_fp8_recipe().float8_per_tensor_scaling() - ): - raise NotImplementedError( - "Comm+GEMM overlap is only supported with FP8 delayed scaling or per-tensor" - " current scaling" - ) - - if input_quantizer is None: - raise ValueError("Missing quantizer for input tensor") - if with_input_all_gather_nccl: - assert not isinstance(inputmat, QuantizedTensor), "All gather of fp8 input is not supported" - input_quantizer.set_usage(rowwise=True, columnwise=False) - inputmat_total, _ = gather_along_first_dim( - inputmat, - tp_group, - quantizer=input_quantizer, - ) - else: - if FP8GlobalStateManager.get_fp8_recipe().float8_per_tensor_scaling() and ub_bulk_dgrad: - # reduce duplicated transpose in `_fix_gathered_fp8_transpose` - input_quantizer.set_usage(rowwise=True, columnwise=False) - else: - input_quantizer.set_usage( - rowwise=True, - columnwise=backward_needs_input, - ) - if not isinstance(inputmat, QuantizedTensor): - inputmat = input_quantizer(inputmat) - own_quantized_input = True - elif backward_needs_input: - inputmat.update_usage(rowwise_usage=True, columnwise_usage=True) - inputmat_total = inputmat - else: - inputmat = cast_if_needed(inp, activation_dtype) - if with_input_all_gather_nccl: - inputmat_total, _ = gather_along_first_dim(inputmat, tp_group) - else: - inputmat_total = inputmat - nvtx_range_pop(f"{nvtx_label}.input_cast_comm") - - # Cast weight to expected dtype - weightmat = weight - if not fp8: - weightmat = cast_if_needed(weightmat, activation_dtype) - else: - if not isinstance(weight, QuantizedTensor): - # Configure quantizer - if weight_quantizer is not None: - columnwise_usage = is_grad_enabled and inp.requires_grad - if not columnwise_usage: - columnwise_usage = ( - is_fp8_activation_recompute_enabled() and not in_fp8_activation_recompute_phase() - ) - weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - - # FP8 cast to workspace buffer - update_workspace = is_first_microbatch is None or is_first_microbatch - weightmat = module.get_weight_workspace( - tensor=weight, - quantizer=weight_quantizer, - cache_name=(None if is_first_microbatch is None else "weight"), - update_workspace=update_workspace, - skip_update_flag=skip_fp8_weight_update, - fsdp_group=fsdp_group, - create_transpose_cache=keep_fp8_weight_transpose_cache, - ) - - # Cast bias to expected dtype - bias_dtype = activation_dtype - if fp8 and activation_dtype == torch.float32: - bias_dtype = torch.bfloat16 - bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias - - # Configure output quantizer - if output_quantizer is not None: - output_quantizer.set_usage(rowwise=True, columnwise=False) - - # Calibrate quantizers if needed - if not fp8 and fp8_calibration: - if input_quantizer is not None: - input_quantizer.calibrate(inputmat_total) - if weight_quantizer is not None: - weight_quantizer.calibrate(weight) - - ub_obj = None - ub_type = None - rs_out = None - out_dtype = activation_dtype - if ub_overlap_rs_fprop: - ub_obj = get_ub(ub_name + "_fprop") - ub_type = tex.CommOverlapType.RS - out_shape = [reduce(multiply_op, inp_shape[:-1]) // tp_world_size, out_features] - rs_out = torch.empty(out_shape, dtype=activation_dtype, device=inputmat_total.device) - - elif ub_overlap_ag_fprop: - ub_obj = get_ub(ub_name + "_fprop") - ub_type = tex.CommOverlapType.AG - if fp8: - assert ub_obj.is_fp8_ubuf(), "AG overlap with FP8 GEMM inputs requires FP8 buffer." - ub_obj.copy_into_buffer(inputmat_total, input_quantizer, local_chunk=True) - inputmat_total = ub_obj.get_buffer(input_quantizer) - - nvtx_range_push(f"{nvtx_label}.gemm") - fprop_gemm_use_split_accumulator = _2X_ACC_FPROP - if fp8: - recipe = FP8GlobalStateManager.get_fp8_recipe() - if hasattr(recipe, "fp8_gemm_fprop"): - fprop_gemm_use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator - - out, *_, rs_out = general_gemm( - weightmat, - inputmat_total, - get_workspace(), - quantization_params=output_quantizer, - out_dtype=out_dtype, - bias=bias, - use_split_accumulator=fprop_gemm_use_split_accumulator, - ub=ub_obj, - ub_type=ub_type, - extra_output=rs_out, - ) - nvtx_range_pop(f"{nvtx_label}.gemm") - - if is_grad_enabled: - saved_inputmat = None - - ctx.backward_input_needs_gather = ( - weight.requires_grad and parallel_mode == "column" and sequence_parallel - ) - - if backward_needs_input: - if own_quantized_input and isinstance(inputmat, QuantizedTensor): - # For sequence parallel in vanilla FP8, rowwise data is - # to gather the input. For MXFP8, columnwise only data - # can be allgathered. - if isinstance(inputmat, MXFP8TensorBase) or not ctx.backward_input_needs_gather: - inputmat.update_usage(rowwise_usage=False) - saved_inputmat = inputmat - - # Weight with column-wise usage is needed for dgrad GEMM while keeping fp8 weight transpose cache. - if inp.requires_grad and keep_fp8_weight_transpose_cache: - if isinstance(weightmat, QuantizedTensor): - weightmat.update_usage(columnwise_usage=True) - - if cpu_offloading: - set_offloading_param(weight, "weight_offloading", True) - set_offloading_param(weightmat, "weight_offloading", True) - if saved_inputmat is not None: - set_offloading_param(saved_inputmat, "activation_offloading", True) - - # Scatter intermediate/activation tensors saved for the backward pass - # NOTE: FSDP sharding is not valid for models initialized with primary Fp8 weights - nvtx_range_push(f"{nvtx_label}.fsdp_scatter") - ctx.fsdp_group = fsdp_group - ctx.fsdp_shapes = _fsdp_scatter_tensors( - fsdp_group, - saved_inputmat, - weightmat if fp8 and not isinstance(weight, QuantizedTensor) else None, - ) - nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") - - if cpu_offloading: - ctx.grad_added_to_main_grad = hasattr(weight, "grad_added_to_main_grad") - - if ctx.grad_added_to_main_grad: - # If you are passing torch.nn.Parameter through the Torch hooks, you will - # get back torch.Tensor. Torch rips off the Parameter wrapper. - # You need to preserve the weight object to have all the attributes user - # sets for the weights. Because of this, it is not recommended to offload - # weights if weights are externally touched outside this module - ctx.weight_object = weight - - # TODO(ksivamani): Check memory usage - tensors_to_save, tensor_objects = prepare_for_saving( - saved_inputmat, - weightmat, - weight, - bias, - ) - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - - ctx.activation_dtype = activation_dtype - ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - ctx.fp8 = fp8 - ctx.input_quantizer = input_quantizer - ctx.grad_output_quantizer = grad_output_quantizer - ctx.grad_input_quantizer = grad_input_quantizer - ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation - if fuse_wgrad_accumulation and weight.requires_grad: - ctx.main_grad = weight.main_grad - - ctx.cpu_offloading = cpu_offloading - ctx.is_first_microbatch = is_first_microbatch - ctx.use_bias = bias is not None - ctx.sequence_parallel = sequence_parallel - ctx.tensor_parallel = tensor_parallel - ctx.inp_shape = inp_shape - ctx.parallel_mode = parallel_mode - ctx.tp_group = tp_group - ctx.ub_overlap_ag = ub_overlap_ag_dgrad - ctx.ub_overlap_rs_dgrad = ub_overlap_rs_dgrad - ctx.ub_bulk_dgrad = ub_bulk_dgrad - ctx.ub_bulk_wgrad = ub_bulk_wgrad - ctx.ub_name = ub_name - ctx.tp_size = tp_size - ctx.requires_dgrad = inp.requires_grad - ctx.requires_wgrad = weight.requires_grad - ctx.reduce_and_update_bwd_fp8_tensors = False - ctx.owns_input = saved_inputmat is not inp - ctx.keep_fp8_weight_transpose_cache = keep_fp8_weight_transpose_cache - if ctx.fp8 and requires_grad(inp, weight, bias): - _first_fp8_module = FP8GlobalStateManager.IS_FIRST_FP8_MODULE - ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase(): - FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module - - # Row Parallel Linear - if ub_overlap_rs_fprop: - out = rs_out - elif parallel_mode == "row": - nvtx_range_push(f"{nvtx_label}.row_parallel_comm") - if sequence_parallel: - out, _ = reduce_scatter_along_first_dim(out, tp_group) - elif tensor_parallel: - out, _ = allreduce(out, tp_group) - nvtx_range_pop(f"{nvtx_label}.row_parallel_comm") - - out = out.view(-1, *inp_shape[1:-1], out_features) - return out - - @staticmethod - def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: - # pylint: disable=missing-function-docstring - - # NVTX label for profiling - nvtx_label = "transformer_engine._Linear.backward" - if ctx.ub_name is not None: - nvtx_label = f"{nvtx_label}.{ctx.ub_name}" - - with torch.cuda.nvtx.range("_Linear_backward"): - if ( - ctx.fp8 - and any( - [ - ctx.ub_overlap_ag, - ctx.ub_overlap_rs_dgrad, - ctx.ub_bulk_dgrad, - ctx.ub_bulk_wgrad, - ] - ) - and (ctx.fp8_recipe is not None) - ): - if not ctx.fp8_recipe.float8_per_tensor_scaling(): - raise NotImplementedError( - "Comm+GEMM overlap is only supported with FP8 delayed scaling or per-tensor" - " current scaling" - ) - - saved_tensors = ctx.saved_tensors - inputmat, weight_fp8, weight, bias = ( # pylint: disable=unbalanced-tuple-unpacking - restore_from_saved(ctx.tensor_objects, saved_tensors) - ) - # Delete the references to tensor objects once they've been consumed - # by the `restore_from_saved` method to construct back the actual tensors. - ctx.tensor_objects = None - - # Since main_grad can be modified inplace, it should not be a part of saved_tensors - main_grad = ( - ctx.main_grad - if weight is not None and ctx.fuse_wgrad_accumulation and ctx.requires_wgrad - else None - ) - - if ctx.cpu_offloading: - if ctx.grad_added_to_main_grad: - weight = ctx.weight_object - if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation: - weight.main_grad = main_grad - - # Gather intermediate/activation tensors if needed - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already - # shards/unshards the base weights so we don't do it ourselves - nvtx_range_push(f"{nvtx_label}.fsdp_gather") - _fsdp_gather_tensors( - ctx.fsdp_group, - ctx.fsdp_shapes, - inputmat, - weight_fp8, - ) - nvtx_range_pop(f"{nvtx_label}.fsdp_gather") - - ctx.ub_obj_gradout = None - ub_obj_dgrad = None - ub_obj_wgrad = None - ub_type_dgrad = None - ub_type_wgrad = None - dgrad_shape = [reduce(multiply_op, ctx.inp_shape[:-1]), ctx.inp_shape[-1]] - rs_out = None - dgrad_bulk = None - if ctx.ub_overlap_ag: - # Overlap grad_output all-gather with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad") - ub_obj_dgrad = ctx.ub_obj_gradout - ub_type_dgrad = tex.CommOverlapType.AG - - elif ctx.ub_overlap_rs_dgrad: - # Overlap dgrad reduce-scatter with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad") - ub_obj_dgrad = ctx.ub_obj_gradout - ub_type_dgrad = tex.CommOverlapType.RS - rs_out = torch.empty(dgrad_shape, dtype=ctx.activation_dtype, device=grad_output.device) - - else: - if ctx.ub_bulk_dgrad: - # Overlap inputmat all-gather with dgrad compute - # NOTE: Copying into communication buffer will always prefer rowwise data, - # and will copy columnwise data if rowwise does not exist. In that case, - # the all-gather will apply to the leading dimension of the transpose, - # which then needs to be interleaved correctly before WGRAD. - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad") - ub_obj_dgrad = ctx.ub_obj_gradout - ub_type_dgrad = tex.CommOverlapType.AG - ub_obj_dgrad.copy_into_buffer(inputmat, ctx.input_quantizer, local_chunk=True) - - if ctx.ub_bulk_wgrad: - # Overlap dgrad reduce-scatter with wgrad compute - ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad") - ub_type_wgrad = tex.CommOverlapType.RS - ub_obj_wgrad.set_buffer_params(ctx.grad_input_quantizer) - dgrad_bulk = ub_obj_wgrad.get_buffer(ctx.grad_input_quantizer) - - # Prepare grad output tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - if ctx.grad_output_quantizer is not None: - # Reduce duplicated transpose, which is performed in grad_output.update_usage - if ctx.ub_overlap_ag and ctx.fp8_recipe.float8_per_tensor_scaling(): - ctx.grad_output_quantizer.set_usage(rowwise=True, columnwise=False) - else: - ctx.grad_output_quantizer.set_usage(rowwise=True, columnwise=True) - nvtx_range_push(f"{nvtx_label}.grad_output_preprocess") - ( - grad_output, - grad_bias, - ) = TransformerEngineBaseModule.grad_output_preprocess( - ctx, - grad_output, - ctx.parallel_mode == "row", - ctx.grad_output_quantizer, - ) - nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess") - - # Prepare input tensor - # Note: Perform tensor-parallel communication if needed - inputmat_total = None - inputmat_total_work = None - if ctx.backward_input_needs_gather and not ctx.ub_bulk_dgrad: - quantizer = None - if ctx.fp8: - quantizer = ctx.input_quantizer - quantizer.set_usage(rowwise=True, columnwise=True) - nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input") - inputmat_total, inputmat_total_work = gather_along_first_dim( - inputmat, - ctx.tp_group, - async_op=True, - quantizer=quantizer, - ) - nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input") - else: - inputmat_total = inputmat - - # Check whether to output wgrad GEMM directly into main grad - if ctx.is_first_microbatch is not None: - accumulate_wgrad_into_param_main_grad = ( - ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch - ) - else: - accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation - - # Compute grad input tensor - dgrad = None - dgrad_work = None - if ctx.requires_dgrad: - - # Update quantizer - if ctx.grad_input_quantizer is not None: - ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) - - if ctx.fp8 and not ctx.keep_fp8_weight_transpose_cache: - create_fp8_weight_transpose_cache(weight_fp8) - - # dgrad GEMM - nvtx_range_push(f"{nvtx_label}.dgrad_gemm") - dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_dgrad"): - dgrad_gemm_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator - - dgrad, *_, rs_out = general_gemm( - weight_fp8, - grad_output, - get_workspace(), - layout="NN", - grad=True, - quantization_params=ctx.grad_input_quantizer, - out=dgrad_bulk, - out_dtype=ctx.activation_dtype, - use_split_accumulator=dgrad_gemm_use_split_accumulator, - ub=ub_obj_dgrad, - ub_type=ub_type_dgrad, - extra_output=rs_out, - bulk_overlap=ctx.ub_bulk_dgrad, - ) - nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") - - if ctx.fp8 and not ctx.keep_fp8_weight_transpose_cache: - clear_fp8_weight_transpose_cache(weight_fp8) - - # Launch tensor-parallel communication - if ctx.ub_overlap_rs_dgrad: - dgrad = rs_out - elif ctx.parallel_mode == "column" and not ctx.ub_bulk_wgrad: - nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad") - if ctx.sequence_parallel: - dgrad, dgrad_work = reduce_scatter_along_first_dim( - dgrad, - ctx.tp_group, - async_op=True, - ) - else: - dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True) - nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad") - - # Compute grad weight tensor - wgrad = None - if ctx.requires_wgrad: - if ctx.ub_bulk_dgrad: - inputmat_total = ub_obj_dgrad.get_buffer(ctx.input_quantizer) - if ctx.fp8: - if inputmat._data is None: - # All-gather executed on columnwise data and result is in rowwise data, - # so we need to fix the interleaving before WGRAD. - inputmat_total = _fix_gathered_fp8_transpose(inputmat_total, ctx.tp_size) - elif not non_tn_fp8_gemm_supported(): - # FP8 GEMM on Hopper only supports TN layout so the gathered input must - # have a valid transpose. - inputmat_total._create_transpose() - - else: - if inputmat_total_work is not None: - # Synchronize tensor-parallel communication - inputmat_total_work.wait() - inputmat_total_work = None - - if isinstance(grad_output, QuantizedTensor): - # This is a no-op if platform supports non-TN FP8 GEMM or the transpose - # already exists. - grad_output.update_usage(rowwise_usage=True, columnwise_usage=True) - - if ctx.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf(): - rs_out = torch.empty(dgrad_shape, dtype=ctx.activation_dtype, device=grad_output.device) - - # wgrad GEMM - # Note: Fuse with bgrad computation if needed - def pre_process(_grad_output_, _input_, async_op=True): - return _grad_output_, _input_, None - - def process_wgrad(main_grad, grad_output, inputmat_total, handle=None): - nvtx_range_push(f"{nvtx_label}.wgrad_gemm") - - wgrad_gemm_use_split_accumulator = _2X_ACC_WGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_wgrad"): - wgrad_gemm_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator - # print(f"debug acc {accumulate_wgrad_into_param_main_grad}") - wgrad, grad_bias_, _, _ = general_gemm( - inputmat_total, - grad_output, - get_workspace(), - layout="NT", - grad=True, - out_dtype=(main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype), - bias=None, - out=main_grad if ctx.fuse_wgrad_accumulation else None, - use_split_accumulator=wgrad_gemm_use_split_accumulator, - accumulate=accumulate_wgrad_into_param_main_grad, - ub=ub_obj_wgrad, - ub_type=ub_type_wgrad, - extra_output=None, - bulk_overlap=False, - ) - - nvtx_range_pop(f"{nvtx_label}.wgrad_gemm") - - # Deallocate input tensor - if ctx.owns_input: - clear_tensor_data(inputmat_total) - # Handle custom DDP from mcore. - if ( - ctx.fuse_wgrad_accumulation - and weight is not None - and hasattr(weight, "grad_added_to_main_grad") - ): - weight.grad_added_to_main_grad = True - if getattr(weight, "zero_out_wgrad", False): - wgrad = torch.zeros( - weight.main_grad.shape, - dtype=weight.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - else: - wgrad = torch.empty( - weight.main_grad.shape, - dtype=weight.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - elif ctx.fuse_wgrad_accumulation: - pass - - insert_wgrad_func_into_cache( - main_grad, - functools.partial(pre_process, grad_output, inputmat_total), - functools.partial(process_wgrad, main_grad), - ) - - # Synchronize tensor parallel communication - if inputmat_total_work is not None: - assert False - inputmat_total_work.wait() - inputmat_total_work = None - if dgrad_work is not None: - assert False - dgrad_work.wait() - dgrad_work = None - - if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): - nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors") - FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) - nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors") - - # Scatter fp8 weight buffers - if ctx.fp8 and not isinstance(weight, QuantizedTensor): - _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8) - - wgrad = None - return ( - wgrad, - dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, - grad_bias, - None, # is_first_microbatch - None, # fp8 - None, # fp8_calibration - None, # input_quantizer - None, # weight_quantizer - None, # output_quantizer - None, # grad_output_quantizer - None, # grad_input_quantizer - None, # fuse_wgrad_accumulation - None, # cpu_offloading - None, # tp_group - None, # tp_size - None, # sequence_parallel - None, # tensor_parallel - None, # activation_dtype - None, # parallel_mode - None, # is_grad_enabled - None, # ub_overlap_rs_fprop - None, # ub_overlap_ag_dgrad - None, # ub_overlap_ag_fprop - None, # ub_overlap_rs_dgrad - None, # ub_bulk_dgrad - None, # ub_bulk_wgrad - None, # ub_name - None, # fp8_output - None, # fsdp_group - None, # module - None, # skip_fp8_weight_update - None, # keep_fp8_weight_transpose_cache - ) diff --git a/primus/backends/megatron/core/extensions/te_group_gemm_patch_wgrad.py b/primus/backends/megatron/core/extensions/te_group_gemm_patch_wgrad.py deleted file mode 100644 index d7b77ce31..000000000 --- a/primus/backends/megatron/core/extensions/te_group_gemm_patch_wgrad.py +++ /dev/null @@ -1,372 +0,0 @@ -# This file was modified for portability to AMDGPU -# Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -import functools -from typing import List, Tuple, Union - -import torch -import transformer_engine_torch as tex -from transformer_engine.pytorch.constants import TE_DType -from transformer_engine.pytorch.cpp_extensions import general_grouped_gemm -from transformer_engine.pytorch.distributed import ( - in_fp8_activation_recompute_phase, - is_fp8_activation_recompute_enabled, -) -from transformer_engine.pytorch.fp8 import FP8GlobalStateManager -from transformer_engine.pytorch.graph import is_graph_capturing -from transformer_engine.pytorch.module.base import ( - _2X_ACC_DGRAD, - _2X_ACC_FPROP, - _2X_ACC_WGRAD, - get_multi_stream_cublas_workspace, -) -from transformer_engine.pytorch.utils import ( - assert_dim_for_fp8_exec, - cast_if_needed, - clear_tensor_data, - requires_grad, -) - -from primus.backends.megatron.core.pipeline_parallel.wgrad_adapter import ( - insert_wgrad_func_into_cache, -) - -try: - from transformer_engine.pytorch.tensor.quantized_tensor import ( - QuantizedTensor, - Quantizer, - prepare_for_saving, - restore_from_saved, - ) -except ModuleNotFoundError: - from transformer_engine.pytorch.quantized_tensor import ( - QuantizedTensor, - Quantizer, - prepare_for_saving, - restore_from_saved, - ) - - -class _GroupedLinearWithWGradSplit(torch.autograd.Function): - """GroupedLinear semi-top level module - Calls custom cuda extensions. - """ - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - m_splits: List[int], - use_bias: bool, - is_first_microbatch: Union[bool, None], - fp8: bool, - fp8_calibration: bool, - input_quantizers: List[Quantizer], - weight_quantizers: List[Quantizer], - output_quantizers: List[Quantizer], - grad_output_quantizers: List[Quantizer], - fuse_wgrad_accumulation: bool, - cpu_offloading: bool, - sequence_parallel: bool, - activation_dtype: torch.dtype, - is_grad_enabled: bool, - module, - skip_fp8_weight_update, - *weights_and_biases, - ) -> torch.Tensor: - assert fuse_wgrad_accumulation, "fuse_wgrad_accumulation need to be true" - assert not use_bias, "zero bubble not support bias yet" - # pylint: disable=missing-function-docstring - num_gemms = len(m_splits) - weights = weights_and_biases[:num_gemms] - biases = weights_and_biases[num_gemms:] - device = inp.device - - # TODO Support MXFP8 # pylint: disable=fixme - if fp8 and FP8GlobalStateManager.get_fp8_recipe().mxfp8(): - raise NotImplementedError("GroupedLinear does not yet support MXFP8") - # TODO Support Float8 Current Scaling # pylint: disable=fixme - if fp8 and FP8GlobalStateManager.get_fp8_recipe().float8_current_scaling(): - raise NotImplementedError("GroupedLinear does not yet support Float8 Current Scaling") - - # Make sure input dimensions are compatible - in_features = weights[0].shape[-1] - assert inp.shape[-1] == in_features, "GEMM not possible" - inputmats = torch.split(inp.view(-1, in_features), m_splits) - if fp8: - assert_dim_for_fp8_exec(*inputmats, *weights) - - # Cast input to expected dtype - inputmats_no_fp8 = [cast_if_needed(mat, activation_dtype) for mat in inputmats] - inputmats = [] - - weight_requires_grad = weights[0].requires_grad - - if input_quantizers[0] is not None: - for input_quantizer in input_quantizers: - input_quantizer.set_usage( - rowwise=True, - columnwise=(is_grad_enabled and weight_requires_grad), - ) - columnwise_usage = is_grad_enabled and inp.requires_grad - if not columnwise_usage: - columnwise_usage = ( - is_fp8_activation_recompute_enabled() and not in_fp8_activation_recompute_phase() - ) - if weight_quantizers[0] is not None: - for weight_quantizer in weight_quantizers: - weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - if output_quantizers[0] is not None: - for output_quantizer in output_quantizers: - output_quantizer.set_usage(rowwise=True, columnwise=False) - - if fp8: - inputmats = tex.fused_multi_quantize( - inputmats_no_fp8, None, input_quantizers, TE_DType[activation_dtype] - ) - weights_fp8 = [] - bias_dtype = torch.bfloat16 if activation_dtype == torch.float32 else activation_dtype - if not isinstance(weights[0], QuantizedTensor): - # FP8 cast to workspace buffer - update_workspace = is_first_microbatch is None or is_first_microbatch - for i in range(num_gemms): - weight_fp8 = module.get_weight_workspace( - tensor=weights[i], - quantizer=weight_quantizers[i], - cache_name=(None if is_first_microbatch is None else f"weight{i}"), - update_workspace=update_workspace, - skip_update_flag=skip_fp8_weight_update, - ) - weights_fp8.append(weight_fp8) - else: - weights_fp8 = weights - - else: - inputmats = inputmats_no_fp8 - bias_dtype = activation_dtype - weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights] - - biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases - - out = torch.empty( - [sum(m_splits), weights_fp8[0].size(0)], - dtype=activation_dtype, - device=device, - ) - - _ = general_grouped_gemm( - weights_fp8, - inputmats, - [out], - activation_dtype, - get_multi_stream_cublas_workspace(), - single_output=True, - m_splits=m_splits, - bias=biases, - use_bias=use_bias, - use_split_accumulator=_2X_ACC_FPROP, - ) - - if fp8_calibration: - for i in range(num_gemms): - # amax of input - for i in range(num_gemms): - input_quantizers[i].calibrate(inputmats[i]) - for i in range(num_gemms): - weight_quantizers[i].calibrate(weights[i]) - - if is_grad_enabled: - - ctx.weights_shape_1 = weights[0].shape[1] - - tensors_to_save, tensor_objects = prepare_for_saving(*inputmats, *weights_fp8, *biases) - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - - ctx.weights_requires_grad = weights[0].requires_grad - if fuse_wgrad_accumulation and ctx.weights_requires_grad: - ctx.main_grads = [weights[i].main_grad for i in range(num_gemms)] - else: - ctx.main_grads = [None] * num_gemms - ctx.device = device - ctx.grad_output_quantizers = grad_output_quantizers - ctx.m_splits = m_splits - ctx.num_gemms = num_gemms - ctx.activation_dtype = activation_dtype - ctx.fp8 = fp8 - ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation - ctx.cpu_offloading = cpu_offloading - ctx.is_first_microbatch = is_first_microbatch - ctx.use_bias = use_bias - ctx.sequence_parallel = sequence_parallel - ctx.inp_shape = inp.shape - ctx.requires_dgrad = inp.requires_grad - ctx.reduce_and_update_bwd_fp8_tensors = False - if ctx.fp8 and requires_grad(inp, weights[0], biases[0]): - ctx.reduce_and_update_bwd_fp8_tensors = ( - ctx.reduce_and_update_bwd_fp8_tensors or FP8GlobalStateManager.is_first_fp8_module() - ) - - # [*, in_features] -> [*, out_features] except first dimension changes for SP - return out.view(-1, *inp.shape[1:-1], out.shape[-1]) - - @staticmethod - def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: - # pylint: disable=missing-function-docstring - with torch.cuda.nvtx.range("_GroupedLinear_backward"): - saved_tensors = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) - N = ctx.num_gemms - inputmats = saved_tensors[:N] - weights = saved_tensors[N : 2 * N] - biases = saved_tensors[2 * N : 3 * N] - main_grads = ctx.main_grads - - # preprocess grad_output - - grad_output = grad_output.contiguous() - grad_output_mats = torch.split(grad_output.view(-1, grad_output.shape[-1]), ctx.m_splits) - grad_output = [None] * ctx.num_gemms - grad_biases = [None] * ctx.num_gemms - if ctx.fp8: - if ctx.use_bias: - for i in range(ctx.num_gemms): - grad_biases[i], grad_output[i] = tex.bgrad_quantize( - grad_output_mats[i], ctx.grad_output_quantizers[i] - ) - else: - grad_output = tex.fused_multi_quantize( - grad_output_mats, - None, - ctx.grad_output_quantizers, - TE_DType[ctx.activation_dtype], - ) - else: - grad_output = grad_output_mats - - if ctx.is_first_microbatch is not None: - accumulate_wgrad_into_param_main_grad = ( - ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch - ) - else: - accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation - - if ctx.requires_dgrad: - dgrad = torch.empty( - (sum(ctx.m_splits), ctx.weights_shape_1), - dtype=ctx.activation_dtype, - device=ctx.device, - ) - - general_grouped_gemm( - weights, - grad_output, - [dgrad], - ctx.activation_dtype, - get_multi_stream_cublas_workspace(), - single_output=True, - layout="NN", - m_splits=ctx.m_splits, - grad=True, - use_split_accumulator=_2X_ACC_DGRAD, - ) - - if ctx.weights_requires_grad: - if ctx.fuse_wgrad_accumulation: - wgrad_list = main_grads - else: - wgrad_list = [ - torch.empty(w.size(), dtype=ctx.activation_dtype, device=ctx.device) for w in weights - ] - - def handle_custom_ddp_from_mcore(w, wgrad): - if ctx.weights_requires_grad: - if ctx.fuse_wgrad_accumulation and hasattr(w, "grad_added_to_main_grad"): - w.grad_added_to_main_grad = True - if getattr(w, "zero_out_wgrad", False): - wgrad = torch.zeros( - w.main_grad.shape, - dtype=w.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - else: - wgrad = torch.empty( - w.main_grad.shape, - dtype=w.dtype, - device=torch.cuda.current_device(), - requires_grad=False, - ) - elif ctx.fuse_wgrad_accumulation: - wgrad = None - else: - wgrad = None - return wgrad - - def pre_process(_grad_output_, _input_, async_op=True): - return _grad_output_, _input_, None - - kargs_dict = { - "out_dtype": ctx.activation_dtype, - "workspaces": get_multi_stream_cublas_workspace(), - "layout": "NT", - "grad": True, - "m_splits": ctx.m_splits, - "use_bias": ctx.use_bias if grad_biases[0] is None else None, - "bias": biases, - "use_split_accumulator": _2X_ACC_WGRAD, - "accumulate": accumulate_wgrad_into_param_main_grad, - } - - def process_wgrad(wgrad_list, kargs_dict, grad_output, inputmats, handle=None): - _, grad_biases_, _ = general_grouped_gemm( - inputmats, grad_output, wgrad_list, **kargs_dict - ) - for i in range(ctx.num_gemms): - if grad_biases[i] is None: - grad_biases[i] = grad_biases_[i] - del grad_biases_ - clear_tensor_data(*inputmats) - wgrad_list = [ - handle_custom_ddp_from_mcore(w, wgrad) for w, wgrad in zip(weights, wgrad_list) - ] - - insert_wgrad_func_into_cache( - wgrad_list, - functools.partial(pre_process, grad_output, inputmats), - functools.partial(process_wgrad, wgrad_list, kargs_dict), - ) - - wgrad_list_return = [None] * ctx.num_gemms - else: - wgrad_list_return = [None] * ctx.num_gemms - - if not ctx.use_bias: - grad_biases = [None] * ctx.num_gemms - - if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): - FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) - - return ( - dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, # is_grad_enabled - None, # is_grad_enabled - *wgrad_list_return, - *grad_biases, - ) diff --git a/primus/backends/megatron/core/optimizer/layer_wise_optimizer.py b/primus/backends/megatron/core/optimizer/layer_wise_optimizer.py deleted file mode 100644 index cc3ea94c8..000000000 --- a/primus/backends/megatron/core/optimizer/layer_wise_optimizer.py +++ /dev/null @@ -1,307 +0,0 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -import logging -from typing import Callable, List, Optional - -import torch -from megatron.core.dist_checkpointing.dict_utils import nested_values -from megatron.core.dist_checkpointing.mapping import ( - LocalNonpersistentObject, - ShardedStateDict, -) -from megatron.core.optimizer.clip_grads import count_zeros_fp32, get_grad_norm_fp32 -from megatron.core.optimizer.optimizer import ( - ChainedOptimizer, - Float16OptimizerWithFloat16Params, - FP32Optimizer, - MegatronOptimizer, -) -from megatron.core.optimizer.optimizer_config import OptimizerConfig -from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.utils import get_pg_rank, get_pg_size -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors - -logger = logging.getLogger(__name__) - - -class LayerWiseDistributedOptimizer(ChainedOptimizer): - """Layer-wise distributed optimizer for Megatron-core models. - - Experimental distributed optimizer wrapper that distributes weight to DP ranks by layer. - Implemented as ChainedOptimizer to support multiple optimizers (e.g. muon + adamW) - When using, keep all megatron distributed-optimizer related options OFF. - - How LayerWiseDistributedOptimizer work: - 1. weights are splited into lists and each rank only keep its shard in its optimizer - 2. Megatron DDP handle allreduce grad, note that each rank have full model and grad - 3. optimizer is already modified so only param belong to this DP rank is updated - 4. grad_norm and zero counting will reduce metrics globally in step function - 5. Do regular update with chained optimizers, modified optimizer only update shard - 6. allgather updated params to every rank - """ - - def __init__( - self, - optimizers: List[MegatronOptimizer], - config: OptimizerConfig, - pg_collection: Optional[ProcessGroupCollection] = None, - init_state_fn_list: Optional[List[Callable]] = None, - ) -> None: - """ - Initialize LayerWiseDistributedOptimizer. - - Args: - optimizers: List of MegatronOptimizers. - config: OptimizerConfig. - pg_collection: ProcessGroupCollection. - init_state_fn_list: List of init state functions. - """ - - self.pg_collection = pg_collection - self.shard_params(optimizers) - if init_state_fn_list: - assert len(init_state_fn_list) == len( - optimizers - ), "init_state_fn_list must be the same length as optimizers if provided" - - # wrap optimizer after sharding to avoid unnecessary master weight creation - # for higher precision, optimizers are wrapped with megatron already - if config.bf16: - # unwrap FP32 optimizer, possibly from reusing get_megatron_optimizer for adam - for i in range(len(optimizers)): - opt = optimizers[i] - if isinstance(opt, Float16OptimizerWithFloat16Params): - raise TypeError("LayerWiseDistributedOptimizer received Float16 optimizer already.") - # unwrap FP32 optimizer from reusing get_megatron_optimizer for adam - if isinstance(opt, FP32Optimizer): - opt = opt.optimizer - optimizers[i] = Float16OptimizerWithFloat16Params( - opt, config, None, init_state_fn_list[i] if init_state_fn_list else None - ) - - super().__init__(optimizers) - - # TODO(kunlun, deyuf): potential future perf optimization - # since allreduce is unchanged and handled by megatron DDP, they're already in - # contiguous gbuf. So instead of shard param by layer randomly, we can shard by - # buf range but keep some "extras" to keep boundary weight not sharded. - # This way each rank do some duplicated work but allgather_v is no longer needed - # All current distopt optimization can also be potentially applied - - def shard_params(self, optimizers): - """Shard all params into lists by rank.""" - # list of parameter are sorted by numel and assigned to ranks in ping-pong style - # example of 4 ranks and 10 parameters p0-p9 after sorting, then dp_cp_params_list will be - # [[p0, p7, p8], [p1, p6, p9], [p2, p5], [p3, p4]] - - # simplify when dp_cp group size is 1 - if get_pg_size(self.pg_collection.dp_cp) == 1: - self.dp_cp_params_list = None - self.expt_dp_params_list = None - return - - dp_cp_idx, expt_dp_idx = 0, 0 - dp_cp_size = get_pg_size(self.pg_collection.dp_cp) - expt_dp_size = get_pg_size(self.pg_collection.expt_dp) - # create ping-pong style loop so memory is more balanced - dp_cp_loop = list(range(dp_cp_size)) + list(range(dp_cp_size))[::-1] - expt_dp_loop = list(range(expt_dp_size)) + list(range(expt_dp_size))[::-1] - self.dp_cp_params_list = [[] for _ in range(dp_cp_size)] - self.expt_dp_params_list = [[] for _ in range(expt_dp_size)] - # get all param groups - param_groups = [] - for optimizer in optimizers: - param_groups += optimizer.param_groups - - # sort param in all groups by param numel and assign to each rank evenly - param_list = [] - for group_index, group in enumerate(param_groups): - for p in group["params"]: - param_list.append((p, group_index)) - param_list.sort(key=lambda x: x[0].numel()) - param_groups_this_rank = [[] for g in param_groups] - - # assign params to rank in ping-pong style loop - for p, group_index in param_list: - if param_groups[group_index].get("is_expert_parallel", False): - if expt_dp_loop[expt_dp_idx] == get_pg_rank(self.pg_collection.expt_dp): - param_groups_this_rank[group_index].append(p) - self.expt_dp_params_list[expt_dp_loop[expt_dp_idx]].append(p) - expt_dp_idx = (expt_dp_idx + 1) % len(expt_dp_loop) - else: - if dp_cp_loop[dp_cp_idx] == get_pg_rank(self.pg_collection.dp_cp): - param_groups_this_rank[group_index].append(p) - self.dp_cp_params_list[dp_cp_loop[dp_cp_idx]].append(p) - dp_cp_idx = (dp_cp_idx + 1) % len(dp_cp_loop) - - # now we modify the group to only handle local params - for groups, params in zip(param_groups, param_groups_this_rank): - groups["params"] = params - - # simplify when expt_dp group size is 1 or expert parallel is off - if expt_dp_size == 1 or len(self.expt_dp_params_list[0]) == 0: - self.expt_dp_params_list = None - - @torch.no_grad() - def allgather_params(self) -> None: - """All-gather updated params from all ranks.""" - - # helper function to flatten local params, allgather, unflatten and copy to model params - def _allgather_helper(params_list, group): - # flatten this rank's params and create empty tensor output list - device = params_list[0][0].device - dtype = params_list[0][0].dtype - rank = get_pg_rank(group) - # for rank without params create empty tensor and participate in allgather - src = ( - _flatten_dense_tensors(params_list[rank]) - if len(params_list[rank]) > 0 - else torch.empty(0, device=device, dtype=dtype) - ) - output_list = [ - torch.empty(sum([p.numel() for p in params]), device=device, dtype=dtype) - for params in params_list - ] - # single all_gather_v to collect all updated params - torch.distributed.all_gather(output_list, src, group=group) - # unflatten and copy gathered params for each rank i - for idx, (flat_params, params) in enumerate(zip(output_list, params_list)): - # skip local params and empty tensors - if len(params) == 0 or idx == rank: - continue - updated_params = _unflatten_dense_tensors(flat_params, params) - for updated_p, model_p in zip(updated_params, params): - model_p.data.copy_(updated_p) - - if self.pg_collection is None: - return - if self.dp_cp_params_list: - _allgather_helper(self.dp_cp_params_list, self.pg_collection.dp_cp) - if self.expt_dp_params_list: - _allgather_helper(self.expt_dp_params_list, self.pg_collection.expt_dp) - - @torch.no_grad() - def broadcast_params(self): - """All rank broadcast updated local params.""" - # Broadcast linear layer weights to all other ranks. Kept as reference test. - if self.dp_cp_params_list is None: - return - for i, params in enumerate(self.dp_cp_params_list): - src_global_rank = torch.distributed.get_global_rank(self.pg_collection.dp_cp, i) - for p in params: - torch.distributed.broadcast(p, src_global_rank, self.pg_collection.dp_cp) - if self.expt_dp_params_list is None: - return - for i, params in enumerate(self.expt_dp_params_list): - src_global_rank = torch.distributed.get_global_rank(self.pg_collection.expt_dp, i) - for p in params: - torch.distributed.broadcast(p, src_global_rank, self.pg_collection.expt_dp) - - @torch.no_grad() - def get_grad_norm(self): - # similar to dist opt, always aggregate globally - grads_for_norm = [] - for optimizer in self.chained_optimizers: - grads_for_norm += optimizer.get_main_grads_for_grad_norm() - grad_norm = get_grad_norm_fp32(grads_for_norm, grad_stats_parallel_group=None) - return grad_norm - - @torch.no_grad() - def count_zeros(self): - params = [] - for optimizer in self.chained_optimizers: - params += optimizer.get_parameters() - return count_zeros_fp32( - params, - grad_stats_parallel_group=None, - use_decoupled_grad=self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8, - ) - - @torch.no_grad() - def step(self): # type: ignore[no-untyped-def] - """step function for layer-wise optimizer.""" - update_successful, grad_norm, num_zeros_in_grad = super().step() - - # All gather updated params. - self.allgather_params() - - return update_successful, grad_norm, num_zeros_in_grad - - # TODO(deyuf): need to improve dist checkpointing design to properly handle this - # fp32_from_fp16_params is list, each sub list could be empty if group is empty - # this breaks dist checkpointing assumption since extract_sharded_base drop list structure - # for now, we convert it to dict with index as key and convert back in load_state_dict - def load_state_dict(self, state_dict): - if len(self.chained_optimizers) == 1: - wrapped_state_dict = {1: state_dict} - else: - wrapped_state_dict = state_dict - for sd in wrapped_state_dict.values(): - if "fp32_from_fp16_params" in sd and isinstance(sd["fp32_from_fp16_params"], dict): - logger.info("[layerwise] converting fp32_from_fp16_params from dict to list") - sd["fp32_from_fp16_params"] = [v for k, v in sorted(sd["fp32_from_fp16_params"].items())] - super().load_state_dict(state_dict) - - def sharded_state_dict( - self, model_sharded_state_dict: ShardedStateDict, is_loading: bool = False, **kwargs - ): - """ - Sharded state dict for torch_dist format checkpointing. - For fixed DP usage only, set replica_id to 0 for all ShardedTensor. - """ - sharded_state_dict = super().sharded_state_dict(model_sharded_state_dict, is_loading, **kwargs) - - # for fixed DP usage only - for sh_base in nested_values(sharded_state_dict): - if hasattr(sh_base, "replica_id"): - assert ( - isinstance(sh_base.replica_id, int) or len(sh_base.replica_id) == 3 - ), f"Expected replica_id as int or (PP, TP, DP), got: {sh_base}" - sh_base.replica_id = ( - 0 if isinstance(sh_base.replica_id, int) else (*sh_base.replica_id[:2], 0) - ) - - # later code assume list but chained optimizer fallback to non-list if there's only one - if len(self.chained_optimizers) == 1: - wrapped_sharded_state_dict = {1: sharded_state_dict} - else: - wrapped_sharded_state_dict = sharded_state_dict - - # Adjust dict rank 0 output correct global metadata into common_dict - for sd in wrapped_sharded_state_dict.values(): - # wrap empty containers into LocalNonpersistentObject so it won't be saved/loaded - # params is already wrapped, we only need to handle fp32_from_fp16_params and state - # more details in load_state_dict comment - if "fp32_from_fp16_params" in sd: - sd["fp32_from_fp16_params"][:] = [ - group if group else LocalNonpersistentObject(group) - for group in sd["fp32_from_fp16_params"] - ] - sd["fp32_from_fp16_params"] = {i: v for i, v in enumerate(sd["fp32_from_fp16_params"])} - # state is a single dict and will be empty if optimizer is fully empty - if not sd["optimizer"]["state"]: - sd["optimizer"]["state"] = LocalNonpersistentObject(sd["optimizer"]["state"]) - # group keys(e.g. 'step') might be missing or not updated - for i, group in enumerate(sd["optimizer"]["param_groups"]): - # keep local param tensor so we only gather metadata - local_params = group.pop("params") - # save whether this group is empty, so we can use non-empty rank for metadata - group["params"] = bool(local_params.unwrap()) - all_rank_groups = [None for _ in range(torch.distributed.get_world_size())] - torch.distributed.all_gather_object(all_rank_groups, group) - # find first non-empty group if it exists - nonempty_rank_group = next((g for g in all_rank_groups if g["params"]), group) - nonempty_rank_group["params"] = local_params - sd["optimizer"]["param_groups"][i] = nonempty_rank_group - return sharded_state_dict - - def save_state_dict_to_file(self, filename: str) -> None: - """Save the parameter state of the optimizer. For torch format only. - Args: - filename: The filename to save the parameter state. - """ - torch.save(super().state_dict(), filename) - - def load_state_dict_from_file(self, filename: str) -> None: - """Load the parameter state of the optimizer. For torch format only.""" - super().load_state_dict(torch.load(filename)) diff --git a/primus/backends/megatron/core/optimizer/moun.py b/primus/backends/megatron/core/optimizer/moun.py deleted file mode 100644 index 5464d6f00..000000000 --- a/primus/backends/megatron/core/optimizer/moun.py +++ /dev/null @@ -1,353 +0,0 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -"""Megatron muon optimizer wrapper to handle tensor-parallel.""" - -import logging -from typing import Any, Callable, Dict, List, Literal, Optional - -import torch -from megatron.core import parallel_state -from megatron.core.optimizer import _get_param_groups, get_megatron_optimizer -from megatron.core.optimizer.optimizer import ( - ChainedOptimizer, - Float16OptimizerWithFloat16Params, - FP32Optimizer, - MegatronOptimizer, -) -from megatron.core.optimizer.optimizer_config import OptimizerConfig - -try: - from megatron.core.optimizer.optimizer_config import ParamKey -except ImportError: - ParamKey = Any - -try: - from megatron.core.optimizer_param_scheduler import ParamGroupOverride -except ImportError: - ParamGroupOverride = Any -from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.transformer.module import MegatronModule -from megatron.core.utils import get_pg_size, log_single_rank -from torch.optim.optimizer import ParamsT - -from .layer_wise_optimizer import LayerWiseDistributedOptimizer - -try: - from emerging_optimizers.orthogonalized_optimizers import ( - OrthogonalizedOptimizer, - get_muon_scale_factor, - ) - from emerging_optimizers.orthogonalized_optimizers.muon_utils import ( - newton_schulz_tp, - ) - - HAVE_EMERGING_OPTIMIZERS = True - -except ImportError: - HAVE_EMERGING_OPTIMIZERS = False - OrthogonalizedOptimizer = object - - -logger = logging.getLogger(__name__) - - -class TensorParallelMuon(OrthogonalizedOptimizer): - """Tensor Parallel Muon optimizer.""" - - def __init__( - self, - params: ParamsT, - lr: float = 3e-4, - momentum_beta: float = 0.95, - use_nesterov: bool = True, - weight_decay: float = 0.01, - use_decoupled_weight_decay: bool = True, - split_qkv: bool = False, - is_qkv_fn: Callable[[torch.Tensor], bool] | None = None, - qkv_split_shapes: tuple[int, int, int] | None = None, - fp32_matmul_prec: str = "medium", - coefficient_type: str = "quintic", - num_ns_steps: int = 5, - scale_mode: str = "spectral", - extra_scale_factor: float = 1.0, - pg_collection: Optional[ProcessGroupCollection] = None, - mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated", - ) -> None: - if num_ns_steps < 1: - raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}") - - def scaled_orthogonalize_fn( - grad: torch.Tensor, - tp_group: torch.distributed.ProcessGroup, - partition_dim: int | None = None, - ) -> torch.Tensor: - # log_single_rank( - # logger, - # logging.DEBUG, - # f"Orthogonalizing grad with {num_ns_steps} steps, {coefficient_type} coefficient, " - # f"{scale_mode} scale mode, extra_scale_factor={extra_scale_factor}", - # ) - size = [grad.size(-2), grad.size(-1)] - if partition_dim: - size[partition_dim] *= get_pg_size(tp_group) - orth_grad = newton_schulz_tp( - grad, - steps=num_ns_steps, - coefficient_type=coefficient_type, - tp_group=tp_group, - partition_dim=partition_dim, - mode="duplicated" if mode == "blockwise" else mode, - ) - scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode) - return orth_grad * scale_factor * extra_scale_factor - - self.pg_collection = pg_collection - self.mode = mode - self.split_qkv = split_qkv - self.is_qkv_fn = is_qkv_fn - self.qkv_split_shapes = qkv_split_shapes - - weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2" - super().__init__( - params, - lr, - momentum_beta, - use_nesterov=use_nesterov, - weight_decay=weight_decay, - weight_decay_method=weight_decay_method, - fp32_matmul_prec=fp32_matmul_prec, - scaled_orthogonalize_fn=scaled_orthogonalize_fn, - ) - - def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor: - """Orthogonalize the momentum. - - Args: - p: The parameter tensor. i is necessary to pass param tensor in addition to momentum - because a lot of information is only available in the param tensor, - attributes for example. - grad: The momentum tensor. - - Returns: - The orthogonalized gradient tensor. - """ - # TODO(deyuf): switch to group - if self.pg_collection: - tp_group = self.pg_collection.expt_tp if getattr(p, "expert_tp", False) else self.pg_collection.tp - else: - tp_group = None - partition_dim = None if self.mode == "blockwise" else getattr(p, "partition_dim", None) - if partition_dim == -1: - # llm-shower use different default value for partition_dim than TE. - # Because -1 is a valid index for ndarray, we decided to not overload it. - partition_dim = None - - if self.split_qkv and self.is_qkv_fn(p): # type: ignore[misc] - # split grouped attention parameters (e.g., QKV, GQA, etc.) - grad_shape = grad.shape - log_single_rank( - logger, - logging.DEBUG, - f"qkv split grad shape {grad_shape}, split shapes {self.qkv_split_shapes}", - ) - num_query_groups = grad_shape[0] // sum(self.qkv_split_shapes) - qkv_grads = torch.split( - grad.view(num_query_groups, sum(self.qkv_split_shapes), -1), - self.qkv_split_shapes, - dim=1, - ) - qkv_grads = [g.reshape(-1, grad_shape[-1]) for g in qkv_grads] - - # Apply Newton-Schulz and scales to each component, concat back - qkv_grads = [ - self.scaled_orthogonalize_fn(g, tp_group, partition_dim).view( - num_query_groups, -1, grad_shape[-1] - ) - for g in qkv_grads - ] - grad = torch.cat(qkv_grads, dim=1).view(grad_shape) - else: - grad = self.scaled_orthogonalize_fn(grad, tp_group, partition_dim) - return grad - - -def get_megatron_muon_optimizer( - config: OptimizerConfig, - model_chunks: List[MegatronModule], - config_overrides: Optional[Dict[ParamKey, ParamGroupOverride]] = None, - use_gloo_process_groups: bool = True, - layer_wise_distributed_optimizer: bool = False, - pg_collection: Optional[ProcessGroupCollection] = None, - dump_param_to_param_group_map: Optional[str] = None, -) -> MegatronOptimizer: - """This function is used to get the muon optimizer for the model chunks. - It is used to get the muon optimizer for the model chunks. - - Args: - config (OptimizerConfig): optimizer configuration object. - model_chunks (List[MegatronModule]): model chunks to get optimizer for. - config_overrides (Optional[Dict[ParamKey, ParamGroupOverride]]): optional dictionary - of optimizer/scheduler overrides for parameter subsets. - use_gloo_process_groups (bool): if false, disable use of Gloo process groups - in underlying Megatron optimizers. - layer_wise_distributed_optimizer (bool): if true, use layer-wise distributed optimizer. - Defaults to False. - """ - assert HAVE_EMERGING_OPTIMIZERS, "Emerging Optimizers is not installed." - - # dist-optim is not supported due to strong coupling with how DDP init grad buffer - # in thoery we can put some weight to use non-dist-muon and rest to dist-adam - # but there are strong dependency and assumption in DDP that prevent it - if config.use_distributed_optimizer: - raise Exception("muon with dist optimizer is not supported.") - - # before this function receive properly created collection - if pg_collection is None: - pg_collection = ProcessGroupCollection.use_mpu_process_groups() - pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True) - pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group() - - log_single_rank(logger, logging.INFO, f"Setting up emerging optimizer with config {config}") - - optimizers = [] - # record list of non/linear params - linear_params = [] - nonlinear_params = [] - - for model_chunk in model_chunks: - # use config to determine qkv split shapes. - # no need to check tp since tp splits by head and this is per head(group) dimension - num_attention_heads = model_chunk.config.num_attention_heads - num_query_groups = model_chunk.config.num_query_groups - kv_channels = model_chunk.config.kv_channels - qkv_split_shapes = [ - num_attention_heads // num_query_groups * kv_channels, - kv_channels, - kv_channels, - ] - for name, param in model_chunk.named_parameters(): - if not param.requires_grad: - continue - # add flag for expert weight so optimizer can figure which tp group it uses - # alternatively, create new param group and save tp_group. this require more - # change in optimizer - if "experts" in name and "shared" not in name: - param.expert_tp = True - # add flag for qkv parameter - # TODO(deyuf): support MLA - if "linear_qkv.weight" in name and len(param.shape) == 2: - param.is_qkv = True - # TODO(deyuf): might not be sufficient for future algorithm. revisit this conditioning - if not getattr(param, "is_embedding_or_output_parameter", False) and not (len(param.shape) == 1): - linear_params.append(param) - else: - nonlinear_params.append(param) - - # freezing nonlinear params and get param groups for muon - for param in nonlinear_params: - param.requires_grad = False - - linear_param_groups = _get_param_groups( - model_chunks=model_chunks, - config=config, - config_overrides=config_overrides, - ) - - optimizer = TensorParallelMuon( - linear_param_groups, - lr=config.lr, - momentum_beta=config.muon_momentum, - use_nesterov=config.muon_use_nesterov, - weight_decay=config.weight_decay, - fp32_matmul_prec=config.muon_fp32_matmul_prec, - num_ns_steps=config.muon_num_ns_steps, - scale_mode=config.muon_scale_mode, - split_qkv=config.muon_split_qkv, - is_qkv_fn=lambda p: getattr(p, "is_qkv", False), - qkv_split_shapes=qkv_split_shapes, - extra_scale_factor=config.muon_extra_scale_factor, - pg_collection=pg_collection, - mode=config.muon_tp_mode, - ) - - # set config here to: - # 1. get adam for rest of layer - # 2. avoid ChainedOptimizer check fail that assert all optimizers are same kind - # side effect is muon optimizer will have wrong name str, i.e. config.optimizer == 'adam' - # TODO(deyuf): allow user to select optimizer mix and relax ChainedOptimizer design - config.optimizer = "adam" - - # Needed for torch_dist ckpt_format, unlike torch ckpt_format - # For other emerging optimizers, need to implement init_state_fn as well - # TODO(boxiangw): Improve usability after optimizer refactor - # TODO(boxiangw): support precision aware optimizer - def muon_init_state_fn(opt, config=None): - for group in opt.param_groups: - for p in group["params"]: - if len(opt.state[p]) == 0: - opt.state[p]["momentum_buffer"] = torch.zeros_like(p.data) - - def adam_init_state_fn(opt, config=None): - for group in opt.param_groups: - for p in group["params"]: - if len(opt.state[p]) == 0: - if config is None or not config.use_precision_aware_optimizer: - opt.state[p]["exp_avg"] = torch.zeros_like(p.data) - opt.state[p]["exp_avg_sq"] = torch.zeros_like(p.data) - else: - opt.initialize_state(p) - - # need to wrap into megatron mix precision optimizer. (only support bf16 w/o loss scale now) - if config.fp16: - raise Exception("muon with fp16 is not supported.") - - reset_config_bf16 = False - if config.bf16: - if layer_wise_distributed_optimizer: - # creating master weight before layerwise sharding will lead to unnecessary master - # weight so here we delay master weight creation into layer_wise unset config.bf16 - # will also result in all optimizers below(adam) to also not be wrapped - config.bf16 = False - reset_config_bf16 = True - else: - # if not using layer_wise wrapper, just create master weight here is fine - optimizer = Float16OptimizerWithFloat16Params(optimizer, config, None, muon_init_state_fn) - else: - optimizer = FP32Optimizer(optimizer, config, muon_init_state_fn) - - optimizers.append(optimizer) - - # done with muon, unfreeze nonlinear and freeze linear - for param in nonlinear_params: - param.requires_grad = True - for param in linear_params: - param.requires_grad = False - - # call original get. linear params will be skipped since they're freezed - chained_adam = get_megatron_optimizer( - config, - model_chunks, - config_overrides=config_overrides, - use_gloo_process_groups=use_gloo_process_groups, - pg_collection=pg_collection, - dump_param_to_param_group_map=dump_param_to_param_group_map, - ) - - # unfreeze everything - for param in linear_params: - param.requires_grad = True - - # chain everything together - optimizers += chained_adam.chained_optimizers - - if layer_wise_distributed_optimizer: - log_single_rank(logger, logging.INFO, "Using LayerWiseDistributedOptimizer for Muon") - if reset_config_bf16: - config.bf16 = True - return LayerWiseDistributedOptimizer( - optimizers, - config, - pg_collection, - init_state_fn_list=[muon_init_state_fn, adam_init_state_fn], - ) - return ChainedOptimizer(optimizers) diff --git a/primus/backends/megatron/core/optimizer/moun_optimizer_config.py b/primus/backends/megatron/core/optimizer/moun_optimizer_config.py deleted file mode 100644 index 101e997a9..000000000 --- a/primus/backends/megatron/core/optimizer/moun_optimizer_config.py +++ /dev/null @@ -1,287 +0,0 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - -from dataclasses import dataclass -from typing import Callable, Optional - -import torch -from megatron.core.utils import is_te_min_version - - -@dataclass -class MounOptimizerConfig: - """Configuration for optimizer.""" - - ############## - # General - ############## - optimizer: str = "adam" - """Optimizer to use (one of Adam, SGD, or Muon).""" - - lr: Optional[float] = None - """Initial learning rate. Depending on decay style and initial warmup, the learning rate at each - iteration would be different. - """ - - min_lr: Optional[float] = None - """Minumum value for learning rate. The scheduler clip values below this threshold.""" - - decoupled_lr: Optional[float] = None - """Separate learning rate for the input and output layer.""" - - decoupled_min_lr: Optional[float] = None - """Minimum value for learning rate for the input and output layer. The scheduler clip values - below this threshold. - """ - - weight_decay: float = 0.01 - """Weight decay coefficient for L2 regularization.""" - - ############## - # Precision - ############## - fp8_recipe: Optional[str] = None - """The type of fp8 recipe will affect the processing logic inside distributed optimizer.""" - - fp16: bool = False - """If true, train with fp16 mixed precision training. Defaults to False.""" - - bf16: bool = False - """If true, train with bf16 mixed precision training. Defaults to False.""" - - reuse_grad_buf_for_mxfp8_param_ag: bool = False - """If true, reuse the grad buffer for param AG when using mxfp8 recipe. Should be - set to True only when fp8_recipe is mxfp8 and fp8_param_gather is True.""" - - params_dtype: torch.dtype = torch.float32 - """dtype used when intializing the weights. Defaults to torch.float32.""" - - use_precision_aware_optimizer: bool = False - """If true, allows optimizer-related tensors (master_param, gradients and optimizer states) - to be set to lower precision. Defaults to False. - """ - - store_param_remainders: bool = True - """If true, store the 16-bit FP32 parameter remainders in the optimizer state, excluding the - 16 bits shared with the BF16 parameters. This lowers GPU memory usage. Defaults to True. - """ - - main_grads_dtype: torch.dtype = torch.float32 - """dtype of main grads when enabling precision-aware-optimizer""" - - main_params_dtype: torch.dtype = torch.float32 - """dtype of main params when enabling precision-aware-optimizer""" - - exp_avg_dtype: torch.dtype = torch.float32 - """dtype of exp_avg when enabling precision-aware-optimizer""" - - exp_avg_sq_dtype: torch.dtype = torch.float32 - """dtype of exp_avg_sq when enabling precision-aware-optimizer""" - - ############### - # Loss scaling - ############### - loss_scale: Optional[float] = None - """Static loss scaling, positive power of 2 values can improve fp16 convergence. If None, - dynamic loss scaling is used. - """ - - initial_loss_scale: float = 2**32 - """Initial loss-scale for dynamic loss scaling.""" - - min_loss_scale: float = 1.0 - """Minimum loss scale for dynamic loss scaling.""" - - loss_scale_window: float = 1000 - """Window over which to raise/lower dynamic scale.""" - - hysteresis: int = 2 - """Hysteresis for dynamic loss scaling.""" - - ############## - # Optimizer - ############## - # Adam - adam_beta1: float = 0.9 - """First coefficient for computing running averages of gradient and its square in Adam - optimizer. - """ - - adam_beta2: float = 0.999 - """Second coefficient for computing running averages of gradient and its square in Adam - optimizer. - """ - - adam_eps: float = 1e-08 - """Term added to the denominator to improve numerical stability in Adam optimizer.""" - - decoupled_weight_decay: bool = True - """If true, decouples weight decay from the gradient update, equivalent to AdamW. If false, - original Adam update rule will be used. Defaults to True. - """ - - # SGD. - sgd_momentum: float = 0.9 - """Momentum factor for SGD optimizer.""" - - # Muon - muon_momentum: float = 0.95 - """The momentum used by the internal SGD.""" - - muon_split_qkv: bool = True - """Whether to split QKV parameters for Muon optimizer.""" - - muon_use_nesterov: bool = False - """Whether to use Nesterov-style momentum in the internal SGD.""" - - muon_scale_mode: str = "spectral" - """The mode to use for the scale factor. Defaults to "spectral".""" - - muon_fp32_matmul_prec: str = "medium" - """The precision to use for the fp32 matmul. Defaults to "medium".""" - - muon_num_ns_steps: int = 5 - """The number of iteration steps to use in the Newton-Schulz iteration.""" - - muon_tp_mode: str = "blockwise" - """How to perform NS calculation for tensor parallel weights. Defaults to "blockwise".""" - - muon_extra_scale_factor: float = 1.0 - """Additional scale factor for the muon update.""" - - ####################### - # Distributed optimizer - ####################### - use_distributed_optimizer: bool = False - """Distribute optimizer state over data-parallel replicas.""" - - overlap_param_gather: bool = False - """If true, overlap param all-gather with forward compute. - This argument is intended to have the same value as the "overlap_param_gather" argument - in the "distributed_data_parallel_config.py" file. In the optimizer, this argument is - only used when "reuse_grad_buf_for_mxfp8_param_ag=True & fp8_param_gather=True". - """ - - overlap_param_gather_with_optimizer_step: bool = False - """If true, overlap param all-gather of first bucket with optimizer step.""" - - ####################### - # Optimizer Offload - ####################### - - optimizer_cpu_offload: bool = False - """If True, offload optimizer states tensor and compute to CPU.""" - - optimizer_offload_fraction: float = 0.0 - """Specifies the fraction of optimizer states to offload from GPU memory to CPU.""" - - use_torch_optimizer_for_cpu_offload: bool = False - """If True, use torch.optim.Optimizer for CPU offload.""" - - overlap_cpu_optimizer_d2h_h2d: bool = False - """ - When set to `True`, this flag enables overlapping of the CPU optimizer - update process with the data transfer operations. This can help improve - overall training efficiency by reducing idle time during data movement, - allowing the optimizer to perform updates while gradients and parameters - are being transferred between devices. - """ - - pin_cpu_grads: bool = True - """If True, pin the optimizer gradients to CPU memory.""" - - pin_cpu_params: bool = True - """If True, pin the optimizer parameters to CPU memory.""" - - ################ - # Miscellaneous - ################ - clip_grad: float = 1.0 - """Gradient clipping based on global L2 norm.""" - - log_num_zeros_in_grad: bool = False - """If true, calculate and log the number of zeros in gradient.""" - - barrier_with_L1_time: bool = False - """If true, use barrier with level 1 time measurements.""" - - timers: Optional[Callable] = None - """Function to get timers.""" - - config_logger_dir: str = "" - """When non-empty, dumps entry-point configs to config_logger_dir""" - - def __post_init__(self): - """Check the validity of the config.""" - - # The following condition is used to avoid repetition in distrib_optimizer.py. - # This is because in distrib_optimizer.py, the process to handle parameters are - # different for different training precision settings. FP8 cases require different - # handling while FP8 delayed scaling is an exception because the Adam optimizer in - # TransformerEngine supports it in the kernel computation. - # This is also the flag to determine the usage of param.grad or param.decoupled_grad - self.use_precision_aware_optimizer_no_fp8_or_ds_fp8 = self.use_precision_aware_optimizer and ( - self.main_params_dtype != torch.float32 - or (self.fp8_recipe is None or self.fp8_recipe == "delayed") - or self.optimizer_cpu_offload - ) - - if self.fp8_recipe == "mxfp8": - if not self.reuse_grad_buf_for_mxfp8_param_ag: - import warnings - - warnings.warn( - "mxfp8 without using reuse_grad_buf_for_mxfp8_param_ag and fp8_param_gather" - "will use significant amount additional GPU memory." - "Setting --reuse-grad-buf-for-mxfp8-param-ag and --fp8-param-gather is " - "recommended for mxfp8 training." - ) - - if self.use_precision_aware_optimizer: - assert self.optimizer == "adam", "--use-precision-aware-optimizer only supported with adam" - assert ( - self.use_distributed_optimizer - ), "--use-precision-aware-optimizer only supported with distributed optimizer" - - if not is_te_min_version("2.1.0"): - self.store_param_remainders = False - - # Only the FusedAdam in TE and HybridDeviceOptimizer supports - # --use-precision-aware-optimizer. - # TODO: Remove this check when apex's FusedAdam is no longer used. - if self.optimizer_cpu_offload: - return - try: - import inspect - - from transformer_engine.pytorch.optimizers import FusedAdam as Adam - - adam_args = inspect.signature(Adam).parameters - arg_names = [ - "master_weight_dtype", - "exp_avg_dtype", - "exp_avg_sq_dtype", - "use_decoupled_grad", - ] - for name in arg_names: - assert name in adam_args, ( - "Current FusedAdam of TE doesn't support --use-precision-aware-optimizer, " - "please update TE version." - ) - except ImportError: - raise RuntimeError( - "--use-precision-aware-optimizer requires FusedAdam from TransformerEngine, " - "but not found." - ) - else: - assert ( - self.main_grads_dtype == torch.float32 - ), "main_grads_dtype can only be fp32 when not using precision-aware optimizer" - assert ( - self.main_params_dtype == torch.float32 - ), "main_params_dtype can only be fp32 when not using precision-aware optimizer" - assert ( - self.exp_avg_dtype == torch.float32 - ), "exp_avg_dtype can only be fp32 when not using precision-aware optimizer" - assert ( - self.exp_avg_sq_dtype == torch.float32 - ), "exp_avg_sq_dtype can only be fp32 when not using precision-aware optimizer" diff --git a/primus/backends/megatron/megatron_base_trainer.py b/primus/backends/megatron/megatron_base_trainer.py index de0d17587..ef0a9b484 100644 --- a/primus/backends/megatron/megatron_base_trainer.py +++ b/primus/backends/megatron/megatron_base_trainer.py @@ -14,6 +14,7 @@ ) from primus.backends.megatron.training.mlflow_setup import upload_mlflow_artifacts from primus.core.trainer.base_trainer import BaseTrainer +from primus.core.utils.env import flush_before_hard_exit from primus.core.utils.module_utils import log_rank_0, warning_rank_0 @@ -127,14 +128,7 @@ def cleanup(self, on_error: bool = False): if exit_fast and not on_error: log_rank_0("[MegatronBaseTrainer] PRIMUS_EXIT_FAST=1 -> os._exit(0)") - # Flush stdout/stderr so the final log lines are not lost. - try: - import sys - - sys.stdout.flush() - sys.stderr.flush() - except Exception: # pragma: no cover - pass + flush_before_hard_exit() os._exit(0) def _finalize_mlflow_artifacts(self): diff --git a/primus/backends/megatron/patches/attention_backend_patches.py b/primus/backends/megatron/patches/attention_backend_patches.py new file mode 100644 index 000000000..3a806e0f9 --- /dev/null +++ b/primus/backends/megatron/patches/attention_backend_patches.py @@ -0,0 +1,87 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Make megatron-core's attention-backend selection ROCm-safe. + +``LanguageModule._set_attention_backend()`` reconciles the chosen +``attention_backend`` with the ``NVTE_*_ATTN`` env vars by *asserting* they are +unset-or-equal to what the backend wants, then setting them. Primus ROCm images +intentionally bake ``NVTE_FLASH_ATTN=0`` / ``NVTE_FUSED_ATTN=1`` (TE flash attn +is unavailable on ROCm; the fused/CK path is used instead), so stock megatron +assert-crashes before training starts on any model that goes through this probe +(Mamba, hybrid, non-Turbo GPT, ...): + +* ``auto`` wants all three = 1, but the baked ``NVTE_FLASH_ATTN=0`` trips it. +* an explicit backend (e.g. ``unfused``) wants ``NVTE_FUSED_ATTN=0``, but the + baked ``NVTE_FUSED_ATTN=1`` trips it. + +On ROCm the baked vars are image defaults, not user intent, so the selected +backend should *win* over them rather than assert against them: + +* ``auto`` -> enable every backend the platform hasn't explicitly disabled + (fill only the *unset* vars, so the baked ``NVTE_FLASH_ATTN=0`` is respected). +* an explicit backend -> force exactly its ``NVTE_*_ATTN`` combination, + overriding whatever the image baked in. +""" + +import os + +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_rank_0 + +_NVTE_ATTN_ENVS = ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN") + + +def _is_rocm(ctx: PatchContext) -> bool: + import torch + + return getattr(torch.version, "hip", None) is not None + + +@register_patch( + "megatron.attention_backend.rocm_safe", + backend="megatron", + phase="before_train", + description="Let attention_backend override ROCm's baked NVTE_*_ATTN instead of asserting", + condition=_is_rocm, +) +def patch_attention_backend(ctx: PatchContext): + from megatron.core.models.common.language_module.language_module import ( + LanguageModule, + ) + from megatron.core.transformer.enums import AttnBackend + + if getattr(LanguageModule, "_primus_attention_backend_patched", False): + return + + original_set_attention_backend = LanguageModule._set_attention_backend + + # (NVTE_FLASH_ATTN, NVTE_FUSED_ATTN, NVTE_UNFUSED_ATTN) for each explicit backend. + explicit_envs = { + AttnBackend.flash: ("1", "0", "0"), + AttnBackend.fused: ("0", "1", "0"), + AttnBackend.unfused: ("0", "0", "1"), + AttnBackend.local: ("0", "0", "0"), + } + + def _set_attention_backend(self): + backend = self.config.attention_backend + if backend == AttnBackend.auto: + for name in _NVTE_ATTN_ENVS: + os.environ.setdefault(name, "1") + return + values = explicit_envs.get(backend) + if values is None: + original_set_attention_backend(self) + return + for name, value in zip(_NVTE_ATTN_ENVS, values): + os.environ[name] = value + + LanguageModule._set_attention_backend = _set_attention_backend + LanguageModule._primus_attention_backend_patched = True + log_rank_0( + "[Patch:megatron.attention_backend.rocm_safe] attention_backend now overrides baked NVTE_*_ATTN" + ) diff --git a/primus/backends/megatron/patches/muon_optimizer_patches.py b/primus/backends/megatron/patches/muon_optimizer_patches.py deleted file mode 100644 index 7ecd2bc19..000000000 --- a/primus/backends/megatron/patches/muon_optimizer_patches.py +++ /dev/null @@ -1,109 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -""" -Megatron Muon Optimizer patches. - -This module patches megatron.training.training.get_megatron_optimizer to -automatically dispatch to get_megatron_muon_optimizer when args.optimizer -contains "muon". Since training.py uses `from megatron.core.optimizer import -get_megatron_optimizer`, we must patch the training module's namespace where -the function is actually used, not megatron.core.optimizer. -""" - -import dataclasses -import inspect - -from primus.core.patches import PatchContext, register_patch -from primus.core.utils.module_utils import log_rank_0 - - -@register_patch( - "megatron.optimizer.muon", - backend="megatron", - phase="before_train", - description="Patch get_megatron_optimizer to dispatch to muon optimizer when optimizer name contains 'muon'.", -) -def patch_get_megatron_optimizer_muon(ctx: PatchContext) -> None: - """ - Patch megatron.training.training.get_megatron_optimizer to delegate to - get_megatron_muon_optimizer when config.optimizer contains "muon". - - We patch the training module (not megatron.core.optimizer) because - training.py imports get_megatron_optimizer into its namespace at import - time; patching the optimizer module would not affect the training module's - local reference. - """ - try: - import megatron.training.training as training_module - except ImportError as e: - log_rank_0(f"[Patch:megatron.optimizer.muon] Skip patch (Megatron not available): {e}") - return - - original_get_megatron_optimizer = training_module.get_megatron_optimizer - original_signature = inspect.signature(original_get_megatron_optimizer) - - if getattr(original_get_megatron_optimizer, "_primus_muon_wrapper", False): - return - - def _get_bound_arg(bound_arguments, name, fallback=None): - if name in bound_arguments: - return bound_arguments[name] - - parameter = original_signature.parameters.get(name) - if parameter and parameter.default is not inspect.Parameter.empty: - return parameter.default - - return fallback - - def _patched_get_megatron_optimizer(*func_args, **func_kwargs): - config = func_kwargs.get("config") - if config is None and func_args: - config = func_args[0] - - optimizer_name = getattr(config, "optimizer", None) - if not optimizer_name or "muon" not in optimizer_name: - return original_get_megatron_optimizer(*func_args, **func_kwargs) - - bound_arguments = original_signature.bind_partial(*func_args, **func_kwargs).arguments - model_chunks = _get_bound_arg(bound_arguments, "model_chunks") - config_overrides = _get_bound_arg(bound_arguments, "config_overrides") - use_gloo_process_groups = _get_bound_arg(bound_arguments, "use_gloo_process_groups", True) - pg_collection = _get_bound_arg(bound_arguments, "pg_collection") - dump_param_to_param_group_map = _get_bound_arg(bound_arguments, "dump_param_to_param_group_map") - - from primus.backends.megatron.core.optimizer.moun import ( - get_megatron_muon_optimizer, - ) - from primus.backends.megatron.core.optimizer.moun_optimizer_config import ( - MounOptimizerConfig, - ) - - args = ctx.extra.get("backend_args", {}) - kwargs = {} - for f in dataclasses.fields(MounOptimizerConfig): - if hasattr(args, f.name): - kwargs[f.name] = getattr(args, f.name) - - moun_config = MounOptimizerConfig(**kwargs) - moun_config.timers = config.timers - - return get_megatron_muon_optimizer( - moun_config, - model_chunks, - config_overrides=config_overrides, - use_gloo_process_groups=use_gloo_process_groups, - layer_wise_distributed_optimizer="dist" in optimizer_name, - pg_collection=pg_collection, - dump_param_to_param_group_map=dump_param_to_param_group_map, - ) - - setattr(_patched_get_megatron_optimizer, "_primus_muon_wrapper", True) - training_module.get_megatron_optimizer = _patched_get_megatron_optimizer - log_rank_0( - "[Patch:megatron.optimizer.muon] Patched get_megatron_optimizer in megatron.training.training " - "to dispatch to muon when optimizer contains 'muon'." - ) diff --git a/primus/configs/models/megatron_bridge/mamba_130M.yaml b/primus/configs/models/megatron_bridge/mamba_130M.yaml new file mode 100644 index 000000000..463be8c69 --- /dev/null +++ b/primus/configs/models/megatron_bridge/mamba_130M.yaml @@ -0,0 +1,2 @@ +recipe: mamba.mamba2 +flavor: mamba2_130m_pretrain_config diff --git a/primus/configs/modules/megatron/trainer_base.yaml b/primus/configs/modules/megatron/trainer_base.yaml index 003e54eec..8e412f038 100755 --- a/primus/configs/modules/megatron/trainer_base.yaml +++ b/primus/configs/modules/megatron/trainer_base.yaml @@ -98,7 +98,9 @@ use_checkpoint_opt_param_scheduler: false warmup: null decoupled_lr: null decoupled_min_lr: null -# muon +# muon (read by Megatron-LM's native get_megatron_optimizer_config() whenever +# args has these attrs -- not consumed by any Primus code, but the only way +# to tune Muon's internals via Primus config since most have no CLI flag) muon_extra_scale_factor: 1.0 muon_scale_mode: "spectral" muon_fp32_matmul_prec: "medium" @@ -107,8 +109,6 @@ muon_tp_mode: "blockwise" muon_use_nesterov: false muon_split_qkv: true muon_momentum: 0.95 -muon_weight_decay: 0.01 -muon_weight_decay_method: "decoupled" optimizer_cpu_offload: false optimizer_offload_fraction: 1.0 # float diff --git a/runner/helpers/hooks/train/posttrain/megatron/01_convert_checkpoints.py b/runner/helpers/hooks/train/posttrain/megatron/01_convert_checkpoints.py index 88683f6a7..06f33fbee 100755 --- a/runner/helpers/hooks/train/posttrain/megatron/01_convert_checkpoints.py +++ b/runner/helpers/hooks/train/posttrain/megatron/01_convert_checkpoints.py @@ -126,6 +126,40 @@ def _prepend_sys_path(*paths: Path): sys.path[:] = original_sys_path +@contextmanager +def _unset_nvte_attention_env(): + """Neutralize the TE attention-backend env vars around AutoBridge model construction. + + This is a general posttrain fix, not a test-only workaround: this hook runs + for every native-Megatron SFT run, in a separate subprocess against + Megatron-Bridge's own bundled Megatron-LM that never sees Primus's + before_train patches (including the ROCm-safe attention_backend one). So it + hits *stock* megatron's ``auto`` validation: ``AutoBridge.import_ckpt`` builds + a plain ``MCoreGPTModel`` whose ``_set_attention_backend()`` asserts the three + ``NVTE_*_ATTN`` vars are unset-or-1, and the ROCm image's baked + ``NVTE_FLASH_ATTN=0`` trips it. + + Checkpoint conversion only reshapes weights -- it computes no attention -- so + the backend is irrelevant here; the only goal is to get past that assert. + Rather than counteract one specific baked value, unset all three for the + duration (stock ``auto`` then accepts them and picks defaults harmlessly) and + restore whatever was there afterwards. This stays correct for any image: if a + future image leaves them unset or sets them to 1, the pop/restore is a no-op; + it never assumes a particular baked value. Mirrors the Flux/diffusion conftest + fix. + """ + names = ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN") + saved = {name: os.environ.pop(name, None) for name in names} + try: + yield + finally: + for name, value in saved.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + def convert_checkpoint(hf_path: str, megatron_path: str): """ Convert HuggingFace checkpoint to Megatron torch_dist format. @@ -138,7 +172,7 @@ def convert_checkpoint(hf_path: str, megatron_path: str): log_info(f" Source: {hf_path}") log_info(f" Target: {megatron_path}") - with _prepend_sys_path(bridge_path, bridge_megatron_path): + with _prepend_sys_path(bridge_path, bridge_megatron_path), _unset_nvte_attention_env(): from megatron.bridge import AutoBridge # Convert using AutoBridge - creates torch_dist format checkpoint diff --git a/skills/backend-patch-explorer/SKILL.md b/skills/backend-patch-explorer/SKILL.md index 638cc68cf..16c97e662 100644 --- a/skills/backend-patch-explorer/SKILL.md +++ b/skills/backend-patch-explorer/SKILL.md @@ -92,7 +92,7 @@ State only what the current code confirms; never invent a flag, patch, or upstre For "how do I use feature X" requests (deepep is one such example), locate it live instead of recalling it: ```bash -rg -ni "" primus/backends primus/configs # e.g. deepep, fp8, zero_bubble, muon +rg -ni "" primus/backends primus/configs # e.g. deepep, fp8, zero_bubble # no rg: grep -rni "" primus/backends primus/configs --include=*.py ``` diff --git a/tests/trainer/fixtures/sft_lora_smoke.jsonl b/tests/trainer/fixtures/sft_lora_smoke.jsonl new file mode 100644 index 000000000..6773bd5dd --- /dev/null +++ b/tests/trainer/fixtures/sft_lora_smoke.jsonl @@ -0,0 +1,8 @@ +{"instruction": "What is the capital of France?", "input": "", "output": "The capital of France is Paris."} +{"instruction": "Translate the following sentence into French.", "input": "I love programming.", "output": "J'aime la programmation."} +{"instruction": "Summarize the given text in one sentence.", "input": "The quick brown fox jumps over the lazy dog. It is a common English pangram used to test typefaces.", "output": "A pangram about a fox jumping over a dog is used to test typefaces."} +{"instruction": "Give three tips for staying healthy.", "input": "", "output": "1. Eat a balanced diet. 2. Exercise regularly. 3. Get enough sleep."} +{"instruction": "Convert the temperature from Celsius to Fahrenheit.", "input": "100 degrees Celsius", "output": "100 degrees Celsius is 212 degrees Fahrenheit."} +{"instruction": "Write a short greeting for a new employee.", "input": "", "output": "Welcome to the team! We're excited to have you on board."} +{"instruction": "Explain what an API is in one sentence.", "input": "", "output": "An API is a set of rules that lets different software programs communicate with each other."} +{"instruction": "Rewrite the sentence to be more formal.", "input": "hey can u send me that file", "output": "Could you please send me that file?"} diff --git a/tests/trainer/test_megatron_trainer.py b/tests/trainer/test_megatron_trainer.py index 20152dc78..857cf2c5e 100644 --- a/tests/trainer/test_megatron_trainer.py +++ b/tests/trainer/test_megatron_trainer.py @@ -83,6 +83,50 @@ def run_script( return run_training_script(tag=tag, cmd=cmd, train_log_path=train_log_path, env=env) +def run_posttrain_script( + ut_name: str, + tag: str, + exp_path: str, + env_override: dict = None, + extra_args: list[str] = None, +): + """Like run_script, but for the "posttrain" suite (SFT/alignment). + + SFT experiment configs declare a `modules.post_trainer` section (instead + of `modules.pre_trainer`), which the CLI only loads via `train posttrain` + (see primus/cli/subcommands/train.py). The "Training completed." marker + that run_training_script asserts on is emitted generically by + PrimusRuntime._run_trainer_lifecycle for any module, so it applies here + unchanged. + """ + shell_entry = "./runner/primus-cli" + env = os.environ.copy() + if env_override: + env.update(env_override) + env["EXP"] = exp_path + + ut_log_path = os.environ.get("UT_LOG_PATH", "ut_out") + train_log_path = os.path.join(ut_log_path, f"log.test_megatron_trainer-{tag}.txt") + env["TRAIN_LOG"] = train_log_path + + cmd = [ + "bash", + shell_entry, + "direct", + "--log_file", + train_log_path, + "--", + "train", + "posttrain", + "--config", + exp_path, + ] + if extra_args: + cmd.extend(extra_args) + + return run_training_script(tag=tag, cmd=cmd, train_log_path=train_log_path, env=env) + + class TestMegatronTrainer(PrimusUT): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -562,6 +606,117 @@ def test_turbo_deepep(self): Dataloader_mp_context_patch_log in stdout ), "Expected dataloader_mp_context patch log not found in stdout" + def test_sdma_allgather_fused_residual_norm(self): + # Neither patch runs in any other case here: SDMA param all-gather + # only activates with ENABLE_SDMA_ALLGATHER=1 (env var, not a --arg) + + # a distributed optimizer; fused residual+RMSNorm needs + # use_turbo_rms_norm=1 + PRIMUS_FUSED_RESIDUAL_NORM_V2=1. Confirmed via + # coverage instrumentation that both patches install and execute. + run_script( + self.__class__.__name__, + "sdma_allgather_fused_residual_norm", + exp_path=f"examples/megatron/configs/{GPU_PLATFORM}/llama3_8B-BF16-pretrain.yaml", + env_override={ + "ENABLE_SDMA_ALLGATHER": "1", + "PRIMUS_FUSED_RESIDUAL_NORM_V2": "1", + }, + extra_args=[ + "--num_layers", + "4", + "--train_iters", + "3", + "--enable_primus_turbo", + "1", + "--use_turbo_attention", + "1", + "--use_turbo_rms_norm", + "1", + "--use_distributed_optimizer", + "1", + "--overlap_param_gather", + "1", + ], + ) + + def test_mamba_370M(self): + # Default `auto` attention backend exercises attention_backend_patches: + # megatron-core's probe must respect the image's baked NVTE_FLASH_ATTN=0. + run_script( + self.__class__.__name__, + "mamba_370M", + exp_path=f"examples/megatron/configs/{GPU_PLATFORM}/mamba_370M-pretrain.yaml", + env_override={}, + extra_args=[ + "--num_layers", + "4", + "--train_iters", + "3", + "--micro_batch_size", + "2", + "--global_batch_size", + "16", + ], + ) + + def test_zebra_llama_1B_hybrid(self): + # Hybrid Mamba+MLA (HybridStack) path. num_layers=8 is the minimum that + # keeps the default hybrid_attention_ratio=0.25 from allocating zero + # attention layers (division by zero). + run_script( + self.__class__.__name__, + "zebra_llama_1B_hybrid", + exp_path=f"examples/megatron/configs/{GPU_PLATFORM}/zebra_llama_1B-pretrain.yaml", + env_override={}, + extra_args=[ + "--num_layers", + "8", + "--train_iters", + "3", + "--micro_batch_size", + "2", + "--global_batch_size", + "16", + ], + ) + + def test_mamba_130M_bridge_pretrain(self): + # Only E2E covering the megatron_bridge backend (mamba/zebra above use + # the megatron backend). extra_args pin a tiny shape so the test doesn't + # depend on the example yaml's sizes. Don't override seq_length: the + # recipe feeds it to both model and dataset but a CLI override reaches + # only the dataset, and Bridge asserts the two match. + run_script( + self.__class__.__name__, + "mamba_130M_bridge_pretrain", + exp_path=f"examples/megatron_bridge/configs/{GPU_PLATFORM}/mamba_130M_pretrain.yaml", + env_override={}, + extra_args=[ + "--train_iters", + "3", + "--micro_batch_size", + "1", + "--global_batch_size", + "8", + ], + ) + + def test_qwen2_sft_lora(self): + # Only E2E covering the "posttrain" suite (MegatronSFTTrainer) and + # peft/*.py; LoRA is enabled in test_megatron_trainer_sft_lora.yaml. + # The posttrain hook HF->Megatron-converts `tokenizer_model` into the + # base checkpoint when pretrained_checkpoint/load are unset, so this + # uses a tiny stand-in checkpoint, not real Qwen2.5-7B weights. + run_posttrain_script( + self.__class__.__name__, + "qwen2_sft_lora", + exp_path="tests/trainer/test_megatron_trainer_sft_lora.yaml", + env_override={}, + # head_dim (2) is below the image's fused-attention CK kernel minimum, + # so pin the unfused path; attention_backend_patches forces it over + # the image's baked NVTE_FUSED_ATTN=1. + extra_args=["--attention_backend", "unfused"], + ) + def test_deepseekv2_lite_uep(self): run_script( self.__class__.__name__, diff --git a/tests/trainer/test_megatron_trainer_sft_lora.yaml b/tests/trainer/test_megatron_trainer_sft_lora.yaml new file mode 100644 index 000000000..3ad8dfc5a --- /dev/null +++ b/tests/trainer/test_megatron_trainer_sft_lora.yaml @@ -0,0 +1,103 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:test-megatron-sft-lora} +workspace: ./output + +modules: + post_trainer: + framework: megatron + config: sft_trainer.yaml + model: qwen2.5_7B.yaml + overrides: + # log + wandb_project: "Primus_test_SFT_LoRA" + disable_wandb: true + stderr_sink_level: DEBUG + + # required to select MegatronSFTTrainer + stage: sft + + # The posttrain hook (01_convert_checkpoints.py) always HF->Megatron- + # converts `tokenizer_model` into the base checkpoint when + # pretrained_checkpoint/load are unset. Point it at a real-but-tiny, + # non-gated Qwen2 checkpoint (2 layers, hidden_size 8, ~2.4M params) so + # the converted checkpoint matches the tiny model config below and the + # test exercises both the base-checkpoint load and the LoRA wrap. + tokenizer_type: HuggingFaceTokenizer + tokenizer_model: "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5" + num_layers: 2 + hidden_size: 8 + ffn_hidden_size: 32 + num_attention_heads: 4 + num_query_groups: 2 + + # 8 GPUs as pure data parallelism, so global_batch_size=8 is the smallest + # value dividing evenly (1 micro-batch/rank/iter). The 8-sample jsonl + # fixture covers one global batch; cyclic dataloader loops for the rest. + train_iters: 3 + micro_batch_size: 1 + global_batch_size: 8 + seq_length: 128 + max_position_embeddings: 128 + + eval_iters: 0 + eval_interval: 100 + + lr: 2.0e-4 + min_lr: 2.0e-5 + lr_warmup_iters: 1 + lr_decay_iters: null + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: false + + # parallel: no TP/PP -- this suite's default launch uses all 8 node + # GPUs as pure data parallelism (see global_batch_size note above). + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: false + + # Disable Turbo attention: the tiny head_dim (hidden_size/heads = 2) is far + # below what fused/flash kernels expect (test_qwen2_sft_lora pins the + # unfused backend). linear_qkv/proj/fc1/fc2 still use TE's non-turbo Linear + # classes, so the PEFT module_matcher path is exercised. + enable_primus_turbo: false + use_turbo_attention: false + + # SFT data: tiny local fixture (alpaca schema) so this test has no + # network dependency on the HuggingFace Hub. + sft_dataset_name: "tests/trainer/fixtures/sft_lora_smoke.jsonl" + sft_conversation_format: "alpaca" + enable_packed_sequences: false + + # Leave pretrained_checkpoint/load unset so the posttrain hook fills + # pretrained_checkpoint from the converted tokenizer_model checkpoint, + # driving MegatronSFTTrainer's pre-wrap base-checkpoint-load branch. + finetune: true + pretrained_checkpoint: null + load: null + save: null + save_interval: 20000 + no_save_optim: null + no_save_rng: null + disable_last_saving: true + ckpt_format: torch + + # LoRA: the activation switch for primus/backends/megatron/peft/*.py. + lora: + enabled: true + dim: 8 + alpha: 8 + dropout: 0.0 + dropout_position: pre + lora_A_init_method: xavier + lora_B_init_method: zero + target_modules: + - linear_qkv + - linear_proj + - linear_fc1 + - linear_fc2 diff --git a/tests/unit_tests/backends/megatron/test_attention_backend_patches.py b/tests/unit_tests/backends/megatron/test_attention_backend_patches.py new file mode 100644 index 000000000..9ff311f61 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_attention_backend_patches.py @@ -0,0 +1,85 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for the ROCm-safe attention_backend patch. + +These verify the patched ``LanguageModule._set_attention_backend`` reconciles the +selected backend with the image's baked ``NVTE_*_ATTN`` env vars the way ROCm +needs (override/respect rather than assert-crash), without needing a GPU. +""" + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("megatron") + +from megatron.core.models.common.language_module.language_module import LanguageModule +from megatron.core.transformer.enums import AttnBackend + +import primus.backends.megatron.patches.attention_backend_patches as patch_mod + +_NVTE = ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN") + + +@pytest.fixture +def patched_set_attention_backend(monkeypatch): + """Apply the patch onto a pristine LanguageModule and restore it afterwards.""" + monkeypatch.setattr(patch_mod, "log_rank_0", lambda *a, **k: None) + original = LanguageModule.__dict__.get("_set_attention_backend") + monkeypatch.delattr(LanguageModule, "_primus_attention_backend_patched", raising=False) + + patch_mod.patch_attention_backend(None) + try: + yield LanguageModule._set_attention_backend + finally: + if original is not None: + LanguageModule._set_attention_backend = original + if hasattr(LanguageModule, "_primus_attention_backend_patched"): + delattr(LanguageModule, "_primus_attention_backend_patched") + + +def _run(monkeypatch, backend, baked): + # Simulate the ROCm image's baked NVTE_*_ATTN before model construction. + for name in _NVTE: + monkeypatch.delenv(name, raising=False) + for name, value in baked.items(): + monkeypatch.setenv(name, value) + + dummy = SimpleNamespace(config=SimpleNamespace(attention_backend=backend)) + LanguageModule._set_attention_backend(dummy) + + import os + + return tuple(os.environ.get(name) for name in _NVTE) + + +# The Primus ROCm images bake NVTE_FLASH_ATTN=0 / NVTE_FUSED_ATTN=1. +_BAKED = {"NVTE_FLASH_ATTN": "0", "NVTE_FUSED_ATTN": "1"} + + +def test_auto_respects_baked_flash_off(patched_set_attention_backend, monkeypatch): + # "auto" must NOT force NVTE_FLASH_ATTN=1 (that is what crashes stock + # megatron); it fills only the unset var, leaving the baked FLASH=0. + assert _run(monkeypatch, AttnBackend.auto, _BAKED) == ("0", "1", "1") + + +def test_unfused_overrides_baked_fused(patched_set_attention_backend, monkeypatch): + # An explicit backend wins over the baked defaults (FUSED 1 -> 0), where + # stock megatron would assert-crash. + assert _run(monkeypatch, AttnBackend.unfused, _BAKED) == ("0", "0", "1") + + +def test_fused_sets_expected_combination(patched_set_attention_backend, monkeypatch): + assert _run(monkeypatch, AttnBackend.fused, _BAKED) == ("0", "1", "0") + + +def test_local_disables_all(patched_set_attention_backend, monkeypatch): + assert _run(monkeypatch, AttnBackend.local, _BAKED) == ("0", "0", "0") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/test_megatron_argument_builder.py b/tests/unit_tests/backends/megatron/test_megatron_argument_builder.py index 8b179d9f0..dcfd4a3ff 100644 --- a/tests/unit_tests/backends/megatron/test_megatron_argument_builder.py +++ b/tests/unit_tests/backends/megatron/test_megatron_argument_builder.py @@ -15,6 +15,7 @@ 5. Integration: complete workflow from config to final namespace """ +import enum from types import SimpleNamespace from unittest.mock import patch @@ -23,6 +24,18 @@ from primus.backends.megatron.argument_builder import MegatronArgBuilder +class _FakeAttnBackend(enum.Enum): + """Stand-in for Megatron's AttnBackend enum (plain int-valued enum).""" + + flash = 1 + fused = 2 + auto = 3 + + +def _fake_enum_type(value: str) -> "_FakeAttnBackend": + return _FakeAttnBackend[value] + + class TestMegatronArgBuilderFiltering: """Test parameter filtering: only Megatron params are accepted.""" @@ -328,5 +341,103 @@ def test_chained_updates(self, mock_load_defaults, mock_dist_info): assert result.hidden_size == 4096 +class TestMegatronArgBuilderEnumCoercion: + """Test that enum override strings are coerced via Megatron's argparse converter. + + Primus bypasses Megatron's argparse, so without coercion an enum arg like + ``attention_backend`` would arrive (and stay) a raw string and silently + break every downstream ``== AttnBackend.x`` comparison. + """ + + @patch("primus.backends.megatron.argument_builder._load_megatron_enum_types") + @patch("primus.backends.megatron.argument_builder._load_megatron_defaults") + def test_string_coerced_to_enum(self, mock_defaults, mock_types): + mock_defaults.return_value = {"attention_backend": _FakeAttnBackend.auto} + mock_types.return_value = {"attention_backend": _fake_enum_type} + + builder = MegatronArgBuilder() + builder.update({"attention_backend": "fused"}) + + assert builder.overrides["attention_backend"] is _FakeAttnBackend.fused + + @patch("primus.backends.megatron.argument_builder._load_megatron_enum_types") + @patch("primus.backends.megatron.argument_builder._load_megatron_defaults") + def test_auto_coerced_to_enum(self, mock_defaults, mock_types): + # Every enum value is coerced consistently, including "auto" (the ROCm + # incompatibility of megatron-core's auto branch is handled by a patch, + # not by leaving this a raw string). + mock_defaults.return_value = {"attention_backend": _FakeAttnBackend.auto} + mock_types.return_value = {"attention_backend": _fake_enum_type} + + builder = MegatronArgBuilder() + builder.update({"attention_backend": "auto"}) + + assert builder.overrides["attention_backend"] is _FakeAttnBackend.auto + + @patch("primus.backends.megatron.argument_builder._load_megatron_enum_types") + @patch("primus.backends.megatron.argument_builder._load_megatron_defaults") + def test_list_coerced_elementwise(self, mock_defaults, mock_types): + mock_defaults.return_value = {"some_enum_list": []} + mock_types.return_value = {"some_enum_list": _fake_enum_type} + + builder = MegatronArgBuilder() + builder.update({"some_enum_list": ["flash", "fused"]}) + + assert builder.overrides["some_enum_list"] == [ + _FakeAttnBackend.flash, + _FakeAttnBackend.fused, + ] + + @patch("primus.backends.megatron.argument_builder._load_megatron_enum_types") + @patch("primus.backends.megatron.argument_builder._load_megatron_defaults") + def test_non_enum_and_none_pass_through(self, mock_defaults, mock_types): + # num_layers is not an enum arg -> no converter -> value left untouched. + mock_defaults.return_value = {"attention_backend": _FakeAttnBackend.auto, "num_layers": 12} + mock_types.return_value = {"attention_backend": _fake_enum_type} + + builder = MegatronArgBuilder() + builder.update({"attention_backend": None, "num_layers": 32}) + + assert builder.overrides["attention_backend"] is None + assert builder.overrides["num_layers"] == 32 + + @patch("primus.backends.megatron.argument_builder._load_megatron_enum_types") + @patch("primus.backends.megatron.argument_builder._load_megatron_defaults") + def test_invalid_value_keeps_raw_without_raising(self, mock_defaults, mock_types): + mock_defaults.return_value = {"attention_backend": _FakeAttnBackend.auto} + mock_types.return_value = {"attention_backend": _fake_enum_type} + + builder = MegatronArgBuilder() + # "bogus" is not a valid enum member -> conversion fails, raw value kept. + builder.update({"attention_backend": "bogus"}) + + assert builder.overrides["attention_backend"] == "bogus" + + +class TestMegatronArgBuilderRealParser: + """Integration checks against the real Megatron argparse (no mocking). + + The coercion tests above mock the enum-type map; these guard the actual + detection wiring so a change in how Megatron declares attention_backend + (its argparse type/choices) can't silently turn coercion back into a no-op. + """ + + def test_attention_backend_detected_as_enum_arg(self): + from primus.backends.megatron.argument_builder import _load_megatron_enum_types + + pytest.importorskip("megatron") + assert "attention_backend" in _load_megatron_enum_types() + + @pytest.mark.parametrize("value", ["fused", "unfused", "auto"]) + def test_attention_backend_string_coerces_to_real_enum(self, value): + pytest.importorskip("megatron") + from megatron.core.transformer.enums import AttnBackend + + builder = MegatronArgBuilder() + builder.update({"attention_backend": value}) + + assert builder.overrides["attention_backend"] is AttnBackend[value] + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/test_megatron_base_trainer.py b/tests/unit_tests/backends/megatron/test_megatron_base_trainer.py index b1d097d2b..5dfeb5c48 100644 --- a/tests/unit_tests/backends/megatron/test_megatron_base_trainer.py +++ b/tests/unit_tests/backends/megatron/test_megatron_base_trainer.py @@ -10,6 +10,7 @@ Tests path resolution, parse_args patching, and setup orchestration. """ +import os import sys import types from pathlib import Path @@ -184,6 +185,26 @@ def test_ensure_megatron_path_already_importable(self, monkeypatch: pytest.Monke assert sys.path == path_before + def test_cleanup_exit_fast_flushes_before_hard_exit(self, monkeypatch: pytest.MonkeyPatch): + """Regression test: PRIMUS_EXIT_FAST=1's os._exit(0) bypasses atexit, so + flush_before_hard_exit() (stdout/stderr + coverage.save()) must run first. + A prior version inlined only the stdout/stderr half, silently dropping E2E + coverage data for that process.""" + trainer = _build_trainer(monkeypatch) + monkeypatch.setattr(trainer, "_finalize_mlflow_artifacts", lambda: None) + monkeypatch.setenv("PRIMUS_EXIT_FAST", "1") + + calls = [] + monkeypatch.setattr( + "primus.backends.megatron.megatron_base_trainer.flush_before_hard_exit", + lambda: calls.append("flush"), + ) + monkeypatch.setattr(os, "_exit", lambda code: calls.append(("exit", code))) + + trainer.cleanup(on_error=False) + + assert calls == ["flush", ("exit", 0)] + def test_patch_parse_args(self, monkeypatch: pytest.MonkeyPatch): """Test that parse_args is patched in both locations.""" trainer = _build_trainer(monkeypatch) diff --git a/tests/unit_tests/backends/megatron/test_muon_optimizer_patches.py b/tests/unit_tests/backends/megatron/test_muon_optimizer_patches.py deleted file mode 100644 index 2d90ae52f..000000000 --- a/tests/unit_tests/backends/megatron/test_muon_optimizer_patches.py +++ /dev/null @@ -1,214 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -import dataclasses -import sys -import types -from types import SimpleNamespace - -import pytest - -from primus.core.patches.context import PatchContext - - -def _install_fake_megatron_training(monkeypatch: pytest.MonkeyPatch): - megatron_mod = types.ModuleType("megatron") - megatron_mod.__path__ = [] - - training_pkg = types.ModuleType("megatron.training") - training_pkg.__path__ = [] - - training_mod = types.ModuleType("megatron.training.training") - original_calls = [] - - def original_get_megatron_optimizer( - config, - model_chunks, - config_overrides=None, - use_gloo_process_groups=True, - pg_collection=None, - dump_param_to_param_group_map=None, - ): - original_calls.append( - { - "config": config, - "model_chunks": model_chunks, - "config_overrides": config_overrides, - "use_gloo_process_groups": use_gloo_process_groups, - "pg_collection": pg_collection, - "dump_param_to_param_group_map": dump_param_to_param_group_map, - } - ) - return "original-result" - - training_mod.get_megatron_optimizer = original_get_megatron_optimizer - training_pkg.training = training_mod - megatron_mod.training = training_pkg - - monkeypatch.setitem(sys.modules, "megatron", megatron_mod) - monkeypatch.setitem(sys.modules, "megatron.training", training_pkg) - monkeypatch.setitem(sys.modules, "megatron.training.training", training_mod) - - return training_mod, original_calls - - -def _install_fake_muon_dependencies(monkeypatch: pytest.MonkeyPatch): - muon_calls = [] - - moun_mod = types.ModuleType("primus.backends.megatron.core.optimizer.moun") - - def fake_get_megatron_muon_optimizer( - config, - model_chunks, - config_overrides=None, - use_gloo_process_groups=True, - layer_wise_distributed_optimizer=False, - pg_collection=None, - dump_param_to_param_group_map=None, - ): - muon_calls.append( - { - "config": config, - "model_chunks": model_chunks, - "config_overrides": config_overrides, - "use_gloo_process_groups": use_gloo_process_groups, - "layer_wise_distributed_optimizer": layer_wise_distributed_optimizer, - "pg_collection": pg_collection, - "dump_param_to_param_group_map": dump_param_to_param_group_map, - } - ) - return "muon-result" - - moun_mod.get_megatron_muon_optimizer = fake_get_megatron_muon_optimizer - - moun_config_mod = types.ModuleType("primus.backends.megatron.core.optimizer.moun_optimizer_config") - - @dataclasses.dataclass - class FakeMounOptimizerConfig: - optimizer: str = "muon" - muon_tp_mode: str = "blockwise" - timers: object = None - - moun_config_mod.MounOptimizerConfig = FakeMounOptimizerConfig - - monkeypatch.setitem( - sys.modules, - "primus.backends.megatron.core.optimizer.moun", - moun_mod, - ) - monkeypatch.setitem( - sys.modules, - "primus.backends.megatron.core.optimizer.moun_optimizer_config", - moun_config_mod, - ) - - return muon_calls - - -def _call_get_megatron_optimizer( - optimizer_fn, - config, - model_chunks, - config_overrides, - use_gloo_process_groups, - pg_collection, - dump_param_to_param_group_map, - positional_config_overrides, -): - args = [config, model_chunks] - kwargs = { - "use_gloo_process_groups": use_gloo_process_groups, - "pg_collection": pg_collection, - "dump_param_to_param_group_map": dump_param_to_param_group_map, - } - if positional_config_overrides: - args.append(config_overrides) - else: - kwargs["config_overrides"] = config_overrides - return optimizer_fn(*args, **kwargs) - - -@pytest.mark.parametrize( - "positional_config_overrides", - [False, True], - ids=["keyword_config_overrides", "positional_config_overrides"], -) -def test_patch_get_megatron_optimizer_muon_matches_runtime_signature( - monkeypatch: pytest.MonkeyPatch, - positional_config_overrides: bool, -): - training_mod, original_calls = _install_fake_megatron_training(monkeypatch) - muon_calls = _install_fake_muon_dependencies(monkeypatch) - - monkeypatch.setattr( - "primus.backends.megatron.patches.muon_optimizer_patches.log_rank_0", - lambda *args, **kwargs: None, - ) - - from primus.backends.megatron.patches.muon_optimizer_patches import ( - patch_get_megatron_optimizer_muon, - ) - - original_fn = training_mod.get_megatron_optimizer - ctx = PatchContext( - backend="megatron", - phase="before_train", - extra={"backend_args": SimpleNamespace(muon_tp_mode="blockwise")}, - ) - - patch_get_megatron_optimizer_muon(ctx) - - assert training_mod.get_megatron_optimizer is not original_fn - - adam_config = SimpleNamespace(optimizer="adam", timers="adam-timer") - result = _call_get_megatron_optimizer( - training_mod.get_megatron_optimizer, - adam_config, - ["chunk-0"], - config_overrides={"dense": "group"}, - use_gloo_process_groups=False, - pg_collection="pg-0", - dump_param_to_param_group_map="dense-map", - positional_config_overrides=positional_config_overrides, - ) - - assert result == "original-result" - assert original_calls == [ - { - "config": adam_config, - "model_chunks": ["chunk-0"], - "config_overrides": {"dense": "group"}, - "use_gloo_process_groups": False, - "pg_collection": "pg-0", - "dump_param_to_param_group_map": "dense-map", - } - ] - - muon_config = SimpleNamespace(optimizer="muon-dist", timers="muon-timer") - result = _call_get_megatron_optimizer( - training_mod.get_megatron_optimizer, - muon_config, - ["chunk-1"], - config_overrides={"sparse": "group"}, - use_gloo_process_groups=False, - pg_collection="pg-1", - dump_param_to_param_group_map="muon-map", - positional_config_overrides=positional_config_overrides, - ) - - assert result == "muon-result" - assert len(muon_calls) == 1 - muon_call = muon_calls[0] - assert {key: value for key, value in muon_call.items() if key != "config"} == { - "model_chunks": ["chunk-1"], - "config_overrides": {"sparse": "group"}, - "use_gloo_process_groups": False, - "layer_wise_distributed_optimizer": True, - "pg_collection": "pg-1", - "dump_param_to_param_group_map": "muon-map", - } - assert muon_call["config"].timers == "muon-timer" - assert muon_call["config"].muon_tp_mode == "blockwise" diff --git a/tests/unit_tests/ci/test_select_tests.py b/tests/unit_tests/ci/test_select_tests.py index 91031635f..1baddec7f 100644 --- a/tests/unit_tests/ci/test_select_tests.py +++ b/tests/unit_tests/ci/test_select_tests.py @@ -4,7 +4,12 @@ # See LICENSE for license information. ############################################################################### -"""Unit tests for tools/ci/select_tests.py (classify-based PR test selection).""" +"""Unit tests for tools/ci/select_tests.py (classify-based E2E suite selection). + +Only E2E selection is covered: the unit-test suite is always run in full (see +select_tests.py's module docstring for why), so there's nothing to select +there anymore. +""" import importlib.util from pathlib import Path @@ -13,54 +18,7 @@ _SPEC = importlib.util.spec_from_file_location("select_tests", _ROOT / "tools/ci/select_tests.py") _MOD = importlib.util.module_from_spec(_SPEC) _SPEC.loader.exec_module(_MOD) -select = _MOD.select_targets - -FULL = ["tests/unit_tests/"] - - -# --- unit-test selection --------------------------------------------------- -def test_empty_runs_full(): - assert select([]) == FULL - - -def test_global_change_runs_full(): - assert select([".github/workflows/ci.yaml"]) == FULL - assert select(["tools/ci/select_tests.py"]) == FULL - assert select(["runner/helpers/x.sh"]) == FULL - assert select(["requirements.txt"]) == FULL - assert select(["primus/core/launcher/initialize.py"]) == FULL - - -def test_non_py_under_primus_runs_full(): - # configs / fixtures can't be localized to a unit dir -> fail-safe. - assert select(["primus/configs/x.yaml"]) == FULL - - -def test_backend_maps_to_its_unit_dir(): - out = select(["primus/backends/megatron/training/global_vars.py"]) - assert "tests/unit_tests/backends/megatron/" in out - assert "tests/unit_tests/megatron/" in out # megatron's extra GPU-operator tests - -def test_backend_without_unit_dir_runs_full(): - # transformer_engine has no tests/unit_tests/backends/transformer_engine/. - assert select(["primus/backends/transformer_engine/x.py"]) == FULL - - -def test_component_maps_to_isomorphic_dir(): - assert select(["primus/core/projection/engine.py"]) == ["tests/unit_tests/core/projection/"] - assert select(["primus/agents/a.py"]) == ["tests/unit_tests/agents/"] - - -def test_changed_unit_test_runs_its_dir(): - assert select(["tests/unit_tests/agents/test_tools.py"]) == ["tests/unit_tests/agents/"] - - -def test_docs_only_runs_full(): - assert select(["README.md", "docs/guide.md"]) == FULL - - -# --- E2E selection (pass a fixed suite set for determinism) ---------------- SUITES = {"megatron", "torchtitan", "maxtext"} @@ -84,14 +42,29 @@ def test_e2e_backend_with_trainer_runs_its_suite(): def test_e2e_backend_without_trainer_runs_all(): - # No trainer suite (bridge / hummingbirdxt / transformer_engine) -> fail-safe all. + # No trainer suite (bridge / hummingbirdxt / transformer_engine / diffusion) -> fail-safe all. assert set(e2e(["primus/backends/megatron_bridge/x.py"])) == SUITES assert set(e2e(["primus/backends/hummingbirdxt/x.py"])) == SUITES assert set(e2e(["primus/backends/transformer_engine/x.py"])) == SUITES + assert set(e2e(["primus/backends/diffusion/x.py"])) == SUITES def test_e2e_component_change_runs_all(): + # Any other primus/ or tests/unit_tests/ change -- not just the + # explicitly-listed GLOBAL_TRIGGERS -- also runs everything: classify() + # maps it to "component", which select_e2e() treats the same as "global". assert set(e2e(["primus/core/trainer/base.py"])) == SUITES + assert set(e2e(["primus/core/launcher/parser.py"])) == SUITES + assert set(e2e(["tests/unit_tests/core/patches/test_patch.py"])) == SUITES + + +def test_e2e_bare_examples_file_runs_all(): + # A bare examples/ (no backend subdir) is shared launcher plumbing + # -- e.g. test_maxtext_trainer.py shells out to examples/run_pretrain.sh + # directly -- so it must not be silently ignored like a docs-only change. + assert set(e2e(["examples/run_pretrain.sh"])) == SUITES + # examples//... is unaffected: still maps to that one backend. + assert e2e(["examples/maxtext/configs/x.yaml"]) == ["maxtext"] def test_e2e_docs_only_runs_none(): diff --git a/tools/ci/coverage_summary.py b/tools/ci/coverage_summary.py index a7477ace8..400136a74 100644 --- a/tools/ci/coverage_summary.py +++ b/tools/ci/coverage_summary.py @@ -9,51 +9,54 @@ Modes (by number of report arguments): 1 report -> single "Coverage" column (e.g. JAX MaxText E2E). 2+ reports -> "Unit" vs "Unit+E2E". The 1st is unit; the rest are E2E reports, - merged per module by taking the max covered lines. Each E2E - report should be a `coverage combine` of unit + that job's E2E - data (line-level). Taking the max across jobs avoids double - counting and lets torch (megatron/torchtitan) and jax (maxtext) - E2E - which cover near-disjoint modules - share one table. - -Layout: top-level groups (core, backends, modules, agents, cli, ...) are bold -rows at the same level, sorted by coverage. core/ and backends/ also get -indented sub-rows per area, sorted by coverage. __init__.py is dropped; -tools/, platforms/ and the top-level pretrain.py entrypoint are excluded (ops -tooling / env abstraction / thin CLI glue, exercised by E2E and shell tests -rather than unit tests). runner/ is bash, covered by the tests/runner/ shell -tests. - -Single-report mode reflects what a partial run (e.g. MaxText E2E) actually -executed: modules with zero covered lines are hidden and the total is computed -over the executed modules only, so the headline number is meaningful instead of -diluted by code that run can never touch. The two-report comparison keeps every -module (unit gives the full denominator). + merged per module by max covered lines (avoids double-counting + torch's near-disjoint megatron/torchtitan E2E). + +Layout: top-level groups are bold rows sorted by coverage; core/backends also +get indented per-area detail rows. The headline percentage gets a tier emoji +(_TIER_THRESHOLDS) and a curated "Notes" column flags widely-shared infra or a +known low-coverage cause (NOTES) -- both are reading aids, not quality gates. + +What counts toward a total vs. what gets its own row are deliberately +decoupled decisions -- see classify(). + +Single-report mode hides untouched modules and totals only the executed ones, +so a partial run (e.g. MaxText E2E) isn't diluted by code it can't reach; the +two-report comparison keeps the full denominator. """ import json import sys from collections import defaultdict -OMIT_MODULES = {"tools", "platforms", "pretrain.py"} +OMIT_MODULES = {"tools", "platforms"} # Top-level groups whose sub-packages are shown as indented detail rows; every -# other group (modules, agents, cli, ...) is a single bold row. +# other group (agents, cli, ...) is a single bold row. DETAILED_GROUPS = ("core", "backends") def classify(path: str): - """Return (group, detail) for a covered file, or None to skip it. - - group is the top-level row key; detail is the sub-row key for - DETAILED_GROUPS (e.g. core/projection), else None. + """Return (group, detail) for a file, or None to drop it (not counted). + + group is the top-level row key; detail is the sub-row key within + DETAILED_GROUPS (e.g. "core/projection"). detail=None means "counted + toward group's total, no row of its own" -- true for every + non-DETAILED_GROUPS group, and for a bare file with no sub-package of its + own (pretrain.py -> folds into "primus (top-level)"; core/base_module.py -> + folds into "core"). Folding is by path *depth*, not filename, so a future + file added the same way folds the same way for free. + + Dropping entirely (OMIT_MODULES above) is the only *policy* exclusion -- + orthogonal to the structural folding here. """ seg = (path[path.find("primus/") :] if "primus/" in path else path).split("/") if seg[-1] == "__init__.py": return None - if len(seg) < 2: - return "(top-level)", None - if len(seg) == 2: # primus/.py, e.g. pretrain.py - return seg[1], None + if len(seg) <= 2: # outside primus/ entirely, or primus/.py directly (no sub-package) + return "primus (top-level)", None if seg[1] in DETAILED_GROUPS: + if len(seg) == 3: # primus//.py directly: no sub-package of its own + return seg[1], None return seg[1], seg[1] + "/" + seg[2] return seg[1], None @@ -62,6 +65,52 @@ def _pct(covered: int, total: int) -> float: return (100.0 * covered / total) if total else 0.0 +# Best-effort visual highlight for the headline column: GitHub strips CSS color +# from Action run summaries, so emoji is the portable substitute. Thresholds +# are a rough reading aid, not a quality gate. +_TIER_THRESHOLDS = ((50.0, "\U0001F7E2"), (25.0, "\U0001F7E1")) # >=50% green, >=25% yellow, else red +_TIER_RED = "\U0001F534" + + +def _tier(pct: float) -> str: + for threshold, emoji in _TIER_THRESHOLDS: + if pct >= threshold: + return emoji + return _TIER_RED + + +# Coverage priority: lower sorts first (ties keep coverage order), pinning the +# most important groups to the top. core (Primus's core) ranks highest. +PRIORITY = {"core": 0} +_DEFAULT_PRIORITY = 1 + + +def _priority(group: str) -> int: + return PRIORITY.get(group, _DEFAULT_PRIORITY) + + +# Curated, best-effort context -- not exhaustive. Flags infra whose coverage +# matters more than its size suggests, and low-coverage areas with a known, +# persistent cause, so it isn't re-litigated every read. Keyed like the table +# (group, or "/"); update alongside any fix or new finding. +#
forces cell wrapping so one long note can't stretch the whole column. +NOTES = { + "primus (top-level)": "loose primus/ modules, no sub-package;
auto-folded by path depth", + "core": "\U0001F511 Primus core;
imported by every run", + "backends/megatron": ( + "100+ patches gated by fp8 / MoE /
zero-bubble-pp / fsdp2 flags;
" + "CI E2E runs only 1-2 configs" + ), + "backends/transformer_engine": "fp8 GEMM / attn-overlap kernels;
only hit when an E2E enables fp8", + "backends/diffusion": "no E2E trainer suite yet
(unit-tested only)", +} + +_LEGEND = ( + "\U0001F7E2 >=50% / \U0001F7E1 >=25% / \U0001F534 <25% (next to module) " + " \u00b7  \U0001F511 Primus core\n" +) + + def _aggregate(report: dict): """Return {group: {detail|group: [covered, statements]}} for kept modules.""" agg = defaultdict(lambda: defaultdict(lambda: [0, 0])) @@ -89,27 +138,34 @@ def e2e_cov(group, key): # modules, so max avoids double counting the shared core code). return max((sa.get(group, {}).get(key, [0])[0] for sa in sas), default=0) if two else 0 - def row(label, cov, stmts, e2e, bold=False): + def row(label, key, cov, stmts, e2e, bold=False): + # Tier dot rides next to the module name (leftmost), so it reads as a + # per-module health mark and leaves the right-aligned % columns clean. + # It goes *after* any   indent so detail-row dots stay indented. + headline = _pct(e2e if two else cov, stmts) + dot = _tier(headline) + " " if stmts else "" if two: vals = [format(stmts, ","), "%.1f%%" % _pct(cov, stmts), "%.1f%%" % _pct(e2e, stmts)] else: vals = [format(cov, ","), format(stmts, ","), "%.1f%%" % _pct(cov, stmts)] w = "**" if bold else "" - return "| " + " | ".join("%s%s%s" % (w, x, w) for x in [label] + vals) + " |" + indent = " " if label.startswith(" ") else "" + rest = label[len(indent) :] + label_cell = "%s%s%s%s%s" % (indent, dot, w, rest, w) + cells = [label_cell] + ["%s%s%s" % (w, x, w) for x in vals] + cells.append(NOTES.get(key, "")) + return "| " + " | ".join(cells) + " |" def group_totals(group): - # Single-report mode counts only executed entries (cov > 0) so a partial - # run isn't diluted by sub-modules it never touched; two-report keeps all. + # Two-report mode keeps every entry (full denominator); single-report + # mode counts only executed ones so a partial run isn't diluted. entries = [(k, v) for k, v in pa[group].items() if two or v[0] > 0] cov = sum(v[0] for _, v in entries) stmts = sum(v[1] for _, v in entries) e2e = sum(e2e_cov(group, k) for k, _ in entries) if two else 0 return cov, stmts, e2e - # Single-report mode hides modules with zero coverage and totals over the - # executed modules only, so a partial run (e.g. MaxText E2E) isn't diluted by - # code it can never touch. The two-report comparison keeps the full denominator. - def group_executed(group): + def group_executed(group): # single-report mode: hide groups a partial run never touched return group_totals(group)[0] > 0 if not two else True groups = [g for g in pa if group_totals(g)[1] > 0 and group_executed(g)] @@ -136,25 +192,27 @@ def group_executed(group): out.append( "_Including all modules (nothing excluded): Unit %.1f%% -> Unit+E2E %.1f%%._\n" % (p_all, s_all) ) - out += ["| Module | Stmts | Unit | Unit+E2E |", "|---|--:|--:|--:|"] + out.append(_LEGEND) + out += ["| Module | Stmts | Unit | Unit+E2E | Notes |", "|---|--:|--:|--:|---|"] else: out.append( "**Total line coverage: %.1f%%** (%s / %s statements; excludes %s)\n" % (_pct(tc, tn), format(tc, ","), format(tn, ","), excl) ) - out += ["| Module | Covered | Stmts | Coverage |", "|---|--:|--:|--:|"] + out.append(_LEGEND) + out += ["| Module | Covered | Stmts | Coverage | Notes |", "|---|--:|--:|--:|---|"] - # Top-level groups, sorted by coverage (desc). - for group in sorted(groups, key=lambda g: -_pct(group_totals(g)[0], group_totals(g)[1])): + for group in sorted(groups, key=lambda g: (_priority(g), -_pct(group_totals(g)[0], group_totals(g)[1]))): cov, stmts, e2e = group_totals(group) - out.append(row("`%s`" % group, cov, stmts, e2e, bold=True)) + out.append(row("`%s`" % group, group, cov, stmts, e2e, bold=True)) if group in DETAILED_GROUPS: - # In single-report mode, hide sub-rows that were never executed. - details = ((k, v) for k, v in pa[group].items() if v[1] > 0 and (two or v[0] > 0)) + # k == group is a folded loose file (see classify()), already + # counted in group_totals() above -- no row of its own here. + details = ((k, v) for k, v in pa[group].items() if k != group and v[1] > 0 and (two or v[0] > 0)) for k, v in sorted(details, key=lambda kv: -_pct(kv[1][0], kv[1][1])): - out.append(row(" `%s`" % k, v[0], v[1], e2e_cov(group, k))) + out.append(row(" `%s`" % k, k, v[0], v[1], e2e_cov(group, k))) - out.append(row("TOTAL", tc, tn, te, bold=True)) + out.append(row("TOTAL", None, tc, tn, te, bold=True)) return "\n".join(out) diff --git a/tools/ci/junit_summary.py b/tools/ci/junit_summary.py index b84936674..121fdda67 100644 --- a/tools/ci/junit_summary.py +++ b/tools/ci/junit_summary.py @@ -25,6 +25,17 @@ COLUMNS = ("tests", "passed", "failed", "errors", "skipped") +def _fmt_time(seconds): + seconds = int(seconds) + hours, rem = divmod(seconds, 3600) + minutes, secs = divmod(rem, 60) + if hours: + return f"{hours}h{minutes:02d}m{secs:02d}s" + if minutes: + return f"{minutes}m{secs:02d}s" + return f"{secs}s" + + def parse_file(path): """Return (label, stats|None, failures) for one JUnit XML file.""" label = os.path.splitext(os.path.basename(path))[0] @@ -69,15 +80,30 @@ def render(reports, title=None): lines.append("| `%s` | _no report_ | | | | | |" % label) continue lines.append( - "| `%s` | %d | %d | %d | %d | %d | %.1fs |" - % (label, st["tests"], st["passed"], st["failed"], st["errors"], st["skipped"], st["time"]) + "| `%s` | %d | %d | %d | %d | %d | %s |" + % ( + label, + st["tests"], + st["passed"], + st["failed"], + st["errors"], + st["skipped"], + _fmt_time(st["time"]), + ) ) for k in total: total[k] += st[k] failures += [(label, *f) for f in fails] lines.append( - "| **TOTAL** | **%d** | **%d** | **%d** | **%d** | **%d** | **%.1fs** |" - % (total["tests"], total["passed"], total["failed"], total["errors"], total["skipped"], total["time"]) + "| **TOTAL** | **%d** | **%d** | **%d** | **%d** | **%d** | **%s** |" + % ( + total["tests"], + total["passed"], + total["failed"], + total["errors"], + total["skipped"], + _fmt_time(total["time"]), + ) ) lines.append("") diff --git a/tools/ci/runtime_summary.py b/tools/ci/runtime_summary.py index dd55ff08c..4d8635f97 100644 --- a/tools/ci/runtime_summary.py +++ b/tools/ci/runtime_summary.py @@ -4,14 +4,32 @@ # See LICENSE for license information. ############################################################################### -"""Render CI stage wall-clock times (a stageseconds TSV) as a Markdown table. +"""Render the *current* job's own step wall-clock times (fetched from the +Actions API) as a Markdown table for the CI run summary. -Complements junit_summary (test time) by surfacing the heavy build/install -stages. A missing/empty TSV renders nothing and exits 0 so the step never fails -the job; stage order (i.e. execution order) is preserved. +Auto-discovers every step from the job's live metadata, so adding, removing, +or renaming a step in ci.yaml needs no matching edit here and no hand-rolled +per-step timing block -- unlike manually timing "just this stage" into a +file, which silently misses whatever nobody remembered to wrap. + +Requires the calling step to export GITHUB_TOKEN with `actions: read` (see +ci.yaml); a missing token, network failure, or running outside Actions all +degrade to "print nothing, exit 0" so this can never fail the job: + + GITHUB_TOKEN=... python tools/ci/runtime_summary.py --title torch >> "$GITHUB_STEP_SUMMARY" + +Steps shorter than --min-seconds (default 5s, almost always a banner `echo` +step) are hidden so the table stays focused. """ import argparse +import json +import os +import urllib.error +import urllib.request +from datetime import datetime, timezone + +API_TIMEOUT_S = 30 def fmt(seconds): @@ -25,44 +43,84 @@ def fmt(seconds): return f"{secs}s" -def parse(path): +def _parse_ts(s): + if not s: + return None + return datetime.strptime(s, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc) + + +def fetch_steps(repo, run_id, job_name, runner_name, token): + """Return the step list for the run's job named `job_name`, preferring the + one running on `runner_name` (in case the name is ever repeated, e.g. a + future matrix), or [] if no such job is found.""" + url = f"https://api.github.com/repos/{repo}/actions/runs/{run_id}/jobs?per_page=100" + req = urllib.request.Request( + url, + headers={ + "Authorization": f"Bearer {token}", + "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28", + }, + ) + with urllib.request.urlopen(req, timeout=API_TIMEOUT_S) as resp: # noqa: S310 (fixed https:// API host) + jobs = json.load(resp).get("jobs", []) + matches = [j for j in jobs if j.get("name") == job_name] + exact = [j for j in matches if j.get("runner_name") == runner_name] + job = (exact or matches or [None])[0] + return job.get("steps", []) if job else [] + + +def step_durations(steps, min_seconds=5.0): + """[(name, seconds), ...] in run order, for steps that actually completed + (excludes the still-running caller itself, any not-yet-run step, and + anything skipped) and cleared the min_seconds noise floor.""" rows = [] - try: - with open(path) as handle: - for line in handle: - parts = line.rstrip("\n").split("\t") - if len(parts) != 2 or not parts[0].strip(): - continue - try: - rows.append((parts[0].strip(), float(parts[1].strip()))) - except ValueError: - continue - except OSError: - return [] + for step in steps: + if step.get("conclusion") == "skipped": + continue + start, end = _parse_ts(step.get("started_at")), _parse_ts(step.get("completed_at")) + if not start or not end: + continue # still running (incl. this very step) or never reached + secs = (end - start).total_seconds() + if secs >= min_seconds: + rows.append((step.get("name", "?"), secs)) return rows -def render(rows, title=None): +def render(steps, title=None, min_seconds=5.0): + rows = step_durations(steps, min_seconds) + if not rows: + return "" suffix = f" - {title}" if title else "" lines = [f"## CI runtime{suffix}\n", "| Stage | Time |", "|---|--:|"] - total = 0.0 - for stage, secs in rows: - lines.append(f"| {stage} | {fmt(secs)} |") - total += secs - lines.append(f"| **TOTAL (timed stages)** | **{fmt(total)}** |") + lines += [f"| {name} | {fmt(secs)} |" for name, secs in rows] + lines.append(f"| **TOTAL (shown stages)** | **{fmt(sum(s for _, s in rows))}** |") return "\n".join(lines) def main(): - ap = argparse.ArgumentParser(description="Render a stageseconds TSV as a Markdown runtime table.") - ap.add_argument("tsv", help="Path to the runtime TSV (stageseconds per line).") + ap = argparse.ArgumentParser(description="Render this job's own step timings as a Markdown table.") ap.add_argument("--title", default=None, help="Optional section title (e.g. torch).") + ap.add_argument( + "--min-seconds", type=float, default=5.0, help="Hide steps shorter than this (default 5s)." + ) args = ap.parse_args() - rows = parse(args.tsv) - if not rows: - return 0 # nothing timed; don't emit an empty table or fail the step - print(render(rows, args.title)) + token = os.environ.get("GITHUB_TOKEN") + repo, run_id, job_name = ( + os.environ.get("GITHUB_REPOSITORY"), + os.environ.get("GITHUB_RUN_ID"), + os.environ.get("GITHUB_JOB"), + ) + if not (token and repo and run_id and job_name): + return 0 # not running in Actions (or token not wired) -- nothing to render + try: + steps = fetch_steps(repo, run_id, job_name, os.environ.get("RUNNER_NAME"), token) + except (urllib.error.URLError, TimeoutError, ValueError, OSError): + return 0 # a reporting step must never fail the job + out = render(steps, args.title, args.min_seconds) + if out: + print(out) return 0 diff --git a/tools/ci/select_tests.py b/tools/ci/select_tests.py index cfff17855..cdff9d014 100644 --- a/tools/ci/select_tests.py +++ b/tools/ci/select_tests.py @@ -4,19 +4,20 @@ # See LICENSE for license information. ############################################################################### -"""Map a PR's changed files (stdin, one per line) to the tests to run. - -Default prints the minimal unit-test paths; --e2e prints the E2E trainer suites -(or "all"). A single classify() decides each path's blast radius and both -selections build on it. Conventions over hard-coded tables: - - unit dirs mirror the source tree (primus/ -> tests/unit_tests/), - resolved by walking up to the nearest existing dir; +"""Map a PR's changed files (stdin, one per line) to the E2E trainer suites to +run (or "all"). A single classify() decides each path's blast radius. +Conventions over hard-coded tables: - E2E suites are auto-discovered from tests/trainer/test__trainer.py; - a backend is named by its dir (primus/backends/ or examples/). -Fail-safe is the only invariant: anything global, unlocatable, or a backend -without a trainer expands to everything -- over-select, never under-select. +Fail-safe is the only invariant: anything global, a non-backend source change, +or a backend without a trainer expands to everything -- over-select, never +under-select. + + git diff --name-only "$BASE" HEAD | python tools/ci/select_tests.py - git diff --name-only "$BASE" HEAD | python tools/ci/select_tests.py [--e2e] +Unit tests are deliberately NOT selected/narrowed here: the whole suite only +takes ~5 minutes (vs. each E2E suite's tens of minutes), so narrowing risked +under-selection for little wall-clock gain. ci.yaml always runs it in full. """ import re @@ -24,28 +25,24 @@ from pathlib import Path ROOT = Path(__file__).resolve().parents[2] -FULL = "tests/unit_tests/" -# The only hard-coded list: changes whose blast radius is the whole repo. Being -# absent here only ever falls back to other fail-safe paths, never to "skip". +# The only hard-coded list: changes whose blast radius is the whole repo. +# Scoped to paths *outside* primus/, tests/unit_tests/, and examples/ -- +# anything under those trees that isn't a recognized backend already falls +# back to "component" below, which select_e2e() treats the same as "global" +# anyway, so listing e.g. primus/core/launcher/ here would be redundant. +# Being absent here only ever falls back to other fail-safe paths, never "skip". GLOBAL_TRIGGERS = ( ".github/", "tools/", "runner/", # the launcher drives all training "pyproject.toml", - "primus/__init__.py", - "primus/core/launcher/", - "primus/core/utils/", - "primus/core/config/", "tests/utils.py", "tests/conftest.py", "tests/unit_tests/conftest.py", "tests/run_unit_tests.py", ) -# megatron's GPU-operator tests aren't path-isomorphic to the backend source. -_BACKEND_EXTRA_UNIT = {"megatron": ("tests/unit_tests/megatron/",)} - _TRAINER_RE = re.compile(r"test_(.+)_trainer") _BACKEND_RE = re.compile(r"(?:primus/backends|examples)/([^/]+)/") @@ -64,20 +61,8 @@ def _is_global(path): return any(path == t or path.startswith(t) for t in GLOBAL_TRIGGERS) -def _nearest_unit_dir(rel): - # rel is source-relative (under primus/ or tests/unit_tests/); walk up to the - # nearest existing tests/unit_tests/<...> dir, or None if none exists. - parts = rel.split("/")[:-1] - while parts: - cand = "tests/unit_tests/" + "/".join(parts) + "/" - if (ROOT / cand).is_dir(): - return cand - parts.pop() - return None - - def classify(path): - """('global', None) | ('backend', name) | ('component', unit_dir|None) | ('ignore', None).""" + """('global', None) | ('backend', name) | ('component', None) | ('ignore', None).""" if _is_global(path): return ("global", None) backend = _BACKEND_RE.match(path) # primus/backends// or examples// @@ -86,43 +71,14 @@ def classify(path): if path.startswith("tests/trainer/"): m = _TRAINER_RE.search(path) return ("backend", m.group(1)) if m else ("ignore", None) - for root in ("primus/", "tests/unit_tests/"): - if path.startswith(root): - if not path.endswith(".py"): - return ("global", None) # non-.py here (configs, fixtures) -> fail-safe - return ("component", _nearest_unit_dir(path[len(root) :])) + # A bare examples/ (no backend subdir, so it didn't match _BACKEND_RE + # above) is shared launcher plumbing, not per-backend -- e.g. maxtext's E2E + # shells out to examples/run_pretrain.sh directly. + if path.startswith("primus/") or path.startswith("tests/unit_tests/") or path.startswith("examples/"): + return ("component", None) # any other source/unit-test/launcher change return ("ignore", None) # docs, README, ... outside the source/test trees -def select_targets(files): - files = [f.strip() for f in files if f.strip()] - if not files: - return [FULL] - targets = [] - - def add(d): - if d and d not in targets: - targets.append(d) - - for path in files: - kind, val = classify(path) - if kind == "global": - return [FULL] - if kind == "backend": - base = f"tests/unit_tests/backends/{val}/" - if not (ROOT / base).is_dir(): - return [FULL] # backend without a unit dir (e.g. transformer_engine) -> safe - add(base) - for extra in _BACKEND_EXTRA_UNIT.get(val, ()): - add(extra) - elif kind == "component": - if val is None: - return [FULL] # couldn't localize a unit dir -> safe - add(val) - # ignore -> skip - return targets or [FULL] - - def select_e2e(files, suites=None): suites = discover_e2e_suites() if suites is None else set(suites) files = [f.strip() for f in files if f.strip()] @@ -146,12 +102,9 @@ def select_e2e(files, suites=None): def main(): files = sys.stdin.read().splitlines() - if "--e2e" in sys.argv[1:]: - suites = discover_e2e_suites() - e2e = select_e2e(files, suites) - print("all" if suites and set(e2e) == suites else " ".join(e2e)) - else: - print(" ".join(select_targets(files))) + suites = discover_e2e_suites() + e2e = select_e2e(files, suites) + print("all" if suites and set(e2e) == suites else " ".join(e2e)) return 0 From 253b33229b5ddacd3584e70600698a4b437dc3a3 Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Thu, 16 Jul 2026 09:41:36 +0800 Subject: [PATCH 034/127] feat: add moe_router_force_load_balancing_type to select force load balancing type (#875) # Description This PR adds a new option `moe_router_force_load_balancing_type` that lets users select how force load balancing is applied in the MoE router. Previously, when `moe_router_force_load_balancing` was enabled, the balancing strategy was implicitly fixed. In practice we need two different behaviors: - A **deterministic** strategy where every expert receives exactly the same, step-invariant number of tokens (constant `M_total`) - The **original** Megatron-LM random-logits balancing behavior. The new option makes this choice explicit and backward compatible: existing behavior is preserved via the `"uniform"` mode, while the new `"even"` mode enables deterministic round-robin token assignment. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Add a new config option `moe_router_force_load_balancing_type` in `primus/configs/models/megatron/primus_megatron_model.yaml`, supporting `"even"` and `"uniform"` (Megatron-LM original behavior). - In `PrimusTopKRouter` (`primus/backends/megatron/core/transformer/moe/router.py`), add `_force_even_routing`, which deterministically assigns each token's top-k slots to experts via a round-robin cycle `(token_idx * topk + k) % num_experts`, so per-expert token counts are exactly balanced and step-invariant. Original top-k probability magnitudes are carried over to the new expert positions to keep `scores` and `routing_map` mutually consistent. This path is only taken when `moe_router_force_load_balancing` is on, the type is `"even"`, and DeepEP is not enabled. - In `PrimusTurboDeepEPTokenDispatcher` (`primus/backends/megatron/core/extensions/primus_turbo.py`), read the new option and only build the even round-robin `token_indices` when the type is `"even"`; the `"uniform"` type keeps `token_indices=None` and relies on the router's upstream random-logits balancing. # Checklist: - [x] The functionality is complete - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes Co-authored-by: Xiaoming-AMD --- .../megatron/core/extensions/primus_turbo.py | 13 ++++- .../megatron/core/transformer/moe/router.py | 49 +++++++++++++++++++ .../megatron/primus_megatron_model.yaml | 7 +++ 3 files changed, 67 insertions(+), 2 deletions(-) diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index 4d7a99c36..51b4f699a 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -1902,6 +1902,12 @@ def __init__( self._comm_manager = self.deepep_dispatcher self.moe_router_force_load_balancing = args.moe_router_force_load_balancing + # "even" -> deterministic round-robin token assignment (constant per-expert + # counts / constant M_total); "uniform" -> Megatron-LM original random-logits + # balancing (handled upstream in the router, token_indices stays None here). + self.moe_router_force_load_balancing_type = getattr( + args, "moe_router_force_load_balancing_type", "uniform" + ) def dispatch_preprocess( self, hidden_states: torch.Tensor, routing_map: torch.Tensor, probs: torch.Tensor @@ -1925,9 +1931,12 @@ def dispatch_preprocess( hidden_states = hidden_states.view(-1, self.config.hidden_size) num_tokens = hidden_states.shape[0] - # when force_load_balancing, we use even token_indices to make sure each expert get same number of tokens + # when force_load_balancing with type "even", we use round-robin token_indices + # to make sure each expert gets the same (deterministic) number of tokens. + # type "uniform" keeps token_indices=None and relies on the router's upstream + # random-logits balancing. token_indices = None - if self.moe_router_force_load_balancing: + if self.moe_router_force_load_balancing and self.moe_router_force_load_balancing_type == "even": token_indices = ( torch.arange(num_tokens * self.config.moe_router_topk, device=hidden_states.device).view( num_tokens, self.config.moe_router_topk diff --git a/primus/backends/megatron/core/transformer/moe/router.py b/primus/backends/megatron/core/transformer/moe/router.py index bec8a232b..a6b0d6bf7 100644 --- a/primus/backends/megatron/core/transformer/moe/router.py +++ b/primus/backends/megatron/core/transformer/moe/router.py @@ -115,4 +115,53 @@ def routing(self, logits: torch.Tensor, **kwargs): # so by the time we get here ``logits`` is already random when # ``args.moe_router_force_load_balancing`` is True. There is nothing # extra to do. + # ``moe_router_force_load_balancing_type`` selects the balancing mode + # (only relevant when ``moe_router_force_load_balancing`` is on): + # * "even" -> deterministic round-robin so per-expert counts are + # constant every step (constant M_total, no autotune/ + # recompile churn); handled by ``_force_even_routing``. + # * "uniform" -> Megatron-LM original random-logits balancing (already + # applied upstream in ``TopKRouter.forward``; nothing to + # do here). + force_load_balancing_type = getattr(args, "moe_router_force_load_balancing_type", "uniform") + if ( + args.moe_router_force_load_balancing + and force_load_balancing_type == "even" + and not getattr(args, "moe_enable_deepep", False) + ): + scores, routing_map = self._force_even_routing(scores, routing_map) + return scores, routing_map + + def _force_even_routing( + self, scores: torch.Tensor, routing_map: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Deterministically assign each token's top-k slots to experts in a + round-robin cycle ``(token_idx * topk + k) % num_experts`` so per-expert + token counts are exactly balanced and step-invariant (constant M_total). + + The original top-k probability magnitudes are carried over to the new + expert positions (per-slot), keeping ``scores`` and ``routing_map`` + mutually consistent so the MoE combine weights stay correct. + """ + num_tokens = routing_map.shape[0] + num_experts = self.config.num_moe_experts + topk = self.topk + device = routing_map.device + + # slot[t, k] = deterministic destination expert for token t's k-th slot. + # topk consecutive experts per token -> distinct while topk <= num_experts. + slot = ( + torch.arange(num_tokens * topk, device=device).view(num_tokens, topk) % num_experts + ) # [num_tokens, topk], int64 + + new_routing_map = torch.zeros_like(routing_map) + new_routing_map.scatter_(1, slot, torch.ones_like(slot, dtype=routing_map.dtype)) + + # Carry the real per-token top-k probability magnitudes to the new slots. + # scores is non-zero only on the real top-k, so topk() extracts exactly them. + topk_vals, _ = torch.topk(scores, topk, dim=1) # [num_tokens, topk] + new_scores = torch.zeros_like(scores) + new_scores.scatter_(1, slot, topk_vals.to(scores.dtype)) + + return new_scores, new_routing_map diff --git a/primus/configs/models/megatron/primus_megatron_model.yaml b/primus/configs/models/megatron/primus_megatron_model.yaml index 0376e6014..a7a12d70b 100644 --- a/primus/configs/models/megatron/primus_megatron_model.yaml +++ b/primus/configs/models/megatron/primus_megatron_model.yaml @@ -13,3 +13,10 @@ router_logit_softcapping: null # float lfm_layer_types: null # list[str] conv_L_cache: 3 # int conv_bias: false # bool + + +# Primus patch option +# Control the force load balancing type for the MoE router. +# If set to "even", the router will force the load balancing to be even. +# If set to "uniform", the router will force the load balancing to be uniform. (Megatron-LM original behavior) +moe_router_force_load_balancing_type: "even" From 962a91d22f577851cacc0fffe45bbee89e1e4bb1 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Thu, 16 Jul 2026 09:05:20 +0300 Subject: [PATCH 035/127] feat(flux): diffusion + Flux pretrain trainers (#819) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/flux` and also merges `feat/flux/trainprim`, `feat/flux/fp8`, `feat/flux/data` — review after all four. ## What this changes The trainers that tie the feature together: the generic diffusion trainer, the Flux pretrain trainer, and the validation-metric evaluator. ## Why it has several parents The diffusion trainer lazily imports the data providers (`feat/flux/data`) and pulls in training primitives (`feat/flux/trainprim`) and fp8 (`feat/flux/fp8`); without the data parent its own tests fail with `ModuleNotFoundError: primus.backends.megatron.data`. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); its tests exercise no mxfp4/compile path and pass on the current CI pin (no turbo-bump dependency). Builds on `feat/flux/flux` + `feat/flux/trainprim` + `feat/flux/fp8` + `feat/flux/data`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion/training tests/unit_tests/backends/megatron/test_chimera_rng_restore.py`. Validated locally on an AMD GPU container: 41 passed. ## Files 11 (diffusion trainer, Flux pretrain trainer, evaluator + tests). --------- Co-authored-by: Flux Split Trial Co-authored-by: Luiza Sayfullina --- .../core/models/diffusion/flux/config.py | 12 + primus/backends/megatron/diffusion_trainer.py | 386 ++++++++ .../megatron/flux_pretrain_trainer.py | 843 ++++++++++++++++++ .../backends/megatron/training/evaluator.py | 16 +- .../training/test_diffusion_trainer.py | 343 +++++++ .../training/test_flux_forward_step_e2e.py | 274 ++++++ .../training/test_flux_model_creation.py | 288 ++++++ .../diffusion/training/test_flux_trainer.py | 219 +++++ .../training/test_forward_step_count_gate.py | 173 ++++ .../test_vae_resample_reproducibility.py | 211 +++++ .../megatron/test_chimera_rng_restore.py | 109 +++ ...st_diffusion_trainer_forward_step_count.py | 161 ++++ 12 files changed, 3033 insertions(+), 2 deletions(-) create mode 100644 primus/backends/megatron/diffusion_trainer.py create mode 100644 primus/backends/megatron/flux_pretrain_trainer.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_flux_trainer.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_forward_step_count_gate.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/training/test_vae_resample_reproducibility.py create mode 100644 tests/unit_tests/backends/megatron/test_chimera_rng_restore.py create mode 100644 tests/unit_tests/backends/megatron/test_diffusion_trainer_forward_step_count.py diff --git a/primus/backends/megatron/core/models/diffusion/flux/config.py b/primus/backends/megatron/core/models/diffusion/flux/config.py index 4d920de51..83446615b 100644 --- a/primus/backends/megatron/core/models/diffusion/flux/config.py +++ b/primus/backends/megatron/core/models/diffusion/flux/config.py @@ -20,6 +20,18 @@ from ..common.config import BaseDiffusionConfig +# Try to import erf_gelu from megatron, fallback to custom implementation +try: + from megatron.core.transformer.utils import erf_gelu +except ImportError: + # Fallback used when Megatron's erf_gelu is unavailable + def erf_gelu(x): + """GELU activation using error function approximation.""" + return 0.5 * x * (1.0 + torch.erf(x / 1.4142135623730951)) + + +__all__ = ["FluxConfig", "openai_gelu_no_jit", "erf_gelu"] + # Custom non-JIT compiled openai_gelu to avoid ROCm bugs def openai_gelu_no_jit(x): diff --git a/primus/backends/megatron/diffusion_trainer.py b/primus/backends/megatron/diffusion_trainer.py new file mode 100644 index 000000000..dd91f3285 --- /dev/null +++ b/primus/backends/megatron/diffusion_trainer.py @@ -0,0 +1,386 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Base diffusion trainer for Megatron-LM. + +This trainer provides a foundation for diffusion model training by overriding +the forward step and dataset provider methods from MegatronPretrainTrainer. +""" + +from abc import abstractmethod + +import torch +import torch.nn.functional as F + +from primus.backends.megatron.megatron_pretrain_trainer import MegatronPretrainTrainer +from primus.core.utils.module_utils import log_rank_0 + + +class DiffusionPretrainTrainer(MegatronPretrainTrainer): + """ + Base trainer for diffusion models. + + This trainer inherits from the backend's MegatronPretrainTrainer (which uses + MegatronBaseTrainer). + + It overrides the forward step and dataset provider methods to support + diffusion-specific training. Subclasses should implement create_scheduler() + and may override create_model() if needed. + + Configuration is accessed via backend_args (set by BaseTrainer.__init__()), + not via module_config.params, to avoid BaseModule dependency. + """ + + def __init__(self, *args, **kwargs): + """ + Initialize diffusion trainer. + + Args: + *args: Positional arguments passed to parent trainer + **kwargs: Keyword arguments passed to parent trainer + (backend_args is extracted from kwargs by BaseTrainer) + """ + super().__init__(*args, **kwargs) + + self._scheduler = None + self._compiled_loss_fn = None + self._forward_step_count = 0 + self._forward_step_count_initialized = False + + # Composition pattern: avoids recreating the provider on each call + use_mock_data = getattr(self.backend_args, "mock_data", False) + + if use_mock_data: + from primus.backends.megatron.data.synthetic_dataset_provider import ( + SyntheticDatasetProvider, + ) + + # Get mock dataset configuration from YAML (if specified) + mock_dataset_config_ns = getattr(self.backend_args, "mock_dataset", None) + mock_dataset_config = self._convert_namespace_to_dict(mock_dataset_config_ns) + + # Determine model type for default dataset selection + model_type_raw = getattr(self.backend_args, "model_type", None) + model_type = ( + model_type_raw + if (isinstance(model_type_raw, str) and model_type_raw != "diffusion_model") + else "flux" + ) + + self.data_provider = SyntheticDatasetProvider( + dataset_config=mock_dataset_config, model_type=model_type + ) + log_rank_0("Created SyntheticDatasetProvider for mock data") + else: + from primus.backends.megatron.data.energon_dataset_provider import ( + EnergonDatasetProvider, + ) + + self.data_provider = EnergonDatasetProvider(task_encoder_factory=lambda: self.get_task_encoder()) + log_rank_0("Created EnergonDatasetProvider for real data") + + log_rank_0(f"{self.__class__.__name__} initialized") + log_rank_0(f"Data provider: {type(self.data_provider).__name__}") + + def _convert_namespace_to_dict(self, ns_obj): + """ + Convert SimpleNamespace (or nested SimpleNamespace) to dict. + + Handles nested namespaces by recursively converting them. + + Args: + ns_obj: SimpleNamespace, dict, or None + + Returns: + dict, original input (if already a dict or primitive), or None + """ + if ns_obj is None: + return None + + if not hasattr(ns_obj, "__dict__"): + return ns_obj # Already a dict or other type + + result = vars(ns_obj) + + # Recursively convert nested namespaces + if "params" in result and hasattr(result["params"], "__dict__"): + result["params"] = vars(result["params"]) + + return result + + def setup(self): + """ + Override setup() to inject diffusion model_provider. + + MegatronBaseTrainer.setup() handles Megatron path, global vars, and parse_args patching. + We only need to set the model_provider here. + """ + # Ensure data_parallel_size is set before parent setup() calls set_primus_global_variables() + # This is required by set_primus_global_variables() + if not hasattr(self.backend_args, "data_parallel_size"): + world_size = getattr(self.backend_args, "world_size", 1) + tensor_model_parallel_size = getattr(self.backend_args, "tensor_model_parallel_size", 1) + pipeline_model_parallel_size = getattr(self.backend_args, "pipeline_model_parallel_size", 1) + context_parallel_size = getattr(self.backend_args, "context_parallel_size", 1) + data_parallel_size = world_size // ( + tensor_model_parallel_size * pipeline_model_parallel_size * context_parallel_size + ) + setattr(self.backend_args, "data_parallel_size", data_parallel_size) + log_rank_0( + f"Computed data_parallel_size={data_parallel_size} from world_size={world_size}, tp={tensor_model_parallel_size}, pp={pipeline_model_parallel_size}, cp={context_parallel_size}" + ) + + # Create and set diffusion model provider + def _diffusion_model_provider( + pre_process=True, post_process=True, vp_stage=None, config=None, pg_collection=None + ): + """ + Model provider wrapper for diffusion models. + + This matches Megatron's model_provider signature but calls + our diffusion-specific create_model() instead of gpt_builder. + + Note: vp_stage, config, and pg_collection are accepted for interface + compatibility with Megatron's model_provider signature but are not + used by diffusion models. + """ + return self.create_model(pre_process=pre_process, post_process=post_process) + + self.model_provider = _diffusion_model_provider + log_rank_0("=" * 80) + log_rank_0("Overridden model_provider to use diffusion model builder") + log_rank_0("=" * 80) + + # Call parent's setup() which handles Megatron initialization + super().setup() + + @abstractmethod + def create_model(self, pre_process=True, post_process=True): + """ + Create diffusion model instance. + + Returns: + Model instance (e.g., Flux) + + Example: + from primus.backends.megatron.core.models.diffusion.flux import Flux + return Flux(config=self.model_config) + """ + raise NotImplementedError("Subclasses must implement create_model()") + + def _get_loss_fn(self): + """Return the loss function, optionally compiled. + + When torch.compile is enabled for the model, the loss function is also + compiled as a standalone region. This fuses the ~10 eager ATen ops + (sub, float casts, mse_loss, mean) into 1-2 kernels and gives the + autograd engine a single CompiledFunctionBackward node instead of + multiple eager backward nodes. + + The compiled function is created lazily on first call and cached. + """ + if self._compiled_loss_fn is not None: + return self._compiled_loss_fn + + from primus.backends.megatron.training.diffusion.loss_computation import ( + compute_flow_matching_loss, + ) + + try: + from megatron.training import get_args + + args = get_args() + compile_enabled = getattr(args, "enable_torch_compile", False) + except Exception: + compile_enabled = False + + if compile_enabled: + import torch + + self._compiled_loss_fn = torch.compile( + compute_flow_matching_loss, + backend="inductor", + fullgraph=False, + ) + log_rank_0("[DiffusionPretrainTrainer] Loss function compiled with torch.compile") + else: + self._compiled_loss_fn = compute_flow_matching_loss + + return self._compiled_loss_fn + + def forward_step(self, data_iterator, model, return_schedule_plan=False): + """ + Forward training step for diffusion models. + + Args: + data_iterator: Data iterator + model: Diffusion model (Flux) + return_schedule_plan: Whether to return schedule plan (for pipeline parallelism) + + Returns: + Tuple of (noise_pred, loss_func_callable) + """ + from primus.backends.megatron.training.diffusion.forward_step import ( + flux_forward_step_func, + ) + + # Skip counter advance in eval. Advancing _forward_step_count on + # validation steps would shift the next training step's per-step seed + # by eval_iters * num_microbatches per --eval-interval window, + # defeating the goal of isolating training RNG from unrelated forward + # passes. Eval forward passes reuse the most recent training counter + # value, so the per-step CUDA reseed is a no-op replay during eval. + if model.training: + self._forward_step_count += 1 + + # Megatron's pattern: forward_step returns model output, loss_func computes loss + noise_pred, clean_latents, noise, loss_mask, metrics, is_validation = flux_forward_step_func( + data_iterator, + model, + scheduler=self.scheduler, + use_guidance_embed=getattr(self, "use_guidance_embed", False), + guidance_scale=getattr(self, "guidance_scale", None), + timestep_sampler=getattr(self, "timestep_sampler", None), + cfg_dropout_prob=getattr(self, "cfg_dropout_prob", 0.0), + empty_t5_encodings=getattr(self, "empty_t5_encodings", None), + empty_clip_encodings=getattr(self, "empty_clip_encodings", None), + vae_scale=getattr(self, "vae_scale", None), + vae_shift=getattr(self, "vae_shift", None), + vae_latent_mode=getattr(self, "vae_latent_mode", "presampled"), + per_step_rng_reseed=getattr(self, "per_step_rng_reseed", False), + step_count=self._forward_step_count, + ) + + # Store values needed for loss computation (will be used by loss function) + self._last_clean_latents = clean_latents + self._last_noise = noise + self._last_loss_mask = loss_mask + + # Store metrics in runtime state + if hasattr(self, "runtime_state") and self.runtime_state: + self.runtime_state.update_metrics(metrics) + else: + log_rank_0("[DiffusionPretrainTrainer] WARNING: runtime_state not available, metrics not stored") + + if is_validation: + + def val_loss_func(output_tensor, non_loss_data=False): + if non_loss_data: + return output_tensor + target = self._last_noise - self._last_clean_latents + loss = F.mse_loss(output_tensor.float(), target.float(), reduction="none") + loss_per_sample = loss.mean(dim=tuple(range(1, loss.ndim))) + loss_sum = loss_per_sample.sum() + sample_count = torch.tensor( + loss_per_sample.numel(), dtype=loss_sum.dtype, device=loss_sum.device + ) + return loss_sum, {"loss": (loss_sum.detach(), sample_count.detach())} + + return noise_pred, val_loss_func + + def diffusion_loss_func(output_tensor, non_loss_data=False): + if non_loss_data: + return output_tensor + + loss = self._get_loss_fn()( + output_tensor, self._last_clean_latents, self._last_noise, self._last_loss_mask + ) + + reporting_metrics = {"reduced_train_loss": loss.detach().clone()} + + return loss, reporting_metrics + + return noise_pred, diffusion_loss_func + + def get_forward_step(self): + """ + Return forward step function for diffusion models. + + Returns a function that wraps self.forward_step() to match + the interface expected by Megatron's pretrain() function. + + On the first call, reconstructs the forward step counter from + checkpoint state (args.iteration * num_microbatches) so that + the per-step RNG seed sequence continues correctly after resume. + """ + + def diffusion_forward_step(data_iterator, model): + if not self._forward_step_count_initialized: + from megatron.core.num_microbatches_calculator import ( + get_num_microbatches, + ) + from megatron.training import get_args + + args = get_args() + # Reconstruct counter from checkpoint iteration. Assumes no + # iterations were skipped (iterations_to_skip is unused in + # diffusion training). + self._forward_step_count = args.iteration * get_num_microbatches() + self._forward_step_count_initialized = True + + return self.forward_step(data_iterator, model, return_schedule_plan=False) + + return diffusion_forward_step + + def get_datasets_provider(self): + """ + Return dataset provider function that delegates to self.data_provider. + + This uses the composition pattern: the provider function delegates + to the injected data_provider instance, avoiding recreation on each call. + """ + + def diffusion_datasets_provider(train_val_test_num_samples, vp_stage=None): + """Delegate to injected data provider.""" + from megatron.training import get_args + + args = get_args() + + return self.data_provider.create_dataloaders( + trainer_config=args, train_val_test_num_samples=train_val_test_num_samples, vp_stage=vp_stage + ) + + # Mark as distributed (required by Megatron) + diffusion_datasets_provider.is_distributed = self.data_provider.is_distributed + + # Set __module__ to help with debugging (point to this module, not pretrain_gpt) + diffusion_datasets_provider.__module__ = __name__ + + log_rank_0( + f"[DiffusionPretrainTrainer] Created datasets provider: {type(self.data_provider).__name__}" + ) + return diffusion_datasets_provider + + @property + def scheduler(self): + """Lazily-initialized diffusion scheduler.""" + if self._scheduler is None: + self._scheduler = self.create_scheduler() + return self._scheduler + + @abstractmethod + def create_scheduler(self): + """ + Create diffusion scheduler. + + Returns: + Scheduler instance (e.g., FlowMatchEulerDiscreteScheduler) + """ + raise NotImplementedError("Subclasses must implement create_scheduler()") + + @abstractmethod + def get_task_encoder(self): + """ + Get task encoder for Energon data pipeline. + + Returns: + TaskEncoder instance (e.g., EncodedDiffusionTaskEncoder) + + This method is called by EnergonDatasetProvider to create the task encoder + for processing WebDataset samples. + """ + raise NotImplementedError("Subclasses must implement get_task_encoder()") diff --git a/primus/backends/megatron/flux_pretrain_trainer.py b/primus/backends/megatron/flux_pretrain_trainer.py new file mode 100644 index 000000000..36da05ce3 --- /dev/null +++ b/primus/backends/megatron/flux_pretrain_trainer.py @@ -0,0 +1,843 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Flux Pretrain Trainer for Primus-Megatron. + +This trainer implements Flux-specific training logic including: + - Flow matching scheduler with dynamic shifting + - Guidance embedding support + - Custom forward step function +""" + +import os + +import numpy as np +import torch +import torch.nn as nn + +from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer +from primus.backends.megatron.training.diffusion.schedulers import ( + FlowMatchEulerDiscreteScheduler, +) +from primus.backends.megatron.training.diffusion.timestep_sampling import ( + create_timestep_sampler, +) +from primus.core.utils.module_utils import log_rank_0 + + +def _restore_chimera_rng_state(args) -> None: + """Restore canonical RNG state after chimera model init. + + Calls Megatron's `_set_random_seed` to restore CPU, CUDA default, and + model-parallel tracker generators to a canonical (per-rank-uniform) state. + Falls back to a manual restore if Megatron's signature changes. + + Raises: + May propagate exceptions other than ImportError/TypeError from + _set_random_seed. The fallback handles ImportError (missing module) + and TypeError (API signature changes) gracefully. + """ + try: + from megatron.training.initialize import _set_random_seed + + _set_random_seed( + args.seed, + args.data_parallel_random_init, + args.te_rng_tracker, + args.inference_rng_tracker, + use_cudagraphable_rng=getattr(args, "enable_cuda_graph", False), + ) + except (ImportError, TypeError) as e: + import logging + + from megatron.core.tensor_parallel import random as tp_random + + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + tp_random.model_parallel_cuda_manual_seed(args.seed) + logging.warning( + "[nemo_chimera_init] _set_random_seed API changed (%s). " + "Restored all RNG generators manually. Verify convergence " + "matches expected behavior.", + e, + ) + + +class FluxPretrainTrainer(DiffusionPretrainTrainer): + """ + Trainer for Flux diffusion model pre-training. + + Flux-specific features: + - Flow matching with Euler discrete scheduler + - Dynamic timestep shifting for variable resolution + - Optional guidance embedding for CFG + + Config access via backend_args: + - Megatron args (guidance_embed, guidance_scale, etc.) via backend_args + - Primus-specific config (torch_compile) via backend_args (overrides section) + """ + + def __init__(self, *args, **kwargs): + """ + Initialize Flux pretrain trainer. + + Args: + *args: Positional arguments passed to parent trainer + **kwargs: Keyword arguments passed to parent trainer + (backend_args is extracted from kwargs by BaseTrainer.__init__()) + """ + super().__init__(*args, **kwargs) + + self._training_rng_seeded = False + + # backend_args is set by BaseTrainer.__init__() + params = self.backend_args + + # Default False: per-step CUDA RNG reseed is an MLPerf-alignment + # feature; flipping it on globally changes the observable RNG + # sequence for every consumer of this trainer. MLPerf-aligned YAMLs + # opt in explicitly via per_step_rng_reseed: true. + self.per_step_rng_reseed = getattr(params, "per_step_rng_reseed", False) + self.nemo_chimera_init = getattr(params, "nemo_chimera_init", False) + + # These come from the 'overrides' section in the YAML + self.use_guidance_embed = getattr(params, "guidance_embed", False) + self.guidance_scale = getattr(params, "guidance_scale", 3.5) + + # Scheduler config + self.num_train_timesteps = getattr(params, "num_train_timesteps", 1000) + self.scheduler_shift = getattr(params, "scheduler_shift", 1.0) + self.use_dynamic_shifting = getattr(params, "use_dynamic_shifting", False) + + # Timestep sampling strategy (MLPerf uses "direct_uniform") + timestep_strategy = getattr(params, "timestep_sampling_strategy", "logit_normal") + self.timestep_sampler = create_timestep_sampler(timestep_strategy) + log_rank_0( + f"Timestep sampling strategy: {timestep_strategy} -> {type(self.timestep_sampler).__name__}" + ) + + # CFG dropout: replace text embeddings with fixed empty encodings at this probability + self.cfg_dropout_prob = getattr(params, "cfg_dropout_prob", 0.0) + self.empty_t5_encodings = None + self.empty_clip_encodings = None + + if self.cfg_dropout_prob > 0.0: + self._init_cfg_dropout(params) + + # VAE latent normalization (matches NVIDIA MLPerf v5.1) + self.vae_scale = getattr(params, "vae_scale", None) + self.vae_shift = getattr(params, "vae_shift", None) + if self.vae_scale is not None: + log_rank_0(f"VAE normalization: scale={self.vae_scale}, shift={self.vae_shift}") + + # VAE latent mode: "presampled" uses stored latents directly, + # "resample" re-draws from (mean, logvar) each step + self.vae_latent_mode = getattr(params, "vae_latent_mode", "presampled") + if self.vae_latent_mode not in ("presampled", "resample"): + raise ValueError( + f"vae_latent_mode must be 'presampled' or 'resample', " f"got '{self.vae_latent_mode}'" + ) + if self.vae_latent_mode == "resample": + if self.vae_scale is None or self.vae_shift is None: + raise ValueError( + "vae_latent_mode='resample' requires vae_scale and vae_shift " + "to be set (e.g. vae_scale: 0.3611, vae_shift: 0.1159)" + ) + log_rank_0(f"VAE latent mode: resample (reparameterization from mean+logvar each step)") + else: + log_rank_0(f"VAE latent mode: presampled (stored latents used directly)") + + log_rank_0(f"Guidance embedding: {self.use_guidance_embed}") + log_rank_0(f"Scheduler shift: {self.scheduler_shift}") + log_rank_0(f"Dynamic shifting: {self.use_dynamic_shifting}") + + def _init_cfg_dropout(self, params): + """ + Initialize CFG dropout with real empty encodings. + + Resolution order for empty encodings: + 1. Explicit ``empty_encodings_path`` from config + 2. ``{data_path}/empty_encodings/`` (generated by EncodedDatasetPipeline) + 3. ``{data_path}/../empty_encodings/`` (MLPerf convention) + + For mock_data runs, falls back to torch.randn(). + For real data, raises FileNotFoundError if no encodings are found. + """ + tp_size = getattr(params, "tensor_model_parallel_size", 1) + if tp_size != 1: + raise ValueError( + f"CFG dropout requires tensor_model_parallel_size=1 (got {tp_size}). " + "Different TP ranks would generate different dropout masks, causing divergent forward passes." + ) + + context_dim = getattr(params, "context_dim", 4096) + vec_in_dim = getattr(params, "vec_in_dim", 768) + + encodings_dir = self._discover_empty_encodings(params) + + if encodings_dir is not None: + t5_path = os.path.join(encodings_dir, "t5_empty.npy") + clip_path = os.path.join(encodings_dir, "clip_empty.npy") + + self.empty_t5_encodings = torch.from_numpy(np.load(t5_path))[0].unsqueeze(1) + self.empty_clip_encodings = torch.from_numpy(np.load(clip_path))[0] + + log_rank_0( + f"CFG dropout: loaded real empty encodings from {encodings_dir}, " + f"t5={self.empty_t5_encodings.shape}, clip={self.empty_clip_encodings.shape}" + ) + elif getattr(params, "mock_data", False): + image_size = getattr(getattr(params, "mock_dataset", None), "params", None) + image_size = getattr(image_size, "image_size", 256) if image_size is not None else 256 + image_tokens = (image_size // 8 // 2) ** 2 + total_seq = getattr(params, "seq_length", 512) + t5_seq_len = total_seq - image_tokens + + self.empty_t5_encodings = torch.randn(t5_seq_len, 1, context_dim) + self.empty_clip_encodings = torch.randn(vec_in_dim) + log_rank_0("CFG dropout: using torch.randn() empty encodings (mock_data mode)") + else: + data_path = getattr(params, "data_path", "") + if isinstance(data_path, list): + data_path = data_path[0] if data_path else "" + raise FileNotFoundError( + f"CFG dropout requires empty T5/CLIP encodings but none were found.\n" + f"Searched locations:\n" + f" 1. empty_encodings_path config key (not set)\n" + f" 2. {data_path}/empty_encodings/\n" + f" 3. {os.path.dirname(str(data_path))}/empty_encodings/\n\n" + f"To fix, either:\n" + f" - Re-run dataset preparation with 'primus data diffusion-encoded' (auto-generates them)\n" + f" - Run: python tools/generate_empty_encodings.py --output_dir \n" + f" and set empty_encodings_path in your YAML config" + ) + + log_rank_0(f"CFG dropout prob: {self.cfg_dropout_prob}") + + @staticmethod + def _discover_empty_encodings(params) -> "str | None": + """Return the first valid empty_encodings directory, or None.""" + + def _has_files(dirpath): + return ( + os.path.isdir(dirpath) + and os.path.isfile(os.path.join(dirpath, "t5_empty.npy")) + and os.path.isfile(os.path.join(dirpath, "clip_empty.npy")) + ) + + explicit = getattr(params, "empty_encodings_path", None) + if explicit and _has_files(str(explicit)): + return str(explicit) + + data_path = getattr(params, "data_path", None) + if data_path is not None: + if isinstance(data_path, list): + data_path = data_path[0] if data_path else None + if data_path: + data_path = str(data_path) + inside = os.path.join(data_path, "empty_encodings") + if _has_files(inside): + log_rank_0(f"CFG dropout: auto-discovered empty encodings at {inside}") + return inside + alongside = os.path.join(os.path.dirname(data_path), "empty_encodings") + if _has_files(alongside): + log_rank_0(f"CFG dropout: auto-discovered empty encodings at {alongside}") + return alongside + + return None + + def forward_step(self, data_iterator, model, return_schedule_plan=False): + """ + Forward step for Flux diffusion training. + + Overrides base class to provide Flux forward step with scheduler and guidance config. + The base class implementation handles the new return signature correctly, so we just + call super() to use it. + + On the first call, sets a per-DP-rank random seed so that each data-parallel + rank independently samples timesteps and noise (matching NeMo's approach). + + Args: + data_iterator: Data iterator + model: Diffusion model (Flux) + return_schedule_plan: Whether to return schedule plan (for pipeline parallelism) + + Returns: + Tuple of (output_tensor, loss_func_callable) + """ + if not self._training_rng_seeded: + from megatron.core import parallel_state + from megatron.training import get_args + + seed = get_args().seed + dp_rank = parallel_state.get_data_parallel_rank() + per_rank_seed = seed + 100 * dp_rank + torch.manual_seed(per_rank_seed) + torch.cuda.manual_seed(per_rank_seed) + self._training_rng_seeded = True + log_rank_0(f"Per-DP-rank training seed: {per_rank_seed} " f"(base={seed}, dp_rank={dp_rank})") + + return super().forward_step(data_iterator, model, return_schedule_plan=return_schedule_plan) + + def create_model(self, pre_process=True, post_process=True): + """ + Create Flux model from YAML configuration. + + Model architecture is loaded from YAML files: + - flux_12b.yaml / flux_535m.yaml for layer counts + - flux_base.yaml for common Flux architecture + + Optionally loads checkpoint after model creation if configured. + + Args: + pre_process: Not used (kept for Megatron model_provider interface compatibility) + post_process: Not used (kept for Megatron model_provider interface compatibility) + + Returns: + Flux model instance + """ + try: + log_rank_0("=" * 80) + log_rank_0(f"Creating Flux model from YAML config") + + from megatron.training import get_args + + from primus.backends.megatron.core.models.diffusion.flux.model import Flux + + config = self._build_flux_config_from_yaml() + + # get_args() is safe here since create_model() is called after setup() completes + # setup() calls set_primus_global_variables() which initializes _GLOBAL_ARGS + args = get_args() + # Log complete FluxConfig at DEBUG level (matches Megatron's argument logging) + self._log_flux_config(config, args) + + # Set torch_compile settings on args for trainer's apply_torch_compile_if_enabled method + # FluxConfig always has these attributes (defined as dataclass fields) + torch_compile_attrs = [ + "enable_torch_compile", + "torch_compile_backend", + "torch_compile_mode", + "torch_compile_fullgraph", + "torch_compile_optimizer", + "torch_compile_optimizer_scope", + ] + for attr in torch_compile_attrs: + if not hasattr(args, attr): + setattr(args, attr, getattr(config, attr)) + + # Backend selection is handled automatically by Flux model based on config.transformer_impl + # Pass backend=None to let get_flux_layer_spec() handle backend selection + backend = None + + # Log which transformer implementation will be used + if config.transformer_impl == "local": + log_rank_0("Using local transformer implementation (NO TransformerEngine dependency)") + else: + log_rank_0("Using TransformerEngine implementation") + log_rank_0( + "Backend will be selected automatically by Flux model based on config.transformer_impl" + ) + + # Chimera init: replicate NeMo's per-rank seed contamination where + # each DP rank initializes non-parallel weights with a different seed. + # The distributed optimizer merges these into a chimera model. + if self.nemo_chimera_init: + import logging + + from megatron.core import parallel_state + + dp_rank = parallel_state.get_data_parallel_rank() + per_rank_seed = args.seed + 100 * dp_rank + torch.manual_seed(per_rank_seed) + torch.cuda.manual_seed(per_rank_seed) + logging.warning( + "[nemo_chimera_init] WARNING: DP weight invariant intentionally broken. " + "Each rank initializes non-parallel layers with a different seed. " + "This replicates NeMo's per-rank seed contamination for convergence " + "parity experiments only. dp_rank=%d, init_seed=%d (base=%d)", + dp_rank, + per_rank_seed, + args.seed, + ) + + # Create Flux model (backend=None lets model select based on config.transformer_impl) + model = Flux(config=config, backend=backend) + + if self.nemo_chimera_init: + _restore_chimera_rng_state(args) + log_rank_0( + f"[nemo_chimera_init] Restored canonical RNG state (seed={args.seed}) after chimera model init" + ) + + # Calculate parameters for logging + total_params = sum(p.numel() for p in model.parameters()) + + log_rank_0( + f"Flux model created: {config.num_joint_layers} joint + {config.num_single_layers} single layers" + ) + log_rank_0(f"Total parameters: {total_params / 1e6:.1f}M") + + log_rank_0("=" * 80) + + return model + except Exception as e: + log_rank_0(f"[ERROR] create_model() failed: {type(e).__name__}: {e}") + import traceback + + log_rank_0(f"[ERROR] Traceback:\n{traceback.format_exc()}") + raise + + def _build_flux_config_from_yaml(self): + """ + Build FluxConfig from YAML configuration + training args. + + All architectural parameters come from YAML files (merged into Megatron args): + - flux_535m.yaml or flux_12b.yaml specifies layer counts + - flux_base.yaml provides common architecture parameters + + Training precision and recomputation settings come from backend_args. + Torch compile settings come from backend_args.torch_compile (Primus-specific, in overrides section, not in Megatron args). + + Returns: + FluxConfig instance with all parameters + """ + from functools import partial + + import torch.nn.functional as F + + from primus.backends.megatron.core.models.diffusion.flux.config import ( + FluxConfig, + erf_gelu, + ) + + _openai_gelu_fused = partial(F.gelu, approximate="tanh") + _openai_gelu_fused.__name__ = "openai_gelu_fused" + + # backend_args is set by BaseTrainer.__init__() + params = self.backend_args + + # All architectural params from params (merged from YAML) + config_params = { + "num_joint_layers": getattr(params, "num_joint_layers", 19), + "num_single_layers": getattr(params, "num_single_layers", 38), + "hidden_size": getattr(params, "hidden_size", 3072), + "num_attention_heads": getattr(params, "num_attention_heads", 24), + "in_channels": getattr(params, "in_channels", 64), + "context_dim": getattr(params, "context_dim", 4096), + "vec_in_dim": getattr(params, "vec_in_dim", 768), + "model_channels": getattr(params, "model_channels", 256), + "guidance_embed": getattr(params, "guidance_embed", False), + # RoPE configuration + "apply_rope_fusion": getattr(params, "apply_rope_fusion", False), + "rotary_interleaved": getattr(params, "rotary_interleaved", True), + # Training hyperparameters stored on FluxConfig + "timestep_sampling_strategy": getattr(params, "timestep_sampling_strategy", "logit_normal"), + "cfg_dropout_prob": getattr(params, "cfg_dropout_prob", 0.0), + # Weight init for Megatron-managed layers (ColumnParallelLinear, RowParallelLinear). + # Xavier uniform matches NeMo's CustomFluxConfig (MLPerf v5.1). + "init_method": nn.init.xavier_uniform_, + "output_layer_init_method": nn.init.xavier_uniform_, + } + + # Activation: YAML "openai_gelu" maps to fused F.gelu(approximate="tanh"); default in + # FluxConfig is the non-JIT Python GELU (ROCm-safe). Both match the tanh GELU approx. + activation_func_name = getattr(params, "activation_func", None) + if activation_func_name is not None: + activation_func_map = { + "erf_gelu": erf_gelu, + "openai_gelu": _openai_gelu_fused, + } + if activation_func_name not in activation_func_map: + raise ValueError( + f"Unknown activation_func '{activation_func_name}'. " + f"Choose from: {list(activation_func_map.keys())}" + ) + func = activation_func_map[activation_func_name] + config_params["activation_func"] = func + log_rank_0(f"Activation function: {activation_func_name} -> {func.__name__}") + + # Transformer implementation (standard TransformerConfig parameter) + # Read from backend_args (overrides section) + transformer_impl = getattr(params, "transformer_impl", "transformer_engine") + config_params["transformer_impl"] = transformer_impl + + # Precision settings from params (not in model YAML) + config_params.update( + { + "bf16": getattr(params, "bf16", True), + "fp16": getattr(params, "fp16", False), + "params_dtype": getattr(params, "params_dtype", torch.float32), + } + ) + + # FP8 settings from params + config_params.update( + { + "fp8": getattr(params, "fp8", None), + "fp8_recipe": getattr(params, "fp8_recipe", "delayed"), + "fp8_margin": getattr(params, "fp8_margin", 0), + "fp8_amax_history_len": getattr(params, "fp8_amax_history_len", 1), + "fp8_amax_compute_algo": getattr(params, "fp8_amax_compute_algo", "most_recent"), + "fp8_wgrad": getattr(params, "fp8_wgrad", True), + "fp8_dot_product_attention": getattr(params, "fp8_dot_product_attention", False), + "fp8_multi_head_attention": getattr(params, "fp8_multi_head_attention", False), + "fp8_scaling_strategy": getattr(params, "fp8_scaling_strategy", "dynamic"), + "fp8_force_nt_layout": getattr(params, "fp8_force_nt_layout", False), + "fp8_reduce_amax": getattr(params, "fp8_reduce_amax", False), + } + ) + + # FP4/MXFP4 settings + fp4_enabled = getattr(params, "fp4", None) + fp4_recipe = getattr(params, "fp4_recipe", None) + if fp4_enabled and not fp4_recipe: + raise ValueError( + "fp4_recipe must be explicitly set in YAML when fp4 is enabled. " + "Use fp4_recipe: 'mxfp4' for AMD (native FP4 GEMM) or 'nvfp4' for NVIDIA." + ) + config_params.update( + { + "fp4": fp4_enabled, + "fp4_recipe": fp4_recipe, + "mxfp4_backward_precision": getattr(params, "mxfp4_backward_precision", "mxfp4"), + } + ) + + # Sensitive layer configuration + config_params.update( + { + "sensitive_layers_enabled": getattr(params, "sensitive_layers_enabled", False), + "sensitive_layers_start": getattr(params, "sensitive_layers_start", 0), + "sensitive_layers_end": getattr(params, "sensitive_layers_end", 0), + "sensitive_layer_precision": getattr(params, "sensitive_layer_precision", "bf16"), + "mxfp4_gradient_stochastic_rounding": getattr( + params, "mxfp4_gradient_stochastic_rounding", False + ), + } + ) + + config_params["use_dual_fp8_output_projection"] = getattr( + params, + "use_dual_fp8_output_projection", + False, + ) + + if ( + config_params["use_dual_fp8_output_projection"] + and config_params.get("fp8_scaling_strategy") == "delayed" + ): + raise ValueError( + "use_dual_fp8_output_projection=True is incompatible with " + "fp8_scaling_strategy='delayed'. DualFP8LinearTensorwiseFunction " + "bypasses delayed-scaling staged amax buffers, causing stale/zero " + "values in the _DelayedScalingRegistry. Use " + "fp8_scaling_strategy='dynamic' or set " + "use_dual_fp8_output_projection=False." + ) + + config_params["use_triton_ops"] = getattr( + params, + "use_triton_ops", + False, + ) + config_params["adaln_plain_ops"] = getattr( + params, + "adaln_plain_ops", + False, + ) + config_params["adaln_always_jit_fuser"] = getattr( + params, + "adaln_always_jit_fuser", + False, + ) + config_params["optimizer_foreach"] = getattr( + params, + "optimizer_foreach", + True, + ) + config_params["overlap_grad_norm"] = getattr( + params, + "overlap_grad_norm", + False, + ) + config_params["use_cpp_fp8_quantize"] = getattr( + params, + "use_cpp_fp8_quantize", + False, + ) + + # FSDP2 prefetch depth + fsdp_prefetch = getattr(params, "fsdp_prefetch_depth", None) + if fsdp_prefetch is not None: + config_params["fsdp_prefetch_depth"] = int(fsdp_prefetch) + + # FSDP2 FP32 optimizer: initialize model in FP32, FSDP2 casts to BF16 + if getattr(params, "use_fsdp2_fp32_param_optimizer", False): + config_params["params_dtype"] = torch.float32 + config_params["pipeline_dtype"] = torch.bfloat16 + log_rank_0("FSDP2 FP32 optimizer: params_dtype=FP32, pipeline=BF16") + + # Recomputation settings from params + config_params.update( + { + "recompute_granularity": getattr(params, "recompute_granularity", None), + "recompute_method": getattr(params, "recompute_method", None), + "recompute_num_layers": getattr(params, "recompute_num_layers", None), + "recompute_modules": getattr(params, "recompute_modules", None), + } + ) + + # Torch compile settings from backend_args (overrides section) + # torch_compile is in overrides section in YAML, accessible via backend_args.torch_compile + torch_compile_config = getattr(self.backend_args, "torch_compile", None) + + if torch_compile_config is not None: + compile_settings = { + "enable_torch_compile": getattr(torch_compile_config, "enable", False), + "torch_compile_backend": getattr(torch_compile_config, "backend", "inductor"), + "torch_compile_mode": getattr(torch_compile_config, "mode", "default"), + "torch_compile_fullgraph": getattr(torch_compile_config, "fullgraph", False), + "torch_compile_optimizer": getattr(torch_compile_config, "compile_optimizer", False), + "torch_compile_optimizer_scope": getattr( + torch_compile_config, "compile_optimizer_scope", "full" + ), + "torch_compile_strategy": getattr(torch_compile_config, "strategy", "per_block"), + "torch_compile_replace_qk_rmsnorm": getattr( + torch_compile_config, "replace_qk_rmsnorm", False + ), + "torch_compile_disable_inductor_cudagraphs": getattr( + torch_compile_config, "disable_inductor_cudagraphs", True + ), + "torch_compile_emulate_precision_casts": getattr( + torch_compile_config, "emulate_precision_casts", True + ), + "torch_compile_fused_ln_modulate": getattr(torch_compile_config, "fused_ln_modulate", True), + } + else: + # Default values if torch_compile section not present + compile_settings = { + "enable_torch_compile": False, + "torch_compile_backend": "inductor", + "torch_compile_mode": "default", + "torch_compile_fullgraph": False, + "torch_compile_optimizer": False, + "torch_compile_optimizer_scope": "full", + "torch_compile_strategy": "per_block", + "torch_compile_replace_qk_rmsnorm": False, + "torch_compile_disable_inductor_cudagraphs": True, + "torch_compile_emulate_precision_casts": True, + "torch_compile_fused_ln_modulate": True, + } + + # Set on FluxConfig + config_params.update(compile_settings) + + return FluxConfig(**config_params) + + def _log_flux_config(self, config, args): + """Print FluxConfig fields (rank 0 only), similar to Megatron argument dumps.""" + import dataclasses + + # Only log on rank 0 + if args.rank != 0: + return + + # Accumulate the full dump into one multi-line message so the banner + # alignment is preserved (log_rank_0 stamps a caller prefix per call). + lines = [] + lines.append("=" * 80) + lines.append("FluxConfig (Model Configuration)") + lines.append("=" * 80) + + # Organize config parameters by category for better readability + categories = { + "Model Architecture": [ + "model_type", + "num_joint_layers", + "num_single_layers", + "num_layers", + "hidden_size", + "num_attention_heads", + "ffn_hidden_size", + ], + "Input/Output Dimensions": [ + "in_channels", + "out_channels", + "patch_size", + "context_dim", + "vec_in_dim", + "model_channels", + ], + "Position Embeddings (RoPE)": ["theta", "axes_dim", "rotary_interleaved", "apply_rope_fusion"], + "Guidance & Diffusion": [ + "guidance_embed", + "guidance_scale", + "cfg_dropout_prob", + "timestep_sampling_strategy", + ], + "Attention & Layers": [ + "add_qkv_bias", + "single_block_bias", + "attention_dropout", + "hidden_dropout", + "bias_dropout_fusion", + ], + "Normalization & Activation": ["activation_func", "layernorm_epsilon", "normalization"], + "Precision & Optimization": [ + "bf16", + "fp16", + "fp32_residual_connection", + "gradient_accumulation_fusion", + "use_dual_fp8_output_projection", + "params_dtype", + "fp8", + "fp8_recipe", + "fp8_scaling_strategy", + "fp8_force_nt_layout", + "fp4", + "fp4_recipe", + "mxfp4_backward_precision", + "mxfp4_gradient_stochastic_rounding", + "sensitive_layers_enabled", + "sensitive_layers_start", + "sensitive_layers_end", + "sensitive_layer_precision", + ], + "Recomputation": [ + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + "recompute_modules", + ], + "Torch Compile": [ + "enable_torch_compile", + "torch_compile_backend", + "torch_compile_mode", + "torch_compile_fullgraph", + "torch_compile_optimizer", + "torch_compile_optimizer_scope", + "torch_compile_strategy", + "torch_compile_replace_qk_rmsnorm", + "torch_compile_disable_inductor_cudagraphs", + "torch_compile_emulate_precision_casts", + "torch_compile_fused_ln_modulate", + ], + "Parallelism": [ + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "sequence_parallel", + "expert_model_parallel_size", + ], + "CUDA Graph": ["enable_cuda_graph", "cuda_graph_scope", "cuda_graph_warmup_steps"], + "Transformer Engine": ["use_te_rng_tracker"], + } + + # Get all dataclass fields + all_fields = {f.name for f in dataclasses.fields(config)} + categorized_fields = set() + + # Print categorized fields + for category, field_names in categories.items(): + # Check if any fields in this category exist + existing_fields = [f for f in field_names if f in all_fields] + if not existing_fields: + continue + + lines.append(f"\n{category}:") + for field_name in existing_fields: + if hasattr(config, field_name): + value = getattr(config, field_name) + # Format value (handle callables specially) + if callable(value) and not isinstance(value, type): + value_str = getattr(value, "__name__", str(value)) + else: + value_str = str(value) + + # Create dots for alignment (match Megatron's style: 48 chars) + dots = "." * (48 - len(field_name)) + lines.append(f" {field_name} {dots} {value_str}") + categorized_fields.add(field_name) + + # Print uncategorized fields (any fields not in our category lists) + uncategorized = all_fields - categorized_fields + if uncategorized: + lines.append("\nOther Configuration:") + for field_name in sorted(uncategorized): + if hasattr(config, field_name): + value = getattr(config, field_name) + if callable(value) and not isinstance(value, type): + value_str = getattr(value, "__name__", str(value)) + else: + value_str = str(value) + dots = "." * (48 - len(field_name)) + lines.append(f" {field_name} {dots} {value_str}") + + # Add special note about apply_rope_fusion mismatch + if hasattr(config, "apply_rope_fusion") and hasattr(args, "apply_rope_fusion"): + if config.apply_rope_fusion != args.apply_rope_fusion: + lines.append("\n" + "!" * 80) + lines.append("IMPORTANT: FluxConfig vs Megatron Args Mismatch") + lines.append("!" * 80) + lines.append(f" FluxConfig.apply_rope_fusion: {config.apply_rope_fusion}") + lines.append(f" Megatron args.apply_rope_fusion: {args.apply_rope_fusion}") + lines.append( + f' → Flux uses FluxConfig value: RoPE fusion is {"ENABLED" if config.apply_rope_fusion else "DISABLED"}' + ) + lines.append(" → Megatron args value is set by position_embedding_type validation") + lines.append( + ' → (Megatron sets apply_rope_fusion=False when position_embedding_type != "rope")' + ) + lines.append("!" * 80) + + lines.append("=" * 80) + lines.append("End of FluxConfig") + lines.append("=" * 80) + + log_rank_0("\n".join(lines)) + + def create_scheduler(self): + """ + Create Flow Matching Euler Discrete Scheduler for Flux. + + Returns: + FlowMatchEulerDiscreteScheduler instance + """ + log_rank_0("Creating Flow Matching scheduler...") + log_rank_0(f" num_train_timesteps: {self.num_train_timesteps}") + log_rank_0(f" shift: {self.scheduler_shift}") + log_rank_0(f" use_dynamic_shifting: {self.use_dynamic_shifting}") + + scheduler = FlowMatchEulerDiscreteScheduler( + num_train_timesteps=self.num_train_timesteps, + shift=self.scheduler_shift, + use_dynamic_shifting=self.use_dynamic_shifting, + base_shift=0.5, # Flux defaults + max_shift=1.15, + base_image_seq_len=256, + max_image_seq_len=4096, + ) + + return scheduler + + def get_task_encoder(self): + """ + Get Flux task encoder for Energon data pipeline. + + Returns: + EncodedDiffusionTaskEncoder for pre-encoded data + """ + from primus.backends.megatron.data.diffusion.task_encoders import ( + EncodedDiffusionTaskEncoder, + ) + + # EnergonDatasetProvider will pass WorkerConfig to Energon directly + # TaskEncoder doesn't need WorkerConfig for pre-encoded data + task_encoder = EncodedDiffusionTaskEncoder(worker_config=None) + + log_rank_0("Created EncodedDiffusionTaskEncoder for pre-encoded Flux data") + return task_encoder diff --git a/primus/backends/megatron/training/evaluator.py b/primus/backends/megatron/training/evaluator.py index 68da99d69..9ea10e438 100644 --- a/primus/backends/megatron/training/evaluator.py +++ b/primus/backends/megatron/training/evaluator.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -137,9 +137,21 @@ def primus_evaluate( log_rank_0("Exiting during evaluation, timelimit reached") return None, None, True - # Compute final average loss across all eval iterations + # DP all-reduce for tuple-path (validation) metrics so that every + # rank sees the same globally-averaged loss. Scalar/legacy metrics + # are NOT all-reduced, matching upstream Megatron's evaluate(). total_loss_dict = {} if is_pipeline_stage_containing_loss(): + from megatron.core import mpu + + dp_group = mpu.get_data_parallel_group(with_context_parallel=True) + for key in total_loss_numerators.keys(): + num = total_loss_numerators[key] + den = total_loss_denominators[key] + if isinstance(num, torch.Tensor) and isinstance(den, torch.Tensor): + torch.distributed.all_reduce(num, group=dp_group) + torch.distributed.all_reduce(den, group=dp_group) + for key in total_loss_numerators.keys(): # Reduce numerator/denominator across data-parallel ranks so the # validation loss is a TRUE global average, identical on every rank. diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py b/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py new file mode 100644 index 000000000..487af4193 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py @@ -0,0 +1,343 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for DiffusionPretrainTrainer. + +Tests initialization, setup, data provider delegation, forward step, +and scheduler lazy initialization. +""" + +import sys +import types +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer + + +def _build_diffusion_trainer(monkeypatch: pytest.MonkeyPatch, backend_args=None, use_mock_data=True): + """Helper to build DiffusionPretrainTrainer with stubbed dependencies.""" + + # Create a concrete subclass for testing (DiffusionPretrainTrainer is abstract) + class ConcreteDiffusionTrainer(DiffusionPretrainTrainer): + def create_model(self, pre_process=True, post_process=True): + return Mock() + + def create_scheduler(self): + return Mock() + + def get_task_encoder(self): + return Mock() + + # Stub out MegatronBaseTrainer.__init__ to avoid real Megatron imports + def dummy_base_init(self, backend_args=None, *args, **kwargs): + self.backend_args = backend_args + self.rank = 0 + self.world_size = 1 + self.local_rank = 0 + self.master_addr = "localhost" + self.master_port = 12345 + + monkeypatch.setattr( + "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.__init__", + dummy_base_init, + ) + + # Silence logging + monkeypatch.setattr( + "primus.backends.megatron.diffusion_trainer.log_rank_0", + lambda *args, **kwargs: None, + ) + + if backend_args is None: + backend_args = SimpleNamespace(mock_data=use_mock_data) + else: + if not hasattr(backend_args, "mock_data"): + backend_args.mock_data = use_mock_data + + return ConcreteDiffusionTrainer(backend_args=backend_args) + + +class TestDiffusionPretrainTrainer: + """Tests for DiffusionPretrainTrainer.""" + + def test_init_with_mock_data_creates_synthetic_provider(self, monkeypatch: pytest.MonkeyPatch): + """Test that initialization with mock_data=True creates SyntheticDatasetProvider.""" + backend_args = SimpleNamespace(mock_data=True, model_type="flux") + trainer = _build_diffusion_trainer(monkeypatch, backend_args) + + assert hasattr(trainer, "data_provider") + assert trainer.data_provider is not None + # Verify it's SyntheticDatasetProvider (check class name) + assert "Synthetic" in type(trainer.data_provider).__name__ + + def test_setup_calculates_data_parallel_size(self, monkeypatch: pytest.MonkeyPatch): + """Test that setup() calculates data_parallel_size when not present.""" + backend_args = SimpleNamespace( + mock_data=True, + world_size=8, + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + context_parallel_size=1, + ) + trainer = _build_diffusion_trainer(monkeypatch, backend_args) + + # Mock parent setup and other dependencies + setup_calls = [] + + def mock_parent_setup(self): + setup_calls.append("parent_setup") + + monkeypatch.setattr( + "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.setup", + mock_parent_setup, + ) + + # Mock set_primus_global_variables + monkeypatch.setattr( + "primus.backends.megatron.training.global_vars.set_primus_global_variables", + lambda args: None, + ) + + trainer.setup() + + # Verify data_parallel_size was calculated + assert hasattr(trainer.backend_args, "data_parallel_size") + assert trainer.backend_args.data_parallel_size == 4 # 8 / (2 * 1 * 1) + assert "parent_setup" in setup_calls + + def test_setup_uses_existing_data_parallel_size(self, monkeypatch: pytest.MonkeyPatch): + """Test that setup() uses existing data_parallel_size if present.""" + backend_args = SimpleNamespace( + mock_data=True, + data_parallel_size=16, + world_size=8, # This would normally calculate to 4, but we have existing value + ) + trainer = _build_diffusion_trainer(monkeypatch, backend_args) + + # Mock parent setup + monkeypatch.setattr( + "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.setup", + lambda self: None, + ) + monkeypatch.setattr( + "primus.backends.megatron.training.global_vars.set_primus_global_variables", + lambda args: None, + ) + + trainer.setup() + + # Verify existing data_parallel_size was not overwritten + assert trainer.backend_args.data_parallel_size == 16 + + def test_setup_sets_model_provider(self, monkeypatch: pytest.MonkeyPatch): + """Test that setup() sets model_provider with correct signature.""" + trainer = _build_diffusion_trainer(monkeypatch) + + # Mock create_model + mock_model = Mock() + trainer.create_model = lambda pre_process=True, post_process=True: mock_model + + # Mock parent setup + monkeypatch.setattr( + "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.setup", + lambda self: None, + ) + monkeypatch.setattr( + "primus.backends.megatron.training.global_vars.set_primus_global_variables", + lambda args: None, + ) + + trainer.setup() + + # Verify model_provider was set + assert hasattr(trainer, "model_provider") + assert callable(trainer.model_provider) + + # Verify model_provider signature matches Megatron's interface + result = trainer.model_provider( + pre_process=True, + post_process=True, + vp_stage=None, + config=None, + pg_collection=None, + ) + assert result is mock_model + + def test_get_datasets_provider_returns_delegating_function(self, monkeypatch: pytest.MonkeyPatch): + """Test that get_datasets_provider() returns a function that delegates to data_provider.""" + trainer = _build_diffusion_trainer(monkeypatch) + + # Replace data_provider with a mock that has is_distributed as an attribute (not property) + mock_dataloaders = [Mock(), Mock(), Mock()] # train, val, test + + class MockDataProvider: + def __init__(self): + self.is_distributed = True + self.create_dataloaders = Mock(return_value=mock_dataloaders) + + mock_data_provider = MockDataProvider() + trainer.data_provider = mock_data_provider + + # Mock get_args + mock_args = SimpleNamespace() + training_mod = types.SimpleNamespace(get_args=lambda: mock_args) + monkeypatch.setitem(sys.modules, "megatron.training", training_mod) + + provider_func = trainer.get_datasets_provider() + + # Verify it's a function + assert callable(provider_func) + + # Verify is_distributed flag is set + assert provider_func.is_distributed is True + + # Call the provider function + train_val_test_num_samples = [100, 10, 10] + result = provider_func(train_val_test_num_samples, vp_stage=None) + + # Verify data_provider.create_dataloaders was called + mock_data_provider.create_dataloaders.assert_called_once_with( + trainer_config=mock_args, + train_val_test_num_samples=train_val_test_num_samples, + vp_stage=None, + ) + + # Verify result matches what data_provider returned + assert result is mock_dataloaders + + def test_forward_step_calls_flux_forward_step_func(self, monkeypatch: pytest.MonkeyPatch): + """Test that forward_step() calls flux_forward_step_func and creates loss function.""" + trainer = _build_diffusion_trainer(monkeypatch) + + # Mock scheduler + mock_scheduler = Mock() + trainer._scheduler = mock_scheduler + + # Mock flux_forward_step_func + mock_noise_pred = Mock() + mock_clean_latents = Mock() + mock_noise = Mock() + mock_loss_mask = Mock() + mock_metrics = {"test_metric": 1.0} + + flux_forward_step_calls = [] + + def mock_flux_forward_step( + data_iterator, + model, + scheduler=None, + use_guidance_embed=False, + guidance_scale=None, + timestep_sampler=None, + cfg_dropout_prob=0.0, + empty_t5_encodings=None, + empty_clip_encodings=None, + **kwargs, + ): + flux_forward_step_calls.append( + (data_iterator, model, scheduler, use_guidance_embed, guidance_scale) + ) + # 6-tuple per sibling commit 362ee36 (is_validation appended). + return ( + mock_noise_pred, + mock_clean_latents, + mock_noise, + mock_loss_mask, + mock_metrics, + False, + ) + + # Patch at the import location (it's imported inside forward_step) + monkeypatch.setattr( + "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func", + mock_flux_forward_step, + ) + + # Mock runtime_state + trainer.runtime_state = Mock() + trainer.runtime_state.update_metrics = Mock() + + mock_data_iterator = Mock() + mock_model = Mock() + + output, loss_func = trainer.forward_step(mock_data_iterator, mock_model) + + # Verify flux_forward_step_func was called + assert len(flux_forward_step_calls) == 1 + assert flux_forward_step_calls[0][0] is mock_data_iterator + assert flux_forward_step_calls[0][1] is mock_model + assert flux_forward_step_calls[0][2] is mock_scheduler + + # Verify values were stored + assert trainer._last_clean_latents is mock_clean_latents + assert trainer._last_noise is mock_noise + assert trainer._last_loss_mask is mock_loss_mask + + # Verify metrics were updated + trainer.runtime_state.update_metrics.assert_called_once_with(mock_metrics) + + # Verify output is noise_pred + assert output is mock_noise_pred + + # Verify loss_func is callable + assert callable(loss_func) + + # Test loss function - mock compute_flow_matching_loss to avoid real computation + # It's imported inside the loss function, so patch at the import location + mock_loss = Mock() + mock_loss.detach.return_value = mock_loss + monkeypatch.setattr( + "primus.backends.megatron.training.diffusion.loss_computation.compute_flow_matching_loss", + lambda *args, **kwargs: mock_loss, + ) + + loss_result = loss_func(mock_noise_pred, non_loss_data=False) + assert len(loss_result) == 2 # (loss, metrics_dict) + # Verify the metrics dict has the expected key + assert "reduced_train_loss" in loss_result[1] + + def test_forward_step_loss_func_with_non_loss_data(self, monkeypatch: pytest.MonkeyPatch): + """Test that loss function returns output_tensor when non_loss_data=True.""" + trainer = _build_diffusion_trainer(monkeypatch) + + # Mock forward_step components + mock_scheduler = Mock() + trainer._scheduler = mock_scheduler + + mock_noise_pred = Mock() + mock_clean_latents = Mock() + mock_noise = Mock() + mock_loss_mask = Mock() + mock_metrics = {} + + # Patch at the import location (it's imported inside forward_step). + # 6-tuple per sibling commit 362ee36 (is_validation appended). + monkeypatch.setattr( + "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func", + lambda *args, **kwargs: ( + mock_noise_pred, + mock_clean_latents, + mock_noise, + mock_loss_mask, + mock_metrics, + False, + ), + ) + + trainer.runtime_state = Mock() + trainer.runtime_state.update_metrics = Mock() + + output, loss_func = trainer.forward_step(Mock(), Mock()) + + # Test loss function with non_loss_data=True + result = loss_func(output, non_loss_data=True) + + # Should return output_tensor directly + assert result is output diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py new file mode 100644 index 000000000..88c440b10 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_forward_step_e2e.py @@ -0,0 +1,274 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +End-to-end tests for flux_forward_step_func. + +Tests the full forward step with a real Flux 535M model on CUDA, +covering presampled, resample, validation, and CFG dropout paths. +""" + +import contextlib +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from primus.backends.megatron.training.diffusion.forward_step import ( + flux_forward_step_func, +) +from primus.backends.megatron.training.diffusion.schedulers.flow_matching import ( + FlowMatchEulerDiscreteScheduler, +) +from tests.unit_tests.backends.megatron.diffusion.constants import ( + CLIP_L_EMBEDDING_DIM, + T5_XXL_EMBEDDING_DIM, + VAE_LATENT_CHANNELS, +) + + +def _patch_parallel_state(): + """Return a stack of mock.patch context managers for parallel state.""" + return [ + patch( + "megatron.core.parallel_state.get_tensor_model_parallel_world_size", + return_value=1, + ), + patch( + "megatron.core.parallel_state.get_pipeline_model_parallel_world_size", + return_value=1, + ), + patch( + "megatron.core.parallel_state.get_data_parallel_rank", + return_value=0, + ), + patch( + "megatron.core.parallel_state.get_tensor_model_parallel_rank", + return_value=0, + ), + patch( + "megatron.training.get_args", + return_value=SimpleNamespace(seed=42), + ), + ] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +class TestFluxForwardStepE2E: + """End-to-end tests for flux_forward_step_func with real model.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state for model tests.""" + + @pytest.fixture + def model(self): + config = FluxConfig.flux_535m() + m = Flux(config).cuda() + m.train() + return m + + @pytest.fixture + def scheduler(self): + return FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000) + + def _make_presampled_batch(self, batch_size=2, height=16, width=16): + return { + "latents": torch.randn(batch_size, VAE_LATENT_CHANNELS, height, width), + "prompt_embeds": torch.randn(batch_size, 32, T5_XXL_EMBEDDING_DIM), + "pooled_prompt_embeds": torch.randn(batch_size, CLIP_L_EMBEDDING_DIM), + } + + def _make_resample_batch(self, batch_size=2, height=16, width=16): + return { + "mean": torch.randn(batch_size, VAE_LATENT_CHANNELS, height, width), + "logvar": torch.randn(batch_size, VAE_LATENT_CHANNELS, height, width), + "prompt_embeds": torch.randn(batch_size, 32, T5_XXL_EMBEDDING_DIM), + "pooled_prompt_embeds": torch.randn(batch_size, CLIP_L_EMBEDDING_DIM), + } + + def test_presampled_path(self, model, scheduler): + """Test forward step with pre-encoded latents.""" + batch = self._make_presampled_batch() + data_iterator = iter([batch]) + + with contextlib.ExitStack() as stack: + for p in _patch_parallel_state(): + stack.enter_context(p) + result = flux_forward_step_func( + data_iterator, + model, + scheduler=scheduler, + step_count=1, + ) + + noise_pred, clean_latents, noise, loss_mask, metrics, is_validation = result + assert noise_pred.shape == (2, VAE_LATENT_CHANNELS, 16, 16) + assert clean_latents.shape == noise_pred.shape + assert noise.shape == noise_pred.shape + assert loss_mask is None + assert not is_validation + assert metrics["batch_size"] == 2 + assert metrics["latent_channels"] == VAE_LATENT_CHANNELS + assert metrics["image_height"] == 16 * 8 + assert metrics["image_width"] == 16 * 8 + + def test_resample_path(self, model, scheduler): + """Test forward step with VAE resample mode.""" + batch = self._make_resample_batch() + # Snapshot mean before the call (the forward step casts and moves it). + original_mean = batch["mean"].clone() + data_iterator = iter([batch]) + + vae_scale = 0.3611 + vae_shift = 0.1159 + + with contextlib.ExitStack() as stack: + for p in _patch_parallel_state(): + stack.enter_context(p) + result = flux_forward_step_func( + data_iterator, + model, + scheduler=scheduler, + vae_latent_mode="resample", + vae_scale=vae_scale, + vae_shift=vae_shift, + step_count=1, + ) + + noise_pred, clean_latents, _, _, _, is_validation = result + assert noise_pred.shape == (2, VAE_LATENT_CHANNELS, 16, 16) + assert not is_validation + + # Verify reparameterization actually fired: clean_latents must NOT + # equal the deterministic scale/shift on `mean` alone (which would + # be the result if vae_eps had been dropped or zeroed). + expected_no_eps = vae_scale * (original_mean.cuda() - vae_shift) + assert not torch.allclose(clean_latents.float(), expected_no_eps.float(), atol=1e-3), ( + "clean_latents matches mean*(scale-shift) — reparameterization " + "noise term `eps * std` appears to have been dropped" + ) + + def test_validation_with_timestep_key(self, model, scheduler): + """Batch with 'timestep' key triggers validation mode.""" + batch = self._make_presampled_batch() + batch["timestep"] = torch.arange(2) + data_iterator = iter([batch]) + + with contextlib.ExitStack() as stack: + for p in _patch_parallel_state(): + stack.enter_context(p) + result = flux_forward_step_func( + data_iterator, + model, + scheduler=scheduler, + step_count=1, + ) + + _, _, _, _, _, is_validation = result + assert is_validation is True + # The forward step writes derived timesteps (timestep / 8.0) into the batch. + assert "timesteps" in batch + assert torch.equal( + batch["timesteps"].float().cpu(), + torch.arange(2).float() / 8.0, + ) + + def test_validation_equidistant_injection(self, model, scheduler): + """model.eval() without timestep key injects equidistant timesteps.""" + model.eval() + batch_size = 8 + batch = self._make_presampled_batch(batch_size=batch_size) + data_iterator = iter([batch]) + + with contextlib.ExitStack() as stack: + for p in _patch_parallel_state(): + stack.enter_context(p) + result = flux_forward_step_func( + data_iterator, + model, + scheduler=scheduler, + step_count=1, + ) + + _, _, _, _, _, is_validation = result + assert is_validation is True + # Equidistant injection: batch["timestep"] = arange(B) % 8, + # batch["timesteps"] = that / 8.0. + assert "timestep" in batch + assert "timesteps" in batch + expected_timestep = torch.arange(batch_size, device="cuda") % 8 + assert torch.equal(batch["timestep"], expected_timestep), ( + f"batch['timestep'] expected {expected_timestep.tolist()}, " f"got {batch['timestep'].tolist()}" + ) + assert torch.allclose( + batch["timesteps"].float(), + expected_timestep.float() / 8.0, + ) + + def test_cfg_dropout_full(self, model, scheduler): + """cfg_dropout_prob=1.0 replaces all text embeddings with empty encodings.""" + batch_size = 2 + seq_len = 32 + batch = self._make_presampled_batch(batch_size=batch_size) + data_iterator = iter([batch]) + + # Use distinguishable empty encodings (not zeros) so we can verify + # the dropout actually substituted them rather than incidentally + # matching the original prompt_embeds. + empty_t5_value = 7.5 + empty_clip_value = -3.25 + empty_t5 = torch.full((seq_len, 1, T5_XXL_EMBEDDING_DIM), empty_t5_value, dtype=torch.float32) + empty_clip = torch.full((CLIP_L_EMBEDDING_DIM,), empty_clip_value, dtype=torch.float32) + + # Capture the inputs the model receives by patching its forward. + captured = {} + original_forward = model.forward + + def capturing_forward(*args, **kwargs): + captured["txt"] = kwargs.get("txt").detach().clone() + captured["y"] = kwargs.get("y").detach().clone() + return original_forward(*args, **kwargs) + + model.forward = capturing_forward + + try: + with contextlib.ExitStack() as stack: + for p in _patch_parallel_state(): + stack.enter_context(p) + result = flux_forward_step_func( + data_iterator, + model, + scheduler=scheduler, + cfg_dropout_prob=1.0, + empty_t5_encodings=empty_t5, + empty_clip_encodings=empty_clip, + step_count=1, + ) + finally: + model.forward = original_forward + + noise_pred, _, _, _, _, is_validation = result + assert noise_pred.shape == (batch_size, VAE_LATENT_CHANNELS, 16, 16) + assert not is_validation + + # txt is in (S, B, C) layout after the transpose. + # With prob=1.0 every row must be the empty encoding. + txt = captured["txt"].float() + assert torch.allclose(txt, torch.full_like(txt, empty_t5_value), atol=1e-2), ( + f"prompt_embeds was not replaced by empty_t5 with cfg_dropout_prob=1.0; " + f"sample value={txt.flatten()[0].item():.3f}, expected {empty_t5_value}" + ) + + y = captured["y"].float() + assert torch.allclose(y, torch.full_like(y, empty_clip_value), atol=1e-2), ( + f"pooled_prompt_embeds was not replaced by empty_clip with prob=1.0; " + f"sample value={y.flatten()[0].item():.3f}, expected {empty_clip_value}" + ) diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py new file mode 100644 index 000000000..9a965ccfc --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_model_creation.py @@ -0,0 +1,288 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for FluxPretrainTrainer model creation. + +Tests high-value Flux behavior in Primus: +- backend_args -> FluxConfig mapping/defaults +- model construction wiring +- torch_compile settings propagation semantics +""" + +import sys +import types +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.flux_pretrain_trainer import FluxPretrainTrainer + + +def _build_flux_trainer(monkeypatch: pytest.MonkeyPatch, backend_args=None): + """Helper to build FluxPretrainTrainer with stubbed dependencies.""" + + # Stub out MegatronBaseTrainer.__init__ to avoid real Megatron imports + def dummy_base_init(self, backend_args=None, *args, **kwargs): + self.backend_args = backend_args + self.rank = 0 + self.world_size = 1 + self.local_rank = 0 + self.master_addr = "localhost" + self.master_port = 12345 + + monkeypatch.setattr( + "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.__init__", + dummy_base_init, + ) + + # Stub out DiffusionPretrainTrainer.__init__ + def dummy_diffusion_init(self, *args, **kwargs): + self.backend_args = kwargs.get("backend_args") + if self.backend_args is None and args: + self.backend_args = args[0] + if self.backend_args is None: + self.backend_args = SimpleNamespace() + self._scheduler = None + self.data_provider = Mock() + + monkeypatch.setattr( + "primus.backends.megatron.diffusion_trainer.DiffusionPretrainTrainer.__init__", + dummy_diffusion_init, + ) + + # Silence logging + monkeypatch.setattr( + "primus.backends.megatron.flux_pretrain_trainer.log_rank_0", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "primus.backends.megatron.diffusion_trainer.log_rank_0", + lambda *args, **kwargs: None, + ) + + if backend_args is None: + backend_args = SimpleNamespace( + mock_data=True, + guidance_embed=False, + guidance_scale=3.5, + num_train_timesteps=1000, + scheduler_shift=1.0, + use_dynamic_shifting=False, + ) + + return FluxPretrainTrainer(backend_args=backend_args) + + +class TestFluxModelCreation: + """Tests for FluxPretrainTrainer model creation.""" + + def test_create_model_calls_build_flux_config(self, monkeypatch: pytest.MonkeyPatch): + """Test that create_model() calls _build_flux_config_from_yaml().""" + trainer = _build_flux_trainer(monkeypatch) + + build_config_calls = [] + built_config = None + + def tracked_build_config(self): + nonlocal built_config + build_config_calls.append(1) + from primus.backends.megatron.core.models.diffusion.flux.config import ( + FluxConfig, + ) + + built_config = FluxConfig.flux_535m() + return built_config + + trainer._build_flux_config_from_yaml = tracked_build_config.__get__(trainer, type(trainer)) + + # Mock Flux model + mock_flux_model = Mock() + mock_flux_model.parameters.return_value = [] + flux_ctor = Mock(return_value=mock_flux_model) + monkeypatch.setattr( + "primus.backends.megatron.core.models.diffusion.flux.model.Flux", + flux_ctor, + ) + + # Mock get_args + mock_args = SimpleNamespace(rank=0) + training_mod = types.SimpleNamespace(get_args=lambda: mock_args) + monkeypatch.setitem(sys.modules, "megatron.training", training_mod) + + result = trainer.create_model() + + # Verify _build_flux_config_from_yaml was called + assert len(build_config_calls) == 1 + flux_ctor.assert_called_once() + assert flux_ctor.call_args.kwargs["backend"] is None + assert flux_ctor.call_args.kwargs["config"] is built_config + assert result is mock_flux_model + + def test_build_flux_config_from_yaml_extracts_parameters(self, monkeypatch: pytest.MonkeyPatch): + """Test that _build_flux_config_from_yaml() extracts parameters from backend_args.""" + backend_args = SimpleNamespace( + mock_data=True, + # Model architecture params + num_joint_layers=24, + num_single_layers=8, + hidden_size=2048, + num_attention_heads=16, + ffn_hidden_size=8192, + # Precision + bf16=True, + fp16=False, + params_dtype=torch.bfloat16, + # Transformer impl + transformer_impl="local", + ) + + trainer = _build_flux_trainer(monkeypatch, backend_args) + + config = trainer._build_flux_config_from_yaml() + + # Verify config attributes were set + assert config.num_joint_layers == 24 + assert config.num_single_layers == 8 + assert config.hidden_size == 2048 + assert config.num_attention_heads == 16 + assert config.ffn_hidden_size == 8192 + assert config.bf16 is True + assert config.fp16 is False + assert config.params_dtype == torch.bfloat16 + assert config.transformer_impl == "local" + + def test_build_flux_config_from_yaml_torch_compile_settings(self, monkeypatch: pytest.MonkeyPatch): + """Test that torch_compile settings are extracted from backend_args.""" + backend_args = SimpleNamespace( + mock_data=True, + torch_compile=SimpleNamespace( + enable=True, + backend="inductor", + mode="reduce-overhead", + fullgraph=True, + ), + ) + + trainer = _build_flux_trainer(monkeypatch, backend_args) + + config = trainer._build_flux_config_from_yaml() + + # Verify torch_compile settings + assert config.enable_torch_compile is True + assert config.torch_compile_backend == "inductor" + assert config.torch_compile_mode == "reduce-overhead" + assert config.torch_compile_fullgraph is True + + def test_create_model_sets_torch_compile_on_args(self, monkeypatch: pytest.MonkeyPatch): + """Test that create_model() sets torch_compile attributes on args.""" + backend_args = SimpleNamespace( + mock_data=True, + torch_compile=SimpleNamespace( + enable=True, + backend="inductor", + mode="reduce-overhead", + fullgraph=True, + ), + ) + + trainer = _build_flux_trainer(monkeypatch, backend_args) + + # Mock Flux model + mock_flux_model = Mock() + mock_flux_model.parameters.return_value = [] + monkeypatch.setattr( + "primus.backends.megatron.core.models.diffusion.flux.model.Flux", + lambda *args, **kwargs: mock_flux_model, + ) + + # Mock get_args + mock_args = SimpleNamespace(rank=0) + training_mod = types.SimpleNamespace(get_args=lambda: mock_args) + monkeypatch.setitem(sys.modules, "megatron.training", training_mod) + + trainer.create_model() + + # Verify torch_compile attributes were set on args + assert mock_args.enable_torch_compile is True + assert mock_args.torch_compile_backend == "inductor" + assert mock_args.torch_compile_mode == "reduce-overhead" + assert mock_args.torch_compile_fullgraph is True + + def test_build_flux_config_from_yaml_fp8_settings(self, monkeypatch: pytest.MonkeyPatch): + """Test that FP8 fields are extracted from backend_args into FluxConfig.""" + backend_args = SimpleNamespace( + mock_data=True, + fp8="e4m3", + fp8_recipe="tensorwise", + fp8_margin=0, + fp8_amax_history_len=1, + fp8_amax_compute_algo="most_recent", + fp8_wgrad=True, + fp8_dot_product_attention=False, + fp8_multi_head_attention=False, + ) + + trainer = _build_flux_trainer(monkeypatch, backend_args) + + config = trainer._build_flux_config_from_yaml() + + assert config.fp8 == "e4m3" + assert config.fp8_recipe == "tensorwise" + assert config.fp8_margin == 0 + assert config.fp8_amax_history_len == 1 + assert config.fp8_amax_compute_algo == "most_recent" + assert config.fp8_wgrad is True + assert config.fp8_dot_product_attention is False + assert config.fp8_multi_head_attention is False + + def test_create_model_does_not_overwrite_existing_torch_compile_attrs( + self, monkeypatch: pytest.MonkeyPatch + ): + """Test that create_model() doesn't overwrite existing torch_compile attributes.""" + backend_args = SimpleNamespace( + mock_data=True, + torch_compile=SimpleNamespace( + enable=True, + backend="inductor", + mode="reduce-overhead", + fullgraph=True, + ), + ) + + trainer = _build_flux_trainer(monkeypatch, backend_args) + + # Mock Flux model + mock_flux_model = Mock() + mock_flux_model.parameters.return_value = [] + monkeypatch.setattr( + "primus.backends.megatron.core.models.diffusion.flux.model.Flux", + lambda *args, **kwargs: mock_flux_model, + ) + + # Mock get_args with existing attributes + mock_args = SimpleNamespace( + rank=0, + enable_torch_compile=False, # Already set + torch_compile_backend="aot_eager", # Already set + ) + training_mod = types.SimpleNamespace(get_args=lambda: mock_args) + monkeypatch.setitem(sys.modules, "megatron.training", training_mod) + + trainer.create_model() + + # Verify existing attributes were NOT overwritten + assert mock_args.enable_torch_compile is False + assert mock_args.torch_compile_backend == "aot_eager" + # New attributes should still be set + assert mock_args.torch_compile_mode == "reduce-overhead" + assert mock_args.torch_compile_fullgraph is True diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_flux_trainer.py b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_trainer.py new file mode 100644 index 000000000..c7c1035d7 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_flux_trainer.py @@ -0,0 +1,219 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for FluxPretrainTrainer. + +Tests initialization, configuration, method overrides, scheduler setup, +and CFG dropout encoding discovery. +""" + +import os +from types import SimpleNamespace + +import numpy as np +import pytest + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.flux_pretrain_trainer import FluxPretrainTrainer + + +def _build_flux_trainer(monkeypatch: pytest.MonkeyPatch, backend_args=None): + """Helper to build FluxPretrainTrainer with stubbed dependencies.""" + + # Stub out MegatronBaseTrainer.__init__ to avoid real Megatron imports + def dummy_init(self, backend_args: any = None): + self.backend_args = backend_args + + monkeypatch.setattr( + "primus.backends.megatron.megatron_base_trainer.MegatronBaseTrainer.__init__", + dummy_init, + ) + + # Silence logging + monkeypatch.setattr( + "primus.backends.megatron.flux_pretrain_trainer.log_rank_0", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "primus.backends.megatron.diffusion_trainer.log_rank_0", + lambda *args, **kwargs: None, + ) + + if backend_args is None: + backend_args = SimpleNamespace( + guidance_embed=False, + guidance_scale=3.5, + num_train_timesteps=1000, + scheduler_shift=1.0, + use_dynamic_shifting=False, + ) + + return FluxPretrainTrainer(backend_args=backend_args) + + +class TestFluxPretrainTrainer: + """Tests for FluxPretrainTrainer.""" + + def test_flux_trainer_scheduler_creation(self, monkeypatch: pytest.MonkeyPatch): + """Test that scheduler is created correctly.""" + backend_args = SimpleNamespace( + guidance_embed=False, + guidance_scale=3.5, + num_train_timesteps=1000, + scheduler_shift=1.0, + use_dynamic_shifting=False, + ) + trainer = _build_flux_trainer(monkeypatch, backend_args) + + # Create scheduler + scheduler = trainer.create_scheduler() + + # Should be FlowMatchEulerDiscreteScheduler + from primus.backends.megatron.training.diffusion.schedulers import ( + FlowMatchEulerDiscreteScheduler, + ) + + assert isinstance(scheduler, FlowMatchEulerDiscreteScheduler) + assert scheduler.num_train_timesteps == 1000 + assert scheduler.shift == 1.0 + assert scheduler.use_dynamic_shifting is False + + def test_flux_trainer_scheduler_with_dynamic_shifting(self, monkeypatch: pytest.MonkeyPatch): + """Test scheduler creation with dynamic shifting enabled.""" + backend_args = SimpleNamespace( + guidance_embed=False, + guidance_scale=3.5, + num_train_timesteps=1000, + scheduler_shift=1.0, + use_dynamic_shifting=True, + base_shift=0.5, + max_shift=1.15, + ) + trainer = _build_flux_trainer(monkeypatch, backend_args) + + scheduler = trainer.create_scheduler() + + assert scheduler.use_dynamic_shifting is True + assert scheduler.base_shift == 0.5 + assert scheduler.max_shift == 1.15 + + def test_flux_trainer_scheduler_lazy_initialization(self, monkeypatch: pytest.MonkeyPatch): + """Test that scheduler is lazily initialized via property.""" + trainer = _build_flux_trainer(monkeypatch) + + # Scheduler should not be created yet + assert trainer._scheduler is None + + # Accessing scheduler property should create it + scheduler1 = trainer.scheduler + assert scheduler1 is not None + assert trainer._scheduler is scheduler1 + + # Accessing again should return same instance + scheduler2 = trainer.scheduler + assert scheduler2 is scheduler1 + + +def _create_encoding_files(directory): + """Create minimal t5_empty.npy and clip_empty.npy in a directory.""" + os.makedirs(directory, exist_ok=True) + np.save(os.path.join(directory, "t5_empty.npy"), np.zeros((1, 256, 4096), dtype=np.float32)) + np.save(os.path.join(directory, "clip_empty.npy"), np.zeros((1, 768), dtype=np.float32)) + + +class TestCFGDropoutDiscovery: + """Tests for CFG dropout empty encoding discovery and error handling.""" + + def test_discover_encodings_inside_data_path(self, monkeypatch, tmp_path): + """Auto-discovers encodings at {data_path}/empty_encodings/.""" + monkeypatch.setattr( + "primus.backends.megatron.flux_pretrain_trainer.log_rank_0", + lambda *a, **kw: None, + ) + enc_dir = tmp_path / "empty_encodings" + _create_encoding_files(str(enc_dir)) + + params = SimpleNamespace(data_path=str(tmp_path)) + result = FluxPretrainTrainer._discover_empty_encodings(params) + + assert result == str(enc_dir) + + def test_discover_encodings_alongside_data_path(self, monkeypatch, tmp_path): + """Falls back to {data_path}/../empty_encodings/ when not inside.""" + monkeypatch.setattr( + "primus.backends.megatron.flux_pretrain_trainer.log_rank_0", + lambda *a, **kw: None, + ) + data_dir = tmp_path / "dataset" + data_dir.mkdir() + enc_dir = tmp_path / "empty_encodings" + _create_encoding_files(str(enc_dir)) + + params = SimpleNamespace(data_path=str(data_dir)) + result = FluxPretrainTrainer._discover_empty_encodings(params) + + assert result == str(enc_dir) + + def test_discover_encodings_explicit_path(self, tmp_path): + """Explicit empty_encodings_path takes priority.""" + explicit_dir = tmp_path / "custom_encodings" + _create_encoding_files(str(explicit_dir)) + + inside_dir = tmp_path / "data" / "empty_encodings" + _create_encoding_files(str(inside_dir)) + + params = SimpleNamespace( + data_path=str(tmp_path / "data"), + empty_encodings_path=str(explicit_dir), + ) + result = FluxPretrainTrainer._discover_empty_encodings(params) + + assert result == str(explicit_dir) + + def test_discover_encodings_returns_none_when_missing(self, tmp_path): + """Returns None when no encoding files exist anywhere.""" + params = SimpleNamespace(data_path=str(tmp_path)) + result = FluxPretrainTrainer._discover_empty_encodings(params) + + assert result is None + + def test_discover_encodings_returns_none_for_partial_files(self, tmp_path): + """Returns None when only one of the two required files exists.""" + enc_dir = tmp_path / "empty_encodings" + enc_dir.mkdir() + np.save(str(enc_dir / "t5_empty.npy"), np.zeros((1, 256, 4096), dtype=np.float32)) + + params = SimpleNamespace(data_path=str(tmp_path)) + result = FluxPretrainTrainer._discover_empty_encodings(params) + + assert result is None + + def test_discover_encodings_handles_list_data_path(self, monkeypatch, tmp_path): + """Handles data_path passed as a list (takes first element).""" + monkeypatch.setattr( + "primus.backends.megatron.flux_pretrain_trainer.log_rank_0", + lambda *a, **kw: None, + ) + enc_dir = tmp_path / "empty_encodings" + _create_encoding_files(str(enc_dir)) + + params = SimpleNamespace(data_path=[str(tmp_path), "/nonexistent"]) + result = FluxPretrainTrainer._discover_empty_encodings(params) + + assert result == str(enc_dir) + + def test_cfg_dropout_raises_when_no_encodings_found(self, monkeypatch, tmp_path): + """Real data without encodings raises FileNotFoundError with instructions.""" + backend_args = SimpleNamespace( + cfg_dropout_prob=0.1, + mock_data=False, + data_path=str(tmp_path), + tensor_model_parallel_size=1, + ) + + with pytest.raises(FileNotFoundError, match="primus data diffusion-encoded"): + _build_flux_trainer(monkeypatch, backend_args) diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step_count_gate.py b/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step_count_gate.py new file mode 100644 index 000000000..8a7d7635c --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_forward_step_count_gate.py @@ -0,0 +1,173 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for the model.training gate on DiffusionPretrainTrainer._forward_step_count. + +Contract under test: when ``model.training`` is False (the trainer is running +an evaluation forward pass), the counter must NOT advance and the per-step +CUDA RNG reseed must NOT fire. Otherwise validation steps shift the training +RNG sequence by ``eval_iters * num_microbatches`` per ``--eval-interval`` +window, defeating the goal of isolating the training RNG from unrelated +forward passes. + +These tests pair with +``tests/unit_tests/backends/megatron/test_diffusion_trainer_forward_step_count.py``, +which covers lazy initialization from checkpoint state and the +``step_count`` flow into ``flux_forward_step_func``. This file is the +single source of truth for the eval-skip contract — if the gate is reverted, +the three ``TestForwardStepCountTrainingGate`` tests below must fail. +""" + +from unittest.mock import Mock, patch + +import torch + +from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer + + +class _ConcreteDiffusionTrainer(DiffusionPretrainTrainer): + """Minimal concrete subclass: abstract methods as no-ops.""" + + def create_model(self, *args, **kwargs): + return None + + def create_scheduler(self, *args, **kwargs): + return None + + def get_task_encoder(self, *args, **kwargs): + return None + + +def _make_trainer(): + """Construct a trainer with ``__init__`` bypassed. + + Sets only the attributes ``forward_step`` reads. Lazy-init is forced to + "already initialized" because that surface is covered by + ``test_diffusion_trainer_forward_step_count.py::TestForwardStepCountLazyInit``. + """ + trainer = _ConcreteDiffusionTrainer.__new__(_ConcreteDiffusionTrainer) + trainer._forward_step_count = 0 + trainer._forward_step_count_initialized = True + trainer._scheduler = None + + class _FakeRuntimeState: + def update_metrics(self, metrics): + pass + + trainer.runtime_state = _FakeRuntimeState() + return trainer + + +def _patch_flux_forward_step_func(): + """Patch the real func with a recorder that captures ``step_count``.""" + captured = [] + + def _recorder(*args, **kwargs): + captured.append(kwargs.get("step_count")) + t = torch.zeros(1) + return t, t, t, None, {}, False + + return ( + patch( + "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func", + side_effect=_recorder, + ), + captured, + ) + + +class TestForwardStepCountTrainingGate: + """Direct coverage of the ``if model.training:`` gate on counter advance.""" + + def test_eval_does_not_advance_counter(self): + trainer = _make_trainer() + patcher, captured = _patch_flux_forward_step_func() + eval_model = Mock() + eval_model.training = False + + with patcher: + trainer.forward_step(data_iterator=None, model=eval_model) + + assert trainer._forward_step_count == 0 + # Counter unchanged → step_count passed to the step func is the pre-call value. + assert captured == [0] + + def test_train_advances_counter_by_one(self): + trainer = _make_trainer() + patcher, captured = _patch_flux_forward_step_func() + train_model = Mock() + train_model.training = True + + with patcher: + trainer.forward_step(data_iterator=None, model=train_model) + + assert trainer._forward_step_count == 1 + assert captured == [1] + + def test_eval_between_train_does_not_shift_sequence(self): + """Eval batches interleaved with training must NOT change the + ``step_count`` that subsequent training steps see.""" + trainer = _make_trainer() + patcher, captured = _patch_flux_forward_step_func() + + train_model = Mock() + train_model.training = True + eval_model = Mock() + eval_model.training = False + + with patcher: + trainer.forward_step(data_iterator=None, model=train_model) # 0 -> 1 + trainer.forward_step(data_iterator=None, model=eval_model) # stays 1 + trainer.forward_step(data_iterator=None, model=eval_model) # stays 1 + trainer.forward_step(data_iterator=None, model=train_model) # 1 -> 2 + + assert trainer._forward_step_count == 2 + # Critical assertion: the 4th training step sees step_count=2, the + # same value it would see in the absence of the two eval calls. If + # the gate were missing, the sequence would be [1, 2, 3, 4]. + assert captured == [1, 1, 1, 2] + + +class TestPerStepReseedFormula: + """Positive control: when the counter advances, the per-step reseed + formula in ``flux_forward_step_func`` is computed correctly. + + This is a focused regression test for the seed expression + ``(seed + 100 * dp_rank) * 10000 + step_count) % (2**63)``. + It deliberately reaches into the underlying step function rather than + going through the trainer so the formula is testable in isolation + from the gate logic.""" + + def test_reseed_uses_step_count_in_formula(self): + from primus.backends.megatron.training.diffusion import forward_step as fwd_mod + + fake_args = Mock() + fake_args.seed = 7 + + train_model = Mock() + train_model.training = True + + with patch.object(torch.cuda, "manual_seed") as mock_seed, patch( + "megatron.training.get_args", return_value=fake_args + ), patch("megatron.core.parallel_state.get_data_parallel_rank", return_value=3): + # flux_forward_step_func runs significant downstream logic that + # would require a full Megatron-Flux harness. We only need to + # verify the reseed call. Suppress the post-reseed failure so the + # mock.assert below still observes the manual_seed call. + try: + fwd_mod.flux_forward_step_func( + data_iterator=None, + model=train_model, + scheduler=Mock(), + per_step_rng_reseed=True, + step_count=42, + ) + except Exception: + # Expected: downstream code needs real data/model. The reseed + # call we care about happens before any of that — see the + # assert below. + pass + + expected_seed = ((7 + 100 * 3) * 10000 + 42) % (2**63) + mock_seed.assert_called_once_with(expected_seed) diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_vae_resample_reproducibility.py b/tests/unit_tests/backends/megatron/diffusion/training/test_vae_resample_reproducibility.py new file mode 100644 index 000000000..c429ede2c --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_vae_resample_reproducibility.py @@ -0,0 +1,211 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for per-step RNG isolation reproducibility. + +Validates that: +- Same step_count + same seed produces bitwise-identical outputs +- Different step_count produces different outputs (RNG isolation works) +""" + +import contextlib +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.core.models.diffusion.flux.config import FluxConfig +from primus.backends.megatron.core.models.diffusion.flux.model import Flux +from primus.backends.megatron.training.diffusion.forward_step import ( + flux_forward_step_func, +) +from primus.backends.megatron.training.diffusion.schedulers.flow_matching import ( + FlowMatchEulerDiscreteScheduler, +) +from tests.unit_tests.backends.megatron.diffusion.constants import ( + CLIP_L_EMBEDDING_DIM, + T5_XXL_EMBEDDING_DIM, + VAE_LATENT_CHANNELS, +) + + +def _patch_parallel_state(seed=42): + """Return a list of mock.patch context managers for parallel state.""" + return [ + patch( + "megatron.core.parallel_state.get_tensor_model_parallel_world_size", + return_value=1, + ), + patch( + "megatron.core.parallel_state.get_pipeline_model_parallel_world_size", + return_value=1, + ), + patch( + "megatron.core.parallel_state.get_data_parallel_rank", + return_value=0, + ), + patch( + "megatron.core.parallel_state.get_tensor_model_parallel_rank", + return_value=0, + ), + patch( + "megatron.training.get_args", + return_value=SimpleNamespace(seed=seed), + ), + ] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +class TestVaeResampleReproducibility: + """Tests per-step RNG isolation via observable output reproducibility.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state for model tests.""" + + @pytest.fixture + def model(self): + config = FluxConfig.flux_535m() + m = Flux(config).cuda() + m.train() + return m + + @pytest.fixture + def scheduler(self): + return FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000) + + def _make_resample_batch(self, seed=99): + """Create a deterministic resample batch.""" + g = torch.Generator().manual_seed(seed) + return { + "mean": torch.randn(2, VAE_LATENT_CHANNELS, 16, 16, generator=g), + "logvar": torch.randn(2, VAE_LATENT_CHANNELS, 16, 16, generator=g), + "prompt_embeds": torch.randn(2, 32, T5_XXL_EMBEDDING_DIM, generator=g), + "pooled_prompt_embeds": torch.randn(2, CLIP_L_EMBEDDING_DIM, generator=g), + } + + def test_same_step_count_produces_identical_output(self, model, scheduler): + """Identical seed + step_count must produce bitwise-identical noise_pred.""" + with contextlib.ExitStack() as stack: + for p in _patch_parallel_state(): + stack.enter_context(p) + + batch1 = self._make_resample_batch(seed=99) + result1 = flux_forward_step_func( + iter([batch1]), + model, + scheduler=scheduler, + vae_latent_mode="resample", + vae_scale=0.3611, + vae_shift=0.1159, + step_count=5, + ) + noise_pred_1 = result1[0].detach().clone() + + batch2 = self._make_resample_batch(seed=99) + result2 = flux_forward_step_func( + iter([batch2]), + model, + scheduler=scheduler, + vae_latent_mode="resample", + vae_scale=0.3611, + vae_shift=0.1159, + step_count=5, + ) + noise_pred_2 = result2[0].detach().clone() + + assert torch.equal(noise_pred_1, noise_pred_2), ( + "Same step_count should produce identical output " + f"(max diff: {(noise_pred_1 - noise_pred_2).abs().max().item():.2e})" + ) + + def test_different_step_count_produces_different_latents(self, model, scheduler): + """Different step_count values should produce different clean_latents. + + Note: We compare clean_latents (index 1) rather than noise_pred (index 0) + because proj_out is zero-initialized (NeMo-aligned init), making noise_pred + always zero until training updates the projection layer. + """ + with contextlib.ExitStack() as stack: + for p in _patch_parallel_state(): + stack.enter_context(p) + + batch1 = self._make_resample_batch(seed=99) + result1 = flux_forward_step_func( + iter([batch1]), + model, + scheduler=scheduler, + vae_latent_mode="resample", + vae_scale=0.3611, + vae_shift=0.1159, + step_count=5, + ) + latents_1 = result1[1].detach().clone() + + batch2 = self._make_resample_batch(seed=99) + result2 = flux_forward_step_func( + iter([batch2]), + model, + scheduler=scheduler, + vae_latent_mode="resample", + vae_scale=0.3611, + vae_shift=0.1159, + step_count=6, + ) + latents_2 = result2[1].detach().clone() + + assert not torch.equal(latents_1, latents_2), ( + "Different step_count values should produce different latents " + "(per-step RNG isolation not working)" + ) + + def test_different_base_seed_produces_different_output(self, model, scheduler): + """Different args.seed (same step_count) must produce different latents. + + Validates the `_per_rank_seed = seed + 100*dp_rank` term in the step + seed formula at forward_step.py:237. A regression that drops `seed` + and uses only `step_count` would cause both runs to produce identical + output (since step_count is fixed at 5). + """ + # Run 1: seed=42 + with contextlib.ExitStack() as stack: + for p in _patch_parallel_state(seed=42): + stack.enter_context(p) + batch1 = self._make_resample_batch(seed=99) + result1 = flux_forward_step_func( + iter([batch1]), + model, + scheduler=scheduler, + vae_latent_mode="resample", + vae_scale=0.3611, + vae_shift=0.1159, + step_count=5, + ) + latents_42 = result1[1].detach().clone() + + # Run 2: seed=43, identical step_count and identical batch contents. + with contextlib.ExitStack() as stack: + for p in _patch_parallel_state(seed=43): + stack.enter_context(p) + batch2 = self._make_resample_batch(seed=99) + result2 = flux_forward_step_func( + iter([batch2]), + model, + scheduler=scheduler, + vae_latent_mode="resample", + vae_scale=0.3611, + vae_shift=0.1159, + step_count=5, + ) + latents_43 = result2[1].detach().clone() + + assert not torch.equal(latents_42, latents_43), ( + "Different args.seed values produced identical output — " + "the per-rank seed term appears to be missing from the step seed formula" + ) diff --git a/tests/unit_tests/backends/megatron/test_chimera_rng_restore.py b/tests/unit_tests/backends/megatron/test_chimera_rng_restore.py new file mode 100644 index 000000000..05b3bc14c --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_chimera_rng_restore.py @@ -0,0 +1,109 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for `_restore_chimera_rng_state` defensive RNG restoration. + +Validates Fix 2: when Megatron's private `_set_random_seed` API drifts +(signature change or removal), the helper falls back to a manual restore +that covers all three RNG generators — CPU default, CUDA default, and the +model-parallel tracker — so chimera training does not fail at startup. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from tests.utils import skip_if_no_cuda + +skip_if_no_cuda() + +from primus.backends.megatron.flux_pretrain_trainer import _restore_chimera_rng_state + + +def _canonical_args(seed=42): + """Build the args namespace the helper reads.""" + return SimpleNamespace( + seed=seed, + data_parallel_random_init=False, + te_rng_tracker=False, + inference_rng_tracker=False, + enable_cuda_graph=False, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +class TestChimeraRngRestoreFallback: + """The fallback path: when `_set_random_seed` raises TypeError or + ImportError, the helper must restore CPU, CUDA default, and the + model-parallel tracker generators manually so training continues.""" + + @pytest.fixture(autouse=True) + def setup_parallel(self, init_parallel_state): + """Initialize parallel state — required because the fallback's + tp_random.model_parallel_cuda_manual_seed needs the tracker set up.""" + + def _verify_all_rngs_match_seed(self, seed): + """Sample from each generator and compare to a freshly-seeded + reference. If the helper's fallback fully restored the state, + the two samples must be bitwise identical for every generator.""" + from megatron.core.tensor_parallel import random as tp_random + + # CPU default generator. + cpu_after = torch.randn(8) + torch.manual_seed(seed) + cpu_ref = torch.randn(8) + assert torch.equal(cpu_after, cpu_ref), "CPU default generator was not restored to seed" + + # CUDA default generator. + cuda_after = torch.randn(8, device="cuda") + torch.cuda.manual_seed(seed) + cuda_ref = torch.randn(8, device="cuda") + assert torch.equal(cuda_after, cuda_ref), "CUDA default generator was not restored to seed" + + # Model-parallel tracker. + with tp_random.get_cuda_rng_tracker().fork(): + tracker_after = torch.randn(8, device="cuda") + tp_random.model_parallel_cuda_manual_seed(seed) + with tp_random.get_cuda_rng_tracker().fork(): + tracker_ref = torch.randn(8, device="cuda") + assert torch.equal(tracker_after, tracker_ref), "Model-parallel tracker was not restored to seed" + + def _contaminate_rngs(self, contamination_seed): + """Set all three generators to a non-canonical seed so the test + observes the helper actively restoring rather than no-op-passing.""" + from megatron.core.tensor_parallel import random as tp_random + + torch.manual_seed(contamination_seed) + torch.cuda.manual_seed(contamination_seed) + tp_random.model_parallel_cuda_manual_seed(contamination_seed) + + def test_typeerror_fallback_restores_all_rngs(self): + """Megatron signature drift (TypeError) must trigger manual restore.""" + seed = 42 + self._contaminate_rngs(contamination_seed=999) + + args = _canonical_args(seed=seed) + with patch( + "megatron.training.initialize._set_random_seed", + side_effect=TypeError("signature changed"), + ): + _restore_chimera_rng_state(args) + + self._verify_all_rngs_match_seed(seed) + + def test_importerror_fallback_restores_all_rngs(self): + """Megatron module removal (ImportError) must trigger manual restore.""" + seed = 42 + self._contaminate_rngs(contamination_seed=999) + + args = _canonical_args(seed=seed) + with patch( + "megatron.training.initialize._set_random_seed", + side_effect=ImportError("module gone"), + ): + _restore_chimera_rng_state(args) + + self._verify_all_rngs_match_seed(seed) diff --git a/tests/unit_tests/backends/megatron/test_diffusion_trainer_forward_step_count.py b/tests/unit_tests/backends/megatron/test_diffusion_trainer_forward_step_count.py new file mode 100644 index 000000000..44f7cc798 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_diffusion_trainer_forward_step_count.py @@ -0,0 +1,161 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the Apache License, Version 2.0. + +""" +Unit tests for DiffusionPretrainTrainer's forward step count. + +Validates Fix 1: the `_forward_step_count` instance variable and its +checkpoint-compatible lazy initialization from `args.iteration * +get_num_microbatches()`. Also verifies the counter flows through to +`flux_forward_step_func` as the `step_count` keyword argument. +""" + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from primus.backends.megatron.diffusion_trainer import DiffusionPretrainTrainer + + +class _ConcreteDiffusionTrainer(DiffusionPretrainTrainer): + """Minimal concrete subclass: implements abstract methods as no-ops. + + These methods are only required for instantiation; the tests never + invoke them. + """ + + def create_model(self, *args, **kwargs): + return None + + def create_scheduler(self, *args, **kwargs): + return None + + def get_task_encoder(self, *args, **kwargs): + return None + + +def _make_trainer(): + """Construct a DiffusionPretrainTrainer with __init__ bypassed. + + Sets only the attributes required by `get_forward_step()` and + `forward_step()`. Other attributes are intentionally left unset so + that tests fail loudly if production code starts depending on them. + """ + trainer = _ConcreteDiffusionTrainer.__new__(_ConcreteDiffusionTrainer) + trainer._forward_step_count = 0 + trainer._forward_step_count_initialized = False + return trainer + + +class TestForwardStepCountLazyInit: + """Tests for lazy initialization of _forward_step_count from checkpoint state.""" + + def test_lazy_init_from_checkpoint_iteration(self): + """First closure call must reconstruct counter as iteration * num_microbatches.""" + trainer = _make_trainer() + trainer.forward_step = lambda data_iterator, model, return_schedule_plan=False: None + + with patch( + "megatron.training.get_args", + return_value=SimpleNamespace(iteration=100), + ), patch( + "megatron.core.num_microbatches_calculator.get_num_microbatches", + return_value=4, + ): + closure = trainer.get_forward_step() + closure(data_iterator=None, model=None) + + assert trainer._forward_step_count == 400 + assert trainer._forward_step_count_initialized is True + + def test_lazy_init_only_runs_once(self): + """Second call must not re-read args.iteration even if it changes.""" + trainer = _make_trainer() + trainer.forward_step = lambda data_iterator, model, return_schedule_plan=False: None + + args_state = SimpleNamespace(iteration=10) + + with patch( + "megatron.training.get_args", + side_effect=lambda: args_state, + ), patch( + "megatron.core.num_microbatches_calculator.get_num_microbatches", + return_value=2, + ): + closure = trainer.get_forward_step() + closure(data_iterator=None, model=None) + assert trainer._forward_step_count == 20 + + args_state.iteration = 999 + closure(data_iterator=None, model=None) + + # Must remain 20: lazy init flag prevents a second reconstruction. + assert trainer._forward_step_count == 20 + assert trainer._forward_step_count_initialized is True + + def test_initial_iteration_zero(self): + """Fresh start (iteration=0) reconstructs counter to 0.""" + trainer = _make_trainer() + trainer.forward_step = lambda data_iterator, model, return_schedule_plan=False: None + + with patch( + "megatron.training.get_args", + return_value=SimpleNamespace(iteration=0), + ), patch( + "megatron.core.num_microbatches_calculator.get_num_microbatches", + return_value=8, + ): + closure = trainer.get_forward_step() + closure(data_iterator=None, model=None) + + assert trainer._forward_step_count == 0 + assert trainer._forward_step_count_initialized is True + + +class TestCounterFlowsToForwardStepFunc: + """Integration test: trainer counter must pass through to flux_forward_step_func.""" + + def test_counter_flows_to_forward_step_func(self): + """forward_step must invoke flux_forward_step_func with step_count == counter.""" + trainer = _make_trainer() + # _scheduler backs the `scheduler` lazy property; setting it pre-empts + # create_scheduler() which would otherwise raise abstract NotImplementedError. + trainer._scheduler = None + + # Provide a fake runtime_state so forward_step doesn't fall into the + # log_rank_0 warning branch (which requires the Primus logger to be + # initialized — outside this test's scope). + class _FakeRuntimeState: + def update_metrics(self, metrics): + pass + + trainer.runtime_state = _FakeRuntimeState() + + recorded_kwargs = {} + + def recording_func(*args, **kwargs): + recorded_kwargs.update(kwargs) + # Return shape matches: (noise_pred, clean_latents, noise, loss_mask, metrics, is_validation) + t = torch.zeros(1) + return t, t, t, None, {}, False + + # Training model: `if model.training:` gate in forward_step requires + # a real attribute (None would AttributeError). Eval-skip behavior is + # covered in test_forward_step_count_gate.py. + train_model = Mock() + train_model.training = True + + with patch( + "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func", + side_effect=recording_func, + ): + # First forward step — counter increments from 0 to 1 + trainer.forward_step(data_iterator=None, model=train_model) + assert recorded_kwargs["step_count"] == 1 + assert trainer._forward_step_count == 1 + + # Second forward step — counter increments to 2 + trainer.forward_step(data_iterator=None, model=train_model) + assert recorded_kwargs["step_count"] == 2 + assert trainer._forward_step_count == 2 From 714156684f5c824239d880a4591abf122576239c Mon Sep 17 00:00:00 2001 From: zirui Date: Thu, 16 Jul 2026 16:43:50 +0800 Subject: [PATCH 036/127] Add flux.1 to diffusion backend (#832) ## Summary This PR adds FLUX.1-dev text-to-image training support on top of the Primus `diffusion` backend introduced in #779. The goal is to make FLUX T2I pretraining runnable through the existing Primus training entrypoint, reusing the diffusion FSDP2 trainer, config system, checkpoint flow, logging, and launch infrastructure. --- ## Motivation FLUX is a primary diffusion workload for text-to-image training. This integration enables: - Unified Primus training interface for FLUX and Wan diffusion models - FLUX.1-dev T2I pretraining with FSDP2 - Support for both precomputed encodings and raw image-text data bring-up - Reuse of existing Primus infra for launch, optimizer, checkpointing, logging, and hooks --- ## Scope ### Included - FLUX.1-dev model wrapper and registration - FLUX training pipeline with flow-matching objective - Precomputed dataset path using T5 encodings, CLIP encodings, and VAE latent statistics - Raw image-text dataset path with frozen T5 / CLIP / autoencoder encoding - Diffusion data registry/generalized dataset processing - MI355X example configs for precomputed and raw FLUX pretraining - Preflight/dependency updates for FLUX assets and datasets - Focused unit tests for config conversion, dataset processing, and tiny FLUX loss computation ### Not included / future work - FLUX sequence parallelism; FLUX currently requires `sp_size=1` - Multi-node validation - Inference pipeline - Full production-scale convergence validation - Kernel-level performance optimization --- ## Current Status ### Completed - [x] FLUX.1-dev model integration - [x] FLUX model/data registration in the diffusion backend - [x] Precomputed encoding dataset support - [x] Raw image-text dataset support - [x] Single-node 8-GPU smoke validation - [x] Focused FLUX backend unit tests - [x] Single node training benchmarks ### In Progress / Follow-up --- ## Testing ### Unit Tests - `python -m pytest tests/unit_tests/backends/diffusion/test_flux_backend.py` - Result: 8 passed ### Single-node Validation Validated on 8x AMD Instinct MI355X with `torchrun --standalone --nproc_per_node=8`: - FLUX precomputed encoding mode: local batch sizes 1 / 8 / 16 / 24 - FLUX raw image-text mode: local batch sizes 1 / 8 / 16 / 24 - Additional 8-step smoke runs for: - Precomputed mode, local batch size 24 - Raw mode, local batch size 16 ## Known Issues / Risks - Multi-node stability has not been validated yet - FLUX sequence parallelism is not implemented; configs should keep `sp_size=1` - Full convergence quality still needs longer training runs with production data - Current ROCm validation used `gradient_checkpointing=false`; enabling it needs more investigation with FSDP2 ## Next Steps 1. Complete multi-node validation 2. Run longer convergence checks on real FLUX pretraining data 3. Add performance benchmark notes for recommended batch sizes 4. Investigate gradient checkpointing and FLUX-specific SP support --- ## Notes This is a WIP draft PR building on the diffusion backend and Wan training support from #779. # FLUX TorchTitan-aligned benchmark summary - Hardware: single node, 8x MI355X - Config: `flux-schnell`, block-level `torch.compile`, TorchTitan-style FSDP wrapping, AdamW beta2=0.95, weight_decay=0.1 - Steps: 100 logged, metrics below use steps 2-100. ## Precomputed Encodings - Dataset: `/data/cc12m_preprocessed`; empty encodings: `/data/empty_encodings`. - Runtime encoders: disabled, matching TorchTitan `flux_schnell_mlperf_preprocessed`. ## 100-step average, wall-clock consistent | local batch/GPU | mean alloc GB | mean reserved GB | max peak GB | mean step_time s | step/s | TPS samples/s/GPU | p95 step_time s | |---:|---:|---:|---:|---:|---:|---:|---:| | 16 | 9.11 | 86.95 | 78.47 | 1.4182 | 0.7051 | 11.28 | 8.1700 | | 32 | 9.09 | 155.58 | 147.15 | 1.2381 | 0.8077 | 25.85 | 3.1700 | | 64 | 9.08 | 294.31 | 284.41 | 1.6671 | 0.5999 | 38.39 | 3.3900 | ## Steady median | local batch/GPU | median step_time s | median step/s | median TPS samples/s/GPU | mean instantaneous TPS | |---:|---:|---:|---:|---:| | 16 | 0.4800 | 2.0833 | 33.33 | 30.64 | | 32 | 0.7700 | 1.2987 | 41.56 | 38.92 | | 64 | 1.3900 | 0.7194 | 46.04 | 43.51 | ## Raw cc12m-test JSONL - Dataset: `local_runs/flux_bench_20260618/cc12m_test_repeat/metadata.local.jsonl` (1024 repeated local samples). - Raw smoke and bs16/32/64 100-step runs all completed. ### 100-Step Average | local batch/GPU | mean alloc GB | mean reserved GB | max peak GB | mean step_time s | wall step/s | wall TPS samples/s/GPU | mean logged TPS samples/s/GPU | p95 step_time s | |---:|---:|---:|---:|---:|---:|---:|---:|---:| | 16 | 19.04 | 93.45 | 88.40 | 0.6321 | 1.5820 | 25.31 | 26.87 | 0.6500 | | 32 | 19.08 | 161.61 | 157.13 | 1.6035 | 0.6236 | 19.96 | 30.88 | 4.4200 | | 64 | 19.03 | 299.44 | 294.38 | 1.8658 | 0.5360 | 34.30 | 34.58 | 2.0400 | ### Steady Median | local batch/GPU | median step_time s | median step/s | median TPS samples/s/GPU | |---:|---:|---:|---:| | 16 | 0.5800 | 1.7241 | 27.59 | | 32 | 0.9600 | 1.0417 | 33.33 | | 64 | 1.8900 | 0.5291 | 33.86 | --------- Co-authored-by: Cursor Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- .gitignore | 4 + examples/diffusion/README.md | 101 +++- .../MI355X/flux.1_schnell_t2i-pretrain.yaml | 58 +++ .../flux.1_schnell_t2i-raw-pretrain.yaml | 69 +++ .../MI355X/wan2.1_t2v_1.3b-posttrain.yaml | 2 +- .../MI355X/wan2.1_t2v_1.3b-pretrain.yaml | 2 +- .../MI355X/wan2.2_ti2v_5b-posttrain.yaml | 2 +- .../MI355X/wan2.2_ti2v_5b-pretrain.yaml | 2 +- primus/backends/diffusion/README.md | 263 ++-------- primus/backends/diffusion/argument_builder.py | 156 +++++- .../backends/diffusion/attention/attention.py | 2 +- primus/backends/diffusion/data/__init__.py | 22 +- primus/backends/diffusion/data/collator.py | 98 ++++ primus/backends/diffusion/data/config.py | 89 ++++ primus/backends/diffusion/data/dataset.py | 291 ++++++++++++ .../diffusion/data/flux_precomputed.py | 340 +++++++++++++ .../diffusion/data/processing_wanvideo.py | 401 ++++++++++++++++ primus/backends/diffusion/data/processor.py | 182 +++++++ .../diffusion/data/registrations/__init__.py | 2 +- .../diffusion/data/registrations/flux.py | 38 ++ .../diffusion/data/registrations/wan.py | 190 +------- .../backends/diffusion/diffusion_adapter.py | 11 +- .../diffusion/diffusion_pretrain_trainer.py | 38 +- primus/backends/diffusion/models/__init__.py | 2 +- .../diffusion/models/flux/__init__.py | 25 + .../backends/diffusion/models/flux/adapter.py | 113 +++++ .../diffusion/models/flux/autoencoder.py | 315 ++++++++++++ .../diffusion/models/flux/conditioner.py | 54 +++ .../models/flux/configuration_flux.py | 14 + .../backends/diffusion/models/flux/layers.py | 211 ++++++++ primus/backends/diffusion/models/flux/math.py | 44 ++ .../backends/diffusion/models/flux/model.py | 167 +++++++ .../diffusion/models/flux/train_pipeline.py | 184 +++++++ .../backends/diffusion/models/flux/utils.py | 49 ++ .../diffusion/models/registrations/flux.py | 234 +++++++++ .../backends/diffusion/models/wan/adapter.py | 2 +- .../diffusion/models/wan/train_pipeline.py | 28 +- primus/backends/diffusion/registry.py | 27 ++ .../diffusion/schedulers/flow_match.py | 5 + .../backends/diffusion/trainers/__init__.py | 2 +- primus/backends/diffusion/trainers/base.py | 38 +- primus/backends/diffusion/trainers/fsdp2.py | 53 ++- primus/backends/diffusion/utils/__init__.py | 4 +- .../backends/diffusion/utils/train_utils.py | 12 - .../models/diffusion/flux.1_dev_t2i.yaml | 15 + .../models/diffusion/flux.1_schnell_t2i.yaml | 15 + .../hooks/train/pretrain/diffusion/prepare.py | 105 +++- .../diffusion/requirements-diffusion.txt | 4 + tests/unit_tests/backends/__init__.py | 0 .../unit_tests/backends/diffusion/__init__.py | 0 .../diffusion/test_flow_match_scheduler.py | 14 + .../backends/diffusion/test_flux_backend.py | 449 ++++++++++++++++++ .../diffusion/test_review_regressions.py | 66 +++ .../unit_tests/backends/megatron/__init__.py | 0 .../backends/torchtitan/__init__.py | 0 tests/unit_tests/cli/__init__.py | 0 tests/unit_tests/core/__init__.py | 0 tests/unit_tests/core/backend/__init__.py | 0 tests/unit_tests/core/launcher/__init__.py | 0 tests/unit_tests/core/patches/__init__.py | 0 tests/unit_tests/core/runtime/__init__.py | 0 tests/unit_tests/core/trainer/__init__.py | 0 tests/unit_tests/core/utils/__init__.py | 0 63 files changed, 4110 insertions(+), 504 deletions(-) create mode 100644 examples/diffusion/configs/MI355X/flux.1_schnell_t2i-pretrain.yaml create mode 100644 examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml create mode 100644 primus/backends/diffusion/data/collator.py create mode 100644 primus/backends/diffusion/data/config.py create mode 100644 primus/backends/diffusion/data/dataset.py create mode 100644 primus/backends/diffusion/data/flux_precomputed.py create mode 100644 primus/backends/diffusion/data/processing_wanvideo.py create mode 100644 primus/backends/diffusion/data/processor.py create mode 100644 primus/backends/diffusion/data/registrations/flux.py create mode 100644 primus/backends/diffusion/models/flux/__init__.py create mode 100644 primus/backends/diffusion/models/flux/adapter.py create mode 100644 primus/backends/diffusion/models/flux/autoencoder.py create mode 100644 primus/backends/diffusion/models/flux/conditioner.py create mode 100644 primus/backends/diffusion/models/flux/configuration_flux.py create mode 100644 primus/backends/diffusion/models/flux/layers.py create mode 100644 primus/backends/diffusion/models/flux/math.py create mode 100644 primus/backends/diffusion/models/flux/model.py create mode 100644 primus/backends/diffusion/models/flux/train_pipeline.py create mode 100644 primus/backends/diffusion/models/flux/utils.py create mode 100644 primus/backends/diffusion/models/registrations/flux.py create mode 100644 primus/configs/models/diffusion/flux.1_dev_t2i.yaml create mode 100644 primus/configs/models/diffusion/flux.1_schnell_t2i.yaml create mode 100644 tests/unit_tests/backends/__init__.py create mode 100644 tests/unit_tests/backends/diffusion/__init__.py create mode 100644 tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py create mode 100644 tests/unit_tests/backends/diffusion/test_flux_backend.py create mode 100644 tests/unit_tests/backends/diffusion/test_review_regressions.py create mode 100644 tests/unit_tests/backends/megatron/__init__.py create mode 100644 tests/unit_tests/backends/torchtitan/__init__.py create mode 100644 tests/unit_tests/cli/__init__.py create mode 100644 tests/unit_tests/core/__init__.py create mode 100644 tests/unit_tests/core/backend/__init__.py create mode 100644 tests/unit_tests/core/launcher/__init__.py create mode 100644 tests/unit_tests/core/patches/__init__.py create mode 100644 tests/unit_tests/core/runtime/__init__.py create mode 100644 tests/unit_tests/core/trainer/__init__.py create mode 100644 tests/unit_tests/core/utils/__init__.py diff --git a/.gitignore b/.gitignore index 59e120234..a6178a7c2 100644 --- a/.gitignore +++ b/.gitignore @@ -13,5 +13,9 @@ local/ output experiment /data/* +!primus/backends/diffusion/data/ +!primus/backends/diffusion/data/** +primus/backends/diffusion/data/**/__pycache__/ +primus/backends/diffusion/data/**/*.pyc pp_simulation_result .cursor/ diff --git a/examples/diffusion/README.md b/examples/diffusion/README.md index 2df6093f7..42ad9a61a 100644 --- a/examples/diffusion/README.md +++ b/examples/diffusion/README.md @@ -1,12 +1,60 @@ -# Wan Examples +# Diffusion Examples -Wan examples exercise the independent PyTorch Diffusion backend under -`primus/backends/diffusion`. For backend details, data/checkpoint layout, and minimal -configs, see `primus/backends/diffusion/README.md`. +This directory contains launch examples for the in-tree `diffusion` backend. -## Data +## Common Launch Env -The default smoke-test dataset is `zirui3/tiny-video-samples` on Hugging Face: +```bash +export NNODES=${NNODES:-1} +export NODE_RANK=${NODE_RANK:-0} +export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} +export MASTER_PORT=${MASTER_PORT:-29500} +export GPUS_PER_NODE=${GPUS_PER_NODE:-8} +``` + +## FLUX.1-schnell Raw Image-Text + +Raw mode loads image-text samples and runs frozen T5, CLIP, and FLUX AE online. +The default `DATASET=cc12m-test` uses the Hugging Face dataset +`zirui3/cc12m-test`, so no dataset preprocessing is required for a smoke test. + +Download the encoders and autoencoder before launching training: + +```bash +huggingface-cli download google/t5-v1_1-xxl \ + --local-dir /models/t5-v1_1-xxl +huggingface-cli download openai/clip-vit-large-patch14 \ + --local-dir /models/clip-vit-large-patch14 +huggingface-cli download black-forest-labs/FLUX.1-dev ae.safetensors \ + --local-dir /models/FLUX.1-dev +``` + +Launch raw training: + +```bash +T5_ENCODER=/models/t5-v1_1-xxl \ +CLIP_ENCODER=/models/clip-vit-large-patch14 \ +VAE_CHECKPOINT=/models/FLUX.1-dev/ae.safetensors \ +MAX_STEPS=10 \ +torchrun \ + --nnodes="$NNODES" --node_rank="$NODE_RANK" \ + --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ + --nproc_per_node="$GPUS_PER_NODE" \ + -m primus.cli.main train pretrain \ + --config examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml +``` + +To use a local WebDataset directory instead, set `DATASET_PATH=/path/to/tars`. +To use the full Hugging Face dataset directly, add `DATASET=cc12m-wds` to the +launch command and omit `DATASET_PATH`. + +To run FLUX.1-dev, use the same training example shape and set the model preset +to `flux.1_dev_t2i.yaml`. FLUX.1-dev has a guidance embedding module; +FLUX.1-schnell does not. + +## Wan Data + +Wan examples use a JSONL metadata file plus a media directory: ```bash huggingface-cli download zirui3/tiny-video-samples \ @@ -22,19 +70,18 @@ Expected layout: data/*.mp4 ``` -## Run - -Set the shared `torchrun` environment first: +Download Wan checkpoints separately and set the model paths used by the selected +config. For Wan2.2 TI2V 5B, the default paths can be overridden with: ```bash -export NNODES=${NNODES:-1} -export NODE_RANK=${NODE_RANK:-0} -export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} -export MASTER_PORT=${MASTER_PORT:-29500} -export GPUS_PER_NODE=${GPUS_PER_NODE:-8} +export PRETRAINED_PATH=/models/Wan2.2-TI2V-5B +export INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B +export TEXT_TOKENIZER=/models/Wan2.2-TI2V-5B/google/umt5-xxl +export TEXT_ENCODER=/models/Wan2.2-TI2V-5B/models_t5_umt5-xxl-enc-bf16.pth +export VAE_CHECKPOINT=/models/Wan2.2-TI2V-5B/Wan2.2_VAE.pth ``` -### Pretrain +## Wan Pretrain ```bash DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ @@ -50,15 +97,16 @@ torchrun \ --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml ``` -Use `SP_SIZE=4` or `SP_SIZE=8` to enable Ulysses sequence parallelism -when the model head count supports it. +Use `SP_SIZE=4` or `SP_SIZE=8` when the model head count supports it. -### Posttrain +## Wan Posttrain ```bash INIT_CHECKPOINT=/models/Wan2.2-TI2V-5B \ DATASET_PATH=/data/tiny-video-samples/meta.jsonl \ DATA_FOLDER=/data/tiny-video-samples/data \ +ATTENTION_BACKEND=flash_attn_aiter \ +SP_SIZE=1 \ MAX_STEPS=10 \ torchrun \ --nnodes="$NNODES" --node_rank="$NODE_RANK" \ @@ -68,7 +116,16 @@ torchrun \ --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml ``` -The MI355X configs use Primus-style override sections such as `training`, -`data`, `parallelism`, `optimizer`, `runtime`, and `metrics`. The diffusion -adapter normalizes those sections into the Wan model/dataset/trainer -arguments at runtime. +## Prepare Check + +Validate configured paths before launching: + +```bash +python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ + --config examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml + +python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ + --config examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml +``` + +On success the hook prints `env.PREPARED=1`. diff --git a/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-pretrain.yaml new file mode 100644 index 000000000..9429fb1f7 --- /dev/null +++ b/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-pretrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:flux.1_schnell_t2i-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + + model: flux.1_schnell_t2i.yaml + overrides: + sink_level: null + file_sink_level: INFO + stderr_sink_level: INFO + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: ${LOCAL_BATCH_SIZE:1} + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/flux.1_schnell_t2i-pretrain} + save_steps: 0 + save_strategy: ${SAVE_STRATEGY:dit_only} + run_name: flux.1_schnell_t2i-pretrain + + data: + dataset_path: ${DATASET_PATH:} + empty_encodings_path: ${EMPTY_ENCODINGS_PATH:} + prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1} + img_size: ${IMG_SIZE:256} + + parallelism: + sp_size: 1 + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: ${LR:2.0e-4} + weight_decay: ${WEIGHT_DECAY:0.1} + adam_beta2: ${ADAM_BETA2:0.95} + + lr_scheduler: + lr_scheduler_type: constant_with_warmup + warmup_steps: ${WARMUP_STEPS:1600} + + runtime: + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} + gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false} + compile_transformer_blocks: ${COMPILE_TRANSFORMER_BLOCKS:true} + fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true} + report_to: none diff --git a/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml b/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml new file mode 100644 index 000000000..3b8798bc8 --- /dev/null +++ b/examples/diffusion/configs/MI355X/flux.1_schnell_t2i-raw-pretrain.yaml @@ -0,0 +1,69 @@ +work_group: ${PRIMUS_TEAM:local} +user_name: ${PRIMUS_USER:local} +exp_name: ${PRIMUS_EXP_NAME:flux.1_schnell_t2i-raw-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +platform: + config: platform_local.yaml + +modules: + pre_trainer: + framework: diffusion + config: pre_trainer.yaml + + model: flux.1_schnell_t2i.yaml + overrides: + sink_level: null + file_sink_level: INFO + stderr_sink_level: INFO + + model: + config: + encoder: + t5_encoder: ${T5_ENCODER:google/t5-v1_1-xxl} + clip_encoder: ${CLIP_ENCODER:openai/clip-vit-large-patch14} + autoencoder: ${VAE_CHECKPOINT:black-forest-labs/FLUX.1-dev/ae.safetensors} + max_t5_length: ${MAX_T5_LENGTH:256} + + metrics: + log_freq: 1 + enable_wandb: false + + training: + local_batch_size: ${LOCAL_BATCH_SIZE:1} + steps: ${MAX_STEPS:50} + num_train_epochs: 1000 + gradient_accumulation_steps: 1 + output_dir: ${OUTPUT_DIR:./output/flux.1_schnell_t2i-raw-pretrain} + save_steps: 0 + save_strategy: ${SAVE_STRATEGY:dit_only} + run_name: flux.1_schnell_t2i-raw-pretrain + + data: + dataset_type: raw + dataset: ${DATASET:cc12m-test} + dataset_format: ${DATASET_FORMAT:webdataset} + dataset_path: ${DATASET_PATH:} + prompt_dropout_prob: ${PROMPT_DROPOUT_PROB:0.1} + img_size: ${IMG_SIZE:256} + skip_low_resolution: false + + parallelism: + sp_size: 1 + dp_replicate: ${DP_REPLICATE:1} + + optimizer: + lr: ${LR:2.0e-4} + weight_decay: ${WEIGHT_DECAY:0.1} + adam_beta2: ${ADAM_BETA2:0.95} + + lr_scheduler: + lr_scheduler_type: constant_with_warmup + warmup_steps: ${WARMUP_STEPS:1600} + + runtime: + attention_backend: ${ATTENTION_BACKEND:flash_attn_aiter} + gradient_checkpointing: ${GRADIENT_CHECKPOINTING:false} + compile_transformer_blocks: ${COMPILE_TRANSFORMER_BLOCKS:true} + fsdp2_reshard_after_forward: ${FSDP2_RESHARD_AFTER_FORWARD:true} + report_to: none diff --git a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml index c7f93c421..cbd877eb8 100644 --- a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml +++ b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml @@ -15,7 +15,7 @@ modules: model: wan2.1_t2v_1.3b_sft.yaml overrides: sink_level: null - file_sink_level: DEBUG + file_sink_level: INFO stderr_sink_level: INFO metrics: diff --git a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml index 5f17780f1..c13a2cf0b 100644 --- a/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml +++ b/examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml @@ -15,7 +15,7 @@ modules: model: wan2.1_t2v_1.3b.yaml overrides: sink_level: null - file_sink_level: DEBUG + file_sink_level: INFO stderr_sink_level: INFO metrics: diff --git a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml index 203aa13ff..94ee0d09e 100644 --- a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml +++ b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml @@ -15,7 +15,7 @@ modules: model: wan2.2_ti2v_5b_sft.yaml overrides: sink_level: null - file_sink_level: DEBUG + file_sink_level: INFO stderr_sink_level: INFO metrics: diff --git a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml index 7eed61bdd..765a454b5 100644 --- a/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml +++ b/examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml @@ -15,7 +15,7 @@ modules: model: wan2.2_ti2v_5b.yaml overrides: sink_level: null - file_sink_level: DEBUG + file_sink_level: INFO stderr_sink_level: INFO metrics: diff --git a/primus/backends/diffusion/README.md b/primus/backends/diffusion/README.md index 9b7db9704..88e29227b 100644 --- a/primus/backends/diffusion/README.md +++ b/primus/backends/diffusion/README.md @@ -1,237 +1,76 @@ # Primus Diffusion Backend -`diffusion` integrates PyTorch diffusion-model training as an independent Primus -backend. Primus provides the config/launch entrypoint, while model, dataset, -attention, and FSDP2 training logic are owned by the in-tree Wan implementation -under `primus/backends/diffusion`. +`diffusion` is an in-tree PyTorch backend for Wan video training and FLUX +text-to-image training. Primus owns config loading and launch; this backend owns +model construction, datasets, attention selection, FSDP2 wrapping, and the +training loop. -All runtime code resolves through the Primus namespace, for example -`primus.backends.diffusion.models` and `primus.backends.diffusion.trainers`, -so Wan training is a first-class part of the Primus `diffusion` backend. +## Supported Scope -Supported scope: +- Models: `wan`, `flux.1-schnell`, and `flux.1-dev`. +- Trainer: FSDP2. +- Wan sequence parallelism: supported through `sp_size`. +- FLUX sequence parallelism: not supported; keep `sp_size: 1`. -- Model implementation: `wan` for Wan2.1 and Wan2.2. -- Trainer: FSDP2 only. -- Sequence parallelism: Ulysses SP via `trainer.args.sp_size`. - -Wan-specific dependencies are kept out of top-level Primus requirements. See -`runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt`. - -## Data - -Wan training reads a jsonl metadata file and a video folder: - -```jsonl -{"prompt": "text prompt", "video": "example.mp4"} -``` - -The default small dataset for smoke tests is the Hugging Face dataset -`zirui3/tiny-video-samples`: +Install backend dependencies with: ```bash -huggingface-cli download zirui3/tiny-video-samples \ - --repo-type dataset \ - --local-dir /data/tiny-video-samples -``` - -This produces: - -```text -/data/tiny-video-samples/ - meta.jsonl - data/*.mp4 -``` - -Use these public config fields to point training at another dataset: - -```yaml -data: - dataset_path: /path/to/meta.jsonl - data_folder: /path/to/videos - video_backend: imageio +pip install -r runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt ``` -## Checkpoints - -Each Wan model is a single Hugging Face repo that already bundles everything Wan -training needs: the DiT weights, the UMT5-XXL text encoder -(`models_t5_umt5-xxl-enc-bf16.pth`), the VAE (`Wan2.1_VAE.pth` / -`Wan2.2_VAE.pth`), and the tokenizer under `google/umt5-xxl`. Download the -model(s) you plan to train: - -| Model | Preset | Hugging Face repo | -| --- | --- | --- | -| Wan2.1-T2V-1.3B | `wan2.1_t2v_1.3b.yaml` | [Wan-AI/Wan2.1-T2V-1.3B](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B) | -| Wan2.2-TI2V-5B | `wan2.2_ti2v_5b.yaml` | [Wan-AI/Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) | -| Wan2.1-T2V-14B | (no shipped preset) | [Wan-AI/Wan2.1-T2V-14B](https://huggingface.co/Wan-AI/Wan2.1-T2V-14B) | - -```bash -# Wan2.1-T2V-1.3B -huggingface-cli download Wan-AI/Wan2.1-T2V-1.3B \ - --local-dir /models/Wan2.1-T2V-1.3B - -# Wan2.2-TI2V-5B -huggingface-cli download Wan-AI/Wan2.2-TI2V-5B \ - --local-dir /models/Wan2.2-TI2V-5B -``` +## Model Presets -A downloaded Wan repo looks like this (the T5 encoder, VAE, and tokenizer are -shipped inside the same repo, so no separate download is required): +Model presets live under `primus/configs/models/diffusion/`. ```text -/models/Wan2.1-T2V-1.3B/ - config.json - diffusion_pytorch_model*.safetensors # DiT weights - models_t5_umt5-xxl-enc-bf16.pth # T5 (UMT5-XXL) text encoder - Wan2.1_VAE.pth # VAE (Wan2.2_VAE.pth in the 5B repo) - google/umt5-xxl/ # tokenizer -``` - -The model presets reference these paths by default. Override any asset with -environment variables: - -```bash -export PRETRAINED_PATH=/models/Wan2.1-T2V-1.3B # pretrain DiT init -export INIT_CHECKPOINT=/models/Wan2.1-T2V-1.3B # post-train/SFT DiT init -export TEXT_TOKENIZER=/models/Wan2.1-T2V-1.3B/google/umt5-xxl -export TEXT_ENCODER=/models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth -export VAE_CHECKPOINT=/models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth -``` - -Pretrain presets use `PRETRAINED_PATH`; post-train SFT presets use -`INIT_CHECKPOINT` for DiT initialization. For Wan2.2-5B use the matching -`/models/Wan2.2-TI2V-5B` directory and `Wan2.2_VAE.pth`. - -## Pre-flight validation - -The diffusion prepare hook validates the prepared assets; it does **not** -download them. Download the dataset and checkpoints above first, then run the -hook to confirm the configured dataset, tokenizer, and DiT/T5/VAE paths exist -before launching distributed training: - -```bash -# pretrain config (validates modules.pre_trainer) -python3 runner/helpers/hooks/train/pretrain/diffusion/prepare.py \ - --config examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml - -# post-train config (validates modules.post_trainer) -python3 runner/helpers/hooks/train/posttrain/diffusion/prepare.py \ - --config examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml -``` - -On success it prints `env.PREPARED=1`. Set `SKIP_PREPARE=1` to bypass the check -for debugging. - -## Launch - -Training runs through the Primus CLI (`primus.cli.main train `) -under `torchrun`. Single-node and multi-node share **one** launch command: the -same `torchrun` invocation runs on every node, parameterized by the standard -rendezvous variables. The defaults below give a single-node 8-GPU run. - -```bash -# distributed knobs (defaults = single node, 8 GPUs) -export NNODES=${NNODES:-1} -export NODE_RANK=${NODE_RANK:-0} -export MASTER_ADDR=${MASTER_ADDR:-127.0.0.1} -export MASTER_PORT=${MASTER_PORT:-29500} -export GPUS_PER_NODE=${GPUS_PER_NODE:-8} - -torchrun \ - --nnodes="$NNODES" --node_rank="$NODE_RANK" \ - --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \ - --nproc_per_node="$GPUS_PER_NODE" \ - -m primus.cli.main train pretrain --config /path/to/wan_config.yaml +wan2.1_t2v_1.3b.yaml +wan2.1_t2v_1.3b_sft.yaml +wan2.2_ti2v_5b.yaml +wan2.2_ti2v_5b_sft.yaml +flux.1_schnell_t2i.yaml +flux.1_dev_t2i.yaml ``` -- **Single node**: run as-is (the defaults above). -- **Multi-node**: run the same command on each node with a shared - `MASTER_ADDR`/`MASTER_PORT` (a routable IP of node rank 0) and a distinct - `NODE_RANK` per node. World size is `NNODES * GPUS_PER_NODE`. -- **Post-train**: identical command with `train posttrain` and a posttrain config. +FLUX.1-schnell and FLUX.1-dev are separate presets. They share the same main +transformer shape, but `flux.1-dev` has a guidance embedding module while +`flux.1-schnell` does not. Select the preset by choosing +`flux.1_schnell_t2i.yaml` or `flux.1_dev_t2i.yaml`. -Useful runtime knobs: +This follows the TorchTitan-style separation where `flux_schnell()` is its own +config/preset and architecture differences such as `guidance_embed=False` are +part of that preset, not a launch-time switch. -- `trainer.args.attention_backend`: defaults to `flash_attn_aiter` for Wan training; use `sdpa` as the portable fallback or baseline. -- `trainer.args.sp_size`: Ulysses sequence parallel size. It must divide the - model attention head count; for example Wan2.1-1.3B supports `sp_size=4` but - not `sp_size=8`. -- `trainer.args.dp_replicate`: data parallel replication size. -- `FIXED_TIMESTEP` and `FIXED_SEED`: optional debug variables for reproducible - loss-alignment checks. +## Public Config Sections -On ROCm clusters, point compiler/cache directories to a large filesystem: +Diffusion examples use Primus override sections that the backend converts into +model, dataset, and trainer args: -```bash -export TMPDIR=/path/to/large/tmp -export TRITON_CACHE_DIR=/path/to/large/cache/triton -export TORCHINDUCTOR_CACHE_DIR=/path/to/large/cache/inductor -export AMD_COMGR_CACHE_DIR=/path/to/large/cache/comgr -``` +```yaml +training: + local_batch_size: 1 + steps: 50 + output_dir: ./output/run -## Primus-Style Configs +data: + dataset_path: /path/to/data -New examples should use Primus-style override sections and let the diffusion -adapter normalize them into Wan args. See: +parallelism: + sp_size: 1 + dp_replicate: 1 -```text -examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-pretrain.yaml -examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-pretrain.yaml -examples/diffusion/configs/MI355X/wan2.1_t2v_1.3b-posttrain.yaml -examples/diffusion/configs/MI355X/wan2.2_ti2v_5b-posttrain.yaml -``` +optimizer: + lr: 2.0e-4 + weight_decay: 0.1 -Minimal Wan2.1-1.3B shape: +lr_scheduler: + lr_scheduler_type: constant_with_warmup + warmup_steps: 1600 -```yaml -work_group: local -user_name: local -exp_name: wan2.1_t2v_1.3b-pretrain -workspace: ./output - -platform: - config: platform_local.yaml - -modules: - pre_trainer: - framework: diffusion - config: pre_trainer.yaml - model: wan2.1_t2v_1.3b.yaml - overrides: - metrics: - log_freq: 1 - enable_wandb: false - training: - local_batch_size: 1 - steps: 100 - gradient_accumulation_steps: 1 - output_dir: ./output/wan2.1_t2v_1.3b-pretrain - save_steps: 0 - data: - dataset_path: /data/tiny-video-samples/meta.jsonl - data_folder: /data/tiny-video-samples/data - text_tokenizer: /models/Wan2.1-T2V-1.3B/google/umt5-xxl - height: 480 - width: 832 - parallelism: - sp_size: 1 - dp_replicate: 1 - runtime: - attention_backend: flash_attn_aiter - report_to: none +runtime: + attention_backend: flash_attn_aiter + gradient_checkpointing: false + report_to: none ``` -Minimal Wan2.2-5B uses the same shape with `model: wan2.2_ti2v_5b.yaml` and -the Wan2.2 tokenizer path: - -```yaml -modules: - pre_trainer: - framework: diffusion - config: pre_trainer.yaml - model: wan2.2_ti2v_5b.yaml - overrides: - data: - text_tokenizer: /models/Wan2.2-TI2V-5B/google/umt5-xxl -``` +`examples/diffusion/README.md` contains data preparation notes and launch +commands for the shipped examples. diff --git a/primus/backends/diffusion/argument_builder.py b/primus/backends/diffusion/argument_builder.py index b6b444ceb..92f6a0a70 100644 --- a/primus/backends/diffusion/argument_builder.py +++ b/primus/backends/diffusion/argument_builder.py @@ -13,8 +13,8 @@ from primus.core.utils.yaml_utils import nested_namespace_to_dict -class WanArgBuilder: - """Build the compact config object consumed by the Wan trainer.""" +class DiffusionArgBuilder: + """Build the compact config object consumed by diffusion trainers.""" DEFAULT_DATASET: dict[str, Any] = { "name": "wan", @@ -45,7 +45,7 @@ class WanArgBuilder: }, }, } - DEFAULT_TRAINER: dict[str, Any] = { + DEFAULT_WAN_TRAINER: dict[str, Any] = { "name": "fsdp2", "args": { "output_dir": "./output/wan", @@ -86,6 +86,68 @@ class WanArgBuilder: }, }, } + DEFAULT_FLUX_DATASET: dict[str, Any] = { + "name": "flux", + "config": { + "dataset_type": "precomputed", + "dataset_format": "hf_dataset", + "dataset_path": "/path/to/flux_precomputed_dataset", + "dataset": None, + "shuffle": True, + "processor_config": { + "processor_name": "flux_precomputed", + "processor_type": "flux_precomputed", + "prompt_dropout_prob": 0.0, + "empty_encodings_path": None, + "img_size": 256, + "skip_low_resolution": True, + }, + }, + } + DEFAULT_FLUX_TRAINER: dict[str, Any] = { + "name": "fsdp2", + "args": { + "output_dir": "./output/flux", + "per_device_train_batch_size": 1, + "per_device_eval_batch_size": 1, + "gradient_accumulation_steps": 1, + "gradient_checkpointing": False, + "attention_backend": "flash_attn_aiter", + "learning_rate": 2.0e-4, + "lr_scheduler_type": "constant_with_warmup", + "warmup_steps": 1600, + "weight_decay": 0.1, + "num_train_epochs": 1, + "max_steps": 100, + "logging_steps": 1, + "save_steps": 0, + "dataloader_num_workers": 4, + "report_to": "none", + "run_name": "flux-schnell-fsdp2", + "bf16": True, + "seed": 10007, + "optim": "adamw_torch", + "adam_beta1": 0.9, + "adam_beta2": 0.95, + "adam_epsilon": 1.0e-8, + "max_grad_norm": 1.0, + "fsdp2_wrap_target": "dit", + "fsdp_transformer_layer_cls_to_wrap": "DoubleStreamBlock,SingleStreamBlock", + "fsdp_module_paths_to_wrap": "img_in,time_in,vector_in,txt_in,final_layer", + "fsdp_module_paths_no_reshard": "final_layer", + "fsdp2_reshard_after_forward": True, + "compile_transformer_blocks": True, + "save_strategy": "dit_only", + "sp_size": 1, + "dp_replicate": 1, + "flow_match_scheduler": { + "shift": 1, + "sigma_min": 0.0, + "extra_one_step": False, + "num_train_timesteps": 1000, + }, + }, + } def __init__(self) -> None: self._params: dict[str, Any] = {} @@ -96,16 +158,18 @@ def update(self, params: Any) -> None: elif isinstance(params, dict): self._params = copy.deepcopy(params) else: - raise TypeError(f"WanArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}") + raise TypeError( + f"DiffusionArgBuilder expects dict or SimpleNamespace, got {type(params).__name__}" + ) def finalize(self) -> SimpleNamespace: params = copy.deepcopy(self._params) if "model" not in params: - raise ValueError("Wan backend config requires a model preset.") + raise ValueError("Diffusion backend config requires a model preset.") for legacy_section in ("dataset", "trainer"): if legacy_section in params: raise ValueError( - f"Wan backend no longer accepts public `{legacy_section}` overrides. " + f"Diffusion backend no longer accepts public `{legacy_section}` overrides. " "Use Primus-style `data`, `training`, `parallelism`, `optimizer`, " "`runtime`, and `metrics` sections instead." ) @@ -138,18 +202,48 @@ def _get_any(source: dict[str, Any], *keys: str) -> Any: return source[key] return None + @staticmethod + def _coerce_bool(value: Any) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + return bool(value) + + def _model_family(self, params: dict[str, Any]) -> str: + model = params.get("model") or {} + if not isinstance(model, dict): + raise TypeError(f"Diffusion model preset must be a dict, got {type(model).__name__}") + name = str(model.get("name", "")).strip().lower() + if not name: + raise ValueError("Diffusion model preset requires `model.name`.") + return name + + def _defaults_for_model(self, model_name: str) -> tuple[dict[str, Any], dict[str, Any]]: + if model_name == "wan": + return self.DEFAULT_DATASET, self.DEFAULT_WAN_TRAINER + if model_name == "flux" or model_name.startswith("flux."): + return self.DEFAULT_FLUX_DATASET, self.DEFAULT_FLUX_TRAINER + raise ValueError(f"Unsupported diffusion model name: {model_name!r}") + def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, Any]: - """Translate concise Primus-style Wan sections into trainer arguments. + """Translate concise Primus-style sections into trainer arguments. Public configs use high-level sections such as `training`, `data`, - `parallelism`, and `runtime`; internal defaults supply the compact - `dataset` and `trainer` objects consumed by the Wan runtime. + `parallelism`, and `runtime`; internal defaults supply compact + `dataset` and `trainer` objects consumed by the diffusion runtime. """ + model_name = self._model_family(params) + default_dataset, default_trainer = self._defaults_for_model(model_name) normalized = { "model": params["model"], - "dataset": copy.deepcopy(self.DEFAULT_DATASET), - "trainer": copy.deepcopy(self.DEFAULT_TRAINER), + "dataset": copy.deepcopy(default_dataset), + "trainer": copy.deepcopy(default_trainer), "stage": params.get("stage", "pretrain"), "primus": params.get("primus", {}), } @@ -161,6 +255,7 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, data = params.get("data") or {} parallelism = params.get("parallelism") or {} optimizer = params.get("optimizer") or {} + lr_scheduler = params.get("lr_scheduler") or {} runtime = params.get("runtime") or {} metrics = params.get("metrics") or {} @@ -171,6 +266,7 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, ("gradient_accumulation_steps",): ("gradient_accumulation_steps",), ("output_dir",): ("output_dir",), ("save_steps",): ("save_steps",), + ("save_strategy",): ("save_strategy",), ("run_name",): ("run_name",), ("num_train_epochs",): ("num_train_epochs",), ("dataloader_num_workers",): ("dataloader_num_workers",), @@ -189,6 +285,14 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, ("text_tokenizer",): ("processor_config", "text_tokenizer"), ("processor_name",): ("processor_config", "processor_name"), ("processor_type",): ("processor_config", "processor_type"), + ("dataset_format",): ("dataset_format",), + ("dataset_type",): ("dataset_type",), + ("dataset",): ("dataset",), + ("shuffle",): ("shuffle",), + ("empty_encodings_path",): ("processor_config", "empty_encodings_path"), + ("prompt_dropout_prob",): ("processor_config", "prompt_dropout_prob"), + ("img_size",): ("processor_config", "img_size"), + ("skip_low_resolution",): ("processor_config", "skip_low_resolution"), } for source_path, target_path in data_map.items(): value = self._get_any(data, *source_path) @@ -233,15 +337,36 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, if value is not None: self._set_nested(trainer_args, target_path, value) + lr_scheduler_map = { + ("lr_scheduler_type",): ("lr_scheduler_type",), + ("warmup_steps",): ("warmup_steps",), + } + for source_path, target_path in lr_scheduler_map.items(): + value = self._get_any(lr_scheduler, *source_path) + if value is not None: + self._set_nested(trainer_args, target_path, value) + runtime_map = { ("attention_backend",): ("attention_backend",), ("report_to",): ("report_to",), ("seed",): ("seed",), + ("bf16",): ("bf16",), + ("fp16",): ("fp16",), + ("gradient_checkpointing",): ("gradient_checkpointing",), + ("compile_transformer_blocks",): ("compile_transformer_blocks",), ("fsdp2_reshard_after_forward",): ("fsdp2_reshard_after_forward",), } for source_path, target_path in runtime_map.items(): value = self._get_any(runtime, *source_path) if value is not None: + if target_path in { + ("bf16",), + ("fp16",), + ("gradient_checkpointing",), + ("compile_transformer_blocks",), + ("fsdp2_reshard_after_forward",), + }: + value = self._coerce_bool(value) self._set_nested(trainer_args, target_path, value) log_freq = metrics.get("log_freq") @@ -259,4 +384,13 @@ def _normalize_primus_style_sections(self, params: dict[str, Any]) -> dict[str, if resume_from_checkpoint is not None: self._set_nested(trainer_args, ("resume_from_checkpoint",), resume_from_checkpoint) + if (model_name == "flux" or model_name.startswith("flux.")) and int( + trainer_args.get("sp_size", 1) + ) != 1: + raise ValueError("FLUX diffusion training currently requires `parallelism.sp_size: 1`.") + return normalized + + +# Backwards-compatible import path for existing Wan-only callers. +WanArgBuilder = DiffusionArgBuilder diff --git a/primus/backends/diffusion/attention/attention.py b/primus/backends/diffusion/attention/attention.py index 4e67b59b9..cbe8467ef 100644 --- a/primus/backends/diffusion/attention/attention.py +++ b/primus/backends/diffusion/attention/attention.py @@ -319,7 +319,7 @@ def attention( resolved = _resolve_flash_version(q.device.type) if resolved is not None: version = resolved - # Only "auto" allows call-site override (useful for debugging). + # Only "auto" allows call-site override. if _ATTENTION_BACKEND == "auto" and fa_version is not None: version = fa_version return flash_attention( diff --git a/primus/backends/diffusion/data/__init__.py b/primus/backends/diffusion/data/__init__.py index 18ae84901..b72452684 100644 --- a/primus/backends/diffusion/data/__init__.py +++ b/primus/backends/diffusion/data/__init__.py @@ -4,4 +4,24 @@ # See LICENSE for license information. ############################################################################### -"""Data builders for the diffusion backend.""" +"""Data modules for the Primus diffusion backend.""" + +from .config import DatasetConfig +from .dataset import WanVideoDataset +from .flux_precomputed import ( + FluxPrecomputedDataset, + FluxPrecomputedProcessor, + FluxRawImageTextDataset, + FluxRawImageTextProcessor, +) +from .processor import WanVideoDataProcessor + +__all__ = [ + "DatasetConfig", + "FluxPrecomputedDataset", + "FluxPrecomputedProcessor", + "FluxRawImageTextDataset", + "FluxRawImageTextProcessor", + "WanVideoDataProcessor", + "WanVideoDataset", +] diff --git a/primus/backends/diffusion/data/collator.py b/primus/backends/diffusion/data/collator.py new file mode 100644 index 000000000..fb6d5fba1 --- /dev/null +++ b/primus/backends/diffusion/data/collator.py @@ -0,0 +1,98 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import collections +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass +class VisionCollator: + processor: Any + + def pad_sequence(self, input_ids, batch_first, padding_value): + if self.processor.tokenizer.padding_side == "left": + input_ids = [torch.flip(_input_ids, [0]) for _input_ids in input_ids] + input_ids = torch.nn.utils.rnn.pad_sequence( + input_ids, batch_first=batch_first, padding_value=padding_value + ) + if self.processor.tokenizer.padding_side == "left": + input_ids = torch.flip(input_ids, [1]) + return input_ids + + def __call__(self, instances: Sequence[dict]) -> dict[str, torch.Tensor]: + if isinstance(instances[0], list): + instances = [inst for instance in instances for inst in instance] + inputs = collections.defaultdict(list) + for instance in instances: + for key, values in instance.items(): + inputs[key].append(values) + + batched_inputs = {} + input_ids = None + if "input_ids" in inputs: + input_ids = inputs.pop("input_ids") + input_ids = self.pad_sequence( + input_ids, + batch_first=True, + padding_value=self.processor.tokenizer.pad_token_id, + ) + batched_inputs["input_ids"] = input_ids + if "labels" in inputs: + labels = inputs.pop("labels") + labels = self.pad_sequence( + labels, + batch_first=True, + padding_value=-100, + ) + batched_inputs["labels"] = labels + + attention_mask = None + if "attention_mask" in inputs: + attention_mask = inputs.pop("attention_mask") + + if input_ids is not None: + batched_inputs["attention_mask"] = input_ids.ne(self.processor.tokenizer.pad_token_id).long() + elif attention_mask is not None: + batched_inputs["attention_mask"] = self.pad_sequence( + attention_mask, + batch_first=True, + padding_value=0, + ) + + # for the other keys + for key, values in inputs.items(): + # Handle scalar/boolean values ( use_audio_in_video) + if isinstance(values[0], bool) or ( + isinstance(values[0], (int, float)) and not isinstance(values[0], torch.Tensor) + ): + batched_inputs[key] = values[0] + else: + batched_inputs[key] = torch.stack(values, dim=0) + return batched_inputs + + @property + def image_token_id(self): + return self.processor.tokenizer.convert_tokens_to_ids(self.processor.image_token) + + +@dataclass +class RawBatchCollator: + """ + A minimal collator that returns raw samples as a list of dicts. + + This is useful when model-specific padding/encoding happens in a separate + batch preparation step (e.g. processor.prepare_batch / trainer.prepare_batch), + keeping the DataLoader and trainer model-agnostic. + """ + + def __call__(self, instances: Sequence[dict]) -> list[dict]: + if isinstance(instances[0], list): + instances = [inst for instance in instances for inst in instance] + return list(instances) diff --git a/primus/backends/diffusion/data/config.py b/primus/backends/diffusion/data/config.py new file mode 100644 index 000000000..e425149ba --- /dev/null +++ b/primus/backends/diffusion/data/config.py @@ -0,0 +1,89 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from typing import Any, Literal + +from pydantic import BaseModel, field_validator + + +class Args(BaseModel): + extra_kwargs: dict[str, Any] = {} + + def to_dict(self): + return self.model_dump() + + def to_json(self): + return self.model_dump_json() + + +class ProcessorConfig(Args): + processor_name: str + processor_type: str + + +class DatasetConfig(Args): + dataset_type: str + data_folder: str + dataset_format: Literal["json", "jsonl", "csv", "yaml", "hf_dataset", "arrow"] + processor_config: dict | ProcessorConfig + + # Dataset configuration + dataset_path: str | None = None # Optional - used for external files + datasets: list[dict] | None = None # Optional - used for inline YAML definitions + shuffle: bool = True + data_seed: int | None = 42 + eval_dataset_path: str | None = None + + # Object storage configuration + object_storage: Literal["azure", "gcs", "none"] | None = "none" + bucket_name: str | None = None + + # Packing configuration + packing: bool | None = False + packing_strategy: str | None = None + packing_length: int | None = 32000 + filter_overlong: bool | None = True + filter_overlong_workers: int | None = 8 + max_length: int | None = None + + # Video configuration + video_sampling_strategy: Literal["fps", "frame_num"] | None = "fps" + video_max_pixels: int | None = 768 * 28 * 28 + video_max_frames: int | None = 768 + video_min_pixels: int | None = 3136 + frame_num: int | None = 64 + fps: int | None = 1 + video_backend: Literal["decord", "qwen_vl_utils", "qwen_omni_utils", "imageio"] | None = "qwen_vl_utils" + + @field_validator( + "video_max_pixels", + "video_max_frames", + "frame_num", + "fps", + "packing_length", + "max_length", + "filter_overlong_workers", + ) + @classmethod + def validate_positive_values(cls, v, info): + """Validate that numeric video and packing parameters are positive.""" + if v is not None and v <= 0: + field_name = info.field_name + raise ValueError(f"{field_name} must be positive, got {v}") + return v + + @field_validator("video_backend") + @classmethod + def validate_video_backend_migration(cls, v): + """Provide migration warning for deprecated torchvision backend.""" + if v == "torchvision": + raise ValueError( + "The 'torchvision' video backend has been removed. " + "Please use 'decord', 'qwen_vl_utils', or 'qwen_omni_utils' instead. " + "Migration guide: If you were using torchvision, 'decord' provides " + "similar functionality with better performance." + ) + return v diff --git a/primus/backends/diffusion/data/dataset.py b/primus/backends/diffusion/data/dataset.py new file mode 100644 index 000000000..c5584dbdc --- /dev/null +++ b/primus/backends/diffusion/data/dataset.py @@ -0,0 +1,291 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Dataset class for loading video data from JSONL/CSV files. +""" + +import json +import os +from io import BytesIO +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +from primus.backends.diffusion.utils.data_utils import smart_nframes +from primus.backends.diffusion.utils.vision_process import strip_file_uri + +from .collator import RawBatchCollator + + +def _get_decord_vr( + video_path: str, + *, + num_threads: int, +): + """Construct a decord.VideoReader.""" + from decord import VideoReader, cpu + + return VideoReader(video_path, ctx=cpu(0), num_threads=num_threads) + + +class BaseDataset(Dataset): + """Base dataset with minimal shared interface.""" + + def __init__(self, config, **kwargs) -> None: + super().__init__() + self.config = config + self.samples = [] + + def __len__(self): + return len(self.samples) + + +class WanVideoDataset(BaseDataset): + """Dataset for WanVideo training from JSONL/CSV.""" + + def __init__(self, processor, config=None): + """ + Initialize WanVideo dataset. + + Args: + processor: WanVideoDataProcessor instance + video_backend: Backend for video loading ('qwen_vl_utils' or 'decord') + """ + if config is None: + raise ValueError("WanVideoDataset requires a dataset config with dataset_path") + if getattr(config, "dataset_path", None) is None: + raise ValueError("WanVideoDataset requires config.dataset_path") + + super().__init__(config) + self.config = config + self.data_path = Path(self.config.dataset_path) + self.processor = processor + + # Load metadata + self.samples = self._load_metadata() + self._sync_processor_video_limits() + + def _sync_processor_video_limits(self): + image_processor = None + if hasattr(self.processor, "processor") and hasattr(self.processor.processor, "image_processor"): + image_processor = self.processor.processor.image_processor + if image_processor is None: + return + if ( + getattr(image_processor, "max_pixels", None) is None + and getattr(self.config, "video_max_pixels", None) is not None + ): + image_processor.max_pixels = self.config.video_max_pixels + + def _load_metadata(self) -> list[dict]: + """Load metadata from JSONL or CSV file.""" + samples = [] + + if self.data_path.suffix == ".jsonl": + with open(self.data_path) as f: + for line in f: + samples.append(json.loads(line.strip())) + elif self.data_path.suffix == ".json": + with open(self.data_path) as f: + samples = json.load(f) + elif self.data_path.suffix == ".csv": + import pandas as pd + + df = pd.read_csv(self.data_path) + samples = df.to_dict("records") + else: + raise ValueError(f"Unsupported file format: {self.data_path=}") + + return samples + + def _load_video_frames(self, video_path: str, data_folder=None, fps: int = 1) -> tuple[np.ndarray, float]: + """Load video frames using the specified backend.""" + if self.config.data_folder is not None: + video_path = os.path.join(self.config.data_folder, video_path) + + if self.config.video_backend == "decord": + return self.load_video_decord(video_path, fps) + elif self.config.video_backend == "qwen_vl_utils": + return self.load_video_qwen_vl_utils(video_path, fps) + elif self.config.video_backend == "imageio": + return self.load_video_imageio(video_path, fps) + else: + raise ValueError(f"Unsupported video backend: {self.config.video_backend}") + + def load_video_imageio(self, video_path, fps): + import imageio + + video_path = strip_file_uri(video_path) + reader = imageio.get_reader(video_path) + + # Sampling Strategy + total_frames = reader.count_frames() + total_frames = int(total_frames) + + if self.config.video_sampling_strategy == "frame_num": + nframes = self.config.frame_num + # Enforce VAE divisibility: (n - 1) % 4 == 0 + actual_nframes = min(nframes, total_frames) + + valid_nframes = (actual_nframes - 1) // 4 * 4 + 1 if actual_nframes > 1 else 1 + + # DiffSynth Sequential Reading + frames = [] + for i, frame in enumerate(reader): + if i >= valid_nframes: + break + frames.append(frame) + + # Stack to numpy (T, H, W, C) + frames = np.array(frames) + sample_fps = fps # Simplification + + reader.close() + else: + reader.close() + raise NotImplementedError("Only frame_num strategy implemented for imageio backend") + + return frames, sample_fps + + def load_video_decord( + self, + video_path: str | list[str] | BytesIO, + fps: int, + ) -> tuple[np.ndarray, float]: + """ + Load video using Decord backend. + + Args: + video_path: Path to video file or BytesIO object + fps: Target frames per second + + Returns: + Tuple of (video frames, sample fps) + """ + # Keep dataset logic simple: use a fixed thread count here. + from decord import VideoReader, cpu + + num_threads = 4 + + if isinstance(video_path, BytesIO): + vr = VideoReader(video_path, ctx=cpu(0), num_threads=num_threads) + elif isinstance(video_path, list): + vr = _get_decord_vr(strip_file_uri(video_path[0]), num_threads=num_threads) + elif isinstance(video_path, str): + vr = _get_decord_vr(strip_file_uri(video_path), num_threads=num_threads) + else: + raise ValueError(f"Unsupported video path type: {type(video_path)}") + + total_frames, video_fps = len(vr), vr.get_avg_fps() + if self.config.video_sampling_strategy == "fps": + nframes = smart_nframes(total_frames, video_fps=video_fps, fps=fps) + # Maintain uniform sampling for FPS strategy + uniform_sampled_frames = np.linspace(0, total_frames - 1, nframes, dtype=int) + elif self.config.video_sampling_strategy == "frame_num": + nframes = self.config.frame_num + # Enforce VAE divisibility: (n - 1) % 4 == 0 + actual_nframes = min(nframes, total_frames) + + valid_nframes = (actual_nframes - 1) // 4 * 4 + 1 if actual_nframes > 1 else 1 + uniform_sampled_frames = np.linspace(0, total_frames - 1, valid_nframes, dtype=int) + else: + raise ValueError(f"Invalid video sampling strategy: {self.config.video_sampling_strategy}") + + frame_idx = uniform_sampled_frames.tolist() + spare_frames = vr.get_batch(frame_idx).asnumpy() + # spare_frames = torch.tensor(spare_frames).permute(0, 3, 1, 2) # Convert to TCHW format + + # Calculate sample_fps + sample_fps = nframes / max(total_frames, 1e-6) * video_fps + + # Return HWC numpy array to match processor expectations + return spare_frames, sample_fps # (frames, height, width, channels) + + def load_video_qwen_vl_utils( + self, + video_path: str, + fps: int, + ) -> tuple[np.ndarray, float]: + """ + Load video using Qwen VL utils. + + Args: + video_path: Path to video file + fps: Target frames per second + + Returns: + Tuple of (video frames, sample fps) + """ + from primus.backends.diffusion.utils.vision_process import fetch_video + + video_dict = { + "type": "video", + "video": f"file://{video_path}", + "min_frames": 1, + "max_pixels": self.config.video_max_pixels, + "max_frames": self.config.video_max_frames, + "min_pixels": self.config.video_min_pixels, + } + + if self.config.video_sampling_strategy == "frame_num": + is_even = self.config.frame_num % 2 == 0 + n_frames = self.config.frame_num if is_even else self.config.frame_num + 1 + video_dict["nframes"] = n_frames + + frames, sample_fps = fetch_video(video_dict, return_video_sample_fps=True) + frames = frames.numpy() + + # Enforce VAE divisibility constraint + actual_n = len(frames) + if actual_n > 1: + valid_n = ((actual_n - 1) // 4) * 4 + 1 + frames = frames[:valid_n] + # else: keep 1 frame (or handle error) + + return frames, sample_fps + elif self.config.video_sampling_strategy == "fps": + video_dict["fps"] = fps + frames, sample_fps = fetch_video(video_dict, return_video_sample_fps=True) + frames = frames.numpy() + + # Also enforce for fps strategy + actual_n = len(frames) + if actual_n > 1: + valid_n = ((actual_n - 1) // 4) * 4 + 1 + frames = frames[:valid_n] + + return frames, sample_fps + else: + raise ValueError(f"Invalid video sampling strategy: {self.config.video_sampling_strategy}") + + def __len__(self) -> int: + return len(self.samples) + + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: + """Get a single sample.""" + sample = self.samples[idx] + + # Load video frames + video_path = sample["video"] + video_frames, fps = self._load_video_frames(video_path) + + # Get prompt + prompt = sample.get("prompt", "") + # Return raw sample. + return { + "video_frames": video_frames, # np.ndarray, typically (T, H, W, C) + "prompt": prompt, + "fps": fps, + "num_frames": int(getattr(self.config, "frame_num", 0) or 0), + "video_path": str(video_path), + } + + def get_collator(self): + # Prefer raw collation; model-specific processing should happen in processor.prepare_batch. + return RawBatchCollator() diff --git a/primus/backends/diffusion/data/flux_precomputed.py b/primus/backends/diffusion/data/flux_precomputed.py new file mode 100644 index 000000000..f2ea9ab0d --- /dev/null +++ b/primus/backends/diffusion/data/flux_precomputed.py @@ -0,0 +1,340 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import io +import json +import math +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +import torch +from PIL import Image, ImageFile +from torch.utils.data import Dataset + +from primus.backends.diffusion.data.collator import RawBatchCollator + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +def _tensor_from_serialized_numpy(value: bytes) -> torch.Tensor: + array = np.load(io.BytesIO(value)) + tensor = torch.from_numpy(array) + if tensor.dtype == torch.uint16: + tensor = tensor.view(torch.bfloat16) + return tensor + + +def _to_tensor(value: Any) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (bytes, bytearray)): + return _tensor_from_serialized_numpy(bytes(value)) + if isinstance(value, np.ndarray): + tensor = torch.from_numpy(value) + if tensor.dtype == torch.uint16: + tensor = tensor.view(torch.bfloat16) + return tensor + return torch.as_tensor(value) + + +class FluxPrecomputedDataset(Dataset): + """Map-style dataset for precomputed FLUX text and VAE encodings.""" + + required_fields = ("t5_encodings", "clip_encodings", "mean", "logvar") + + def __init__(self, dataset_path: str): + if not dataset_path: + raise ValueError("FLUX precomputed dataset requires `dataset_path`.") + if not os.path.isdir(dataset_path): + raise FileNotFoundError(f"FLUX precomputed dataset directory not found: {dataset_path}") + from datasets import load_from_disk + + self.dataset_path = dataset_path + self.dataset = load_from_disk(dataset_path) + + def __len__(self) -> int: + return len(self.dataset) + + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: + sample = self.dataset[idx] + missing = [field for field in self.required_fields if field not in sample] + if missing: + raise KeyError(f"FLUX precomputed sample missing fields: {missing}") + return {field: _to_tensor(sample[field]) for field in self.required_fields} + + def get_collator(self): + return RawBatchCollator() + + +class FluxRawImageTextDataset(Dataset): + """Map-style raw image-text dataset for online FLUX encoding.""" + + def __init__( + self, + *, + dataset_path: str | None, + dataset_format: str = "webdataset", + dataset_name: str | None = None, + data_folder: str | None = None, + ): + dataset_name = dataset_name or None + dataset_path, dataset_format = self._resolve_dataset(dataset_name, dataset_path, dataset_format) + self.dataset_path = dataset_path + self.dataset_format = dataset_format + self.dataset_name = dataset_name + self.data_folder = data_folder + self._records: list[dict[str, Any]] | None = None + + if dataset_format == "jsonl": + path = Path(dataset_path) + if not path.is_file(): + raise FileNotFoundError(f"FLUX raw jsonl metadata not found: {dataset_path}") + self._records = [json.loads(line) for line in path.read_text().splitlines() if line.strip()] + self.dataset = None + elif dataset_format == "hf_dataset": + if not os.path.isdir(dataset_path): + raise FileNotFoundError(f"FLUX raw HF dataset directory not found: {dataset_path}") + from datasets import load_from_disk + + self.dataset = load_from_disk(dataset_path) + elif dataset_format == "hf_repo": + from datasets import load_dataset + + self.dataset = load_dataset(dataset_path, split="train") + elif dataset_format == "webdataset": + if not os.path.isdir(dataset_path): + raise FileNotFoundError(f"FLUX raw webdataset directory not found: {dataset_path}") + from datasets import load_dataset + + self.dataset = load_dataset( + "webdataset", + split="train", + data_dir=dataset_path, + data_files={"train": "*.tar"}, + ) + else: + raise ValueError("FLUX raw dataset_format must be one of: jsonl, hf_dataset, hf_repo, webdataset") + + @staticmethod + def _resolve_dataset( + dataset_name: str | None, + dataset_path: str | None, + dataset_format: str, + ) -> tuple[str, str]: + if dataset_name == "cc12m-test": + return (dataset_path, dataset_format) if dataset_path else ("zirui3/cc12m-test", "hf_repo") + if dataset_name == "cc12m-wds": + return dataset_path or "pixparse/cc12m-wds", "hf_repo" + if not dataset_path: + raise ValueError("FLUX raw dataset requires either `dataset` or `dataset_path`.") + return dataset_path, dataset_format + + def __len__(self) -> int: + if self._records is not None: + return len(self._records) + return len(self.dataset) + + @staticmethod + def _prompt_from_sample(sample: dict[str, Any]) -> str: + for key in ("txt", "caption", "prompt", "text"): + if key in sample: + value = sample[key] + if isinstance(value, list): + value = value[0] + return str(value) + raise KeyError("FLUX raw sample requires one of: txt, caption, prompt, text") + + def _image_from_jsonl_sample(self, sample: dict[str, Any]) -> Image.Image: + image_key = "image" if "image" in sample else "jpg" if "jpg" in sample else "png" + image_path = Path(str(sample[image_key])) + if self.data_folder and not image_path.is_absolute(): + image_path = Path(self.data_folder) / image_path + return Image.open(image_path).convert("RGB") + + @staticmethod + def _image_from_dataset_sample(sample: dict[str, Any]) -> Image.Image: + if "image" in sample: + image = sample["image"] + elif "jpg" in sample: + image = sample["jpg"] + elif "png" in sample: + image = sample["png"] + else: + raise KeyError("FLUX raw dataset sample requires image, jpg, or png") + if isinstance(image, Image.Image): + return image.convert("RGB") + if isinstance(image, (bytes, bytearray)): + return Image.open(io.BytesIO(image)).convert("RGB") + return image.convert("RGB") + + def __getitem__(self, idx: int) -> dict[str, Any]: + if self._records is not None: + sample = self._records[idx] + image = self._image_from_jsonl_sample(sample) + else: + sample = self.dataset[idx] + image = self._image_from_dataset_sample(sample) + return {"image": image, "prompt": self._prompt_from_sample(sample)} + + def get_collator(self): + return RawBatchCollator() + + +@dataclass +class FluxPrecomputedProcessor: + config: dict[str, Any] + + def __post_init__(self): + self.prompt_dropout_prob = float(self.config.get("prompt_dropout_prob", 0.0) or 0.0) + self.empty_encodings_path = self.config.get("empty_encodings_path") + self.img_size = int(self.config.get("img_size", 256) or 256) + self._empty_t5: torch.Tensor | None = None + self._empty_clip: torch.Tensor | None = None + + def build(self): + if self.prompt_dropout_prob > 0.0: + self._load_empty_encodings() + + @staticmethod + def _normalize_empty_encoding(tensor: torch.Tensor) -> torch.Tensor: + if tensor.ndim >= 1 and tensor.shape[0] == 1: + return tensor[0] + return tensor + + def _check_empty_encoding_shapes(self, t5_encodings: torch.Tensor, clip_encodings: torch.Tensor) -> None: + # The empty encodings are broadcast-assigned into the per-sample encodings + # for dropped prompts, so their trailing (non-batch) shape must match the + # dataset encodings exactly. Fail fast with a clear message instead of a + # cryptic in-place assignment shape error mid-training. + assert self._empty_t5 is not None and self._empty_clip is not None + if self._empty_t5.shape != t5_encodings.shape[1:]: + raise ValueError( + "FLUX prompt dropout: empty T5 encoding shape " + f"{tuple(self._empty_t5.shape)} does not match the per-sample T5 encoding shape " + f"{tuple(t5_encodings.shape[1:])}. Regenerate t5_empty.npy with a matching " + "sequence length / hidden size." + ) + if self._empty_clip.shape != clip_encodings.shape[1:]: + raise ValueError( + "FLUX prompt dropout: empty CLIP encoding shape " + f"{tuple(self._empty_clip.shape)} does not match the per-sample CLIP encoding shape " + f"{tuple(clip_encodings.shape[1:])}. Regenerate clip_empty.npy with a matching shape." + ) + + def _load_empty_encodings(self): + if self._empty_t5 is not None and self._empty_clip is not None: + return + if not self.empty_encodings_path: + raise ValueError("FLUX prompt dropout requires `data.empty_encodings_path`.") + t5_path = os.path.join(self.empty_encodings_path, "t5_empty.npy") + clip_path = os.path.join(self.empty_encodings_path, "clip_empty.npy") + if not os.path.isfile(t5_path) or not os.path.isfile(clip_path): + raise FileNotFoundError( + "FLUX empty encodings must contain t5_empty.npy and clip_empty.npy " + f"under {self.empty_encodings_path}" + ) + self._empty_t5 = self._normalize_empty_encoding(torch.from_numpy(np.load(t5_path))) + self._empty_clip = self._normalize_empty_encoding(torch.from_numpy(np.load(clip_path))) + + @staticmethod + def _collate_raw(batch: Any) -> dict[str, torch.Tensor]: + if isinstance(batch, dict): + return {key: _to_tensor(value) for key, value in batch.items()} + if not isinstance(batch, Sequence) or not batch: + raise ValueError(f"FLUX prepare_batch expected non-empty sequence, got {type(batch).__name__}") + + keys = ("t5_encodings", "clip_encodings", "mean", "logvar") + collated: dict[str, torch.Tensor] = {} + for key in keys: + collated[key] = torch.stack([_to_tensor(sample[key]) for sample in batch], dim=0) + return collated + + def prepare_batch( + self, *, batch: Any, device: torch.device, dtype: torch.dtype + ) -> dict[str, torch.Tensor]: + tensors = self._collate_raw(batch) + for key, value in tensors.items(): + tensors[key] = value.to(device=device, dtype=dtype, non_blocking=True) + + if self.prompt_dropout_prob > 0.0: + self._load_empty_encodings() + assert self._empty_t5 is not None and self._empty_clip is not None + self._check_empty_encoding_shapes(tensors["t5_encodings"], tensors["clip_encodings"]) + bsz = tensors["t5_encodings"].shape[0] + drop_mask = torch.rand((bsz,), device=device) < self.prompt_dropout_prob + if drop_mask.any(): + tensors["t5_encodings"][drop_mask] = self._empty_t5.to(device=device, dtype=dtype) + tensors["clip_encodings"][drop_mask] = self._empty_clip.to(device=device, dtype=dtype) + + return tensors + + +@dataclass +class FluxRawImageTextProcessor: + config: dict[str, Any] + + def __post_init__(self): + self.prompt_dropout_prob = float(self.config.get("prompt_dropout_prob", 0.0) or 0.0) + self.img_size = int(self.config.get("img_size", 256) or 256) + self.skip_low_resolution = bool(self.config.get("skip_low_resolution", True)) + + def build(self): + return + + def _process_image(self, image: Image.Image) -> torch.Tensor | None: + width, height = image.size + if self.skip_low_resolution and (width < self.img_size or height < self.img_size): + return None + + if width == self.img_size and height == self.img_size: + resized = image + elif width >= height: + new_width, new_height = math.ceil(self.img_size / height * width), self.img_size + image = image.resize((new_width, new_height), Image.Resampling.BICUBIC) + left = int(torch.randint(0, new_width - self.img_size + 1, (1,)).item()) + resized = image.crop((left, 0, left + self.img_size, self.img_size)) + else: + new_width, new_height = self.img_size, math.ceil(self.img_size / width * height) + image = image.resize((new_width, new_height), Image.Resampling.BICUBIC) + lower = int(torch.randint(0, new_height - self.img_size + 1, (1,)).item()) + resized = image.crop((0, lower, self.img_size, lower + self.img_size)) + + if resized.mode != "RGB": + resized = resized.convert("RGB") + np_img = np.array(resized).transpose((2, 0, 1)) + return torch.tensor(np_img).float() / 255.0 * 2.0 - 1.0 + + def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) -> dict[str, Any]: + if isinstance(batch, dict): + batch = [batch] + if not isinstance(batch, Sequence) or not batch: + raise ValueError( + f"FLUX raw prepare_batch expected non-empty sequence, got {type(batch).__name__}" + ) + + images: list[torch.Tensor] = [] + prompts: list[str] = [] + for sample in batch: + image = self._process_image(sample["image"]) + if image is None: + continue + prompt = str(sample.get("prompt", "")) + if self.prompt_dropout_prob > 0.0 and torch.rand(1).item() < self.prompt_dropout_prob: + prompt = "" + images.append(image) + prompts.append(prompt) + + if not images: + raise ValueError("FLUX raw batch contained no usable images after preprocessing") + return { + "image": torch.stack(images, dim=0).to(device=device, dtype=dtype, non_blocking=True), + "prompts": prompts, + } diff --git a/primus/backends/diffusion/data/processing_wanvideo.py b/primus/backends/diffusion/data/processing_wanvideo.py new file mode 100644 index 000000000..ce1be8a32 --- /dev/null +++ b/primus/backends/diffusion/data/processing_wanvideo.py @@ -0,0 +1,401 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import torch +from PIL import Image +from transformers import AutoTokenizer +from transformers.image_processing_utils import BaseImageProcessor +from transformers.image_utils import ImageInput +from transformers.tokenization_utils_base import PreTokenizedInput, TextInput +from transformers.utils import TensorType, logging + +logger = logging.get_logger(__name__) + + +class WanVideoImageProcessor(BaseImageProcessor): + """ + Image/Video processor for WanVideo models. + + Args: + do_resize: Whether to resize the image/video frames. + size: Target size for resizing. + do_center_crop: Whether to center crop. + crop_size: Size for center cropping. + do_normalize: Whether to normalize pixel values. + image_mean: Mean values for normalization. + image_std: Standard deviation values for normalization. + do_convert_rgb: Whether to convert to RGB. + """ + + model_input_names = ["pixel_values"] + + def __init__( + self, + do_resize: bool = True, + size: Optional[Dict[str, int]] = None, + do_center_crop: bool = True, + crop_size: Optional[Dict[str, int]] = None, + do_normalize: bool = True, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + do_convert_rgb: bool = True, + max_pixels: Optional[int] = None, + height_division_factor: int = 16, + width_division_factor: int = 16, + **kwargs, + ): + super().__init__(**kwargs) + + self.do_resize = do_resize + self.size = size + self.do_center_crop = do_center_crop + self.crop_size = crop_size + self.do_normalize = do_normalize + self.image_mean = image_mean or [0.5, 0.5, 0.5] + self.image_std = image_std or [0.5, 0.5, 0.5] + self.do_convert_rgb = do_convert_rgb + self.max_pixels = max_pixels + self.height_division_factor = height_division_factor + self.width_division_factor = width_division_factor + + def resize( + self, + image: np.ndarray, + size: Dict[str, int], + **kwargs, + ) -> np.ndarray: + """Resize image or video frame.""" + from PIL import Image as PILImage + + image = PILImage.fromarray(image.astype(np.uint8)) + image = image.resize((size["width"], size["height"]), PILImage.LANCZOS) + return np.array(image) + + def _extract_hw(self, image: Union[np.ndarray, Image.Image]) -> Tuple[int, int]: + if isinstance(image, Image.Image): + width, height = image.size + return height, width + if isinstance(image, np.ndarray): + if image.ndim == 3: + return image.shape[0], image.shape[1] + if image.ndim == 2: + return image.shape[0], image.shape[1] + raise ValueError("Unsupported image type for size extraction.") + + def _dynamic_target_size(self, image: Union[np.ndarray, Image.Image]) -> Dict[str, int]: + height, width = self._extract_hw(image) + if self.max_pixels is not None and width * height > self.max_pixels: + scale = (width * height / self.max_pixels) ** 0.5 + height = int(height / scale) + width = int(width / scale) + height = height // self.height_division_factor * self.height_division_factor + width = width // self.width_division_factor * self.width_division_factor + height = max(self.height_division_factor, height) + width = max(self.width_division_factor, width) + return {"height": height, "width": width} + + def _resolve_sizes( + self, + image: Union[np.ndarray, Image.Image], + size: Optional[Dict[str, int]] = None, + crop_size: Optional[Dict[str, int]] = None, + ) -> Tuple[Dict[str, int], Dict[str, int]]: + size = size if size is not None else self.size + crop_size = crop_size if crop_size is not None else self.crop_size + if size is None and crop_size is None: + size = self._dynamic_target_size(image) + crop_size = dict(size) + elif size is None: + size = dict(crop_size) + elif crop_size is None: + crop_size = dict(size) + size = { + "height": self._ceil_to_factor(size["height"], self.height_division_factor), + "width": self._ceil_to_factor(size["width"], self.width_division_factor), + } + crop_size = { + "height": self._ceil_to_factor(crop_size["height"], self.height_division_factor), + "width": self._ceil_to_factor(crop_size["width"], self.width_division_factor), + } + return size, crop_size + + @staticmethod + def _ceil_to_factor(value: int, factor: int) -> int: + if value % factor == 0: + return value + return (value + factor - 1) // factor * factor + + def center_crop( + self, + image: np.ndarray, + crop_size: Dict[str, int], + **kwargs, + ) -> np.ndarray: + """Center crop image or video frame.""" + h, w = image.shape[:2] + crop_h, crop_w = crop_size["height"], crop_size["width"] + + top = (h - crop_h) // 2 + left = (w - crop_w) // 2 + if image.ndim == 3: + return image[top : top + crop_h, left : left + crop_w, :] + else: + return image[top : top + crop_h, left : left + crop_w] + + def normalize( + self, + image: np.ndarray, + mean: List[float], + std: List[float], + **kwargs, + ) -> np.ndarray: + """Normalize image or video frame.""" + image = image.astype(np.float32) / 255.0 + mean = np.array(mean).reshape(1, 1, -1) + std = np.array(std).reshape(1, 1, -1) + return (image - mean) / std + + def preprocess( + self, + images: ImageInput, + do_resize: Optional[bool] = None, + size: Optional[Dict[str, int]] = None, + do_center_crop: Optional[bool] = None, + crop_size: Optional[Dict[str, int]] = None, + do_normalize: Optional[bool] = None, + image_mean: Optional[Union[float, List[float]]] = None, + image_std: Optional[Union[float, List[float]]] = None, + do_convert_rgb: Optional[bool] = None, + num_frames: Optional[int] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + **kwargs, + ) -> Dict[str, Any]: + """ + Preprocess images or video frames. + + Args: + images: Input images or video frames. + return_tensors: Type of tensors to return ("pt" for PyTorch). + + Returns: + Dictionary with preprocessed pixel values. + """ + do_resize = do_resize if do_resize is not None else self.do_resize + size = size if size is not None else self.size + do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop + crop_size = crop_size if crop_size is not None else self.crop_size + do_normalize = do_normalize if do_normalize is not None else self.do_normalize + image_mean = image_mean if image_mean is not None else self.image_mean + image_std = image_std if image_std is not None else self.image_std + do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb + + # Handle single image or list of images (video frames) + if not isinstance(images, list): + images = [images] + + processed_images = [] + for image in images: + if isinstance(image, Image.Image): + image = np.array(image) + + # Handle 4D tensor case (batch of video frames) + if isinstance(image, np.ndarray) and image.ndim == 4: + # Process each frame in the batch + batch_processed_frames = [] + for i in range(image.shape[0]): + frame = image[i] # Get single frame + if frame.ndim == 3 and (frame.shape[0] == 3 or frame.shape[0] == 1): + frame = np.transpose(frame, (1, 2, 0)) # (C, H, W) -> (H, W, C) + + # Convert to RGB if needed + if do_convert_rgb and frame.shape[-1] != 3: + if len(frame.shape) == 2: # Grayscale + frame = np.stack([frame] * 3, axis=-1) + elif frame.shape[-1] == 4: # RGBA + frame = frame[..., :3] + + size_, crop_size_ = self._resolve_sizes(frame, size=size, crop_size=crop_size) + if do_resize: + frame = self.resize(frame, size_) + if do_center_crop: + frame = self.center_crop(frame, crop_size_) + + # Normalize + if do_normalize: + frame = self.normalize(frame, image_mean, image_std) + + batch_processed_frames.append(frame) + + # Stack frames back into 4D tensor + processed_image = np.stack(batch_processed_frames, axis=0) + else: + # Handle single image case (3D or 2D) + # Convert to RGB if needed + if do_convert_rgb and image.shape[-1] != 3: + if len(image.shape) == 2: # Grayscale + image = np.stack([image] * 3, axis=-1) + elif image.shape[-1] == 4: # RGBA + image = image[..., :3] + + size_, crop_size_ = self._resolve_sizes(image, size=size, crop_size=crop_size) + if do_resize: + image = self.resize(image, size_) + if do_center_crop: + image = self.center_crop(image, crop_size_) + + # Normalize + if do_normalize: + image = self.normalize(image, image_mean, image_std) + + processed_image = image + + processed_images.append(processed_image) + + # Stack frames for video + processed_images = np.stack(processed_images, axis=0) # B, T, H, W, C + + # Temporal Handling (Interpolate or Truncate) + if num_frames is not None: + current_frames = processed_images.shape[1] + if current_frames > num_frames: + logger.info(f"Truncating video frames from {current_frames} to {num_frames}") + processed_images = processed_images[:, :num_frames, ...] + elif current_frames < num_frames: + logger.info(f"Interpolating video frames from {current_frames} to {num_frames}") + # Interpolate requires (B, C, T, H, W) or (B, C, H, W) - we have (B, T, H, W, C) + # Permute to (B, C, T, H, W) for interpolate + vid_tensor = torch.from_numpy(processed_images).permute(0, 4, 1, 2, 3) + + # Interpolate + vid_tensor = torch.nn.functional.interpolate( + vid_tensor, + size=(num_frames, vid_tensor.shape[3], vid_tensor.shape[4]), + mode="trilinear", + align_corners=False, + ) + + # Permute back to (B, T, H, W, C) and convert to numpy + processed_images = vid_tensor.permute(0, 2, 3, 4, 1).numpy() + + # Convert to tensor if requested + if return_tensors == "pt": + processed_images = torch.from_numpy(processed_images) + # Rearrange to (B, C, T, H, W) for video (since input was B, T, H, W, C) + if processed_images.ndim == 5: + processed_images = processed_images.permute(0, 4, 1, 2, 3) + elif processed_images.ndim == 4: + # (T, H, W, C) -> (C, T, H, W) if it was list of frames + processed_images = processed_images.permute(3, 0, 1, 2) + # Add batch dimension + # processed_images = processed_images.unsqueeze(0) + + return {"pixel_values": processed_images} + + +class WanVideoProcessor: + """ + Processor for WanVideo models, combining image/video processing and text tokenization. + + Args: + image_processor: Image/video processor instance. + tokenizer: Text tokenizer instance. + """ + + attributes = ["tokenizer", "image_processor"] + valid_kwargs = [ + "chat_template", + ] + tokenizer_class = "AutoTokenizer" + image_processor_class = "WanVideoImageProcessor" + + def __init__(self, image_processor=None, tokenizer=None, max_text_length: Optional[int] = None, **kwargs): + if image_processor is None: + image_processor = WanVideoImageProcessor(**kwargs) + if tokenizer is None: + # Default to T5 tokenizer for text encoding + try: + logger.info("Loading default tokenizer: google/umt5-xxl") + tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl") + except Exception: + logger.warning("Could not load default tokenizer, using None") + tokenizer = None + + self.image_processor = image_processor + self.tokenizer = tokenizer + self.max_text_length = max_text_length + + def __call__( + self, + text: Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None, + images: Optional[ImageInput] = None, + videos: Optional[ImageInput] = None, + return_tensors: Optional[Union[str, TensorType]] = None, + **kwargs, + ) -> Dict[str, Any]: + """ + Process text and image/video inputs. + + Args: + text: Text input(s) to tokenize. + images: Image input(s) to process. + videos: Video frames to process. + return_tensors: Type of tensors to return. + + Returns: + Dictionary with processed inputs. + """ + if text is None and images is None and videos is None: + raise ValueError("You must provide either text, images, or videos.") + + data = {} + + # Process text + if text is not None and self.tokenizer is not None: + max_length = kwargs.pop("max_length", None) + if max_length is None: + max_length = self.max_text_length or getattr(self.tokenizer, "model_max_length", 256) + text_inputs = self.tokenizer( + text, + return_tensors=return_tensors, + padding=True, + truncation=True, + max_length=max_length, + **kwargs, + ) + data.update(text_inputs) + + # Process images or video + if images is not None or videos is not None: + image_inputs = self.image_processor( + images or videos, + return_tensors=return_tensors, + **kwargs, + ) + data.update(image_inputs) + + return data + + def batch_decode(self, *args, **kwargs): + """Delegate to tokenizer's batch_decode method.""" + if self.tokenizer is None: + raise ValueError("No tokenizer available for decoding.") + return self.tokenizer.batch_decode(*args, **kwargs) + + def decode(self, *args, **kwargs): + """Delegate to tokenizer's decode method.""" + if self.tokenizer is None: + raise ValueError("No tokenizer available for decoding.") + return self.tokenizer.decode(*args, **kwargs) + + @property + def model_input_names(self): + """Get model input names from components.""" + tokenizer_input_names = self.tokenizer.model_input_names if self.tokenizer else [] + image_processor_input_names = self.image_processor.model_input_names + return list(set(tokenizer_input_names + image_processor_input_names)) diff --git a/primus/backends/diffusion/data/processor.py b/primus/backends/diffusion/data/processor.py new file mode 100644 index 000000000..6434d24a6 --- /dev/null +++ b/primus/backends/diffusion/data/processor.py @@ -0,0 +1,182 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +Processes video and text data for training. +""" + +from collections.abc import Sequence +from typing import Any + +import torch + + +class WanVideoDataProcessor: + """Standalone data processor for WanVideo training.""" + + def __init__(self, config, model_id=None): + self.config = config + self.model_id = model_id + self.processor = None + self.tokenizer = None + + def apply_prompt_template(self, hf_messages: str) -> str: + """Apply prompt template for WanVideo.""" + # WanVideo uses direct prompts without special formatting. + # Keep backward compatibility: + # - old path passed `hf_messages` list[dict] + # - new path passes plain prompt string + if isinstance(hf_messages, str): + return hf_messages + try: + return hf_messages[0]["content"][1]["text"] + except Exception: + # Best-effort fallback + return str(hf_messages) + + def save_pretrained(self, save_directory: str): + pass + + def build(self): + """Initialize the processor and tokenizer.""" + if self.processor is not None and self.tokenizer is not None: + return + + from transformers import AutoTokenizer + + from primus.backends.diffusion.data.processing_wanvideo import ( + WanVideoProcessor as WanVideoModelProcessor, + ) + + wanvideo_kwargs = self.config.get("extra_kwargs", {}) + max_text_length = self.config.get("max_text_length") + if max_text_length is not None: + wanvideo_kwargs.setdefault("max_text_length", max_text_length) + + # Load tokenizer if specified + if self.config.get("text_tokenizer", None) is not None: + self.tokenizer = AutoTokenizer.from_pretrained(self.config.get("text_tokenizer")) + else: + self.tokenizer = None + self.processor = WanVideoModelProcessor(**wanvideo_kwargs, tokenizer=self.tokenizer) + + if self.tokenizer is None: + self.tokenizer = self.processor.tokenizer + + def _normalize_raw_batch(self, batch: Any) -> tuple[list[str], list[Any], int | None]: + """ + Normalize dataloader output into: + - prompts: list[str] + - frames_list: list[Any] (each item is typically np.ndarray (T,H,W,C)) + - num_frames: Optional[int] (only if consistent across samples) + """ + if isinstance(batch, dict): + prompts = batch.get("prompt") + frames_list = batch.get("video_frames") + num_frames = batch.get("num_frames", None) + + if prompts is None or frames_list is None: + raise ValueError("prepare_batch(dict) requires keys: 'prompt' and 'video_frames'") + if isinstance(prompts, str): + prompts = [prompts] + if not isinstance(prompts, (list, tuple)): + raise TypeError(f"'prompt' must be str or list[str], got {type(prompts)}") + if not isinstance(frames_list, (list, tuple)): + raise TypeError(f"'video_frames' must be list, got {type(frames_list)}") + + if isinstance(num_frames, (list, tuple)): + nfs = [int(x) for x in num_frames if x] + num_frames = nfs[0] if nfs and all(int(x) == int(nfs[0]) for x in nfs) else None + elif num_frames: + num_frames = int(num_frames) + else: + num_frames = None + + return list(prompts), list(frames_list), num_frames + + # RawBatchCollator returns list[dict] + if not isinstance(batch, (list, tuple)) or not batch: + raise ValueError(f"prepare_batch expected non-empty list/tuple, got {type(batch)}") + + prompts = [str(ex.get("prompt", "")) for ex in batch] + frames_list = [ex.get("video_frames") for ex in batch] + + # Prefer per-example num_frames if present and consistent; else None. + nfs = [ex.get("num_frames") for ex in batch if ex.get("num_frames")] + num_frames = nfs[0] if nfs and all(int(x) == int(nfs[0]) for x in nfs) else None + num_frames = int(num_frames) if num_frames else None + + return prompts, frames_list, num_frames + + def _tokenize_prompts(self, prompts: Sequence[str]) -> dict[str, torch.Tensor]: + formatted_prompts = [self.apply_prompt_template(p) for p in prompts] + return self.tokenizer( + formatted_prompts, + return_tensors="pt", + padding=self.config.get("padding_strategy", "max_length"), + truncation=True, + max_length=self.config.get("max_text_length", 512), + ) + + def _preprocess_videos(self, frames_list: Sequence[Any], num_frames: int | None) -> torch.Tensor: + if any(v is None for v in frames_list): + raise ValueError("prepare_batch got None in 'video_frames'") + video_inputs = self.processor.image_processor.preprocess( + list(frames_list), + num_frames=num_frames, + return_tensors="pt", + ) + if "pixel_values" not in video_inputs: + raise KeyError("image_processor.preprocess must return dict with key 'pixel_values'") + pixel_values = video_inputs["pixel_values"] + if not isinstance(pixel_values, torch.Tensor): + raise TypeError(f"pixel_values must be torch.Tensor, got {type(pixel_values)}") + if pixel_values.ndim != 5: + raise ValueError(f"Expected pixel_values shape [B,C,T,H,W], got {tuple(pixel_values.shape)}") + return pixel_values + + def _assemble_model_batch( + self, *, pixel_values: torch.Tensor, text_inputs: dict[str, torch.Tensor] + ) -> dict[str, Any]: + if "input_ids" not in text_inputs or "attention_mask" not in text_inputs: + raise KeyError("tokenizer output must include 'input_ids' and 'attention_mask'") + + out: dict[str, Any] = { + "video": pixel_values, # [B,C,T,H,W] + "input_ids": text_inputs["input_ids"], + "attention_mask": text_inputs["attention_mask"], + } + + # Legacy convenience fields (some models expect these keys). + out.setdefault("num_frames", int(pixel_values.shape[2])) + out.setdefault("height", int(pixel_values.shape[3])) + out.setdefault("width", int(pixel_values.shape[4])) + return out + + def prepare_batch(self, *, batch: Any, device: torch.device, dtype: torch.dtype) -> dict[str, Any]: + """ + Convert raw dataloader batch into model inputs. + + Expected raw batch from `RawBatchCollator`: + - batch: list[dict] with keys: + - video_frames: np.ndarray (T,H,W,C) or equivalent + - prompt: str + - num_frames: int (optional) + + Output matches current training pipelines: + - video: Tensor [B, C, T, H, W] + - input_ids: Tensor [B, L] + - attention_mask: Tensor [B, L] + - (optional) height/width/num_frames for legacy wan/wan_new paths + """ + if self.processor is None or self.tokenizer is None: + # Lazy init to keep behavior robust across call sites. + self.build() + + prompts, frames_list, num_frames = self._normalize_raw_batch(batch) + text_inputs = self._tokenize_prompts(prompts) + pixel_values = self._preprocess_videos(frames_list, num_frames=num_frames) + return self._assemble_model_batch(pixel_values=pixel_values, text_inputs=text_inputs) diff --git a/primus/backends/diffusion/data/registrations/__init__.py b/primus/backends/diffusion/data/registrations/__init__.py index b5a3705a2..4573e0a65 100644 --- a/primus/backends/diffusion/data/registrations/__init__.py +++ b/primus/backends/diffusion/data/registrations/__init__.py @@ -4,4 +4,4 @@ # See LICENSE for license information. ############################################################################### -"""Registered diffusion dataset builders.""" +"""Dataset registrations.""" diff --git a/primus/backends/diffusion/data/registrations/flux.py b/primus/backends/diffusion/data/registrations/flux.py new file mode 100644 index 000000000..234868970 --- /dev/null +++ b/primus/backends/diffusion/data/registrations/flux.py @@ -0,0 +1,38 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from primus.backends.diffusion.data.flux_precomputed import ( + FluxPrecomputedDataset, + FluxPrecomputedProcessor, + FluxRawImageTextDataset, + FluxRawImageTextProcessor, +) +from primus.backends.diffusion.utils.log import logger + + +def build_flux_dataset(dataset_config: dict): + processor_config = dataset_config.get("processor_config", {}) or {} + dataset_type = str(dataset_config.get("dataset_type", "precomputed")).lower() + if dataset_type == "raw": + processor = FluxRawImageTextProcessor(processor_config) + processor.build() + dataset = FluxRawImageTextDataset( + dataset_path=dataset_config.get("dataset_path"), + dataset_format=dataset_config.get("dataset_format", "webdataset"), + dataset_name=dataset_config.get("dataset"), + data_folder=dataset_config.get("data_folder"), + ) + logger.info(f"Built FLUX raw image-text dataset with {len(dataset)} samples") + elif dataset_type == "precomputed": + processor = FluxPrecomputedProcessor(processor_config) + processor.build() + dataset = FluxPrecomputedDataset(dataset_config["dataset_path"]) + logger.info(f"Built FLUX precomputed dataset with {len(dataset)} samples") + else: + raise ValueError("FLUX dataset_type must be either 'precomputed' or 'raw'") + return dataset, processor diff --git a/primus/backends/diffusion/data/registrations/wan.py b/primus/backends/diffusion/data/registrations/wan.py index 76492cb2b..f832f352a 100644 --- a/primus/backends/diffusion/data/registrations/wan.py +++ b/primus/backends/diffusion/data/registrations/wan.py @@ -4,173 +4,25 @@ # See LICENSE for license information. ############################################################################### -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any - -import numpy as np -import torch -from PIL import Image -from torch.utils.data import Dataset -from torchvision.transforms import functional as F -from transformers import AutoTokenizer - -from primus.backends.diffusion.utils.vision_process import fetch_video, strip_file_uri - - -class WanVideoProcessor: - """Tokenize prompts and normalize video tensors for Wan training.""" - - def __init__(self, config: dict[str, Any]): - self.config = config - self.max_text_length = int(config.get("max_text_length", 512)) - tokenizer_path = config.get("text_tokenizer") - if not tokenizer_path: - raise ValueError("Wan dataset processor requires `text_tokenizer`.") - trust_remote_code = bool(config.get("trust_remote_code", False)) - self.tokenizer = AutoTokenizer.from_pretrained( - tokenizer_path, - trust_remote_code=trust_remote_code, - ) - - extra_kwargs = config.get("extra_kwargs", {}) or {} - size = extra_kwargs.get("size", {}) or {} - self.height = int(size.get("height", 480)) - self.width = int(size.get("width", 832)) - self.image_mean = torch.tensor(extra_kwargs.get("image_mean", [0.5, 0.5, 0.5])).view(3, 1, 1, 1) - self.image_std = torch.tensor(extra_kwargs.get("image_std", [0.5, 0.5, 0.5])).view(3, 1, 1, 1) - - def tokenize(self, prompt: str) -> dict[str, torch.Tensor]: - encoded = self.tokenizer( - prompt, - max_length=self.max_text_length, - padding="max_length", - truncation=True, - return_tensors="pt", - ) - return { - "input_ids": encoded["input_ids"].squeeze(0).long(), - "attention_mask": encoded["attention_mask"].squeeze(0).long(), - } - - def normalize_video(self, video_tchw: torch.Tensor) -> torch.Tensor: - if video_tchw.ndim != 4: - raise ValueError(f"Expected video tensor [T,C,H,W], got shape={tuple(video_tchw.shape)}") - video_tchw = video_tchw[:, :3].float() - if video_tchw.max() > 2: - video_tchw = video_tchw / 255.0 - video_tchw = F.resize(video_tchw, [self.height, self.width], antialias=True) - video_cthw = video_tchw.permute(1, 0, 2, 3).contiguous() - return (video_cthw - self.image_mean) / self.image_std - - def prepare_batch( - self, *, batch: dict[str, Any], device: torch.device, dtype: torch.dtype - ) -> dict[str, Any]: - return batch - - -class WanVideoDataset(Dataset): - def __init__(self, config: dict[str, Any], processor: WanVideoProcessor): - self.config = config - self.processor = processor - self.dataset_path = Path(config.get("dataset_path", "")) - self.data_folder = Path(config.get("data_folder", "")) - self.frame_num = int(config.get("frame_num", 81)) - self.video_backend = str(config.get("video_backend", "imageio")).lower() - - if not self.dataset_path.exists(): - raise FileNotFoundError(f"Wan dataset metadata not found: {self.dataset_path}") - with self.dataset_path.open(encoding="utf-8") as f: - self.samples = [json.loads(line) for line in f if line.strip()] - if not self.samples: - raise ValueError(f"Wan dataset metadata is empty: {self.dataset_path}") - - def __len__(self) -> int: - return len(self.samples) - - def _resolve_video_path(self, value: str) -> str: - if value.startswith(("http://", "https://", "file://")): - return value - path = Path(value) - if not path.is_absolute() and self.data_folder: - path = self.data_folder / path - return str(path) - - def _sample_video(self, video_tchw: torch.Tensor) -> torch.Tensor: - total_frames = int(video_tchw.shape[0]) - if total_frames <= 0: - raise ValueError("Video contains no frames.") - idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long() - return video_tchw[idx] - - def _read_video_imageio(self, path: str) -> torch.Tensor: - import imageio.v3 as iio - - path = strip_file_uri(path) - frames = [] - for frame in iio.imiter(path): - image = Image.fromarray(frame).convert("RGB") - frames.append(torch.from_numpy(np.asarray(image).copy())) - if not frames: - raise ValueError(f"Video contains no frames: {path}") - return torch.stack(frames).permute(0, 3, 1, 2) - - def _read_video_decord(self, path: str) -> torch.Tensor: - import decord - - path = strip_file_uri(path) - vr = decord.VideoReader(path) - total_frames = len(vr) - idx = torch.linspace(0, total_frames - 1, self.frame_num).round().long().tolist() - return torch.from_numpy(vr.get_batch(idx).asnumpy()).permute(0, 3, 1, 2) - - def _read_video(self, path: str) -> torch.Tensor: - if self.video_backend == "imageio": - video = self._sample_video(self._read_video_imageio(path)) - elif self.video_backend == "decord": - video = self._read_video_decord(path) - else: - video = fetch_video( - { - "video": path, - "nframes": self.frame_num, - "resized_height": self.processor.height, - "resized_width": self.processor.width, - } - ) - return self.processor.normalize_video(video) - - def __getitem__(self, index: int) -> dict[str, Any]: - sample = self.samples[index] - prompt = sample.get("prompt") or sample.get("text") or sample.get("caption") or "" - video_key = sample.get("video") or sample.get("video_path") - if not video_key: - raise KeyError(f"Wan dataset sample missing `video`: index={index}") - - item = self.processor.tokenize(str(prompt)) - item["video"] = self._read_video(self._resolve_video_path(str(video_key))) - if "seed" in sample: - item["seed"] = int(sample["seed"]) - return item - - @staticmethod - def get_collator(): - def collate(samples: list[dict[str, Any]]) -> dict[str, Any]: - batch = { - "video": torch.stack([sample["video"] for sample in samples]), - "input_ids": torch.stack([sample["input_ids"] for sample in samples]), - "attention_mask": torch.stack([sample["attention_mask"] for sample in samples]), - } - if any("seed" in sample for sample in samples): - batch["seed"] = torch.tensor([sample.get("seed", 0) for sample in samples], dtype=torch.long) - return batch - - return collate - - -def build_wan_dataset(config: dict[str, Any]): - processor = WanVideoProcessor(config.get("processor_config", {}) or {}) - dataset = WanVideoDataset(config, processor) +"""Register Wan dataset builder.""" + +from primus.backends.diffusion.data import ( + DatasetConfig, + WanVideoDataProcessor, + WanVideoDataset, +) +from primus.backends.diffusion.utils.log import logger + + +def build_wan_dataset(dataset_config: dict): + processor_config = dataset_config["processor_config"] + processor = WanVideoDataProcessor(processor_config) + processor.build() + logger.info("Built data processor") + + dataset = WanVideoDataset( + processor=processor, + config=DatasetConfig(**dataset_config), + ) + logger.info(f"Built dataset with {len(dataset)} samples") return dataset, processor diff --git a/primus/backends/diffusion/diffusion_adapter.py b/primus/backends/diffusion/diffusion_adapter.py index 344c293b3..878ea3fea 100644 --- a/primus/backends/diffusion/diffusion_adapter.py +++ b/primus/backends/diffusion/diffusion_adapter.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Any -from primus.backends.diffusion.argument_builder import WanArgBuilder +from primus.backends.diffusion.argument_builder import DiffusionArgBuilder from primus.core.backend.backend_adapter import BackendAdapter from primus.core.utils.module_utils import log_rank_0 @@ -42,17 +42,18 @@ def setup_backend_path(self, backend_path=None) -> str: return resolved_str def convert_config(self, params: Any): - builder = WanArgBuilder() + builder = DiffusionArgBuilder() builder.update(params) - wan_args = builder.finalize() + diffusion_args = builder.finalize() # convert_config is also called by the standalone prepare hook, where the # Primus logger may not be initialized yet; guard the informational log. try: - log_rank_0("[Primus:DiffusionAdapter] Converted Primus module params -> Wan args") + model_name = getattr(getattr(diffusion_args, "model", {}), "get", lambda *_: None)("name") + log_rank_0(f"[Primus:DiffusionAdapter] Converted Primus module params -> {model_name} args") except Exception: # Standalone prepare hooks may run before the Primus logger is bound. pass - return wan_args + return diffusion_args def load_trainer_class(self, stage: str = "pretrain"): if stage in ("pretrain", "posttrain", "sft"): diff --git a/primus/backends/diffusion/diffusion_pretrain_trainer.py b/primus/backends/diffusion/diffusion_pretrain_trainer.py index 1b51b6a4c..218fc7019 100644 --- a/primus/backends/diffusion/diffusion_pretrain_trainer.py +++ b/primus/backends/diffusion/diffusion_pretrain_trainer.py @@ -9,17 +9,18 @@ import importlib.util from typing import Any +from primus.core.base_module import BaseModule from primus.core.trainer.base_trainer import BaseTrainer from primus.core.utils.module_utils import log_rank_0 from primus.core.utils.yaml_utils import nested_namespace_to_dict -class DiffusionPretrainTrainer(BaseTrainer): - """Primus lifecycle wrapper for Wan diffusion training.""" +class DiffusionPretrainTrainer(BaseTrainer, BaseModule): + """Primus lifecycle wrapper for diffusion backend training.""" - def __init__(self, backend_args: Any): - super().__init__(backend_args=backend_args) - self.wan_trainer = None + def __init__(self, backend_args: Any, *args, **kwargs): + super().__init__(backend_args=backend_args, *args, **kwargs) + self.diffusion_trainer = None @staticmethod def _as_dict(value: Any) -> dict: @@ -47,7 +48,18 @@ def setup(self): ) if importlib.util.find_spec(package) is None ] - video_backend = (dataset_cfg.get("config", {}) or {}).get("video_backend") + dataset_config = dataset_cfg.get("config", {}) or {} + video_backend = dataset_config.get("video_backend") + if dataset_cfg.get("name") == "flux": + for package in ("datasets", "huggingface_hub", "sentencepiece"): + if importlib.util.find_spec(package) is None: + missing.append(package) + if ( + dataset_config.get("dataset_type") == "raw" + and dataset_config.get("dataset_format") == "webdataset" + and importlib.util.find_spec("webdataset") is None + ): + missing.append("webdataset") if video_backend == "imageio" and importlib.util.find_spec("imageio") is None: missing.append("imageio") if video_backend == "decord" and importlib.util.find_spec("decord") is None: @@ -55,7 +67,7 @@ def setup(self): if missing: raise RuntimeError( "Diffusion backend missing required Python packages: " - f"{', '.join(missing)}. Install the Wan diffusion training extras first." + f"{', '.join(missing)}. Install the diffusion training extras first." ) if attention_backend: @@ -99,7 +111,7 @@ def init(self): ) model = get_model_builder(model_name)(model_config) dataset, processor = get_dataset_builder(dataset_name)(dataset_config) - self.wan_trainer = get_trainer_builder(trainer_name)( + self.diffusion_trainer = get_trainer_builder(trainer_name)( model=model, dataset=dataset, processor=processor, @@ -107,11 +119,15 @@ def init(self): ) def train(self): - if self.wan_trainer is None: + if self.diffusion_trainer is None: raise RuntimeError("DiffusionPretrainTrainer.init() must be called before train().") - self.wan_trainer.train() - self.wan_trainer.save_model() + self.diffusion_trainer.train() + self.diffusion_trainer.save_model() + + def run(self, *args, **kwargs): + """Compatibility hook for BaseModule; TrainRuntime drives lifecycle phases.""" + self.train() def cleanup(self, on_error: bool = False): try: diff --git a/primus/backends/diffusion/models/__init__.py b/primus/backends/diffusion/models/__init__.py index 2798dc3fb..7503f134c 100644 --- a/primus/backends/diffusion/models/__init__.py +++ b/primus/backends/diffusion/models/__init__.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Wan model exports for the Primus Wan backend.""" +"""Model exports for the Primus diffusion backend.""" from .wan import WanForTraining diff --git a/primus/backends/diffusion/models/flux/__init__.py b/primus/backends/diffusion/models/flux/__init__.py new file mode 100644 index 000000000..04ec5ff51 --- /dev/null +++ b/primus/backends/diffusion/models/flux/__init__.py @@ -0,0 +1,25 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from primus.backends.diffusion.models.flux.adapter import FluxForTraining +from primus.backends.diffusion.models.flux.model import ( + Flux, + FluxParams, + flux_1_dev_params, + flux_1_schnell_params, +) +from primus.backends.diffusion.models.flux.train_pipeline import ( + FluxFlowMatchTrainPipeline, +) + +__all__ = [ + "Flux", + "FluxForTraining", + "FluxFlowMatchTrainPipeline", + "FluxParams", + "flux_1_dev_params", + "flux_1_schnell_params", +] diff --git a/primus/backends/diffusion/models/flux/adapter.py b/primus/backends/diffusion/models/flux/adapter.py new file mode 100644 index 000000000..428306ce7 --- /dev/null +++ b/primus/backends/diffusion/models/flux/adapter.py @@ -0,0 +1,113 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn + +from primus.backends.diffusion.models.flux.train_pipeline import ( + FluxFlowMatchTrainPipeline, +) +from primus.backends.diffusion.models.interface import GenAIModel + + +@dataclass +class FluxConfigShim: + raw: dict + + def save_pretrained(self, save_directory: str): + import json + import os + + os.makedirs(save_directory, exist_ok=True) + with open(os.path.join(save_directory, "flux_config.json"), "w") as f: + json.dump(self.raw, f, indent=2, sort_keys=True) + + def to_dict(self): + return self.raw + + +class FluxForTraining(GenAIModel, nn.Module): + """Thin Primus adapter around the FLUX DiT backbone.""" + + def __init__( + self, + *, + dit: nn.Module, + train_pipeline: FluxFlowMatchTrainPipeline, + model_config: Any, + autoencoder: nn.Module | None = None, + t5_encoder: nn.Module | None = None, + clip_encoder: nn.Module | None = None, + raw_config: dict | None = None, + trainable_modules: str | None = None, + ): + super().__init__() + self.dit = dit + self.autoencoder = autoencoder + self.t5_encoder = t5_encoder + self.clip_encoder = clip_encoder + self.train_pipeline = train_pipeline + self.model_config = model_config + self.trainable_modules = trainable_modules + self.config = FluxConfigShim(raw=raw_config or {}) + + @property + def device(self): + return next(self.parameters()).device + + @property + def dtype(self): + return next(self.parameters()).dtype + + def freeze_except(self): + mode = ( + self.trainable_modules or getattr(self.model_config, "trainable_modules", None) or "dit" + ).lower() + + def freeze(module: nn.Module): + for param in module.parameters(): + param.requires_grad_(False) + + def unfreeze(module: nn.Module): + for param in module.parameters(): + param.requires_grad_(True) + + freeze(self) + if mode in ("dit", "diffusion", "backbone"): + unfreeze(self.dit) + elif mode == "all": + unfreeze(self) + else: + unfreeze(self.dit) + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + if hasattr(self.dit, "gradient_checkpointing"): + self.dit.gradient_checkpointing = True + + def forward(self, *args, **kwargs): + if len(args) >= 1 and isinstance(args[0], dict): + scheduler = args[1] if len(args) >= 2 else kwargs.get("scheduler") + return self.forward_train(args[0], scheduler=scheduler) + raise TypeError("FluxForTraining.forward expects (batch_dict, scheduler)") + + def forward_train(self, batch: dict[str, Any], scheduler: Any = None) -> dict[str, torch.Tensor]: + del scheduler + return self.train_pipeline.compute_loss( + dit=self.dit, + autoencoder=self.autoencoder, + t5_encoder=self.t5_encoder, + clip_encoder=self.clip_encoder, + batch=batch, + model_config=self.model_config, + ) + + def forward_inference(self, batch: dict[str, Any], **kwargs): + raise NotImplementedError("FLUX inference pipeline is not wired for Primus training yet") diff --git a/primus/backends/diffusion/models/flux/autoencoder.py b/primus/backends/diffusion/models/flux/autoencoder.py new file mode 100644 index 000000000..9ff6f3e56 --- /dev/null +++ b/primus/backends/diffusion/models/flux/autoencoder.py @@ -0,0 +1,315 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import torch +from einops import rearrange +from safetensors.torch import load_file as load_sft +from torch import Tensor, nn + + +@dataclass +class AutoEncoderParams: + resolution: int = 256 + in_channels: int = 3 + ch: int = 128 + out_ch: int = 3 + ch_mult: list[int] | None = None + num_res_blocks: int = 2 + z_channels: int = 16 + scale_factor: float = 0.3611 + shift_factor: float = 0.1159 + + def __post_init__(self): + if self.ch_mult is None: + self.ch_mult = [1, 2, 4, 4] + + +def swish(x: Tensor) -> Tensor: + return x * torch.sigmoid(x) + + +class AttnBlock(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1) + self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + bsz, channels, height, width = q.shape + q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous() + k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous() + v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous() + h_ = nn.functional.scaled_dot_product_attention(q, k, v) + return rearrange(h_, "b 1 (h w) c -> b c h w", h=height, w=width, c=channels, b=bsz) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x: Tensor) -> Tensor: + h = self.conv1(swish(self.norm1(x))) + h = self.conv2(swish(self.norm2(h))) + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class Downsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) + + def forward(self, x: Tensor) -> Tensor: + return self.conv(nn.functional.pad(x, (0, 1, 0, 1), mode="constant", value=0)) + + +class Upsample(nn.Module): + def __init__(self, in_channels: int): + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor) -> Tensor: + x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") + return self.conv(x) + + +class Encoder(nn.Module): + def __init__( + self, + resolution: int, + in_channels: int, + ch: int, + ch_mult: list[int], + num_res_blocks: int, + z_channels: int, + ): + super().__init__() + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1) + + in_ch_mult = (1,) + tuple(ch_mult) + self.down = nn.ModuleList() + block_in = self.ch + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + down.downsample = Downsample(block_in) + self.down.append(down) + + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: Tensor) -> Tensor: + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions - 1: + hs.append(self.down[i_level].downsample(hs[-1])) + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + return self.conv_out(swish(self.norm_out(h))) + + +class Decoder(nn.Module): + def __init__( + self, + ch: int, + out_ch: int, + ch_mult: list[int], + num_res_blocks: int, + in_channels: int, + resolution: int, + z_channels: int, + ): + super().__init__() + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + block_in = ch * ch_mult[self.num_resolutions - 1] + self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in) + self.up.insert(0, up) + + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1) + + def forward(self, z: Tensor) -> Tensor: + upscale_dtype = next(self.up.parameters()).dtype + h = self.conv_in(z) + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h).to(upscale_dtype) + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + return self.conv_out(swish(self.norm_out(h))) + + +class DiagonalGaussian(nn.Module): + def __init__(self, sample: bool = True, chunk_dim: int = 1): + super().__init__() + self.sample = sample + self.chunk_dim = chunk_dim + + def forward(self, z: Tensor) -> Tensor: + mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim) + if not self.sample: + return mean + std = torch.exp(0.5 * logvar) + return mean + std * torch.randn_like(mean) + + +class AutoEncoder(nn.Module): + def __init__(self, params: AutoEncoderParams, sample_z: bool = True): + super().__init__() + self.params = params + self.encoder = Encoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + ch_mult=params.ch_mult or [1, 2, 4, 4], + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.decoder = Decoder( + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + out_ch=params.out_ch, + ch_mult=params.ch_mult or [1, 2, 4, 4], + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.reg = DiagonalGaussian(sample=sample_z) + self.scale_factor = params.scale_factor + self.shift_factor = params.shift_factor + + def encode(self, x: Tensor) -> Tensor: + z = self.reg(self.encoder(x)) + return self.scale_factor * (z - self.shift_factor) + + def decode(self, z: Tensor) -> Tensor: + z = z / self.scale_factor + self.shift_factor + return self.decoder(z) + + def forward(self, x: Tensor) -> Tensor: + return self.decode(self.encode(x)) + + +def load_autoencoder( + ckpt_path: str | None, + params: AutoEncoderParams | None = None, + *, + dtype: torch.dtype = torch.bfloat16, + sample_z: bool = True, +) -> AutoEncoder: + ae = AutoEncoder(params or AutoEncoderParams(), sample_z=sample_z) + if ckpt_path: + ckpt_path = _resolve_checkpoint_path(ckpt_path, default_filename="ae.safetensors") + state = load_sft(ckpt_path, device="cpu") + missing, unexpected = ae.load_state_dict(state, strict=False) + if missing: + raise ValueError(f"FLUX autoencoder checkpoint missing {len(missing)} keys; first={missing[:3]}") + if unexpected: + raise ValueError( + f"FLUX autoencoder checkpoint has {len(unexpected)} unexpected keys; first={unexpected[:3]}" + ) + return ae.to(dtype=dtype) + + +def _resolve_checkpoint_path(path_or_repo_file: str, *, default_filename: str) -> str: + path = Path(path_or_repo_file).expanduser() + if path.exists(): + return str(path) + if path_or_repo_file.startswith(("/", "./", "../", "~")): + raise FileNotFoundError(f"FLUX checkpoint path not found: {path}") + + parts = path_or_repo_file.split("/") + if len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")): + raise FileNotFoundError(f"FLUX checkpoint path not found: {path}") + if len(parts) >= 3: + repo_id = "/".join(parts[:2]) + filename = "/".join(parts[2:]) + elif len(parts) == 2: + repo_id = path_or_repo_file + filename = default_filename + else: + raise FileNotFoundError( + f"FLUX checkpoint not found locally and is not a HF repo path: {path_or_repo_file}" + ) + + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id=repo_id, filename=filename) diff --git a/primus/backends/diffusion/models/flux/conditioner.py b/primus/backends/diffusion/models/flux/conditioner.py new file mode 100644 index 000000000..736bf225d --- /dev/null +++ b/primus/backends/diffusion/models/flux/conditioner.py @@ -0,0 +1,54 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +import torch +from torch import Tensor, nn +from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer + + +class HFEmbedder(nn.Module): + def __init__(self, version: str, max_length: int, **hf_kwargs): + super().__init__() + self.is_clip = "clip" in version.lower() or version.startswith("openai") + self.max_length = max_length + self.output_key = "pooler_output" if self.is_clip else "last_hidden_state" + + if self.is_clip: + self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length) + self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs) + else: + self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length) + self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs) + + self.hf_module = self.hf_module.eval().requires_grad_(False) + for param in self.hf_module.parameters(): + if not param.is_contiguous(): + param.data = param.data.contiguous() + + @property + def device(self) -> torch.device: + return next(self.hf_module.parameters()).device + + def forward(self, text: list[str]) -> Tensor: + batch_encoding = self.tokenizer( + text, + truncation=True, + max_length=self.max_length, + return_length=False, + return_overflowing_tokens=False, + padding="max_length", + return_tensors="pt", + ) + outputs = self.hf_module( + input_ids=batch_encoding["input_ids"].to(self.device), + attention_mask=batch_encoding["attention_mask"].to(self.device), + output_hidden_states=False, + ) + return outputs[self.output_key] diff --git a/primus/backends/diffusion/models/flux/configuration_flux.py b/primus/backends/diffusion/models/flux/configuration_flux.py new file mode 100644 index 000000000..6808fcac7 --- /dev/null +++ b/primus/backends/diffusion/models/flux/configuration_flux.py @@ -0,0 +1,14 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from types import SimpleNamespace + + +class FluxTrainingConfig(SimpleNamespace): + def to_dict(self) -> dict: + return dict(self.__dict__) diff --git a/primus/backends/diffusion/models/flux/layers.py b/primus/backends/diffusion/models/flux/layers.py new file mode 100644 index 000000000..475be057b --- /dev/null +++ b/primus/backends/diffusion/models/flux/layers.py @@ -0,0 +1,211 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch +from einops import rearrange +from torch import Tensor, nn + +from primus.backends.diffusion.models.flux.math import attention, rope + + +class EmbedND(nn.Module): + def __init__(self, dim: int, theta: int, axes_dim: list[int]): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: Tensor) -> Tensor: + n_axes = ids.shape[-1] + emb = torch.cat([rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], dim=-3) + return emb.unsqueeze(1) + + +def timestep_embedding(t: Tensor, dim: int, max_period: int = 10000, time_factor: float = 1000.0) -> Tensor: + t = time_factor * t + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + if torch.is_floating_point(t): + embedding = embedding.to(t) + return embedding + + +class MLPEmbedder(nn.Module): + def __init__(self, in_dim: int, hidden_dim: int): + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) + + def forward(self, x: Tensor) -> Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class RMSNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.scale = nn.Parameter(torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + x_dtype = x.dtype + x = x.float() + rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6) + return (x * rrms).to(dtype=x_dtype) * self.scale + + +class QKNorm(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.query_norm = RMSNorm(dim) + self.key_norm = RMSNorm(dim) + + def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]: + q = self.query_norm(q) + k = self.key_norm(k) + return q.to(v), k.to(v) + + +class SelfAttention(nn.Module): + def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.norm = QKNorm(head_dim) + self.proj = nn.Linear(dim, dim) + + def forward(self, x: Tensor, pe: Tensor) -> Tensor: + qkv = self.qkv(x) + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + x = attention(q, k, v, pe=pe) + return self.proj(x) + + +@dataclass +class ModulationOut: + shift: Tensor + scale: Tensor + gate: Tensor + + +class Modulation(nn.Module): + def __init__(self, dim: int, double: bool): + super().__init__() + self.is_double = double + self.multiplier = 6 if double else 3 + self.lin = nn.Linear(dim, self.multiplier * dim, bias=True) + + def forward(self, vec: Tensor) -> tuple[ModulationOut, ModulationOut | None]: + out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1) + return ModulationOut(*out[:3]), ModulationOut(*out[3:]) if self.is_double else None + + +class DoubleStreamBlock(nn.Module): + def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False): + super().__init__() + mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.num_heads = num_heads + self.hidden_size = hidden_size + self.img_mod = Modulation(hidden_size, double=True) + self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) + self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.img_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_hidden_dim, hidden_size, bias=True), + ) + self.txt_mod = Modulation(hidden_size, double=True) + self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) + self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.txt_mlp = nn.Sequential( + nn.Linear(hidden_size, mlp_hidden_dim, bias=True), + nn.GELU(approximate="tanh"), + nn.Linear(mlp_hidden_dim, hidden_size, bias=True), + ) + + def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]: + img_mod1, img_mod2 = self.img_mod(vec) + txt_mod1, txt_mod2 = self.txt_mod(vec) + + img_modulated = (1 + img_mod1.scale) * self.img_norm1(img) + img_mod1.shift + img_qkv = self.img_attn.qkv(img_modulated) + img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + img_q, img_k = self.img_attn.norm(img_q, img_k, img_v) + + txt_modulated = (1 + txt_mod1.scale) * self.txt_norm1(txt) + txt_mod1.shift + txt_qkv = self.txt_attn.qkv(txt_modulated) + txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v) + + q = torch.cat((txt_q, img_q), dim=2) + k = torch.cat((txt_k, img_k), dim=2) + v = torch.cat((txt_v, img_v), dim=2) + attn = attention(q, k, v, pe=pe) + txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :] + + img = img + img_mod1.gate * self.img_attn.proj(img_attn) + img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift) + txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn) + txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift) + return img, txt + + +class SingleStreamBlock(nn.Module): + def __init__( + self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0, qk_scale: float | None = None + ): + super().__init__() + self.hidden_dim = hidden_size + self.num_heads = num_heads + head_dim = hidden_size // num_heads + self.scale = qk_scale or head_dim**-0.5 + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim) + self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size) + self.norm = QKNorm(head_dim) + self.hidden_size = hidden_size + self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.mlp_act = nn.GELU(approximate="tanh") + self.modulation = Modulation(hidden_size, double=False) + + def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor: + mod, _ = self.modulation(vec) + x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift + qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1) + q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) + q, k = self.norm(q, k, v) + attn = attention(q, k, v, pe=pe) + output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2)) + return x + mod.gate * output + + +class LastLayer(nn.Module): + def __init__(self, hidden_size: int, patch_size: int, out_channels: int): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) + + def forward(self, x: Tensor, vec: Tensor) -> Tensor: + shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1) + x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :] + return self.linear(x) diff --git a/primus/backends/diffusion/models/flux/math.py b/primus/backends/diffusion/models/flux/math.py new file mode 100644 index 000000000..829d2f395 --- /dev/null +++ b/primus/backends/diffusion/models/flux/math.py @@ -0,0 +1,44 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +import torch +from einops import rearrange +from torch import Tensor + +from primus.backends.diffusion.attention import attention as backend_attention + + +def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor: + q, k = apply_rope(q, k, pe) + x = backend_attention( + q=rearrange(q, "B H L D -> B L H D"), + k=rearrange(k, "B H L D -> B L H D"), + v=rearrange(v, "B H L D -> B L H D"), + dtype=q.dtype if q.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16, + ) + return rearrange(x, "B L H D -> B L (H D)") + + +def rope(pos: Tensor, dim: int, theta: int) -> Tensor: + if dim % 2 != 0: + raise ValueError(f"RoPE dimension must be even, got {dim}") + scale = torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device) / dim + omega = 1.0 / (theta**scale) + out = torch.einsum("...n,d->...nd", pos, omega) + out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1) + return rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2).float() + + +def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]: + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk) diff --git a/primus/backends/diffusion/models/flux/model.py b/primus/backends/diffusion/models/flux/model.py new file mode 100644 index 000000000..ffa4fa108 --- /dev/null +++ b/primus/backends/diffusion/models/flux/model.py @@ -0,0 +1,167 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### +# +# Adapted from Black Forest Labs FLUX official implementation. + +from __future__ import annotations + +from dataclasses import asdict, dataclass + +import torch +from torch import Tensor, nn + +from primus.backends.diffusion.models.flux.layers import ( + DoubleStreamBlock, + EmbedND, + LastLayer, + MLPEmbedder, + SingleStreamBlock, + timestep_embedding, +) + + +@dataclass +class FluxParams: + in_channels: int + out_channels: int + vec_in_dim: int + context_in_dim: int + hidden_size: int + mlp_ratio: float + num_heads: int + depth: int + depth_single_blocks: int + axes_dim: list[int] + theta: int + qkv_bias: bool + guidance_embed: bool + + def to_dict(self) -> dict: + return asdict(self) + + +def flux_1_dev_params(**overrides) -> FluxParams: + values = { + "in_channels": 64, + "out_channels": 64, + "vec_in_dim": 768, + "context_in_dim": 4096, + "hidden_size": 3072, + "mlp_ratio": 4.0, + "num_heads": 24, + "depth": 19, + "depth_single_blocks": 38, + "axes_dim": [16, 56, 56], + "theta": 10000, + "qkv_bias": True, + "guidance_embed": True, + } + values.update(overrides) + return FluxParams(**values) + + +def flux_1_schnell_params(**overrides) -> FluxParams: + values = flux_1_dev_params(guidance_embed=False).to_dict() + values.update(overrides) + return FluxParams(**values) + + +class Flux(nn.Module): + """Transformer model for flow matching on packed latent sequences.""" + + def __init__(self, params: FluxParams): + super().__init__() + self.params = params + self.in_channels = params.in_channels + self.out_channels = params.out_channels + if params.hidden_size % params.num_heads != 0: + raise ValueError( + f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}" + ) + pe_dim = params.hidden_size // params.num_heads + if sum(params.axes_dim) != pe_dim: + raise ValueError(f"Got axes_dim={params.axes_dim}, expected sum={pe_dim}") + + self.hidden_size = params.hidden_size + self.num_heads = params.num_heads + self.gradient_checkpointing = False + self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim) + self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True) + self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) + self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size) + self.guidance_in = ( + MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if params.guidance_embed else nn.Identity() + ) + self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size) + self.double_blocks = nn.ModuleList( + [ + DoubleStreamBlock( + self.hidden_size, + self.num_heads, + mlp_ratio=params.mlp_ratio, + qkv_bias=params.qkv_bias, + ) + for _ in range(params.depth) + ] + ) + self.single_blocks = nn.ModuleList( + [ + SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio) + for _ in range(params.depth_single_blocks) + ] + ) + self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels) + + def _checkpoint_double(self, block: nn.Module, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor): + import torch.utils.checkpoint as checkpoint_utils + + return checkpoint_utils.checkpoint(block, img, txt, vec, pe, use_reentrant=False) + + def _checkpoint_single(self, block: nn.Module, img: Tensor, vec: Tensor, pe: Tensor): + import torch.utils.checkpoint as checkpoint_utils + + return checkpoint_utils.checkpoint(block, img, vec, pe, use_reentrant=False) + + def forward( + self, + img: Tensor, + img_ids: Tensor, + txt: Tensor, + txt_ids: Tensor, + timesteps: Tensor, + y: Tensor, + guidance: Tensor | None = None, + ) -> Tensor: + if img.ndim != 3 or txt.ndim != 3: + raise ValueError("Input img and txt tensors must have 3 dimensions.") + + img = self.img_in(img) + vec = self.time_in(timestep_embedding(timesteps, 256)) + if self.params.guidance_embed: + if guidance is None: + raise ValueError("FLUX guidance-distilled models require a guidance tensor.") + vec = vec + self.guidance_in(timestep_embedding(guidance, 256)) + vec = vec + self.vector_in(y) + txt = self.txt_in(txt) + + ids = torch.cat((txt_ids, img_ids), dim=1) + pe = self.pe_embedder(ids) + + use_checkpoint = self.training and self.gradient_checkpointing + for block in self.double_blocks: + if use_checkpoint: + img, txt = self._checkpoint_double(block, img, txt, vec, pe) + else: + img, txt = block(img=img, txt=txt, vec=vec, pe=pe) + + img = torch.cat((txt, img), 1) + for block in self.single_blocks: + if use_checkpoint: + img = self._checkpoint_single(block, img, vec, pe) + else: + img = block(img, vec=vec, pe=pe) + img = img[:, txt.shape[1] :, ...] + return self.final_layer(img, vec) diff --git a/primus/backends/diffusion/models/flux/train_pipeline.py b/primus/backends/diffusion/models/flux/train_pipeline.py new file mode 100644 index 000000000..98ac2d077 --- /dev/null +++ b/primus/backends/diffusion/models/flux/train_pipeline.py @@ -0,0 +1,184 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F + +from primus.backends.diffusion.models.flux.utils import ( + create_position_encoding_for_latents, + generate_latent_from_mean_logvar, + pack_latents, +) + + +@dataclass +class FluxFlowMatchTrainPipelineConfig: + autoencoder_scale_factor: float = 0.3611 + autoencoder_shift_factor: float = 0.1159 + guidance: float | None = None + + +class FluxFlowMatchTrainPipeline: + """Flow-matching loss for FLUX using precomputed or online encodings.""" + + def __init__(self, cfg: FluxFlowMatchTrainPipelineConfig | None = None): + self.cfg = cfg or FluxFlowMatchTrainPipelineConfig() + + @staticmethod + def _require_module(module: torch.nn.Module | None, name: str) -> torch.nn.Module: + if module is None: + raise ValueError(f"FLUX raw image-text training requires `{name}` to be configured.") + return module + + @staticmethod + def _align_module_dtype(module: torch.nn.Module, *, device: torch.device, dtype: torch.dtype) -> None: + """Move a frozen encoder to the target device/dtype only when it differs. + + `nn.Module.to(dtype=...)` casts floating-point params/buffers only, leaving + integer buffers (e.g. token position ids) untouched. We inspect the first + floating-point parameter to avoid re-casting on every step. + """ + current = next((p for p in module.parameters() if p.is_floating_point()), None) + if current is None: + module.to(device=device) + return + if current.dtype != dtype or current.device != device: + module.to(device=device, dtype=dtype) + + def _prepare_precomputed( + self, + *, + batch: dict[str, Any], + device: torch.device, + dtype: torch.dtype, + model_config: Any, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + t5_encodings = batch["t5_encodings"].to(device=device, dtype=dtype, non_blocking=True) + clip_encodings = batch["clip_encodings"].to(device=device, dtype=dtype, non_blocking=True) + mean = batch["mean"].to(device=device, dtype=dtype, non_blocking=True) + logvar = batch["logvar"].to(device=device, dtype=dtype, non_blocking=True) + + latents = generate_latent_from_mean_logvar(mean, logvar) + scale = float(getattr(model_config, "autoencoder_scale_factor", self.cfg.autoencoder_scale_factor)) + shift = float(getattr(model_config, "autoencoder_shift_factor", self.cfg.autoencoder_shift_factor)) + labels = (latents - shift) * scale + return labels, t5_encodings, clip_encodings + + def _prepare_raw( + self, + *, + batch: dict[str, Any], + autoencoder: torch.nn.Module | None, + t5_encoder: torch.nn.Module | None, + clip_encoder: torch.nn.Module | None, + device: torch.device, + dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ae = self._require_module(autoencoder, "model.config.encoder.autoencoder") + t5 = self._require_module(t5_encoder, "model.config.encoder.t5_encoder") + clip = self._require_module(clip_encoder, "model.config.encoder.clip_encoder") + image = batch["image"].to(device=device, dtype=dtype, non_blocking=True) + prompts = batch.get("prompts") + if not isinstance(prompts, list): + raise ValueError("FLUX raw batch requires `prompts` as list[str].") + + # The frozen encoders are built in bf16 regardless of the training compute + # dtype; align them with the DiT dtype so running them on `image`/inputs of + # `dtype` does not raise a dtype-mismatch error (e.g. for fp32 runs). + self._align_module_dtype(ae, device=device, dtype=dtype) + self._align_module_dtype(t5, device=device, dtype=dtype) + self._align_module_dtype(clip, device=device, dtype=dtype) + ae.eval() + t5.eval() + clip.eval() + with torch.no_grad(): + labels = ae.encode(image).to(device=device, dtype=dtype) + t5_encodings = t5(prompts).to(device=device, dtype=dtype) + clip_encodings = clip(prompts).to(device=device, dtype=dtype) + return labels, t5_encodings, clip_encodings + + def compute_loss( + self, + *, + dit: torch.nn.Module, + batch: dict[str, Any], + model_config: Any, + autoencoder: torch.nn.Module | None = None, + t5_encoder: torch.nn.Module | None = None, + clip_encoder: torch.nn.Module | None = None, + ) -> dict[str, torch.Tensor]: + if batch.get("sp_group") is not None: + raise ValueError("FLUX diffusion training currently requires `parallelism.sp_size: 1`.") + + device = next(dit.parameters()).device + dtype = next(dit.parameters()).dtype + if "image" in batch: + labels, t5_encodings, clip_encodings = self._prepare_raw( + batch=batch, + autoencoder=autoencoder, + t5_encoder=t5_encoder, + clip_encoder=clip_encoder, + device=device, + dtype=dtype, + ) + else: + required = ("t5_encodings", "clip_encodings", "mean", "logvar") + missing = [key for key in required if key not in batch] + if missing: + raise ValueError(f"FLUX precomputed batch is missing required keys: {missing}") + labels, t5_encodings, clip_encodings = self._prepare_precomputed( + batch=batch, + device=device, + dtype=dtype, + model_config=model_config, + ) + + bsz = labels.shape[0] + noise = torch.randn_like(labels) + timesteps = torch.rand((bsz,), device=device, dtype=dtype) + sigmas = timesteps.view(-1, 1, 1, 1) + noisy_latents = (1 - sigmas) * labels + sigmas * noise + target = noise - labels + + _, _, latent_height, latent_width = noisy_latents.shape + # Position ids are integer grid indices consumed by RoPE; build them in + # float32 (independent of the model compute dtype) so that indices remain + # exactly representable. bf16 only represents integers up to 256 exactly, + # which would silently corrupt positions for larger latent grids. + img_ids = create_position_encoding_for_latents( + bsz, + latent_height, + latent_width, + position_dim=3, + device=device, + dtype=torch.float32, + ) + txt_ids = torch.zeros(bsz, t5_encodings.shape[1], 3, device=device, dtype=torch.float32) + noisy_latents = pack_latents(noisy_latents) + target = pack_latents(target) + + guidance_value = getattr(model_config, "guidance", self.cfg.guidance) + guidance = ( + None + if guidance_value is None + else torch.full((bsz,), float(guidance_value), device=device, dtype=dtype) + ) + pred = dit( + img=noisy_latents, + img_ids=img_ids, + txt=t5_encodings, + txt_ids=txt_ids, + y=clip_encodings, + timesteps=timesteps, + guidance=guidance, + ) + loss = F.mse_loss(pred.float(), target.float(), reduction="sum") / target.numel() + return {"loss": loss, "log_metrics": {"mse": loss.detach()}} diff --git a/primus/backends/diffusion/models/flux/utils.py b/primus/backends/diffusion/models/flux/utils.py new file mode 100644 index 000000000..00d0a2acf --- /dev/null +++ b/primus/backends/diffusion/models/flux/utils.py @@ -0,0 +1,49 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import torch +from torch import Tensor + +PATCH_HEIGHT = 2 +PATCH_WIDTH = 2 +LATENT_CHANNELS = 16 +IMAGE_LATENT_SIZE_RATIO = 8 + + +def generate_latent_from_mean_logvar(mean: Tensor, logvar: Tensor) -> Tensor: + return mean + torch.exp(0.5 * logvar) * torch.randn_like(mean) + + +def create_position_encoding_for_latents( + bsz: int, + latent_height: int, + latent_width: int, + position_dim: int = 3, + *, + device: torch.device | None = None, + dtype: torch.dtype | None = None, +) -> Tensor: + height = latent_height // PATCH_HEIGHT + width = latent_width // PATCH_WIDTH + position_encoding = torch.zeros(height, width, position_dim, device=device, dtype=dtype) + position_encoding[:, :, 1] = torch.arange(height, device=device, dtype=dtype).unsqueeze(1) + position_encoding[:, :, 2] = torch.arange(width, device=device, dtype=dtype).unsqueeze(0) + return position_encoding.view(1, height * width, position_dim).repeat(bsz, 1, 1) + + +def pack_latents(x: Tensor) -> Tensor: + bsz, channels, latent_height, latent_width = x.shape + if latent_height % PATCH_HEIGHT != 0 or latent_width % PATCH_WIDTH != 0: + raise ValueError( + "FLUX latents must have height and width divisible by 2, " f"got shape={tuple(x.shape)}" + ) + height = latent_height // PATCH_HEIGHT + width = latent_width // PATCH_WIDTH + x = x.unfold(2, PATCH_HEIGHT, PATCH_HEIGHT).unfold(3, PATCH_WIDTH, PATCH_WIDTH) + x = x.permute(0, 2, 3, 1, 4, 5) + return x.reshape(bsz, height * width, channels * PATCH_HEIGHT * PATCH_WIDTH) diff --git a/primus/backends/diffusion/models/registrations/flux.py b/primus/backends/diffusion/models/registrations/flux.py new file mode 100644 index 000000000..93a57be60 --- /dev/null +++ b/primus/backends/diffusion/models/registrations/flux.py @@ -0,0 +1,234 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from __future__ import annotations + +import glob +import os +from typing import Any + +import torch +from safetensors.torch import load_file as safe_load_file + +from primus.backends.diffusion.models.flux.adapter import FluxForTraining +from primus.backends.diffusion.models.flux.autoencoder import ( + AutoEncoderParams, + load_autoencoder, +) +from primus.backends.diffusion.models.flux.conditioner import HFEmbedder +from primus.backends.diffusion.models.flux.configuration_flux import FluxTrainingConfig +from primus.backends.diffusion.models.flux.model import ( + Flux, + flux_1_dev_params, + flux_1_schnell_params, +) +from primus.backends.diffusion.models.flux.train_pipeline import ( + FluxFlowMatchTrainPipeline, + FluxFlowMatchTrainPipelineConfig, +) +from primus.backends.diffusion.utils.log import logger +from primus.backends.diffusion.utils.train_utils import count_parameters + +_FLUX_PRESET_ALIASES = { + "flux-schnell": "flux-schnell", + "flux.1-schnell": "flux-schnell", + "flux1-schnell": "flux-schnell", + "flux-dev": "flux-dev", + "flux.1-dev": "flux-dev", + "flux1-dev": "flux-dev", +} + + +def _strip_known_prefixes(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + prefixes = ("module.", "dit.", "model.") + out: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + stripped = key + changed = True + while changed: + changed = False + for prefix in prefixes: + if stripped.startswith(prefix): + stripped = stripped[len(prefix) :] + changed = True + out[stripped] = value + return out + + +def _load_state_dict(path: str) -> dict[str, torch.Tensor]: + if path.endswith(".safetensors"): + return dict(safe_load_file(path)) + obj = torch.load(path, map_location="cpu") + if isinstance(obj, dict) and "model" in obj and isinstance(obj["model"], dict): + obj = obj["model"] + if not isinstance(obj, dict): + raise ValueError(f"Unsupported checkpoint format at {path}") + return obj + + +def _candidate_weight_files(path: str, *, default_filename: str) -> list[str]: + if os.path.isfile(path): + return [path] + if not os.path.exists(path): + resolved = _resolve_hf_checkpoint(path, default_filename=default_filename) + if resolved: + return [resolved] + candidates: list[str] = [] + for fname in ( + "flux1-schnell.safetensors", + "flux1-dev.safetensors", + "dit_model.safetensors", + "model.safetensors", + ): + candidate = os.path.join(path, fname) + if os.path.exists(candidate): + candidates.append(candidate) + if not candidates: + candidates = sorted(glob.glob(os.path.join(path, "*.safetensors"))) + if not candidates: + candidates = sorted(glob.glob(os.path.join(path, "*.bin"))) + return candidates + + +def _resolve_hf_checkpoint(path_or_repo_file: str, *, default_filename: str) -> str | None: + if path_or_repo_file.startswith(("/", "./", "../", "~")): + return None + parts = path_or_repo_file.split("/") + if len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")): + return None + if len(parts) < 2: + return None + if len(parts) >= 3: + repo_id = "/".join(parts[:2]) + filename = "/".join(parts[2:]) + else: + repo_id = path_or_repo_file + filename = default_filename + from huggingface_hub import hf_hub_download + + return hf_hub_download(repo_id=repo_id, filename=filename) + + +def _load_flux_weights(dit: torch.nn.Module, pretrained_path: str, *, default_filename: str) -> None: + candidates = _candidate_weight_files(pretrained_path, default_filename=default_filename) + if not candidates: + raise FileNotFoundError(f"No FLUX DiT weights found under {pretrained_path}") + + merged: dict[str, torch.Tensor] = {} + for ckpt in candidates: + merged.update(_strip_known_prefixes(_load_state_dict(ckpt))) + + result = dit.load_state_dict(merged, strict=False) + logger.info( + "Loaded FLUX DiT weights. " + f"files={len(candidates)} missing={len(result.missing_keys)} unexpected={len(result.unexpected_keys)}" + ) + + +def _build_flux_dit(params) -> Flux: + local_rank = os.environ.get("LOCAL_RANK") + use_cuda = local_rank is not None and torch.cuda.is_available() + device = torch.device(f"cuda:{local_rank}") if use_cuda else torch.device("cpu") + old_dtype = torch.get_default_dtype() + try: + if use_cuda: + torch.set_default_dtype(torch.bfloat16) + with torch.device(device): + dit = Flux(params) + finally: + torch.set_default_dtype(old_dtype) + return dit + + +def build_flux_model(model_config: dict[str, Any]): + """ + Build a FLUX model from the selected model preset. + + `model_preset` is injected by the registry from `model.name` for Primus + configs such as `flux.1-dev` and `flux.1-schnell`. + """ + cfg_dict: dict[str, Any] = dict(model_config.get("config", {}) or {}) + preset_name = str(model_config.get("model_preset") or cfg_dict.get("model_preset") or "flux.1-schnell") + preset = _FLUX_PRESET_ALIASES.get(preset_name.lower(), preset_name) + + params_overrides = dict(cfg_dict.get("params", {}) or {}) + if preset == "flux-dev": + params = flux_1_dev_params(**params_overrides) + elif preset == "flux-schnell": + params = flux_1_schnell_params(**params_overrides) + else: + raise ValueError( + "Unsupported FLUX model_preset=" + f"{preset_name!r}; expected one of: 'flux.1-dev', 'flux.1-schnell'" + ) + dit = _build_flux_dit(params) + + pretrained_path = model_config.get("load_from_pretrained_path") or model_config.get("pretrained_path") + if pretrained_path: + logger.info(f"Loading FLUX DiT weights from {pretrained_path}") + default_filename = "flux1-dev.safetensors" if preset == "flux-dev" else "flux1-schnell.safetensors" + _load_flux_weights(dit, pretrained_path, default_filename=default_filename) + + encoder_cfg = dict(model_config.get("encoder", {}) or cfg_dict.get("encoder", {}) or {}) + dtype = torch.bfloat16 + t5_encoder = None + clip_encoder = None + autoencoder = None + if encoder_cfg.get("t5_encoder"): + t5_encoder = HFEmbedder( + str(encoder_cfg["t5_encoder"]), + max_length=int(encoder_cfg.get("max_t5_length", 256)), + torch_dtype=dtype, + ) + if encoder_cfg.get("clip_encoder"): + clip_encoder = HFEmbedder( + str(encoder_cfg["clip_encoder"]), + max_length=int(encoder_cfg.get("max_clip_length", 77)), + torch_dtype=dtype, + ) + if encoder_cfg.get("autoencoder"): + ae_params = AutoEncoderParams( + resolution=int(encoder_cfg.get("resolution", 256)), + scale_factor=float(cfg_dict.get("autoencoder_scale_factor", 0.3611)), + shift_factor=float(cfg_dict.get("autoencoder_shift_factor", 0.1159)), + ) + autoencoder = load_autoencoder( + str(encoder_cfg["autoencoder"]), + ae_params, + dtype=dtype, + sample_z=bool(encoder_cfg.get("sample_z", True)), + ) + + training_cfg = FluxTrainingConfig( + model_preset=preset, + trainable_modules=cfg_dict.get("trainable_modules", "dit"), + guidance=None if not params.guidance_embed else float(cfg_dict.get("guidance", 1.0)), + autoencoder_scale_factor=float(cfg_dict.get("autoencoder_scale_factor", 0.3611)), + autoencoder_shift_factor=float(cfg_dict.get("autoencoder_shift_factor", 0.1159)), + ) + pipeline = FluxFlowMatchTrainPipeline( + FluxFlowMatchTrainPipelineConfig( + autoencoder_scale_factor=training_cfg.autoencoder_scale_factor, + autoencoder_shift_factor=training_cfg.autoencoder_shift_factor, + guidance=training_cfg.guidance, + ) + ) + model = FluxForTraining( + dit=dit, + train_pipeline=pipeline, + model_config=training_cfg, + autoencoder=autoencoder, + t5_encoder=t5_encoder, + clip_encoder=clip_encoder, + raw_config={ + "model_config": model_config, + "flux_params": params.to_dict(), + }, + trainable_modules=training_cfg.trainable_modules, + ) + total_params, trainable_params = count_parameters(model) + logger.info(f"Built FLUX model: total={total_params:,} trainable={trainable_params:,}") + return model diff --git a/primus/backends/diffusion/models/wan/adapter.py b/primus/backends/diffusion/models/wan/adapter.py index 995177d6e..1d453c5f9 100644 --- a/primus/backends/diffusion/models/wan/adapter.py +++ b/primus/backends/diffusion/models/wan/adapter.py @@ -28,7 +28,7 @@ class WanConfigShim: raw: dict def save_pretrained(self, save_directory: str): - # Best-effort: keep minimal JSON for debugging/reproducibility. + # Best-effort: keep minimal JSON for reproducibility. import json import os diff --git a/primus/backends/diffusion/models/wan/train_pipeline.py b/primus/backends/diffusion/models/wan/train_pipeline.py index 3f8a4d1b8..db755638e 100644 --- a/primus/backends/diffusion/models/wan/train_pipeline.py +++ b/primus/backends/diffusion/models/wan/train_pipeline.py @@ -55,14 +55,6 @@ def _encode_prompt( @staticmethod def _get_seed_from_env_or_batch(batch: Dict[str, Any]) -> Optional[int]: - # Keep parity with existing scripts/wan_new behavior. - import os - - if os.environ.get("FIXED_SEED"): - try: - return int(os.environ["FIXED_SEED"]) - except ValueError as exc: - raise ValueError(f"Invalid FIXED_SEED value: {os.environ['FIXED_SEED']}") from exc seed = batch.get("seed", None) if seed is None: return None @@ -137,24 +129,8 @@ def _crop_latents(latents: torch.Tensor, *, spatial_size: tuple[int, int]) -> to @staticmethod def _select_timestep(scheduler: Any, device: torch.device) -> torch.Tensor: - """ - Match `wan_new` timestep selection: - - If FIXED_TIMESTEP is set, use that discrete index into [0, num_train_timesteps). - - Else uniform randint over [0, num_train_timesteps). - """ - import os - - if os.environ.get("FIXED_TIMESTEP"): - try: - fixed_step = int(os.environ["FIXED_TIMESTEP"]) - except ValueError as exc: - raise ValueError(f"Invalid FIXED_TIMESTEP value: {os.environ['FIXED_TIMESTEP']}") from exc - max_step = int(scheduler.num_train_timesteps) - 1 - fixed_step = max(0, min(fixed_step, max_step)) - timestep_id = torch.tensor([fixed_step], device=device) - else: - timestep_id = torch.randint(0, int(scheduler.num_train_timesteps), (1,), device=device) - + """Sample one training timestep uniformly.""" + timestep_id = torch.randint(0, int(scheduler.num_train_timesteps), (1,), device=device) # scheduler.timesteps live on CPU in this repo; `wan_new` indexes with cpu tensor. timestep = scheduler.timesteps[timestep_id.cpu()].float() return timestep.to(device=device) diff --git a/primus/backends/diffusion/registry.py b/primus/backends/diffusion/registry.py index b0a728827..03bb08f3c 100644 --- a/primus/backends/diffusion/registry.py +++ b/primus/backends/diffusion/registry.py @@ -23,6 +23,29 @@ def _build_wan_dataset(dataset_config: dict): return build_wan_dataset(dataset_config) +def _build_flux_model(model_config: dict): + from primus.backends.diffusion.models.registrations.flux import build_flux_model + + return build_flux_model(model_config) + + +def _build_flux_preset_model(preset: str): + def _builder(model_config: dict): + from primus.backends.diffusion.models.registrations.flux import build_flux_model + + config = dict(model_config) + config["model_preset"] = preset + return build_flux_model(config) + + return _builder + + +def _build_flux_dataset(dataset_config: dict): + from primus.backends.diffusion.data.registrations.flux import build_flux_dataset + + return build_flux_dataset(dataset_config) + + def _build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict): from primus.backends.diffusion.trainers.fsdp2 import build_fsdp2_trainer @@ -35,9 +58,13 @@ def _build_fsdp2_trainer(*, model, dataset, processor, trainer_args: dict): MODEL_BUILDERS: Dict[str, Callable[[dict], Any]] = { + "flux": _build_flux_model, + "flux.1-dev": _build_flux_preset_model("flux.1-dev"), + "flux.1-schnell": _build_flux_preset_model("flux.1-schnell"), "wan": _build_wan_model, } DATASET_BUILDERS: Dict[str, Callable[[dict], Tuple[Any, Any]]] = { + "flux": _build_flux_dataset, "wan": _build_wan_dataset, } TRAINER_BUILDERS: Dict[str, Callable[..., Any]] = { diff --git a/primus/backends/diffusion/schedulers/flow_match.py b/primus/backends/diffusion/schedulers/flow_match.py index ec5519574..cdd3f2f9b 100644 --- a/primus/backends/diffusion/schedulers/flow_match.py +++ b/primus/backends/diffusion/schedulers/flow_match.py @@ -59,6 +59,11 @@ def set_timesteps( if dynamic_shift_len is not None else self.exponential_shift_mu ) + if mu is None: + raise ValueError( + "`exponential_shift=True` requires either `dynamic_shift_len` " + "or `exponential_shift_mu`." + ) self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1)) else: self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) diff --git a/primus/backends/diffusion/trainers/__init__.py b/primus/backends/diffusion/trainers/__init__.py index d4df6355e..f1fbeeb37 100644 --- a/primus/backends/diffusion/trainers/__init__.py +++ b/primus/backends/diffusion/trainers/__init__.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Trainer registrations for the Primus Wan backend.""" +"""Trainer registrations for the Primus diffusion backend.""" from .fsdp2 import build_fsdp2_trainer diff --git a/primus/backends/diffusion/trainers/base.py b/primus/backends/diffusion/trainers/base.py index c5d5227fd..f47537096 100644 --- a/primus/backends/diffusion/trainers/base.py +++ b/primus/backends/diffusion/trainers/base.py @@ -46,6 +46,14 @@ def create_lr_scheduler(optimizer, scheduler_type, warmup_steps, total_steps): if total_steps <= 0: return torch.optim.lr_scheduler.LambdaLR(optimizer, lambda step: 1.0) + warmup_steps = int(warmup_steps) + if warmup_steps > total_steps: + logger.warning( + f"Warmup steps ({warmup_steps}) exceed total steps ({total_steps}). " + f"Adjusting warmup steps to {total_steps}." + ) + warmup_steps = total_steps + def linear_warmup(step): if warmup_steps == 0: return 1.0 @@ -139,8 +147,6 @@ def __init__( seed = self.args.get("seed") if seed is not None: set_seed(int(seed)) - if os.environ.get("FIXED_SEED"): - set_seed(int(os.environ["FIXED_SEED"])) # --- Gradient Checkpointing --- if self.args.get("gradient_checkpointing", False): @@ -240,10 +246,27 @@ def _grad_sync_context(self, is_update_step: bool): def _clip_grad_norm(self) -> float: """Clip gradient norm. Returns the total norm value.""" - if self.max_grad_norm > 0: - norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm) - return norm.item() if isinstance(norm, torch.Tensor) else float(norm) - return 0.0 + if self.max_grad_norm <= 0: + return 0.0 + + parameters = [p for p in self.model.parameters() if p.grad is not None] + if not parameters: + return 0.0 + + grads = [p.grad for p in parameters] + norm = torch.nn.utils.get_total_norm(grads, norm_type=2.0, foreach=True) + dtensor_cls = None + try: + from torch.distributed.tensor import DTensor + + dtensor_cls = DTensor + except (ImportError, RuntimeError) as exc: + logger.debug(f"Skipping DTensor grad-norm conversion: {exc}") + if dtensor_cls is not None and isinstance(norm, dtensor_cls): + norm = norm.full_tensor() + + torch.nn.utils.clip_grads_with_norm_(parameters, self.max_grad_norm, norm, foreach=True) + return norm.item() if isinstance(norm, torch.Tensor) else float(norm) def _save_checkpoint(self): """Save checkpoint at save_steps intervals. Override for custom strategies.""" @@ -385,6 +408,9 @@ def _infer_batch_size_from_sequences(self, value) -> int | None: return None def _infer_local_batch_size(self, batch) -> int: + if isinstance(batch, (list, tuple)): + return len(batch) + tensor_batch_size = self._infer_batch_size_from_tensors(batch) if tensor_batch_size is not None: return tensor_batch_size diff --git a/primus/backends/diffusion/trainers/fsdp2.py b/primus/backends/diffusion/trainers/fsdp2.py index 1fa79b177..958ea4132 100644 --- a/primus/backends/diffusion/trainers/fsdp2.py +++ b/primus/backends/diffusion/trainers/fsdp2.py @@ -116,7 +116,11 @@ def _apply_fsdp2(self): pass # fallback to full mesh wrap_target = str(self.args.get("fsdp2_wrap_target", "") or "").strip() - wrap_root = self._get_module_by_path(self.model, wrap_target) if wrap_target else self.model + wrap_root = ( + self._get_module_by_path(self.model, wrap_target, option_name="fsdp2_wrap_target") + if wrap_target + else self.model + ) # Pre-cast to bf16 before FSDP wrapping so parameters are stored in bf16, # matching DiffSynth/DeepSpeed bf16 behavior. @@ -130,7 +134,32 @@ def _apply_fsdp2(self): logger.info("FSDP2: world_size=1; skipping composable FSDP wrapping.") return + if bool(self.args.get("compile_transformer_blocks", False)): + self._compile_transformer_blocks(wrap_root) + reshard_after_forward = bool(self.args.get("fsdp2_reshard_after_forward", True)) + no_reshard_paths = { + item.strip() + for item in str(self.args.get("fsdp_module_paths_no_reshard", "") or "").split(",") + if item.strip() + } + + module_paths_spec = str(self.args.get("fsdp_module_paths_to_wrap", "") or "") + module_paths = [item.strip() for item in module_paths_spec.split(",") if item.strip()] + wrapped_paths = 0 + for module_path in module_paths: + module = self._get_module_by_path(wrap_root, module_path, option_name="fsdp_module_paths_to_wrap") + fully_shard( + module, + mesh=fsdp_mesh, + reshard_after_forward=False if module_path in no_reshard_paths else reshard_after_forward, + mp_policy=mp_policy, + ) + wrapped_paths += 1 + if self.rank == 0 and wrapped_paths: + logger.info( + "FSDP2: wrapped explicit module paths under " f"'{wrap_target or ''}': {module_paths}" + ) # Wrap transformer blocks first for optimal memory management layer_cls_spec = self.args.get("fsdp_transformer_layer_cls_to_wrap") @@ -180,18 +209,30 @@ def _apply_fsdp2(self): if self.rank == 0: logger.info(f"FSDP2: applied fully_shard to '{wrap_target or ''}' with mp={mp_dtype}") + def _compile_transformer_blocks(self, root: torch.nn.Module) -> None: + compiled = 0 + for attr in ("double_blocks", "single_blocks"): + blocks = getattr(root, attr, None) + if blocks is None: + continue + for idx, block in enumerate(blocks): + blocks[idx] = torch.compile(block, fullgraph=True) + compiled += 1 + if self.rank == 0: + logger.info(f"FSDP2: compiled {compiled} FLUX transformer blocks with torch.compile") + @staticmethod - def _get_module_by_path(root: torch.nn.Module, path: str) -> torch.nn.Module: + def _get_module_by_path(root: torch.nn.Module, path: str, *, option_name: str) -> torch.nn.Module: """Resolve a dot-separated attribute path on a module.""" cur = root if not path: return cur for part in path.split("."): if not hasattr(cur, part): - raise ValueError(f"fsdp2_wrap_target='{path}' is invalid: missing attribute '{part}'") + raise ValueError(f"{option_name}='{path}' is invalid: missing attribute '{part}'") cur = getattr(cur, part) if not isinstance(cur, torch.nn.Module): - raise ValueError(f"fsdp2_wrap_target='{path}' did not resolve to a torch.nn.Module") + raise ValueError(f"{option_name}='{path}' did not resolve to a torch.nn.Module") return cur # ------------------------------------------------------------------ # @@ -370,6 +411,10 @@ def _save_dit(self, save_path: str) -> None: def save_model(self): """Save final model using the configured strategy.""" + if self.save_strategy in ("none", "skip", "disabled"): + if self.rank == 0: + logger.info(f"Skipping final checkpoint save (strategy={self.save_strategy})") + return if self.save_strategy == "dit_only": save_path = os.path.join(self.output_dir, "dit_model.safetensors") self._save_dit(save_path) diff --git a/primus/backends/diffusion/utils/__init__.py b/primus/backends/diffusion/utils/__init__.py index 3d7aa5fe7..b4ecc5389 100644 --- a/primus/backends/diffusion/utils/__init__.py +++ b/primus/backends/diffusion/utils/__init__.py @@ -4,7 +4,7 @@ # See LICENSE for license information. ############################################################################### -"""Utility package for the Primus Wan backend. +"""Utility package for the Primus diffusion backend. -Heavy vision helpers are imported lazily by the qwen_vl_utils video path. +Heavy vision helpers are imported lazily by the video dataset path. """ diff --git a/primus/backends/diffusion/utils/train_utils.py b/primus/backends/diffusion/utils/train_utils.py index 315ee9e67..219fb6957 100644 --- a/primus/backends/diffusion/utils/train_utils.py +++ b/primus/backends/diffusion/utils/train_utils.py @@ -67,18 +67,6 @@ def get_memory(unit=1e9): return allocated / unit, reserved / unit, max_alloc / unit -def print_cuda_memory(prefix="", unit=1e9): - allocated = torch.cuda.memory_allocated() / unit - reserved = torch.cuda.memory_reserved() / unit - max_alloc = torch.cuda.max_memory_allocated() / unit - print( - f"{prefix} " - f"allocated={allocated:.2f}GB, " - f"reserved={reserved:.2f}GB, " - f"max_alloc={max_alloc:.2f}GB" - ) - - @contextmanager def init_weights_on_device(device=torch.device("meta"), include_buffers: bool = False): diff --git a/primus/configs/models/diffusion/flux.1_dev_t2i.yaml b/primus/configs/models/diffusion/flux.1_dev_t2i.yaml new file mode 100644 index 000000000..983a956bb --- /dev/null +++ b/primus/configs/models/diffusion/flux.1_dev_t2i.yaml @@ -0,0 +1,15 @@ +model: + name: flux.1-dev + config: + load_from_pretrained_path: ${PRETRAINED_PATH:} + config: + trainable_modules: dit + guidance: ${FLUX_GUIDANCE:1.0} + autoencoder_scale_factor: 0.3611 + autoencoder_shift_factor: 0.1159 + # Required only for raw image-text data. + encoder: + t5_encoder: ${T5_ENCODER:} + clip_encoder: ${CLIP_ENCODER:} + autoencoder: ${VAE_CHECKPOINT:} + max_t5_length: ${MAX_T5_LENGTH:512} diff --git a/primus/configs/models/diffusion/flux.1_schnell_t2i.yaml b/primus/configs/models/diffusion/flux.1_schnell_t2i.yaml new file mode 100644 index 000000000..f018bdf90 --- /dev/null +++ b/primus/configs/models/diffusion/flux.1_schnell_t2i.yaml @@ -0,0 +1,15 @@ +model: + name: flux.1-schnell + config: + load_from_pretrained_path: ${PRETRAINED_PATH:} + config: + trainable_modules: dit + guidance: ${FLUX_GUIDANCE:1.0} + autoencoder_scale_factor: 0.3611 + autoencoder_shift_factor: 0.1159 + # Required only for raw image-text data. + encoder: + t5_encoder: ${T5_ENCODER:} + clip_encoder: ${CLIP_ENCODER:} + autoencoder: ${VAE_CHECKPOINT:} + max_t5_length: ${MAX_T5_LENGTH:256} diff --git a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py index b23e6a5ae..4bdd711bd 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/prepare.py +++ b/runner/helpers/hooks/train/pretrain/diffusion/prepare.py @@ -7,7 +7,6 @@ from __future__ import annotations import argparse -import os import sys from pathlib import Path from typing import Any @@ -71,6 +70,42 @@ def _require_path(path: str | None, description: str, *, kind: str = "any") -> N _log(f"{description}: {resolved}") +def _require_optional_path(path: str | None, description: str, *, kind: str = "any") -> None: + if _is_placeholder(path): + _log(f"{description}: not configured; skipping optional initialization") + return + _require_path(path, description, kind=kind) + + +def _looks_like_local_path(value: str) -> bool: + if value.startswith(("/", "./", "../", "~")): + return True + parts = value.split("/") + return len(parts) == 2 and parts[-1].endswith((".safetensors", ".bin", ".pt", ".pth", ".ckpt")) + + +def _looks_like_hf_reference(value: str) -> bool: + if _looks_like_local_path(value): + return False + parts = value.split("/") + return len(parts) >= 2 and all(part for part in parts[:2]) + + +def _validate_local_or_hf_id(value: str | None, description: str) -> None: + if _is_placeholder(value): + _fail(f"{description} is not configured: {value!r}") + value = str(value) + path = Path(value).expanduser() + if path.exists(): + _log(f"{description}: {path}") + elif _looks_like_local_path(value): + _fail(f"{description} path not found: {path}") + elif _looks_like_hf_reference(value): + _log(f"{description}: {value} (assuming Hugging Face reference)") + else: + _fail(f"{description} is neither an existing path nor a Hugging Face reference: {value!r}") + + def validate_diffusion_config(config_path: Path, module_name: str | None = None) -> None: cfg = load_primus_config(config_path) selected_module = _select_module_name(cfg, module_name) @@ -87,16 +122,62 @@ def validate_diffusion_config(config_path: Path, module_name: str | None = None) dataset_cfg = dataset.get("config", {}) processor_cfg = dataset_cfg.get("processor_config", {}) - encoder_cfg = model.get("config", {}).get("encoder", {}) or model.get("encoder", {}) - - _require_path(dataset_cfg.get("dataset_path"), "dataset metadata", kind="file") - _require_path(dataset_cfg.get("data_folder"), "dataset media folder", kind="dir") - _require_path(processor_cfg.get("text_tokenizer"), "text tokenizer", kind="dir") - model_cfg = model.get("config", {}) - _require_path(model_cfg.get("load_from_pretrained_path"), "DiT initialization checkpoint") - _require_path(encoder_cfg.get("t5_encoder"), "text encoder checkpoint", kind="file") - _require_path(encoder_cfg.get("autoencoder"), "VAE checkpoint", kind="file") + model_name = model.get("name") + + if model_name == "wan": + encoder_cfg = model_cfg.get("encoder", {}) or model.get("encoder", {}) + _require_path(dataset_cfg.get("dataset_path"), "dataset metadata", kind="file") + _require_path(dataset_cfg.get("data_folder"), "dataset media folder", kind="dir") + _require_path(processor_cfg.get("text_tokenizer"), "text tokenizer", kind="dir") + _require_path(model_cfg.get("load_from_pretrained_path"), "DiT initialization checkpoint") + _require_path(encoder_cfg.get("t5_encoder"), "text encoder checkpoint", kind="file") + _require_path(encoder_cfg.get("autoencoder"), "VAE checkpoint", kind="file") + elif str(model_name).startswith("flux"): + dataset_type = str(dataset_cfg.get("dataset_type", "precomputed")).lower() + if dataset_type == "raw": + dataset_name = dataset_cfg.get("dataset") + dataset_format = str(dataset_cfg.get("dataset_format", "webdataset")).lower() + if dataset_name == "cc12m-test": + if _is_placeholder(dataset_cfg.get("dataset_path")): + _log("FLUX raw image-text dataset: zirui3/cc12m-test (Hugging Face dataset)") + elif dataset_format == "hf_repo": + _validate_local_or_hf_id(dataset_cfg.get("dataset_path"), "FLUX raw Hugging Face dataset") + else: + kind = "file" if dataset_format == "jsonl" else "dir" + _require_path(dataset_cfg.get("dataset_path"), "FLUX raw image-text dataset", kind=kind) + elif dataset_name == "cc12m-wds": + if _is_placeholder(dataset_cfg.get("dataset_path")): + _log("FLUX raw image-text dataset: pixparse/cc12m-wds (Hugging Face dataset)") + else: + _validate_local_or_hf_id(dataset_cfg.get("dataset_path"), "FLUX raw Hugging Face dataset") + elif dataset_format == "hf_repo": + _validate_local_or_hf_id(dataset_cfg.get("dataset_path"), "FLUX raw Hugging Face dataset") + else: + kind = "file" if dataset_format == "jsonl" else "dir" + _require_path(dataset_cfg.get("dataset_path"), "FLUX raw image-text dataset", kind=kind) + encoder_cfg = model_cfg.get("encoder", {}) or model.get("encoder", {}) + _validate_local_or_hf_id(encoder_cfg.get("t5_encoder"), "FLUX T5 encoder") + _validate_local_or_hf_id(encoder_cfg.get("clip_encoder"), "FLUX CLIP encoder") + _validate_local_or_hf_id(encoder_cfg.get("autoencoder"), "FLUX autoencoder checkpoint") + elif dataset_type == "precomputed": + _require_path(dataset_cfg.get("dataset_path"), "FLUX precomputed dataset", kind="dir") + prompt_dropout_prob = float(processor_cfg.get("prompt_dropout_prob", 0.0) or 0.0) + if prompt_dropout_prob > 0.0: + empty_dir = processor_cfg.get("empty_encodings_path") + _require_path(empty_dir, "FLUX empty encodings directory", kind="dir") + _require_path(str(Path(empty_dir) / "t5_empty.npy"), "FLUX empty T5 encoding", kind="file") + _require_path( + str(Path(empty_dir) / "clip_empty.npy"), "FLUX empty CLIP encoding", kind="file" + ) + else: + _fail(f"unsupported FLUX dataset_type: {dataset_type!r}") + _require_optional_path( + model_cfg.get("load_from_pretrained_path"), + "FLUX DiT initialization checkpoint", + ) + else: + _fail(f"unsupported diffusion model for prepare hook: {model_name!r}") _log( "validated " @@ -118,10 +199,6 @@ def main() -> None: parser.add_argument("--module_name", type=str, default=None, help="Override module name to validate") args, _unknown = parser.parse_known_args() - if os.environ.get("SKIP_PREPARE") == "1": - _log("SKIP_PREPARE=1; skipping validation") - return - if args.backend_path: _fail("diffusion is an in-tree backend and does not support --backend_path") diff --git a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt index 38fabbeef..50e712cd7 100644 --- a/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt +++ b/runner/helpers/hooks/train/pretrain/diffusion/requirements-diffusion.txt @@ -1,6 +1,10 @@ einops +datasets +webdataset +huggingface_hub loguru safetensors +sentencepiece numpy pillow packaging diff --git a/tests/unit_tests/backends/__init__.py b/tests/unit_tests/backends/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/backends/diffusion/__init__.py b/tests/unit_tests/backends/diffusion/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py b/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py new file mode 100644 index 000000000..9a97445dd --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_flow_match_scheduler.py @@ -0,0 +1,14 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +import pytest + +from primus.backends.diffusion.schedulers.flow_match import FlowMatchScheduler + + +def test_exponential_shift_requires_mu_or_dynamic_shift_len(): + with pytest.raises(ValueError, match="exponential_shift=True"): + FlowMatchScheduler(exponential_shift=True) diff --git a/tests/unit_tests/backends/diffusion/test_flux_backend.py b/tests/unit_tests/backends/diffusion/test_flux_backend.py new file mode 100644 index 000000000..83f520c3b --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_flux_backend.py @@ -0,0 +1,449 @@ +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from primus.backends.diffusion.argument_builder import DiffusionArgBuilder +from primus.backends.diffusion.attention import ( + get_attention_backend, + set_attention_backend, +) +from primus.backends.diffusion.data.flux_precomputed import ( + FluxPrecomputedProcessor, + FluxRawImageTextDataset, + FluxRawImageTextProcessor, +) +from primus.backends.diffusion.models.flux.adapter import FluxForTraining +from primus.backends.diffusion.models.flux.conditioner import HFEmbedder +from primus.backends.diffusion.models.flux.math import apply_rope +from primus.backends.diffusion.models.flux.math import attention as flux_attention +from primus.backends.diffusion.models.flux.math import rope +from primus.backends.diffusion.models.registrations.flux import build_flux_model +from primus.backends.diffusion.trainers.fsdp2 import FSDP2Trainer + + +def _finalize(params: dict): + builder = DiffusionArgBuilder() + builder.update(params) + return builder.finalize() + + +def test_flux_argument_builder_selects_flux_defaults(): + args = _finalize( + { + "model": {"name": "flux.1-dev", "config": {}}, + "training": {"steps": 7, "local_batch_size": 3}, + "data": { + "dataset_path": "/tmp/precomputed", + "dataset_type": "precomputed", + "empty_encodings_path": "/tmp/empty", + "prompt_dropout_prob": 0.25, + }, + "lr_scheduler": { + "lr_scheduler_type": "constant_with_warmup", + "warmup_steps": 11, + }, + } + ) + + assert args.model["name"] == "flux.1-dev" + assert args.dataset["name"] == "flux" + assert args.dataset["config"]["dataset_path"] == "/tmp/precomputed" + assert args.dataset["config"]["processor_config"]["empty_encodings_path"] == "/tmp/empty" + assert args.dataset["config"]["processor_config"]["prompt_dropout_prob"] == 0.25 + assert args.trainer["args"]["max_steps"] == 7 + assert args.trainer["args"]["per_device_train_batch_size"] == 3 + assert args.trainer["args"]["lr_scheduler_type"] == "constant_with_warmup" + assert args.trainer["args"]["warmup_steps"] == 11 + assert args.trainer["args"]["attention_backend"] == "flash_attn_aiter" + assert args.trainer["args"]["fsdp_transformer_layer_cls_to_wrap"] == "DoubleStreamBlock,SingleStreamBlock" + assert args.trainer["args"]["compile_transformer_blocks"] is True + + +def test_flux_argument_builder_maps_raw_dataset_type(): + args = _finalize( + { + "model": {"name": "flux.1-dev", "config": {}}, + "data": { + "dataset_type": "raw", + "dataset_format": "webdataset", + "dataset_path": "/tmp/cc12m_test", + "prompt_dropout_prob": 0.1, + "img_size": 128, + }, + } + ) + + assert args.dataset["config"]["dataset_type"] == "raw" + assert args.dataset["config"]["dataset_format"] == "webdataset" + assert args.dataset["config"]["processor_config"]["img_size"] == 128 + + +def test_flux_raw_dataset_name_defaults(): + path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-test", None, "webdataset") + assert path == "zirui3/cc12m-test" + assert fmt == "hf_repo" + + path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-test", "/tmp/cc12m_test", "webdataset") + assert path == "/tmp/cc12m_test" + assert fmt == "webdataset" + + path, fmt = FluxRawImageTextDataset._resolve_dataset("cc12m-wds", None, "webdataset") + assert path == "pixparse/cc12m-wds" + assert fmt == "hf_repo" + + +def test_flux_argument_builder_rejects_sequence_parallelism(): + with pytest.raises(ValueError, match="sp_size"): + _finalize( + { + "model": {"name": "flux.1-dev", "config": {}}, + "parallelism": {"sp_size": 2}, + } + ) + + +def test_flux_attention_dispatch_matches_sdpa_layout(): + previous_backend = get_attention_backend() + set_attention_backend("sdpa") + try: + torch.manual_seed(7) + q = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16) + k = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16) + v = torch.randn(2, 3, 5, 4, dtype=torch.bfloat16) + pos = torch.arange(5, dtype=torch.float32).repeat(2, 1) + pe = rope(pos, dim=4, theta=10000).unsqueeze(1) + + actual = flux_attention(q, k, v, pe=pe) + q_rope, k_rope = apply_rope(q, k, pe) + expected = torch.nn.functional.scaled_dot_product_attention(q_rope, k_rope, v) + expected = expected.transpose(1, 2).flatten(2) + + torch.testing.assert_close(actual, expected) + finally: + set_attention_backend(previous_backend) + + +def test_flux_hf_embedder_passes_attention_mask(): + class FakeTokenizer: + def __call__(self, text, **kwargs): + assert kwargs["padding"] == "max_length" + return { + "input_ids": torch.tensor([[1, 2, 0], [3, 0, 0]], dtype=torch.long), + "attention_mask": torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long), + } + + class FakeTextModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.param = torch.nn.Parameter(torch.zeros(())) + self.seen_attention_mask = None + + def forward(self, *, input_ids, attention_mask, output_hidden_states): + self.seen_attention_mask = attention_mask + return {"last_hidden_state": input_ids.float().unsqueeze(-1)} + + embedder = HFEmbedder.__new__(HFEmbedder) + torch.nn.Module.__init__(embedder) + embedder.tokenizer = FakeTokenizer() + embedder.hf_module = FakeTextModel() + embedder.output_key = "last_hidden_state" + embedder.max_length = 3 + + output = embedder(["short", "x"]) + + torch.testing.assert_close( + embedder.hf_module.seen_attention_mask, + torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long), + ) + assert output.shape == (2, 3, 1) + + +def test_flux_forward_uses_positional_scheduler(): + model = FluxForTraining( + dit=torch.nn.Identity(), + train_pipeline=object(), + model_config=object(), + ) + scheduler = object() + captured = {} + + def forward_train(batch, scheduler=None): + captured["batch"] = batch + captured["scheduler"] = scheduler + return {"loss": torch.tensor(0.0)} + + model.forward_train = forward_train + batch = {"x": torch.tensor(1)} + + output = model(batch, scheduler) + + assert output["loss"].item() == 0.0 + assert captured == {"batch": batch, "scheduler": scheduler} + + +def test_fsdp2_compile_transformer_blocks_replaces_modules(monkeypatch): + class CompiledBlock(torch.nn.Module): + def __init__(self, original): + super().__init__() + self.original = original + + root = torch.nn.Module() + root.double_blocks = torch.nn.ModuleList([torch.nn.Identity(), torch.nn.ReLU()]) + root.single_blocks = torch.nn.ModuleList([torch.nn.Sigmoid()]) + compiled_inputs = [] + + def fake_compile(module, *, fullgraph): + assert fullgraph is True + compiled_inputs.append(module) + return CompiledBlock(module) + + monkeypatch.setattr(torch, "compile", fake_compile) + trainer = FSDP2Trainer.__new__(FSDP2Trainer) + trainer.rank = 1 + + trainer._compile_transformer_blocks(root) + + assert len(compiled_inputs) == 3 + assert all(isinstance(block, CompiledBlock) for block in root.double_blocks) + assert all(isinstance(block, CompiledBlock) for block in root.single_blocks) + + +def test_flux_precomputed_processor_stacks_and_drops_empty_encodings(tmp_path): + empty_dir = tmp_path / "empty" + empty_dir.mkdir() + np.save(empty_dir / "t5_empty.npy", np.zeros((1, 3, 8), dtype=np.float32)) + np.save(empty_dir / "clip_empty.npy", np.zeros((1, 4), dtype=np.float32)) + + processor = FluxPrecomputedProcessor( + { + "prompt_dropout_prob": 1.0, + "empty_encodings_path": str(empty_dir), + } + ) + batch = [ + { + "t5_encodings": torch.ones(3, 8), + "clip_encodings": torch.ones(4), + "mean": torch.zeros(1, 2, 2), + "logvar": torch.zeros(1, 2, 2), + }, + { + "t5_encodings": torch.ones(3, 8), + "clip_encodings": torch.ones(4), + "mean": torch.zeros(1, 2, 2), + "logvar": torch.zeros(1, 2, 2), + }, + ] + + out = processor.prepare_batch(batch=batch, device=torch.device("cpu"), dtype=torch.float32) + + assert out["t5_encodings"].shape == (2, 3, 8) + assert out["clip_encodings"].shape == (2, 4) + assert torch.count_nonzero(out["t5_encodings"]) == 0 + assert torch.count_nonzero(out["clip_encodings"]) == 0 + + +def test_flux_precomputed_processor_rejects_mismatched_empty_encoding(tmp_path): + empty_dir = tmp_path / "empty" + empty_dir.mkdir() + # Empty T5 encoding has sequence length 5, but the batch samples use length 3. + np.save(empty_dir / "t5_empty.npy", np.zeros((1, 5, 8), dtype=np.float32)) + np.save(empty_dir / "clip_empty.npy", np.zeros((1, 4), dtype=np.float32)) + + processor = FluxPrecomputedProcessor( + { + "prompt_dropout_prob": 1.0, + "empty_encodings_path": str(empty_dir), + } + ) + batch = [ + { + "t5_encodings": torch.ones(3, 8), + "clip_encodings": torch.ones(4), + "mean": torch.zeros(1, 2, 2), + "logvar": torch.zeros(1, 2, 2), + } + ] + + with pytest.raises(ValueError, match="empty T5 encoding shape"): + processor.prepare_batch(batch=batch, device=torch.device("cpu"), dtype=torch.float32) + + +def test_flux_raw_processor_prepares_images_and_prompts(): + from PIL import Image + + processor = FluxRawImageTextProcessor( + {"img_size": 8, "prompt_dropout_prob": 1.0, "skip_low_resolution": False} + ) + image = Image.fromarray(np.full((6, 10, 3), 127, dtype=np.uint8)) + + out = processor.prepare_batch( + batch=[{"image": image, "prompt": "a test image"}], + device=torch.device("cpu"), + dtype=torch.float32, + ) + + assert out["image"].shape == (1, 3, 8, 8) + assert out["prompts"] == [""] + + +def test_tiny_flux_model_computes_precomputed_loss(): + model = build_flux_model( + { + "config": { + "model_preset": "flux.1-dev", + "guidance": 1.0, + "params": { + "in_channels": 4, + "out_channels": 4, + "vec_in_dim": 4, + "context_in_dim": 8, + "hidden_size": 12, + "num_heads": 2, + "depth": 1, + "depth_single_blocks": 1, + "axes_dim": [2, 2, 2], + }, + } + } + ) + batch = { + "t5_encodings": torch.randn(2, 3, 8), + "clip_encodings": torch.randn(2, 4), + "mean": torch.randn(2, 1, 2, 2), + "logvar": torch.zeros(2, 1, 2, 2), + } + + outputs = model.forward_train(batch) + + assert outputs["loss"].ndim == 0 + assert torch.isfinite(outputs["loss"]) + + +def test_tiny_flux_schnell_model_computes_without_guidance(): + model = build_flux_model( + { + "config": { + "model_preset": "flux.1-schnell", + "params": { + "in_channels": 4, + "out_channels": 4, + "vec_in_dim": 4, + "context_in_dim": 8, + "hidden_size": 12, + "num_heads": 2, + "depth": 1, + "depth_single_blocks": 1, + "axes_dim": [2, 2, 2], + }, + } + } + ) + batch = { + "t5_encodings": torch.randn(2, 3, 8), + "clip_encodings": torch.randn(2, 4), + "mean": torch.randn(2, 1, 2, 2), + "logvar": torch.zeros(2, 1, 2, 2), + } + + outputs = model.forward_train(batch) + + assert model.dit.params.guidance_embed is False + assert model.model_config.guidance is None + assert outputs["loss"].ndim == 0 + assert torch.isfinite(outputs["loss"]) + + +def test_flux_position_ids_are_float32_regardless_of_model_dtype(): + model = build_flux_model( + { + "config": { + "model_preset": "flux.1-dev", + "guidance": 1.0, + "params": { + "in_channels": 4, + "out_channels": 4, + "vec_in_dim": 4, + "context_in_dim": 8, + "hidden_size": 12, + "num_heads": 2, + "depth": 1, + "depth_single_blocks": 1, + "axes_dim": [2, 2, 2], + }, + } + } + ) + # Force a low-precision compute dtype; position ids must stay float32 so RoPE + # grid indices are not corrupted (bf16 only represents integers up to 256). + model.dit = model.dit.to(dtype=torch.bfloat16) + + captured = {} + original_forward = model.dit.forward + + def capturing_forward(*args, **kwargs): + captured["img_ids_dtype"] = kwargs["img_ids"].dtype + captured["txt_ids_dtype"] = kwargs["txt_ids"].dtype + return original_forward(*args, **kwargs) + + model.dit.forward = capturing_forward + model.forward_train( + { + "t5_encodings": torch.randn(2, 3, 8, dtype=torch.bfloat16), + "clip_encodings": torch.randn(2, 4, dtype=torch.bfloat16), + "mean": torch.randn(2, 1, 2, 2, dtype=torch.bfloat16), + "logvar": torch.zeros(2, 1, 2, 2, dtype=torch.bfloat16), + } + ) + + assert captured["img_ids_dtype"] == torch.float32 + assert captured["txt_ids_dtype"] == torch.float32 + + +def test_tiny_flux_model_computes_raw_loss_with_dummy_encoders(): + class DummyAutoencoder(torch.nn.Module): + def encode(self, image): + return image[:, :1, :2, :2] + + class DummyT5(torch.nn.Module): + def forward(self, prompts): + return torch.zeros(len(prompts), 3, 8) + + class DummyClip(torch.nn.Module): + def forward(self, prompts): + return torch.zeros(len(prompts), 4) + + model = build_flux_model( + { + "config": { + "model_preset": "flux.1-dev", + "guidance": 1.0, + "params": { + "in_channels": 4, + "out_channels": 4, + "vec_in_dim": 4, + "context_in_dim": 8, + "hidden_size": 12, + "num_heads": 2, + "depth": 1, + "depth_single_blocks": 1, + "axes_dim": [2, 2, 2], + }, + } + } + ) + model.autoencoder = DummyAutoencoder() + model.t5_encoder = DummyT5() + model.clip_encoder = DummyClip() + + outputs = model.forward_train( + { + "image": torch.randn(2, 3, 4, 4), + "prompts": ["cat", "dog"], + } + ) + + assert outputs["loss"].ndim == 0 + assert torch.isfinite(outputs["loss"]) diff --git a/tests/unit_tests/backends/diffusion/test_review_regressions.py b/tests/unit_tests/backends/diffusion/test_review_regressions.py new file mode 100644 index 000000000..041a9b24c --- /dev/null +++ b/tests/unit_tests/backends/diffusion/test_review_regressions.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from primus.backends.diffusion.data.collator import VisionCollator +from primus.backends.diffusion.data.dataset import WanVideoDataset +from primus.backends.diffusion.diffusion_pretrain_trainer import ( + DiffusionPretrainTrainer, +) + + +def test_vision_collator_uses_provided_attention_mask_without_input_ids(): + class Tokenizer: + padding_side = "right" + pad_token_id = 0 + + processor = SimpleNamespace(tokenizer=Tokenizer()) + collator = VisionCollator(processor=processor) + + batch = collator( + [ + {"attention_mask": torch.tensor([1, 1]), "pixel_values": torch.ones(2)}, + {"attention_mask": torch.tensor([1]), "pixel_values": torch.zeros(2)}, + ] + ) + + torch.testing.assert_close(batch["attention_mask"], torch.tensor([[1, 1], [1, 0]])) + torch.testing.assert_close(batch["pixel_values"], torch.tensor([[1.0, 1.0], [0.0, 0.0]])) + + +def test_wan_video_dataset_requires_config_with_dataset_path(): + with pytest.raises(ValueError, match="dataset config with dataset_path"): + WanVideoDataset(processor=object()) + + with pytest.raises(ValueError, match="config.dataset_path"): + WanVideoDataset(processor=object(), config=SimpleNamespace(dataset_path=None)) + + +def test_diffusion_setup_reports_missing_flux_dependencies(monkeypatch): + missing = {"datasets", "huggingface_hub", "sentencepiece", "webdataset"} + + def fake_find_spec(package): + return None if package in missing else object() + + monkeypatch.setattr("importlib.util.find_spec", fake_find_spec) + trainer = DiffusionPretrainTrainer.__new__(DiffusionPretrainTrainer) + trainer.backend_args = SimpleNamespace( + trainer={"args": {}}, + dataset={ + "name": "flux", + "config": { + "dataset_type": "raw", + "dataset_format": "webdataset", + }, + }, + ) + + with pytest.raises(RuntimeError) as exc_info: + trainer.setup() + + message = str(exc_info.value) + for package in sorted(missing): + assert package in message diff --git a/tests/unit_tests/backends/megatron/__init__.py b/tests/unit_tests/backends/megatron/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/backends/torchtitan/__init__.py b/tests/unit_tests/backends/torchtitan/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/cli/__init__.py b/tests/unit_tests/cli/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/__init__.py b/tests/unit_tests/core/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/backend/__init__.py b/tests/unit_tests/core/backend/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/launcher/__init__.py b/tests/unit_tests/core/launcher/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/patches/__init__.py b/tests/unit_tests/core/patches/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/runtime/__init__.py b/tests/unit_tests/core/runtime/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/trainer/__init__.py b/tests/unit_tests/core/trainer/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/core/utils/__init__.py b/tests/unit_tests/core/utils/__init__.py new file mode 100644 index 000000000..e69de29bb From 0aad8d274673e249996e427d6dae390081082acc Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Fri, 17 Jul 2026 11:37:50 +0800 Subject: [PATCH 037/127] feat: add use_turbo_autotune flag and refine moe_router_force_load_balancing_type flag (#880) # Description This PR introduces a new `use_turbo_autotune` flag to enable Primus-Turbo auto-tuning, and refines the MoE router force load balancing configuration by moving the `moe_router_force_load_balancing_type` flag to a more appropriate module config. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [x] Code refactoring ## Changes Please list the changes introduced in this PR: - Add `use_turbo_autotune` flag to `primus_turbo.yaml` (default `false`), which sets the `PRIMUS_TURBO_AUTO_TUNE=1` environment variable during ROCm arg validation to enable Primus-Turbo auto-tuning. - Refine the `moe_router_force_load_balancing_type` flag by moving it from `primus_megatron_model.yaml` to `primus_megatron_module.yaml`. - Update `docs/03-configuration-reference/megatron-parameters.md` to document the `moe_router_force_load_balancing_type` flag and clarify Turbo GEMM flags. # Checklist: - [ ] The functionality is complete - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --- docs/03-configuration-reference/megatron-parameters.md | 5 ++++- .../backends/megatron/patches/args/rocm_arg_validation.py | 6 ++++++ primus/configs/models/megatron/primus_megatron_model.yaml | 7 ------- .../configs/modules/megatron/primus_megatron_module.yaml | 5 +++++ primus/configs/modules/megatron/primus_turbo.yaml | 4 ++++ 5 files changed, 19 insertions(+), 8 deletions(-) diff --git a/docs/03-configuration-reference/megatron-parameters.md b/docs/03-configuration-reference/megatron-parameters.md index 25c880542..97e0bfc04 100644 --- a/docs/03-configuration-reference/megatron-parameters.md +++ b/docs/03-configuration-reference/megatron-parameters.md @@ -597,6 +597,8 @@ models: | `disable_primus_topk_router` | `false` | *Primus:* disable Primus top-k router patch. | | `moe_router_force_load_balancing` | `false` | *Primus:* force load-balanced routing. | | `use_deprecated_20241209_moe_layer` | `false` | *Primus:* legacy MoE layer implementation. | +| `moe_router_force_load_balancing_type` | `even` | *Primus:* Control the force load balancing type for the MoE router. Choices: even, uniform. | + ### 10.7 Logit softcapping (Primus / Grok-style) @@ -625,7 +627,8 @@ models: | `use_sink_attention` | `false` | GPT-OSS-style learned sink attention. | | `sink_sliding_window` | `0` | Sliding-window size for sink attention (GPT-OSS uses `128`). | | `sink_window_even_layers_only` | `true` | Apply the sliding window only to even layers (GPT-OSS pattern). | -| `use_turbo_parallel_linear` | `false` | Turbo parallel linear layers. | +| `use_turbo_gemm` | `false` | Active Turbo GEMM flag for Dense paths. | +| `use_turbo_parallel_linear` | *(removed)* | Removed—use `use_turbo_gemm`. Passing this key now raises an assertion error (`use_turbo_parallel_linear has been removed; please use use_turbo_gemm instead`). | | `use_turbo_grouped_gemm` | `false` | Active Turbo grouped GEMM flag for MoE paths. | | `use_turbo_grouped_mlp` | *(removed)* | Removed—use `use_turbo_grouped_gemm`. Passing this key now raises an assertion error (`use_turbo_grouped_mlp has been removed; please use use_turbo_grouped_gemm instead`). | | `moe_use_fused_router_with_aux_score` | `false` | Fused MoE router with auxiliary scores. | diff --git a/primus/backends/megatron/patches/args/rocm_arg_validation.py b/primus/backends/megatron/patches/args/rocm_arg_validation.py index e68251f8d..95aff1b15 100644 --- a/primus/backends/megatron/patches/args/rocm_arg_validation.py +++ b/primus/backends/megatron/patches/args/rocm_arg_validation.py @@ -98,6 +98,12 @@ def validate_fsdp2_optimizer_exclusivity(args) -> None: def validate_args_on_rocm(args): + # Primus-Turbo auto-tuning + use_turbo_autotune = getattr(args, "use_turbo_autotune", False) + if use_turbo_autotune: + # NOTE: Set PRIMUS_TURBO_AUTO_TUNE to 1 to enable turbo auto-tuning. + os.environ["PRIMUS_TURBO_AUTO_TUNE"] = "1" + # Deterministic mode if args.deterministic_mode: # NOTE: Some environment variables affect deterministic mode on ROCm. Need to do extra check. diff --git a/primus/configs/models/megatron/primus_megatron_model.yaml b/primus/configs/models/megatron/primus_megatron_model.yaml index a7a12d70b..0376e6014 100644 --- a/primus/configs/models/megatron/primus_megatron_model.yaml +++ b/primus/configs/models/megatron/primus_megatron_model.yaml @@ -13,10 +13,3 @@ router_logit_softcapping: null # float lfm_layer_types: null # list[str] conv_L_cache: 3 # int conv_bias: false # bool - - -# Primus patch option -# Control the force load balancing type for the MoE router. -# If set to "even", the router will force the load balancing to be even. -# If set to "uniform", the router will force the load balancing to be uniform. (Megatron-LM original behavior) -moe_router_force_load_balancing_type: "even" diff --git a/primus/configs/modules/megatron/primus_megatron_module.yaml b/primus/configs/modules/megatron/primus_megatron_module.yaml index 838ca07a4..74f8d7179 100644 --- a/primus/configs/modules/megatron/primus_megatron_module.yaml +++ b/primus/configs/modules/megatron/primus_megatron_module.yaml @@ -96,3 +96,8 @@ recompute_layer_ids: null # int list; global layer ids, range from 0 to (num_lay # dataloader dataloader_mp_context: null # "forkserver" | "spawn" | "fork" | null + +# Control the force load balancing type for the MoE router. +# If set to "even", the router will force the load balancing to be even. +# If set to "uniform", the router will force the load balancing to be uniform. (Megatron-LM original behavior) +moe_router_force_load_balancing_type: "even" diff --git a/primus/configs/modules/megatron/primus_turbo.yaml b/primus/configs/modules/megatron/primus_turbo.yaml index bad3403d6..db0b1e66c 100644 --- a/primus/configs/modules/megatron/primus_turbo.yaml +++ b/primus/configs/modules/megatron/primus_turbo.yaml @@ -2,6 +2,10 @@ # main control flag enable_primus_turbo: false +# ===== AutoTune ===== +# enable turbo auto-tuning +use_turbo_autotune: false + # ===== Attention ===== # operator switch use_turbo_attention: false From dd989744f16d4914352ff06757b89323ee1c290a Mon Sep 17 00:00:00 2001 From: RuibinCheung Date: Fri, 17 Jul 2026 11:55:23 +0800 Subject: [PATCH 038/127] feat: remove extra htod when enable turbo grouped gemm (#878) # Description When `use_turbo_grouped_gemm` is enabled, the MoE All-to-All token dispatcher still copies `tokens_per_expert` from device to host (D2H) inside `_maybe_dtoh_and_synchronize`. However, PrimusTurbo grouped gemm consumes `tokens_per_expert` directly on the GPU, so this D2H copy is redundant. This extra copy introduces an unnecessary device-to-host transfer and stream synchronization overhead on the critical path. This PR adds a `before_train` patch that replaces `MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize` so that `tokens_per_expert` is kept on-device when PrimusTurbo grouped gemm is enabled. All other splits (`input_splits`, `output_splits`, `output_splits_tp`, `num_out_tokens`, `num_global_tokens_per_local_expert`) are still moved to CPU, and the DtoH stream synchronization behavior is unchanged. Fixes # (issue) ## Type of change - [ ] Documentation change (change only to the documentation, either a fix or a new content) - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Infra/Build change - [ ] Code refactoring ## Changes Please list the changes introduced in this PR: - Add `primus/backends/megatron/patches/moe_patches/moe_alltoall_dtoh_patches.py`, registering the `megatron.moe_alltoall_dtoh_turbo_grouped_gemm` patch. - The patch is applied at the `before_train` phase and is only active when `use_turbo_grouped_gemm` is set (guarded by a `condition`). - Override `MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize` to skip the `tokens_per_expert` D2H copy while keeping all other splits' CPU migration and the DtoH stream sync unchanged. # Checklist: - [x] The functionality is complete - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes --- .../moe_patches/moe_alltoall_dtoh_patches.py | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 primus/backends/megatron/patches/moe_patches/moe_alltoall_dtoh_patches.py diff --git a/primus/backends/megatron/patches/moe_patches/moe_alltoall_dtoh_patches.py b/primus/backends/megatron/patches/moe_patches/moe_alltoall_dtoh_patches.py new file mode 100644 index 000000000..17a75a67b --- /dev/null +++ b/primus/backends/megatron/patches/moe_patches/moe_alltoall_dtoh_patches.py @@ -0,0 +1,85 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Primus MoE All-to-All dispatcher D2H patch. + +Patches ``MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize`` so that when +``use_turbo_grouped_gemm`` is enabled, ``tokens_per_expert`` is kept on-device +(PrimusTurbo grouped gemm consumes it on the GPU) instead of being copied to the +host. All other splits are still moved to CPU and the stream sync is unchanged. +""" + +import torch + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +@register_patch( + "megatron.moe_alltoall_dtoh_turbo_grouped_gemm", + backend="megatron", + phase="before_train", + description=( + "Skip tokens_per_expert D2H copy in MoEAlltoAllTokenDispatcher " + "when PrimusTurbo grouped gemm is enabled" + ), + condition=lambda ctx: getattr(get_args(ctx), "use_turbo_grouped_gemm", False), +) +def patch_moe_alltoall_dtoh(ctx: PatchContext): + """Replace ``MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize``.""" + from megatron.core.transformer.moe import token_dispatcher + + cls = token_dispatcher.MoEAlltoAllTokenDispatcher + + def _maybe_dtoh_and_synchronize(self, point, tokens_per_expert=None): + """ + Move all possible GPU tensors to CPU and make a synchronization at the expected point. + """ + maybe_move_tensor_to_cpu = token_dispatcher.maybe_move_tensor_to_cpu + + if not self.drop_and_pad: + if point == self.cuda_dtoh_point: + # Move all possible GPU tensors to CPU at self.cuda_dtoh_point. + on_side_stream = torch.cuda.current_stream() != self.cuda_dtoh_stream + if on_side_stream: + self.cuda_dtoh_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.cuda_dtoh_stream): + # NOTE: PrimusTurbo grouped gemm consumes tokens_per_expert on device, + # so keep it on the GPU and skip the D2H copy when enabled. + # tokens_per_expert = maybe_move_tensor_to_cpu( + # tokens_per_expert, record_stream=on_side_stream + # ) + self.input_splits = maybe_move_tensor_to_cpu( + self.input_splits, as_numpy=True, record_stream=on_side_stream + ) + self.output_splits = maybe_move_tensor_to_cpu( + self.output_splits, as_numpy=True, record_stream=on_side_stream + ) + self.output_splits_tp = maybe_move_tensor_to_cpu( + self.output_splits_tp, as_numpy=True, record_stream=on_side_stream + ) + self.num_out_tokens = maybe_move_tensor_to_cpu( + self.num_out_tokens, record_stream=on_side_stream + ) + if self.num_local_experts > 1 and not self.config.moe_permute_fusion: + self.num_global_tokens_per_local_expert = maybe_move_tensor_to_cpu( + self.num_global_tokens_per_local_expert, record_stream=on_side_stream + ) + self.d2h_event = self.cuda_dtoh_stream.record_event() + + if point == self.cuda_sync_point: + # Synchronize with the DtoH stream at self.cuda_sync_point. + self.d2h_event.synchronize() + + return tokens_per_expert + + cls._maybe_dtoh_and_synchronize = _maybe_dtoh_and_synchronize + + log_rank_0( + "[Patch:megatron.moe_alltoall_dtoh_turbo_grouped_gemm] Patched " + "MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize " + ) From b28ec7c1a36fe2fe6d738962516a059310891365 Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Fri, 17 Jul 2026 11:57:02 +0800 Subject: [PATCH 039/127] fix(mlperf): dataset prep no longer deletes sibling files in data_dir (#879) ## Problem MLPerf llama2-70b post-training failed with: ``` ValueError: Invalid pretrained checkpoint directory found: /data/megatron_checkpoints/Llama-2-70b-hf ``` The dataset prep script `download_dataset.py` ran `find ! -name '*.parquet' -exec rm -rf` over the whole `data_dir` and hashed the whole dir, assuming it holds only the dataset. But the converted checkpoint lives in the same `data_dir`, so it got deleted (and the same blanket delete can wipe unrelated data on a shared/mounted dir). ## Fix - `download_dataset.py`: move only the dataset parquets and remove only HF's own artifacts; verify via `hash_files`. - `convert_dataset.py`: verify via `hash_files` over the produced `.npy` (no longer walks the large checkpoint). - `dataset_hash.py`: replace `hash_directory` with `hash_files` (hash an explicit file list). Modified NVIDIA-origin (Apache-2.0) files carry an added AMD 'Modifications Copyright' line. ## Verification Reproduced the failure, then confirmed the fix end-to-end in a container: prep completes, the checkpoint is preserved, training loads it and proceeds; dataset hashes unchanged. pre-commit passes. --- .../mlperf_llama2_70b/convert_dataset.py | 8 ++++++-- .../recipes/mlperf_llama2_70b/dataset_hash.py | 17 +++++++++-------- .../mlperf_llama2_70b/download_dataset.py | 15 ++++++++++++--- 3 files changed, 27 insertions(+), 13 deletions(-) diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py index 2a8f96c0f..146ba3210 100644 --- a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/convert_dataset.py @@ -1,3 +1,5 @@ +# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -23,7 +25,7 @@ import numpy as np import pandas as pd -from dataset_hash import hash_directory +from dataset_hash import hash_files def convert(data_dir, split): @@ -54,5 +56,7 @@ def transform_row(row): executable="/bin/bash", check=True, ) - directory_hash = hash_directory(args.data_dir) + # Verify only the produced files, so other content in data_dir does not + # affect the hash. + directory_hash = hash_files([f"{args.data_dir}/train.npy", f"{args.data_dir}/validation.npy"]) print(f"Succesfully converted dataset with hash {directory_hash}") diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py index 7f53eecdc..05d3b969c 100644 --- a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/dataset_hash.py @@ -1,3 +1,5 @@ +# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,7 +15,6 @@ # limitations under the License. import hashlib -import os from concurrent.futures import ThreadPoolExecutor @@ -29,15 +30,15 @@ def hash_file_md5(file_path, chunk_size=4194304): # Default chunk size 4MB. return md5_hash.hexdigest() -def hash_directory(path): - """Hashes all files in a directory in parallel using MD5.""" +def hash_files(file_paths): + """Hashes an explicit list of files in parallel using MD5. + + Only the given files are considered, so unrelated content in the same + directory (e.g. checkpoints) does not affect the result. + """ hashes = [] with ThreadPoolExecutor() as executor: - futures = [ - executor.submit(hash_file_md5, os.path.join(root, file)) - for root, dirs, files in os.walk(path) - for file in files - ] + futures = [executor.submit(hash_file_md5, p) for p in file_paths] for future in futures: file_hash = future.result() if file_hash: diff --git a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py index 0d885e07b..fe6473dde 100644 --- a/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py +++ b/primus/backends/megatron_bridge/recipes/mlperf_llama2_70b/download_dataset.py @@ -1,3 +1,5 @@ +# Modifications Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,6 +15,7 @@ # limitations under the License. import argparse +import glob import subprocess import sys from pathlib import Path @@ -21,7 +24,7 @@ if str(_RECIPE_DIR) not in sys.path: sys.path.insert(0, str(_RECIPE_DIR)) -from dataset_hash import hash_directory +from dataset_hash import hash_files from huggingface_hub import snapshot_download parser = argparse.ArgumentParser() @@ -37,14 +40,20 @@ max_workers=16, repo_type="dataset", ) + +# Move the dataset parquets up and remove only HF's own artifacts. Do NOT +# blanket-delete data_dir: it may hold a sibling checkpoint (the old +# `find ! -name '*.parquet' -exec rm -rf` wiped it). subprocess.run( - f"mv {args.data_dir}/data/* {args.data_dir}/ && find {args.data_dir} -mindepth 1 ! -name '*.parquet' -exec rm -rf {{}} +", + f"mv {args.data_dir}/data/*.parquet {args.data_dir}/ && rm -rf {args.data_dir}/data {args.data_dir}/.cache", shell=True, executable="/bin/bash", check=True, ) -directory_hash = hash_directory(args.data_dir) +# Verify only the downloaded parquets, so other content in data_dir does not +# affect the hash. +directory_hash = hash_files(sorted(glob.glob(f"{args.data_dir}/*.parquet"))) assert ( directory_hash == "682a5f40b790a56751bf8303554efc08" ), f"Expected hash 682a5f40b790a56751bf8303554efc08, but got {directory_hash}" From 4999e928bf22913dc1903e528a8b1e675516511c Mon Sep 17 00:00:00 2001 From: wenxie-amd Date: Fri, 17 Jul 2026 16:43:12 +0800 Subject: [PATCH 040/127] Add DeepSeek-V4 training support (model, attention/MoE kernels, Muon, FP8/FP4, projection toolkit) (#882) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary This PR lands end-to-end **DeepSeek-V4** pretraining on the Megatron backend (ROCm / MI355X-CDNA4), including the V4 architecture, a family of sparse-MLA attention kernels (Triton / Gluon / FlyDSL / Primus-Turbo), Turbo MoE, the Muon optimizer path, FP8/FP4 recipes, a performance-projection toolkit + site, runnable examples, and an extensive unit-test suite. It also re-enables the torch/jax unit-test CI jobs. All DeepSeek-V4 assets live under `examples/deepseek-v4/` and `tests/unit_tests/megatron/transformer/deepseek_v4/`; core model/kernel code is under `primus/backends/megatron/...`. **Scope:** 226 files changed (203 new), ~71.4k insertions. New feature, additive — no existing model paths change behavior by default. ## What's included ### Model (`model_type=deepseek_v4`) - Hybrid attention (Dense / HCA / CSA) with attention sink + interleaved dual-RoPE, HyperConnections, and the Compressor/Indexer CSA selector. - MoE with hash + `sqrtsoftplus` routers, pre-mul clamped SwiGLU, and Multi-Token Prediction (MTP) via the upstream `MultiTokenPredictionBlock`. - V4 config schema + model/pretrain YAMLs; spec-based build (strict, no `nn.Linear` fallback); V4-aware `num_floating_point_operations` closed form. ### Attention kernels (selectable sparse-MLA backends) - `triton` (production launchers), `triton_v2`, `gluon` / `gluon_v2` / `gluon_v3`, `flydsl_v1`, and the integrated Primus-Turbo native-FlyDSL backend (`turbo`). - Architecture-aware tuned defaults; split CSA fwd + atomic-free V4/CSA bwd; dual-RoPE bf16 cast fix. - Microbenchmark: `examples/deepseek-v4/benchmark/bench_v4_attention.py` compares all backends. ### MoE / Primus-Turbo + elementwise fusions - Turbo DeepEP dispatcher, Turbo grouped GEMM, sync-free MoE. - Triton fusions: `stack_grouped_weight`, interleaved partial RoPE, Sinkhorn, HyperConnection tail, Indexer tail, router post-logits. ### Optimizer (Muon) - Muon / dist_muon via `emerging_optimizers` (auto-installed by a gated hook) with DeepSeek-V4 hybrid Newton-Schulz, batched grouped-expert orthogonalization, and AdamW chained for embedding/head/RMSNorm. ### Precision - FP8 (E4M3; `tensorwise` / `mxfp8` ue8m0 microscaling) and FP4 (MXFP4) recipes; FP8 Indexer QK path; clamped SwiGLU for FP8 outlier robustness. ### Performance-projection toolkit - `examples/deepseek-v4/projection/` — analytic FLOPs/memory/timeline projection + static site, plus a `deploy-projection.yml` GitHub Pages workflow. ### Examples & tests - Runners under `examples/deepseek-v4/`: `run_deepseek_v4.sh`, `run_deepseek_v4_flash*.sh`, `run_deepseek_v4_pro_muon*.sh`, `run_dsv4_projection_1gpu.sh`. - Comprehensive unit tests under `tests/unit_tests/megatron/transformer/deepseek_v4/` (attention fwd/bwd parity, dispatch, MoE, MTP, routers, compiled Sinkhorn, kernel fusions, etc.). ### CI - Re-enabled `run-unittest-torch`, `run-unittest-jax`, and `coverage-summary` jobs. - Pinned Primus-Turbo to the FlyDSL sparse-MLA commit; SHA-pinned all `deploy-projection.yml` actions; removed the unused `-v25.09-ainic` docker image build. ## Validation (single node, 8× MI355X / gfx950) - **Flash EP=8 proxy** (`run_deepseek_v4_flash_proxy.sh`, `triton_v2`): 20/20 iterations, ~334 ms/iter, ~800 TFLOP/s/GPU, loss descending, 0 NaN. - **Attention benchmark**: all in-tree backends run across V4-Flash/Pro × cr∈{0,4,128}. - **Pro + Muon** (`run_deepseek_v4_pro_muon.sh`): 10/10 iterations, 0 NaN, loss descending (the default 4-layer config is at the single-node HBM ceiling for 384 experts + Muon fp32 states; validated at reduced depth). ## Notes for reviewers / migration - **Turbo flag rename** (aligns with main): `use_turbo_parallel_linear` → `use_turbo_gemm`, `use_turbo_grouped_mlp` → `use_turbo_grouped_gemm`. `validate_args_on_rocm` now rejects the old names. - **Invocation**: run the example scripts from the repo root (they use `./primus-cli` + cwd-relative config paths). - **`emerging_optimizers`** is pinned (submodule + install hook) to a commit whose `OrthogonalizedOptimizer` matches Megatron-LM `muon.py` (`use_nesterov`). - Dev-only working notes are **not** included in this PR (kept out of the tree). --------- Co-authored-by: Cursor Co-authored-by: yanyuqin Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: lihuan Co-authored-by: JohnQinAMD Co-authored-by: RuibinCheung --- .github/workflows/ci.yaml | 34 +- .github/workflows/deploy-projection.yml | 66 + .github/workflows/docker/Dockerfile | 8 +- .../workflows/docker/Dockerfile_v25.09_ainic | 106 - .gitignore | 10 + .../benchmark/bench_v4_attention.py | 592 +++++ .../benchmark/bench_v4_attention_results.md | 97 + examples/deepseek-v4/projection/README.md | 93 + .../projection/design/01-overview.md | 91 + .../projection/design/02-assumptions.md | 108 + .../projection/design/03-json-schema.md | 117 + .../projection/design/04-projection-math.md | 193 ++ .../projection/design/05-deployment.md | 58 + .../projection/design/06-calibration.md | 73 + .../design/07-iteration-timeline.md | 165 ++ .../deepseek_v4_layer_trace-projection.sh | 252 ++ .../deepseek-v4/projection/site/assets/app.js | 1567 +++++++++++ .../projection/site/assets/style.css | 333 +++ .../projection/site/data/flash.json | 1727 ++++++++++++ .../deepseek-v4/projection/site/data/pro.json | 1750 +++++++++++++ .../deepseek-v4/projection/site/index.html | 134 + .../projection/tools/gen_mock_data.py | 208 ++ .../projection/tools/kernel_module_map.py | 147 ++ .../projection/tools/parse_trace.py | 651 +++++ .../deepseek-v4/projection/tools/v4_flops.py | 250 ++ examples/deepseek-v4/run_deepseek_v4.sh | 314 +++ examples/deepseek-v4/run_deepseek_v4_flash.sh | 64 + .../run_deepseek_v4_flash_proxy.sh | 280 ++ .../deepseek-v4/run_deepseek_v4_pro_muon.sh | 361 +++ .../run_deepseek_v4_pro_muon_1gpu.sh | 498 ++++ .../deepseek-v4/run_dsv4_projection_1gpu.sh | 158 ++ .../deepseek_v4_flash-BF16-pretrain.yaml | 138 + .../deepseek_v4_flash-FP8-pretrain.yaml | 166 ++ .../llama2_70b_lora_mlperf_posttrain.yaml | 2 +- .../MI355X/llama3.1_8B-pretrain-FP4.yaml | 4 +- examples/run_pretrain.sh | 11 +- .../core/extensions/_triton/__init__.py | 14 + .../extensions/_triton/multi_tensor_add.py | 207 ++ .../_triton/stack_grouped_weight.py | 413 +++ .../megatron/core/extensions/primus_turbo.py | 170 +- .../transformer_engine_spec_provider.py | 211 +- primus/backends/megatron/core/fp8_utils.py | 10 + .../core/fusions/fused_bias_swiglu.py | 342 +++ .../core/fusions/fused_pad_routing_map.py | 114 + .../core/models/deepseek_v4/__init__.py | 52 + .../core/models/deepseek_v4/build_context.py | 102 + .../models/deepseek_v4/deepseek_v4_block.py | 1190 +++++++++ .../deepseek_v4/deepseek_v4_builders.py | 301 +++ .../deepseek_v4/deepseek_v4_layer_specs.py | 697 +++++ .../models/deepseek_v4/deepseek_v4_model.py | 323 +++ .../deepseek_v4/deepseek_v4_mtp_layer.py | 189 ++ .../deepseek_v4/deepseek_v4_mtp_specs.py | 201 ++ .../deepseek_v4_transformer_config.py | 214 ++ .../core/transformer/clamped_swiglu.py | 192 ++ .../megatron/core/transformer/compressor.py | 198 ++ .../core/transformer/deepseek_v4_attention.py | 1652 ++++++++++++ .../megatron/core/transformer/dual_rope.py | 362 +++ .../megatron/core/transformer/experts.py | 64 + .../core/transformer/hyper_connection.py | 468 ++++ .../megatron/core/transformer/indexer.py | 422 +++ .../core/transformer/local_rmsnorm.py | 103 + .../core/transformer/moe/_triton/__init__.py | 0 .../transformer/moe/_triton/v4_router_post.py | 548 ++++ .../core/transformer/moe/shared_experts.py | 91 + .../core/transformer/moe/v4_hash_router.py | 270 ++ .../megatron/core/transformer/moe/v4_moe.py | 633 +++++ .../core/transformer/moe/v4_topk_router.py | 275 ++ .../core/transformer/sliding_window_kv.py | 81 + .../v4_attention_kernels/README.md | 102 + .../v4_attention_kernels/__init__.py | 219 ++ .../v4_attention_kernels/_eager/__init__.py | 18 + .../v4_attention_kernels/_eager/reference.py | 329 +++ .../_flydsl_v0_deprecated/__init__.py | 169 ++ .../_flydsl_v0_deprecated/kernels/__init__.py | 0 .../kernels/v4_attention_bwd_flydsl_mqa.py | 1270 +++++++++ .../kernels/v4_attention_fwd_flydsl_csa.py | 219 ++ .../kernels/v4_attention_fwd_flydsl_mqa.py | 143 + .../v4_csa_attention_bwd_flydsl_mqa.py | 399 +++ .../kernels/v4_csa_bwd_dq_kernel.py | 665 +++++ .../kernels/v4_csa_bwd_full_kernel.py | 664 +++++ .../kernels/v4_csa_fwd_kernel.py | 748 ++++++ .../kernels/v4_hca_bwd_dkv_pool_kernel.py | 926 +++++++ .../kernels/v4_hca_bwd_dq_pool_kernel.py | 1404 ++++++++++ .../kernels/v4_sla_bwd_dkv_kernel.py | 1113 ++++++++ .../kernels/v4_sla_bwd_dq_kernel.py | 1284 +++++++++ .../kernels/v4_sla_bwd_kernel.py | 203 ++ .../kernels/v4_sla_fwd_kernel.py | 1387 ++++++++++ .../_flydsl_v1/__init__.py | 36 + .../_flydsl_v1/dsa_bwd_dq_flydsl_kernel.py | 592 +++++ .../_flydsl_v1/dsa_bwd_v4_flydsl.py | 184 ++ .../_flydsl_v1/dsa_fwd_v4_flydsl.py | 151 ++ .../_flydsl_v1/dsa_fwd_v4_flydsl_kernel.py | 650 +++++ .../_gluon_dsa/__init__.py | 37 + .../_gluon_dsa/_dsa_bwd_gather.py | 617 +++++ .../_gluon_dsa/_dsa_bwd_preprocess.py | 50 + .../_gluon_dsa/dsa_bwd_dkv_interm.py | 229 ++ .../_gluon_dsa/dsa_bwd_dq.py | 502 ++++ .../_gluon_dsa/dsa_bwd_v4_gluon.py | 166 ++ .../_gluon_dsa/dsa_fwd_v4_gluon.py | 625 +++++ .../_gluon_v2/__init__.py | 31 + .../_gluon_v2/dsa_bwd_dkv_interm_gluon.py | 237 ++ .../_gluon_v2/dsa_bwd_dq_gluon.py | 523 ++++ .../_gluon_v2/dsa_bwd_v4_gluon.py | 175 ++ .../_gluon_v2/dsa_fwd_v4_gluon.py | 722 +++++ .../_gluon_v3/__init__.py | 31 + .../_gluon_v3/aiter_lse_fwd.py | 172 ++ .../_gluon_v3/aiter_mla_gluon.py | 999 +++++++ .../_gluon_v3/dsa_bwd_dkv_interm_gluon.py | 237 ++ .../_gluon_v3/dsa_bwd_dq_gluon.py | 523 ++++ .../_gluon_v3/dsa_bwd_v4_gluon.py | 177 ++ .../_gluon_v3/dsa_fwd_v4_gluon.py | 733 ++++++ .../_tilelang/__init__.py | 388 +++ .../v4_attention_autograd_tilelang.py | 147 ++ .../_tilelang/v4_attention_bwd_tilelang.py | 427 +++ .../_tilelang/v4_attention_fwd_tilelang.py | 480 ++++ .../_triton_common/__init__.py | 20 + .../_triton_common/compressor_pool.py | 283 ++ .../_triton_common/hc_collapse.py | 232 ++ .../_triton_common/hc_expand.py | 287 ++ .../_triton_common/hc_glue.py | 622 +++++ .../_triton_common/indexer_score.py | 503 ++++ .../_triton_common/indexer_score_post.py | 468 ++++ .../_triton_common/rmsnorm.py | 366 +++ .../rope_interleaved_partial.py | 778 ++++++ .../_triton_common/sinkhorn.py | 520 ++++ .../_triton_v0_deprecated/__init__.py | 23 + .../_triton_v0_deprecated/v4_csa_attention.py | 238 ++ .../v4_csa_attention_bwd.py | 488 ++++ .../v4_csa_attention_fwd.py | 407 +++ .../_triton_v1/__init__.py | 29 + .../_triton_v1/_v4_attn_tuning.py | 65 + .../_triton_v1/v4_attention.py | 298 +++ .../_triton_v1/v4_attention_bwd.py | 2086 +++++++++++++++ .../_triton_v1/v4_attention_fwd.py | 494 ++++ .../_triton_v1/v4_csa_attention.py | 175 ++ .../_triton_v1/v4_csa_attention_bwd.py | 2321 +++++++++++++++++ .../_triton_v1/v4_csa_attention_fwd.py | 1040 ++++++++ .../_triton_v2/__init__.py | 24 + .../_triton_v2/_amd_knobs.py | 44 + .../_triton_v2/dsa_bwd_kernels.py | 337 +++ .../_triton_v2/dsa_bwd_v4_triton.py | 229 ++ .../_triton_v2/dsa_fwd_v4_triton.py | 206 ++ .../_turbo_flydsl/__init__.py | 41 + .../v4_csa_attention_flydsl.py | 38 + .../v4_csa_attention_gluon.py | 33 + .../v4_csa_attention_gluon_v2.py | 38 + .../v4_csa_attention_gluon_v3.py | 26 + .../v4_csa_attention_triton.py | 37 + .../v4_csa_attention_turbo_flydsl.py | 39 + .../v4_sparse_mla_adapter.py | 295 +++ .../megatron/megatron_pretrain_trainer.py | 13 +- .../patches/deepseek_v4_flops_patches.py | 910 +++++++ .../patches/deepseek_v4_get_batch_patches.py | 271 ++ .../patches/deepseek_v4_pp_shape_patches.py | 162 ++ .../emerging_optimizers_log_level_patches.py | 58 + .../patches/fused_pad_routing_map_patches.py | 68 + .../backends/megatron/patches/mla_patches.py | 21 +- .../patches/moe_alltoall_dtoh_patches.py | 105 + .../patches/triton_autotune_print_patches.py | 66 + .../megatron/training/tokenizer/tokenizer.py | 1 + .../models/megatron/deepseek_v4_base.yaml | 117 + .../models/megatron/deepseek_v4_flash.yaml | 54 + .../models/megatron/deepseek_v4_pro.yaml | 46 + .../megatron/primus_megatron_module.yaml | 2 +- .../projection/module_profilers/attention.py | 30 +- .../projection/module_profilers/moe_mlp.py | 25 +- .../module_profilers/transformer_layer.py | 63 +- .../core/projection/module_profilers/utils.py | 184 +- .../performance_projection/projection.py | 23 +- primus/core/utils/import_utils.py | 18 +- rccl_avg_workaround/.gitignore | 2 + rccl_avg_workaround/sitecustomize.py | 143 + .../01_install_emerging_optimizers.sh | 58 + .../backends/megatron/test_compressor_pool.py | 58 + .../test_deepseek_v4_flops_patches.py | 647 +++++ .../megatron/test_rope_arange_cache.py | 47 + .../configs/test_deepseek_v4_yaml.py | 335 +++ tests/unit_tests/conftest.py | 37 + .../megatron/cco/test_tp_overlap.py | 17 +- .../transformer/deepseek_v4/__init__.py | 0 .../transformer/deepseek_v4/conftest.py | 80 + .../deepseek_v4/test_clamped_swiglu.py | 163 ++ .../deepseek_v4/test_deepseek_v4_attention.py | 872 +++++++ .../test_fused_hc_collapse_triton.py | 160 ++ .../deepseek_v4/test_fused_rmsnorm_triton.py | 233 ++ .../test_fused_rope_from_positions.py | 116 + .../deepseek_v4/test_hc_glue_triton.py | 267 ++ .../deepseek_v4/test_indexer_tail_triton.py | 297 +++ .../deepseek_v4/test_rope_triton.py | 434 +++ .../deepseek_v4/test_router_post_triton.py | 267 ++ .../deepseek_v4/test_sinkhorn_triton.py | 358 +++ .../test_v4_backend_import_gating.py | 122 + .../deepseek_v4/test_v4_core_attention.py | 500 ++++ .../deepseek_v4/test_v4_fp8_indexer.py | 140 + .../test_v4_gluon_dsa_attention.py | 300 +++ .../deepseek_v4/test_v4_gluon_v2_attention.py | 277 ++ .../deepseek_v4/test_v4_gluon_v3_attention.py | 322 +++ .../transformer/deepseek_v4/test_v4_moe.py | 504 ++++ .../transformer/deepseek_v4/test_v4_mtp.py | 361 +++ .../deepseek_v4/test_v4_routers.py | 381 +++ .../test_v4_turbo_deepep_dispatcher.py | 397 +++ .../test_v4_turbo_flydsl_attention.py | 293 +++ .../deepseek_v4/test_v4_v4_attention_bwd.py | 439 ++++ .../deepseek_v4/test_v4_v4_attention_fwd.py | 409 +++ .../test_v4_v4_csa_in_kernel_gather.py | 280 ++ .../deepseek_v4/v4_attention_shapes.py | 260 ++ .../deepseek_v4/v4_attention_test_utils.py | 221 ++ tools/backend_gap_report/build_site_bundle.py | 7 + 208 files changed, 66961 insertions(+), 248 deletions(-) create mode 100644 .github/workflows/deploy-projection.yml delete mode 100644 .github/workflows/docker/Dockerfile_v25.09_ainic create mode 100644 examples/deepseek-v4/benchmark/bench_v4_attention.py create mode 100644 examples/deepseek-v4/benchmark/bench_v4_attention_results.md create mode 100644 examples/deepseek-v4/projection/README.md create mode 100644 examples/deepseek-v4/projection/design/01-overview.md create mode 100644 examples/deepseek-v4/projection/design/02-assumptions.md create mode 100644 examples/deepseek-v4/projection/design/03-json-schema.md create mode 100644 examples/deepseek-v4/projection/design/04-projection-math.md create mode 100644 examples/deepseek-v4/projection/design/05-deployment.md create mode 100644 examples/deepseek-v4/projection/design/06-calibration.md create mode 100644 examples/deepseek-v4/projection/design/07-iteration-timeline.md create mode 100755 examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh create mode 100644 examples/deepseek-v4/projection/site/assets/app.js create mode 100644 examples/deepseek-v4/projection/site/assets/style.css create mode 100644 examples/deepseek-v4/projection/site/data/flash.json create mode 100644 examples/deepseek-v4/projection/site/data/pro.json create mode 100644 examples/deepseek-v4/projection/site/index.html create mode 100755 examples/deepseek-v4/projection/tools/gen_mock_data.py create mode 100644 examples/deepseek-v4/projection/tools/kernel_module_map.py create mode 100755 examples/deepseek-v4/projection/tools/parse_trace.py create mode 100644 examples/deepseek-v4/projection/tools/v4_flops.py create mode 100755 examples/deepseek-v4/run_deepseek_v4.sh create mode 100644 examples/deepseek-v4/run_deepseek_v4_flash.sh create mode 100755 examples/deepseek-v4/run_deepseek_v4_flash_proxy.sh create mode 100755 examples/deepseek-v4/run_deepseek_v4_pro_muon.sh create mode 100755 examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh create mode 100755 examples/deepseek-v4/run_dsv4_projection_1gpu.sh create mode 100644 examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml create mode 100644 examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml create mode 100644 primus/backends/megatron/core/extensions/_triton/__init__.py create mode 100644 primus/backends/megatron/core/extensions/_triton/multi_tensor_add.py create mode 100644 primus/backends/megatron/core/extensions/_triton/stack_grouped_weight.py create mode 100644 primus/backends/megatron/core/fusions/fused_bias_swiglu.py create mode 100644 primus/backends/megatron/core/fusions/fused_pad_routing_map.py create mode 100644 primus/backends/megatron/core/models/deepseek_v4/__init__.py create mode 100644 primus/backends/megatron/core/models/deepseek_v4/build_context.py create mode 100644 primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py create mode 100644 primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_builders.py create mode 100644 primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py create mode 100644 primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py create mode 100644 primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_layer.py create mode 100644 primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_specs.py create mode 100644 primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py create mode 100644 primus/backends/megatron/core/transformer/clamped_swiglu.py create mode 100644 primus/backends/megatron/core/transformer/compressor.py create mode 100644 primus/backends/megatron/core/transformer/deepseek_v4_attention.py create mode 100644 primus/backends/megatron/core/transformer/dual_rope.py create mode 100644 primus/backends/megatron/core/transformer/hyper_connection.py create mode 100644 primus/backends/megatron/core/transformer/indexer.py create mode 100644 primus/backends/megatron/core/transformer/local_rmsnorm.py create mode 100644 primus/backends/megatron/core/transformer/moe/_triton/__init__.py create mode 100644 primus/backends/megatron/core/transformer/moe/_triton/v4_router_post.py create mode 100644 primus/backends/megatron/core/transformer/moe/shared_experts.py create mode 100644 primus/backends/megatron/core/transformer/moe/v4_hash_router.py create mode 100644 primus/backends/megatron/core/transformer/moe/v4_moe.py create mode 100644 primus/backends/megatron/core/transformer/moe/v4_topk_router.py create mode 100644 primus/backends/megatron/core/transformer/sliding_window_kv.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/README.md create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/reference.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_bwd_flydsl_mqa.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_csa.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_mqa.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_attention_bwd_flydsl_mqa.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_dq_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_full_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_fwd_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dkv_pool_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dq_pool_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dkv_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dq_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_fwd_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_dq_flydsl_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_v4_flydsl.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl_kernel.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_gather.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_preprocess.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dkv_interm.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dq.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_v4_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_fwd_v4_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dkv_interm_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dq_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_v4_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_fwd_v4_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_lse_fwd.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_mla_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dkv_interm_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dq_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_v4_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_fwd_v4_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_autograd_tilelang.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_bwd_tilelang.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_fwd_tilelang.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/compressor_pool.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_collapse.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_expand.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_glue.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rmsnorm.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/sinkhorn.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_bwd.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_fwd.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/_v4_attn_tuning.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_bwd.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_fwd.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_bwd.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_fwd.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/_amd_knobs.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_kernels.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_fwd_v4_triton.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/__init__.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_flydsl.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v2.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v3.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_triton.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_turbo_flydsl.py create mode 100644 primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py create mode 100644 primus/backends/megatron/patches/deepseek_v4_flops_patches.py create mode 100644 primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py create mode 100644 primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py create mode 100644 primus/backends/megatron/patches/emerging_optimizers_log_level_patches.py create mode 100644 primus/backends/megatron/patches/fused_pad_routing_map_patches.py create mode 100644 primus/backends/megatron/patches/moe_alltoall_dtoh_patches.py create mode 100644 primus/backends/megatron/patches/triton_autotune_print_patches.py create mode 100644 primus/configs/models/megatron/deepseek_v4_base.yaml create mode 100644 primus/configs/models/megatron/deepseek_v4_flash.yaml create mode 100644 primus/configs/models/megatron/deepseek_v4_pro.yaml create mode 100644 rccl_avg_workaround/.gitignore create mode 100644 rccl_avg_workaround/sitecustomize.py create mode 100644 runner/helpers/hooks/train/pretrain/megatron/01_install_emerging_optimizers.sh create mode 100644 tests/unit_tests/backends/megatron/test_compressor_pool.py create mode 100644 tests/unit_tests/backends/megatron/test_deepseek_v4_flops_patches.py create mode 100644 tests/unit_tests/backends/megatron/test_rope_arange_cache.py create mode 100644 tests/unit_tests/configs/test_deepseek_v4_yaml.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/__init__.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/conftest.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_clamped_swiglu.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_deepseek_v4_attention.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_hc_collapse_triton.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rmsnorm_triton.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rope_from_positions.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_hc_glue_triton.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_indexer_tail_triton.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_rope_triton.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_router_post_triton.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_sinkhorn_triton.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_backend_import_gating.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_core_attention.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_fp8_indexer.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_dsa_attention.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v2_attention.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v3_attention.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_moe.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_routers.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_deepep_dispatcher.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_bwd.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_fwd.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_csa_in_kernel_gather.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_shapes.py create mode 100644 tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_test_utils.py diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index bda09abab..f65542997 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -23,7 +23,7 @@ permissions: contents: read env: - PRIMUS_TURBO_COMMIT: a04a233cbfb468dbe21600cbf9db70953428b25c # feat: force use nt layout gemm in bwd (#386) + PRIMUS_TURBO_COMMIT: 56c789e58f72aaf733b7715f1536be1ed33b69a1 # dev/kyle/flydsl_attn_deepseekv4 (dsv4 sparse-MLA attention + cr=4 fixes + flydsl kernels packaging __init__.py fix) PRIMUS_TURBO_AITER_COMMIT: 0f3c58e6edb6754940bcf9fd5f09ccb6f389f52e # AITER v0.1.14.post1 (tag commit) — required by Primus-Turbo main aiter_utils.py ROCSHMEM_COMMIT: 17ff985c026f9f97f85068647e863ab541dd5645 # Update version to 3.2.0 for 7.2.0 rocm release (#351) (#355) UCCL_COMMIT: 5afb4117893c58cc0c8557d9286336141a301053 # [EP]: fix fp8 error of internode_ll on amd gfx950 arch. (#710) @@ -154,25 +154,6 @@ jobs: docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-ainic # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }} - echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-v25.09-ainic" - start_time=$(date +%s) - mkdir -p $GITHUB_WORKSPACE/.github/workflows/docker/ainic - cp /apps/tas/0_public/primus_docker_ci/ainic/ainic_bundle_1.117.5-a-56.tar.gz $GITHUB_WORKSPACE/.github/workflows/docker/ainic/ || { echo "Error: Failed to copy ainic bundle"; exit 1; } - docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile_v25.09_ainic \ - --network=host \ - -t tasimage/primus:${{env.IMAGE_TAG}}-v25.09-ainic \ - --build-arg AINIC_BUNDLE_PATH=ainic \ - --build-arg PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT} \ - $GITHUB_WORKSPACE/.github/workflows/docker - end_time=$(date +%s) - elapsed=$((end_time - start_time)) - echo "⏱️ [build primus docker-v25.09-ainic] Total elapsed time: ${elapsed} seconds" - - docker tag tasimage/primus:${{env.IMAGE_TAG}}-v25.09-ainic docker.io/tasimage/primus:${{env.IMAGE_TAG}}-v25.09-ainic - docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }} - docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}-v25.09-ainic - # docker login -u rocmshared -p ${{ secrets.ROCM_DOCKER_HUB_TOKEN }} - echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}-jax" start_time=$(date +%s) docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile \ @@ -216,8 +197,6 @@ jobs: # echo "> Docker cleanup local images" # docker rmi tasimage/primus:${{env.IMAGE_TAG}} - # docker rmi tasimage/primus:${{env.IMAGE_TAG}}-v25.09-ainic - # docker rmi tasimage/primus:${{env.IMAGE_TAG}}-v25.10-ainic # docker rmi tasimage/primus:${{env.IMAGE_TAG}}-jax echo "> build-docker success" @@ -249,6 +228,7 @@ jobs: - run: echo "Begin AITER + Primus-Turbo Install." - name: Install AITER run: | + : > "$RUNNER_TEMP/runtime.tsv" # reset the CI runtime log for this job echo "✅ [Uninstall old aiter] started at: $(date)" pip3 uninstall aiter amd-aiter -y || true rm -rf /tmp/aiter || true @@ -265,6 +245,7 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [Build aiter] ended at: $(date)" echo "⏱️ [Build aiter] Total elapsed time: ${elapsed} seconds" + echo -e "Build aiter\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" - name: Install Primus-Turbo run: | rm -rf /tmp/Primus-Turbo || true @@ -280,6 +261,7 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [Pip install requirements] ended at: $(date)" echo "⏱️ [Pip install requirements] Total elapsed time: ${elapsed} seconds" + echo -e "primus-turbo: pip install requirements\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" start_time=$(date +%s) echo "✅ [build primus-turbo] started at: $(date)" pip3 install --no-build-isolation -e . -v @@ -287,6 +269,7 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [build primus-turbo] ended at: $(date)" echo "⏱️ [build primus-turbo] Total elapsed time: ${elapsed} seconds" + echo -e "primus-turbo: build/install\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" - run: echo "🎉 Begin Primus Unit Test." - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 with: @@ -340,6 +323,8 @@ jobs: run: | echo "Running Primus Core tests..." # Note: The tests `test_fp8_te_linear` and `test_te_linear` are temporarily skipped due to intermittent failures. + # Note: `test_limit_layers_moe_with_dense_keeps_two` is temporarily disabled pending a fix to + # `_limit_layers_for_projection` dense-layer handling (num_layers collapses to 1 instead of 2). # Note HSA_NO_SCRATCH_RECLAIM=1 must be set to avoid RCCL perf hit (TAS-8N Node), rocm ver:70125424 export HSA_NO_SCRATCH_RECLAIM=1 mkdir -p "${GITHUB_WORKSPACE}/test-reports" @@ -349,7 +334,8 @@ jobs: --deselect=tests/unit_tests/megatron/cco/test_tp_overlap.py::TPOverlapTestCase::test_fp8_te_linear \ --deselect=tests/unit_tests/megatron/cco/test_tp_overlap.py::TPOverlapTestCase::test_te_linear \ --deselect=tests/unit_tests/megatron/transformer/moe/test_token_dispatcher.py::TestFlexDispatcher::test_forward_backward \ - --deselect=tests/unit_tests/megatron/transformer/moe/test_token_dispatcher.py::TestFlexDispatcher::test_capacity_forward_backward + --deselect=tests/unit_tests/megatron/transformer/moe/test_token_dispatcher.py::TestFlexDispatcher::test_capacity_forward_backward \ + --deselect=tests/unit_tests/core/projection/test_performance_projection.py::test_limit_layers_moe_with_dense_keeps_two - name: Snapshot unit-test coverage if: always() continue-on-error: true @@ -524,6 +510,7 @@ jobs: echo "Primus-Turbo dir: /tmp/Primus-Turbo" git config --global --add safe.directory /tmp/Primus-Turbo cd /tmp/Primus-Turbo + : > "$RUNNER_TEMP/runtime.tsv" # reset the CI runtime log for this job start_time=$(date +%s) echo "✅ [Pip install requirements] started at: $(date)" mkdir -p ${PRIMUS_WORKDIR}/primus-cache @@ -532,6 +519,7 @@ jobs: elapsed=$((end_time - start_time)) echo "✅ [Pip install requirements] ended at: $(date)" echo "⏱️ [Pip install requirements] Total elapsed time: ${elapsed} seconds" + echo -e "pip install/upgrade\t${elapsed}" >> "$RUNNER_TEMP/runtime.tsv" start_time=$(date +%s) echo "✅ [build primus-turbo] started at: $(date)" end_time=$(date +%s) diff --git a/.github/workflows/deploy-projection.yml b/.github/workflows/deploy-projection.yml new file mode 100644 index 000000000..e9fd62bdc --- /dev/null +++ b/.github/workflows/deploy-projection.yml @@ -0,0 +1,66 @@ +name: Deploy DeepSeek-V4 Projection + +# Publishes the DeepSeek-V4 performance-projection site to GitHub Pages from the +# current development branch (no need to merge to main). The site is served at +# https://.github.io//deepseek-v4-projection/ +# NOTE: a repo has a single Pages site, so this deployment and the main +# backend-gap dashboard deployment overwrite each other (last run wins). If the +# github-pages environment restricts deployment branches, allow this branch in +# repo Settings -> Environments -> github-pages -> Deployment branches. + +on: + workflow_dispatch: + push: + branches: + - dev/tas/deepseek-v4 + paths: + - "examples/deepseek-v4/projection/site/**" + - ".github/workflows/deploy-projection.yml" + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: deepseek-v4-projection-pages + cancel-in-progress: true + +jobs: + build: + runs-on: ubuntu-latest + steps: + - name: Check out repository + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + + - name: Assemble Pages bundle + run: | + set -euo pipefail + mkdir -p _site/deepseek-v4-projection + cp -r examples/deepseek-v4/projection/site/. _site/deepseek-v4-projection/ + cat > _site/index.html <<'HTML' + + + + DeepSeek-V4 Projection + DeepSeek-V4 Performance Projection + HTML + + - name: Configure Pages + uses: actions/configure-pages@45bfe0192ca1faeb007ade9deae92b16b8254a0d # v6.0.0 + + - name: Upload Pages artifact + uses: actions/upload-pages-artifact@fc324d3547104276b827a68afc52ff2a11cc49c9 # v5.0.0 + with: + path: _site + + deploy: + needs: build + runs-on: ubuntu-latest + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@cd2ce8fcbc39b97be8ca5fce6e763baed58fa128 # v5.0.0 diff --git a/.github/workflows/docker/Dockerfile b/.github/workflows/docker/Dockerfile index 84f7c5f39..d25e6c77b 100644 --- a/.github/workflows/docker/Dockerfile +++ b/.github/workflows/docker/Dockerfile @@ -77,11 +77,11 @@ RUN rm -rf /opt/Primus-Turbo # Install Triton # --------------------------------------------------------------------------- RUN cd /opt && \ - git clone -b release/3.7.x https://github.com/triton-lang/triton.git && \ + git clone https://github.com/triton-lang/triton.git && \ cd triton && \ - git checkout ${TRITON_COMMIT} && \ - pip3 install ninja cmake && \ - pip3 install --no-build-isolation -v . + git checkout 09500db9 && \ + pip3 install -r python/requirements.txt && \ + MAX_JOBS=96 pip3 install --no-build-isolation --force-reinstall --no-deps . RUN rm -rf /opt/triton diff --git a/.github/workflows/docker/Dockerfile_v25.09_ainic b/.github/workflows/docker/Dockerfile_v25.09_ainic deleted file mode 100644 index d8c5934a7..000000000 --- a/.github/workflows/docker/Dockerfile_v25.09_ainic +++ /dev/null @@ -1,106 +0,0 @@ -# Base image -FROM docker.io/rocm/megatron-lm:v25.9_gfx950 - -# Specify the commit of Primus-Turbo when building: docker build --build-arg PRIMUS_TURBO_COMMIT=xxx .) -ARG PRIMUS_TURBO_COMMIT -ARG AINIC_BUNDLE_PATH - -# Install basic dependencies -RUN apt-get update - -# Clone and install the Primus-Turbo -WORKDIR /opt -RUN mkdir -p /opt && cd /opt && \ - git clone https://github.com/AMD-AGI/Primus-Turbo.git && \ - cd Primus-Turbo && \ - git checkout ${PRIMUS_TURBO_COMMIT} && \ - git submodule update --init --recursive && \ - pip3 install -r requirements.txt && \ - GPU_ARCHS="gfx942;gfx950" pip3 install --no-build-isolation . - -RUN apt-get install --reinstall binutils -y && apt-get install numactl -y - -WORKDIR /opt -ENV WORKDIR=/opt -ENV ROCM_PATH=/opt/rocm - -RUN apt-get update && \ - apt-get install jq dpkg-dev kmod xz-utils \ - libfmt-dev libboost-all-dev \ - libibverbs-dev ibverbs-utils infiniband-diags -y - -# =============================== Build AINIC Driver =============================== -# WARNING: Please ensure the following environment variables are correctly set: -# WARNING: 1. PATH: /usr/sbin must be included. -# WARNING: 2. LD_LIBRARY_PATH: /usr/lib must be included. -# WARNING: If these paths are missing, tools and libraries may not function correctly. -# INFO: Installation completed successfully - -COPY ${AINIC_BUNDLE_PATH}/ainic_bundle_1.117.5-a-56.tar.gz ${WORKDIR} -RUN cd ${WORKDIR} && \ - echo "Building ainic bundle... current directory: ${WORKDIR}" && \ - tar zxf ainic_bundle_1.117.5-a-56.tar.gz && \ - cd ainic_bundle_1.117.5-a-56 && \ - tar zxf host_sw_pkg.tar.gz && \ - cd host_sw_pkg && \ - ./install.sh --domain=user -y 2>&1 | tee log_install.txt && \ - cd ${WORKDIR} && \ - apt-get install -y ./amd/ainic/deb-repo/libionic*.deb - -# =============================== Test AINIC Driver =============================== -# ibv_devices -# rdma link -# ethtool -i enp9s0 -# ibv_devinfo -vv | grep GID - -# =============================== Build UCX =============================== -RUN cd ${WORKDIR} && wget https://github.com/openucx/ucx/releases/download/v1.18.0/ucx-1.18.0.tar.gz && \ - mkdir -p ucx-1.18.0 && \ - tar -zxf ucx-1.18.0.tar.gz -C ucx-1.18.0 --strip-components=1 && \ - cd ucx-1.18.0 && mkdir build && cd build && \ - ../configure --prefix=${WORKDIR}/ucx-1.18.0/install --with-rocm=${ROCM_PATH} 2>&1 | tee log_ucx_configure.txt && \ - make -j 16 2>&1 | tee log_ucx_build.txt && \ - make install && \ - cd ${WORKDIR} - -ENV UCX_INSTALL_DIR=${WORKDIR}/ucx-1.18.0/install - -# =============================== Build MPI =============================== -RUN cd ${WORKDIR} && \ - wget https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-4.1.6.tar.gz && \ - mkdir -p ompi-4.1.6 && \ - tar -zxf openmpi-4.1.6.tar.gz -C ompi-4.1.6 --strip-components=1 && \ - cd ompi-4.1.6 && mkdir build && cd build && \ - ../configure --prefix=${WORKDIR}/ompi-4.1.6/install --with-ucx=${UCX_INSTALL_DIR} \ - --disable-oshmem --disable-mpi-fortran 2>&1 | tee log_mpi_configure.txt && \ - make -j 16 2>&1 | tee log_mpi_build.txt && \ - make install && \ - cd ${WORKDIR} - -ENV MPI_PATH=${WORKDIR}/ompi-4.1.6/install - -# =============================== Build RCCL =============================== -RUN cd ${WORKDIR} && \ - git clone https://github.com/ROCm/rccl.git && \ - cd rccl && git checkout drop/2025-08 && \ - ./install.sh -l --prefix build/ --disable-mscclpp \ - --disable-msccl-kernel --amdgpu_targets="gfx950" 2>&1 | tee log_rccl_install.txt && \ - cd ${WORKDIR} - -ENV RCCL_HOME=${WORKDIR}/rccl - -# =============================== Build AMD ANP =============================== - -RUN cd ${WORKDIR} && git clone https://github.com/rocm/amd-anp.git && \ -cd amd-anp && git checkout tags/v1.1.0-5 && \ -sed -i '5a CFLAGS += --offload-arch=gfx950' ./Makefile && head -10 ./Makefile && \ -make -j 16 RCCL_BUILD=${RCCL_HOME}/build/release \ - MPI_INCLUDE=${MPI_PATH}/include/ \ - MPI_LIB_PATH=${MPI_PATH}/lib/ \ - ROCM_PATH=${ROCM_PATH} 2>&1 | tee log_amd_anp_build.txt - -# Set the default working directory -WORKDIR /opt - -# check the installed Primus-Turbo package -RUN python3 -m pip show primus-turbo || true diff --git a/.gitignore b/.gitignore index a6178a7c2..c64ac79e6 100644 --- a/.gitignore +++ b/.gitignore @@ -18,4 +18,14 @@ experiment primus/backends/diffusion/data/**/__pycache__/ primus/backends/diffusion/data/**/*.pyc pp_simulation_result + +# Allow the projection site's breakdown JSON (consumed by the static site / +# GitHub Pages) despite the generic `data` ignore above. +!examples/deepseek-v4/projection/site/data/ +!examples/deepseek-v4/projection/site/data/*.json + +*.log +*.nohup +*.zip +.triton_cache_shared/ .cursor/ diff --git a/examples/deepseek-v4/benchmark/bench_v4_attention.py b/examples/deepseek-v4/benchmark/bench_v4_attention.py new file mode 100644 index 000000000..45cd76cde --- /dev/null +++ b/examples/deepseek-v4/benchmark/bench_v4_attention.py @@ -0,0 +1,592 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 attention kernel benchmark — all backends in one table. + +Unified benchmark of the **forward** and **backward** V4 attention kernels for +every backend, so there is a single place to compare them (one row per backend, +``ms | TFLOP/s`` cells; TFLOP/s over the useful work so all rows are comparable): + +* ``triton`` — PRODUCTION Triton (separate K/V; pool/SWA/HCA launchers) +* ``gluon`` — fused single-latent (K==V) sparse-MLA, hand-tuned gfx950 +* ``triton_v2`` — fused single-latent sparse-MLA in plain Triton (tl.dot / MFMA) +* ``flydsl_v1`` — fused single-latent sparse-MLA in native FlyDSL MFMA (fwd + bwd) +* ``turbo_flydsl`` — extracted Primus-Turbo sparse_mla_v2 native FlyDSL MFMA +* ``_turbo_flydsl`` — the INTEGRATED in-tree Primus-Turbo backend via the turbo API + (``primus_turbo.flydsl.attention``); the ``turbo`` model backend + (``use_v4_attention_backend`` / ``use_v4_csa_attention_backend = "turbo"``) + +The legacy ``_flydsl_v0_deprecated`` gathered-CSA backend (scalarized GEMV) is +NOT benchmarked: it has known correctness issues and depends on the +/workspace/FlyDSL-amd source tree. + +The fused ``gluon`` / ``triton_v2`` / ``flydsl_v1`` backends share ONE kernel-pair +API and are timed on IDENTICAL V4-form inputs (zero rope pad + [local ++ pool] +kv + [SWA window ++ pool] topk). Each backend is guarded; unavailable ones are +simply skipped. + +Benchmarks for the two production model sizes across all three layer kinds: + +* ``cr=0`` — dense / sliding-window (SWA-only) attention +* ``cr=4`` — CSA (local SWA + sparse top-k from the compressed pool) +* ``cr=128`` — HCA (local SWA + full compressed pool, joint softmax) + +at the real attention shapes (``seq_len=4096``, ``mbs=1``, bf16, sink on, +``swa_window=128``): + +* V4-Flash: H=64, head_dim=512, index_topk=512 +* V4-Pro: H=128, head_dim=512, index_topk=1024 + +Backend detail: + +* ``triton`` (ALL crs): PRODUCTION launchers used by ``DeepseekV4Attention`` — + cr=0/128 ``_launch_v4_attention_fwd``/``_bwd`` (dense/HCA); cr=4 + ``_launch_v4_csa_attention_pool_fwd``/``_pool_bwd`` (split FWD + segreduce + BWD in-kernel gather; NOT the legacy gathered API, ~30-260x slower). +* ``gluon`` / ``triton_v2`` / ``flydsl_v1`` (ALL crs): fused single-latent + sparse-MLA (``sparse_mla_{fwd,bwd}_v4_*``); the layer kind is just a + different TOPK (cr=0: swa 128; cr=4: 128+sparse; cr=128: 128+pool). + +Effective TFLOP/s uses ``2*T*H*TOPK*(D_V+D_V)`` (useful work over head_dim=512), +BWD = 2.5x FWD, the SAME formula for every backend so rows are directly +comparable (the fused backends' zero-rope-pad overhead shows up in ms). + +Run inside the dev container (gfx950 / MI355X): + + python examples/deepseek-v4/benchmark/bench_v4_attention.py + python examples/deepseek-v4/benchmark/bench_v4_attention.py --variant pro --cr 4 +""" + +from __future__ import annotations + +import argparse +import math +import os +import sys +from typing import Tuple + +import torch + +# NOTE: the legacy `_flydsl_v0_deprecated` CSA backend (scalarized GEMV) is +# intentionally NOT imported/benchmarked here — it has known correctness issues +# and depends on the /workspace/FlyDSL-amd source tree. The native FlyDSL MFMA +# backend is `_flydsl_v1` (benchmarked below as `flydsl_v1`). +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import ( + _launch_v4_attention_bwd, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import ( + _launch_v4_attention_fwd, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_csa_attention_bwd import ( + _launch_v4_csa_attention_pool_bwd, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_csa_attention_fwd import ( + _launch_v4_csa_attention_pool_fwd, +) + +# Fused single-latent (K==V) sparse-MLA backends. All share ONE kernel-pair +# API (fwd(q, kv, topk, attn_sink, kv_lora_rank, scale) -> (o, lse); bwd -> +# (dq, dkv, d_sink)) so they are timed on identical V4-form inputs. Each is +# guarded so the benchmark still runs where a backend is unavailable. +_SPARSE_MLA_BACKENDS = {} # name -> (fwd_fn, bwd_fn) +try: + from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_dsa import ( + sparse_mla_bwd_v4_gluon, + sparse_mla_fwd_v4_gluon, + ) + + _SPARSE_MLA_BACKENDS["gluon"] = (sparse_mla_fwd_v4_gluon, sparse_mla_bwd_v4_gluon) +except Exception: # noqa: BLE001 + pass +try: + from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v2 import ( + sparse_mla_bwd_v4_triton, + sparse_mla_fwd_v4_triton, + ) + + _SPARSE_MLA_BACKENDS["triton_v2"] = (sparse_mla_fwd_v4_triton, sparse_mla_bwd_v4_triton) +except Exception: # noqa: BLE001 + pass +try: + from primus.backends.megatron.core.transformer.v4_attention_kernels._flydsl_v1 import ( + sparse_mla_bwd_v4_flydsl, + sparse_mla_fwd_v4_flydsl, + ) + + _SPARSE_MLA_BACKENDS["flydsl_v1"] = (sparse_mla_fwd_v4_flydsl, sparse_mla_bwd_v4_flydsl) +except Exception: # noqa: BLE001 + pass +# _turbo_flydsl: the INTEGRATED Primus-Turbo sparse-MLA backend, via the turbo API +# (primus_turbo.flydsl.attention). Not an agent/workspace extraction — this is the +# in-tree `_turbo_flydsl` backend (enabled in the model via +# use_v4_attention_backend / use_v4_csa_attention_backend = "turbo"). Requires an +# installed primus_turbo carrying primus_turbo.flydsl.attention. +try: + from primus.backends.megatron.core.transformer.v4_attention_kernels._turbo_flydsl import ( + sparse_mla_bwd_v4_turbo_flydsl, + sparse_mla_fwd_v4_turbo_flydsl, + ) + + _SPARSE_MLA_BACKENDS["_turbo_flydsl"] = ( + sparse_mla_fwd_v4_turbo_flydsl, + sparse_mla_bwd_v4_turbo_flydsl, + ) +except Exception: # noqa: BLE001 + pass +# gluon_v2: our aiter-gluon-inspired backend (guarded; skipped until present). +try: + from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v2 import ( + sparse_mla_bwd_v4_gluon_v2, + sparse_mla_fwd_v4_gluon_v2, + ) + + _SPARSE_MLA_BACKENDS["gluon_v2"] = (sparse_mla_fwd_v4_gluon_v2, sparse_mla_bwd_v4_gluon_v2) +except Exception: # noqa: BLE001 + pass +# gluon_v3: active gfx950 optimization campaign backend. +try: + from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v3 import ( + sparse_mla_bwd_v4_gluon_v3, + sparse_mla_fwd_v4_gluon_v3, + ) + + _SPARSE_MLA_BACKENDS["gluon_v3"] = (sparse_mla_fwd_v4_gluon_v3, sparse_mla_bwd_v4_gluon_v3) +except Exception: # noqa: BLE001 + pass +# aiter PR#3833 gluon sparse-MLA prefill (fwd-only) — OPTIONAL reference for +# comparison. Lives under agent/workspace; set PRIMUS_V4_AITER_DIR to override. +# Guarded so the benchmark is unchanged when the extracted adapter is absent. +try: + _aiter_dir = os.environ.get( + "PRIMUS_V4_AITER_DIR", + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "agent", + "workspace", + "aiter_dsv4_prefill_20260706", + ), + ) + _aiter_dir = os.path.abspath(_aiter_dir) + if _aiter_dir not in sys.path: + sys.path.insert(0, _aiter_dir) + from aiter_v4_adapter import sparse_mla_fwd_v4_aiter + + _SPARSE_MLA_BACKENDS["aiter_gluon"] = (sparse_mla_fwd_v4_aiter, None) +except Exception: # noqa: BLE001 + pass + +_GLUON_AVAIL = "gluon" in _SPARSE_MLA_BACKENDS + +_VARIANTS = { + "flash": dict(H=64, index_topk=512), + "pro": dict(H=128, index_topk=1024), +} +_HEAD_DIM = 512 +_SWA_WINDOW = 128 +_ROPE_DIM = 64 # sparse-MLA d_qk = kv_lora_rank (=_HEAD_DIM) + rope +_CR_NAME = {0: "SWA/dense", 4: "CSA", 128: "HCA"} + + +# --------------------------------------------------------------------------- +# Input builders +# --------------------------------------------------------------------------- + + +def _common_inputs(B: int, H: int, S: int, D: int, seed: int = 0): + """q + MQA single-latent K/V (full + [.,1,.,.]) + sink + dout, all bf16.""" + g = torch.Generator(device="cuda").manual_seed(seed) + dev, dt = "cuda", torch.bfloat16 + q = torch.randn(B, H, S, D, generator=g, device=dev, dtype=dt) + k_mqa = torch.randn(B, 1, S, D, generator=g, device=dev, dtype=dt) + v_mqa = torch.randn(B, 1, S, D, generator=g, device=dev, dtype=dt) + k_full = k_mqa.expand(B, H, S, D).contiguous() + v_full = v_mqa.expand(B, H, S, D).contiguous() + sink = torch.randn(H, generator=g, device=dev, dtype=torch.float32) * 0.1 + dout = torch.randn(B, H, S, D, generator=g, device=dev, dtype=dt) + return q, k_mqa, v_mqa, k_full, v_full, sink, dout + + +def _csa_sparse(B: int, H: int, S: int, D: int, K: int, P: int, g: torch.Generator): + """CSA sparse branch: compressed pool + per-query top-K indices, plus the + pre-gathered equivalent (FlyDSL) sharing the same valid/invalid pattern.""" + dev, dt = "cuda", torch.bfloat16 + pool = torch.randn(B, P, D, generator=g, device=dev, dtype=dt) + valid = torch.rand(B, S, K, generator=g, device=dev) > 0.25 + topk_idxs = torch.randint(0, P, (B, S, K), generator=g, device=dev, dtype=torch.int32) + topk_idxs = torch.where(valid, topk_idxs, torch.full_like(topk_idxs, -1)) + safe = topk_idxs.clamp(min=0).long() + gathered = torch.gather( + pool.unsqueeze(1).expand(B, S, P, D), dim=2, index=safe.unsqueeze(-1).expand(B, S, K, D) + ) + gathered = gathered * valid.unsqueeze(-1).to(dt) + sparse_mask = torch.where( + valid, + torch.zeros((), dtype=dt, device=dev), + torch.tensor(float("-inf"), dtype=dt, device=dev), + ) + return pool, topk_idxs, gathered, sparse_mask + + +def _hca_mask(S: int, P: int, ratio: int, device, dtype): + """HCA pool-only additive causal mask [S, P]: pool slot s visible to query t + iff (s+1)*ratio - 1 <= t (matches DeepseekV4Attention._hca_extra_mask).""" + t = torch.arange(S, device=device).unsqueeze(1) + s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * ratio - 1 + return torch.where(s_end <= t, 0.0, float("-inf")).to(dtype) + + +def _build_gluon_v4form(*, cr: int, H: int, S: int, D: int, K: int, P: int, W: int, seed: int = 0): + """Gluon inputs in the **V4 form** (matches the training adapter): + + The 512 V4 latent (RoPE baked in-place) is the gluon "lora" with a zero rope + pad (kernel needs D_ROPE>0); kv = [local latent ++ pool] and topk = [SWA + window ++ (cr=4: sparse pool top-k | cr=128: full causal pool | cr=0: none)]. + ``topk`` is padded to a multiple of 64 so the gluon bwd dKV tiling (TILE_K=64) + is valid (notably HCA: 128+32=160 -> 192). scale = 1/sqrt(D) (V4, over 512). + """ + g = torch.Generator(device="cuda").manual_seed(seed) + dev, dt = "cuda", torch.bfloat16 + latent = torch.randn(S, D, generator=g, device=dev, dtype=dt) + q512 = torch.randn(S, H, D, generator=g, device=dev, dtype=dt) + z_q = torch.zeros(S, H, _ROPE_DIM, device=dev, dtype=dt) + q_g = torch.cat([q512, z_q], dim=-1).contiguous() # [S, H, D+rope_pad] + sink = torch.randn(H, generator=g, device=dev, dtype=torch.float32) * 0.1 + do = torch.randn(S, H, D, generator=g, device=dev, dtype=dt) + + ti = torch.arange(S, device=dev).view(S, 1) + win = ti - W + 1 + torch.arange(W, device=dev).view(1, W) # [S, W] local token idx + win = torch.where(win >= 0, win, torch.full_like(win, -1)) + + if cr == 0: + kv512 = latent.unsqueeze(1) # [S, 1, D] + topk = win + else: + pool = torch.randn(P, D, generator=g, device=dev, dtype=dt) + kv512 = torch.cat([latent, pool], dim=0).unsqueeze(1) # [S+P, 1, D] + if cr == 4: + sp = torch.randint(0, P, (S, K), generator=g, device=dev) + pool_topk = S + sp + else: # cr == 128: HCA full causal pool + ps = torch.arange(P, device=dev).view(1, P) + vis = ((ps + 1) * cr - 1) <= ti # [S, P] + pool_topk = torch.where(vis, S + ps, torch.full_like(ps.expand(S, P), -1)) + topk = torch.cat([win, pool_topk], dim=1) + + tk = topk.shape[1] + pad = ((tk + 63) // 64) * 64 - tk + if pad > 0: + topk = torch.cat([topk, torch.full((S, pad), -1, device=dev, dtype=topk.dtype)], dim=1) + topk_g = topk.to(torch.int32).contiguous() + + z_kv = torch.zeros(kv512.shape[0], 1, _ROPE_DIM, device=dev, dtype=dt) + kv_g = torch.cat([kv512, z_kv], dim=-1).contiguous() + return q_g, kv_g, topk_g, sink, do + + +# --------------------------------------------------------------------------- +# Timing helpers +# --------------------------------------------------------------------------- + + +def _time_ms(fn, *, warmup: int, iters: int) -> Tuple[float, float]: + """Return (median_ms, mean_ms) over ``iters`` timed launches.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + times = [] + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for _ in range(iters): + start.record() + fn() + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + times.sort() + return times[len(times) // 2], sum(times) / len(times) + + +def _safe_time(fn, *, warmup: int, iters: int): + """Time ``fn``; on failure return (None, short_error_string).""" + if fn is None: + return None, None + try: + med, _ = _time_ms(fn, warmup=warmup, iters=iters) + return med, None + except Exception as exc: # noqa: BLE001 - benchmark must survive one cell failing + msg = str(exc).strip().splitlines()[-1] if str(exc).strip() else type(exc).__name__ + return None, f"{type(exc).__name__}: {msg}" + + +def _call_fwd(fn): + """Call a fwd launcher once for the bwd's (out, lse); (None, None) on failure.""" + if fn is None: + return None, None + try: + return fn() + except Exception: # noqa: BLE001 + return None, None + + +def _tflops(flops: float, med_ms) -> float: + if med_ms is None or med_ms <= 0: + return float("nan") + return flops / (med_ms * 1e-3) / 1e12 + + +def _cell(med, flops) -> str: + if med is None: + return f"{'FAIL':>18s}" + return f"{med:9.2f} | {_tflops(flops, med):6.1f}" + + +# --------------------------------------------------------------------------- +# Per-(variant, cr) benchmark +# --------------------------------------------------------------------------- + + +def _bench_cr(variant: str, cr: int, *, B: int, S: int, warmup: int, iters: int): + cfg = _VARIANTS[variant] + H = cfg["H"] + D = _HEAD_DIM + scale = 1.0 / math.sqrt(D) + g = torch.Generator(device="cuda").manual_seed(11) + q, k_mqa, v_mqa, k_full, v_full, sink, dout = _common_inputs(B, H, S, D) + + fwd_triton = bwd_triton = None + + if cr == 0: + topk_eff = _SWA_WINDOW + glu_K, glu_P = 0, 0 + extra = "" + + def fwd_triton(): + return _launch_v4_attention_fwd( + q, + k_full, + v_full, + sink=sink, + swa_window=_SWA_WINDOW, + additive_mask=None, + scale=scale, + hca_local_seqlen=0, + ) + + out_t, lse_t = _call_fwd(fwd_triton) + + def bwd_triton(): + return _launch_v4_attention_bwd( + q, + k_full, + v_full, + out_t, + dout, + lse_t, + sink=sink, + swa_window=_SWA_WINDOW, + additive_mask=None, + scale=scale, + hca_local_seqlen=0, + ) + + elif cr == 4: + P = max(S // 4, 1) + K = min(cfg["index_topk"], P) + topk_eff = _SWA_WINDOW + K + glu_K, glu_P = K, P + extra = f" K_topk={K} P={P}" + pool, topk_idxs, gathered, sparse_mask = _csa_sparse(B, H, S, D, K, P, g) + + def fwd_triton(): + return _launch_v4_csa_attention_pool_fwd( + q, + k_full, + v_full, + pool, + topk_idxs, + sink=sink, + swa_window=_SWA_WINDOW, + scale=scale, + ) + + out_t, lse_t = _call_fwd(fwd_triton) + + def bwd_triton(): + return _launch_v4_csa_attention_pool_bwd( + q, + k_full, + v_full, + pool, + topk_idxs, + out_t, + dout, + lse_t, + sink=sink, + swa_window=_SWA_WINDOW, + scale=scale, + ) + + elif cr == 128: + P = max(S // cr, 1) + topk_eff = _SWA_WINDOW + P + glu_K, glu_P = 0, P + extra = f" P={P}" + pool_bh = torch.randn(B, H, P, D, generator=g, device="cuda", dtype=torch.bfloat16) + k_hca = torch.cat([k_full, pool_bh], dim=2).contiguous() + v_hca = torch.cat([v_full, pool_bh], dim=2).contiguous() + hca_mask = _hca_mask(S, P, cr, "cuda", torch.bfloat16) + + def fwd_triton(): + return _launch_v4_attention_fwd( + q, + k_hca, + v_hca, + sink=sink, + swa_window=_SWA_WINDOW, + additive_mask=hca_mask, + scale=scale, + hca_local_seqlen=S, + ) + + out_t, lse_t = _call_fwd(fwd_triton) + + def bwd_triton(): + return _launch_v4_attention_bwd( + q, + k_hca, + v_hca, + out_t, + dout, + lse_t, + sink=sink, + swa_window=_SWA_WINDOW, + additive_mask=hca_mask, + scale=scale, + hca_local_seqlen=S, + ) + + else: + raise ValueError(f"unsupported cr={cr}") + + # Fused single-latent (K==V) sparse-MLA backends (gluon / triton_v2 / + # turbo_flydsl): all on the SAME V4-form inputs (zero rope pad + [local ++ pool] + # kv + [SWA window ++ pool] topk). Time each by swapping the kernel pair. + sparse_fb = {} # name -> (fwd_closure, bwd_closure) + if _SPARSE_MLA_BACKENDS: + gq, gkv, gtopk_idx, gsink, gdo = _build_gluon_v4form( + cr=cr, H=H, S=S, D=D, K=glu_K, P=glu_P, W=_SWA_WINDOW + ) + gscale = 1.0 / math.sqrt(D) # V4 scale (score over 512) + for _name, (_fwd_k, _bwd_k) in _SPARSE_MLA_BACKENDS.items(): + fwd_c = ( + lambda fk: (lambda: fk(gq, gkv, gtopk_idx, attn_sink=gsink, kv_lora_rank=D, scale=gscale)) + )(_fwd_k) + _out, _lse = _call_fwd(fwd_c) + if _out is None: # fwd unavailable (e.g. flydsl_v2 kernel WIP) -> skip bwd + bwd_c = None + else: + bwd_c = ( + lambda bk, o, l: ( + lambda: bk( + gq, gkv, o, gdo, gtopk_idx, l, attn_sink=gsink, kv_lora_rank=D, scale=gscale + ) + ) + )(_bwd_k, _out, _lse) + sparse_fb[_name] = (fwd_c, bwd_c) + + # Effective FLOPs over the USEFUL work (score+value both over head_dim=512, + # TOPK = real key count); BWD = 2.5x FWD. Same formula for ALL backends so + # TFLOP/s is comparable (the sparse-MLA zero-rope-pad overhead shows up in + # ms, not in counted FLOPs). + T = B * S + fwd_flop = 2.0 * T * H * topk_eff * (D + D) + flops = {"fwd": fwd_flop, "bwd": 2.5 * fwd_flop} + + # Backend table: native production Triton (separate K/V), then the fused + # sparse-MLA backends (legacy `_flydsl_v0` gathered CSA is excluded). + backends = [("triton", fwd_triton, bwd_triton)] + for _name in ( + "gluon", + "triton_v2", + "gluon_v2", + "gluon_v3", + "flydsl_v1", + "_turbo_flydsl", + "aiter_gluon", + ): + if _name in sparse_fb: + backends.append((_name, sparse_fb[_name][0], sparse_fb[_name][1])) + + print( + f"\n=== V4-{variant.upper()} cr={cr} ({_CR_NAME[cr]}) | B={B} H={H} S={S} D={D}" + f"{extra} TOPK_eff={topk_eff} swa={_SWA_WINDOW} sink=on bf16 ===\n" + f" FWD GFLOP={fwd_flop / 1e9:.1f} (useful, over head_dim={D}) " + f"(cells: ms | TFLOP/s; *_v2/gluon = V4-form fused single-latent)", + flush=True, + ) + print( + f" {'backend':12s} {'fwd (ms|TF)':>18s} {'bwd (ms|TF)':>18s}", + flush=True, + ) + rows = [] + for name, fwd_fn, bwd_fn in backends: + fwd_med, fwd_err = _safe_time(fwd_fn, warmup=warmup, iters=iters) + bwd_med, bwd_err = _safe_time(bwd_fn, warmup=warmup, iters=iters) + print(f" {name:12s} {_cell(fwd_med, flops['fwd'])} {_cell(bwd_med, flops['bwd'])}", flush=True) + for op, err in (("fwd", fwd_err), ("bwd", bwd_err)): + if err: + print(f" {name} {op} error: {err}", flush=True) + rows.append((name, fwd_med, bwd_med)) + return rows + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--variant", + choices=["flash", "pro", "both"], + default="both", + help="model size to benchmark (default: both)", + ) + parser.add_argument( + "--cr", + choices=["0", "4", "128", "all"], + default="all", + help="compress ratio / layer kind to benchmark (default: all)", + ) + parser.add_argument("--seq", type=int, default=4096, help="sequence length (default 4096)") + parser.add_argument("--mbs", type=int, default=1, help="micro batch size / B (default 1)") + parser.add_argument("--warmup", type=int, default=10, help="warmup launches (default 10)") + parser.add_argument("--iters", type=int, default=30, help="timed launches (default 30)") + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA / HIP device required for this benchmark") + + # Production-optimal Triton config (matches attention_perf.md P57 defaults): + # split CSA FWD (monolithic OFF), split + segreduce BWD. + os.environ.setdefault("PRIMUS_V4_CSA_FWD_FORCE_MONOLITHIC", "0") + os.environ.setdefault("PRIMUS_V4_ATTN_BWD_USE_SPLIT", "1") + os.environ.setdefault("PRIMUS_V4_CSA_BWD_SEGREDUCE", "1") + + torch.backends.cuda.matmul.allow_tf32 = True + variants = ["flash", "pro"] if args.variant == "both" else [args.variant] + crs = [0, 4, 128] if args.cr == "all" else [int(args.cr)] + + print( + f"device={torch.cuda.get_device_name(0)} torch={torch.__version__} " + f"seq={args.seq} mbs={args.mbs} warmup={args.warmup} iters={args.iters}", + flush=True, + ) + for v in variants: + for cr in crs: + _bench_cr(v, cr, B=args.mbs, S=args.seq, warmup=args.warmup, iters=args.iters) + + +if __name__ == "__main__": + main() diff --git a/examples/deepseek-v4/benchmark/bench_v4_attention_results.md b/examples/deepseek-v4/benchmark/bench_v4_attention_results.md new file mode 100644 index 000000000..751c4ab68 --- /dev/null +++ b/examples/deepseek-v4/benchmark/bench_v4_attention_results.md @@ -0,0 +1,97 @@ +# DeepSeek-V4 Attention — Backend Performance + +Full forward + backward benchmark of every V4 attention backend, produced by +[`bench_v4_attention.py`](./bench_v4_attention.py). + +## Setup + +- **GPU**: AMD Instinct MI355X (gfx950), single GPU (`smci355-ccs-aus-n04-25`) +- **Container**: `dev_primus_wenx` +- **Torch**: `2.10.0+git94c6e04`, Triton `3.7.0`, FlyDSL `0.2.2` +- **Primus-Turbo**: `dev/kyle/flydsl_attn_deepseekv4` @ `350ec3f` (native-FlyDSL sparse-MLA v2) +- **Config**: `seq_len=4096`, `mbs=1`, bf16, sink **on**, `swa_window=128` +- **Models**: V4-Flash (`H=64`, index_topk=512), V4-Pro (`H=128`, index_topk=1024) +- **Layer kinds**: `cr=0` dense/SWA, `cr=4` CSA, `cr=128` HCA +- **Timing**: `--warmup 10 --iters 30`, median latency +- **Cell format**: `latency ms | TFLOP/s` + +Raw log: `agent/workspace/_bench_all_final.log` + +## Backends + +| Backend | Description | +|---------|-------------| +| `triton` | Production separate-K/V Triton dense/SWA/HCA + split CSA pool kernels | +| `gluon` | 1st-gen fused single-latent Gluon sparse-MLA (`_gluon_dsa`) | +| `triton_v2` | Fused single-latent sparse-MLA in plain Triton | +| `gluon_v2` | 2nd-gen Gluon sparse-MLA baseline | +| `gluon_v3` | Optimized Gluon sparse-MLA: Round-9 CSA formula-pack + aiter Gluon LSE fwd route (benchmark-only; not wired for training) | +| `flydsl_v1` | In-tree native FlyDSL MFMA sparse-MLA backend | +| `turbo_flydsl` | Extracted Primus-Turbo `sparse_mla_v2` (June `optimize/...` branch) | +| `_turbo_flydsl` | **Integrated** Primus-Turbo native-FlyDSL sparse-MLA via the turbo API (`primus_turbo.flydsl.attention`); the `turbo` model backend | +| `aiter_gluon` | aiter Gluon sparse-MLA prefill reference, fwd-only | + +`aiter_gluon` has no backward implementation, so bwd is shown as `—`. + +## Forward + +`latency ms | TFLOP/s`; **bold** is fastest latency in the row. + +| variant | cr | triton | gluon | triton_v2 | gluon_v2 | gluon_v3 | flydsl_v1 | turbo_flydsl | _turbo_flydsl | aiter_gluon | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| flash | 0 | 0.46 \| 151.0 | 0.31 \| 222.1 | 0.30 \| 230.0 | 0.28 \| 247.8 | 0.28 \| 248.3 | 0.45 \| 153.3 | 0.30 \| 228.7 | **0.20 \| 335.5** | 0.47 \| 145.2 | +| flash | 4 | 1.47 \| 234.3 | 0.89 \| 386.1 | 0.87 \| 397.1 | 0.73 \| 469.0 | 0.66 \| 523.6 | 1.37 \| 250.9 | 0.72 \| 476.6 | **0.53 \| 651.7** | 0.83 \| 413.4 | +| flash | 128 | 0.75 \| 114.7 | 0.38 \| 226.3 | 0.38 \| 223.9 | 0.33 \| 263.9 | 0.33 \| 263.2 | 0.58 \| 147.3 | 0.35 \| 245.5 | **0.22 \| 384.2** | 0.52 \| 164.7 | +| pro | 0 | 0.86 \| 159.5 | 0.57 \| 239.8 | 0.58 \| 236.2 | 0.51 \| 269.1 | 0.51 \| 269.0 | 1.05 \| 131.0 | 0.55 \| 251.9 | **0.38 \| 357.9** | 0.79 \| 173.1 | +| pro | 4 | 4.44 \| 278.3 | 2.91 \| 425.5 | 2.78 \| 444.3 | 2.36 \| 525.0 | 1.92 \| 645.1 | 4.82 \| 256.6 | 2.09 \| 591.0 | **1.41 \| 878.1** | 2.16 \| 573.4 | +| pro | 128 | 1.46 \| 117.7 | 0.72 \| 239.4 | 0.72 \| 238.6 | 0.61 \| 281.2 | 0.61 \| 280.9 | 1.20 \| 143.5 | 0.63 \| 270.7 | **0.43 \| 395.5** | 0.88 \| 195.2 | + +## Backward + +`latency ms | TFLOP/s`; **bold** is fastest latency in the row. + +| variant | cr | triton | gluon | triton_v2 | gluon_v2 | gluon_v3 | flydsl_v1 | turbo_flydsl | _turbo_flydsl | aiter_gluon | +|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| flash | 0 | 2.09 \| 82.2 | 1.20 \| 143.4 | 1.16 \| 148.4 | 1.13 \| 152.6 | 1.13 \| 152.0 | 2.20 \| 78.3 | 1.38 \| 124.6 | **0.67 \| 257.8** | — | +| flash | 4 | 5.18 \| 165.8 | 5.18 \| 166.0 | 5.93 \| 144.9 | 4.81 \| 178.5 | 3.99 \| 215.1 | 6.16 \| 139.5 | 3.94 \| 218.1 | **2.55 \| 336.9** | — | +| flash | 128 | 2.86 \| 75.0 | 1.63 \| 131.4 | 1.67 \| 128.8 | 1.54 \| 139.3 | 1.54 \| 139.2 | 2.77 \| 77.5 | 1.84 \| 117.0 | **0.78 \| 274.9** | — | +| pro | 0 | 4.02 \| 85.4 | 1.82 \| 189.0 | 1.81 \| 190.1 | 1.71 \| 201.4 | 1.70 \| 202.2 | 5.69 \| 60.3 | 2.09 \| 164.5 | **1.29 \| 267.2** | — | +| pro | 4 | 15.10 \| 204.8 | 13.48 \| 229.3 | 10.74 \| 287.8 | 8.52 \| 362.9 | 8.52 \| 362.9 | 30.45 \| 101.5 | 9.41 \| 328.7 | **6.32 \| 489.3** | — | +| pro | 128 | 5.55 \| 77.3 | 2.42 \| 177.4 | 2.47 \| 174.2 | 2.27 \| 189.4 | 2.27 \| 189.4 | 6.94 \| 61.9 | 2.91 \| 147.6 | **1.49 \| 288.2** | — | + +## `_turbo_flydsl` (integrated turbo) vs `gluon_v3` (best in-tree) + +`_turbo_flydsl` is the fastest backend on **every** fwd and bwd cell. Speedup over +`gluon_v3` (TFLOP/s ratio): + +| variant | cr | FWD speedup | BWD speedup | +|---|---:|---:|---:| +| flash | 0 | 1.35× | 1.70× | +| flash | 4 | 1.24× | 1.57× | +| flash | 128 | 1.46× | 1.98× | +| pro | 0 | 1.33× | 1.32× | +| pro | 4 | 1.36× | 1.35× | +| pro | 128 | 1.41× | 1.52× | +| **mean** | | **~1.36×** | **~1.57×** | + +## Summary + +- **`_turbo_flydsl`** (the integrated Primus-Turbo native-FlyDSL backend, selectable + in the model via `use_v4_attention_backend = turbo`) is the fastest backend across + all six shapes in both directions — **~1.36× fwd** and **~1.57× bwd** over the best + in-tree backend (`gluon_v3`), and larger margins over `triton_v2`/`gluon_v2`. +- The biggest wins are the CSA `cr=4` forward (pro cr=4: `1.41 ms` / 878 TF vs + gluon_v3 `1.92 ms` / 645 TF) and the backward across the board (flash cr=128 bwd + `0.78 ms` / 275 TF ≈ **2×** gluon_v3's `1.54 ms` / 139 TF). The fully-native FlyDSL + backward is the headline improvement. +- `_turbo_flydsl` also clearly beats the older extracted `turbo_flydsl` (June branch), + chiefly on the backward (e.g. pro cr=4 bwd `6.32 ms` vs `9.41 ms`). +- `gluon_v3` remains the strongest **in-tree** backend but is benchmark-only (not wired + for training); `gluon_v2` is the wired gluon training backend. + +## Reproduce + +```bash +PYTHONPATH= python examples/deepseek-v4/benchmark/bench_v4_attention.py \ + --variant both --cr all --warmup 10 --iters 30 +``` diff --git a/examples/deepseek-v4/projection/README.md b/examples/deepseek-v4/projection/README.md new file mode 100644 index 000000000..6c5b1319e --- /dev/null +++ b/examples/deepseek-v4/projection/README.md @@ -0,0 +1,93 @@ +# DeepSeek-V4 Performance Projection + +A trace-driven performance projection toolkit + static website for DeepSeek-V4 +(Flash / Pro) training on AMD Instinct GPUs (MI355X measured, MI455X projected). + +## Idea in one paragraph + +We profile a **single transformer layer** of a given compression-ratio (`cr`) +type on **MI355X**, extract a clean forward / backward **time + TFLOPs breakdown** +per module (attention sub-modules, MoE sub-modules), plus the model's non-layer +parts (embedding, output/logits, loss) and the optimizer step. We emit one JSON +per model variant. A static website then loads that JSON and, given a target GPU +and a distributed strategy (PP / VPP / EP / DP / CP), reconstructs the full model +(real layer count + `cr` schedule), models the PP bubble, EP dispatch/combine, +recompute, and the optimizer, and derives **iteration time, TFLOP/s/GPU and +tokens/s/GPU** step by step. Page 1 is the MI355X projection; page 2 scales the +breakdown to MI455X by theoretical ratios. + +## Pipeline + +``` +run profiling script (per cr) -> chrome trace JSON (rank 0) + | | + | v + | tools/parse_trace.py (+ kernel/module map) + v | + one trace per cr type {0,4,128} v + breakdown JSON (site/data/.json) + | + v + static site (site/index.html) + - model config view + - per-cr fwd/bwd breakdown tables + - GPU + parallelism controls + - step-by-step iter-time / TFLOPs / tok/s derivation + - iteration timeline (3 levels: layer / PP ranks / schedule) + - MI355X page + MI455X scaled page +``` + +## Directory layout + +``` +examples/deepseek-v4/projection/ + README.md # this file + design/ # methodology, assumptions, schema, math (the spec) + 01-overview.md + 02-assumptions.md + 03-json-schema.md + 04-projection-math.md + 07-iteration-timeline.md # 3-level iteration-time composition view + script/ # profiling launchers (one trace per cr) + deepseek_v4_layer_trace-projection.sh + tools/ # trace -> breakdown JSON + parse_trace.py + kernel_module_map.py + site/ # static website (no build step) + index.html + assets/app.js + assets/style.css + data/.json # breakdown JSON consumed by the site +``` + +## Quick start + +1. Profile each `cr` on MI355X (run inside the training container): + + ```bash + # one trace per cr type; 1 layer, seq 4096, adam + dist-opt, GA=2, no recompute + CR=0 ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh + CR=4 ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh + CR=128 ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh + ``` + +2. Build the breakdown JSON from the three traces: + + ```bash + python3 examples/deepseek-v4/projection/tools/parse_trace.py \ + --model pro \ + --trace cr0= \ + --trace cr4= \ + --trace cr128= \ + --out examples/deepseek-v4/projection/site/data/pro.json + ``` + +3. Open the site locally: + + ```bash + python3 -m http.server -d examples/deepseek-v4/projection/site 8000 + # http://localhost:8000/?model=pro + ``` + +See `design/` for the full methodology and the exact assumptions baked into the +projection math. diff --git a/examples/deepseek-v4/projection/design/01-overview.md b/examples/deepseek-v4/projection/design/01-overview.md new file mode 100644 index 000000000..eed63aac8 --- /dev/null +++ b/examples/deepseek-v4/projection/design/01-overview.md @@ -0,0 +1,91 @@ +# 01 — Overview & methodology + +## Goal + +Project DeepSeek-V4 (Flash / Pro) training throughput (iteration time, +TFLOP/s/GPU, tokens/s/GPU) for a real multi-hundred-GPU job, from a small set of +**single-layer** MI355X traces, then scale the result to MI455X by theoretical +hardware ratios. + +## Why single-layer, per-cr traces + +DeepSeek-V4 layers come in three attention flavours selected by the per-layer +`compress_ratio` (`cr`): + +| cr | attention branch | +|------|-----------------------------------------------| +| 0 | dense + sliding-window attention (SWA) | +| 4 | CSA (compressed sparse attention, top-k via Indexer) | +| 128 | HCA (compressed attention, full pool visibility) | + +The MoE block is **identical across all `cr`** (cr only changes attention). So +three single-layer traces — one per cr — fully characterise the per-layer cost, +and the full model is `Σ over layers (attention[cr_of_layer] + moe)`. + +Profiling a *single* layer per cr (instead of a full model) keeps: +- **memory** within the MI355X budget at the production `seq=4096`, and +- **attribution clean**: with one cr in the trace, the dense attention kernels + (which are shared across cr types in a multi-cr run) belong unambiguously to + that cr. + +## What we measure vs what we model + +**Measured on MI355X (from traces):** +- per-module forward / backward time for one layer of each cr, +- TFLOPs for the three compute-bound kernel classes (`gemm`, `grouped_gemm`, + `attn`); everything else is treated as memory-bound, +- embedding, output/logits, loss (the non-layer parts), taken once, +- the optimizer-step cost per unit parameter (for scaling). + +**Modeled on top (in the website):** +- full layer count and exact `cr` schedule (Flash 43L, Pro 61L), +- PP / VPP partitioning -> per-stage critical path + pipeline bubble, +- EP dispatch/combine cost (no overlap in current stack), +- gradient-accumulation (GA) and DP behaviour (DP/PP comm assumed hidden), +- activation recompute (add a forward pass to the recomputed layers' backward), +- optimizer step scaled to per-rank parameter count for the target sharding, +- MI455X = MI355X breakdown scaled by compute / memory-bandwidth ratios. + +## Capture configuration (the projection trace) + +The profiling script (`script/deepseek_v4_layer_trace-projection.sh`) deliberately +differs from `run_deepseek_v4_pro_muon.sh`: + +| knob | projection value | why | +|-------------------|------------------|-----| +| `seq_length` | 4096 | production per-microbatch token count | +| `num_layers` | 1 | clean per-layer attribution; fits memory at seq 4096 | +| `compress_ratios` | `[CR]` | one cr per trace | +| optimizer | `adam` + distributed optimizer | compute is optimizer-independent; dist-opt = zero1 | +| overlap grad/param | off | num_layers=1 breaks Megatron's chained param-sync; compute already clean | +| `global_batch_size` | `2 * DP` | GA=2 (see below) | +| recompute | off (`recompute_num_layers 0`) | capture pure fwd / pure bwd | +| profiler | on, `with_stack=True`, window iter 6->7 | map kernels -> nn.module | + +### Why overlap-off + GA=2 + take min + +We capture with the distributed optimizer's comm overlap **off**. Two reasons: + +- with `num_layers=1`, Megatron's chained param-gather sync trips an assertion + (`param_and_grad_buffer.start_param_sync`) when overlap is on; and +- with overlap off, **every microbatch's compute is already overlap-free** — + exactly the clean per-kernel time we want. In a real run GA is large so the + vast majority of microbatches are clean anyway, and the projection assumes DP + comm is hidden (A2), so there is no need to measure the contaminated overlap. + +GA=2 (two microbatches) plus a late steady profiler window lets the parser +compute per-microbatch time as `sum(kernel_durations) / num_microbatches`. +This keeps scalar control-flow stalls (for example Indexer top-k syncs) that the +full-model calibration shows are real per-layer costs, while avoiding warm-up +and autotune iterations. + +## Non-overlap assumptions in the current stack + +- **No MoE deepep-comm + grouped-gemm overlap.** dispatch + grouped_gemm + + combine are summed directly. +- **EP dispatch/combine has no overlap** with compute; counted in full. +- **DP / PP comm assumed hidden** (only the PP bubble remains). Optimistic; first + thing to revisit at calibration time. + +See `02-assumptions.md` for the complete list, and `04-projection-math.md` for +the formulas. diff --git a/examples/deepseek-v4/projection/design/02-assumptions.md b/examples/deepseek-v4/projection/design/02-assumptions.md new file mode 100644 index 000000000..c86e751d8 --- /dev/null +++ b/examples/deepseek-v4/projection/design/02-assumptions.md @@ -0,0 +1,108 @@ +# 02 — Assumptions (single source of truth) + +Every assumption baked into the projection. When a projection number looks off, +start here. + +## Optimizer / DP + +- **A1.** Production optimizer modeled = **AdamW + distributed optimizer (zero1)**. + Muon is out of scope. NOTE on **capture**: the trace itself is taken with + `use_distributed_optimizer=False` (+ fp32 states) because with dist-opt ON the + ROCm Kineto profiler drops the compute GPU kernels for pure dense(cr=0)/HCA + (cr=128) layers (CSA cr=4 is unaffected). dist-opt does not change the fwd/bwd + compute, so this only affects which kernels Kineto records; the optimizer step + is modeled analytically regardless (A3). +- **A2.** DP communication (param all-gather / grad reduce-scatter) is **fully + hidden** behind compute at large GA. We do not add a DP comm term. *(Optimistic; + primary calibration target.)* +- **A3.** The optimizer step is a per-iteration term, **not** multiplied by GA or + replicated per PP microbatch. It scales with **per-rank optimizer parameter + count**: full-model params are first averaged over PP/TP ownership, then + sharded over DP under ZeRO-1. CP does not shard parameters. The modeled Adam + traffic uses the full mixed-precision read/write cost per parameter and is + treated as memory-bound. + +## Pipeline / parallelism + +- **A4.** PP point-to-point comm is **hidden**; only the pipeline **bubble** + remains. CP/TP comm not modeled in v1 (TP=1, CP=1 in the V4 release configs). +- **A5.** Pipeline bubble fraction uses 1F1B: `(PP-1)/GA`; interleaved VPP divides + it by the VPP degree: `(PP-1)/(GA*VPP)`. +- **A6.** Per-stage time is the **sum of the specific cr-type layers** mapped to + that stage; the iteration critical path is driven by the **max** (slowest) + stage, plus embedding on stage 0 and output/loss on the last stage. + +## EP / MoE + +- **A7.** EP dispatch/combine has **no overlap** with compute (current stack) and + is counted in full, every microbatch. +- **A8.** No MoE deepep-comm + grouped-gemm overlap; MoE = dispatch + grouped_gemm + + combine summed. +- **A9.** EP is intra-node only (e.g. EP=8 within an 8-GPU node). Cross-node EP is + out of scope for v1; if EP spans nodes the dispatch/combine cost model must + change (RDMA, different bandwidth). +- **A10.** MoE per-layer cost is `cr`-independent; the three single-cr traces + must agree on it (cross-check). The site uses one MoE breakdown for all layers. + +## Trace capture / attribution + +- **A11.** One trace per cr (`0`, `4`, `128`), **1 layer**, `seq=4096`, + `recompute off`, profiler window iter 6->7 (post warmup/autotune). +- **A12.** Capture with comm-overlap **off** (`num_layers=1` breaks Megatron's + chained param sync, and compute is already clean without overlap). GA=2; clean + per-kernel time = `min` over launches grouped by `(module, phase, shape)`, + removing residual jitter. +- **A13.** Kernel -> nn.module attribution uses `with_stack=True`: GPU kernels are + linked to their launching CPU op via trace flow events, and the module is read + from the CPU op's python call stack. **fwd/bwd phase** is determined by, in + priority order: (1) a `_fwd_`/`_bwd_` tag in the kernel name; (2) for linked + kernels, whether the launching CPU op's timestamp lies inside an + `autograd::engine::evaluate_function` interval (= backward); (3) for unlinked + kernels (no `External id`), whether the kernel's GPU timestamp lies in the + backward GPU-time window reconstructed from the linked backward kernels. The + old rule (default-to-forward when unlinked) leaked backward compute -- incl. + the MoE dgrad/wgrad grouped GEMMs -- into forward; see `design/06`. One-off + device stalls billed to a compute kernel (> `_MAX_PLAUSIBLE_LAUNCH_US`) are + dropped as artifacts and reported in `provenance.dropped_stall_us_per_mb`. +- **A14.** Only `gemm`, `grouped_gemm`, `attn` kernels get a TFLOPs number; all + other kernels are memory-bound (TFLOPs = null) and contribute time only. +- **A15.** Embedding / output-logits / loss are taken **once** (from any single + trace), not triple-counted across the three cr traces. + +## Recompute + +- **A16.** Traces are captured with recompute **off** (pure fwd, pure bwd). At + projection time, a recomputed layer's backward gets `+1 forward` of that layer + added back. Recompute selection (#layers / which) is a site control. + +## MTP + +- **A17.** MTP is modeled analytically when `mtp_num_layers > 0`. Each MTP + depth uses `mtp_compress_ratios` (default cr=4, matching the current Flash + Megatron FLOPs anchor), plus the per-depth `eh_proj`, extra logits/loss, and + HyperHead FLOPs. Timing is an approximation: the MTP inner layer reuses the + measured layer time for that cr, while `eh_proj` is scaled from output-GEMM + throughput. A dedicated MTP trace remains the next calibration step. + +## Scope exclusions (v1) + +- **A18.** No cross-node EP; no TP; no CP comm modeling. +- **A19.** FP8 / MXFP8 not modeled; BF16 only. + +## MI355X -> MI455X scaling + +- **A20.** Compute-bound kernels (`gemm`, `grouped_gemm`, `attn`) scale by the + **peak-TFLOPs ratio** (BF16) `t_mi355 / t_mi455`. +- **A21.** Memory-bound kernels (everything else, incl. optimizer) scale by the + **HBM-bandwidth ratio** `bw_mi355 / bw_mi455`. +- **A22.** A single tunable **efficiency factor** (default 1.0) multiplies the + scaled compute time to account for MFU differences on new HW (flat peak ratio + is optimistic). +- **A23.** Comm (EP/DP/PP) is not rescaled in v1 (intra-node EP only; DP/PP + hidden). + +## Validation + +- **A24.** Self-consistency: configured to the trace scenario (PP=1, EP=8, + measured GA, single node) the model must reproduce the measured single-node + iteration time. Multi-node calibration is deferred. diff --git a/examples/deepseek-v4/projection/design/03-json-schema.md b/examples/deepseek-v4/projection/design/03-json-schema.md new file mode 100644 index 000000000..15f6e66fd --- /dev/null +++ b/examples/deepseek-v4/projection/design/03-json-schema.md @@ -0,0 +1,117 @@ +# 03 — Breakdown JSON schema + +One JSON per model variant (`flash` / `pro`), written to +`site/data/.json`, consumed by the website. All times in **microseconds +(us)**; all FLOP counts are per single layer / per single invocation at the +captured `seq` and `micro_batch_size`. + +## Top level + +```jsonc +{ + "schema_version": 1, + "model": "pro", // "flash" | "pro" + "generated_at": "2026-06-18T11:00:00Z", + "provenance": { + "commit": "dac0a60c", + "host": "smci355-ccs-aus-n06-25", + "container": "dev_primus_wenx", + "traces": { "cr0": "", "cr4": "", "cr128": "" } + }, + + "capture": { // how the trace was taken (the unit) + "gpu": "MI355X", + "seq_length": 4096, + "micro_batch_size": 1, + "tokens_per_microbatch": 4096, // seq_length * micro_batch_size + "ep": 8, + "ga_for_capture": 2, + "optimizer": "adam", + "distributed_optimizer": true, + "recompute": "off" + }, + + "model_config": { // shown on the site's config panel + "num_layers": 61, + "hidden_size": 7168, + "num_attention_heads": 128, + "kv_channels": 512, + "num_experts": 384, + "moe_router_topk": 6, + "moe_ffn_hidden_size": 3072, + "index_topk": 1024, + "vocab_size": 129280, + "compress_ratios": [128,128,4, /* ... */ ,0], + "cr_layer_counts": { "0": 1, "4": 29, "128": 31 } // derived from compress_ratios + }, + + "hardware": { // for MI355->MI455 scaling (A20-A23) + "MI355X": { "peak_tflops_bf16": 2300, "hbm_bandwidth_gbps": 8000 }, + "MI455X": { "peak_tflops_bf16": 4600, "hbm_bandwidth_gbps": 16000 } + }, + + "layers": { // per-cr breakdown + "0": { "attention": , "moe": }, + "4": { "attention": , "moe": }, + "128": { "attention": , "moe": } + }, + + "non_layer": { // taken once (A15) + "embedding": , + "output": , // final norm + lm_head/logits + "loss": + }, + + "optimizer": { // per-iteration term (A3) + "type": "adam", + "measured_params": 123456789, // params updated in the 1-layer trace, this rank + "time_us": 850.0, // measured optimizer-step time for those params + "bytes_per_param": 18, // bf16 master-param-remainder adam state bytes + "class": "memory_bound" + }, + + "comm": { // per-microbatch EP cost (A7-A9), per layer + "ep_dispatch_us": 0.0, + "ep_combine_us": 0.0 + } +} +``` + +## `` object + +A phase-split list of modules. Each module is one row group; the site renders +forward left-to-right and backward right-to-left. + +```jsonc +{ + "forward": [ , ... ], + "backward": [ , ... ] +} +``` + +## `` object + +```jsonc +{ + "module": "attn.core", // logical module name (from kernel/module map) + "time_us": 740.0, // clean min-grouped time for this module/phase + "class": "compute_bound", // "compute_bound" | "memory_bound" + "flop_class": "attn", // "gemm" | "grouped_gemm" | "attn" | null + "flops": 1.84e12, // total FLOPs for this module/phase, or null + "tflops": 740.0, // achieved TFLOP/s = flops / time_s / 1e12 (null if memory_bound) + "kernels": [ // optional: contributing kernels (debug / drill-down) + { "name": "_v4_attention_fwd_kernel", "time_us": 500.0, "launches": 1 } + ] +} +``` + +## Notes + +- `time_us` is always the **clean** (min-grouped, overlap-free) time per A12. +- `tflops` is present only when `flop_class != null` (A14). +- The site computes everything else (full model, PP/EP/DP, MI455) from this JSON; + the JSON itself is hardware-MI355X, single-layer, single-microbatch ground + truth + static config. +- `cr_layer_counts` is derived from `compress_ratios` by the parser so the site + doesn't re-parse the schedule. +- For Flash, `model` = `"flash"`, `cr_layer_counts` e.g. `{"0":3,"4":20,"128":20}`. diff --git a/examples/deepseek-v4/projection/design/04-projection-math.md b/examples/deepseek-v4/projection/design/04-projection-math.md new file mode 100644 index 000000000..fbe192bcf --- /dev/null +++ b/examples/deepseek-v4/projection/design/04-projection-math.md @@ -0,0 +1,193 @@ +# 04 — Projection math + +This is the exact derivation the website implements (`site/assets/app.js`). All +inputs come from the breakdown JSON (`03-json-schema.md`); all assumptions are in +`02-assumptions.md`. Times in seconds unless noted. + +Notation: +- `cr ∈ {0, 4, 128}`, `n[cr]` = number of layers of that cr (`cr_layer_counts`). +- A `` has `forward` / `backward` module lists. Define + `T(bd, phase) = Σ_module bd[phase][m].time` (sum of clean module times). + +## Step 0 — per-layer and non-layer base times (MI355X) + +For each cr: +``` +layer_fwd[cr] = T(attention[cr], forward) + T(moe, forward) +layer_bwd[cr] = T(attention[cr], backward) + T(moe, backward) +``` +EP dispatch/combine are included as memory-bound module rows inside `moe` +(A7/A8), so they are already in these sums — do not add `comm` again (the `comm` +field is informational for the UI). + +Non-layer (taken once, A15): +``` +emb_fwd = T(embedding, forward); emb_bwd = T(embedding, backward) +out_fwd = T(output, forward) + T(loss, forward) +out_bwd = T(output, backward) + T(loss, backward) +``` + +MTP (if `mtp_num_layers > 0`, A17): +``` +mtp_inner_fwd = layer_fwd[mtp_cr] ; mtp_inner_bwd = layer_bwd[mtp_cr] +mtp_out_fwd = out_fwd ; mtp_out_bwd = out_bwd +mtp_eh_time ≈ output_time * (mtp_eh_proj_flops / output_flops) + +mtp_fwd = mtp_num_layers * (mtp_inner_fwd + mtp_out_fwd) + mtp_eh_time / 3 +mtp_bwd = mtp_num_layers * (mtp_inner_bwd + mtp_out_bwd) + 2 * mtp_eh_time / 3 +``` +The current Flash Megatron FLOPs anchor uses `mtp_cr=4`; a future dedicated MTP +trace can replace the `eh_proj` and inner-layer approximation. + +## Step 0b — manual layer-time mode (optional, UI-only) + +The site exposes a **layer-timing mode** toggle in the controls panel: + +- `trace` (default): `layer_fwd/bwd[cr]` come from the breakdown JSON exactly as + in Step 0. +- `manual`: the user types `layer_fwd[cr]` / `layer_bwd[cr]` directly (one + fwd/bwd pair per `cr ∈ {0,4,128}`, in µs) and those values replace the + trace-derived per-layer times. This is a what-if calculator: you supply the + per-layer cost and the site reuses Steps 1-5 unchanged to derive the full-model + iteration time, bubble, optimizer, tokens/s and TFLOP/s. + +Rules baked into the implementation: + +- **Granularity is per-`cr`** (same as the data model — `cr` only changes + attention, the MoE block is shared), not per physical layer. The `cr` schedule, + PP/VPP layout and recompute still expand these per-cr times to the full model. +- **Per-GPU storage.** A hand-entered time already targets one GPU, so manual + values are stored separately for MI355X and MI455X and the + MI355→MI455 scaling (Step 6) is **bypassed** in manual mode. Switching the GPU + tab edits that GPU's own set. +- **Prefill from trace.** Entering manual mode (or switching GPU within it) + seeds any unset field with the current trace-derived value, so toggling never + changes the result until you actually edit a number. Unset/blank fields keep + falling back to trace. +- **Time only, FLOPs stay analytic.** Manual input overrides time but not FLOPs; + `TFLOP/s/GPU` continues to use the V4 analytic model FLOPs (Step 5), so it + stays meaningful. +- **Scope.** Manual covers the three per-cr decoder layers **and** the non-layer + parts — embedding (PP stage 0), output / loss / MTP (last PP stage) — each as a + fwd/bwd pair. Any field left unset falls back to its trace-derived value, so you + can override just the parts you care about. FLOPs stay analytic regardless. + +## Step 1 — recompute (A16) + +If a layer is recomputed, its backward replays one forward: +``` +layer_bwd_eff[cr] = layer_bwd[cr] + recompute_factor[cr] * layer_fwd[cr] +``` +`recompute_factor` ∈ {0,1} per layer. The site exposes `none`, `full`, and +`first-n`; `first-n` adds one forward replay to the first N decoder layers owned +by each physical PP stage, matching the common Megatron +`recompute_num_layers=N` block pattern. + +## Step 2 — map layers to PP stages / VPP chunks (A6) + +Inputs: `PP`, `VPP`, optional `pipeline_layout`. Total model chunks +`C = PP * VPP`. If `pipeline_layout` is provided, parse Megatron-style `t` / +`t*N` stage specs (for example `Et*10|t*11|t*11|t*11mL`) and assign virtual +chunk `k` to device `k mod PP`. Otherwise build the ordered layer list from +`compress_ratios`, slice it into `C` contiguous chunks, and use the same +`k mod PP` mapping. The UI validates that an explicit layout has exactly +`PP*VPP` stages and exactly `num_layers` decoder layers; invalid layouts block +projection instead of silently falling back. For device `d`: +``` +Df[d] = Σ_{chunks on d} Σ_{layer in chunk} layer_fwd[cr(layer)] +Db[d] = Σ_{chunks on d} Σ_{layer in chunk} layer_bwd_eff[cr(layer)] +``` +Add non-layer parts to their devices: +``` +Df[0] += emb_fwd ; Db[0] += emb_bwd +Df[PP-1] += out_fwd ; Db[PP-1] += out_bwd +Df[PP-1] += mtp_fwd ; Db[PP-1] += mtp_bwd +``` +Critical device: +``` +Df_crit = max_d Df[d] ; Db_crit = max_d Db[d] +``` +(Using per-device max is an upper-bound for imbalanced stages; for a balanced +schedule it is exact.) + +## Step 3 — pipeline iteration time (A4/A5) + +With gradient accumulation `GA` microbatches and interleaved VPP, the steady +1F1B time on the critical device is: +``` +pipe_compute = (GA + (PP - 1) / VPP) * (Df_crit + Db_crit) +``` +- `GA * (Df_crit + Db_crit)` is the steady throughput term; +- `(PP-1)/VPP * (Df_crit + Db_crit)` is the bubble (fraction `(PP-1)/(GA*VPP)`). +- PP=1 ⇒ `pipe_compute = GA * (Df_crit + Db_crit)` (no bubble). + +`GA = GBS / (DP * MBS)`. The site takes `GBS`, `MBS`, `DP` as inputs (or derives +`DP = world_size / (PP * TP * CP)` with EP ⊆ DP). + +## Step 4 — optimizer step (A1/A3) + +Per-iteration, once, zero1-sharded, memory-bound: +``` +local_model_params = total_params / (PP * TP) +per_rank_opt_params = local_model_params / DP +opt_bytes = per_rank_opt_params * bytes_per_param +opt_time = opt_bytes / hbm_bandwidth / opt_efficiency +``` +`total_params` is computed from `model_config` for the full model (dense + +expert params + untied embedding/output). PP/TP determine the average local +model-parameter ownership; CP does not shard parameters. EP is represented in +the full expert count and cancels with the data-replica count for ZeRO-1 +optimizer sharding, so the average optimizer shard is `total/(PP*TP*DP)`. +`bytes_per_param` is the full Adam mixed-precision step traffic (default 30B: +reads + writes), and `opt_efficiency` is tunable. The measured +`optimizer.time_us` carried in the JSON is displayed as a sanity reference. + +## Step 5 — totals (A2/A4: DP & PP comm hidden) + +``` +iter_time = pipe_compute + opt_time +``` +Throughput: +``` +tokens_per_iter = GBS * seq_length (= GA * DP * MBS * seq) +tokens_per_s = tokens_per_iter / iter_time +tokens_per_s_per_gpu = tokens_per_s / world_size +``` +TFLOP/s/GPU (matmul-flops convention, A14): per-microbatch model compute FLOPs +``` +F_mb = Σ_cr n[cr] * (flops_fwd[cr] + flops_bwd[cr]) + nonlayer_flops + where flops_*[cr] = Σ_module (module.flops or 0) over the cr breakdown + + mtp_inner + mtp_eh_proj + mtp_extra_logits + mtp_hc_head +flops_per_iter = F_mb * GA * DP (all microbatches, all DP replicas) +TFLOP_s_per_gpu = flops_per_iter / iter_time / world_size / 1e12 +``` +`tokens_per_s_per_gpu` is the headline metric (independent of FLOP convention); +`TFLOP/s/GPU` is reported for comparison with Primus' own logging. + +## Step 6 — MI455X scaling (A20-A23) + +Rescale every module time before re-running Steps 0-5: +``` +ratio_compute = peak_tflops_bf16[MI355] / peak_tflops_bf16[MI455] +ratio_memory = hbm_bandwidth[MI355] / hbm_bandwidth[MI455] + +time'(module) = module.time * ratio_compute / compute_efficiency if compute_bound + = module.time * ratio_memory if memory_bound +``` +`compute_efficiency` (default 1.0) is the MFU knob (A22). FLOPs are unchanged +(same math), so MI455 `tflops` rises by `1/ratio_compute * compute_efficiency`. +Optimizer scales by `ratio_memory`. Comm not rescaled (A23). + +## Step 7 — self-consistency check (A24) + +Configure `PP=1, VPP=1, EP=8, DP=1, MBS=1, GA=GA_capture` and confirm the model's +`iter_time` matches the measured single-node iteration time within tolerance. The +site shows this check on the MI355X page when capture metadata is present. + +## Worked control set (website inputs) + +GPU page (MI355X / MI455X), then: `world_size`, `PP`, `VPP`, `EP`, `DP` (or +derive), `CP`, `TP`, `MBS`, `GBS` (or `GA`), recompute mode, and the tunables +`opt_efficiency`, `compute_efficiency`, `bytes_per_param`. Every intermediate +(`layer_fwd/bwd`, `Df/Db` per stage, `pipe_compute`, bubble %, `opt_time`, +`iter_time`, `tokens/s/gpu`, `TFLOP/s/gpu`) is displayed step by step. diff --git a/examples/deepseek-v4/projection/design/05-deployment.md b/examples/deepseek-v4/projection/design/05-deployment.md new file mode 100644 index 000000000..3ca252ada --- /dev/null +++ b/examples/deepseek-v4/projection/design/05-deployment.md @@ -0,0 +1,58 @@ +# 05 — Deployment (GitHub Pages) + +The repository already publishes a single GitHub Pages site from `main` via +`.github/workflows/deploy-backend-gap-dashboard.yml`, which builds a bundle with +`tools/backend_gap_report/build_site_bundle.py` and deploys it with +`actions/deploy-pages`. A repo can only serve one Pages site, so the projection +site is published as a **subpath of that same bundle** rather than as a separate +deployment. + +## How it is wired + +`build_site_bundle.py` copies `examples/deepseek-v4/projection/site/` into the bundle at +`deepseek-v4-projection/` (after the backend-gap bundle is built, before +validation). The projection site uses only relative asset/data paths +(`./assets/...`, `./data/...`), so it works unchanged under a subpath. + +Result URL: + +``` +https://.github.io//deepseek-v4-projection/?model=pro +https://.github.io//deepseek-v4-projection/?model=flash +``` + +## Triggering a deploy + +The Pages workflow triggers on pushes to `main` touching its `paths:` list +(currently `docs/backend-gap/**`, `docs/weekly_reports/**`, +`docs/monthly_reports/**`, `tools/backend_gap_report/**`, and the workflow file). + +- The change to `tools/backend_gap_report/build_site_bundle.py` in this work is + itself under a watched path, so the **first** merge to `main` will rebuild and + publish the projection site automatically. +- To make **projection-only** changes (new `site/data/*.json`, site tweaks) also + auto-deploy, add the projection path to the workflow's `paths:` on `main`: + + ```yaml + # .github/workflows/deploy-backend-gap-dashboard.yml (on: push: paths:) + - "examples/deepseek-v4/projection/site/**" + ``` + +- Or trigger manually: the workflow has `workflow_dispatch` (Run workflow button). + +No `.nojekyll` is needed: the bundle is uploaded as a Pages artifact and served +directly (no Jekyll processing), and no asset path starts with `_`. + +## Local preview + +```bash +python3 -m http.server -d examples/deepseek-v4/projection/site 8011 +# http://localhost:8011/?model=pro +``` + +## Bundle-build smoke (optional, needs pandoc/weasyprint for the backend-gap PDFs) + +```bash +python3 tools/backend_gap_report/build_site_bundle.py --output-dir /tmp/primus-site +ls /tmp/primus-site/deepseek-v4-projection/ # index.html, assets/, data/ +``` diff --git a/examples/deepseek-v4/projection/design/06-calibration.md b/examples/deepseek-v4/projection/design/06-calibration.md new file mode 100644 index 000000000..72373ef03 --- /dev/null +++ b/examples/deepseek-v4/projection/design/06-calibration.md @@ -0,0 +1,73 @@ +# 06 — Calibration (single-node, measured) + +The projection is anchored to a real single-node run so its iteration time and +TFLOP/s line up with what Megatron reports. + +## FLOPs: ported V4 closed-form (exact) + +`tools/v4_flops.py` ports Megatron's `deepseek_v4_flops_patches` closed form. +Self-test against the measured flash 16-layer run (GBS64, seq4096): + +| component | analytic (TFLOP/gb) | measured (TFLOP/gb) | +|---|---:|---:| +| attn_qkv_o | 10766.4 | 10766.4 | +| attn_scores | 2476.0 | 2476.0 | +| compressor | 527.8 | 527.8 | +| indexer | 1650.9 | 1650.9 | +| moe | 17838.5 | 17818.7 | +| logits | 832.9 | 833.7 | +| **TOTAL** | **34112** | **34093** (0.05%) | + +The site uses these analytic FLOPs (per cr-layer, B=1, capture seq) for TFLOP/s, +matching Megatron's convention (fwd+bwd × FMA = 6×, recompute excluded). The +breakdown JSON carries them in `analytic_flops`. + +## Iteration time: single-layer → full-model bias + +Measured single-node anchor (`script/_calibrate_flash.sh`, full flash): + +| knob | value | +|---|---| +| layers / cr | 16, cr=[0×3, 4×6, 128×7] | +| parallel | PP1 / EP8 / DP8 (world 8), TP1/CP1 | +| GBS / GA / MBS / seq | 64 / 8 / 1 / 4096 | +| recompute | full (uniform, 1) | +| optimizer | adam + distributed optimizer | + +**Measured**: iter ≈ 6665 ms, 636 TFLOP/s/GPU, ~4917 tokens/s/GPU. + +**Projection (raw, calibFactor=1.0)**: iter 7177 ms (+7.7%), 586 TFLOP/s/GPU +(−7.9%). The per-layer time captured from the single-layer profile runs ~7-8% +high vs a layer inside the full model (single-layer capture has no neighbour- +layer overlap / cache reuse, and per-launch overhead is a larger share). This is +a systematic, near-constant bias. + +**Calibration**: a single `calibFactor = 0.93` on the pipeline compute time +brings it in line: + +| metric | measured | projection (calibFactor 0.93) | +|---|---:|---:| +| iter time | 6665 ms | ~6680 ms (+0.2%) | +| TFLOP/s/GPU | 636 | ~630 (−1%) | +| tokens/s/GPU | 4917 | ~4900 (−0.4%) | + +`calibFactor` is a site control (default 0.93). + +## Caveats + +- `calibFactor` is from one anchor (flash, 16L, PP1). Pro / other parallel + layouts may want a slightly different value; re-anchor with another + `_calibrate_*` run if precision matters. +- analytic FLOPs are evaluated at the capture seq (4096); changing seq in the + UI does not re-derive them (attention FLOPs are seq-dependent). +- Optimizer step is analytic (per-rank params / HBM-BW); DP/PP comm assumed + hidden (A2/A4). + +## Reproduce + +```bash +# measured anchor (single node, full model) +LAYERS=16 GBS=64 bash examples/deepseek-v4/projection/script/_calibrate_flash.sh # (helper; not committed) +# analytic flops self-test +python3 examples/deepseek-v4/projection/tools/v4_flops.py +``` diff --git a/examples/deepseek-v4/projection/design/07-iteration-timeline.md b/examples/deepseek-v4/projection/design/07-iteration-timeline.md new file mode 100644 index 000000000..1dcc38386 --- /dev/null +++ b/examples/deepseek-v4/projection/design/07-iteration-timeline.md @@ -0,0 +1,165 @@ +# 07 — Iteration timeline (3-level composition view) + +A visual, drill-down view of **how one training iteration's time is composed**, +layered from a single layer up to the whole pipeline schedule. It reuses the +projection math in `04-projection-math.md` (same controls, same per-layer times) +and adds only rendering + one pipeline-schedule simulator. Nothing here changes +the headline `iteration time` / `tokens/s/GPU` numbers; it explains them. + +All times are microseconds (µs) unless noted, scaled to the active GPU tab by +Step 6 (`rowScaledTime`) exactly like the rest of the site. + +## Why three levels + +The iteration time is built bottom-up: + +``` +module time --(sum per phase)--> layer fwd/bwd (Level 1) +layer times --(map to PP/VPP)--> per-device chunk cost (Level 2) +device costs --(1F1B schedule)--> iteration timeline (Level 3) +``` + +Each level answers one question: + +- **Level 1 — where does a single layer's time go?** (attn / mlp / a2a) +- **Level 2 — how is work distributed across pipeline ranks?** (layer granularity) +- **Level 3 — how do the ranks overlap in time, and where is the bubble?** + +## Module → category mapping (Level 1 granularity) + +The default minimum granularity is three categories plus an explicit +"unattributed" bucket, derived from the `module` field (`03-json-schema.md`): + +| category | modules | meaning | +|----------|---------|---------| +| `attn` | `attn.norm`, `attn.proj`, `attn.core`, `attn.indexer`, `attn.misc` | attention (varies by `cr`) | +| `mlp` | `moe.grouped_gemm`, `moe.shared_expert`, `moe.router` | expert / shared-expert compute (cr-independent) | +| `a2a` | `moe.dispatch`, `moe.combine` | EP all-to-all (captured at EP=8) | +| `misc` | any `*.misc` / unattributed | casts, scalar, control-flow; kept visible | + +Rules: + +- `attn.misc` counts as `attn` (it is attention-local unattributed time), while a + standalone `misc` category is only used if a non-attn/non-moe module is + unmapped. In practice every row maps to `attn` or `mlp`/`a2a`; the `misc` + bucket surfaces `*.misc` share the same way `unattributedShare` does today. +- Because MoE is identical across `cr` (A10), only three representative layers are + shown: `cr=0`, `cr=4`, `cr=128`. "Same params/category → show one" reduces to + "one per cr". +- Optional drill-down expands `attn` into `core/indexer/proj/norm` and `mlp` into + `grouped_gemm/shared_expert/router`. + +## Level 2 — per-device chunk composition + +Granularity is **one physical decoder layer**. The projection already maps every +layer to a `PP*VPP` chunk and a device (`04` Step 2). Level 2 exposes that map: + +- Each device (PP rank) is a row; within it, chunks (VPP virtual chunks) are laid + out in schedule order, each chunk a run of layers coloured by `cr`. +- Recomputed layers (their backward replays one forward, `04` Step 1) are marked + (hatched / outlined) because they cost extra in `Db`. +- Non-layer parts are drawn on their owning device: `embedding` on device 0, + `output`+`loss`+`MTP` on the last device. +- **Dedup:** devices with an identical ordered signature + `(cr-list, recompute-flags, hasEmb, hasOut, hasMtp)` are drawn once, annotated + `×N ranks (d..d)`. +- The **critical device** (`max Df` / `max Db`, the one that sets the pipeline + critical path) is highlighted; each row shows its `Df`/`Db` and share of the + critical stage, making load imbalance (the bubble's root cause) visible. + +## Level 3 — pipeline schedule (Megatron-2 Figure 4 style) + +A Gantt chart: y-axis = device, x-axis = time. Forward cells one colour, backward +another; VPP virtual chunks distinguished by lightness; bubbles are gaps. + +Reference: Narayanan et al., "Efficient Large-Scale Language Model Training on +GPU Clusters Using Megatron-LM", arXiv:2104.04473, Figure 4 (default 1F1B on top, +interleaved below). + +### Colour scheme + +- forward = `--accent` (#4f8cff), backward = `--accent-2` (#36c08f). +- VPP chunk index shifts lightness (chunk 0 lightest → deeper for higher chunks), + mirroring the paper's light/dark model-chunk shading. +- bubble = empty (optionally a faint diagonal hatch) with the fraction labelled. + +### Schedule simulator + +The site's analytic pipe time is +`pipe_compute = (GA + (PP-1)/VPP) * (Df_crit + Db_crit)` (`04` Step 3). To *draw* +the schedule we simulate 1F1B (optionally interleaved) microbatch ordering and +let bubbles emerge from the gaps. + +Inputs: `PP`, `VPP`, `GA`, and per-virtual-chunk forward/backward cost. Each +virtual chunk `k` (device `k % PP`, vpp `⌊k/PP⌋`) uses the **exact** per-chunk +sums from `schedule.chunks` (`f_chunk[k]`, `b_chunk[k]`), with the non-layer +parts folded into the first (`embedding`) and last (`output`/`loss`/`MTP`) virtual +chunk so the drawn per-device length stays consistent with `Df`/`Db`. The +non-interleaved view collapses to one chunk per device with `f=Df[d]`, `b=Db[d]`. +Large `GA` is capped to `TL_VIS_GA_CAP` (48) microbatches for display only. + +Algorithm (interleaved 1F1B, `VPP` model chunks per device, `k mod PP` device +map): + +``` +num_chunks = PP * VPP +num_warmup = min(GA, (PP - 1 - device) * ... ) # standard interleaved warmup +events[device] = ordered list of {kind:'F'|'B', mb, chunk, start, dur} +- time advances per device; a device starts an op when both its own timeline and + the producing/consuming neighbour dependency are satisfied (F flows forward + along devices, B flows backward), with p2p comm assumed zero (A4). +``` + +The simulator returns per-device event lists with `start`/`dur`; the drawn +iteration length `max_device(last_end)` is compared against the analytic +`pipe_compute` (pre-`calibFactor`). For a balanced schedule they agree; a +mismatch beyond tolerance is surfaced as a warning rather than hidden. The +analytic number remains the official iteration time; the Gantt chart is the +visualization and is scaled so its total width equals the analytic pipe time. + +### Interleaving follows VPP + +There is no separate interleaved/non-interleaved toggle: the schedule is driven +directly by the `VPP` control. `VPP=1` renders plain 1F1B (Figure 4 top); +`VPP>1` renders interleaved 1F1B (Figure 4 bottom), which shrinks the bubble from +`(PP-1)/GA` to `(PP-1)/(GA*VPP)` (A5). Change `VPP` in the projection controls to +compare. + +## UI: layout + cross-level linkage + +- **Layout toggle.** *Tabbed* shows one level at a time (L1/L2/L3 buttons); + *Stacked (all + link)* renders all three top-to-bottom with section headings. +- **Drill-down linkage** (works in both layouts, most visible when stacked): + - click a **cr** in Level 1 (or a layer cell in Level 2) → highlights that cr's + layers in Level 2 and dims the rest; Level 1 emphasises the matching row. + - click a **PP rank** in Level 2 (or a device row in Level 3) → highlights the + linked rank in Levels 2 and 3, and Level 1 highlights the cr types present on + that rank. + - a "Clear" bar removes the selection. +- **Export.** Level 3's Gantt has *Export SVG* / *Export PNG* buttons; the + serializer resolves the theme CSS variables and paints a background so the file + is self-contained (good for slides / the paper-style figure). +- **Fit + zoom.** Level 3 fits the whole schedule in the panel at 1× (no + scrollbar) for an at-a-glance overview. A zoom slider (1×–10×) widens the time + axis so the per-cell microbatch numbers become readable; above 1× the chart + scrolls horizontally. Zoom stretches only the time axis (font/row height fixed). +- **Cell tooltip.** Hovering a Gantt cell shows `compute µs` (the op's + duration) and `starts @ ms` (its start time measured from the iteration + start). The x-axis is that same wall-clock time; gaps between cells are bubble. + +## Consistency / self-checks + +- Level 1 category sums per cr == `layerTimes(cr).fwd/bwd` (no time lost in + categorization). +- Level 2 per-device `Df/Db` == `project().Df/Db` (same mapping, just detailed). +- Level 3 simulated iteration length ≈ analytic `pipe_compute` (balanced case); + otherwise warn. +- All three levels react live to the shared projection controls (GPU tab, PP/VPP, + GA/GBS/MBS/DP, recompute, manual layer-timing mode). + +## Scope / caveats (inherit from 02-assumptions) + +- `a2a` is captured at EP=8 intra-node; other EP values are not re-modeled (A7-A9). +- p2p PP comm is hidden; only the bubble is shown (A4). +- seq is fixed at the capture value (4096). +- MI455X tab rescales module times per Step 6 before all three levels are drawn. diff --git a/examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh b/examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh new file mode 100755 index 000000000..9a3d16a29 --- /dev/null +++ b/examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh @@ -0,0 +1,252 @@ +#!/bin/bash +############################################################################### +# DeepSeek-V4 single-layer, single-cr profiling launcher for the perf +# *projection* pipeline (examples/deepseek-v4/projection). +# +# Produces ONE chrome trace for ONE compression-ratio (cr) layer type, captured +# under production-representative conditions so the trace can be turned into a +# clean per-module forward/backward breakdown (see design/01-overview.md): +# +# * seq_length = 4096 (production per-microbatch token count) +# * num_layers = 1 (clean per-layer attribution; fits memory @4096) +# * compress_ratios = [CR] (one cr per trace; CR in {0,4,128}) +# * optimizer = adam, NON-distributed, fp32 states (dist-opt ON makes the +# ROCm Kineto profiler +# drop dense/HCA compute +# kernels; NOT muon) +# * overlap_grad_reduce/param_gather = False (num_layers=1 breaks +# Megatron's chained +# param-sync assert; and +# with overlap off every +# microbatch's compute is +# already clean) +# * GA = 2 => GBS = 2 * DP * MBS (parser averages the +# steady profiler window +# per microbatch) +# * recompute = OFF (capture pure fwd / pure bwd) +# * profiler ON, with_stack=True, window iter 6 -> 7 (kernel -> nn.module) +# +# Usage (inside the training container, repo root): +# CR=0 ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh +# CR=4 ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh +# CR=128 ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh +# MODEL=flash CR=4 ./examples/deepseek-v4/projection/script/deepseek_v4_layer_trace-projection.sh +# +# The trace lands under: +# output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/tensorboard/*.pt.trace.json +############################################################################### +set -euo pipefail +set -x + +export HF_TOKEN="${HF_TOKEN:-}" + +# ---------- What to profile ------------------------------------------------- +export MODEL=${MODEL:-pro} # pro | flash +export CR=${CR:-4} # 0 | 4 | 128 (single cr per trace) + +# ---------- Model: DeepSeek-V4 (pro or flash) ------------------------------ +# Widths (hidden/heads/kv) come from the model yaml via PRIMUS_MODEL; we only +# override the MoE / indexer shape knobs the runner exposes, per variant. +export EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml} +case "$MODEL" in + pro) + export PRIMUS_MODEL=${PRIMUS_MODEL:-deepseek_v4_pro} + export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-384} + export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6} + export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-3072} + export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-1024} + ;; + flash) + export PRIMUS_MODEL=${PRIMUS_MODEL:-deepseek_v4_flash} + export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-256} + export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6} + export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-2048} + export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-512} + ;; + *) + echo "[ERROR] MODEL must be 'pro' or 'flash', got '$MODEL'"; exit 1 ;; +esac +# Megatron aux-loss-free expert bias needs sigmoid; V4 uses sqrtsoftplus -> off. +export PRIMUS_MOE_ENABLE_EXPERT_BIAS=${PRIMUS_MOE_ENABLE_EXPERT_BIAS:-False} + +# ---------- Single layer per cr (or 3-layer mix) at production seq --------- +# Pure dense (cr=0) / HCA (cr=128) single layers get CUDA-graph/stream-captured +# (compute hidden from the trace). CR=mix runs a 3-layer [0,4,128] block: the +# dynamic CSA (cr=4) layer keeps the block out of graph capture, so all three +# attention types are visible in one trace and split by kernel name. +export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-4096} +export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}} +case "$CR" in + 0|4|128) export PRIMUS_COMPRESS_RATIOS="[$CR]"; export PRIMUS_TOTAL_LAYERS=1 ;; + mix) export PRIMUS_COMPRESS_RATIOS="[0,4,128]"; export PRIMUS_TOTAL_LAYERS=3 ;; + *) echo "[ERROR] CR must be 0, 4, 128 or mix, got '$CR'"; exit 1 ;; +esac + +# ---------- Single-node EP=8 (intra-node; DP=8) ---------------------------- +export PRIMUS_TP=${PRIMUS_TP:-1} +export PRIMUS_PP=${PRIMUS_PP:-1} +export PRIMUS_EP=${PRIMUS_EP:-8} +export MBS=${MBS:-1} +# DP = world/(TP*PP). On one 8-GPU node with TP=PP=1 => DP=8. +export DP=${DP:-8} +# GA = 2 so the schedule is F1 B1 F2 B2: the clean (overlap-free) forward is F2 +# and the clean backward is B1. The parser averages the steady profiler window +# per microbatch, keeping scalar control-flow stalls that survive full-model +# calibration. +export GBS=${GBS:-$((2 * DP * MBS))} + +# ---------- Optimizer: adam + distributed optimizer (zero1) ---------------- +# Compute (fwd/bwd) is optimizer-independent; we use adam (a) to model the +# zero1 DP-comm overlap that GA=2 isolates, and (b) to dodge Muon's fp32 +# optimizer-state memory blow-up so seq=4096 fits. The optimizer step itself is +# modeled separately in the projection (design/04-projection-math.md Step 4). +export OPTIMIZER=${OPTIMIZER:-adam} +# CRITICAL: distributed optimizer (zero1) MUST be off. With it on, the ROCm +# Kineto profiler silently drops the compute GPU kernels for pure dense (cr=0) +# and HCA (cr=128) layers (only optimizer/comm/elementwise get recorded); +# turning it off makes every cr's kernels visible. dist-opt has no bearing on +# the captured fwd/bwd compute, which is what the projection needs (the +# optimizer step is modeled analytically, design/04 Step 4). +export USE_DISTRIBUTED_OPTIMIZER=${USE_DISTRIBUTED_OPTIMIZER:-False} +# Overlap OFF. With num_layers=1, Megatron's chained param-gather sync trips an +# assertion (param_and_grad_buffer.start_param_sync). More importantly, with +# overlap off every microbatch's compute is already overlap-free — exactly the +# clean per-kernel time we want; the parser averages the steady profiler window +# per microbatch. (Production DP comm is assumed hidden in the projection anyway; A2.) +export PRIMUS_OVERLAP_GRAD_REDUCE=${PRIMUS_OVERLAP_GRAD_REDUCE:-False} +export PRIMUS_OVERLAP_PARAM_GATHER=${PRIMUS_OVERLAP_PARAM_GATHER:-False} +# Distributed optimizer (zero1) is the default. The Kineto trace drops the +# compute GPU kernels for pure dense(cr=0)/HCA(cr=128) layers ONLY when the +# distributed optimizer is on; turning it off makes all kernels visible (but +# then precision-aware optimizer must be off too -> fp32 optimizer states). +DISTOPT_ARGS=(--use_distributed_optimizer "$USE_DISTRIBUTED_OPTIMIZER") +if [ "$USE_DISTRIBUTED_OPTIMIZER" = "False" ]; then + DISTOPT_ARGS+=(--use_precision_aware_optimizer False --main_grads_dtype fp32 --exp_avg_dtype fp32 --exp_avg_sq_dtype fp32) +fi + +# ---------- Perf knobs (production V4 attn backends + Turbo MoE) ------------ +export ENABLE_PRIMUS_TURBO=${ENABLE_PRIMUS_TURBO:-True} +export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False} +export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-True} +export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-True} +export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-triton_v1} +export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v1} +export USE_V4_FP8_INDEXER=${USE_V4_FP8_INDEXER:-True} +export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-True} +export PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU=${PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU:-True} +export PRIMUS_V4_ATTN_BWD_USE_SPLIT=${PRIMUS_V4_ATTN_BWD_USE_SPLIT:-1} +export PRIMUS_V4_CSA_BWD_SEGREDUCE=${PRIMUS_V4_CSA_BWD_SEGREDUCE:-1} +export PRIMUS_STACK_GROUPED_WEIGHT_TRITON=${PRIMUS_STACK_GROUPED_WEIGHT_TRITON:-1} +export PRIMUS_ROPE_TRITON=${PRIMUS_ROPE_TRITON:-1} +export PRIMUS_SINKHORN_TRITON=${PRIMUS_SINKHORN_TRITON:-1} +export PRIMUS_HC_TRITON=${PRIMUS_HC_TRITON:-1} +export PRIMUS_INDEXER_TRITON=${PRIMUS_INDEXER_TRITON:-1} +export PRIMUS_INDEXER_TRITON_FULL=${PRIMUS_INDEXER_TRITON_FULL:-0} +export PRIMUS_V4_ROUTER_TRITON=${PRIMUS_V4_ROUTER_TRITON:-1} + +TURBO_DEEPEP_CLI_ARGS=() +if [ "$USE_TURBO_DEEPEP" = "True" ]; then + export TURBO_DEEPEP_NUM_CU=${TURBO_DEEPEP_NUM_CU:-80} + export TURBO_DEEPEP_USE_COMM_STREAM=${TURBO_DEEPEP_USE_COMM_STREAM:-False} + export MOE_ROUTER_DTYPE=${MOE_ROUTER_DTYPE:-fp32} + export MOE_SHARED_EXPERT_OVERLAP=${MOE_SHARED_EXPERT_OVERLAP:-False} + TURBO_DEEPEP_CLI_ARGS=( + --turbo_deepep_num_cu "$TURBO_DEEPEP_NUM_CU" + --turbo_deepep_use_comm_stream "$TURBO_DEEPEP_USE_COMM_STREAM" + --moe_router_dtype "$MOE_ROUTER_DTYPE" + --moe_shared_expert_overlap "$MOE_SHARED_EXPERT_OVERLAP" + ) +fi + +export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-1} +# Disable the loss NaN/Inf validation (check_for_nan_in_loss_and_grad). Its +# torch.isnan/isinf checks in loss_func force a device->host sync once per +# microbatch; in a 1-layer capture (nothing to overlap) the profiler bills that +# sync wait as a multi-ms "stall" kernel under the loss/rerun frames, which then +# pollutes per-layer attribution and gets multiplied by the layer count in the +# projection. It is a per-step validation cost, not per-layer compute, so we +# turn it off for a clean capture (the projection models it separately if needed). +# Profiler window must land in the STEADY state: the first ~9 iters are warm-up +# (kernel autotune / hipBLASLt + Triton compilation) and have noisy, inflated +# per-iter times (and inflated comm-kernel durations as ranks desync on the +# autotuning rank); from ~iter 10 onward the iteration time is stable. Default +# to a late, multi-step window so the captured steps are clean; all three are +# env-overridable. +export TRAIN_ITERS=${TRAIN_ITERS:-22} +export PROFILE_STEP_START=${PROFILE_STEP_START:-16} +export PROFILE_STEP_END=${PROFILE_STEP_END:-19} + +# ---------- Profiler: trace with python stacks for kernel->module ---------- +export PROFILE=True +# Optional GPU-only trace (drop CPU activity). With CPU activity on, pure dense +# (cr=0) / HCA (cr=128) layers' compute kernels can vanish from the trace +# (profiler/stream-capture interaction); GPU-only capture brings them back, at +# the cost of CPU-side module/stack attribution (fall back to kernel-name rules). +PROFILER_ACTIVITY_ARGS=() +if [ "${DISABLE_PROFILER_CPU:-False}" = "True" ]; then + PROFILER_ACTIVITY_ARGS=(--disable_profiler_activity_cpu True) +fi +export BACKEND_PATH=${BACKEND_PATH:-"$(pwd)/third_party/Megatron-LM"} +export PRIMUS_TEAM=${PRIMUS_TEAM:-amd} +export PRIMUS_USER=${PRIMUS_USER:-tas-mi355x-$(date +%Y%m%d)} +export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-projection_${MODEL}_cr${CR}_seq${PRIMUS_SEQ_LENGTH}_ep${PRIMUS_EP}} + +if [ ! -d "$BACKEND_PATH" ] || [ -z "$(ls -A "$BACKEND_PATH" 2>/dev/null)" ]; then + echo "[ERROR] BACKEND_PATH does not exist or is empty: $BACKEND_PATH" + echo "Run: git submodule update --init --recursive" + exit 1 +fi + +mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" + +./primus-cli direct \ + -- train pretrain --config "$EXP" \ + --backend_path "$BACKEND_PATH" \ + --num_layers "$PRIMUS_TOTAL_LAYERS" \ + --train_iters "$TRAIN_ITERS" \ + --lr_warmup_iters 0 \ + --lr_decay_iters "$TRAIN_ITERS" \ + --micro_batch_size "$MBS" \ + --global_batch_size "$GBS" \ + --seq_length "$PRIMUS_SEQ_LENGTH" \ + --max_position_embeddings "$PRIMUS_MAX_POSITION_EMBEDDINGS" \ + --rope_type rope \ + --tensor_model_parallel_size "$PRIMUS_TP" \ + --pipeline_model_parallel_size "$PRIMUS_PP" \ + --expert_model_parallel_size "$PRIMUS_EP" \ + --num_experts "$PRIMUS_NUM_EXPERTS" \ + --moe_router_topk "$PRIMUS_MOE_TOPK" \ + --moe_router_enable_expert_bias "$PRIMUS_MOE_ENABLE_EXPERT_BIAS" \ + --moe_ffn_hidden_size "$PRIMUS_MOE_FFN_HIDDEN_SIZE" \ + --index_topk "$PRIMUS_INDEX_TOPK" \ + --v4_grouped_experts_support_clamped_swiglu "$PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU" \ + --compress_ratios "$PRIMUS_COMPRESS_RATIOS" \ + --mtp_num_layers 0 \ + --mock_data True \ + --optimizer "$OPTIMIZER" \ + "${DISTOPT_ARGS[@]}" \ + --enable_primus_turbo "$ENABLE_PRIMUS_TURBO" \ + --use_turbo_attention "$USE_TURBO_ATTENTION" \ + --use_v4_attention_backend "$USE_V4_ATTENTION_BACKEND" \ + --use_v4_csa_attention_backend "$USE_V4_CSA_ATTENTION_BACKEND" \ + --use_v4_fp8_indexer "$USE_V4_FP8_INDEXER" \ + --use_v4_compiled_sinkhorn "$USE_V4_COMPILED_SINKHORN" \ + --use_turbo_deepep "$USE_TURBO_DEEPEP" \ + "${TURBO_DEEPEP_CLI_ARGS[@]}" \ + --use_turbo_grouped_gemm "$TURBO_USE_GROUPED_MLP" \ + --moe_use_legacy_grouped_gemm False \ + --fp8 null \ + --fp8_recipe null \ + --recompute_num_layers 0 \ + --check_for_nan_in_loss_and_grad False \ + --overlap_grad_reduce "$PRIMUS_OVERLAP_GRAD_REDUCE" \ + --overlap_param_gather "$PRIMUS_OVERLAP_PARAM_GATHER" \ + --disable_last_saving True \ + --disable_wandb True \ + --disable_tensorboard False \ + --profile True \ + --use_pytorch_profiler True \ + "${PROFILER_ACTIVITY_ARGS[@]}" \ + --profile_step_start "$PROFILE_STEP_START" \ + --profile_step_end "$PROFILE_STEP_END" \ + 2>&1 | tee "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/log_node_${NODE_RANK:-0}.txt" diff --git a/examples/deepseek-v4/projection/site/assets/app.js b/examples/deepseek-v4/projection/site/assets/app.js new file mode 100644 index 000000000..b1b0b32bd --- /dev/null +++ b/examples/deepseek-v4/projection/site/assets/app.js @@ -0,0 +1,1567 @@ +"use strict"; + +// DeepSeek-V4 performance projection — static, no-build. Implements the math in +// design/04-projection-math.md. All breakdown times are microseconds (us) for +// one microbatch (seq from capture); the projection scales to a full model run. + +const STATE = { + data: null, + gpu: "MI355X", + controls: null, + // iteration-timeline view (design/07): active level (1|2|3). Level 3 + // interleaving follows the VPP control directly (no separate toggle). + tlLevel: 1, + // layout: "tabs" (one level via buttons) or "stacked" (all three + linkage). + tlView: "stacked", + // Level 3 Gantt horizontal zoom (1 = fit-to-view, no scrollbar; >1 widens the + // chart and reveals per-cell microbatch numbers, with horizontal scroll). + tlZoom: 1, + // cross-level selection for drill-down highlight. cr: focus a compression + // ratio; devices: focus one or more PP ranks. Mutually exclusive. + tlSel: { cr: null, devices: null }, +}; + +const $ = (sel) => document.querySelector(sel); +const el = (tag, attrs = {}, ...kids) => { + const n = document.createElement(tag); + for (const [k, v] of Object.entries(attrs)) { + if (k === "class") n.className = v; + else if (k === "html") n.innerHTML = v; + else n.setAttribute(k, v); + } + for (const kid of kids) n.append(kid?.nodeType ? kid : document.createTextNode(kid ?? "")); + return n; +}; +const fmt = (x, d = 1) => + x == null || !isFinite(x) ? "—" : Number(x).toLocaleString(undefined, { maximumFractionDigits: d, minimumFractionDigits: d }); +const fmtInt = (x) => (x == null || !isFinite(x) ? "—" : Math.round(x).toLocaleString()); + +// --------------------------------------------------------------------------- +// Load +// --------------------------------------------------------------------------- +function modelFromQuery() { + const m = new URLSearchParams(location.search).get("model"); + return m === "flash" || m === "pro" ? m : "pro"; +} + +async function loadModel(model) { + const res = await fetch(`./data/${model}.json`, { cache: "no-store" }); + if (!res.ok) throw new Error(`failed to load data/${model}.json (${res.status})`); + return res.json(); +} + +// --------------------------------------------------------------------------- +// Controls +// --------------------------------------------------------------------------- +function defaultControls(data) { + const hw = data.hardware || {}; + const m355 = hw.MI355X || { peak_tflops_bf16: 2500, hbm_bandwidth_gbps: 8000 }; + const m455 = hw.MI455X || { peak_tflops_bf16: 5000, hbm_bandwidth_gbps: 16000 }; + const isPro = data.model === "pro"; + const optBytes = data.optimizer?.bytes_per_param && data.optimizer.bytes_per_param !== 18 + ? data.optimizer.bytes_per_param : 30; + return { + world: isPro ? 256 : 32, + pp: isPro ? 16 : 4, vpp: 1, ep: 8, tp: 1, cp: 1, + mbs: 1, gbs: isPro ? 1024 : 256, + recompute: isPro ? "full" : "first-n", + recomputeLayers: isPro ? 0 : 3, + ppLayout: data.model_config?.pipeline_layout || "", + optEff: 0.7, computeEff: 1.0, calibFactor: 0.91, bytesPerParam: optBytes, + peak355: m355.peak_tflops_bf16, bw355: m355.hbm_bandwidth_gbps, + peak455: m455.peak_tflops_bf16, bw455: m455.hbm_bandwidth_gbps, + // Modeling mode: "trace" (derive per-layer fwd/bwd from the breakdown JSON) + // or "manual" (user types per-cr fwd/bwd directly). Manual values are stored + // per GPU because a hand-entered time already targets a specific GPU, so the + // MI355->MI455 scaling is bypassed in manual mode. + modelMode: "trace", + man: { + MI355X: emptyManual(), + MI455X: emptyManual(), + }, + }; +} + +// Per-cr manual fwd/bwd holders (µs). null = "not set yet" -> falls back to the +// trace-derived value, so toggling into manual mode never changes the result +// until the user actually edits a field. +function emptyManual() { + return { + f0: null, b0: null, f4: null, b4: null, f128: null, b128: null, + // non-layer + MTP overrides (per iteration, one device) + emb_f: null, emb_b: null, out_f: null, out_b: null, loss_f: null, loss_b: null, mtp_f: null, mtp_b: null, + }; +} + +const MANUAL_CR_KEYS = { "0": ["f0", "b0"], "4": ["f4", "b4"], "128": ["f128", "b128"] }; +// Manually-overridable non-layer parts: key prefix, label, and the JSON section +// (mtp is synthesized, not a non_layer entry). +const MANUAL_NONLAYER = [ + { key: "emb", label: "embedding", which: "embedding" }, + { key: "out", label: "output", which: "output" }, + { key: "loss", label: "loss", which: "loss" }, + { key: "mtp", label: "MTP", which: "mtp" }, +]; +const MANUAL_NONLAYER_KEYS = MANUAL_NONLAYER.flatMap((n) => [`${n.key}_f`, `${n.key}_b`]); + +const CONTROL_DEFS = [ + { key: "world", label: "World size (GPUs)", kind: "int" }, + { key: "pp", label: "PP (pipeline)", kind: "int" }, + { key: "vpp", label: "VPP (interleave)", kind: "int" }, + { key: "ep", label: "EP (expert)", kind: "int" }, + { key: "dp", label: "DP (derived)", kind: "ro" }, + { key: "tp", label: "TP (tensor)", kind: "int" }, + { key: "cp", label: "CP (context)", kind: "int" }, + { key: "mbs", label: "Micro batch size", kind: "int" }, + { key: "gbs", label: "Global batch size", kind: "int" }, + { key: "recompute", label: "Recompute", kind: "sel" }, + { key: "recomputeLayers", label: "Recompute layers", kind: "int" }, + { key: "ppLayout", label: "PP layout", kind: "txt", full: true }, + { key: "bytesPerParam", label: "Optim bytes/param", kind: "int" }, + { key: "calibFactor", label: "Calibration factor", kind: "f" }, + { key: "optEff", label: "Optim efficiency", kind: "f" }, + { key: "computeEff", label: "MI455 compute eff", kind: "f" }, + { key: "peak355", label: "MI355 peak TFLOPs", kind: "int" }, + { key: "bw355", label: "MI355 HBM GB/s", kind: "int" }, + { key: "peak455", label: "MI455 peak TFLOPs", kind: "int" }, + { key: "bw455", label: "MI455 HBM GB/s", kind: "int" }, +]; + +// DP is derived from the user-set world size: DP = world / (PP*TP*CP). EP is a +// sub-grouping of DP (EP <= DP) and does not multiply world size. +const derivedDP = (c) => { + const denom = c.pp * c.tp * c.cp; + if (!Number.isFinite(denom) || denom <= 0 || !Number.isFinite(c.world)) return NaN; + return c.world / denom; +}; + +function parseControlValue(input, kind) { + if (kind === "sel" || kind === "txt") return input.value; + if (kind === "int") return input.value.trim() === "" ? NaN : Number(input.value); + if (kind === "f") return input.value.trim() === "" ? NaN : Number(input.value); + return input.value; +} + +const isPositiveInt = (x) => Number.isInteger(x) && x > 0; +const isNonNegativeInt = (x) => Number.isInteger(x) && x >= 0; +const isPositiveNumber = (x) => Number.isFinite(x) && x > 0; +const isNonNegativeNumber = (x) => Number.isFinite(x) && x >= 0; + +function renderControls() { + const grid = $("#controls-grid"); + grid.innerHTML = ""; + const c = STATE.controls; + for (const def of CONTROL_DEFS) { + const { key, label, kind } = def; + const field = el("div", { class: "field" }); + if (def.full) field.classList.add("field--full"); + field.append(el("span", {}, label)); + let input; + if (kind === "sel") { + input = el("select", { id: `ctl-${key}` }); + for (const opt of ["none", "full", "first-n"]) { + const o = el("option", { value: opt }, opt); + if (c[key] === opt) o.selected = true; + input.append(o); + } + } else if (kind === "ro") { + const dp = derivedDP(c); + input = el("input", { id: `ctl-${key}`, value: Number.isFinite(dp) ? dp : "—", disabled: "true" }); + } else if (kind === "txt") { + input = el("input", { id: `ctl-${key}`, type: "text", value: c[key] || "" }); + } else { + input = el("input", { id: `ctl-${key}`, type: "number", value: c[key], step: kind === "f" ? "0.05" : "1" }); + } + if (kind !== "ro") { + input.addEventListener("change", () => { + c[key] = parseControlValue(input, kind); + renderAll(); + }); + } + field.append(input); + grid.append(field); + } +} + +// Prefill the active GPU's manual fields from the trace-derived times so that +// switching into manual mode (or switching GPU while in manual mode) starts from +// the current baseline instead of empty boxes. Already-set fields are kept. +function prefillManual(gpu) { + const c = STATE.controls; + const m = c.man[gpu]; + const lt = {}; + for (const cr of ["0", "4", "128"]) { + const [fk, bk] = MANUAL_CR_KEYS[cr]; + const t = layerTimes(STATE.data, cr, gpu, c); + lt[cr] = t; + if (!isPositiveNumber(m[fk])) m[fk] = Math.round(t.fwd); + if (!isPositiveNumber(m[bk])) m[bk] = Math.round(t.bwd); + } + // non-layer parts (embedding / output / loss). Skip seeding fields whose trace + // value is 0 (e.g. output/loss backward) — leaving them unset shows the "0" + // placeholder and falls back to trace, instead of pinning an explicit 0. + const seed = (key, val) => { if (!isPositiveNumber(m[key]) && Math.round(val) > 0) m[key] = Math.round(val); }; + for (const which of ["embedding", "output", "loss"]) { + const key = NONLAYER_KEY[which]; + seed(`${key}_f`, nonLayer(STATE.data, which, "forward", gpu, c)); + seed(`${key}_b`, nonLayer(STATE.data, which, "backward", gpu, c)); + } + // MTP (only when the model uses it) + if ((STATE.data.model_config.mtp_num_layers || 0) > 0) { + const base = mtpTimes(STATE.data, gpu, c, lt); + seed("mtp_f", base.fwd); + seed("mtp_b", base.bwd); + } +} + +function renderModeSwitch() { + document.querySelectorAll(".mode-tab").forEach((t) => { + t.classList.toggle("is-active", t.dataset.mode === STATE.controls.modelMode); + }); +} + +function renderManualGrid() { + const host = $("#manual-grid"); + if (!host) return; + host.hidden = STATE.controls.modelMode !== "manual"; + host.innerHTML = ""; + if (host.hidden) return; + const c = STATE.controls, gpu = STATE.gpu; + const counts = STATE.data.model_config.cr_layer_counts || {}; + const m = c.man[gpu]; + const header = el("div", { class: "manual-grid__header" }); + header.append(el("p", { class: "muted manual-grid__hint" }, + `Per-layer and non-layer fwd/bwd (µs) for ${gpu}. Set values override the trace; placeholders show the current trace baseline. Embedding is on PP stage 0; output / loss / MTP are on the last PP stage.`)); + const resetBtn = el("button", { class: "manual-reset" }, "Restore defaults"); + resetBtn.addEventListener("click", () => { + c.man[gpu] = emptyManual(); + prefillManual(gpu); + renderAll(); + }); + header.append(resetBtn); + host.append(header); + + // one fwd/bwd input row + const makeRow = (labelNode, fk, bk, traceF, traceB) => { + const row = el("div", { class: "manual-row" }); + row.append(labelNode); + for (const [label, key, traceVal] of [["fwd", fk, traceF], ["bwd", bk, traceB]]) { + const field = el("div", { class: "field" }); + field.append(el("span", {}, `${label} µs`)); + const input = el("input", { + id: `man-${gpu}-${key}`, type: "number", step: "1", min: "0", + placeholder: String(Math.round(traceVal || 0)), + }); + if (isPositiveNumber(m[key])) input.value = m[key]; + input.addEventListener("change", () => { + const v = input.value.trim(); + m[key] = v === "" ? null : Number(v); + renderAll(); + }); + field.append(input); + row.append(field); + } + return row; + }; + + const rows = el("div", { class: "manual-rows" }); + for (const cr of ["0", "4", "128"]) { + if (!(counts[cr] > 0)) continue; + const [fk, bk] = MANUAL_CR_KEYS[cr]; + const t = layerTimes(STATE.data, cr, gpu, c); + const lab = el("span", { class: `manual-row__lab cr-tag cr-${cr}` }, `cr=${cr} ×${counts[cr] || 0}`); + rows.append(makeRow(lab, fk, bk, t.fwd, t.bwd)); + } + host.append(rows); + + // non-layer + MTP rows + host.append(el("p", { class: "muted manual-grid__subhead" }, "Non-layer (once per iteration)")); + const nlRows = el("div", { class: "manual-rows" }); + const lt = {}; + for (const cr of ["0", "4", "128"]) lt[cr] = layerTimes(STATE.data, cr, gpu, c); + for (const nl of MANUAL_NONLAYER) { + if (nl.which === "mtp" && !((STATE.data.model_config.mtp_num_layers || 0) > 0)) continue; + let tF, tB; + if (nl.which === "mtp") { + const base = mtpTimes(STATE.data, gpu, c, lt); + tF = base.fwd; tB = base.bwd; + } else { + tF = nonLayer(STATE.data, nl.which, "forward", gpu, c); + tB = nonLayer(STATE.data, nl.which, "backward", gpu, c); + } + const lab = el("span", { class: "manual-row__lab manual-row__lab--nl" }, nl.label); + nlRows.append(makeRow(lab, `${nl.key}_f`, `${nl.key}_b`, tF, tB)); + } + host.append(nlRows); +} + +// --------------------------------------------------------------------------- +// Hardware scaling (Step 6) +// --------------------------------------------------------------------------- +function rowScaledTime(row, gpu, c) { + if (gpu === "MI355X") return row.time_us; + const compute = row.class === "compute_bound"; + if (compute) return row.time_us * (c.peak355 / c.peak455) / c.computeEff; + return row.time_us * (c.bw355 / c.bw455); +} +function rowTflops(row, scaledTimeUs) { + if (!row.flops || !scaledTimeUs) return null; + return row.flops / (scaledTimeUs * 1e-6) / 1e12; +} + +const sumTime = (list, gpu, c) => list.reduce((a, r) => a + rowScaledTime(r, gpu, c), 0); +const sumFlops = (list) => list.reduce((a, r) => a + (r.flops || 0), 0); + +// --------------------------------------------------------------------------- +// Per-layer base (Step 0/1) +// --------------------------------------------------------------------------- +function layerTimes(data, cr, gpu, c) { + const L = data.layers[cr]; + if (!L) return { fwd: 0, bwd: 0, fFlops: 0, bFlops: 0 }; + const aF = L.attention.forward, aB = L.attention.backward; + const mF = L.moe.forward, mB = L.moe.backward; + let fwd = sumTime(aF, gpu, c) + sumTime(mF, gpu, c); + let bwd = sumTime(aB, gpu, c) + sumTime(mB, gpu, c); + let fFlops = sumFlops(aF) + sumFlops(mF); + let bFlops = sumFlops(aB) + sumFlops(mB); + return { fwd, bwd, fFlops, bFlops }; +} + +// Effective per-layer time used by the projection. In "manual" mode a set field +// overrides the trace-derived time for the active GPU; unset fields fall back to +// trace. FLOPs always stay analytic/trace-derived (manual only overrides time), +// so TFLOP/s/GPU remains meaningful. +function effectiveLayerTimes(data, cr, gpu, c) { + const trace = layerTimes(data, cr, gpu, c); + if (c.modelMode !== "manual") return trace; + const m = (c.man && c.man[gpu]) || {}; + const [fk, bk] = MANUAL_CR_KEYS[cr] || []; + return { + fwd: isPositiveNumber(m[fk]) ? m[fk] : trace.fwd, + bwd: isPositiveNumber(m[bk]) ? m[bk] : trace.bwd, + fFlops: trace.fFlops, + bFlops: trace.bFlops, + }; +} + +function expandLayoutRepeats(layout) { + let out = layout; + let prev; + do { + prev = out; + out = out.replace(/\(([^()]+)\)\*(\d+)/g, (_m, body, count) => body.repeat(Number(count))); + } while (out !== prev); + return out; +} + +function parsePipelineLayout(layout, numLayers, chunks) { + const raw = String(layout ?? "").trim(); + if (!raw) return { ok: true, stages: null, counts: [], normalized: "", message: "empty layout; using balanced fallback" }; + const normalized = expandLayoutRepeats(raw.replace(/^['"]|['"]$/g, "")); + if (/[()]/.test(normalized)) { + return { ok: false, stages: null, counts: [], normalized, message: "unsupported nested or malformed repeat expression" }; + } + const specs = normalized.split("|").map((x) => x.trim()).filter(Boolean); + if (specs.length !== chunks) { + return { + ok: false, + stages: null, + counts: specs.map((spec) => [...spec.matchAll(/[tT](?:\*(\d+))?/g)].reduce((a, m) => a + Number(m[1] || 1), 0)), + normalized, + message: `layout has ${specs.length} stages, expected PP*VPP=${chunks}`, + }; + } + let nextLayer = 0; + const out = []; + const counts = []; + for (const spec of specs) { + const layers = []; + for (const m of spec.matchAll(/[tT](?:\*(\d+))?/g)) { + const n = m[1] ? Number(m[1]) : 1; + for (let i = 0; i < n && nextLayer < numLayers; i++) layers.push(nextLayer++); + } + counts.push(layers.length); + out.push(layers); + } + if (nextLayer !== numLayers) { + return { + ok: false, + stages: null, + counts, + normalized, + message: `layout maps ${nextLayer} decoder layers, expected ${numLayers}`, + }; + } + return { ok: true, stages: out, counts, normalized, message: "layout applied" }; +} + +function recomputeLayer(c, stageOrdinal) { + if (c.recompute === "full") return true; + if (c.recompute === "first-n") return stageOrdinal < Math.max(0, c.recomputeLayers || 0); + return false; +} + +function validateControls(data, c, gpu = STATE.gpu) { + const errors = []; + const warnings = []; + const ints = ["world", "pp", "vpp", "ep", "tp", "cp", "mbs", "gbs", "bytesPerParam", "peak355", "bw355", "peak455", "bw455"]; + for (const key of ints) { + if (!isPositiveInt(c[key])) errors.push(`${key} must be a positive integer.`); + } + if (!isNonNegativeInt(c.recomputeLayers)) errors.push("recomputeLayers must be a non-negative integer."); + for (const key of ["calibFactor", "optEff", "computeEff"]) { + if (!isPositiveNumber(c[key])) errors.push(`${key} must be a positive number.`); + } + if (!["none", "full", "first-n"].includes(c.recompute)) errors.push(`unsupported recompute mode: ${c.recompute}`); + + const denom = c.pp * c.tp * c.cp; + const dp = derivedDP(c); + if (isPositiveInt(c.world) && isPositiveInt(denom) && c.world % denom !== 0) { + errors.push(`world must be divisible by PP*TP*CP (${denom}); got world=${c.world}.`); + } + if (Number.isFinite(dp) && isPositiveInt(c.gbs) && isPositiveInt(c.mbs) && !Number.isInteger(c.gbs / (dp * c.mbs))) { + errors.push(`GA must be an integer: GBS / (DP*MBS) = ${c.gbs} / (${dp}*${c.mbs}).`); + } + if (Number.isFinite(dp) && isPositiveInt(c.ep) && c.ep > dp) { + errors.push(`EP must be <= DP; got EP=${c.ep}, DP=${dp}.`); + } + + const chunks = c.pp * c.vpp; + const layout = parsePipelineLayout(c.ppLayout, data.model_config.compress_ratios.length, chunks); + if (!layout.ok) errors.push(`PP layout invalid: ${layout.message}.`); + + if (c.modelMode === "manual") { + const man = (c.man && c.man[gpu]) || {}; + for (const cr of ["0", "4", "128"]) { + const count = data.model_config.cr_layer_counts?.[cr] || 0; + if (!count) continue; // cr not used by this model; ignore its inputs + for (const k of MANUAL_CR_KEYS[cr]) { + const v = man[k]; + if (v != null && v !== "" && !isNonNegativeNumber(Number(v))) { + errors.push(`manual ${k} (cr=${cr}) must be a non-negative number (µs).`); + } + } + } + for (const k of MANUAL_NONLAYER_KEYS) { + const v = man[k]; + if (v != null && v !== "" && !isNonNegativeNumber(Number(v))) { + errors.push(`manual ${k} must be a non-negative number (µs).`); + } + } + warnings.push(`Manual mode for ${gpu}: per-cr layer, non-layer (embedding/output/loss) and MTP fwd/bwd you set override the trace; unset fields fall back to trace.`); + } + + const captureEp = data.capture?.ep; + if (captureEp && c.ep !== captureEp) { + warnings.push(`EP=${c.ep} is only partially modeled; traces captured EP=${captureEp}, so MoE dispatch/combine are not re-estimated.`); + } else { + warnings.push(`EP is a consistency control only; captured MoE dispatch/combine are reused from EP=${captureEp || 8}.`); + } + if (c.tp !== 1 || c.cp !== 1) { + warnings.push("TP/CP values affect derived DP/GA/optimizer only; TP/CP layer compute and communication are not re-modeled."); + } + if (gpu === "MI355X") warnings.push("MI455 peak/bandwidth/compute-eff controls affect only the MI455X tab."); + warnings.push(`Sequence length is fixed at captured seq=${data.capture.seq_length}; changing seq is not currently exposed.`); + + return { errors, warnings, layout, dp }; +} + +function nonLayer(data, which, phase, gpu, c) { + const bd = data.non_layer[which]; + return bd ? sumTime(bd[phase], gpu, c) : 0; +} + +// Non-layer time with manual override (embedding / output / loss). Falls back to +// the trace time when the field is unset or not in manual mode. +const NONLAYER_KEY = { embedding: "emb", output: "out", loss: "loss" }; +function effectiveNonLayer(data, which, phase, gpu, c) { + const trace = nonLayer(data, which, phase, gpu, c); + if (c.modelMode !== "manual") return trace; + const m = (c.man && c.man[gpu]) || {}; + const v = m[`${NONLAYER_KEY[which]}_${phase === "forward" ? "f" : "b"}`]; + return isPositiveNumber(v) ? v : trace; +} +function nonLayerFlops(data, which, phase) { + const bd = data.non_layer[which]; + return bd ? sumFlops(bd[phase]) : 0; +} + +function mtpTimes(data, gpu, c, lt) { + const cfg = data.model_config; + const depth = cfg.mtp_num_layers || 0; + if (!depth) return { fwd: 0, bwd: 0, ehUs: 0 }; + + const cr = String((cfg.mtp_compress_ratios && cfg.mtp_compress_ratios[0]) || 0); + const inner = lt[cr] || lt["0"]; + const innerBwd = inner.bwd + (c.recompute === "full" ? inner.fwd : 0); + const outF = nonLayer(data, "output", "forward", gpu, c) + nonLayer(data, "loss", "forward", gpu, c); + const outB = nonLayer(data, "output", "backward", gpu, c) + nonLayer(data, "loss", "backward", gpu, c); + + // No MTP trace row exists yet. Approximate eh_proj as GEMM-like work scaled + // from the measured output projection time, then split fwd/bwd as 1:2. + const af = data.analytic_flops || {}; + const mtp = af.mtp || {}; + const outFlops = af.output_flops || 0; + const outUs = outF + outB; + const ehUs = outFlops > 0 ? outUs * ((mtp.eh_proj_flops || 0) / outFlops) : 0; + + return { + fwd: depth * (inner.fwd + outF) + ehUs / 3, + bwd: depth * (innerBwd + outB) + (2 * ehUs) / 3, + ehUs, + }; +} + +// MTP time with manual override. Falls back to the analytic estimate. +function effectiveMtp(data, gpu, c, lt) { + const base = mtpTimes(data, gpu, c, lt); + if (c.modelMode !== "manual" || !((data.model_config.mtp_num_layers || 0) > 0)) return base; + const m = (c.man && c.man[gpu]) || {}; + return { + fwd: isPositiveNumber(m.mtp_f) ? m.mtp_f : base.fwd, + bwd: isPositiveNumber(m.mtp_b) ? m.mtp_b : base.bwd, + ehUs: base.ehUs, + }; +} + +// --------------------------------------------------------------------------- +// Param estimate (Step 4) +// --------------------------------------------------------------------------- +function estimateParams(cfg) { + // Prefer the exact V4 count emitted by parse_trace (MLA low-rank attention + + // MoE + tied-free embedding/output). Fall back to a crude estimate. + if (cfg.total_params) return cfg.total_params; + const h = cfg.hidden_size, exp = cfg.num_experts, mff = cfg.moe_ffn_hidden_size; + const sff = cfg.moe_shared_expert_intermediate_size || mff, V = cfg.vocab_size, L = cfg.num_layers; + const perLayer = 4 * h * h + exp * 3 * h * mff + 3 * h * sff; + return L * perLayer + 2 * V * h; +} + +function estimateOptimizerParams(data, c, dp) { + const cfg = data.model_config; + const totalParams = estimateParams(cfg); + const pp = Math.max(1, c.pp); + const tp = Math.max(1, c.tp); + const dpSize = Math.max(1, dp); + + // total_params is a full-model count. CP does not shard weights; EP is already + // represented in the full expert count and cancels out with the DP replica + // count for ZeRO-1 optimizer sharding, so the average rank owns total/(PP*TP) + // model params and updates total/(PP*TP*DP) params per optimizer step. + const localModelParams = totalParams / (pp * tp); + const perRankParams = localModelParams / dpSize; + const measuredUs = data.optimizer?.time_us ?? null; + return { totalParams, localModelParams, perRankParams, measuredUs }; +} + +// --------------------------------------------------------------------------- +// Pipeline mapping + projection (Steps 2-5) +// --------------------------------------------------------------------------- +function project(data, gpu, c, validation = validateControls(data, c, gpu)) { + if (validation.errors.length) return null; + const cfg = data.model_config; + const crs = cfg.compress_ratios; + const L = crs.length; + const dp = validation.dp; + const world = c.world; + const lt = {}; + for (const cr of ["0", "4", "128"]) lt[cr] = effectiveLayerTimes(data, cr, gpu, c); + + // Step 2: assign layers to PP*VPP chunks -> devices + const C = c.pp * c.vpp; + const perChunk = Math.ceil(L / C); + const Df = new Array(c.pp).fill(0), Db = new Array(c.pp).fill(0); + const layoutInfo = validation.layout; + const layoutStages = layoutInfo.ok ? layoutInfo.stages : null; + const stageOrdinals = new Array(c.pp).fill(0); + // Per virtual-chunk detail (k = 0..C-1): device = k % PP, vpp index = floor(k/PP). + // Kept for the iteration-timeline view (design/07); does not affect Df/Db math. + const chunks = []; + for (let k = 0; k < C; k++) { + chunks.push({ chunk: k, device: k % c.pp, vpp: Math.floor(k / c.pp), layers: [], fwd: 0, bwd: 0 }); + } + const addLayer = (chunkIdx, i) => { + const dev = chunkIdx % c.pp; + const t = lt[String(crs[i])] || { fwd: 0, bwd: 0 }; + const recompute = recomputeLayer(c, stageOrdinals[dev]++); + const effBwd = t.bwd + (recompute ? t.fwd : 0); + Df[dev] += t.fwd; + Db[dev] += effBwd; + const ch = chunks[chunkIdx]; + ch.layers.push({ globalIdx: i, cr: crs[i], recompute, fwd: t.fwd, bwd: effBwd }); + ch.fwd += t.fwd; + ch.bwd += effBwd; + }; + if (layoutStages) { + layoutStages.forEach((layers, chunk) => { + for (const i of layers) addLayer(chunk, i); + }); + } else { + for (let i = 0; i < L; i++) addLayer(Math.floor(i / perChunk), i); + } + // non-layer parts on first / last device (manual-overridable) + const embFwd = effectiveNonLayer(data, "embedding", "forward", gpu, c); + const embBwd = effectiveNonLayer(data, "embedding", "backward", gpu, c); + Df[0] += embFwd; + Db[0] += embBwd; + const last = c.pp - 1; + const outFwd = effectiveNonLayer(data, "output", "forward", gpu, c) + effectiveNonLayer(data, "loss", "forward", gpu, c); + const outBwd = effectiveNonLayer(data, "output", "backward", gpu, c) + effectiveNonLayer(data, "loss", "backward", gpu, c); + Df[last] += outFwd; + Db[last] += outBwd; + const mtp = effectiveMtp(data, gpu, c, lt); + Df[last] += mtp.fwd; + Db[last] += mtp.bwd; + + const critF = Math.max(...Df), critB = Math.max(...Db); + + // Assemble per-device schedule detail (design/07). Non-layer parts are attached + // to their owning device so the timeline can render them; Df/Db already include + // them above. + const critFdev = Df.indexOf(critF), critBdev = Db.indexOf(critB); + const devices = []; + for (let d = 0; d < c.pp; d++) { + devices.push({ + device: d, + chunks: chunks.filter((ch) => ch.device === d).sort((a, b) => a.vpp - b.vpp), + Df: Df[d], Db: Db[d], + hasEmb: d === 0, hasOut: d === last, hasMtp: d === last && mtp.fwd > 0, + embFwd: d === 0 ? embFwd : 0, embBwd: d === 0 ? embBwd : 0, + outFwd: d === last ? outFwd : 0, outBwd: d === last ? outBwd : 0, + mtpFwd: d === last ? mtp.fwd : 0, mtpBwd: d === last ? mtp.bwd : 0, + isCritF: d === critFdev, isCritB: d === critBdev, + }); + } + const schedule = { chunks, devices, C, critFdev, critBdev }; + + // Step 3: pipeline compute time (us) ; GA = gbs/(dp*mbs) + // calibFactor corrects the single-layer-capture -> full-model bias (~0.93, + // from the flash 16-layer single-node calibration; see design/06). + const ga = c.gbs / (dp * c.mbs); + const pipeUs = (ga + (c.pp - 1) / c.vpp) * (critF + critB) * c.calibFactor; + const bubbleFrac = (c.pp - 1) / c.vpp / (ga + (c.pp - 1) / c.vpp); + + // Step 4: optimizer (memory-bound, zero1-sharded over DP; CP does not shard params) + const optParams = estimateOptimizerParams(data, c, dp); + const totalParams = optParams.totalParams; + const perRankParams = optParams.perRankParams; + const bw = (gpu === "MI355X" ? c.bw355 : c.bw455) * 1e9; // bytes/s + const optTimeS = (perRankParams * c.bytesPerParam) / bw / c.optEff; + const optUs = optTimeS * 1e6; + + // Step 5: totals + const iterUs = pipeUs + optUs; + const iterS = iterUs * 1e-6; + const seq = data.capture.seq_length; + const tokIter = c.gbs * seq; + const tokS = tokIter / iterS; + const tokSgpu = tokS / world; + + // FLOPs/iter — Megatron-convention V4 analytic model FLOPs (independent of + // recompute; recompute adds time, not model flops). Falls back to breakdown + // gemm flops if analytic_flops is absent. + const counts = cfg.cr_layer_counts; + let fMb = 0; + const af = data.analytic_flops; + if (af && af.per_cr_layer_flops) { + for (const cr of ["0", "4", "128"]) fMb += (counts[cr] || 0) * (af.per_cr_layer_flops[cr] || 0); + fMb += af.output_flops || 0; + if (af.mtp) { + fMb += (af.mtp.inner_layer_flops || 0) + + (af.mtp.eh_proj_flops || 0) + + (af.mtp.extra_logits_flops || 0) + + (af.mtp.hc_head_flops || 0); + } + } else { + for (const cr of ["0", "4", "128"]) fMb += (counts[cr] || 0) * (lt[cr].fFlops + lt[cr].bFlops); + fMb += nonLayerFlops(data, "output", "forward") + nonLayerFlops(data, "output", "backward"); + } + const flopsIter = fMb * ga * dp; + const tflopsGpu = flopsIter / iterS / world / 1e12; + + return { + lt, Df, Db, critF, critB, ga, pipeUs, bubbleFrac, world, dp, totalParams, + layoutApplied: Boolean(layoutStages), + layoutCounts: layoutInfo.counts, + layoutMessage: layoutInfo.message, + localModelParams: optParams.localModelParams, perRankParams, measuredOptUs: optParams.measuredUs, + mtp, optUs, iterUs, tokIter, tokS, tokSgpu, flopsIter, tflopsGpu, seq, + schedule, + }; +} + +// --------------------------------------------------------------------------- +// Rendering +// --------------------------------------------------------------------------- +function renderConfig() { + const cfg = STATE.data.model_config; + const grid = $("#config-grid"); + grid.innerHTML = ""; + const items = [ + ["Model", STATE.data.model.toUpperCase()], + ["Layers", cfg.num_layers], + ["Hidden", cfg.hidden_size], + ["Attn heads", cfg.num_attention_heads], + ["Experts", cfg.num_experts], + ["Router top-k", cfg.moe_router_topk], + ["MoE FFN", cfg.moe_ffn_hidden_size], + ["Index top-k", cfg.index_topk], + ["Vocab", cfg.vocab_size], + ["MTP depths", cfg.mtp_num_layers || 0], + ["Capture seq", STATE.data.capture.seq_length], + ["cr=0 layers", cfg.cr_layer_counts["0"]], + ["cr=4 layers", cfg.cr_layer_counts["4"]], + ["cr=128 layers", cfg.cr_layer_counts["128"]], + ]; + for (const [k, v] of items) { + const kv = el("div", { class: "kv" }); + kv.append(el("b", {}, k), el("span", {}, String(v))); + grid.append(kv); + } + // cr schedule strip + const strip = $("#cr-schedule"); + strip.innerHTML = ""; + const row = el("div", { class: "cr-schedule" }); + for (const cr of cfg.compress_ratios) row.append(el("span", { class: `cr-cell cr-${cr}`, title: `cr=${cr}` }, "x")); + strip.append(row); + strip.append(el("div", { class: "cr-legend", html: + 'cr=0 dense+SWAcr=4 CSAcr=128 HCA' })); +} + +function breakdownBlock(title, fwdList, bwdList) { + const c = STATE.controls, gpu = STATE.gpu; + const block = el("div", { class: "bd-block" }); + block.append(el("h3", {}, title)); + const scroll = el("div", { class: "bd-scroll" }); + const table = el("table", { class: "bd" }); + + const fwd = fwdList.map((r) => ({ r, t: rowScaledTime(r, gpu, c), phase: "F" })); + const bwd = bwdList.map((r) => ({ r, t: rowScaledTime(r, gpu, c), phase: "B" })).reverse(); + const cols = [...fwd, ...bwd]; + const dividerIdx = fwd.length; + + const head = el("tr"); + head.append(el("th", { class: "rowlab" }, "Module")); + cols.forEach((col, i) => { + const th = el("th", { class: i === dividerIdx ? "divider" : "" }); + th.append(el("div", {}, col.r.module.replace(/^(attn|moe)\./, ""))); + th.append(el("div", { class: "phase-tag" }, col.phase)); + head.append(th); + }); + table.append(head); + + const timeRow = el("tr"); + timeRow.append(el("td", { class: "rowlab" }, "Time µs")); + cols.forEach((col, i) => { + const cls = (col.r.class === "compute_bound" ? "cell-compute" : "cell-memory") + (i === dividerIdx ? " divider" : ""); + timeRow.append(el("td", { class: cls }, fmt(col.t, 0))); + }); + table.append(timeRow); + + const tfRow = el("tr"); + tfRow.append(el("td", { class: "rowlab" }, "TFLOP/s (kernel)")); + cols.forEach((col, i) => { + const tf = rowTflops(col.r, col.t); + tfRow.append(el("td", { class: i === dividerIdx ? "divider" : "" }, tf ? fmt(tf, 0) : "—")); + }); + table.append(tfRow); + + scroll.append(table); + block.append(scroll); + return block; +} + +function renderValidation(validation) { + const host = $("#validation-panel"); + if (!host) return; + host.innerHTML = ""; + if (!validation.errors.length && !validation.warnings.length) { + host.hidden = true; + return; + } + host.hidden = false; + if (validation.errors.length) { + const box = el("div", { class: "validation validation--error" }); + box.append(el("b", {}, "Errors")); + const ul = el("ul"); + for (const msg of validation.errors) ul.append(el("li", {}, msg)); + box.append(ul); + host.append(box); + } + if (validation.warnings.length) { + const box = el("div", { class: "validation validation--warn" }); + box.append(el("b", {}, "Warnings")); + const ul = el("ul"); + for (const msg of validation.warnings) ul.append(el("li", {}, msg)); + box.append(ul); + host.append(box); + } +} + +function renderBreakdown() { + $("#breakdown-gpu").textContent = `· ${STATE.gpu}`; + const panel = $("#breakdown-panel"); + if (panel) panel.classList.toggle("is-muted", STATE.controls?.modelMode === "manual"); + const note = $("#breakdown-manual-note"); + if (note) note.hidden = STATE.controls?.modelMode !== "manual"; + const host = $("#breakdown-blocks"); + host.innerHTML = ""; + const d = STATE.data; + host.append(breakdownBlock("Embedding", d.non_layer.embedding.forward, d.non_layer.embedding.backward)); + for (const cr of ["0", "4", "128"]) { + const L = d.layers[cr]; + host.append(breakdownBlock( + `Layer cr=${cr} — attention`, L.attention.forward, L.attention.backward)); + host.append(breakdownBlock( + `Layer cr=${cr} — MoE`, L.moe.forward, L.moe.backward)); + } + host.append(breakdownBlock("Output (norm + lm_head)", d.non_layer.output.forward, d.non_layer.output.backward)); + host.append(breakdownBlock("Loss", d.non_layer.loss.forward, d.non_layer.loss.backward)); +} + +function step(label, value, sub) { + const s = el("div", { class: "step" }); + s.append(el("span", { class: "num" }, value)); + s.append(el("b", {}, label)); + if (sub) { s.append(document.createElement("br")); s.append(el("small", {}, sub)); } + return s; +} + +function renderResults(validation) { + const c = STATE.controls, gpu = STATE.gpu, d = STATE.data; + $("#results-gpu").textContent = `· ${gpu}`; + const p = project(d, gpu, c, validation); + + const head = $("#results-headline"); + head.innerHTML = ""; + const steps = $("#results-steps"); + steps.innerHTML = ""; + if (!p) { + const errs = (validation && validation.errors) || []; + const summary = errs.length === 0 ? "Invalid controls" : `${errs.length} error${errs.length > 1 ? "s" : ""}`; + head.append(el("div", { class: "metric metric--error" }, el("b", {}, "Projection blocked"), el("span", {}, summary))); + if (errs.length) { + for (const msg of errs) { + const s = el("div", { class: "step step--error" }); + s.append(el("b", {}, "⚠ Validation error")); + s.append(document.createElement("br")); + s.append(el("small", {}, msg)); + steps.append(s); + } + } else { + steps.append(step("Fix validation errors", "", "Projection is not recomputed while errors are present.")); + } + return; + } + const mk = (label, val, primary) => { + const m = el("div", { class: "metric" + (primary ? " metric--primary" : "") }); + m.append(el("b", {}, label), el("span", {}, val)); + return m; + }; + head.append(mk("tokens/s/GPU", fmtInt(p.tokSgpu), true)); + head.append(mk("TFLOP/s/GPU", fmt(p.tflopsGpu, 0))); + head.append(mk("Iteration time", `${fmt(p.iterUs / 1000, 1)} ms`)); + head.append(mk("Pipeline bubble", `${fmt(p.bubbleFrac * 100, 1)} %`)); + + const ltDesc = ["0", "4", "128"].map((cr) => + `cr${cr}: F ${fmt(p.lt[cr].fwd, 0)} / B ${fmt(p.lt[cr].bwd, 0)} µs`).join(" · "); + const recomputeDesc = c.recompute === "first-n" ? `first ${c.recomputeLayers} layers/stage` : (c.recompute === "full" ? "recompute on" : "no recompute"); + const sourceTag = c.modelMode === "manual" ? "manual" : "trace"; + steps.append(step(`Per-layer fwd/bwd (µs, ${sourceTag}, ${recomputeDesc})`, "", ltDesc)); + steps.append(step("Critical PP stage (µs)", `F ${fmt(p.critF, 0)} + B ${fmt(p.critB, 0)}`, + `max over ${c.pp} stages; layout ${p.layoutApplied ? "applied" : "balanced fallback"} (${p.layoutMessage}); counts=[${p.layoutCounts.join(", ")}]; per-device fwd=[${p.Df.map((x) => fmt(x / 1000, 1)).join(", ")}] ms`)); + if ((d.model_config.mtp_num_layers || 0) > 0) { + steps.append(step("MTP on last stage", `F ${fmt(p.mtp.fwd, 0)} + B ${fmt(p.mtp.bwd, 0)} µs`, + `${d.model_config.mtp_num_layers} depth(s), inner cr=${(d.model_config.mtp_compress_ratios || [0])[0]}, eh_proj estimated from output GEMM throughput`)); + } + steps.append(step("GA (microbatches)", fmt(p.ga, 0), `GA = GBS ${c.gbs} / (DP ${p.dp} × MBS ${c.mbs})`)); + steps.append(step("Pipeline compute / iter", `${fmt(p.pipeUs / 1000, 2)} ms`, + `(GA + (PP−1)/VPP) × (F+B)_crit ; bubble ${fmt(p.bubbleFrac * 100, 1)}%`)); + const optHint = p.measuredOptUs + ? `; one-layer trace optimizer ref ${fmt(p.measuredOptUs / 1000, 2)} ms` + : ""; + steps.append(step("Optimizer step / iter", `${fmt(p.optUs / 1000, 2)} ms`, + `zero1: ${fmtInt(p.perRankParams / 1e6)}M optim params/rank (${fmtInt(p.localModelParams / 1e6)}M local model params) × ${c.bytesPerParam}B / HBM-BW / eff ${c.optEff}${optHint}`)); + steps.append(step("Iteration time", `${fmt(p.iterUs / 1000, 2)} ms`, "pipeline compute + optimizer (DP/PP comm assumed hidden)")); + steps.append(step("World size", fmtInt(p.world), `PP ${c.pp} × TP ${c.tp} × CP ${c.cp} × DP ${p.dp} (derived); EP ${c.ep} ≤ DP`)); + steps.append(step("Tokens / iter", fmtInt(p.tokIter), `GBS ${c.gbs} × seq ${p.seq}`)); + steps.append(step("tokens/s/GPU", fmtInt(p.tokSgpu), `${fmtInt(p.tokS)} tok/s ÷ ${p.world} GPUs`)); + steps.append(step("TFLOP/s/GPU", fmt(p.tflopsGpu, 0), "V4 analytic model FLOPs (Megatron convention)")); + + // self-consistency hint + const cap = d.capture; + if (cap && cap.measured_iter_time_ms) { + steps.append(step("Self-consistency (measured)", `${fmt(cap.measured_iter_time_ms, 1)} ms`, + "set PP=1,VPP=1,DP=1,EP=8,MBS=1,GBS=2 to compare against capture")); + } +} + +// --------------------------------------------------------------------------- +// Iteration timeline (design/07): 3-level composition view +// --------------------------------------------------------------------------- +const TL_CATS = ["attn", "mlp", "a2a", "misc"]; +const TL_CAT_LABEL = { attn: "attn", mlp: "mlp", a2a: "a2a (dispatch+combine)", misc: "misc/unattrib" }; +const CR_HEX = { "0": "#6b4a2d", "4": "#2d6b4a", "128": "#2d4a6b" }; + +// Map a module name to one of the Level-1 categories (design/07). +function moduleCategory(module) { + const m = String(module || ""); + if (m.startsWith("attn.")) return "attn"; + if (m === "moe.dispatch" || m === "moe.combine") return "a2a"; + if (/(^|\.)misc$|unattrib/.test(m)) return "misc"; + return "mlp"; // moe.grouped_gemm / shared_expert / router and any other moe.* +} + +// Per-cr forward/backward time split into categories, scaled to the active GPU. +// In manual mode the trace composition is rescaled so the bar total matches the +// effective (manual-overridden) per-layer time used by the projection. +function categoryBreakdown(data, cr, gpu, c) { + const out = { forward: { attn: 0, mlp: 0, a2a: 0, misc: 0 }, backward: { attn: 0, mlp: 0, a2a: 0, misc: 0 } }; + const L = data.layers[cr]; + if (!L) return out; + for (const bucket of [L.attention, L.moe]) { + for (const phase of ["forward", "backward"]) { + for (const r of bucket[phase]) out[phase][moduleCategory(r.module)] += rowScaledTime(r, gpu, c); + } + } + if (c.modelMode === "manual") { + const eff = effectiveLayerTimes(data, cr, gpu, c); + for (const phase of ["forward", "backward"]) { + const sum = TL_CATS.reduce((a, k) => a + out[phase][k], 0); + const target = phase === "forward" ? eff.fwd : eff.bwd; + if (sum > 0 && target > 0) for (const k of TL_CATS) out[phase][k] *= target / sum; + } + } + return out; +} + +// True when, in manual mode, the user has actually changed this cr's fwd/bwd away +// from the trace-derived baseline (comparison is rounded to µs so the prefilled +// baseline itself does not count as an edit). +function layerManualEdited(data, cr, gpu, c) { + if (c.modelMode !== "manual") return false; + const trace = layerTimes(data, cr, gpu, c); + const eff = effectiveLayerTimes(data, cr, gpu, c); + return Math.round(eff.fwd) !== Math.round(trace.fwd) || Math.round(eff.bwd) !== Math.round(trace.bwd); +} + +// Small hex lighten/darken for VPP chunk shading (Level 3). +function shadeHex(hex, amt) { + const n = parseInt(hex.slice(1), 16); + const clamp = (x) => Math.max(0, Math.min(255, Math.round(x))); + const r = clamp(((n >> 16) & 255) + amt), g = clamp(((n >> 8) & 255) + amt), b = clamp((n & 255) + amt); + return `#${((r << 16) | (g << 8) | b).toString(16).padStart(6, "0")}`; +} + +// Group PP devices with an identical ordered layer/recompute/non-layer signature +// so identical middle stages are drawn once (design/07 Level 2 dedup). +function dedupDevices(devices) { + const groups = []; + const byKey = new Map(); + for (const dev of devices) { + const sig = JSON.stringify({ + layers: dev.chunks.map((ch) => ch.layers.map((l) => `${l.cr}${l.recompute ? "r" : ""}`)), + emb: dev.hasEmb, out: dev.hasOut, mtp: dev.hasMtp, + }); + if (byKey.has(sig)) byKey.get(sig).members.push(dev.device); + else { + const g = { rep: dev, members: [dev.device] }; + byKey.set(sig, g); + groups.push(g); + } + } + return groups; +} + +// 1F1B (optionally interleaved) schedule simulator (design/07 Level 3). Returns +// per-device event lists with start/dur (µs) plus the drawn iteration length. +const TL_VIS_GA_CAP = 48; +function simulateSchedule(p, c, { interleaved }) { + const PP = c.pp; + const VPP = interleaved ? Math.max(1, c.vpp) : 1; + const C = PP * VPP; + const gaFull = Math.round(p.ga); + const GA = Math.min(gaFull, TL_VIS_GA_CAP); + const capped = gaFull > GA; + + // per-virtual-chunk fwd/bwd durations (k = 0..C-1, device = k % PP) + const fdur = new Array(C).fill(0), bdur = new Array(C).fill(0); + if (VPP === 1) { + for (let d = 0; d < PP; d++) { fdur[d] = p.Df[d]; bdur[d] = p.Db[d]; } + } else { + for (const ch of p.schedule.chunks) { fdur[ch.chunk] = ch.fwd; bdur[ch.chunk] = ch.bwd; } + // fold non-layer parts into first/last virtual chunk so the drawn length is + // consistent with the per-device Df/Db (which include them). + const dev0 = p.schedule.devices[0], devL = p.schedule.devices[PP - 1]; + fdur[0] += dev0.embFwd; bdur[0] += dev0.embBwd; + fdur[C - 1] += devL.outFwd + devL.mtpFwd; bdur[C - 1] += devL.outBwd + devL.mtpBwd; + } + + const total = GA * VPP; + const group = PP * VPP; + const fwdChunkMb = (i) => { + const inGroup = i % group; + const v = Math.floor(inGroup / PP); + const m = Math.floor(i / group) * PP + (inGroup % PP); + return { v, m }; + }; + + const ops = []; // per device ordered op list + for (let d = 0; d < PP; d++) { + const warmup = Math.min( + VPP === 1 ? PP - 1 - d : (PP - d - 1) * 2 + (VPP - 1) * PP, + total, + ); + const list = []; + for (let i = 0; i < warmup; i++) { + const { v, m } = fwdChunkMb(i); + list.push({ kind: "F", m, k: v * PP + d }); + } + let fptr = warmup, bptr = 0; + const n1f1b = total - warmup; + for (let i = 0; i < n1f1b; i++) { + const f = fwdChunkMb(fptr++); + list.push({ kind: "F", m: f.m, k: f.v * PP + d }); + const b = fwdChunkMb(bptr++); + list.push({ kind: "B", m: b.m, k: (VPP - 1 - b.v) * PP + d }); + } + for (let i = 0; i < warmup; i++) { + const b = fwdChunkMb(bptr++); + list.push({ kind: "B", m: b.m, k: (VPP - 1 - b.v) * PP + d }); + } + ops.push(list); + } + + // ASAP scheduling: forward flows k-1 -> k, backward flows k+1 -> k (p2p hidden). + const endF = new Map(), endB = new Map(); + const kf = (m, k) => `${m}:${k}`; + const ptr = new Array(PP).fill(0); + const free = new Array(PP).fill(0); + const events = Array.from({ length: PP }, () => []); + let remaining = ops.reduce((a, l) => a + l.length, 0); + let guard = remaining * 4 + 16; + while (remaining > 0 && guard-- > 0) { + let progressed = false; + for (let d = 0; d < PP; d++) { + if (ptr[d] >= ops[d].length) continue; + const op = ops[d][ptr[d]]; + let dep = 0, ready = true; + if (op.kind === "F") { + if (op.k > 0) { + const e = endF.get(kf(op.m, op.k - 1)); + if (e == null) ready = false; else dep = e; + } + } else { + if (op.k < C - 1) { + const e = endB.get(kf(op.m, op.k + 1)); + if (e == null) ready = false; else dep = e; + } else { + const e = endF.get(kf(op.m, op.k)); + if (e == null) ready = false; else dep = e; + } + } + if (!ready) continue; + const dur = op.kind === "F" ? fdur[op.k] : bdur[op.k]; + const start = Math.max(free[d], dep); + const end = start + dur; + events[d].push({ kind: op.kind, m: op.m, k: op.k, vpp: Math.floor(op.k / PP), start, dur }); + free[d] = end; + (op.kind === "F" ? endF : endB).set(kf(op.m, op.k), end); + ptr[d]++; remaining--; progressed = true; + } + if (!progressed) break; // dependency stall guard + } + const drawnUs = Math.max(0, ...free); + return { events, drawnUs, GA, capped, PP, VPP, C }; +} + +// --- cross-level selection (drill-down linkage) --- +function tlSelActive() { + return STATE.tlSel.cr != null || (STATE.tlSel.devices && STATE.tlSel.devices.length); +} +function tlSelectCr(cr) { + STATE.tlSel = STATE.tlSel.cr === cr ? { cr: null, devices: null } : { cr, devices: null }; + renderTimeline(); +} +function tlSelectDevices(devices) { + const same = STATE.tlSel.devices && STATE.tlSel.devices.length === devices.length && STATE.tlSel.devices.every((d, i) => d === devices[i]); + STATE.tlSel = same ? { cr: null, devices: null } : { cr: null, devices }; + renderTimeline(); +} +function tlClearSel() { + STATE.tlSel = { cr: null, devices: null }; + renderTimeline(); +} +// Compression ratios "active" under the current selection (drives L1 highlight). +function tlActiveCrSet(p) { + if (STATE.tlSel.cr != null) return new Set([String(STATE.tlSel.cr)]); + if (STATE.tlSel.devices && STATE.tlSel.devices.length) { + const s = new Set(); + for (const d of STATE.tlSel.devices) { + for (const ch of p.schedule.devices[d].chunks) for (const l of ch.layers) s.add(String(l.cr)); + } + return s; + } + return null; +} +const tlDevSet = () => (STATE.tlSel.devices && STATE.tlSel.devices.length ? new Set(STATE.tlSel.devices) : null); + +function renderTimeline() { + const host = $("#tl-body"); + if (!host) return; + $("#timeline-gpu").textContent = `· ${STATE.gpu}`; + const stacked = STATE.tlView === "stacked"; + document.querySelectorAll(".tl-view").forEach((t) => t.classList.toggle("is-active", (t.dataset.view === "stacked") === stacked)); + document.querySelectorAll(".tl-tab").forEach((t) => t.classList.toggle("is-active", Number(t.dataset.level) === STATE.tlLevel)); + const tabs = $("#tl-tabs"); + if (tabs) tabs.style.display = stacked ? "none" : ""; + host.innerHTML = ""; + const validation = validateControls(STATE.data, STATE.controls, STATE.gpu); + const p = project(STATE.data, STATE.gpu, STATE.controls, validation); + if (!p) { + host.append(el("p", { class: "tl-warn" }, "Timeline unavailable: fix the validation errors above.")); + return; + } + + // selection / linkage indicator: a fixed-position floating card appended to + // (NOT into #tl-body) so it is fully outside the timeline's flow and + // toggling a selection never shifts the page layout at all. + document.getElementById("tl-selbar-float")?.remove(); + if (tlSelActive()) { + const bar = el("div", { id: "tl-selbar-float", class: "tl-selbar" }); + const txt = el("div", { class: "tl-selbar__txt" }); + const desc = STATE.tlSel.cr != null + ? `Focused on cr=${STATE.tlSel.cr}` + : `Focused on PP rank(s) ${STATE.tlSel.devices.join(", ")}`; + txt.append(el("b", {}, desc)); + txt.append(el("span", { class: "muted" }, stacked ? " — highlighted across all levels" : " — highlighted; switch to Stacked to see all levels")); + bar.append(txt); + const btn = el("button", { class: "tl-clear" }, "Clear"); + btn.addEventListener("click", tlClearSel); + bar.append(btn); + document.body.append(bar); + } + + if (stacked) { + const section = (title, sub, fn) => { + const sec = el("div", { class: "tl-section" }); + const h = el("h3", { class: "tl-section__title" }, title); + if (sub) h.append(el("span", { class: "tl-section__sub" }, sub)); + sec.append(h); + const body = el("div", {}); + fn(body); + sec.append(body); + host.append(sec); + }; + section("Level 1 · single layer", "attn / mlp / a2a — click a cr to link", (b) => renderTimelineL1(b, p)); + section("Level 2 · pipeline ranks", "layer granularity — click a rank or a layer to link", (b) => renderTimelineL2(b, p)); + section("Level 3 · pipeline schedule", "1F1B / interleaved — click a device to link", (b) => renderTimelineL3(b, p)); + } else if (STATE.tlLevel === 1) renderTimelineL1(host, p); + else if (STATE.tlLevel === 2) renderTimelineL2(host, p); + else renderTimelineL3(host, p); +} + +// --- gantt export (SVG / PNG) --- +function downloadBlob(blob, filename) { + const url = URL.createObjectURL(blob); + const a = el("a", { href: url, download: filename }); + document.body.append(a); + a.click(); + a.remove(); + setTimeout(() => URL.revokeObjectURL(url), 2000); +} +function serializeGantt(svg) { + const clone = svg.cloneNode(true); + const cs = getComputedStyle(document.documentElement); + const bg = (cs.getPropertyValue("--bg").trim() || "#0f1218"); + let markup = new XMLSerializer().serializeToString(clone); + for (const v of ["--panel-2", "--border", "--muted", "--text", "--bg"]) { + markup = markup.split(`var(${v})`).join(cs.getPropertyValue(v).trim() || "#888"); + } + if (!markup.includes("xmlns=")) markup = markup.replace("]*>)/, `$1`); + return markup; +} +function exportGanttSvg(svg) { + downloadBlob(new Blob([serializeGantt(svg)], { type: "image/svg+xml" }), `dsv4-pp-schedule-${STATE.gpu}.svg`); +} +function exportGanttPng(svg) { + const markup = serializeGantt(svg); + const vb = svg.viewBox.baseVal, scale = 2; + const img = new Image(); + img.onload = () => { + const canvas = el("canvas"); + canvas.width = vb.width * scale; + canvas.height = vb.height * scale; + const ctx = canvas.getContext("2d"); + ctx.scale(scale, scale); + ctx.drawImage(img, 0, 0); + canvas.toBlob((b) => downloadBlob(b, `dsv4-pp-schedule-${STATE.gpu}.png`)); + }; + img.src = "data:image/svg+xml;base64," + btoa(unescape(encodeURIComponent(markup))); +} + +function catLegend() { + return el("div", { class: "tl-legend", html: + 'attn' + + 'mlp (experts)' + + 'a2a (dispatch+combine)' + + 'misc/unattributed' }); +} + +function stackedTrack(phaseObj, maxTotal) { + const track = el("div", { class: "tl-bar__track" }); + const total = TL_CATS.reduce((a, k) => a + phaseObj[k], 0); + for (const cat of TL_CATS) { + const t = phaseObj[cat]; + if (t <= 0) continue; + const w = maxTotal > 0 ? (t / maxTotal) * 100 : 0; + const seg = el("div", { + class: `tl-seg tl-c-${cat}`, + style: `width:${w}%`, + title: `${TL_CAT_LABEL[cat]}: ${fmt(t, 0)} µs (${fmt(total > 0 ? (t / total) * 100 : 0, 0)}%)`, + }, w > 6 ? cat : ""); + track.append(seg); + } + return { track, total }; +} + +// Single-segment bar used when a cr's layer time is manually set: no module split +// is possible, so show only the total fwd/bwd. +function aggregateTrack(timeUs, maxTotal, label) { + const track = el("div", { class: "tl-bar__track" }); + const w = maxTotal > 0 ? (timeUs / maxTotal) * 100 : 0; + track.append(el("div", { + class: "tl-seg tl-c-manual", style: `width:${w}%`, + title: `${label}: ${fmt(timeUs, 0)} µs (manual per-layer)`, + }, w > 6 ? "manual" : "")); + return { track, total: timeUs }; +} + +function renderTimelineL1(host, p) { + const d = STATE.data, gpu = STATE.gpu, c = STATE.controls; + host.append(el("p", { class: "tl-note" }, + "One representative layer per compression ratio (MoE is cr-independent, so only attention differs). Forward and backward split into attn / mlp / a2a; bar length is comparable across cr.")); + const crs = ["0", "4", "128"].filter((cr) => (d.model_config.cr_layer_counts?.[cr] || 0) > 0); + const crSet = tlActiveCrSet(p); + // Per-cr: whether the layer time is manually overridden (then no module split), + // the module breakdown, and the fwd/bwd totals used for the shared scale. + const info = {}; + let maxTotal = 0; + let anyManual = false; + for (const cr of crs) { + const edited = layerManualEdited(d, cr, gpu, c); + const bd = edited ? null : categoryBreakdown(d, cr, gpu, c); + const eff = effectiveLayerTimes(d, cr, gpu, c); + const fT = edited ? eff.fwd : TL_CATS.reduce((a, k) => a + bd.forward[k], 0); + const bT = edited ? eff.bwd : TL_CATS.reduce((a, k) => a + bd.backward[k], 0); + info[cr] = { edited, bd, fT, bT }; + anyManual = anyManual || edited; + maxTotal = Math.max(maxTotal, fT, bT); + } + for (const cr of crs) { + const { edited, bd, fT, bT } = info[cr]; + const isSel = STATE.tlSel.cr === cr; + const dim = crSet && !crSet.has(cr); + const row = el("div", { class: "tl-l1-row tl-clickable" + (isSel ? " is-sel" : "") + (dim ? " tl-dim" : "") }); + row.addEventListener("click", () => tlSelectCr(cr)); + const lab = el("div", { class: "tl-l1-lab" }); + lab.append(el("b", {}, `cr=${cr}`), el("small", {}, `×${d.model_config.cr_layer_counts[cr]} layers${edited ? " · manual" : ""}`)); + row.append(lab); + const bars = el("div", { class: "tl-bars" }); + for (const [tag, phase, tot] of [["fwd", "forward", fT], ["bwd", "backward", bT]]) { + const bar = el("div", { class: "tl-bar" }); + bar.append(el("span", { class: "tl-bar__tag" }, tag)); + const { track, total } = edited ? aggregateTrack(tot, maxTotal, `${tag} layer`) : stackedTrack(bd[phase], maxTotal); + bar.append(track); + bar.append(el("span", { class: "tl-bar__total" }, `${fmt(total, 0)} µs`)); + bars.append(bar); + } + row.append(bars); + host.append(row); + } + if (anyManual) { + host.append(el("div", { class: "tl-legend", html: + 'manual per-layer total (no module split)' })); + host.append(el("p", { class: "tl-note" }, + "A cr marked “manual” uses a hand-entered whole-layer fwd/bwd time, so it cannot be split into attn / mlp / a2a — only the total is shown. Switch layer timing back to “Trace-derived”, or click “Restore defaults”, to see the per-module breakdown again.")); + } + host.append(catLegend()); +} + +function renderTimelineL2(host, p) { + const c = STATE.controls; + host.append(el("p", { class: "tl-note" }, + "Each pipeline rank's layers (coloured by cr; hatched = recomputed). Identical ranks are drawn once. The critical stage (max fwd/bwd, sets the pipeline critical path) is outlined.")); + const groups = dedupDevices(p.schedule.devices); + const maxTotal = Math.max(1, ...p.schedule.devices.map((dv) => dv.Df + dv.Db)); + const selDevices = tlDevSet(); + const selCr = STATE.tlSel.cr != null ? String(STATE.tlSel.cr) : null; + for (const g of groups) { + const dev = g.rep; + const isCrit = dev.isCritF || dev.isCritB; + const isSelGroup = selDevices && g.members.some((m) => selDevices.has(m)); + const dimRow = (selDevices && !isSelGroup) || (selCr && !g.members.some((m) => p.schedule.devices[m].chunks.some((ch) => ch.layers.some((l) => String(l.cr) === selCr)))); + const row = el("div", { class: "tl-l2-row tl-clickable" + (isCrit ? " is-critical" : "") + (isSelGroup ? " is-sel" : "") + (dimRow ? " tl-dim" : "") }); + row.addEventListener("click", () => tlSelectDevices(g.members)); + const lab = el("div", { class: "tl-l2-lab" }); + const members = g.members; + const rangeTxt = members.length > 1 ? `PP ranks ${members[0]}–${members[members.length - 1]} (×${members.length})` : `PP rank ${members[0]}`; + lab.append(el("b", {}, rangeTxt)); + lab.append(el("small", {}, `${dev.chunks.reduce((a, ch) => a + ch.layers.length, 0)} layers · ${dev.chunks.length} vpp chunk(s)${isCrit ? " · critical" : ""}`)); + row.append(lab); + + const total = dev.Df + dev.Db; + const strip = el("div", { class: "tl-strip", style: `width:${(total / maxTotal) * 100}%` }); + if (dev.hasEmb) { + const t = dev.embFwd + dev.embBwd; + strip.append(el("div", { class: "tl-cell tl-cell--emb", style: `width:${(t / total) * 100}%`, title: `embedding F ${fmt(dev.embFwd, 0)} / B ${fmt(dev.embBwd, 0)} µs` })); + } + dev.chunks.forEach((ch, ci) => { + if (ci > 0 || dev.hasEmb) strip.append(el("div", { class: "tl-chunk-gap" })); + for (const l of ch.layers) { + const t = l.fwd + l.bwd; + const cellSel = selCr && String(l.cr) === selCr; + const cellDim = selCr && String(l.cr) !== selCr; + const cell = el("div", { + class: "tl-cell tl-clickable" + (l.recompute ? " tl-cell--recompute" : "") + (cellSel ? " tl-cell--sel" : "") + (cellDim ? " tl-dim" : ""), + style: `width:${(t / total) * 100}%; background:${CR_HEX[String(l.cr)] || "#555"}`, + title: `layer #${l.globalIdx} cr=${l.cr}${l.recompute ? " (recompute)" : ""} · F ${fmt(l.fwd, 0)} / B ${fmt(l.bwd, 0)} µs`, + }); + cell.addEventListener("click", (e) => { e.stopPropagation(); tlSelectCr(String(l.cr)); }); + strip.append(cell); + } + }); + if (dev.hasOut) { + const t = dev.outFwd + dev.outBwd + dev.mtpFwd + dev.mtpBwd; + strip.append(el("div", { class: "tl-chunk-gap" })); + strip.append(el("div", { class: "tl-cell tl-cell--out", style: `width:${(t / total) * 100}%`, title: `output/loss${dev.hasMtp ? "+MTP" : ""} F ${fmt(dev.outFwd + dev.mtpFwd, 0)} / B ${fmt(dev.outBwd + dev.mtpBwd, 0)} µs` })); + } + row.append(strip); + row.append(el("div", { class: "tl-l2-meta" }, `Df ${fmt(dev.Df / 1000, 2)} / Db ${fmt(dev.Db / 1000, 2)} ms`)); + host.append(row); + } + host.append(el("div", { class: "cr-legend", html: + 'cr=0cr=4cr=128' + + 'embeddingoutput/loss/MTP' })); +} + +function renderTimelineL3(host, p) { + const c = STATE.controls; + // Interleaving is driven directly by the VPP control (no separate toggle): + // VPP=1 is plain 1F1B, VPP>1 is interleaved 1F1B. + const interleaved = c.vpp > 1; + + const controls = el("div", { class: "tl-controls" }); + controls.append(el("span", { class: "mode-switch__label" }, `Schedule · 1F1B${interleaved ? ` (interleaved, VPP=${c.vpp})` : ""}`)); + + // Zoom slider: 1x fits the whole schedule in view (no scrollbar); zooming in + // widens the chart so the per-cell microbatch numbers become readable. + const zoomWrap = el("span", { class: "tl-zoom-wrap" }); + zoomWrap.append(el("span", { class: "mode-switch__label" }, "Zoom")); + const zoom = el("input", { type: "range", min: "1", max: "10", step: "0.5", value: String(STATE.tlZoom), class: "tl-zoom" }); + const zlab = el("span", { class: "tl-zoom-val" }, `${STATE.tlZoom}×`); + zoom.addEventListener("input", () => { zlab.textContent = `${zoom.value}×`; }); + zoom.addEventListener("change", () => { STATE.tlZoom = Number(zoom.value); renderTimeline(); }); + zoomWrap.append(zoom, zlab); + controls.append(zoomWrap); + host.append(controls); + if (!interleaved) { + host.append(el("p", { class: "tl-note" }, "VPP=1 → plain 1F1B. Set VPP>1 in the controls to interleave the schedule and shrink the pipeline bubble.")); + } + + const sim = simulateSchedule(p, c, { interleaved }); + const PP = sim.PP; + const rowH = 30, gap = 6, padL = 54, padT = 8, padB = 26; + // Zoom widens the coordinate system itself (more px per µs; font size stays + // fixed so cells become readable) instead of CSS-scaling the SVG. At 1x the + // chart is sized to the actual right-content width so it fits with no + // scrollbar; >1x overflows and scrolls horizontally, undistorted. + const avail = Math.max(600, (($("#tl-body")?.clientWidth) || 1100) - 16); + const plotW = Math.round((avail - padL) * STATE.tlZoom); + const width = padL + plotW + 8; + const scale = sim.drawnUs > 0 ? plotW / sim.drawnUs : 0; + const height = padT + PP * (rowH + gap) + padB; + + const svg = document.createElementNS("http://www.w3.org/2000/svg", "svg"); + svg.setAttribute("class", "tl-gantt"); + svg.setAttribute("viewBox", `0 0 ${width} ${height}`); + svg.setAttribute("width", String(width)); + svg.setAttribute("height", String(height)); + const mkEl = (tag, attrs, text) => { + const n = document.createElementNS("http://www.w3.org/2000/svg", tag); + for (const [k, v] of Object.entries(attrs)) n.setAttribute(k, v); + if (text != null) n.textContent = text; + return n; + }; + + const selDevices = tlDevSet(); + for (let d = 0; d < PP; d++) { + const y = padT + d * (rowH + gap); + const isSel = selDevices ? selDevices.has(d) : false; + const g = mkEl("g", { class: "tl-devrow" }); + if (selDevices) g.setAttribute("opacity", isSel ? "1" : "0.3"); + g.append(mkEl("text", { x: padL - 8, y: y + rowH / 2 + 4, "text-anchor": "end", fill: "#e6eaf2", "font-size": "11" }, `dev ${d}`)); + const bgRect = mkEl("rect", { + x: padL, y, width: plotW, height: rowH, rx: "3", + fill: isSel ? "var(--panel)" : "var(--panel-2)", + stroke: isSel ? "#36c08f" : "var(--border)", "stroke-width": isSel ? "2" : "1", + style: "cursor:pointer", + }); + bgRect.addEventListener("click", () => tlSelectDevices([d])); + g.append(bgRect); + for (const ev of sim.events[d]) { + const x = padL + ev.start * scale; + const w = Math.max(1, ev.dur * scale); + const base = ev.kind === "F" ? "#4f8cff" : "#36c08f"; + const fill = shadeHex(base, ev.vpp * -26); + const rect = mkEl("rect", { + x, y: y + 2, width: w, height: rowH - 4, rx: "2", + fill, class: ev.kind === "F" ? "tl-fwd" : "tl-bwd", + }); + rect.append(mkEl("title", {}, `${ev.kind === "F" ? "Forward" : "Backward"} · microbatch ${ev.m}${sim.VPP > 1 ? ` · vpp chunk ${ev.vpp}` : ""} · compute ${fmt(ev.dur, 0)} µs · starts @ ${fmt(ev.start / 1000, 2)} ms (from iter start)`)); + g.append(rect); + if (w >= 8) g.append(mkEl("text", { x: x + w / 2, y: y + rowH / 2 + 3, "text-anchor": "middle", fill: "#000000", "font-size": "9" }, String(ev.m))); + } + svg.append(g); + } + // time axis ticks + const yb = padT + PP * (rowH + gap); + for (let i = 0; i <= 4; i++) { + const tx = padL + (plotW * i) / 4; + svg.append(mkEl("text", { x: tx, y: yb + 16, "text-anchor": "middle", fill: "#8d97a8", "font-size": "10" }, `${fmt((sim.drawnUs * i) / 4 / 1000, 1)} ms`)); + } + + const toolbar = el("div", { class: "tl-export" }); + const svgBtn = el("button", { class: "mode-tab" }, "Export SVG"); + const pngBtn = el("button", { class: "mode-tab" }, "Export PNG"); + svgBtn.addEventListener("click", () => exportGanttSvg(svg)); + pngBtn.addEventListener("click", () => exportGanttPng(svg)); + toolbar.append(svgBtn, pngBtn); + host.append(toolbar); + + const wrap = el("div", { class: "tl-gantt-wrap" }); + wrap.append(svg); + host.append(wrap); + + // legend + axis/tooltip explanation + self-check vs analytic pipe time + host.append(el("div", { class: "tl-legend", html: + 'forwardbackward' + + (sim.VPP > 1 ? 'lighter→darker = VPP chunk 0→' + (sim.VPP - 1) + '' : '') + + 'gaps = pipeline bubble (idle)' })); + host.append(el("p", { class: "tl-note tl-axis-note" }, + "X-axis = wall-clock time measured from the start of the pipeline (ms). Each cell is one microbatch's forward or backward on that device; the number inside is the microbatch index. Hover a cell to read its compute duration in µs (the “… µs” before @) and its start time in ms from the iteration start (the “… ms” after @).")); + + const analyticPipeUs = (p.ga + (c.pp - 1) / sim.VPP) * (p.critF + p.critB); // pre-calibFactor, matches the drawn schedule's VPP + const diff = analyticPipeUs > 0 ? Math.abs(sim.drawnUs - analyticPipeUs) / analyticPipeUs : 0; + const summary = el("div", { class: "tl-summary" }); + const gaTxt = sim.capped ? `${sim.GA} of ${Math.round(p.ga)} microbatches (capped for display)` : `${sim.GA} microbatches`; + summary.append(el("div", {}, `Drawn iteration (pipeline compute): ${fmt(sim.drawnUs / 1000, 2)} ms over ${gaTxt}; PP=${c.pp}, VPP=${interleaved ? c.vpp : 1}. Analytic bubble fraction ${fmt(p.bubbleFrac * 100, 1)}%.`)); + if (!sim.capped) { + const cls = diff > 0.08 ? "tl-warn" : ""; + summary.append(el("div", { class: cls }, + `Self-check vs analytic (GA + (PP−1)/VPP)·(F+B)_crit = ${fmt(analyticPipeUs / 1000, 2)} ms → ${fmt(diff * 100, 1)}% ${diff > 0.08 ? "difference (imbalanced stages; analytic uses per-device max)" : "match"}. Official iteration time uses the analytic value × calibFactor.`)); + } + host.append(summary); +} + +function renderAll() { + const validation = validateControls(STATE.data, STATE.controls, STATE.gpu); + // keep the derived-DP readonly box in sync + const w = STATE.controls; + const roDp = document.getElementById("ctl-dp"); + if (roDp) roDp.value = Number.isFinite(validation.dp) ? validation.dp : "—"; + renderValidation(validation); + renderModeSwitch(); + renderManualGrid(); + renderConfig(); + renderBreakdown(); + renderResults(validation); + renderTimeline(); +} + +// --------------------------------------------------------------------------- +// Bootstrap +// --------------------------------------------------------------------------- +async function init(model) { + try { + STATE.data = await loadModel(model); + STATE.controls = defaultControls(STATE.data); + $("#mock-badge").hidden = !(STATE.data.provenance && STATE.data.provenance.mock); + $("#model-select").value = STATE.data.model; + renderControls(); + renderAll(); + } catch (e) { + const err = $("#error-state"); + err.hidden = false; + err.textContent = String(e); + } +} + +globalThis.DSV4Projection = { + defaultControls, + derivedDP, + expandLayoutRepeats, + parsePipelineLayout, + validateControls, + effectiveLayerTimes, + project, + moduleCategory, + categoryBreakdown, + dedupDevices, + simulateSchedule, +}; + +if (typeof document !== "undefined") { + $("#model-select").addEventListener("change", (e) => { + const m = e.target.value; + const u = new URL(location); + u.searchParams.set("model", m); + history.replaceState(null, "", u); + init(m); + }); + + document.querySelectorAll(".tab").forEach((tab) => { + tab.addEventListener("click", () => { + document.querySelectorAll(".tab").forEach((t) => t.classList.remove("is-active")); + tab.classList.add("is-active"); + STATE.gpu = tab.dataset.gpu; + if (STATE.controls?.modelMode === "manual") prefillManual(STATE.gpu); + renderAll(); + }); + }); + + document.querySelectorAll(".mode-tab").forEach((tab) => { + tab.addEventListener("click", () => { + const mode = tab.dataset.mode; + if (!STATE.controls || STATE.controls.modelMode === mode) return; + STATE.controls.modelMode = mode; + if (mode === "manual") prefillManual(STATE.gpu); + renderAll(); + }); + }); + + document.querySelectorAll(".tl-tab").forEach((tab) => { + tab.addEventListener("click", () => { + STATE.tlLevel = Number(tab.dataset.level); + renderTimeline(); + }); + }); + + document.querySelectorAll(".tl-view").forEach((tab) => { + tab.addEventListener("click", () => { + STATE.tlView = tab.dataset.view; + renderTimeline(); + }); + }); + + // Re-fit the schedule Gantt to the content width when the window is resized. + let _tlResizeTimer; + window.addEventListener("resize", () => { + clearTimeout(_tlResizeTimer); + _tlResizeTimer = setTimeout(() => { if (STATE.data) renderTimeline(); }, 150); + }); + + init(modelFromQuery()); +} diff --git a/examples/deepseek-v4/projection/site/assets/style.css b/examples/deepseek-v4/projection/site/assets/style.css new file mode 100644 index 000000000..8d9d9d805 --- /dev/null +++ b/examples/deepseek-v4/projection/site/assets/style.css @@ -0,0 +1,333 @@ +:root { + --bg: #0f1218; + --panel: #171c26; + --panel-2: #1e2532; + --accent: #4f8cff; + --accent-2: #36c08f; + --text: #e6eaf2; + --muted: #8d97a8; + --border: #2a3342; + --warn: #e0a03a; + --compute: #2d4a6b; + --memory: #3a3346; +} + +* { box-sizing: border-box; } + +body { + margin: 0; + font-family: ui-sans-serif, system-ui, -apple-system, "Segoe UI", Roboto, sans-serif; + background: var(--bg); + color: var(--text); + line-height: 1.45; +} + +.hero { + background: linear-gradient(160deg, #1b2330, #11151d); + border-bottom: 1px solid var(--border); + padding: 28px 0; +} +.hero__inner, .layout { max-width: 1600px; margin: 0 auto; padding: 0 24px; } +.eyebrow { color: var(--accent); font-weight: 600; letter-spacing: .08em; text-transform: uppercase; font-size: 12px; margin: 0 0 4px; } +.hero h1 { margin: 0 0 8px; font-size: 28px; } +.hero__summary { color: var(--muted); max-width: 760px; margin: 0 0 16px; } +.hero__row { display: flex; align-items: center; gap: 14px; margin-bottom: 14px; } + +.tabs { display: flex; gap: 8px; } +.tab { + background: var(--panel); color: var(--text); border: 1px solid var(--border); + padding: 9px 16px; border-radius: 8px 8px 0 0; cursor: pointer; font-size: 14px; +} +.tab.is-active { background: var(--accent); border-color: var(--accent); color: #fff; font-weight: 600; } + +/* Sticky switcher bar: model select + GPU tabs stay pinned to the top so the + user can switch without scrolling back up. */ +.switcher { + position: sticky; + top: 0; + z-index: 50; + background: rgba(15, 18, 24, .9); + backdrop-filter: blur(10px); + border-bottom: 1px solid var(--border); +} +.switcher__inner { + max-width: 1600px; + margin: 0 auto; + padding: 10px 24px; + display: flex; + align-items: center; + gap: 14px; + flex-wrap: wrap; +} +.switcher .tabs { margin-left: auto; } +.switcher .tab { border-radius: 8px; } + +.layout { padding-top: 22px; padding-bottom: 60px; display: flex; flex-direction: column; gap: 18px; } + +/* Split layout: sticky Projection-controls sidebar on the left (scrolls on its + own), everything else in a scrolling content column on the right. Lets you + tweak a control and watch the relevant panel on the right update in place. */ +.layout--split { flex-direction: row; align-items: flex-start; } +.sidebar { + flex: 0 0 300px; + position: sticky; + top: 72px; + max-height: calc(100vh - 88px); + overflow-y: auto; + overscroll-behavior: contain; /* wheel inside the panel does not scroll the page */ +} +.content { flex: 1 1 auto; min-width: 0; display: flex; flex-direction: column; gap: 18px; } + +.panel { background: var(--panel); border: 1px solid var(--border); border-radius: 12px; padding: 18px 20px; } +.panel--accent { border-color: #2d4f7a; } +.panel h2 { margin: 0 0 12px; font-size: 18px; } +.panel__head { margin-bottom: 12px; } +.muted { color: var(--muted); font-weight: 400; font-size: 13px; } +.two-col { display: grid; grid-template-columns: 360px 1fr; gap: 18px; align-items: start; } + +.badge { font-size: 11px; padding: 3px 8px; border-radius: 999px; font-weight: 700; letter-spacing: .04em; } +.badge--warn { background: var(--warn); color: #1a1206; } + +select, input { + background: var(--panel-2); color: var(--text); border: 1px solid var(--border); + border-radius: 7px; padding: 7px 9px; font-size: 14px; +} +/* Drop the native number spinner buttons: in the narrow sidebar they eat ~17px + and clip long values (e.g. 15824 shown as "1582"). Values are typed, not + stepped, so the arrows add no value. */ +input[type="number"] { -moz-appearance: textfield; appearance: textfield; } +input[type="number"]::-webkit-outer-spin-button, +input[type="number"]::-webkit-inner-spin-button { -webkit-appearance: none; margin: 0; } +/* Slightly tighter type + padding for the sidebar grids so 5-6 digit numbers + fit their column without truncation. */ +.controls-grid input, .manual-row input { font-size: 13px; padding-left: 8px; padding-right: 8px; } +.field { display: flex; flex-direction: column; gap: 4px; } +.field--inline { flex-direction: row; align-items: center; gap: 8px; } +.field > span { font-size: 12px; color: var(--muted); } + +.kv-grid { display: grid; grid-template-columns: repeat(auto-fill, minmax(180px, 1fr)); gap: 10px; } +.kv { background: var(--panel-2); border: 1px solid var(--border); border-radius: 8px; padding: 8px 10px; } +.kv b { display: block; font-size: 12px; color: var(--muted); font-weight: 500; } +.kv span { font-size: 15px; font-variant-numeric: tabular-nums; } + +.cr-schedule { margin-top: 12px; display: flex; flex-wrap: wrap; gap: 3px; } +.cr-cell { width: 16px; height: 16px; border-radius: 3px; font-size: 0; } +.cr-0 { background: #6b4a2d; } .cr-4 { background: #2d6b4a; } .cr-128 { background: #2d4a6b; } +.cr-legend { display: flex; gap: 14px; margin-top: 8px; font-size: 12px; color: var(--muted); } +.cr-legend i { display: inline-block; width: 12px; height: 12px; border-radius: 3px; margin-right: 5px; vertical-align: -1px; } + +.bd-block { margin-bottom: 16px; } +.bd-block h3 { font-size: 14px; margin: 0 0 6px; } +.bd-scroll { overflow-x: auto; } +table.bd { border-collapse: collapse; font-size: 12px; min-width: 100%; } +table.bd th, table.bd td { border: 1px solid var(--border); padding: 5px 8px; text-align: right; white-space: nowrap; font-variant-numeric: tabular-nums; } +table.bd th { background: var(--panel-2); color: var(--muted); font-weight: 600; } +table.bd td.rowlab, table.bd th.rowlab { text-align: left; color: var(--muted); position: sticky; left: 0; background: var(--panel); } +.cell-compute { background: rgba(79,140,255,.10); } +.cell-memory { background: rgba(160,140,200,.07); } +.divider { border-left: 2px solid var(--accent) !important; } +.phase-tag { font-size: 10px; color: var(--accent-2); } + +/* Layer-timing mode switch (trace vs manual), styled like the GPU tabs. */ +.mode-switch { display: flex; align-items: center; gap: 10px; margin-bottom: 12px; flex-wrap: wrap; } +.mode-switch__label { font-size: 12px; color: var(--muted); } +.mode-switch__tabs { display: flex; gap: 6px; } +.mode-tab { + background: var(--panel-2); color: var(--text); border: 1px solid var(--border); + padding: 6px 12px; border-radius: 7px; cursor: pointer; font-size: 13px; +} +.mode-tab.is-active { background: var(--accent); border-color: var(--accent); color: #fff; font-weight: 600; } + +.manual-grid { margin-bottom: 14px; } +.manual-grid__header { display: flex; align-items: flex-start; justify-content: space-between; gap: 12px; } +.manual-grid__hint { margin: 0 0 10px; flex: 1; } +.manual-reset { + background: var(--panel-2); color: var(--text); border: 1px solid var(--border); + padding: 5px 12px; border-radius: 7px; cursor: pointer; font-size: 12px; white-space: nowrap; +} +.manual-reset:hover { border-color: var(--accent); } +.manual-rows { display: flex; flex-direction: column; gap: 8px; } +.manual-grid__subhead { margin: 12px 0 6px; font-size: 11px; text-transform: uppercase; letter-spacing: .05em; } +.manual-row__lab--nl { + display: inline-block; padding: 3px 7px; border-radius: 6px; align-self: center; + background: var(--panel-2); border: 1px solid var(--border); color: var(--text); font-size: 12px; +} +.manual-row { + display: grid; grid-template-columns: 92px 1fr 1fr; gap: 10px; align-items: end; +} +.manual-row__lab { font-size: 12px; padding-bottom: 8px; } +.cr-tag { + display: inline-block; padding: 3px 7px; border-radius: 6px; color: #e6eaf2; + font-variant-numeric: tabular-nums; align-self: center; padding-bottom: 3px; +} +.manual-row .cr-tag { margin-bottom: 0; } +.manual-row .field { min-width: 0; } +.manual-row input { width: 100%; min-width: 0; } + +#breakdown-panel.is-muted #breakdown-blocks { opacity: .5; } +.manual-note { margin: 8px 0 0; color: var(--warn); } + +.controls-grid { display: grid; grid-template-columns: 1fr 1fr; gap: 10px; } +.controls-grid .field--full { grid-column: 1 / -1; } +/* number inputs have an intrinsic min-width that overflows the 1fr column and + overlaps the results panel; force them to fit their grid cell. */ +.controls-grid .field { min-width: 0; } +.controls-grid input, .controls-grid select { width: 100%; min-width: 0; } + +.validation { + margin-top: 10px; + border-radius: 8px; + padding: 9px 11px; + font-size: 12px; +} +.validation + .validation { margin-top: 8px; } +.validation b { display: block; margin-bottom: 4px; } +.validation ul { margin: 0; padding-left: 18px; } +.validation--error { background: #3a1c1c; border: 1px solid #6b2d2d; color: #ffb4b4; } +.validation--warn { background: #352913; border: 1px solid #6a4d1c; color: #ffd28a; } + +.headline { display: flex; flex-wrap: wrap; gap: 14px; margin-bottom: 16px; } +.metric { background: var(--panel-2); border: 1px solid var(--border); border-radius: 10px; padding: 12px 16px; min-width: 150px; } +.metric b { display: block; font-size: 12px; color: var(--muted); } +.metric span { font-size: 22px; font-weight: 700; font-variant-numeric: tabular-nums; } +.metric.metric--primary { border-color: var(--accent-2); } +.metric.metric--primary span { color: var(--accent-2); } +.metric.metric--error { border-color: #6b2d2d; } +.metric.metric--error span { color: #ffb4b4; font-size: 18px; } + +.steps { display: flex; flex-direction: column; gap: 8px; } +.step { background: var(--panel-2); border: 1px solid var(--border); border-radius: 8px; padding: 9px 12px; font-size: 13px; } +.step b { color: var(--accent); } +.step .num { float: right; font-variant-numeric: tabular-nums; color: var(--text); } +.step small { color: var(--muted); } +.step--error { background: #2a1616; border-color: #6b2d2d; } +.step--error b { color: #ffb4b4; } +.step--error small { color: #ffd0d0; } + +.state--error { background: #3a1c1c; border: 1px solid #6b2d2d; color: #ffb4b4; padding: 12px 16px; border-radius: 8px; } +.site-footer { border-top: 1px solid var(--border); padding: 20px 24px; color: var(--muted); font-size: 13px; max-width: 1600px; margin: 0 auto; } +code { background: var(--panel-2); padding: 1px 5px; border-radius: 4px; } + +@media (max-width: 880px) { + .two-col { grid-template-columns: 1fr; } + .layout--split { flex-direction: column; } + .sidebar { position: static; max-height: none; flex-basis: auto; width: 100%; overflow: visible; } +} + +/* --------------------------------------------------------------------------- + Iteration timeline (design/07): 3-level composition view + --------------------------------------------------------------------------- */ +.tl-tabs { display: flex; gap: 6px; margin-top: 10px; flex-wrap: wrap; } +.tl-tab { + background: var(--panel-2); color: var(--text); border: 1px solid var(--border); + padding: 6px 14px; border-radius: 7px; cursor: pointer; font-size: 13px; +} +.tl-tab.is-active { background: var(--accent); border-color: var(--accent); color: #fff; font-weight: 600; } + +.tl-note { color: var(--muted); font-size: 12px; margin: 0 0 12px; } +.tl-controls { display: flex; align-items: center; gap: 10px; margin-bottom: 12px; flex-wrap: wrap; } + +/* category colours (Level 1 / shared legend) */ +.tl-legend { display: flex; gap: 16px; flex-wrap: wrap; margin: 10px 0 0; font-size: 12px; color: var(--muted); } +.tl-legend i { display: inline-block; width: 12px; height: 12px; border-radius: 3px; margin-right: 5px; vertical-align: -1px; } +.tl-c-attn { background: #4f8cff; } +.tl-c-mlp { background: #36c08f; } +.tl-c-a2a { background: #e0a03a; } +.tl-c-misc { background: #8d97a8; } +.tl-c-emb { background: #a06cd5; } +.tl-c-out { background: #d5675a; } +.tl-c-manual { background: #9a6cff; } + +/* Level 1: per-cr bidirectional stacked bars */ +.tl-l1-row { display: grid; grid-template-columns: 120px 1fr; gap: 12px; align-items: center; margin-bottom: 14px; } +.tl-l1-lab { font-size: 13px; } +.tl-l1-lab small { display: block; color: var(--muted); font-size: 11px; } +.tl-bars { display: flex; flex-direction: column; gap: 6px; } +.tl-bar { display: flex; align-items: center; gap: 8px; } +.tl-bar__tag { width: 34px; font-size: 11px; color: var(--muted); text-align: right; } +.tl-bar__track { position: relative; flex: 1; height: 26px; background: var(--panel-2); border: 1px solid var(--border); border-radius: 6px; overflow: hidden; display: flex; } +.tl-seg { height: 100%; display: flex; align-items: center; justify-content: center; font-size: 10px; color: #08101c; overflow: hidden; white-space: nowrap; cursor: default; } +.tl-seg.tl-c-misc, .tl-seg.tl-c-out, .tl-seg.tl-c-manual { color: #f5f7fb; } +.tl-bar__total { width: 70px; font-size: 12px; font-variant-numeric: tabular-nums; color: var(--text); } + +/* Level 2: per-device layer strips */ +.tl-l2-row { display: grid; grid-template-columns: 150px 1fr 120px; gap: 12px; align-items: center; margin-bottom: 8px; } +.tl-l2-lab { font-size: 12px; } +.tl-l2-lab small { display: block; color: var(--muted); font-size: 11px; } +.tl-strip { display: flex; height: 22px; border: 1px solid var(--border); border-radius: 5px; overflow: hidden; background: var(--panel-2); } +.tl-cell { height: 100%; min-width: 2px; border-right: 1px solid rgba(0,0,0,.25); } +.tl-cell:last-child { border-right: none; } +.tl-cell--recompute { background-image: repeating-linear-gradient(45deg, transparent, transparent 3px, rgba(255,255,255,.35) 3px, rgba(255,255,255,.35) 5px); } +.tl-chunk-gap { width: 3px; background: var(--bg); } +.tl-cell--emb { background: #a06cd5 !important; } +.tl-cell--out { background: #d5675a !important; } +.tl-l2-meta { font-size: 11px; color: var(--muted); font-variant-numeric: tabular-nums; } +.tl-l2-row.is-critical .tl-strip { outline: 2px solid var(--accent-2); outline-offset: 1px; } +.tl-l2-row.is-critical .tl-l2-lab { color: var(--accent-2); } + +/* Level 3: pipeline schedule Gantt */ +.tl-gantt-wrap { overflow-x: auto; } +.tl-gantt { display: block; } +/* text colours are set via inline `fill` attributes in app.js (a stylesheet + `fill` here would override those SVG presentation attributes, e.g. forcing the + in-cell microbatch numbers grey instead of black). */ +.tl-gantt rect.tl-fwd { stroke: rgba(0,0,0,.35); stroke-width: .5; } +.tl-gantt rect.tl-bwd { stroke: rgba(0,0,0,.35); stroke-width: .5; } +.tl-gantt rect.tl-bubble { fill: transparent; } +.tl-summary { font-size: 12px; color: var(--muted); margin-top: 10px; font-variant-numeric: tabular-nums; } +.tl-warn { color: var(--warn); } + +/* layout switch + level tabs on one row */ +.tl-head-row { display: flex; align-items: center; gap: 18px; flex-wrap: wrap; margin-top: 10px; } +.tl-viewswitch { display: flex; align-items: center; gap: 8px; } +.tl-view { + background: var(--panel-2); color: var(--text); border: 1px solid var(--border); + padding: 6px 12px; border-radius: 7px; cursor: pointer; font-size: 13px; +} +.tl-view.is-active { background: var(--accent-2); border-color: var(--accent-2); color: #06231a; font-weight: 600; } + +/* stacked sections */ +.tl-section { border-top: 1px solid var(--border); padding-top: 14px; margin-top: 14px; } +.tl-section:first-of-type { border-top: none; padding-top: 0; margin-top: 4px; } +.tl-section__title { font-size: 14px; margin: 0 0 10px; } +.tl-section__sub { font-size: 12px; color: var(--muted); font-weight: 400; margin-left: 8px; } + +/* linkage selection indicator — floating card, does not affect page layout */ +.tl-selbar { + position: fixed; + right: 24px; + bottom: 24px; + z-index: 60; + display: flex; align-items: center; gap: 14px; + background: rgba(23, 28, 38, .97); + border: 1px solid var(--accent-2); + box-shadow: 0 8px 26px rgba(0, 0, 0, .5); + border-radius: 10px; + padding: 10px 14px; font-size: 12px; color: var(--text); + max-width: 360px; + backdrop-filter: blur(6px); +} +.tl-selbar__txt { line-height: 1.35; } +.tl-clear { + background: var(--panel-2); color: var(--text); border: 1px solid var(--border); + padding: 5px 14px; border-radius: 7px; cursor: pointer; font-size: 12px; white-space: nowrap; +} +.tl-clear:hover { border-color: var(--accent-2); } +@media (max-width: 640px) { .tl-selbar { left: 16px; right: 16px; bottom: 16px; max-width: none; } } + +/* linkage highlight states (shared) */ +.tl-clickable { cursor: pointer; } +.tl-dim { opacity: .32; transition: opacity .12s; } +.tl-l1-row.is-sel, .tl-l2-row.is-sel .tl-strip { outline: 2px solid var(--accent); outline-offset: 2px; border-radius: 6px; } +.tl-l1-row.is-sel { background: rgba(79,140,255,.08); border-radius: 8px; } +.tl-cell--sel { outline: 2px solid #fff; outline-offset: -2px; } + +/* gantt export toolbar */ +.tl-export { display: flex; gap: 8px; margin-bottom: 8px; } +.tl-export .mode-tab { background: var(--panel-2); } + +/* gantt zoom slider */ +.tl-zoom-wrap { display: inline-flex; align-items: center; gap: 8px; margin-left: 6px; } +.tl-zoom { width: 160px; accent-color: var(--accent); } +.tl-zoom-val { font-size: 12px; color: var(--muted); font-variant-numeric: tabular-nums; min-width: 30px; } +.tl-axis-note { margin: 8px 0 0; } diff --git a/examples/deepseek-v4/projection/site/data/flash.json b/examples/deepseek-v4/projection/site/data/flash.json new file mode 100644 index 000000000..31a744e44 --- /dev/null +++ b/examples/deepseek-v4/projection/site/data/flash.json @@ -0,0 +1,1727 @@ +{ + "schema_version": 1, + "model": "flash", + "analytic_flops": { + "per_cr_layer_flops": { + "0": 28775232307200, + "4": 36583482851328, + "128": 29596695134208 + }, + "output_flops": 13013750906880, + "mtp": { + "num_layers": 1, + "compress_ratio": 4, + "inner_layer_flops": 36583482851328, + "eh_proj_flops": 824633720832, + "extra_logits_flops": 13013750906880, + "hc_head_flops": 3221225472 + }, + "seq": 4096, + "note": "per-layer and MTP (B=1) Megatron-convention FLOPs at capture seq; site multiplies by GA x DP" + }, + "generated_at": "2026-06-30T13:59:46Z", + "provenance": { + "traces": { + "0": "/apps/tas/0_public/data/traces/dsv4_projection/projection_flash_cr0_seq4096_ep8/tensorboard/primus-megatron-exp[projection_flash_cr0_seq4096_ep8]-rank[0].1782149271209635673.pt.trace.json", + "4": "/apps/tas/0_public/data/traces/dsv4_projection/projection_flash_cr4_seq4096_ep8/tensorboard/primus-megatron-exp[projection_flash_cr4_seq4096_ep8]-rank[0].1782148745954636194.pt.trace.json", + "128": "/apps/tas/0_public/data/traces/dsv4_projection/projection_flash_cr128_seq4096_ep8/tensorboard/primus-megatron-exp[projection_flash_cr128_seq4096_ep8]-rank[0].1782149357925690324.pt.trace.json" + }, + "graphed_crs_estimated_from_cr4": [], + "dropped_stall_us_per_mb": { + "0": 0.0, + "4": 19280.7, + "128": 0.0 + }, + "note": "cr in graphed_crs_estimated_from_cr4 were CUDA-graph/stream-captured (compute not visible in trace); their breakdown is copied from cr=4 as an estimate. Re-run those cr with graph capture disabled for exact numbers. dropped_stall_us_per_mb: per-mb layer-compute kernel time dropped as implausible one-off device stalls (> _MAX_PLAUSIBLE_LAUNCH_US)." + }, + "capture": { + "gpu": "MI355X", + "seq_length": 4096, + "micro_batch_size": 1, + "tokens_per_microbatch": 4096, + "ep": 8, + "ga_for_capture": 2, + "optimizer": "adam", + "distributed_optimizer": false, + "recompute": "off", + "measured_iter_time_ms": 4076.0, + "measured_anchor": { + "config": "PP8/VPP1/EP8/TP1, world 64, MBS1/GBS128, seq4096, full recompute, layout Et*4|t*5|(t*6|)*5,t*4mL (8-node MI355X)", + "iter_ms": 4076.0, + "tflops_gpu": 722, + "tok_s_gpu": 2010, + "calib_factor": 0.87, + "tolerance": "<0.1%" + } + }, + "model_config": { + "num_layers": 43, + "hidden_size": 4096, + "num_attention_heads": 64, + "kv_channels": 512, + "num_experts": 256, + "moe_router_topk": 6, + "moe_ffn_hidden_size": 2048, + "moe_shared_expert_intermediate_size": 2048, + "index_topk": 512, + "vocab_size": 129280, + "mtp_num_layers": 1, + "mtp_compress_ratios": [ + 4 + ], + "pipeline_layout": "Et*10|t*11|t*11|t*11mL", + "compress_ratios": [ + 0, + 0, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 0 + ], + "cr_layer_counts": { + "0": 3, + "4": 20, + "128": 20 + }, + "total_params": 290419900416 + }, + "hardware": { + "MI355X": { + "peak_tflops_bf16": 2500.0, + "hbm_bandwidth_gbps": 8000.0 + }, + "MI455X": { + "peak_tflops_bf16": 10000.0, + "hbm_bandwidth_gbps": 19600.0 + } + }, + "layers": { + "0": { + "attention": { + "forward": [ + { + "module": "attn.misc", + "time_us": 2768.835, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", + "time_us": 600.524 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor", + "time_us": 252.844 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 190.953 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 151.259 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", + "time_us": 136.028 + } + ] + }, + { + "module": "attn.proj", + "time_us": 732.222, + "class": "compute_bound", + "flop_class": "gemm", + "flops": 449360953344.0, + "tflops": 613.7, + "kernels": [ + { + "name": "Cijk_Ailk_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x16x32_MI16x16x1_SN_LDSB1_AFC1", + "time_us": 260.835 + }, + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", + "time_us": 226.171 + }, + { + "name": "Custom_Cijk_Alik_Bljk_BBS_BH_MT256x256x64_MI16x16x1_UserArgs_shortname1_gfx950", + "time_us": 96.207 + }, + { + "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT32x64x128_MI16x16x1_SN_LDSB1_AFC1_AF", + "time_us": 56.9 + }, + { + "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT64x64x128_MI16x16x1_SN_LDSB1_AFC1_AF", + "time_us": 36.157 + }, + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x128x64_MI16x16x1_SN_LDSB1_AFC", + "time_us": 29.943 + } + ] + }, + { + "module": "attn.norm", + "time_us": 496.029, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 93.8 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", + "time_us": 92.809 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", + "time_us": 83.298 + }, + { + "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp(HIP_vector_type(int const*, int*, int ", + "time_us": 111.084 + }, + { + "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", + "time_us": 19.553 + } + ] + }, + { + "module": "moe.combine", + "time_us": 321.999, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::combine(hip_bfloat", + "time_us": 231.632 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 90.367 + } + ] + }, + { + "module": "moe.shared_expert", + "time_us": 112.521, + "class": "compute_bound", + "flop_class": "gemm", + "flops": 68719476736.0, + "tflops": 610.7, + "kernels": [ + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", + "time_us": 78.11 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 21.219 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 5.396 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor", + "time_us": 8.02 + }, + { + "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc", + "time_us": 6.869 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", + "time_us": 5.616 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 4.652 + } + ] + } + ], + "backward": [ + { + "module": "moe.grouped_gemm", + "time_us": 1916.666, + "class": "compute_bound", + "flop_class": "grouped_gemm", + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", + "time_us": 115.553 + }, + { + "name": "void primus_turbo::compute_grouped_gemm_args(hip_bfloat", + "time_us": 238.449 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 59.93 + } + ] + }, + { + "module": "moe.dispatch", + "time_us": 242.331, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024>(HIP_vector_type(int const*, int", + "time_us": 20.566 + } + ] + }, + { + "module": "moe.router", + "time_us": 44.476, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "_sinkhorn_bwd_kernel", + "time_us": 44.476 + } + ] + } + ] + } + }, + "4": { + "attention": { + "forward": [ + { + "module": "attn.misc", + "time_us": 2773.299, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", + "time_us": 603.359 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor", + "time_us": 247.801 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 192.373 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 150.529 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", + "time_us": 136.255 + } + ] + }, + { + "module": "attn.indexer", + "time_us": 797.919, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 152.083 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 98.513 + }, + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x32_MI16x16x1_SN_LDSB1_AFC", + "time_us": 86.043 + }, + { + "name": "void at::native::reduce_kernel<128, 4, at::native::ReduceOp(hip_bfloat", + "time_us": 230.235 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 100.447 + } + ] + }, + { + "module": "moe.dispatch", + "time_us": 279.761, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024>(HIP_vector_type(int const*, int*, int ", + "time_us": 27.653 + }, + { + "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", + "time_us": 19.776 + } + ] + }, + { + "module": "moe.shared_expert", + "time_us": 112.854, + "class": "compute_bound", + "flop_class": "gemm", + "flops": 68719476736.0, + "tflops": 608.9, + "kernels": [ + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", + "time_us": 78.307 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 21.372 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 5.393 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor", + "time_us": 8.086 + }, + { + "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc", + "time_us": 6.905 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", + "time_us": 5.669 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 4.696 + } + ] + } + ], + "backward": [ + { + "module": "moe.grouped_gemm", + "time_us": 1896.364, + "class": "compute_bound", + "flop_class": "grouped_gemm", + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", + "time_us": 114.644 + }, + { + "name": "void primus_turbo::compute_grouped_gemm_variable_k_args(HIP_vector_type(int const*, int", + "time_us": 178.088 + } + ] + }, + { + "module": "moe.combine", + "time_us": 357.633, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::combine(hip_bfloat", + "time_us": 237.535 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 120.097 + } + ] + }, + { + "module": "moe.router", + "time_us": 44.726, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "_sinkhorn_bwd_kernel", + "time_us": 44.726 + } + ] + } + ] + } + }, + "128": { + "attention": { + "forward": [ + { + "module": "attn.misc", + "time_us": 2784.297, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", + "time_us": 605.113 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::AUnaryFunctor", + "time_us": 248.877 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 191.946 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 151.909 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", + "time_us": 136.028 + } + ] + }, + { + "module": "attn.norm", + "time_us": 851.58, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::(anonymous namespace)::CatArrayBatchedCopy(HIP_vector_type(int const*, int*, int ", + "time_us": 60.447 + }, + { + "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", + "time_us": 20.186 + } + ] + }, + { + "module": "moe.combine", + "time_us": 309.929, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::combine(hip_bfloat", + "time_us": 227.382 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 82.547 + } + ] + }, + { + "module": "moe.shared_expert", + "time_us": 112.514, + "class": "compute_bound", + "flop_class": "gemm", + "flops": 68719476736.0, + "tflops": 610.8, + "kernels": [ + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", + "time_us": 77.993 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 21.266 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 5.43 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor", + "time_us": 14.203 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16tofloat32_", + "time_us": 10.482 + }, + { + "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc", + "time_us": 7.482 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", + "time_us": 6.356 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 4.886 + } + ] + } + ], + "backward": [ + { + "module": "moe.grouped_gemm", + "time_us": 1888.822, + "class": "compute_bound", + "flop_class": "grouped_gemm", + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", + "time_us": 113.277 + }, + { + "name": "void primus_turbo::compute_grouped_gemm_variable_k_args(hip_bfloat", + "time_us": 233.195 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 110.397 + } + ] + }, + { + "module": "moe.dispatch", + "time_us": 225.508, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024>(HIP_vector_type(int const*, int", + "time_us": 4.29 + } + ] + }, + { + "module": "moe.router", + "time_us": 44.606, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "_sinkhorn_bwd_kernel", + "time_us": 44.606 + } + ] + } + ] + } + } + }, + "non_layer": { + "embedding": { + "forward": [ + { + "module": "embedding", + "time_us": 7.84, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::vectorized_gather_kernel<16, long>(char*, char*, long*, int, lo", + "time_us": 7.84 + } + ] + } + ], + "backward": [ + { + "module": "embedding", + "time_us": 86.591, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::(anonymous namespace)::compute_grad_weight", + "time_us": 33.993 + }, + { + "name": "void rocprim::ROCPRIM_400200_NS::detail::trampoline_kernel(lon", + "time_us": 21.21 + }, + { + "name": "void rocprim::ROCPRIM_400200_NS::detail::init_lookback_scan_state_kernel(long*, ", + "time_us": 2.08 + }, + { + "name": "void at::native::(anonymous namespace)::krn_partials_per_segment(long*, lo", + "time_us": 2.023 + } + ] + } + ] + }, + "output": { + "forward": [ + { + "module": "output", + "time_us": 1553.935, + "class": "compute_bound", + "flop_class": "gemm", + "flops": 2171374403584.0, + "tflops": 1397.3, + "kernels": [ + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", + "time_us": 1529.095 + }, + { + "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT16x16x512_MI16x16x1_SN_LDSB1_AFC1_AF", + "time_us": 24.84 + } + ] + } + ], + "backward": [] + }, + "loss": { + "forward": [ + { + "module": "loss", + "time_us": 368.303, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "cross_entropy_kernel", + "time_us": 235.982 + }, + { + "name": "online_softmax_kernel", + "time_us": 132.321 + } + ] + } + ], + "backward": [ + { + "module": "loss", + "time_us": 197.392, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "element_mul_kernel", + "time_us": 197.392 + } + ] + } + ] + } + }, + "optimizer": { + "type": "adam", + "measured_params": null, + "time_us": 18459.0, + "bytes_per_param": 30, + "class": "memory_bound", + "note": "Adam mixed-precision step traffic in bytes/param; measured one-layer optimizer-step kernel time is a sanity reference" + }, + "comm": { + "ep_dispatch_us": null, + "ep_combine_us": null, + "note": "EP dispatch/combine are memory_bound rows inside moe; informational" + } +} diff --git a/examples/deepseek-v4/projection/site/data/pro.json b/examples/deepseek-v4/projection/site/data/pro.json new file mode 100644 index 000000000..8cfd34cb4 --- /dev/null +++ b/examples/deepseek-v4/projection/site/data/pro.json @@ -0,0 +1,1750 @@ +{ + "schema_version": 1, + "model": "pro", + "analytic_flops": { + "per_cr_layer_flops": { + "0": 76885870510080, + "4": 95420801875968, + "128": 78425716948992 + }, + "output_flops": 22774064087040, + "mtp": { + "num_layers": 1, + "compress_ratio": 4, + "inner_layer_flops": 95420801875968, + "eh_proj_flops": 2525440770048, + "extra_logits_flops": 22774064087040, + "hc_head_flops": 5637144576 + }, + "seq": 4096, + "note": "per-layer and MTP (B=1) Megatron-convention FLOPs at capture seq; site multiplies by GA x DP" + }, + "generated_at": "2026-06-30T13:59:43Z", + "provenance": { + "traces": { + "0": "/apps/tas/wenx/workspace/Primus-deepseek-v4/output/amd/tas-mi355x-20260618/projection_pro_cr0_seq4096_ep8/tensorboard/primus-megatron-exp[projection_pro_cr0_seq4096_ep8]-rank[0].1781792532762897620.pt.trace.json", + "4": "/apps/tas/wenx/workspace/Primus-deepseek-v4/output/amd/tas-mi355x-20260618/projection_pro_cr4_seq4096_ep8/tensorboard/primus-megatron-exp[projection_pro_cr4_seq4096_ep8]-rank[0].1781792606762968569.pt.trace.json", + "128": "/apps/tas/wenx/workspace/Primus-deepseek-v4/output/amd/tas-mi355x-20260618/projection_pro_cr128_seq4096_ep8/tensorboard/primus-megatron-exp[projection_pro_cr128_seq4096_ep8]-rank[0].1781792680203624431.pt.trace.json" + }, + "graphed_crs_estimated_from_cr4": [], + "dropped_stall_us_per_mb": { + "0": 36705.7, + "4": 0.0, + "128": 0.0 + }, + "note": "cr in graphed_crs_estimated_from_cr4 were CUDA-graph/stream-captured (compute not visible in trace); their breakdown is copied from cr=4 as an estimate. Re-run those cr with graph capture disabled for exact numbers. dropped_stall_us_per_mb: per-mb layer-compute kernel time dropped as implausible one-off device stalls (> _MAX_PLAUSIBLE_LAUNCH_US)." + }, + "capture": { + "gpu": "MI355X", + "seq_length": 4096, + "micro_batch_size": 1, + "tokens_per_microbatch": 4096, + "ep": 8, + "ga_for_capture": 2, + "optimizer": "adam", + "distributed_optimizer": false, + "recompute": "off", + "measured_iter_time_ms": null, + "measured_anchor": null + }, + "model_config": { + "num_layers": 61, + "hidden_size": 7168, + "num_attention_heads": 128, + "kv_channels": 512, + "num_experts": 384, + "moe_router_topk": 6, + "moe_ffn_hidden_size": 3072, + "moe_shared_expert_intermediate_size": 3072, + "index_topk": 1024, + "vocab_size": 129280, + "mtp_num_layers": 1, + "mtp_compress_ratios": [ + 4 + ], + "pipeline_layout": "", + "compress_ratios": [ + 128, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 4, + 128, + 0 + ], + "cr_layer_counts": { + "0": 1, + "4": 29, + "128": 31 + }, + "total_params": 1597579198464 + }, + "hardware": { + "MI355X": { + "peak_tflops_bf16": 2500.0, + "hbm_bandwidth_gbps": 8000.0 + }, + "MI455X": { + "peak_tflops_bf16": 10000.0, + "hbm_bandwidth_gbps": 19600.0 + } + }, + "layers": { + "0": { + "attention": { + "forward": [ + { + "module": "attn.misc", + "time_us": 7402.198, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", + "time_us": 1437.355 + }, + { + "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp", + "time_us": 709.616 + }, + { + "name": "void at::native::unrolled_elementwise_kernel(hip_", + "time_us": 386.653 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 41.18 + } + ] + }, + { + "module": "moe.dispatch", + "time_us": 422.925, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type(int const*, int*, int ", + "time_us": 20.72 + }, + { + "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", + "time_us": 19.396 + } + ] + }, + { + "module": "moe.shared_expert", + "time_us": 252.402, + "class": "compute_bound", + "flop_class": "gemm", + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT192x256x64_MI16x16x1_SN_LDSB0_AFC", + "time_us": 141.417 + }, + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", + "time_us": 71.463 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 21.462 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 7.096 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor", + "time_us": 14.06 + }, + { + "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc", + "time_us": 7.365 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 6.666 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", + "time_us": 5.559 + } + ] + } + ], + "backward": [ + { + "module": "moe.grouped_gemm", + "time_us": 5933.968, + "class": "compute_bound", + "flop_class": "grouped_gemm", + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", + "time_us": 423.883 + }, + { + "name": "void primus_turbo::compute_grouped_gemm_variable_k_args(hip_", + "time_us": 394.596 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 192.271 + } + ] + }, + { + "module": "moe.dispatch", + "time_us": 449.116, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type(int const*, int", + "time_us": 73.65 + } + ] + }, + { + "module": "moe.router", + "time_us": 44.656, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "_sinkhorn_bwd_kernel", + "time_us": 44.656 + } + ] + } + ] + } + }, + "4": { + "attention": { + "forward": [ + { + "module": "attn.misc", + "time_us": 7425.221, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", + "time_us": 1447.247 + }, + { + "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp", + "time_us": 709.102 + }, + { + "name": "void at::native::unrolled_elementwise_kernel(HIP_vector_type(int const*, int*, int ", + "time_us": 82.547 + }, + { + "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", + "time_us": 19.383 + } + ] + }, + { + "module": "moe.combine", + "time_us": 470.04, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::combine(hip_", + "time_us": 385.316 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 84.724 + } + ] + }, + { + "module": "moe.shared_expert", + "time_us": 254.469, + "class": "compute_bound", + "flop_class": "gemm", + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT192x256x64_MI16x16x1_SN_LDSB0_AFC", + "time_us": 143.571 + }, + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", + "time_us": 71.73 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 21.249 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 6.96 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor", + "time_us": 13.973 + }, + { + "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc", + "time_us": 7.535 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 6.562 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", + "time_us": 5.852 + } + ] + } + ], + "backward": [ + { + "module": "moe.grouped_gemm", + "time_us": 5908.193, + "class": "compute_bound", + "flop_class": "grouped_gemm", + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", + "time_us": 425.279 + }, + { + "name": "void primus_turbo::compute_grouped_gemm_variable_k_args(HIP_vector_type(int const*, int", + "time_us": 204.325 + } + ] + }, + { + "module": "moe.combine", + "time_us": 503.433, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::combine(hip_", + "time_us": 394.109 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 109.324 + } + ] + }, + { + "module": "moe.router", + "time_us": 44.866, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "_sinkhorn_bwd_kernel", + "time_us": 44.866 + } + ] + } + ] + } + }, + "128": { + "attention": { + "forward": [ + { + "module": "attn.misc", + "time_us": 7410.422, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::bfloat16_copy_kern", + "time_us": 1443.783 + }, + { + "name": "void at::native::reduce_kernel<512, 1, at::native::ReduceOp", + "time_us": 708.356 + }, + { + "name": "void at::native::unrolled_elementwise_kernel(hip_", + "time_us": 387.153 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 127.304 + } + ] + }, + { + "module": "moe.dispatch", + "time_us": 448.916, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type(int const*, int*, int ", + "time_us": 45.283 + }, + { + "name": "void primus_turbo::deep_ep::layout::get_dispatch_layout<256, 4, 8>(long const*, ", + "time_us": 19.383 + } + ] + }, + { + "module": "moe.shared_expert", + "time_us": 254.795, + "class": "compute_bound", + "flop_class": "gemm", + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT192x256x64_MI16x16x1_SN_LDSB0_AFC", + "time_us": 143.427 + }, + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", + "time_us": 71.993 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 8, at::native::gpu_kernel", + "time_us": 21.433 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::(anonymous namespa", + "time_us": 7.026 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<8, at::native::BinaryFunctor", + "time_us": 13.92 + }, + { + "name": "void at::native::_scatter_gather_elementwise_kernel<256, 4, at::native::_cuda_sc", + "time_us": 7.152 + }, + { + "name": "void at::native::vectorized_elementwise_kernel<4, at::native::(anonymous namespa", + "time_us": 6.402 + }, + { + "name": "void at::native::elementwise_kernel_manual_unroll<128, 4, at::native::gpu_kernel", + "time_us": 5.779 + } + ] + } + ], + "backward": [ + { + "module": "moe.grouped_gemm", + "time_us": 5906.347, + "class": "compute_bound", + "flop_class": "grouped_gemm", + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void ck_tile::kentry<1, ck_tile::GroupedGemmKernel(std::bfl", + "time_us": 423.59 + }, + { + "name": "void primus_turbo::compute_grouped_gemm_variable_k_args(hip_", + "time_us": 394.309 + }, + { + "name": "void primus_turbo::deep_ep::intranode::cached_notify_combine<8>(void**, int*, in", + "time_us": 84.787 + } + ] + }, + { + "module": "moe.dispatch", + "time_us": 408.969, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void primus_turbo::deep_ep::intranode::dispatch<8, 1024, true>(HIP_vector_type(int const*, int", + "time_us": 33.353 + } + ] + }, + { + "module": "moe.router", + "time_us": 45.489, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "_sinkhorn_bwd_kernel", + "time_us": 45.489 + } + ] + } + ] + } + } + }, + "non_layer": { + "embedding": { + "forward": [ + { + "module": "embedding", + "time_us": 14.566, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::vectorized_gather_kernel<16, long>(char*, char*, long*, int, lo", + "time_us": 14.566 + } + ] + } + ], + "backward": [ + { + "module": "embedding", + "time_us": 129.501, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "void at::native::(anonymous namespace)::compute_grad_weight", + "time_us": 60.12 + }, + { + "name": "void at::native::(anonymous namespace)::sum_and_scatter(lon", + "time_us": 38.046 + }, + { + "name": "void rocprim::ROCPRIM_400200_NS::detail::trampoline_kernel(long*, lo", + "time_us": 2.026 + }, + { + "name": "void at::native::(anonymous namespace)::krn_partial_segment_offset(long*, ", + "time_us": 2.023 + } + ] + } + ] + }, + "output": { + "forward": [ + { + "module": "output", + "time_us": 2688.195, + "class": "compute_bound", + "flop_class": "gemm", + "flops": 3799905206272.0, + "tflops": 1413.6, + "kernels": [ + { + "name": "Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0", + "time_us": 2634.481 + }, + { + "name": "Cijk_Alik_Bljk_S_B_Bias_HA_S_SAV_UserArgs_MT16x16x512_MI16x16x1_SN_LDSB1_AFC1_AF", + "time_us": 53.713 + } + ] + } + ], + "backward": [] + }, + "loss": { + "forward": [ + { + "module": "loss", + "time_us": 350.465, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "cross_entropy_kernel", + "time_us": 222.888 + }, + { + "name": "online_softmax_kernel", + "time_us": 127.577 + } + ] + } + ], + "backward": [ + { + "module": "loss", + "time_us": 195.848, + "class": "memory_bound", + "flop_class": null, + "flops": null, + "tflops": null, + "kernels": [ + { + "name": "element_mul_kernel", + "time_us": 195.848 + } + ] + } + ] + } + }, + "optimizer": { + "type": "adam", + "measured_params": null, + "time_us": 87037.8, + "bytes_per_param": 30, + "class": "memory_bound", + "note": "Adam mixed-precision step traffic in bytes/param; measured one-layer optimizer-step kernel time is a sanity reference" + }, + "comm": { + "ep_dispatch_us": null, + "ep_combine_us": null, + "note": "EP dispatch/combine are memory_bound rows inside moe; informational" + } +} diff --git a/examples/deepseek-v4/projection/site/index.html b/examples/deepseek-v4/projection/site/index.html new file mode 100644 index 000000000..47c29c364 --- /dev/null +++ b/examples/deepseek-v4/projection/site/index.html @@ -0,0 +1,134 @@ + + + + + + DeepSeek-V4 Performance Projection + + + +
+
+

Primus · DeepSeek-V4

+

Training Performance Projection

+

+ Trace-driven single-layer breakdown on MI355X, scaled to a full-model, + multi-GPU projection. Page 1 is measured-MI355X; page 2 scales to MI455X + by theoretical hardware ratios. +

+
+
+ +
+
+ + + +
+
+ +
+ + +
+ + +
+

Model configuration

+
+
+
+ +
+
+

Per-layer breakdown

+

+ Forward (left→right) then backward (right→left). Time in µs / one + microbatch (seq 4096). The TFLOP/s row is per-kernel achieved + (gemm / grouped_gemm / attn only); the headline TFLOP/s/GPU below uses + the V4 analytic model FLOPs instead. +

+ +
+
+
+ +
+

Projected throughput

+
+
+
+ +
+
+

Iteration timeline

+

+ How one iteration's time is composed, bottom-up: a single layer + (attn / mlp / a2a) → each pipeline rank's chunks (layer granularity) → + the whole 1F1B pipeline schedule. Reacts live to the controls on the left. + See design/07-iteration-timeline.md. +

+
+
+ Layout +
+ + +
+
+
+ + + +
+
+
+
+
+
+
+ +
+

+ Generated from examples/deepseek-v4/projection/. Methodology & + assumptions in design/. Numbers are only as good as the input + trace; see the assumptions list before citing. +

+
+ + + + diff --git a/examples/deepseek-v4/projection/tools/gen_mock_data.py b/examples/deepseek-v4/projection/tools/gen_mock_data.py new file mode 100755 index 000000000..1b223a619 --- /dev/null +++ b/examples/deepseek-v4/projection/tools/gen_mock_data.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +"""Generate MOCK breakdown JSON for site development / demo. + +Numbers are placeholders in the right *order of magnitude*, seeded from the +published P57 single-layer attention micro-bench (V4-Flash widths: B=1, H=64, +Sq=4096, D=512) and the P40 EP=8 MoE/kernel attribution. They are NOT measured +ground truth — replace with `parse_trace.py` output once real traces exist +(every file is marked provenance.mock = true). + +Usage: + python3 examples/deepseek-v4/projection/tools/gen_mock_data.py +""" + +from __future__ import annotations + +import json +from datetime import datetime, timezone +from pathlib import Path + +OUT_DIR = Path(__file__).resolve().parent.parent / "site" / "data" + +PRO_COMPRESS = [128, 128] + [4 if i % 2 == 0 else 128 for i in range(2, 60)] + [0] +FLASH_COMPRESS = [0, 0] + [4 if i % 2 == 0 else 128 for i in range(2, 42)] + [0] + +MODEL_CONFIGS = { + "pro": { + "num_layers": 61, + "hidden_size": 7168, + "num_attention_heads": 128, + "kv_channels": 512, + "num_experts": 384, + "moe_router_topk": 6, + "moe_ffn_hidden_size": 3072, + "moe_shared_expert_intermediate_size": 3072, + "index_topk": 1024, + "vocab_size": 129280, + "compress_ratios": PRO_COMPRESS, + }, + "flash": { + "num_layers": 43, + "hidden_size": 4096, + "num_attention_heads": 64, + "kv_channels": 512, + "num_experts": 256, + "moe_router_topk": 6, + "moe_ffn_hidden_size": 2048, + "moe_shared_expert_intermediate_size": 2048, + "index_topk": 512, + "vocab_size": 129280, + "compress_ratios": FLASH_COMPRESS, + }, +} + +# MI355X: BF16 matrix 2.5 PFLOPS, HBM3E 8 TB/s (AMD product page). +# MI455X (MI400): HBM4 19.6 TB/s; BF16 dense not officially published — +# estimated ~10 PFLOPS (half of the 20 PFLOPS FP8 spec). +HARDWARE = { + "MI355X": {"peak_tflops_bf16": 2500.0, "hbm_bandwidth_gbps": 8000.0}, + "MI455X": {"peak_tflops_bf16": 10000.0, "hbm_bandwidth_gbps": 19600.0}, +} + + +def row(module, time_us, flop_class=None, tflops=None): + flops = (tflops * time_us * 1e6) if (flop_class and tflops) else None + return { + "module": module, + "time_us": round(time_us, 1), + "class": "compute_bound" if flop_class else "memory_bound", + "flop_class": flop_class, + "flops": flops, + "tflops": tflops, + "kernels": [], + } + + +# Base (flash) attention core fwd/bwd by cr, from P57 micro-bench (ms -> us). +ATTN_CORE = { + "0": {"fwd": 500.0, "bwd": 2080.0}, + "4": {"fwd": 1430.0, "bwd": 5110.0}, + "128": {"fwd": 570.0, "bwd": 2810.0}, +} + + +def attention_breakdown(cr, s): + """s = linear scale factor vs flash widths.""" + fwd = [ + row("attn.qkv_proj", 220 * s, "gemm", 480), + row("attn.core", ATTN_CORE[cr]["fwd"] * s, "attn", 210), + row("attn.rope", 35 * s), + row("attn.o_proj", 160 * s, "gemm", 470), + row("attn.norm", 25 * s), + ] + bwd = [ + row("attn.qkv_proj", 440 * s, "gemm", 480), + row("attn.core", ATTN_CORE[cr]["bwd"] * s, "attn", 180), + row("attn.rope", 45 * s), + row("attn.o_proj", 320 * s, "gemm", 470), + row("attn.norm", 30 * s), + ] + if cr == "4": # CSA uses the Indexer/Compressor + fwd.insert(2, row("attn.indexer", 300 * s, "gemm", 300)) + bwd.insert(2, row("attn.indexer", 600 * s, "gemm", 300)) + return {"forward": fwd, "backward": bwd} + + +def moe_breakdown(sg, sc): + """sg = grouped-gemm scale, sc = comm/act scale vs flash.""" + fwd = [ + row("moe.router", 55 * sc), + row("moe.dispatch", 820 * sc), + row("moe.grouped_gemm", 1850 * sg, "grouped_gemm", 430), + row("moe.act", 110 * sc), + row("moe.shared_expert", 210 * sg, "gemm", 460), + row("moe.combine", 990 * sc), + ] + bwd = [ + row("moe.router", 75 * sc), + row("moe.dispatch", 900 * sc), + row("moe.grouped_gemm", 3700 * sg, "grouped_gemm", 430), + row("moe.act", 150 * sc), + row("moe.shared_expert", 420 * sg, "gemm", 460), + row("moe.combine", 1050 * sc), + ] + return {"forward": fwd, "backward": bwd} + + +def non_layer(s): + return { + "embedding": {"forward": [row("embedding", 120 * s)], "backward": [row("embedding", 60 * s)]}, + "output": { + "forward": [row("output", 900 * s, "gemm", 300)], + "backward": [row("output", 1800 * s, "gemm", 300)], + }, + "loss": {"forward": [row("loss", 80)], "backward": [row("loss", 60)]}, + } + + +def cr_counts(compress): + c = {"0": 0, "4": 0, "128": 0} + for x in compress: + c[str(x)] += 1 + return c + + +def build(model): + cfg = MODEL_CONFIGS[model] + # scale factors vs flash baseline + if model == "pro": + s_attn, s_grouped, s_comm, s_out = 1.75, 2.0, 1.3, 1.75 + opt_params, opt_time = 2_300_000_000, 1600.0 + else: + s_attn, s_grouped, s_comm, s_out = 1.0, 1.0, 1.0, 1.0 + opt_params, opt_time = 1_000_000_000, 850.0 + + layers = { + cr: {"attention": attention_breakdown(cr, s_attn), "moe": moe_breakdown(s_grouped, s_comm)} + for cr in ("0", "4", "128") + } + return { + "schema_version": 1, + "model": model, + "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), + "provenance": { + "mock": True, + "note": "MOCK data (P57/P40 order-of-magnitude); replace with parse_trace.py output", + }, + "capture": { + "gpu": "MI355X", + "seq_length": 4096, + "micro_batch_size": 1, + "tokens_per_microbatch": 4096, + "ep": 8, + "ga_for_capture": 2, + "optimizer": "adam", + "distributed_optimizer": True, + "recompute": "off", + "measured_iter_time_ms": None, + }, + "model_config": {**cfg, "cr_layer_counts": cr_counts(cfg["compress_ratios"])}, + "hardware": HARDWARE, + "layers": layers, + "non_layer": non_layer(s_out), + "optimizer": { + "type": "adam", + "measured_params": opt_params, + "time_us": opt_time, + "bytes_per_param": 18, + "class": "memory_bound", + }, + "comm": { + "ep_dispatch_us": None, + "ep_combine_us": None, + "note": "EP dispatch/combine included in moe rows; informational", + }, + } + + +def main(): + OUT_DIR.mkdir(parents=True, exist_ok=True) + for model in ("pro", "flash"): + doc = build(model) + path = OUT_DIR / f"{model}.json" + path.write_text(json.dumps(doc, indent=2)) + print(f"[gen_mock_data] wrote {path}") + + +if __name__ == "__main__": + main() diff --git a/examples/deepseek-v4/projection/tools/kernel_module_map.py b/examples/deepseek-v4/projection/tools/kernel_module_map.py new file mode 100644 index 000000000..fd6ddbcea --- /dev/null +++ b/examples/deepseek-v4/projection/tools/kernel_module_map.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +"""Kernel / nn.module -> logical-module + flop-class mapping for the V4 +projection breakdown. + +Two independent classifications are provided: + +1. ``module_from_stack(stack)`` — primary: derive the logical module from the + python call stack captured by ``with_stack=True`` (matches nn.module class + names appearing in the stack frames). This is the accurate path (A13). + +2. ``module_from_kernel(name)`` — fallback: derive the logical module purely from + the GPU kernel name when no usable stack is attached. + +``flop_class_from_kernel(name)`` returns the compute-bound FLOP class +(``gemm`` / ``grouped_gemm`` / ``attn``) or ``None`` (memory-bound, A14). + +The rules are intentionally data-driven so they can be extended as kernels are +renamed. Order matters: the first matching rule wins. +""" + +from __future__ import annotations + +# --- logical module taxonomy (kept small per the design) -------------------- +# attention sub-modules: +# attn.qkv_proj, attn.rope, attn.indexer, attn.core, attn.o_proj, attn.norm +# moe sub-modules: +# moe.router, moe.dispatch, moe.grouped_gemm, moe.act, moe.combine, +# moe.shared_expert +# non-layer: embedding, output, loss + +# (substring, logical_module) — matched against the python call stack text. +# nn.module class names / function names seen in V4 forward stacks. +STACK_MODULE_RULES: list[tuple[str, str]] = [ + ("Indexer", "attn.indexer"), + ("Compressor", "attn.indexer"), + ("apply_rotary", "attn.rope"), + ("rope", "attn.rope"), + ("linear_qkv", "attn.qkv_proj"), + ("q_layernorm", "attn.norm"), + ("k_layernorm", "attn.norm"), + ("linear_proj", "attn.o_proj"), + ("o_proj", "attn.o_proj"), + ("core_attention", "attn.core"), + ("DeepseekV4Attention", "attn.core"), + ("MLASelfAttention", "attn.core"), + ("SelfAttention", "attn.core"), + ("input_layernorm", "attn.norm"), + ("pre_mlp_layernorm", "moe.router"), + ("TopKRouter", "moe.router"), + ("router", "moe.router"), + ("sinkhorn", "moe.router"), + ("shared_expert", "moe.shared_expert"), + ("token_dispatch", "moe.dispatch"), + ("dispatch", "moe.dispatch"), + ("combine", "moe.combine"), + ("GroupedMLP", "moe.grouped_gemm"), + ("SequentialMLP", "moe.grouped_gemm"), + ("grouped", "moe.grouped_gemm"), + ("experts", "moe.grouped_gemm"), + ("activation", "moe.act"), + ("swiglu", "moe.act"), + ("MoELayer", "moe.grouped_gemm"), + ("word_embeddings", "embedding"), + ("embedding", "embedding"), + ("output_layer", "output"), + ("lm_head", "output"), + ("loss", "loss"), + ("cross_entropy", "loss"), +] + +# (substring, logical_module) — matched against the GPU kernel name (fallback). +KERNEL_MODULE_RULES: list[tuple[str, str]] = [ + ("_v4_csa_attention", "attn.core"), + ("_v4_attention", "attn.core"), + ("_hc_compute", "attn.core"), + ("hc_compute", "attn.core"), + ("_indexer", "attn.indexer"), + ("indexer", "attn.indexer"), + ("_compressor", "attn.indexer"), + ("compressor", "attn.indexer"), + ("apply_rope", "attn.rope"), + ("rotary", "attn.rope"), + ("rope", "attn.rope"), + ("_sinkhorn", "moe.router"), + ("sinkhorn", "moe.router"), + ("_v4_router", "moe.router"), + ("deep_ep::", "moe.dispatch"), # refined to dispatch/combine below + ("dispatch", "moe.dispatch"), + ("combine", "moe.combine"), + ("GroupedGemm", "moe.grouped_gemm"), + ("_stack_grouped_weight", "moe.grouped_gemm"), + ("group_gemm", "moe.grouped_gemm"), + ("swiglu", "moe.act"), + ("embedding", "embedding"), + ("cross_entropy", "loss"), + ("nll_loss", "loss"), +] + +# (substring, flop_class) — matched against the GPU kernel name. +# Order matters: grouped GEMM must be checked before generic GEMM. +FLOP_CLASS_RULES: list[tuple[str, str]] = [ + ("GroupedGemmKernel", "grouped_gemm"), + ("grouped_gemm", "grouped_gemm"), + ("group_gemm", "grouped_gemm"), + ("_v4_csa_attention", "attn"), + ("_v4_attention", "attn"), + ("attention_fwd", "attn"), + ("attention_bwd", "attn"), + # generic dense GEMM kernels (hipBLASLt / rocBLAS / CK tile / Triton matmul) + ("GemmKernel", "gemm"), + ("Cijk_", "gemm"), + ("gemm", "gemm"), + ("matmul", "gemm"), +] + + +def module_from_stack(stack: str | None) -> str | None: + """Return the logical module from a python call-stack string, or None.""" + if not stack: + return None + for needle, module in STACK_MODULE_RULES: + if needle in stack: + return module + return None + + +def module_from_kernel(name: str) -> str: + """Return the logical module from a kernel name (fallback).""" + lowered = name + for needle, module in KERNEL_MODULE_RULES: + if needle in lowered: + if module == "moe.dispatch" and "combine" in lowered: + return "moe.combine" + return module + return "other" + + +def flop_class_from_kernel(name: str) -> str | None: + """Return 'gemm' | 'grouped_gemm' | 'attn', or None for memory-bound.""" + for needle, klass in FLOP_CLASS_RULES: + if needle in name: + return klass + return None + + +def is_compute_bound(name: str) -> bool: + return flop_class_from_kernel(name) is not None diff --git a/examples/deepseek-v4/projection/tools/parse_trace.py b/examples/deepseek-v4/projection/tools/parse_trace.py new file mode 100755 index 000000000..dad6291f7 --- /dev/null +++ b/examples/deepseek-v4/projection/tools/parse_trace.py @@ -0,0 +1,651 @@ +#!/usr/bin/env python3 +"""Turn per-cr chrome traces into a projection breakdown JSON. + +Input: one rank-0 PyTorch/Kineto chrome trace per compression-ratio (cr), each +captured by ``script/deepseek_v4_layer_trace-projection.sh`` (1 layer, seq 4096, +GA=2, recompute off, overlap off, profiler window iter 6->7). + +Output: a single ``.json`` matching ``design/03-json-schema.md``. + +Attribution (validated against the real ROCm/Kineto trace): + * GPU kernels (cat=="kernel") link to their launching CPU op via the shared + ``External id`` arg. Optimizer (``multi_tensor_apply``) and DP-comm + (``nccl``) kernels carry no External id and are classified by name. + * The module comes from the enclosing ``nn.Module: _n`` python_function + events (with_stack) on the CPU op's thread; fwd/bwd from "Backward"/ + "autograd" in the CPU op name or an enclosing frame. + * Clean per-call time = ``min`` over launches grouped by + ``(phase, module, kernel, input-dims)`` (overlap is off, so this just + removes warm-up/jitter); calls-per-microbatch is treated as 1 for a single + captured layer. + * Compute-bound FLOP class from the kernel name; GEMM FLOPs from input dims. +""" + +from __future__ import annotations + +import argparse +import bisect +import json +from collections import defaultdict +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +# Merge gap (us) used when reconstructing the backward GPU-time window from the +# linked backward kernels, so an unlinked kernel sitting in a small stall +# between two backward kernels is still counted as backward. +_BWD_WINDOW_GAP_US = 150.0 + +# A single layer-compute GPU kernel launch at seq=4096/B=1 realistically tops +# out around ~10 ms (e.g. the V4 attention bwd or a grouped GEMM). Some captures +# bill a one-off device-side stall / busy-wait to an ordinary elementwise kernel +# (observed: 3 launches of ~145 ms attributed to aten::mul/add_/div in the cr=0 +# pro trace, 17x the cr=4/128 value). Those are capture artifacts, not per-layer +# cost, so layer-compute launches above this cap are dropped (and reported in +# provenance). Optimizer/comm kernels are routed out before this cap applies. +_MAX_PLAUSIBLE_LAUNCH_US = 50_000.0 + +from kernel_module_map import flop_class_from_kernel, module_from_kernel +from v4_flops import FB_FMA, layer_fmac, model_total_params, mtp_flops, nonlayer_fmac + +PRO_COMPRESS = [128, 128] + [4 if i % 2 == 0 else 128 for i in range(2, 60)] + [0] +FLASH_COMPRESS = [0, 0] + [4 if i % 2 == 0 else 128 for i in range(2, 42)] + [0] + +MODEL_CONFIGS: dict[str, dict[str, Any]] = { + "pro": { + "num_layers": 61, + "hidden_size": 7168, + "num_attention_heads": 128, + "kv_channels": 512, + "num_experts": 384, + "moe_router_topk": 6, + "moe_ffn_hidden_size": 3072, + "moe_shared_expert_intermediate_size": 3072, + "index_topk": 1024, + "vocab_size": 129280, + "mtp_num_layers": 1, + "mtp_compress_ratios": [4], + "pipeline_layout": "", + "compress_ratios": PRO_COMPRESS, + }, + "flash": { + "num_layers": 43, + "hidden_size": 4096, + "num_attention_heads": 64, + "kv_channels": 512, + "num_experts": 256, + "moe_router_topk": 6, + "moe_ffn_hidden_size": 2048, + "moe_shared_expert_intermediate_size": 2048, + "index_topk": 512, + "vocab_size": 129280, + "mtp_num_layers": 1, + "mtp_compress_ratios": [4], + "pipeline_layout": "Et*10|t*11|t*11|t*11mL", + "compress_ratios": FLASH_COMPRESS, + }, +} + +# MI355X: BF16 matrix 2.5 PFLOPS, HBM3E 8 TB/s (AMD product page). MI455X +# (MI400): HBM4 19.6 TB/s; BF16 dense not officially published — estimated +# ~10 PFLOPS (half of the 20 PFLOPS FP8 spec). The site can override these. +DEFAULT_HARDWARE = { + "MI355X": {"peak_tflops_bf16": 2500.0, "hbm_bandwidth_gbps": 8000.0}, + "MI455X": {"peak_tflops_bf16": 10000.0, "hbm_bandwidth_gbps": 19600.0}, +} + +# Measured multi-node anchors used to set/validate the site's calibFactor. The +# projection, configured to the anchor's parallel layout, must reproduce these +# (see design/06-calibration.md). flash: real 8-node MI355X run (image pr-768, +# commit 2c9b). pro: production multi-node anchor still TODO (needs >=8 idle +# nodes); until then pro reuses flash's cross-model calibFactor. +MODEL_ANCHORS: dict[str, dict[str, Any]] = { + "flash": { + "measured_iter_time_ms": 4076.0, + "measured_anchor": { + "config": "PP8/VPP1/EP8/TP1, world 64, MBS1/GBS128, seq4096, full recompute, " + "layout Et*4|t*5|(t*6|)*5,t*4mL (8-node MI355X)", + "iter_ms": 4076.0, + "tflops_gpu": 722, + "tok_s_gpu": 2010, + "calib_factor": 0.87, + "tolerance": "<0.1%", + }, + }, +} + +ALL_MODULES = ( + "attn.proj", + "attn.core", + "attn.indexer", + "attn.norm", + "moe.router", + "moe.dispatch", + "moe.grouped_gemm", + "moe.shared_expert", + "moe.combine", + "embedding", + "output", + "loss", + "other", +) + + +def _arg(ev: dict, *keys: str) -> Any: + args = ev.get("args") or {} + for k in keys: + if k in args: + return args[k] + return None + + +def _dims_key(dims: Any) -> str: + if dims is None: + return "" + try: + return json.dumps(dims, separators=(",", ":")) + except TypeError: + return str(dims) + + +def _merge_intervals(intervals: list[tuple[float, float]], gap: float = 0.0) -> list[tuple[float, float]]: + """Merge [start, end] intervals; bridge neighbours separated by <= gap.""" + out: list[tuple[float, float]] = [] + for s, e in sorted(intervals): + if out and s <= out[-1][1] + gap: + out[-1] = (out[-1][0], max(out[-1][1], e)) + else: + out.append((s, e)) + return out + + +def _make_membership(merged: list[tuple[float, float]]): + """Return a fn(ts)->bool: is ts inside one of the merged intervals.""" + starts = [s for s, _ in merged] + + def _in(ts: float | None) -> bool: + if ts is None or not merged: + return False + i = bisect.bisect_right(starts, ts) - 1 + return i >= 0 and merged[i][0] <= ts <= merged[i][1] + + return _in + + +def _is_opt_or_comm_kernel(name: str) -> bool: + low = name.lower() + return "multi_tensor_apply" in name or "fusedadam" in low or "adamw" in low or "nccl" in low + + +def _gemm_flops(dims_key: str) -> float | None: + if not dims_key: + return None + try: + dims = json.loads(dims_key) + except json.JSONDecodeError: + return None + shapes = [d for d in dims if isinstance(d, list) and len(d) >= 2 and all(isinstance(x, int) for x in d)] + if len(shapes) >= 2: + a, b = shapes[0], shapes[1] + m, k = a[-2], a[-1] + k2, n = b[-2], b[-1] + if k == k2: + batch = 1 + for x in a[:-2]: + batch *= x + return 2.0 * batch * m * n * k + return None + + +def _resolve_module(kname: str, cpu_name: str, anc_classes: set[str], flop_class: str | None) -> str: + """Logical module for a kernel given its name, CPU op name, and enclosing + nn.Module classes. Returns '__optimizer__' / '__dpcomm__' for non-layer + buckets handled separately. Priority: kernel name > cpu-op name > enclosing + nn.Module (forward only; backward ops aren't inside module forward ranges).""" + n = kname + low = n.lower() + cn = cpu_name or "" + + def has(*xs: str) -> bool: + return any(any(x in a for a in anc_classes) for x in xs) + + def coarse(module: str) -> str: + if module in ("attn.qkv_proj", "attn.o_proj"): + return "attn.proj" + if module in ("attn.rope",): + return "attn.norm" + if module == "moe.act": + return "moe.grouped_gemm" + return module + + # 1) kernel-name rules (most reliable; present for fwd and bwd) + if "multi_tensor_apply" in n or "fusedadam" in low or "adamw" in low: + return "__optimizer__" + if "nccl" in low: + return "__dpcomm__" + if "cross_entropy" in low or "online_softmax" in low: + return "loss" + if "deep_ep" in n or "deepep" in low: + return "moe.combine" if "combine" in low else "moe.dispatch" + if "_v4_csa" in n or "_v4_attention" in n or "_hc_" in n: + return "attn.core" + if "_sinkhorn" in n or "_v4_router" in n: + return "moe.router" + if ( + "GroupedGemm" in n + or "_grouped" in n + or "group_gemm" in low + or "grouped_gemm" in low + or "grouped_variable" in n + ): + return "moe.grouped_gemm" + + # 2) cpu-op (autograd Function / aten) name rules — needed for backward + if "Attention" in cn or "CSAPool" in cn or "MLA" in cn: + return "attn.core" + if "Indexer" in cn or "Compressor" in cn: + return "attn.indexer" + if "Sinkhorn" in cn or "Router" in cn: + return "moe.router" + if "RMSNorm" in cn or "LayerNorm" in cn or "layer_norm" in cn.lower(): + return "attn.norm" + if "crossentropy" in cn.lower() or "cross_entropy" in cn.lower() or "nll_loss" in cn.lower(): + return "loss" + if "embedding" in cn.lower(): + return "output" if flop_class == "gemm" else "embedding" + if "LinearWithGradAccumulation" in cn or cn in ("aten::mm", "aten::addmm", "aten::matmul", "aten::bmm"): + if has("Embedding") or ( + has("DeepseekV4Model") + and not has( + "DeepseekV4Attention", "DeepseekV4HybridLayer", "Compressor", "Indexer", "MLP", "Expert" + ) + ): + return "output" + if has("Compressor", "Indexer"): + return "attn.indexer" + if has("MLP", "Expert"): + return "moe.grouped_gemm" + return "attn.proj" + + # 3) enclosing nn.Module (forward only) + if has("Compressor", "Indexer"): + return "attn.indexer" + if has("SharedExpert"): + return "moe.shared_expert" + if has("GroupedMLP", "SequentialMLP", "GroupedExperts", "Experts"): + return "moe.grouped_gemm" + if has("Router"): + return "moe.router" + if has("DeepseekV4Attention", "MLASelfAttention", "SelfAttention"): + return "attn.proj" if flop_class == "gemm" else "attn.norm" + if has("Embedding"): + return "output" if flop_class == "gemm" else "embedding" + fallback = coarse(module_from_kernel(kname)) + return fallback if fallback != "other" else "attn.misc" + + +def parse_trace(path: Path, ga: int = 2): + payload = json.loads(path.read_text()) + events = payload.get("traceEvents", []) + + # index cpu ops by External id + cpu_by_extid: dict[Any, dict] = {} + # interesting python_function intervals per tid: (ts, end, name) + pf_by_tid: dict[Any, list[tuple[float, float, str]]] = defaultdict(list) + # `autograd::engine::evaluate_function` CPU intervals precisely bracket the + # backward pass; a kernel whose launching CPU op falls inside one is bwd. + eval_intervals: list[tuple[float, float]] = [] + kernels: list[dict] = [] + num_steps = 0 # real captured training steps (ProfilerStep events, dur>10ms) + + for ev in events: + cat = (ev.get("cat") or "").lower() + ph = ev.get("ph") + if ph == "X" and ev.get("name", "").startswith("ProfilerStep") and (ev.get("dur") or 0) > 10000: + num_steps += 1 + if cat == "cpu_op" and ph == "X": + ext = _arg(ev, "External id") + if ext is not None and ext not in cpu_by_extid: + cpu_by_extid[ext] = ev + if ev.get("name", "").startswith("autograd::engine::evaluate_function"): + ets = ev.get("ts") + if ets is not None: + eval_intervals.append((ets, ets + (ev.get("dur") or 0))) + elif cat == "python_function" and ph == "X": + name = ev.get("name", "") + if name.startswith("nn.Module:") or "ackward" in name or "autograd" in name: + ts = ev.get("ts") + dur = ev.get("dur") or 0 + if ts is not None: + pf_by_tid[ev.get("tid")].append((ts, ts + dur, name)) + elif cat == "kernel" and ph == "X" and ev.get("dur") is not None: + kernels.append(ev) + + for tid in pf_by_tid: + pf_by_tid[tid].sort(key=lambda t: t[0]) + + def enclosing(cpu: dict) -> set[str]: + """Return enclosing nn.Module class names for a cpu op (forward only; + backward ops live under the autograd engine, not module forward ranges).""" + tid = cpu.get("tid") + ts = cpu.get("ts") + end = ts + (cpu.get("dur") or 0) + classes: set[str] = set() + for pts, pend, name in pf_by_tid.get(tid, ()): + if pts > ts: + break + if pend >= end and name.startswith("nn.Module:"): + classes.add(name.split(":", 1)[1].strip().rsplit("_", 1)[0]) + return classes + + # Per-microbatch time = (sum of all launches over the whole profiler window) + # / num_mb, where num_mb = num_steps * GA. We keep the SUM (not a min over + # launches): the single-layer capture serializes each layer's work, and the + # measured flash-16L anchor (design/06) shows this serial per-layer time is + # the right estimate (calibFactor ~0.93 against 6665 ms). A min-over-launches + # rule would assume the scalar control-flow stalls (e.g. the Indexer top-k + # device syncs) are fully hidden in the full model; the anchor refutes that + # (it would need calibFactor ~1.3), so they are kept as real per-layer cost. + # + # Subgroups are keyed by (kernel, input-dims) only so we can (a) classify a + # row as compute_bound iff FLOP-classed kernels dominate its time (>50%) -- + # a stray flop-classed kernel can no longer flip a memory-bound aggregate -- + # and (b) sum dim-derived GEMM FLOPs correctly. + num_mb = max(1, num_steps * ga) + + # ---- phase classification (forward vs backward) ---------------------- + # The backward pass is identified by the autograd engine: a kernel is + # backward iff its launching CPU op was issued inside an + # `autograd::engine::evaluate_function` interval (these precisely bracket + # the backward). A `_fwd_`/`_bwd_` tag in the kernel name (V4 Triton kernels + # encode it) wins. Unlinked kernels (no External id -> no CPU op; common for + # fused/elementwise launches Kineto fails to flow-link) are assigned by + # whether their GPU timestamp lands inside the backward GPU-time window + # rebuilt from the linked backward kernels. This replaces the old "default + # to forward" rule, which systematically leaked backward compute -- incl. + # the MoE dgrad/wgrad grouped GEMMs -- into the forward breakdown. + bwd_eval = _merge_intervals(eval_intervals) + in_bwd_eval = _make_membership(bwd_eval) + + phase_by_idx: list[str | None] = [None] * len(kernels) + _bwd_windows: list[tuple[float, float]] = [] + for i, ev in enumerate(kernels): + name = ev.get("name", "") + ln = name.lower() + if "_fwd" in ln and "_bwd" not in ln: + phase_by_idx[i] = "forward" + continue + if "_bwd" in ln: + phase_by_idx[i] = "backward" + if not _is_opt_or_comm_kernel(name) and ev["dur"] <= _MAX_PLAUSIBLE_LAUNCH_US: + _bwd_windows.append((ev["ts"], ev["ts"] + ev["dur"])) + continue + ext = _arg(ev, "External id") + cpu = cpu_by_extid.get(ext) if ext is not None else None + if cpu is not None: + ph = "backward" if in_bwd_eval(cpu.get("ts")) else "forward" + phase_by_idx[i] = ph + if ( + ph == "backward" + and not _is_opt_or_comm_kernel(name) + and ev["dur"] <= _MAX_PLAUSIBLE_LAUNCH_US + ): + _bwd_windows.append((ev["ts"], ev["ts"] + ev["dur"])) + # else: unlinked -> decided below from the GPU-side backward window + + in_bwd_gpu = _make_membership(_merge_intervals(_bwd_windows, _BWD_WINDOW_GAP_US)) + for i, ev in enumerate(kernels): + if phase_by_idx[i] is None: + phase_by_idx[i] = "backward" if in_bwd_gpu(ev.get("ts")) else "forward" + + # (phase, module) -> { (kernel_name, dims_key): {"durs": [...], "flop_class", "flop_per_launch"} } + groups: dict[tuple[str, str], dict] = defaultdict(dict) + optimizer_us = 0.0 + dpcomm_us = 0.0 + dropped_stall_us = 0.0 + + for i, ev in enumerate(kernels): + name = ev.get("name", "") + dur = float(ev["dur"]) + ext = _arg(ev, "External id") + cpu = cpu_by_extid.get(ext) if ext is not None else None + cpu_name = cpu.get("name", "") if cpu else "" + anc = enclosing(cpu) if cpu else set() + flop_class = flop_class_from_kernel(name) + module = _resolve_module(name, cpu_name, anc, flop_class) + if module == "__optimizer__": + optimizer_us += dur + continue + if module == "__dpcomm__": + dpcomm_us += dur + continue + # Drop one-off device-side stalls billed to a layer-compute kernel (see + # _MAX_PLAUSIBLE_LAUNCH_US); they are capture artifacts, not per-layer cost. + if dur > _MAX_PLAUSIBLE_LAUNCH_US: + dropped_stall_us += dur + continue + phase = phase_by_idx[i] + dims_key = _dims_key(_arg(cpu, "Input Dims")) if cpu else "" + flop_per_launch = _gemm_flops(dims_key) if (cpu and flop_class == "gemm") else None + sub = groups[(phase, module)].setdefault( + (name, dims_key), + {"durs": [], "flop_class": flop_class, "flop_per_launch": flop_per_launch}, + ) + sub["durs"].append(dur) + + optimizer_us /= max(1, num_steps) # optimizer runs once per training iteration + dpcomm_us /= num_mb + dropped_stall_us /= num_mb # report per-microbatch, like the breakdown rows + + out = {b: {"forward": {}, "backward": {}} for b in ("attention", "moe", "embedding", "output", "loss")} + for (phase, module), subs in groups.items(): + time_us = 0.0 + compute_us = 0.0 + flops = 0.0 + has_flops = False + kern: dict[str, float] = defaultdict(float) + class_time: dict[str, float] = defaultdict(float) + for (name, _dims), sub in subs.items(): + n = len(sub["durs"]) + per_mb = sum(sub["durs"]) / num_mb + time_us += per_mb + kern[name[:80]] += per_mb + if sub["flop_class"]: + compute_us += per_mb + class_time[sub["flop_class"]] += per_mb + if sub["flop_per_launch"]: + flops += sub["flop_per_launch"] * n / num_mb + has_flops = True + compute = compute_us > 0.5 * time_us if time_us > 0 else False + flop_class = max(class_time, key=class_time.get) if (compute and class_time) else None + flops_out = flops if (compute and has_flops) else None + time_s = time_us * 1e-6 + tflops = (flops_out / time_s / 1e12) if (flops_out and time_s > 0) else None + kern_list = sorted( + ({"name": n, "time_us": round(t, 3)} for n, t in kern.items()), + key=lambda k: -k["time_us"], + )[:6] + entry = { + "module": module, + "time_us": round(time_us, 3), + "class": "compute_bound" if compute else "memory_bound", + "flop_class": flop_class, + "flops": flops_out, + "tflops": round(tflops, 1) if tflops else None, + "kernels": kern_list, + } + # Logical bucket. Unattributed scalar kernels are labelled `attn.misc` + # by _resolve_module because the V4 traces show they are dominated by + # attention-side hyper-connection / Indexer control-flow work; this keeps + # the MoE bucket strictly the cr-independent router/dispatch/grouped_gemm + # /combine/shared_expert set (A10). + if module.startswith("attn."): + bucket = "attention" + elif module.startswith("moe."): + bucket = "moe" + elif module in ("embedding", "output", "loss"): + bucket = module + else: + bucket = "moe" + out[bucket][phase][module] = entry + return out, optimizer_us, dpcomm_us, dropped_stall_us + + +def _lists(bd: dict) -> dict: + return {p: sorted(bd[p].values(), key=lambda r: -r["time_us"]) for p in ("forward", "backward")} + + +def cr_layer_counts(compress: list[int]) -> dict[str, int]: + counts = {"0": 0, "4": 0, "128": 0} + for c in compress: + counts[str(c)] = counts.get(str(c), 0) + 1 + return counts + + +def _fwd_total(buckets: dict) -> float: + return sum(r["time_us"] for r in buckets["attention"]["forward"].values()) + sum( + r["time_us"] for r in buckets["moe"]["forward"].values() + ) + + +def build(model: str, traces: dict[str, Path], ga: int = 2) -> dict[str, Any]: + cfg = MODEL_CONFIGS[model] + per_cr, opt_us = {}, [] + dropped_stall = {} + for cr, path in traces.items(): + buckets, o, _dp, stall = parse_trace(path, ga) + per_cr[cr] = buckets + opt_us.append(o) + dropped_stall[cr] = round(stall, 1) + + # Some cr layers (pure dense cr=0 / HCA cr=128) get CUDA-graph / stream- + # captured, so their compute kernels are not individually visible in the + # trace (only optimizer/comm/elementwise appear). Fall back to the eager + # cr=4 sample for those so the full-model projection isn't zeroed; flag it. + ref = ( + "4" + if "4" in per_cr and _fwd_total(per_cr["4"]) >= 1000 + else max(per_cr, key=lambda c: _fwd_total(per_cr[c])) + ) + graphed = [] + for cr in list(per_cr): + if cr != ref and _fwd_total(per_cr[cr]) < 1000: + graphed.append(cr) + per_cr[cr] = per_cr[ref] + + src = per_cr[ref] + layers = {cr: {"attention": _lists(b["attention"]), "moe": _lists(b["moe"])} for cr, b in per_cr.items()} + optimizer_us = round(sum(opt_us) / len(opt_us), 1) if opt_us else None + + # Analytic V4 closed-form FLOPs (ported from Megatron's flops patch) so the + # site's TFLOP/s matches a real run. Per layer, per microbatch (B=1), at the + # capture seq; Megatron-convention (fwd+bwd, FMA -> x6). + cap_seq = 4096 + mtp_num_layers = int(cfg.get("mtp_num_layers", 0) or 0) + mtp_cr = int((cfg.get("mtp_compress_ratios") or [0])[0]) + mtp = mtp_flops(model, cap_seq, mtp_num_layers, mtp_cr) + analytic_flops = { + "per_cr_layer_flops": { + cr: sum(layer_fmac(model, int(cr), cap_seq).values()) * FB_FMA for cr in ("0", "4", "128") + }, + "output_flops": nonlayer_fmac(model, cap_seq)["logits"] * FB_FMA, + "mtp": { + "num_layers": mtp_num_layers, + "compress_ratio": mtp_cr, + "inner_layer_flops": mtp["inner_layer"], + "eh_proj_flops": mtp["eh_proj"], + "extra_logits_flops": mtp["extra_logits"], + "hc_head_flops": mtp["hc_head"], + }, + "seq": cap_seq, + "note": "per-layer and MTP (B=1) Megatron-convention FLOPs at capture seq; site multiplies by GA x DP", + } + + return { + "schema_version": 1, + "model": model, + "analytic_flops": analytic_flops, + "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), + "provenance": { + "traces": {cr: str(p) for cr, p in traces.items()}, + "graphed_crs_estimated_from_cr4": graphed, + "dropped_stall_us_per_mb": dropped_stall, + "note": ( + "cr in graphed_crs_estimated_from_cr4 were CUDA-graph/stream-captured " + "(compute not visible in trace); their breakdown is copied from cr=4 as an " + "estimate. Re-run those cr with graph capture disabled for exact numbers. " + "dropped_stall_us_per_mb: per-mb layer-compute kernel time dropped as " + "implausible one-off device stalls (> _MAX_PLAUSIBLE_LAUNCH_US)." + ), + }, + "capture": { + "gpu": "MI355X", + "seq_length": 4096, + "micro_batch_size": 1, + "tokens_per_microbatch": 4096, + "ep": 8, + "ga_for_capture": 2, + "optimizer": "adam", + "distributed_optimizer": False, + "recompute": "off", + "measured_iter_time_ms": MODEL_ANCHORS.get(model, {}).get("measured_iter_time_ms"), + "measured_anchor": MODEL_ANCHORS.get(model, {}).get("measured_anchor"), + }, + "model_config": { + **cfg, + "cr_layer_counts": cr_layer_counts(cfg["compress_ratios"]), + "total_params": model_total_params(model, cfg["num_layers"], mtp_num_layers), + }, + "hardware": DEFAULT_HARDWARE, + "layers": layers, + "non_layer": {k: _lists(src[k]) for k in ("embedding", "output", "loss")}, + "optimizer": { + "type": "adam", + "measured_params": None, + "time_us": optimizer_us, + "bytes_per_param": 30, + "class": "memory_bound", + "note": "Adam mixed-precision step traffic in bytes/param; measured one-layer optimizer-step kernel time is a sanity reference", + }, + "comm": { + "ep_dispatch_us": None, + "ep_combine_us": None, + "note": "EP dispatch/combine are memory_bound rows inside moe; informational", + }, + } + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Build projection breakdown JSON from per-cr traces.") + p.add_argument("--model", required=True, choices=sorted(MODEL_CONFIGS)) + p.add_argument("--trace", action="append", default=[], metavar="cr=PATH") + p.add_argument( + "--ga", type=int, default=2, help="gradient-accumulation at capture (GBS/(DP*MBS)); default 2" + ) + p.add_argument("--out", required=True, type=Path) + return p.parse_args() + + +def main() -> int: + args = parse_args() + traces: dict[str, Path] = {} + for spec in args.trace: + if "=" not in spec: + raise SystemExit(f"--trace must be cr=PATH, got: {spec}") + cr, path = spec.split("=", 1) + traces[cr.replace("cr", "")] = Path(path) + if not traces: + raise SystemExit("at least one --trace cr=PATH is required") + for cr, path in traces.items(): + if not path.exists(): + raise SystemExit(f"trace not found for cr={cr}: {path}") + + doc = build(args.model, traces, args.ga) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(doc, indent=2) + "\n") + print(f"[parse_trace] wrote {args.out} (model={args.model}, crs={sorted(traces)})") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/deepseek-v4/projection/tools/v4_flops.py b/examples/deepseek-v4/projection/tools/v4_flops.py new file mode 100644 index 000000000..34329866e --- /dev/null +++ b/examples/deepseek-v4/projection/tools/v4_flops.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python3 +"""DeepSeek-V4 closed-form analytic FLOPs, ported from +``primus/backends/megatron/patches/deepseek_v4_flops_patches.py`` so the +projection's TFLOP/s matches what Megatron reports on a real run. + +Returns per-component FMAC (multiply-only, pre fwd+bwd/FMA expansion) for ONE +layer of a given cr at batch_size=1. Multiply by FB_FMA (=6) for Megatron- +convention FLOPs (fwd 1 + bwd 2, times FMA 2). + +Validated against the measured flash 16-layer run (TOTAL 34093 TFLOP/global- +batch; per-component within rounding) — see __main__ self-test. +""" + +from __future__ import annotations + +FB_FMA = 6 # _FORWARD_BACKWARD_FACTOR(3) * _FMA_FACTOR(2) +SWIGLU = 3 # gate+up+down collapsed expansion factor +DEFAULT_MTP_LAYERS = {"pro": 1, "flash": 1} + +# Per-model architecture params (from primus/configs/models/megatron/*.yaml + +# deepseek_v4_base.yaml). Shared: index_head_dim=128, index_n_heads=64, +# attn_sliding_window=128, hc_mult=4. +MODEL_PARAMS = { + "pro": dict( + hidden=7168, + heads=128, + head_dim=512, + q_lora=1536, + o_lora=1024, + o_groups=16, + moe_ffn=3072, + shared_ffn=3072, + topk=6, + experts=384, + index_topk=1024, + vocab=129280, + ), + "flash": dict( + hidden=4096, + heads=64, + head_dim=512, + q_lora=1024, + o_lora=1024, + o_groups=8, + moe_ffn=2048, + shared_ffn=2048, + topk=6, + experts=256, + index_topk=512, + vocab=129280, + ), +} +SHARED = dict(index_head_dim=128, index_n_heads=64, swa_window=128, hc_mult=4) + + +def _local_visible_pairs(swa, s): + if swa <= 0 or swa >= s: + return s * (s + 1) // 2 + return swa * s - swa * (swa - 1) // 2 + + +def _pool_visible_pairs(cr, s): + if cr <= 0 or s <= 0: + return 0 + n = s // cr + if n == 0: + return 0 + return cr * n * (n - 1) // 2 + n * (s - cr * n + 1) + + +def _visible_pairs(swa, cr, index_topk, s): + local = _local_visible_pairs(swa, s) + if cr == 0: + return local + pool = max(1, s // cr) + if cr == 128: + return local + _pool_visible_pairs(cr, s) + if cr == 4: + sparse = min(index_topk if index_topk else pool, pool) + return local + sparse * s + return local + pool * s + + +def _attn_qkv_o(s_eff, p): + n_d = p["heads"] * p["head_dim"] + qkv = p["hidden"] * p["q_lora"] + p["q_lora"] * n_d + p["hidden"] * p["head_dim"] + if p["o_lora"] > 0: + o_proj = n_d * p["o_lora"] + (p["o_groups"] * p["o_lora"]) * p["hidden"] + else: + o_proj = n_d * p["hidden"] + return s_eff * (qkv + o_proj) + + +def _attn_scores(s_eff, cr, p): + pairs = _visible_pairs(SHARED["swa_window"], cr, p["index_topk"], s_eff) + return 2 * p["heads"] * p["head_dim"] * pairs + + +def _compressor(s_eff, cr, p): + if cr == 0: + return 0 + coff = 2 if cr == 4 else 1 + return 2 * s_eff * p["hidden"] * (coff * p["head_dim"]) + + +def _indexer(s_eff, cr, p): + if cr != 4: + return 0 + ihd, inh = SHARED["index_head_dim"], SHARED["index_n_heads"] + pool = max(1, s_eff // cr) + dq_rank = ihd + proj = p["hidden"] * dq_rank + dq_rank * (inh * ihd) + p["hidden"] * inh + proj += 2 * p["hidden"] * (2 * ihd) # mini-compressor + return s_eff * proj + s_eff * inh * pool * ihd + + +def _moe(s_eff, p): + router = p["hidden"] * p["experts"] + routed = p["topk"] * SWIGLU * p["hidden"] * p["moe_ffn"] + shared = SWIGLU * p["hidden"] * p["shared_ffn"] if p["shared_ffn"] > 0 else 0 + return s_eff * (router + routed + shared) + + +def _hc_mixer(s, p): + hc = SHARED["hc_mult"] + n_d = hc * p["hidden"] + return 2 * s * n_d * ((2 + hc) * hc) + + +def _hc_head(s, p, mtp_num_layers): + hc = SHARED["hc_mult"] + n_d = hc * p["hidden"] + return (1 + mtp_num_layers) * s * n_d * hc + + +def _mtp_eh_proj(s, p, mtp_num_layers): + return mtp_num_layers * s * (2 * p["hidden"]) * p["hidden"] + + +def layer_fmac(model: str, cr: int, seq: int) -> dict[str, float]: + """Per-layer FMAC components (batch_size=1) for one cr layer.""" + p = MODEL_PARAMS[model] + s_eff = seq * SHARED["hc_mult"] + return { + "attn_qkv_o": _attn_qkv_o(s_eff, p), + "attn_scores": _attn_scores(s_eff, cr, p), + "compressor": _compressor(s_eff, cr, p), + "indexer": _indexer(s_eff, cr, p), + "moe": _moe(s_eff, p), + "hc": _hc_mixer(seq, p), + } + + +def nonlayer_fmac(model: str, seq: int, mtp_num_layers: int = 0) -> dict[str, float]: + p = MODEL_PARAMS[model] + return {"logits": (1 + mtp_num_layers) * seq * p["hidden"] * p["vocab"]} + + +def mtp_fmac(model: str, seq: int, mtp_num_layers: int = 1, mtp_cr: int = 4) -> dict[str, float]: + """Extra FMAC components for MTP depths, batch_size=1. + + V4 MTP reuses a full V4 inner layer per depth; the current Flash Megatron + FLOPs anchor reports a CSA-style MTP inner layer (cr=4). + """ + if mtp_num_layers <= 0: + return { + "inner_layer": 0, + "eh_proj": 0, + "extra_logits": 0, + "hc_head": _hc_head(seq, MODEL_PARAMS[model], 0), + } + p = MODEL_PARAMS[model] + inner = sum(layer_fmac(model, mtp_cr, seq).values()) * mtp_num_layers + main_logits = seq * p["hidden"] * p["vocab"] + return { + "inner_layer": inner, + "eh_proj": _mtp_eh_proj(seq, p, mtp_num_layers), + "extra_logits": main_logits * mtp_num_layers, + "hc_head": _hc_head(seq, p, mtp_num_layers), + } + + +def model_total_params(model: str, num_layers: int, mtp_num_layers: int = 0) -> int: + """Approximate total parameter count (for optimizer-step sizing). Uses the + same V4 MLA low-rank attention shapes as the FLOPs formula (q/o LoRA + single + latent KV) instead of the crude 4*h^2, plus MoE experts + shared + router and + the tied-free embedding/output. MTP adds one full V4 inner layer plus the + 2H->H eh_proj per depth; logits reuse the output layer weights.""" + p = MODEL_PARAMS[model] + n_d = p["heads"] * p["head_dim"] + attn = ( + p["hidden"] * p["q_lora"] + + p["q_lora"] * n_d + + p["hidden"] * p["head_dim"] + + n_d * p["o_lora"] + + p["o_groups"] * p["o_lora"] * p["hidden"] + ) + moe = ( + p["experts"] * SWIGLU * p["hidden"] * p["moe_ffn"] + + SWIGLU * p["hidden"] * p["shared_ffn"] + + p["hidden"] * p["experts"] + ) + return int( + (num_layers + mtp_num_layers) * (attn + moe) + + mtp_num_layers * 2 * p["hidden"] * p["hidden"] + + 2 * p["vocab"] * p["hidden"] + ) + + +# Map analytic components to projection module names (per layer). +def module_flops(model: str, cr: int, seq: int) -> dict[str, float]: + """Megatron-convention FLOPs (×FB_FMA) per module, per layer, batch_size=1.""" + f = layer_fmac(model, cr, seq) + return { + "attn.proj": f["attn_qkv_o"] * FB_FMA, + "attn.core": f["attn_scores"] * FB_FMA, + "attn.indexer": (f["compressor"] + f["indexer"]) * FB_FMA, + "attn.norm": f["hc"] * FB_FMA, + "moe.grouped_gemm": f["moe"] * FB_FMA, + } + + +def output_flops(model: str, seq: int) -> float: + return nonlayer_fmac(model, seq)["logits"] * FB_FMA + + +def mtp_flops(model: str, seq: int, mtp_num_layers: int = 1, mtp_cr: int = 4) -> dict[str, float]: + f = mtp_fmac(model, seq, mtp_num_layers, mtp_cr) + return {k: v * FB_FMA for k, v in f.items()} + + +def _self_test() -> None: + """Self-test against measured flash 16L (GBS64): cr [0x3,4x6,128x7].""" + sched = [0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0] + seq_t, batch = 4096, 64 + comp = {k: 0.0 for k in ("attn_qkv_o", "attn_scores", "compressor", "indexer", "moe", "hc")} + for cr_t in sched: + for k, v in layer_fmac("flash", cr_t, seq_t).items(): + comp[k] += v + logits = nonlayer_fmac("flash", seq_t)["logits"] + tot = (sum(comp.values()) + logits) * FB_FMA * batch / 1e12 + print("flash 16L analytic vs measured (TFLOP/global-batch):") + for k, v in comp.items(): + print(f" {k:12s} = {v*FB_FMA*batch/1e12:9.1f}") + print(f" {'logits':12s} = {logits*FB_FMA*batch/1e12:9.1f}") + print(f" TOTAL = {tot:9.1f} (measured 34093.4)") + + +if __name__ == "__main__": + _self_test() diff --git a/examples/deepseek-v4/run_deepseek_v4.sh b/examples/deepseek-v4/run_deepseek_v4.sh new file mode 100755 index 000000000..eefd55e80 --- /dev/null +++ b/examples/deepseek-v4/run_deepseek_v4.sh @@ -0,0 +1,314 @@ +#!/bin/bash +set -euo pipefail +set -x + +_RUN_START_SEC=$(date +%s) +_RUN_START_TS=$(date '+%Y-%m-%d %H:%M:%S') +_print_run_elapsed() { + local _end_sec _end_ts _elapsed _exit=$1 + _end_sec=$(date +%s) + _end_ts=$(date '+%Y-%m-%d %H:%M:%S') + _elapsed=$((_end_sec - _RUN_START_SEC)) + echo "----------------------------------------" + echo "run_deepseek_v4.sh wall time" + echo " start: ${_RUN_START_TS}" + echo " end: ${_end_ts}" + echo " elapsed: ${_elapsed}s ($((_elapsed / 60))m $((_elapsed % 60))s)" + echo " exit: ${_exit}" +} +trap '_print_run_elapsed $?' EXIT + +export HF_TOKEN="${HF_TOKEN:-}" +export WANDB_API_KEY="${WANDB_API_KEY:-your_wandb_api_key}" + +export NNODES=${NNODES:-1} +export TRAIN_ITERS=${TRAIN_ITERS:-20} + +export DOCKER_IMAGE=${DOCKER_IMAGE:-"tasimage/primus:pr-715-ainic"} +export SLURM_PARTITION=Compute-DCPT +export SLURM_NODELIST="${SLURM_NODELIST:-smci355-ccs-aus-n01-21,smci355-ccs-aus-n01-33,smci355-ccs-aus-n02-21,smci355-ccs-aus-n02-25,smci355-ccs-aus-n02-29,smci355-ccs-aus-n02-33,smci355-ccs-aus-n03-33,smci355-ccs-aus-n04-21,smci355-ccs-aus-n04-25,smci355-ccs-aus-n04-29,smci355-ccs-aus-n04-33,smci355-ccs-aus-n05-21,smci355-ccs-aus-n05-29,smci355-ccs-aus-n05-33,smci355-ccs-aus-n06-25,smci355-ccs-aus-n06-33,smci355-ccs-aus-n10-29}" +export MASTER_PORT=${MASTER_PORT:-29500} + +export USING_AINIC=${USING_AINIC:-1} +export NCCL_IB_HCA="ionic_0:1,ionic_1:1,ionic_2:1,ionic_3:1,ionic_4:1,ionic_5:1,ionic_6:1,ionic_7:1" +# "fenic" is the cluster RDMA NIC. Prefer it when present (cluster / multi-node), +# but fall back to a real local interface (lo) so single-node / direct in-container +# smoke runs work out of the box — otherwise gloo aborts with +# "Unable to find address for: fenic". Still override-guarded: pass +# GLOO_SOCKET_IFNAME / NCCL_SOCKET_IFNAME explicitly to force a specific NIC. +if [ -d /sys/class/net/fenic ]; then + _PRIMUS_DEFAULT_IFNAME=fenic +else + _PRIMUS_DEFAULT_IFNAME=lo +fi +export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-$_PRIMUS_DEFAULT_IFNAME} +export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-$_PRIMUS_DEFAULT_IFNAME} +export NCCL_IB_GID_INDEX=1 +export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-1} +export NVTE_CK_USES_BWD_V3=${NVTE_CK_USES_BWD_V3:-1} + +# Phase-7 fixed knobs for single-node bring-up. +export MBS=${MBS:-1} +export GBS=${GBS:-$((16 * NNODES * MBS))} +export PRIMUS_TP=${PRIMUS_TP:-1} +export PRIMUS_PP=${PRIMUS_PP:-1} +export PRIMUS_EP=${PRIMUS_EP:-8} + +# Keep this smoke config lightweight for quick bring-up. +export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-8} +export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-128} +export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-128} +export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-8} +export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-2} +export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-512} +export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-8} +export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-"[0,0,4,4,4,4,4,0]"} +export PRIMUS_MOE_ENABLE_EXPERT_BIAS=${PRIMUS_MOE_ENABLE_EXPERT_BIAS:-False} +export PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU=${PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU:-True} +export PROFILE=${PROFILE:-False} +export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False} +export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-False} +export LEGACY_GG=${LEGACY_GG:-False} +# Plan-3 P22 / P23: PrimusTurbo gate (must be on for turbo attention / +# turbo deepep to take effect; enable_primus_turbo gates the +# `before_train` patches that re-bind the spec provider). +export ENABLE_PRIMUS_TURBO=${ENABLE_PRIMUS_TURBO:-False} +if [ "$USE_TURBO_ATTENTION" = "True" ] || [ "${USE_TURBO_DEEPEP:-False}" = "True" ]; then + ENABLE_PRIMUS_TURBO=True +fi +export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-False} + +if [ "$TURBO_USE_GROUPED_MLP" = "True" ]; then + export PRIMUS_BIAS_SWIGLU_FUSION=True +fi + +# Plan-3 P23: Turbo DeepEP-related knobs. Only emit these CLI flags +# when USE_TURBO_DEEPEP=True so non-deepep runs don't carry unrelated +# overrides. Best-practice CU count: 64 (or 80) for EP=8, 32 for +# EP>=16 — the EP>=16 cap is asserted by +# `primus/modules/trainer/megatron/utils.py:527`. DeepEP itself +# requires `moe_router_dtype=fp32` and forbids +# `moe_shared_expert_overlap=True` (both are already V4-Flash YAML +# defaults; we pin them via CLI defensively so a stray YAML override +# or future config edit cannot flip them out from under the Turbo +# path mid-run). +TURBO_DEEPEP_CLI_ARGS=() +if [ "$USE_TURBO_DEEPEP" = "True" ]; then + if [ "${PRIMUS_EP:-1}" -ge 16 ]; then + _DEFAULT_TURBO_DEEPEP_NUM_CU=32 + else + _DEFAULT_TURBO_DEEPEP_NUM_CU=80 + fi + export TURBO_DEEPEP_NUM_CU=${TURBO_DEEPEP_NUM_CU:-$_DEFAULT_TURBO_DEEPEP_NUM_CU} + export TURBO_DEEPEP_USE_COMM_STREAM=${TURBO_DEEPEP_USE_COMM_STREAM:-False} + export MOE_ROUTER_DTYPE=${MOE_ROUTER_DTYPE:-fp32} + export MOE_SHARED_EXPERT_OVERLAP=${MOE_SHARED_EXPERT_OVERLAP:-False} + TURBO_DEEPEP_CLI_ARGS=( + --turbo_deepep_num_cu "$TURBO_DEEPEP_NUM_CU" + --turbo_deepep_use_comm_stream "$TURBO_DEEPEP_USE_COMM_STREAM" + --moe_router_dtype "$MOE_ROUTER_DTYPE" + --moe_shared_expert_overlap "$MOE_SHARED_EXPERT_OVERLAP" + ) +fi + +export PRECISION_TYPE=${PRECISION_TYPE:-BF16} +# Honor an incoming FP8 / FP8_RECIPE env (e.g. FP8_RECIPE=mxfp8); default null +# so non-FP8 runs are unchanged. (Previously these were hard-set to null, +# which silently clobbered a caller-provided recipe.) +export FP8=${FP8:-null} +export FP8_RECIPE=${FP8_RECIPE:-null} + +# ---------- Optimizer selection (adam default; muon = DeepSeek-V4 recipe) ---- +# OPTIMIZER=adam (default): unchanged behaviour (BF16 precision-aware AdamW +# from the EXP yaml); overlap_grad_reduce / overlap_param_gather stay ON. +# OPTIMIZER=muon: Primus distributed-Muon path (primus .../optimizer/moun.py). +# Megatron asserts plain `muon` is incompatible with distributed optimizer + +# grad/param overlap, so we force them OFF and switch optimizer states to +# fp32 (Muon does not support the precision-aware optimizer). The +# Newton-Schulz coefficient set auto-selects 'deepseekv4' (8 aggressive + 2 +# stable) for V4 configs inside get_megatron_muon_optimizer. Requires the +# emerging_optimizers package -> we set PRIMUS_INSTALL_EMERGING_OPTIMIZERS so +# the in-container install hook (runner/.../01_install_emerging_optimizers.sh) +# provisions it. +export OPTIMIZER=${OPTIMIZER:-adam} +export PRIMUS_OVERLAP_GRAD_REDUCE=${PRIMUS_OVERLAP_GRAD_REDUCE:-True} +export PRIMUS_OVERLAP_PARAM_GATHER=${PRIMUS_OVERLAP_PARAM_GATHER:-True} +OPTIMIZER_CLI_ARGS=() +if [ "$OPTIMIZER" = "muon" ] || [ "$OPTIMIZER" = "dist_muon" ]; then + export PRIMUS_INSTALL_EMERGING_OPTIMIZERS=${PRIMUS_INSTALL_EMERGING_OPTIMIZERS:-1} + export MUON_MOMENTUM=${MUON_MOMENTUM:-0.95} + export MUON_EXTRA_SCALE_FACTOR=${MUON_EXTRA_SCALE_FACTOR:-0.18} + # Both plain muon (Megatron asserts) and dist_muon (LayerWiseDistributed- + # Optimizer docstring: "keep all megatron distributed-optimizer related + # options OFF"; it manages its own param all-gather, so DDP + # overlap_param_gather double-drives start_param_sync -> crash) need the + # DDP grad/param overlap OFF. + PRIMUS_OVERLAP_GRAD_REDUCE=False + PRIMUS_OVERLAP_PARAM_GATHER=False + OPTIMIZER_CLI_ARGS=( + --optimizer "$OPTIMIZER" + --muon_momentum "$MUON_MOMENTUM" + --muon_extra_scale_factor "$MUON_EXTRA_SCALE_FACTOR" + --use_distributed_optimizer False + --use_precision_aware_optimizer False + --main_grads_dtype fp32 + --exp_avg_dtype fp32 + --exp_avg_sq_dtype fp32 + ) +fi + +# DeepSeek-V4 attention backend selection (unified string selectors). Default +# triton_v2 (production default; fastest V4 sparse-MLA path). These are +# V4-only; no effect on other model types. +# USE_V4_ATTENTION_BACKEND (dense cr=0 / HCA cr=128): eager|triton_v1|triton_v2|gluon +# USE_V4_CSA_ATTENTION_BACKEND (CSA cr=4): eager|triton_v0|triton_v1|triton_v2|gluon|flydsl_v0 +# gluon is gfx950/CDNA4-only (lazily imported; asserts arch when selected). +# use_turbo_attention (when core_attention is built) still wins for the dense path. +export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-triton_v2} +export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v2} + +# Plan-9: FP8 (E4M3) Indexer QK path (CSA selector). Default OFF; flip with +# USE_V4_FP8_INDEXER=True. Passed as a CLI override so it reliably reaches the +# in-container config regardless of env propagation. +export USE_V4_FP8_INDEXER=${USE_V4_FP8_INDEXER:-False} + +# Plan-5 P29 (RESCOPED): wrap sinkhorn_normalize in HyperMixer with a +# cached torch.compile build. Default OFF here; the proxy script +# (run_deepseek_v4_flash_proxy.sh) flips it ON. After G32 + G33b are +# green, the default flips to True for the V4-Flash configs. +export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-False} + +# Plan-4 P27: TP-side guard for the V4 Triton kernels. +# The dense / HCA / CSA kernels operate on the local head slice (each +# rank only sees H/TP query heads) so TP-sharded execution is correct +# by construction (no in-kernel collective comm needed). Plan-4 unit +# tests / smoke gates exercise TP=1 only; emit a soft warning when a +# user enables the kernels at TP>1 so any TP-related regression is +# easy to attribute. TP=1 is the V4-Flash / V4-Pro release default +# (release configs use PP+EP for parallelism, never TP). +if echo "$USE_V4_ATTENTION_BACKEND $USE_V4_CSA_ATTENTION_BACKEND" | grep -q "triton" && [ "${PRIMUS_TP:-1}" -gt 1 ]; then + echo "[WARN] Plan-4 V4 Triton kernels enabled at PRIMUS_TP=${PRIMUS_TP}>1; this combination is not covered by Plan-4 unit tests / smoke gates (G28..G30 ran TP=1 only). Functionally the kernels operate per-rank on the local H/TP head slice, so this should work, but treat any TP>1 regression as a Plan-4 follow-up." +fi + +if [ "$PRECISION_TYPE" = "FP8" ]; then + # Default to the paper's ue8m0 microscaling (e4m3 + mxfp8); honor explicit + # FP8 / FP8_RECIPE overrides. Sentinel-aware because "null" is non-empty, so + # a plain ${FP8:-...} would keep the off-sentinel instead of defaulting. + [ "$FP8" = "null" ] && export FP8=e4m3 + [ "$FP8_RECIPE" = "null" ] && export FP8_RECIPE=mxfp8 +fi + +# ---------- MXFP8 + FP8 param-gather (Muon path; Megatron #4987 analogue) ---- +# Plan-9: combine the distributed-Muon (LayerWise) path with an MXFP8 forward +# recipe + FP8 parameter all-gather. Enable with FP8_PARAM_GATHER=True (best +# paired with OPTIMIZER=dist_muon + PRECISION_TYPE=FP8 FP8_RECIPE=mxfp8). +# MXFP8 on ROCm/TE requires NVTE_ROCM_ENABLE_MXFP8=1; the mxfp8 param-AG path +# is most memory-efficient with --reuse-grad-buf-for-mxfp8-param-ag. NOTE: +# Megatron auto-disables --fp8-param-gather on TE>=2.0.0 (falls back to a +# bf16/all_gather), so on such containers this exercises the MXFP8 forward + +# dist-Muon path with param-gather requested-but-possibly-downgraded. +export FP8_PARAM_GATHER=${FP8_PARAM_GATHER:-False} +FP8_PARAM_GATHER_CLI_ARGS=() +if [ "$FP8_PARAM_GATHER" = "True" ]; then + export NVTE_ROCM_ENABLE_MXFP8=${NVTE_ROCM_ENABLE_MXFP8:-1} + export REUSE_GRAD_BUF_FOR_MXFP8_PARAM_AG=${REUSE_GRAD_BUF_FOR_MXFP8_PARAM_AG:-True} + FP8_PARAM_GATHER_CLI_ARGS=(--fp8_param_gather True) + if [ "$REUSE_GRAD_BUF_FOR_MXFP8_PARAM_AG" = "True" ] && [ "$FP8_RECIPE" = "mxfp8" ]; then + FP8_PARAM_GATHER_CLI_ARGS+=(--reuse_grad_buf_for_mxfp8_param_ag True) + fi +fi + +PP_LAYOUT_ARGS=() +if [ -n "${PRIMUS_PP_LAYOUT:-}" ]; then + PP_LAYOUT_ARGS=(--pipeline_model_parallel_layout "$PRIMUS_PP_LAYOUT") +fi + +PRIMUS_RECOMPUTE_LAYERS=${PRIMUS_RECOMPUTE_LAYERS:-0} + +export EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml} +export BACKEND_PATH=${BACKEND_PATH:-"$(pwd)/third_party/Megatron-LM"} +export PRIMUS_TEAM=${PRIMUS_TEAM:-amd} +export PRIMUS_USER=${PRIMUS_USER:-tas-mi355x-$(date +%Y%m%d)} +export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_smoke_${PRECISION_TYPE}_MBS${MBS}_GBS${GBS}_PP${PRIMUS_PP}_EP${PRIMUS_EP}} + +if [ ! -d "$BACKEND_PATH" ] || [ -z "$(ls -A "$BACKEND_PATH" 2>/dev/null)" ]; then + echo "[ERROR] BACKEND_PATH does not exist or is empty: $BACKEND_PATH" + echo "Run: git submodule update --init --recursive" + exit 1 +fi + +mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" + +export PRIMUS_EXIT_FAST=1 + +# Launcher: slurm (default, multi-node cluster) or direct (single-node, already +# inside the container — e.g. local smoke on one box). PRIMUS_LAUNCHER=direct +# drops the SLURM/srun + docker-image wrap that 'direct' doesn't use. +export PRIMUS_LAUNCHER=${PRIMUS_LAUNCHER:-slurm} +if [ "$PRIMUS_LAUNCHER" = "direct" ]; then + LAUNCHER_ARGS=(direct) +else + LAUNCHER_ARGS=(slurm -N "$NNODES") + [ -n "${SLURM_PARTITION:-}" ] && LAUNCHER_ARGS+=(--partition="${SLURM_PARTITION}") + [ -n "${SLURM_NODELIST:-}" ] && LAUNCHER_ARGS+=(--nodelist="${SLURM_NODELIST}") + LAUNCHER_ARGS+=(-- --image "${DOCKER_IMAGE}" --clean -- --numa) +fi + +./primus-cli "${LAUNCHER_ARGS[@]}" \ + -- train pretrain --config "$EXP" \ + --manual_gc True \ + --manual_gc_interval 100 \ + --pp_warmup "${PP_WARMUP:-True}" \ + "${PP_LAYOUT_ARGS[@]}" \ + --moe_router_force_load_balancing True \ + --log_avg_skip_iterations 3 \ + --backend_path "$BACKEND_PATH" \ + --num_layers "$PRIMUS_TOTAL_LAYERS" \ + --train_iters "$TRAIN_ITERS" \ + --lr_warmup_iters 0 \ + --lr_decay_iters "$TRAIN_ITERS" \ + --micro_batch_size "$MBS" \ + --global_batch_size "$GBS" \ + --seq_length "$PRIMUS_SEQ_LENGTH" \ + --max_position_embeddings "$PRIMUS_MAX_POSITION_EMBEDDINGS" \ + --rope_type rope \ + --tensor_model_parallel_size "$PRIMUS_TP" \ + --pipeline_model_parallel_size "$PRIMUS_PP" \ + --expert_model_parallel_size "$PRIMUS_EP" \ + --num_experts "$PRIMUS_NUM_EXPERTS" \ + --moe_router_topk "$PRIMUS_MOE_TOPK" \ + --moe_router_enable_expert_bias "$PRIMUS_MOE_ENABLE_EXPERT_BIAS" \ + --moe_ffn_hidden_size "$PRIMUS_MOE_FFN_HIDDEN_SIZE" \ + --index_topk "$PRIMUS_INDEX_TOPK" \ + --v4_grouped_experts_support_clamped_swiglu "$PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU" \ + --compress_ratios "$PRIMUS_COMPRESS_RATIOS" \ + --mtp_num_layers "${MTP_NUM_LAYERS:-0}" \ + --mock_data True \ + --enable_primus_turbo "$ENABLE_PRIMUS_TURBO" \ + --use_turbo_attention "$USE_TURBO_ATTENTION" \ + --use_v4_attention_backend "$USE_V4_ATTENTION_BACKEND" \ + --use_v4_csa_attention_backend "$USE_V4_CSA_ATTENTION_BACKEND" \ + --use_v4_fp8_indexer "$USE_V4_FP8_INDEXER" \ + --use_v4_compiled_sinkhorn "$USE_V4_COMPILED_SINKHORN" \ + --use_turbo_deepep "$USE_TURBO_DEEPEP" \ + "${TURBO_DEEPEP_CLI_ARGS[@]}" \ + --use_turbo_grouped_gemm "$TURBO_USE_GROUPED_MLP" \ + --moe_use_legacy_grouped_gemm "$LEGACY_GG" \ + "${OPTIMIZER_CLI_ARGS[@]}" \ + --fp8 "$FP8" \ + --fp8_recipe "$FP8_RECIPE" \ + "${FP8_PARAM_GATHER_CLI_ARGS[@]}" \ + --recompute_num_layers "$PRIMUS_RECOMPUTE_LAYERS" \ + --recompute_granularity full \ + --recompute_method block \ + --overlap_grad_reduce "$PRIMUS_OVERLAP_GRAD_REDUCE" \ + --overlap_param_gather "$PRIMUS_OVERLAP_PARAM_GATHER" \ + --disable_last_saving True \ + --disable_wandb True \ + --disable_tensorboard True \ + --profile "$PROFILE" \ + --use_pytorch_profiler "$PROFILE" \ + --profile_step_end 7 \ + --profile_step_start 6 \ + --bias_swiglu_fusion "$PRIMUS_BIAS_SWIGLU_FUSION" \ + 2>&1 | tee "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/log_node_${NODE_RANK:-0}.txt" diff --git a/examples/deepseek-v4/run_deepseek_v4_flash.sh b/examples/deepseek-v4/run_deepseek_v4_flash.sh new file mode 100644 index 000000000..5c5a0daa9 --- /dev/null +++ b/examples/deepseek-v4/run_deepseek_v4_flash.sh @@ -0,0 +1,64 @@ +#!/bin/bash + +set -euo pipefail + +export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-43} +export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-256} +export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6} +export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-2048} +export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-512} +export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-'"[0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]"'} +export MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-1} + +export NNODES=${NNODES:-8} + +if [ "$NNODES" -eq 8 ]; then + export PRIMUS_TP=${PRIMUS_TP:-1} + export PRIMUS_PP=${PRIMUS_PP:-8} + export PRIMUS_EP=${PRIMUS_EP:-8} + export PRIMUS_RECOMPUTE_LAYERS=0 + if [ "$MTP_NUM_LAYERS" -eq 1 ]; then + export PRIMUS_PP_LAYOUT='"Et*4|t*5|(t*6|)*5,t*4mL"' + else + export PRIMUS_PP_LAYOUT='"Et*4|t*5|(t*6|)*5,t*4L"' + fi +elif [ "$NNODES" -eq 4 ]; then + export PRIMUS_TP=${PRIMUS_TP:-1} + export PRIMUS_PP=${PRIMUS_PP:-4} + export PRIMUS_EP=${PRIMUS_EP:-8} + export PRIMUS_RECOMPUTE_LAYERS=3 + if [ "$MTP_NUM_LAYERS" -eq 1 ]; then + export PRIMUS_PP_LAYOUT='"Et*10|t*11|t*11|t*11mL"' + else + export PRIMUS_PP_LAYOUT='"Et*10|t*11|t*11|t*11L"' + fi +fi + +export MBS=${MBS:-1} +export GBS=${GBS:-$((64 * NNODES * MBS))} +export TRAIN_ITERS=${TRAIN_ITERS:-10} + +export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-4096} +export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}} + +export USE_V4_FP8_INDEXER=${USE_V4_FP8_INDEXER:-True} +export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-triton_v2} +export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v2} +export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-True} +export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-True} +export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-True} +export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False} +export PRIMUS_V4_ATTN_BWD_USE_SPLIT=${PRIMUS_V4_ATTN_BWD_USE_SPLIT:-1} +export PRIMUS_V4_CSA_BWD_SEGREDUCE=${PRIMUS_V4_CSA_BWD_SEGREDUCE:-1} +export PRIMUS_STACK_GROUPED_WEIGHT_TRITON=${PRIMUS_STACK_GROUPED_WEIGHT_TRITON:-1} +export PRIMUS_ROPE_TRITON=${PRIMUS_ROPE_TRITON:-1} +export PRIMUS_SINKHORN_TRITON=${PRIMUS_SINKHORN_TRITON:-1} +export PRIMUS_HC_TRITON=${PRIMUS_HC_TRITON:-1} +export PRIMUS_INDEXER_TRITON=${PRIMUS_INDEXER_TRITON:-1} +export PRIMUS_INDEXER_TRITON_FULL=${PRIMUS_INDEXER_TRITON_FULL:-0} +export PRIMUS_V4_ROUTER_TRITON=${PRIMUS_V4_ROUTER_TRITON:-1} +export PROFILE=${PROFILE:-False} +export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_flash_proxy_pp${PRIMUS_PP}_ep${PRIMUS_EP}_seq${PRIMUS_SEQ_LENGTH}} + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +exec "${SCRIPT_DIR}/run_deepseek_v4.sh" diff --git a/examples/deepseek-v4/run_deepseek_v4_flash_proxy.sh b/examples/deepseek-v4/run_deepseek_v4_flash_proxy.sh new file mode 100755 index 000000000..1552fd6ff --- /dev/null +++ b/examples/deepseek-v4/run_deepseek_v4_flash_proxy.sh @@ -0,0 +1,280 @@ +#!/bin/bash +############################################################################### +# DeepSeek-V4 Flash perf PROXY runner (latest config: plan-5 P32 final). +# +# Wraps `run_deepseek_v4.sh` with a V4-Flash production-shape proxy: +# +# - num_layers 8 (vs production 43; PROXY) +# - hidden_size 4096 (full V4-Flash; from yaml) +# - num_heads 64 (full V4-Flash; from yaml) +# - head_dim 512 (full V4-Flash; from yaml) +# - num_experts 256 (full V4-Flash; PROXY-friendly +# 32 experts/rank at EP=8) +# - moe_router_topk 6 (full V4-Flash) +# - moe_ffn_hidden 2048 (full V4-Flash) +# - index_topk 512 (full V4-Flash CSA top-K) +# - compress_ratios [0,0,4,128,4,128,4,0] (8-layer slice exercising +# every layer kind: 3 cr=0, +# 3 cr=4, 2 cr=128) +# - parallel TP=1 PP=1 EP=8 (single-node 8 GPU) +# - seq_length 4096 (default) (V4 pretrain target; +# set PRIMUS_SEQ_LENGTH=2048 +# / 1024 / 512 on the +# command line if OOM) +# +# Plan-5 perf knobs default ON: +# - USE_V4_ATTENTION_BACKEND (cr ∈ {0, 128} dense/HCA backend; default triton_v2) +# - USE_V4_CSA_ATTENTION_BACKEND (cr == 4 CSA backend; default triton_v2) +# - USE_TURBO_DEEPEP (PrimusTurboDeepEPTokenDispatcher) +# - TURBO_USE_GROUPED_MLP (Turbo grouped-GEMM MoE expert path) +# - USE_V4_COMPILED_SINKHORN (P29: torch.compile-fused Sinkhorn, +# kills the 7.6 s aten::sum fp32 reduce +# that dominated the P28 baseline) +# +# Plan-5 P32 final attention-kernel knobs (also default ON in code; surfaced +# here for visibility / easy A/B): +# - PRIMUS_V4_ATTN_BWD_USE_SPLIT (atomic-free split V4 attention BWD: +# dQ kernel + dK/dV kernel, each writes +# its own disjoint tiles via tl.store +# instead of atomic_add on a shared buf) +# - PRIMUS_V4_CSA_BWD_SEGREDUCE (atomic-free CSA pool BWD via per-visit +# partial buffer + sorted inverse-index +# segmented reduction into dpool) +# +# These two relied on the **plan-5 P32 dual-RoPE bf16 cast fix** in +# `apply_interleaved_partial_rope` (`primus/backends/megatron/core/transformer/ +# dual_rope.py`) to actually win in the proxy: pre-fix, cos/sin from +# `position_ids.float() * inv_freq` was fp32, so `bf16 * fp32 = fp32` +# silently upcast Q / K leaving RoPE — every V4 attention kernel paid 2x +# HBM traffic and ran the slow fp32-specialised Triton binary, inflating +# kernel times 1.8-7x in the proxy and masking the split / segreduce wins. +# The one-line cast of cos/sin to `x.dtype` after the unsqueeze lets the +# microbench-optimal kernels also win end-to-end. See +# `deepseek-v4/develop/progress/p32/p32-summary.md` for the full +# diagnostic walk-through. +# +# USE_TURBO_ATTENTION stays OFF — Turbo would take precedence over the V4 +# Triton dense path in `DeepseekV4Attention.forward` (plan-4 P27 dispatch +# precedence: turbo > v4_triton > eager for cr ∈ {0, 128}). +# +# Steady-state perf (P32 final, mi355-gpu-8 / dev_primus_wenx_693, +# iter 10 of 10, ${VAR:-DEFAULT} only): +# +# iter time : 603 ms / iter (vs P28 baseline 8837 ms; 14.64x) +# TFLOP/s/GPU : 1134 (vs P28 baseline 77.5) +# HBM peak / rank : ~170 GiB +# +# Every override is `${VAR:-DEFAULT}`-guarded, so the caller can flip any +# knob via `PRIMUS_SEQ_LENGTH=2048 ./run_deepseek_v4_flash_proxy.sh` etc. +# without editing the script. +# +# Usage: +# ./run_deepseek_v4_flash_proxy.sh # 10-iter smoke +# TRAIN_ITERS=20 ./run_deepseek_v4_flash_proxy.sh # longer warmup pass +# PRIMUS_V4_ATTN_BWD_USE_SPLIT=0 ./run_deepseek_v4_flash_proxy.sh # fall back to +# # monolithic V4 BWD +# PRIMUS_V4_CSA_BWD_SEGREDUCE=0 ./run_deepseek_v4_flash_proxy.sh # fall back to +# # gather+atomic CSA BWD +# PRIMUS_V4_DIAG_TIME=1 ./run_deepseek_v4_flash_proxy.sh # dump per-call +# # cuda.Event timings for +# # v4_attention (rank 0) +# +# Profile is intentionally OFF in this script (this is the SMOKE / perf +# runner, not the trace capture). For chrome-trace capture use +# `deepseek-v4/develop/progress/p32/run_baseline_trace_ep8_p32_final.sh` +# (mirrors the plan-4 P25 / plan-3 P23 profile-script pattern). +############################################################################### +set -euo pipefail + +# ---------- V4-Flash production widths (8-layer proxy slice) ---------------- +export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-8} +export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-256} +export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6} +export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-2048} +export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-512} +# 8-layer slice — every V4 attention layer kind exercised: +# layer 0 / 1 : cr=0 (dense + SWA + sink) +# layer 2 : cr=4 (CSA) +# layer 3 : cr=128 (HCA) +# layer 4 : cr=4 (CSA) +# layer 5 : cr=128 (HCA) +# layer 6 : cr=4 (CSA) +# layer 7 : cr=0 (dense + SWA + sink) -- V4-Flash production has +# cr=0 first/last layer +export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-"[0,0,4,128,4,128,4,0]"} + +# ---------- Single-node EP=8 ------------------------------------------------ +export PRIMUS_TP=${PRIMUS_TP:-1} +export PRIMUS_PP=${PRIMUS_PP:-1} +export PRIMUS_EP=${PRIMUS_EP:-8} + +# DP=8 with TP=1 PP=1 EP=8 on 8 GPUs (EP shards experts within DP group). +# GBS=8, MBS=1 -> 1 microbatch / DP rank / iter. Profiling-friendly: +# minimises iter-to-iter variance + keeps activation memory bounded. +export MBS=${MBS:-1} +export GBS=${GBS:-8} + +# ---------- Production seq length target ------------------------------------ +# The CSA wrapper-side gather (plan-4 P26) materialises +# [B, H, Sq, K_topk, D] = [1, 64, Sq, 512, 512] * 2 bytes per microbatch +# in HBM. At Sq=4096 that is 64 GiB / microbatch on top of the 256-expert +# MoE state (~12 GiB / rank for 8 layers) + KV cache + activations + +# optimizer state — likely OOMs at MI355X (192 GiB HBM). The plan-5 P28 +# task is to CALIBRATE this value (try 4096 -> 2048 -> 1024 -> 512) and +# document the chosen value in `develop/profile/profile-baseline-ep8-*`. +# Plan-5 P31 (in-kernel `topk_idxs` gather) is the structural fix that +# eventually lets this default reach 4096. +export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-4096} +export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}} + +# ---------- Plan-5 perf knobs (all five ON) --------------------------------- +export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-triton_v2} +export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v2} +export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-True} +export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-True} +# Plan-5 P29 (RESCOPED): torch.compile-fused HyperMixer Sinkhorn. Kills +# the 7.6 s aten::sum fp32 reduce (87.3 % of step time in the P28 +# baseline trace). Default ON in the proxy after G32 + G33b are green. +export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-True} + +# Turbo attention OFF — would take precedence over V4 Triton dense path +# in DeepseekV4Attention.forward (plan-4 P27 dispatch precedence: +# turbo > v4_triton > eager for cr ∈ {0, 128} +# v4_triton_csa > eager for cr == 4 ). +export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False} + +# ---------- Plan-5 P32 final attention-kernel knobs (split + segreduce) ----- +# Both default ON in the kernel code post-RoPE-fix; surface them here so +# the proxy script self-documents the P32 final perf recipe and so a quick +# A/B fallback is a single env-var flip. See header for the full root-cause +# write-up. +export PRIMUS_V4_ATTN_BWD_USE_SPLIT=${PRIMUS_V4_ATTN_BWD_USE_SPLIT:-1} +export PRIMUS_V4_CSA_BWD_SEGREDUCE=${PRIMUS_V4_CSA_BWD_SEGREDUCE:-1} + +# ---------- Plan-6 elemwise-fusion knobs (default ON; A/B with =0) ---------- +# Plan-6 P40 close-out (2026-05-15): each plan-6 phase that wins the EP=8 +# proxy A/B adds its env knob here as default ON, mirroring the plan-5 P32 +# final precedent above. Phases that microbench-win but proxy-noise-lose +# (P38, P39) ship as default OFF with the kernel checked in for future +# tuning. Cumulative plan-6 win at this composition (P34..P37 ON, P38/P39 +# OFF): **-92.7 ms / iter (-15.4 %) vs plan-5 P32 final**; steady-state +# iter time 510.6 ms / 524.9 TFLOP/s/GPU (peak HBM 172.3 GiB / rank). +# See `deepseek-v4/develop/perf/proxy_ep8.md` row `P40 final` and +# `progress/p40/p40-summary.md` for the full close-out write-up. +# +# The kernel code already defaults each to "1" / "0" appropriately; this +# block makes the runner script self-document the recipe and lets users +# flip individual fusions for A/B without editing source. +# +# P34 — stack_grouped_weight Triton FWD/BWD fusion in +# PrimusTurboGroupedMLP._stack_grouped_linear_weight. +# EP=8 proxy A/B win: 580.65 -> 530.85 ms / iter, -49.8 ms (-8.6%); +# TFLOP/s/GPU 463.2 -> 507.2, +9.5%; lm_loss bit-identical (pure +# layout transform). Default ON since 29baf151 (2026-05-14). +export PRIMUS_STACK_GROUPED_WEIGHT_TRITON=${PRIMUS_STACK_GROUPED_WEIGHT_TRITON:-1} + +# P35 — apply_interleaved_partial_rope Triton FWD/BWD fusion in +# dual_rope.py. Collapses the 9-op eager chain (slice + reshape + +# four broadcast muls + stack + reshape + cat) into one Triton +# kernel that does a single contiguous write with the rotation +# baked in. +# EP=8 proxy A/B win: 531.7 -> 526.7 ms / iter, -5.0 ms (-0.94%); +# TFLOP/s/GPU 507.1 -> 513.3, +1.2%; lm_loss bit-identical (pure +# analytic rotation). Default ON since landing (2026-05-14). +export PRIMUS_ROPE_TRITON=${PRIMUS_ROPE_TRITON:-1} + +# P36 — sinkhorn_normalize Triton FWD/BWD fusion in +# hyper_connection.py. Replaces the plan-5 P29 ``torch.compile`` +# cached Sinkhorn body with a hand-rolled Triton kernel that runs +# the 1 + 2*(n_iters - 1) alternating row/col normalize trajectory +# in registers per row of the leading axis (V4-Flash uses K=4). +# Microbench at V4-Flash K=4 (B=1, S=4096): +# FWD 0.045 ms (vs eager 0.600 ms = 13.4x; vs P29 compiled 0.270 +# ms = 6.0x) +# BWD 0.105 ms (vs eager 1.520 ms = 14.5x; vs P29 compiled 0.628 +# ms = 6.0x) +# The compiled-region overhead (`Torch-Compiled Region` ~21 ms / 16 +# calls + `CompiledFunctionBackward` ~41 ms / 16 calls) is removed +# entirely. Default ON since landing (2026-05-14). +export PRIMUS_SINKHORN_TRITON=${PRIMUS_SINKHORN_TRITON:-1} + +# P37 — HyperConnection compute_weights elemwise tail Triton fusion in +# hyper_connection.HyperMixer.compute_weights. Fuses the 3 slices + +# 3 fused-multiply-adds + 2 sigmoid + 1 softmax + 2 eps adds (the +# post-_packed_logits, pre-Sinkhorn chain) into one FWD + one BWD +# Triton kernel. Microbench at V4-Flash K=4 (B=1, S=4096): +# FWD 0.044 ms (vs eager 0.102 ms = 2.34x) +# BWD 0.276 ms (vs eager 0.405 ms = 1.47x) +# The matmul inside _packed_logits stays as F.linear; collapse / expand +# (matmul-adjacent) stay eager too -- they are not net wins as +# separate Triton kernels. Default ON since landing (2026-05-14). +export PRIMUS_HC_TRITON=${PRIMUS_HC_TRITON:-1} + +# P41 — Indexer.forward post-einsum tail Triton fusion. Re-attempt +# of P38 that keeps the cuBLAS / hipBLASLt einsum eager and fuses +# only the bandwidth-bound tail (`relu + mul(w_i) + sum(H) + +# causal_mask`). +# +# Plan-8 P57 close-out 2 (2026-05-15): default flipped to ON. +# Microbench at V4-Flash widths is a clear positive +# (FWD 4.30x / BWD 1.63x); the EP=8 proxy A/B (10-iter smoke) +# showed ~0.2 ms / iter aggregate gain within the ±1 ms noise band +# (small but consistently positive). We default ON so the +# bandwidth-bound tail is fused by default; set +# PRIMUS_INDEXER_TRITON=0 to revert to the eager body. +# +# The env knob was re-purposed at P41: it now controls the +# post-einsum tail path. Legacy P38 full-fuse path lives behind +# PRIMUS_INDEXER_TRITON_FULL (default OFF, kept in tree for small- +# shape paths and future tuning). +export PRIMUS_INDEXER_TRITON=${PRIMUS_INDEXER_TRITON:-1} +export PRIMUS_INDEXER_TRITON_FULL=${PRIMUS_INDEXER_TRITON_FULL:-0} + +# P39 — V4 Router post-logits Triton FWD/BWD fusion (shared by topk + +# hash router). +# +# Plan-8 P57 close-out 2 (2026-05-15): default flipped to ON. +# Microbench at V4-Flash widths (N=4096, E=256, K=8) wins on V4's +# production `sqrtsoftplus` score function (1.56x FWD / 1.22x BWD). +# P39 / P43 EP=8 proxy A/B was inside the proxy noise band, so the +# conservative landing posture left it default-OFF; for P57 R2 we +# default ON to keep the microbench-positive kernel on the production +# path. Set PRIMUS_V4_ROUTER_TRITON=0 to revert to the eager body. +export PRIMUS_V4_ROUTER_TRITON=${PRIMUS_V4_ROUTER_TRITON:-1} + +# ---------- Precision: FP8 training (paper ue8m0 microscaling) -------------- +# V4-Flash trains in FP8 by default. PRECISION_TYPE=FP8 makes run_deepseek_v4.sh +# emit --fp8 e4m3 --fp8_recipe mxfp8 (mxfp8 = the paper's ue8m0 microscaling, +# E8M0 block scale, native on MI355X/CDNA4). The FP4 expert / FP4-Indexer path +# is not yet wired in the Primus V4 integration ("Phase 2"), so experts run FP8 +# here (FP8-everywhere) rather than FP4 — the closest supported step. FP8 is +# outlier-sensitive, hence the clamped SwiGLU (swiglu_limit) in the EXP yaml. +# A/B back to BF16 with PRECISION_TYPE=BF16; override the recipe via FP8_RECIPE. +export PRECISION_TYPE=${PRECISION_TYPE:-FP8} +export EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml} + +# ---------- Profile OFF in the proxy smoke runner --------------------------- +# This script is the steady-state perf / smoke runner — kineto profiling +# stays OFF to avoid contaminating the iter timer with profiler-collection +# overhead. For chrome-trace capture use +# `deepseek-v4/develop/progress/p32/run_baseline_trace_ep8_p32_final.sh`. +export PROFILE=${PROFILE:-False} + +# ---------- Bookkeeping ----------------------------------------------------- +# Distinguish the proxy run output dir from the smoke run output dir so the +# trace-capture script + the smoke run land side-by-side without clobbering +# each other's logs. +export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_flash_proxy_pp${PRIMUS_PP}_ep${PRIMUS_EP}_seq${PRIMUS_SEQ_LENGTH}} + +# ---------- Launcher: single-node in-container by default ------------------- +# This proxy is the single-node 8-GPU smoke/perf runner, normally invoked from +# INSIDE the training container. run_deepseek_v4.sh defaults PRIMUS_LAUNCHER to +# `slurm` (which needs `srun` from a SLURM allocation and would fail in a bare +# container). Default to `direct` here so the proxy just torchruns locally; +# override with PRIMUS_LAUNCHER=slurm when launching from a cluster login node. +export PRIMUS_LAUNCHER=${PRIMUS_LAUNCHER:-direct} + +# Defer to run_deepseek_v4.sh for the actual training launch — every +# CLI flag and the primus-cli invocation lives there. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +exec "${SCRIPT_DIR}/run_deepseek_v4.sh" diff --git a/examples/deepseek-v4/run_deepseek_v4_pro_muon.sh b/examples/deepseek-v4/run_deepseek_v4_pro_muon.sh new file mode 100755 index 000000000..cea44a134 --- /dev/null +++ b/examples/deepseek-v4/run_deepseek_v4_pro_muon.sh @@ -0,0 +1,361 @@ +#!/bin/bash +############################################################################### +# DeepSeek-V4 *Pro* single-node bring-up with the Muon optimizer. +# +# Follows the DeepSeek-V4 paper Pro recipe (§4.2.1 architecture + §4.2.2 +# training setup) as closely as a single 8x288GB node allows: +# 1. Model : deepseek_v4_pro (61L / d7168 / 384 experts — paper §4.2.1). +# Selected via PRIMUS_MODEL, consumed by the +# `model: ${PRIMUS_MODEL:...}.yaml` line in the EXP yaml. +# Widths come from primus/configs/models/megatron/ +# deepseek_v4_pro.yaml; we only override the shape knobs the +# runner exposes (layers/experts/topk/ffn/index_topk/ +# compress_ratios) so the runner's smoke defaults don't win. +# 2. Reduced : 61 layers + seq 4096 do NOT fit one node, so cut depth +# to fit (PRIMUS_TOTAL_LAYERS) and seq (PRIMUS_SEQ_LENGTH); full/ +# uniform recompute on. (CSA-layer gather scales with seq.) +# 3. Optimizer : Muon (paper §4.2.2: Muon for matrices, AdamW for embedding/ +# pred-head/RMSNorm — in-tree ChainedOptimizer). Plain `muon` +# requires overlap_{grad_reduce,param_gather}=False and +# use_distributed_optimizer=False (Megatron asserts). +# 4. Training : paper §4.2.2 hyperparameters — momentum 0.95, update-RMS +# setup 0.18 (muon_extra_scale_factor), AdamW eps 1e-20, LR +# 2.0e-4→2.0e-5, balance-loss 1e-4, and (crucially) a LARGE +# batch via gradient accumulation toward the paper's 94.4M +# tokens/step. The batch is what amortizes the fixed Muon +# Newton-Schulz cost: at GBS=8 (accum 1) NS looks like ~97% +# of GEMM (a starved-batch artifact, NOT a Muon bug); at the +# paper batch it falls to the reported ~1-3%. +# +# Single-node integration gaps vs the paper (not config — Primus V4 TODO): +# - MTP depth 1 (MultiTokenPredictionLayer unsupported) -> MTP_NUM_LAYERS=0 +# - expert-bias/noaux_tc (Megatron needs sigmoid; V4 uses sqrtsoftplus) -> off +# - muon_weight_decay (0.01 vs paper 0.1) / mtp_loss_scaling are yaml-only. +# +# Usage: +# # paper-faithful single-node run (validated: ~890 TFLOP/s/GPU, Muon ~1% GPU): +# PRIMUS_TOTAL_LAYERS=2 PRIMUS_COMPRESS_RATIOS="[128,128]" \ +# PRIMUS_SEQ_LENGTH=4096 GBS=256 ./run_deepseek_v4_pro_muon.sh +# PRIMUS_TOTAL_LAYERS=4 PRIMUS_SEQ_LENGTH=512 GBS=8 \ +# ./run_deepseek_v4_pro_muon.sh # cheap validation +# OPTIMIZER=adam ./run_deepseek_v4_pro_muon.sh # A/B vs AdamW +# PRECISION_TYPE=BF16 ./run_deepseek_v4_pro_muon.sh # A/B vs BF16 (fp8 is default-on) +# PROFILE=True DISABLE_TENSORBOARD=False ... # capture 1-step trace +# +# Precision: FP8 training is ON by default (FP8=e4m3, FP8_RECIPE=tensorwise). +# Paper recipe is ue8m0/mxfp8 but it's not runnable on this gfx950 build (see the +# "FP8 training" block below); tensorwise gives the paper's fp8 layout on the +# weight GEMMs. PRECISION_TYPE=BF16 to A/B back to bf16. +############################################################################### +set -euo pipefail +set -x + +export HF_TOKEN="${HF_TOKEN:-}" + +export NNODES=${PET_NNODES:-1} +export TRAIN_ITERS=${TRAIN_ITERS:-10} +export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-1} + +# ---------- Model: DeepSeek-V4 Pro ----------------------------------------- +export PRIMUS_MODEL=${PRIMUS_MODEL:-deepseek_v4_pro} +export EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml} + +# ---------- Pro production widths (paper §4.2.1) ---------------------------- +export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-384} +export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6} +# Megatron only supports aux-loss-free expert bias with the sigmoid score +# function; V4 uses sqrtsoftplus, so disable expert bias (matches the working +# run_deepseek_v4.sh smoke; balancing falls back to seq_aux_loss). +export PRIMUS_MOE_ENABLE_EXPERT_BIAS=${PRIMUS_MOE_ENABLE_EXPERT_BIAS:-False} +export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-3072} +export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-1024} + +# ---------- Reduced depth + seq to fit single node ------------------------- +# MEASURED CEILING (chi2774, 8x288GB, 2026-05-20): with full Pro width +# (384 experts) + Muon, *4 layers @ seq 512 = 268.8 GB/rank (93%)* is about +# the single-node max. The binding cost is weights + Muon's fp32 optimizer +# states (Muon forces use_precision_aware_optimizer=False, so states cannot +# be bf16) — NOT activations, so lowering seq does not buy more layers. +# 5+ layers OOMs; for more depth use fewer experts or multi-node. +export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-4} +export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-512} +export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}} + +# Per-layer compression schedule, length == PRIMUS_TOTAL_LAYERS, mirroring the +# Pro pattern: first two HCA(128), then CSA(4)/HCA(128) interleaved, last +# dense+SWA(0). (Pro full yaml: idx0,1=128; idx>=2 even=4 / odd=128; last=0.) +export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-$(python3 - "$PRIMUS_TOTAL_LAYERS" <<'PY' +import sys +n=int(sys.argv[1]) +r=[] +for i in range(n): + if i<2: r.append(128) + elif i==n-1: r.append(0) + else: r.append(4 if i%2==0 else 128) +print("["+",".join(map(str,r))+"]") +PY +)} + +# ---------- Single-node EP=8 ----------------------------------------------- +export PRIMUS_TP=${PRIMUS_TP:-1} +export PRIMUS_PP=${PRIMUS_PP:-1} +export PRIMUS_EP=${PRIMUS_EP:-8} +export MBS=${MBS:-1} +# Paper Pro batch = 94.4M tokens/step (batch-size schedule). On one node we +# approach that regime via GRADIENT ACCUMULATION: grad_accum = GBS/(MBS*DP), +# DP=8 here. Large GBS is what amortizes the fixed Muon Newton-Schulz cost over +# many fwd/bwd microbatches (GBS=8 ⇒ accum=1 ⇒ NS runs every step ⇒ NS looks +# like ~97% of GEMM; that was a starved-batch artifact, NOT a Muon bug). At +# GBS=512 (accum 64) × seq 4096 = ~2.1M tokens/step the optimizer share drops +# toward the paper's reported 1-3%. +export GBS=${GBS:-512} + +# ---------- Optimizer: Muon (paper §4.2.2, values for BOTH Flash & Pro) ------ +export OPTIMIZER=${OPTIMIZER:-muon} +# The Primus Muon path (primus/backends/megatron/core/optimizer/moun.py) needs +# the emerging_optimizers package, which is NOT bundled in the container. Set +# PRIMUS_INSTALL_EMERGING_OPTIMIZERS so the in-container install hook +# (runner/.../01_install_emerging_optimizers.sh) provisions the pinned commit. +# Gated on a Muon optimizer so an OPTIMIZER=adam A/B run pays nothing. +if [ "$OPTIMIZER" = "muon" ] || [ "$OPTIMIZER" = "dist_muon" ]; then + export PRIMUS_INSTALL_EMERGING_OPTIMIZERS=${PRIMUS_INSTALL_EMERGING_OPTIMIZERS:-1} +fi +export MUON_MOMENTUM=${MUON_MOMENTUM:-0.95} # paper momentum 0.95 +# Paper: "rescale the RMS of each update matrix to 0.18 for reutilization of the +# AdamW learning rate." With scale_mode=spectral the realized update RMS ≈ +# extra_scale_factor (orth_grad RMS≈1/√max(m,n), times spectral scale √max(m,n)), +# so 0.18 maps directly here. Megatron default is 1.0 (≈5.5× too large vs paper). +export MUON_EXTRA_SCALE_FACTOR=${MUON_EXTRA_SCALE_FACTOR:-0.18} +# Newton-Schulz hardening knobs (matter for mxfp8, where NS can diverge on the +# quant-noised gradient). num_ns_steps = NS iterations (more = better convergence +# on ill-conditioned input); fp32_matmul_prec = precision of the NS matmuls +# ("medium" = tf32-ish, "high" = full fp32 — full precision keeps a near-σ=1 +# input from being pushed past the quintic's stable region by rounding error). +export MUON_NUM_NS_STEPS=${MUON_NUM_NS_STEPS:-5} +export MUON_FP32_MATMUL_PREC=${MUON_FP32_MATMUL_PREC:-medium} +# NOTE: muon_weight_decay is yaml-only (trainer_base.yaml=0.01); paper=0.1. +# No CLI flag, so it stays 0.01 here unless overridden in an EXP yaml. +# Muon hard requirements (Megatron arguments.py:1422): +export USE_DISTRIBUTED_OPTIMIZER=${USE_DISTRIBUTED_OPTIMIZER:-False} +export USE_PRECISION_AWARE_OPTIMIZER=${USE_PRECISION_AWARE_OPTIMIZER:-False} + +# ---------- Paper §4.2.2 Pro training hyperparameters ----------------------- +export LR=${LR:-2.0e-4} # Pro peak LR (Flash exp yaml had 1e-5) +export MIN_LR=${MIN_LR:-2.0e-5} # Pro end LR +export ADAM_EPS=${ADAM_EPS:-1.0e-20} # paper AdamW eps (NOTE: needs decimal point — Primus parses "1e-20" as a string) +export MOE_AUX_LOSS_COEFF=${MOE_AUX_LOSS_COEFF:-0.0001} # paper balance-loss weight +# Paper MTP depth = 1, but the Primus V4 integration does NOT yet support the +# MTP layer ("Unsupported mtp_model_layer submodules type ... when instantiating +# MultiTokenPredictionLayer"), so default 0 here. Set =1 once V4 MTP lands. +export MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-0} + +# ---------- Perf knobs (V4 Triton attn + Turbo MoE; same family as proxy) -- +export ENABLE_PRIMUS_TURBO=${ENABLE_PRIMUS_TURBO:-True} +export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False} +export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-True} +export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-True} +# Primus Sync-Free MoE (eliminates the DeepEP host busy-wait on the variable +# per-expert token counts): 0=off, 1=fused router/permute, 2=+no CPU busy-wait +# (turbo deepep + grouped mlp), 3=fully sync-free (+fused act). Stage >=2 needs +# use_turbo_grouped_gemm=True. Auto-enables the required sub-flags. Default 0. +export TURBO_SYNC_FREE_MOE_STAGE=${TURBO_SYNC_FREE_MOE_STAGE:-0} +# Phase 1b: route the dense/attention projections (q_down/kv/o_a etc.) through +# Primus-Turbo linears so they pick up the mxfp8 (CK) path under the fp8 context. +# Default OFF (attention stays bf16, the validated baseline). Set True to enable +# fp8 attention/dense projections. Requires TP=1 and fp8_recipe in {tensorwise, +# blockwise,mxfp8}; the MLA monkey-patch is auto-skipped for V4 (see mla_patches). +export USE_TURBO_PARALLEL_LINEAR=${USE_TURBO_PARALLEL_LINEAR:-False} +# Per-module recipe (paper): routed experts in MXFP4 while the rest of the layer +# stays FP8. Works under the global FP8 recipe (no --fp4/--fp8 conflict): the +# PrimusTurbo grouped MLP routes expert GEMMs through native FP4 (hipBLASLt). +# Default OFF. When on, force the hipBLASLt FP4 backend (no AITER). +export MOE_EXPERTS_FP4=${MOE_EXPERTS_FP4:-False} +if [ "$MOE_EXPERTS_FP4" = "True" ]; then + export PRIMUS_TURBO_GEMM_BACKEND=${PRIMUS_TURBO_GEMM_BACKEND:-FP4:HIPBLASLT} +fi +# Phase 5 (paper): CSA-indexer QK score in FP4. Rounds q_i and K^{IComp} to +# MXFP4 before the QK product (STE backward); w_i + ReLU/sum tail stay BF16. +# Read directly by the Indexer via PRIMUS_INDEXER_FP4. Default OFF. +export INDEXER_FP4=${INDEXER_FP4:-False} +if [ "$INDEXER_FP4" = "True" ]; then + export PRIMUS_INDEXER_FP4=1 + # The indexer QK now runs a REAL MXFP4 gemm (pt.ops.gemm_fp4) — force the + # hipBLASLt FP4 backend (no AITER), same as the MXFP4 expert path. + export PRIMUS_TURBO_GEMM_BACKEND=${PRIMUS_TURBO_GEMM_BACKEND:-FP4:HIPBLASLT} +fi +# MXFP8 expert-weight caching: expert weights are constant within an optimizer +# step, so re-quantizing them every microbatch + recompute forward (the large +# _mxfp8_quant_weight_fwd kernel) is redundant. When on, PrimusTurboGroupedMLP +# prequantizes once per step and reuses the fp8 buffers (loss-neutral, faster). +# Costs extra bytes/param resident — watch HBM at depth. Only affects the +# mxfp8 (MX_BLOCKWISE) grouped path. Default OFF. +export CACHE_MXFP8_WEIGHT=${CACHE_MXFP8_WEIGHT:-False} +if [ "$CACHE_MXFP8_WEIGHT" = "True" ]; then + export PRIMUS_TURBO_CACHE_MXFP8_WEIGHT=1 +fi +# FP8 attention projections (paper recipe): route q-up / o-proj through the fp8 +# turbo linear instead of the bf16 gather/scatter native path. Only valid at +# TP=1 (gather/scatter are no-ops there); for TP>1 the turbo linear rejects +# gather_output/scatter-input and it stays bf16. Default OFF. +export V4_FP8_ATTN_PROJ=${V4_FP8_ATTN_PROJ:-False} +if [ "$V4_FP8_ATTN_PROJ" = "True" ]; then + export PRIMUS_V4_FP8_ATTN_PROJ=1 +fi +export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-triton_v2} +export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v2} +export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-False} +export PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU=${PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU:-True} + +# ---------- FP8 training (paper §4.x quantization / techblog §9.6) ---------- +# Paper recipe: "FP4 + FP8 Mixed" — MoE experts + CSA-Indexer QK in FP4 (MXFP4), +# EVERYTHING ELSE in FP8, all with the **ue8m0** microscaling scale format. +# On this stack the ue8m0 path is `fp8_recipe=mxfp8` → TE MXFP8BlockScaling → +# Primus-Turbo MX_BLOCKWISE granularity with scale_dtype=E8M0 (fp8_utils.py:148), +# i.e. the paper's exact scaling format, native on MI355X/CDNA4. +# +# Integration gap vs the paper (NOT config — Primus V4 TODO, develop techblog +# item 10 "Phase 2 FP4/FP8 Mixed"): the FP4 expert / FP4-Indexer path is not yet +# wired in V4, so experts run at FP8 here (FP8 everywhere) rather than FP4. This +# is the closest supported step toward the paper recipe. FP8 is highly outlier- +# sensitive, which is why the paper pairs it with clamped SwiGLU (swiglu_limit, +# already on above via PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU=True). +# Precision toggle — shared interface with run_deepseek_v4.sh / the flash proxy: +# PRECISION_TYPE=FP8 (default) -> e4m3 + tensorwise; BF16 -> fp8 off. +# FP8 / FP8_RECIPE still override directly (e.g. FP8_RECIPE=blockwise, or FP8=null). +export PRECISION_TYPE=${PRECISION_TYPE:-FP8} +if [ "$PRECISION_TYPE" = "FP8" ]; then + export FP8=${FP8:-e4m3} # forward fp8 format (paper E4M3); "hybrid" = E4M3 fwd / E5M2 bwd + # Paper recipe is ue8m0 microscaling (mxfp8), but mxfp8 is NOT runnable on this + # gfx950 build (turbo grouped-GEMM has no MX path; TE ROCm MXFP8 needs K%128==0, + # V4 has a K=224 proj). `tensorwise` is the working recipe — paper fp8 layout, + # non-ue8m0 scale. (other: blockwise [TE-ROCm unsupported] / delayed) + export FP8_RECIPE=${FP8_RECIPE:-tensorwise} + # mxfp8 (paper ue8m0) + Muon: TRAINS, but shows a transient early-training + # grad-norm spike. Earlier 8-iter runs only caught the spike and mislabelled it + # a divergence; a 40-iter run shows it SELF-HEALS and the loss descends cleanly. + # What's happening: mxfp8 quant noise ill-conditions the gradient early (random + # init), so Muon's Newton-Schulz update norm spikes for ~10-20 iters, then + # settles as the model organizes. RAW grads stay ~0.99 throughout; only the + # NS-orthogonalized UPDATE norm spikes (Muon-specific: Adam@same-config is flat). + # It is NOT a kernel bug (both MX GEMMs ~4% correct on REAL E2E inputs via + # capture-replay; error zero-mean). + # FIX (verified, L4/64-expert/GBS512/seq128, 40 iters): + # no warmup -> loss 12->0.80, grad-norm peak ~2.4e5 (settles to ~35) + # warmup=10 -> loss 12->0.44, grad-norm peak ~2.1e4 (12x lower), no NaN + # So LR warmup tames the transient AND improves the loss — and it is what the + # paper does. We therefore AUTO-ENABLE warmup for mxfp8 (default 10; override + # LR_WARMUP_ITERS). Two more NS-hardening knobs are exposed if needed: + # MUON_NUM_NS_STEPS (more iters) and MUON_FP32_MATMUL_PREC=high. tensorwise + # stays the conservative default (smooth per-tensor scale, no transient). + # NOTE: validated at reduced depth/width; full 61-layer/384-expert is a + # separate multi-GPU confirmation. + if [ "$FP8_RECIPE" = "mxfp8" ]; then + export LR_WARMUP_ITERS=${LR_WARMUP_ITERS:-10} + echo "[INFO] FP8_RECIPE=mxfp8 + Muon: expect a transient early grad-norm spike" >&2 + echo " (self-heals; LR warmup auto-set to ${LR_WARMUP_ITERS} to damp it)." >&2 + fi +else + export FP8=${FP8:-null} # PRECISION_TYPE=BF16 -> disable fp8 + export FP8_RECIPE=${FP8_RECIPE:-null} +fi + +TURBO_DEEPEP_CLI_ARGS=() +if [ "$USE_TURBO_DEEPEP" = "True" ]; then + export TURBO_DEEPEP_NUM_CU=${TURBO_DEEPEP_NUM_CU:-80} + export TURBO_DEEPEP_USE_COMM_STREAM=${TURBO_DEEPEP_USE_COMM_STREAM:-False} + export MOE_ROUTER_DTYPE=${MOE_ROUTER_DTYPE:-fp32} + export MOE_SHARED_EXPERT_OVERLAP=${MOE_SHARED_EXPERT_OVERLAP:-False} + TURBO_DEEPEP_CLI_ARGS=( + --turbo_deepep_num_cu "$TURBO_DEEPEP_NUM_CU" + --turbo_deepep_use_comm_stream "$TURBO_DEEPEP_USE_COMM_STREAM" + --moe_router_dtype "$MOE_ROUTER_DTYPE" + --moe_shared_expert_overlap "$MOE_SHARED_EXPERT_OVERLAP" + ) +fi + +export PROFILE=${PROFILE:-False} +# Profiler writes the chrome trace via tensorboard_trace_handler(args.tensorboard_dir), +# so tensorboard must be enabled for a trace run. Default True (smoke); set +# DISABLE_TENSORBOARD=False together with PROFILE=True to capture a trace. +export DISABLE_TENSORBOARD=${DISABLE_TENSORBOARD:-True} +export BACKEND_PATH=${BACKEND_PATH:-"$(pwd)/third_party/Megatron-LM"} +export PRIMUS_TEAM=${PRIMUS_TEAM:-amd} +export PRIMUS_USER=${PRIMUS_USER:-tas-mi355x-$(date +%Y%m%d)} +export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_pro_muon_L${PRIMUS_TOTAL_LAYERS}_seq${PRIMUS_SEQ_LENGTH}_ep${PRIMUS_EP}} + +if [ ! -d "$BACKEND_PATH" ] || [ -z "$(ls -A "$BACKEND_PATH" 2>/dev/null)" ]; then + echo "[ERROR] BACKEND_PATH does not exist or is empty: $BACKEND_PATH" + echo "Run: git submodule update --init --recursive" + exit 1 +fi + +mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" + +./primus-cli direct \ + -- train pretrain --config "$EXP" \ + --backend_path "$BACKEND_PATH" \ + --manual_gc True \ + --manual_gc_interval 100 \ + --num_layers "$PRIMUS_TOTAL_LAYERS" \ + --train_iters "$TRAIN_ITERS" \ + --lr_warmup_iters "${LR_WARMUP_ITERS:-0}" \ + --lr_decay_iters "$TRAIN_ITERS" \ + --micro_batch_size "$MBS" \ + --global_batch_size "$GBS" \ + --lr "$LR" \ + --min_lr "$MIN_LR" \ + --adam_eps "$ADAM_EPS" \ + --moe_aux_loss_coeff "$MOE_AUX_LOSS_COEFF" \ + --seq_length "$PRIMUS_SEQ_LENGTH" \ + --max_position_embeddings "$PRIMUS_MAX_POSITION_EMBEDDINGS" \ + --rope_type rope \ + --tensor_model_parallel_size "$PRIMUS_TP" \ + --pipeline_model_parallel_size "$PRIMUS_PP" \ + --expert_model_parallel_size "$PRIMUS_EP" \ + --num_experts "$PRIMUS_NUM_EXPERTS" \ + --moe_router_topk "$PRIMUS_MOE_TOPK" \ + --moe_router_force_load_balancing "${MOE_FORCE_LOAD_BALANCE:-False}" \ + --moe_router_enable_expert_bias "$PRIMUS_MOE_ENABLE_EXPERT_BIAS" \ + --moe_ffn_hidden_size "$PRIMUS_MOE_FFN_HIDDEN_SIZE" \ + --index_topk "$PRIMUS_INDEX_TOPK" \ + --v4_grouped_experts_support_clamped_swiglu "$PRIMUS_V4_GROUPED_EXPERTS_SUPPORT_CLAMPED_SWIGLU" \ + --compress_ratios "$PRIMUS_COMPRESS_RATIOS" \ + --mtp_num_layers "$MTP_NUM_LAYERS" \ + --mock_data True \ + --optimizer "$OPTIMIZER" \ + --muon_momentum "$MUON_MOMENTUM" \ + --muon_extra_scale_factor "$MUON_EXTRA_SCALE_FACTOR" \ + --muon_num_ns_steps "$MUON_NUM_NS_STEPS" \ + --muon_fp32_matmul_prec "$MUON_FP32_MATMUL_PREC" \ + --use_distributed_optimizer "$USE_DISTRIBUTED_OPTIMIZER" \ + --use_precision_aware_optimizer "$USE_PRECISION_AWARE_OPTIMIZER" \ + --main_grads_dtype fp32 \ + --exp_avg_dtype fp32 \ + --exp_avg_sq_dtype fp32 \ + --enable_primus_turbo "$ENABLE_PRIMUS_TURBO" \ + --use_turbo_attention "$USE_TURBO_ATTENTION" \ + --use_v4_attention_backend "$USE_V4_ATTENTION_BACKEND" \ + --use_v4_csa_attention_backend "$USE_V4_CSA_ATTENTION_BACKEND" \ + --use_v4_compiled_sinkhorn "$USE_V4_COMPILED_SINKHORN" \ + --use_turbo_deepep "$USE_TURBO_DEEPEP" \ + --turbo_sync_free_moe_stage "$TURBO_SYNC_FREE_MOE_STAGE" \ + "${TURBO_DEEPEP_CLI_ARGS[@]}" \ + --use_turbo_grouped_gemm "$TURBO_USE_GROUPED_MLP" \ + --use_turbo_gemm "$USE_TURBO_PARALLEL_LINEAR" \ + --moe_experts_fp4 "$MOE_EXPERTS_FP4" \ + --moe_use_legacy_grouped_gemm False \ + --fp8 "$FP8" \ + --fp8_recipe "$FP8_RECIPE" \ + --recompute_num_layers 1 \ + --recompute_granularity full \ + --recompute_method uniform \ + --overlap_grad_reduce False \ + --overlap_param_gather False \ + --disable_last_saving True \ + --disable_wandb True \ + --disable_tensorboard "$DISABLE_TENSORBOARD" \ + --profile "$PROFILE" \ + --use_pytorch_profiler "$PROFILE" \ + --profile_step_end 7 \ + --profile_step_start 6 \ + 2>&1 | tee "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/log_node_${NODE_RANK:-0}.txt" diff --git a/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh b/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh new file mode 100755 index 000000000..1ac99e498 --- /dev/null +++ b/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh @@ -0,0 +1,498 @@ +#!/bin/bash +############################################################################### +# DeepSeek-V4 *Pro* + Muon single-GPU bring-up on mi455 / gfx1250 (1 GPU). +# +# Single-GPU sibling of run_deepseek_v4_flash_proxy_1gpu.sh, for the Pro model. +# The upstream run_deepseek_v4_pro_muon.sh uses `primus-cli direct`, which +# assumes it is ALREADY inside the 8x288GB MI355X container at EP=8. This host +# is one gfx1250 box with no SLURM and one GPU, so this script instead wraps +# the SAME examples/run_pretrain.sh entrypoint in a local `docker run` (the +# validated gfx1250 docker + TransformerEngine recipe from +# ../../mi450/Primus/run_dsv3_proxy_4L.sh), selects the Pro model via +# PRIMUS_MODEL, and scales it down to a MINIMUM single-GPU proxy: +# +# - model deepseek_v4_pro hidden 7168 / 128 heads / head_dim +# 512 / o_groups 16 (full Pro widths +# from deepseek_v4_pro.yaml) +# - parallel TP=1 PP=1 EP=1 (single GPU; no SLURM, no DeepEP) +# - num_layers 4 (vs production 61; MINIMUM slice +# that still exercises every V4 +# attention layer kind) +# - compress_ratios [128,128,4,0] Pro pattern (first two HCA, then +# CSA, last dense+SWA) -> covers +# HCA cr=128 / CSA cr=4 / dense cr=0 +# - num_experts 48 topk 1 (production 384/topk6 div by 8 = +# the per-rank shape of the EP=8 +# production run; topk ceil(6/8)=1. +# ~48 experts x 66M x 4L + Muon fp32 +# states + seq-4096 activations -> 273 +# GB peak (measured), fits the 432 GiB +# card; full 384 will NOT fit.) +# - moe_ffn_hidden 3072 (full Pro MoE width; from yaml) +# - seq_length 4096 (raised from 512: at GBS=8 the fixed +# Muon Newton-Schulz cost otherwise +# dominates GPU time; 4096 +# tokens/microbatch amortizes it to a +# representative profile. CSA gather +# scales w/ seq. Set 512 for a fast smoke.) +# - index_topk 64 (CSA top-K; <= cr=4 pool seq/4, i.e. +# 4096/4=1024 at the default seq) +# - precision FP8 e4m3 + tensorwise (paper fp8 LAYOUT, per-tensor scale; +# mxfp8/ue8m0 is GUARDED off — diverges +# on this build. CK-free TE path. +# PRECISION_TYPE=BF16 for a BF16 A/B.) +# +# Optimizer: Muon (paper §4.2.2), same recipe as the upstream pro_muon runner: +# momentum 0.95, update-RMS scale 0.18, AdamW eps 1e-20, LR 2.0e-4->2.0e-5, +# balance-loss 1e-4. Muon hard-requires use_distributed_optimizer=False + +# use_precision_aware_optimizer=False (so optimizer states stay fp32 — this +# is the binding memory cost, NOT activations). Set OPTIMIZER=adam to A/B. +# NOTE: at the tiny single-GPU GBS the Newton-Schulz cost looks huge as a % +# of GEMM (starved-batch artifact, not a Muon bug); raise GBS to amortize. +# +# Correctness-first defaults (this is a "does V4-Pro train at all on 1 gfx1250 +# GPU" bring-up, not a perf push). Eager attention; Turbo/DeepEP/tilelang/ +# compiled-Sinkhorn/plan-6 Triton fusions all OFF; stock hipBLASLt; profiler +# OFF. Every knob is ${VAR:-DEFAULT}-guarded for command-line A/B. +# +# REQUIRED gfx1250 fix kept ON (single-GPU-safe): RCCL all_reduce(op=AVG) hangs +# even at world_size=1 on this build and Megatron's MoE aux-loss reduce uses +# AVG; sitecustomize on PYTHONPATH rewrites AVG -> SUM/world_size. See +# rccl_avg_workaround/sitecustomize.py. +# +# Usage: +# ./run_deepseek_v4_pro_muon_1gpu.sh # 10-iter smoke +# OPTIMIZER=adam ./run_deepseek_v4_pro_muon_1gpu.sh # A/B vs AdamW +# PRIMUS_TOTAL_LAYERS=6 ./run_deepseek_v4_pro_muon_1gpu.sh # deeper slice (watch HBM) +# PRIMUS_NUM_EXPERTS=8 PRIMUS_MOE_TOPK=2 ./run_deepseek_v4_pro_muon_1gpu.sh # tiny MoE +# HIP_VISIBLE_DEVICES=3 ./run_deepseek_v4_pro_muon_1gpu.sh # pin a card +############################################################################### +set -eo pipefail + +export DOCKER_IMAGE=${DOCKER_IMAGE:-registry-sc-harbor.amd.com/framework/therock-npi@sha256:feba897e2a32a2465b8b296ed2662b2ad6136b5f1cf6f6c2716a3674aafc30f3} +# Repo root: this script lives under examples/deepseek-v4/, so resolve two levels +# up. All paths below (TE_DIR, rccl_avg_workaround, PRIMUS_PATH) are repo-root-relative. +SCRIPT_DIR=$(realpath -m "$(dirname "$0")/../..") +export TE_DIR=${TE_DIR:-$(realpath -m "$SCRIPT_DIR/../../mi450/TransformerEngine")} +export TE_WHEEL_DIR=${TE_WHEEL_DIR:-$(realpath -m "$SCRIPT_DIR/../../mi450/dist/feba897")} + +# ---------- Attention backend env (TE side) -------------------------------- +export NVTE_FUSED_ATTN=1 +export NVTE_FUSED_ATTN_CK=0 +export NVTE_FUSED_ATTN_AOTRITON=1 +export NVTE_USE_CK_GEMM=0 +export NVTE_FLASH_ATTN=0 + +# ---------- hipBLASLt: STOCK by default; opt-in TUNED (PRIMUS_TUNED_HIPBLASLT=1) - +# Stock is the safe default: the older feba897 tuned bundle DEADLOCKED on a +# backward-FP8 GSU split-K kernel on this host (GPU wedge -> node reboot, +# 2026-06-10). PRIMUS_TUNED_HIPBLASLT=1 opts into a freshly built GridBased +# gfx1250 tuned library (qwen3/dsv3 tuned; swept clean on dsv4 fwd shapes) via +# LD_PRELOAD + HIPBLASLT_TENSILE_LIBPATH. This is a GUARDED EXPERIMENT: run with +# a watchdog and expect a possible node reboot if the backward path still wedges. +# NOTE: never export an EMPTY HIPBLASLT_TENSILE_LIBPATH into the container — a +# missing path breaks even stock hipBLASLt ("Cannot read TensileLibrary..."). +export PRIMUS_TUNED_HIPBLASLT=${PRIMUS_TUNED_HIPBLASLT:-0} +export HBL_TUNED_RELEASE=${HBL_TUNED_RELEASE:-/home/yanyuqin/hipblaslt/rocm-libraries/projects/hipblaslt/build/release} +if [ "$PRIMUS_TUNED_HIPBLASLT" = "1" ]; then + if [ ! -f "$HBL_TUNED_RELEASE/library/libhipblaslt.so.1" ]; then + echo "[hipblaslt] ERROR: tuned lib not found at $HBL_TUNED_RELEASE/library/libhipblaslt.so.1" >&2 + exit 1 + fi + # LD_PRELOAD / LD_LIBRARY_PATH are injected INSIDE the container (below) so the + # image's own rocm/torch lib paths are preserved (prepend, not override). + export HIPBLASLT_TENSILE_LIBPATH="$HBL_TUNED_RELEASE/Tensile/library/gfx1250" + echo "[hipblaslt] TUNED (opt-in): libpath=$HIPBLASLT_TENSILE_LIBPATH, LD_PRELOAD=libhipblaslt.so.1 (GUARDED: watch for wedge)" +else + unset HIPBLASLT_DIR HIPBLASLT_LD_PRELOAD HIPBLASLT_TENSILE_LIBPATH + echo "[hipblaslt] STOCK (container built-in gfx1250 catalog)" +fi + +# ---------- REQUIRED gfx1250 RCCL AVG->SUM workaround ----------------------- +export PYTHONPATH="$SCRIPT_DIR/rccl_avg_workaround:${PYTHONPATH:-}" +# Real primus_turbo imports flydsl at import time; put FLYDSL_PKG_DIR on PYTHONPATH. +export FLYDSL_PKG_DIR=${FLYDSL_PKG_DIR:-} +if [ -n "$FLYDSL_PKG_DIR" ] && [ -d "$FLYDSL_PKG_DIR/flydsl" ]; then + export PYTHONPATH="$FLYDSL_PKG_DIR:$PYTHONPATH" +fi + +# SDMA OFF on this host (run 3 debugging, 2026-06-10): an SDMA H2D copy +# intermittently never signals completion — py-spy --native showed the trainer +# pinned in rocr BusyWaitSignal under a trivial `torch.tensor(n, device=dev)` +# in Megatron get_batch, GPU 0%, dmesg clean. Killing the stuck proc then +# leaves MES unrecoverable (recovery disabled) -> node reboot. Blit-kernel +# copies are slower but don't use the flaky SDMA queues. This may ALSO be the +# true cause of the run-1 "permute autotune wedge" (same stuck-queue +# signature; permute fusion possibly innocent). +export HSA_ENABLE_SDMA=${HSA_ENABLE_SDMA:-0} + +# ---------- Distributed / NCCL: single GPU, loopback only ------------------- +export HSA_NO_SCRATCH_RECLAIM=1 +export NCCL_IB_DISABLE=1 +export NCCL_P2P_DISABLE=1 +export NCCL_IB_HCA= +export NCCL_SOCKET_IFNAME=lo +export GLOO_SOCKET_IFNAME=lo +export RCCL_DISABLE_AMDSMI=1 +export NCCL_AMDSMI_DISABLE=1 +export USING_AINIC=0 + +export GPUS_PER_NODE=1 +export NNODES=1 +export PYTHONUNBUFFERED=1 + +# REQUIRED gfx1250 workaround for V4-Pro (default ON). The Pro model build +# (hidden_size 7168, NOT a multiple of 4096) leaves a memory layout that wedges +# the process's first high-priority MES queue creation at iter-1 get_batch +# -> deadlock -> node reboot (debugged 2026-06-11; root cause = MES queue +# creation vs non-4096-aligned allocation layout). AMD_SERIALIZE_COPY=3 alone +# prevents it (kernels stay async; small iter-time cost on the eager +# proxy). Bisected: KERNEL serialize + LAUNCH_BLOCKING are NOT needed, so they +# default off. Set AMD_SERIALIZE_COPY=0 only to re-demonstrate the hang. +export AMD_SERIALIZE_COPY=${AMD_SERIALIZE_COPY:-3} +export AMD_SERIALIZE_KERNEL=${AMD_SERIALIZE_KERNEL:-0} +export HIP_LAUNCH_BLOCKING=${HIP_LAUNCH_BLOCKING:-0} + +# ---------- Model: DeepSeek-V4 Pro (selected via the EXP yaml model: line) -- +export PRIMUS_MODEL=${PRIMUS_MODEL:-deepseek_v4_pro} + +# ---------- Pro MINIMUM single-GPU proxy shape ------------------------------ +export PRIMUS_TP=${PRIMUS_TP:-1} +export PRIMUS_PP=${PRIMUS_PP:-1} +export PRIMUS_EP=${PRIMUS_EP:-1} +# 3 layers (down from 4): the 4-layer model sits near the 432GB gfx1250's +# capacity and OOMs at iter 2 (Muon keeps fp32 optimizer states). Dropping +# the last layer frees the headroom for a warm step. +export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-3} +export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-48} +export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-1} +export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-3072} +export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-64} +export PRIMUS_SEQ_LENGTH=${PRIMUS_SEQ_LENGTH:-4096} +export PRIMUS_MAX_POSITION_EMBEDDINGS=${PRIMUS_MAX_POSITION_EMBEDDINGS:-${PRIMUS_SEQ_LENGTH}} +export MBS=${MBS:-1} +export GBS=${GBS:-8} +export TRAIN_ITERS=${TRAIN_ITERS:-10} + +# Per-layer compression schedule, length == PRIMUS_TOTAL_LAYERS, mirroring the +# Pro pattern (first two HCA(128), then CSA(4)/HCA(128) interleaved, last +# dense+SWA(0)). Pure-bash generator (no host python needed). +gen_pro_compress_ratios() { + local n=$1 i r=() + for ((i = 0; i < n; i++)); do + if (( i < 2 )); then r+=(128) + elif (( i == n-1 )); then r+=(0) + elif (( i % 2 == 0)); then r+=(4) + else r+=(128) + fi + done + local IFS=, + echo "[${r[*]}]" +} +export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-$(gen_pro_compress_ratios "$PRIMUS_TOTAL_LAYERS")} + +# ---------- Optimizer: Muon (paper §4.2.2) --------------------------------- +export OPTIMIZER=${OPTIMIZER:-muon} +export MUON_MOMENTUM=${MUON_MOMENTUM:-0.95} +export MUON_EXTRA_SCALE_FACTOR=${MUON_EXTRA_SCALE_FACTOR:-0.18} +export USE_DISTRIBUTED_OPTIMIZER=${USE_DISTRIBUTED_OPTIMIZER:-False} +export USE_PRECISION_AWARE_OPTIMIZER=${USE_PRECISION_AWARE_OPTIMIZER:-False} +export LR=${LR:-2.0e-4} +export MIN_LR=${MIN_LR:-2.0e-5} +export ADAM_EPS=${ADAM_EPS:-1.0e-20} # needs a decimal point — Primus parses "1e-20" as a string +export MOE_AUX_LOSS_COEFF=${MOE_AUX_LOSS_COEFF:-0.0001} +export MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-0} # V4 MTP layer not yet supported in-tree +# Pro uses sqrtsoftplus; Megatron only supports aux-loss-free expert bias with +# sigmoid, so disable expert bias (balancing falls back to seq_aux_loss). +export PRIMUS_MOE_ENABLE_EXPERT_BIAS=${PRIMUS_MOE_ENABLE_EXPERT_BIAS:-False} + +# Muon needs fp32 optimizer states (precision-aware off forces this anyway). +OPT_DTYPE_ARGS="--main_grads_dtype fp32 --exp_avg_dtype fp32 --exp_avg_sq_dtype fp32" + +# ---------- Perf knobs: V4 attention backends ON; turbo paths OFF ---------- +# V4 attention backend (replaces the unfused/eager path). Covers the dense + +# HCA layers (compress_ratio in {0, 128}) via USE_V4_ATTENTION_BACKEND and the +# CSA layers (compress_ratio == 4) via USE_V4_CSA_ATTENTION_BACKEND. +# Validated on gfx1250 after the WMMA tile-floor fix (06ae5214). +export USE_V4_ATTENTION_BACKEND=${USE_V4_ATTENTION_BACKEND:-triton_v2} +export USE_V4_CSA_ATTENTION_BACKEND=${USE_V4_CSA_ATTENTION_BACKEND:-triton_v2} +export USE_TURBO_ATTENTION=${USE_TURBO_ATTENTION:-False} +export USE_TURBO_DEEPEP=${USE_TURBO_DEEPEP:-False} +export TURBO_USE_GROUPED_MLP=${TURBO_USE_GROUPED_MLP:-False} +# Projections: the FP8 yaml sets use_turbo_gemm=true (PrimusTurboLinear); +# turbo-free here -> TELinear. Override off so no turbo GEMM is invoked. +export USE_TURBO_PARALLEL_LINEAR=${USE_TURBO_PARALLEL_LINEAR:-False} +export USE_V4_COMPILED_SINKHORN=${USE_V4_COMPILED_SINKHORN:-False} +export PRIMUS_STACK_GROUPED_WEIGHT_TRITON=${PRIMUS_STACK_GROUPED_WEIGHT_TRITON:-0} +# RoPE Triton: default ON. Trace (2026-06-25, L3) attributed 960 kernels / 513 tiny +# (<5us) / 38.6 ms to the eager rotary-embedding path — a launch-bound fusion target. +export PRIMUS_ROPE_TRITON=${PRIMUS_ROPE_TRITON:-1} +# Sinkhorn Triton fused FWD/BWD: default ON. The eager Sinkhorn-Knopp loop +# (n_iters=20) launches ~18,600 tiny sum/add/div kernels per step (5,616 on the +# fwd side alone); the Triton path emits exactly 1 fwd + 1 bwd kernel per call. +# Measured 2026-06-25 (0612, L3, FP8): total GPU events 80,962 -> 62,340, sinkhorn +# GPU kernels 5,616 -> 48, warm step ~2,890 -> ~2,797 ms (+3.2%), 0 NaN / loss +# bit-identical. Falls back to eager when the shape/device is unsupported. Set =0 to A/B. +export PRIMUS_SINKHORN_TRITON=${PRIMUS_SINKHORN_TRITON:-1} +# HyperConnection mHC Triton: default ON. The mHC HyperMixer glue (pre/post/comb +# projections + scales), separate from the already-fused HC-expand and sinkhorn. +# Trace (2026-06-25, L3): 1,320 kernels / 872 tiny (<5us) / 52 ms — top remaining +# launch-bound target after sinkhorn. +export PRIMUS_HC_TRITON=${PRIMUS_HC_TRITON:-1} +# CSA indexer Triton: kept OFF — inert at L3 (compress_ratios [128,128,0] has NO CSA +# layer, so the indexer never runs). Enable only with a CSA layer (>=4 layers). +export PRIMUS_INDEXER_TRITON=${PRIMUS_INDEXER_TRITON:-0} +export PRIMUS_INDEXER_TRITON_FULL=${PRIMUS_INDEXER_TRITON_FULL:-0} +# V4 MoE router Triton: default ON. Trace (2026-06-25, L3): 432 kernels / 208 tiny / +# 5.3 ms — marginal, but launch-bound and correctness-neutral. +export PRIMUS_V4_ROUTER_TRITON=${PRIMUS_V4_ROUTER_TRITON:-1} + +export ENABLE_PRIMUS_TURBO=False +if [ "$USE_TURBO_ATTENTION" = "True" ] || [ "$USE_TURBO_DEEPEP" = "True" ] || [ "$TURBO_USE_GROUPED_MLP" = "True" ]; then + ENABLE_PRIMUS_TURBO=True +fi + +# MoE permute fusion OFF for Pro on gfx1250: the Triton permute_with_mask_map +# BACKWARD autotune wedges the GPU stream at Pro shapes (48 experts / hidden +# 7168) — cuda.synchronize inside triton do_bench never returns (debugged +# 2026-06-10 via py-spy; flash shapes 32 experts / hidden 4096 autotune fine). +# Eager permute is the safe path; flip =True to retry after a triton fix. +export MOE_PERMUTE_FUSION=${MOE_PERMUTE_FUSION:-False} + +export PROFILE=${PROFILE:-False} +# PyTorch profiler writes the chrome trace via tensorboard_trace_handler( +# args.tensorboard_dir), so the tensorboard dir MUST be enabled to get a trace. +# Default tensorboard off, but auto-enable it whenever PROFILE=True so a +# profiled run actually produces a trace. profile window = steps [START,END); +# need TRAIN_ITERS > PROFILE_STEP_END. +export DISABLE_TENSORBOARD=${DISABLE_TENSORBOARD:-True} +if [ "$PROFILE" = "True" ]; then export DISABLE_TENSORBOARD=False; fi +export PROFILE_STEP_START=${PROFILE_STEP_START:-6} +export PROFILE_STEP_END=${PROFILE_STEP_END:-7} +export PRIMUS_TEAM=${PRIMUS_TEAM:-amd} +export PRIMUS_USER=${PRIMUS_USER:-gfx1250-1gpu} +export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_pro_muon_1gpu_L${PRIMUS_TOTAL_LAYERS}_E${PRIMUS_NUM_EXPERTS}_seq${PRIMUS_SEQ_LENGTH}} + +PRIMUS_PATH="$SCRIPT_DIR" +DATA_PATH="${PRIMUS_PATH}/data" +mkdir -p "$DATA_PATH" + +EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml} +LOG=${LOG:-deepseek-v4-pro-muon-1gpu.log} + +# ---------- FP8 training (matches upstream run_deepseek_v4_pro_fp8_paper.sh) -- +# PRECISION_TYPE=FP8 (default) -> FP8=e4m3, FP8_RECIPE=tensorwise. This is the +# paper's fp8 LAYOUT (all weight GEMMs in fp8: MoE expert GEMMs via TEGroupedMLP, +# attention QKV/O + dense proj via TELinear; attention core QK^T/softmax*V stays +# BF16; mHC/Sinkhorn fp32; embedding/head/RMSNorm/router BF16; optimizer fp32) — +# but with a per-TENSOR scale instead of the paper's ue8m0 microscale. +# +# CK-free by construction: this 1gpu launcher already has TURBO_USE_GROUPED_MLP= +# False + USE_TURBO_DEEPEP=False, so experts route to TEGroupedMLP (hipBLASLt), +# not PrimusTurboGroupedMLP (ck_grouped_gemm). NVTE_ROCM_ENABLE_MXFP8=1 is set +# by examples/run_pretrain.sh. +# +# WHY NOT mxfp8 (paper ue8m0): two blockers on this build, both upstream- +# root-caused. (1) TE-ROCm MXFP8 asserts GEMM K % 128 == 0 (rocm_gemm.hip:1529) +# and V4 has non-128 K dims (e.g. K=224, K=32) -> errors out. (2) Even if it ran, +# mxfp8's e8m0 per-block quant noise AMPLIFIES MULTIPLICATIVELY through backward +# depth -> divergence. tensorwise's smooth per-tensor fp32 scale is stable +# (upstream: loss 12 -> 0.82 at full depth). So mxfp8 is GUARDED below. +# +# The earlier FP8 no-op (decoder skipped the fp8 context) is fixed upstream +# (commit b662c40b) and lives in the mounted repo, so FP8 now actually engages. +# A/B back to BF16 with PRECISION_TYPE=BF16 (or FP8=null). +export NVTE_ROCM_ENABLE_MXFP8=${NVTE_ROCM_ENABLE_MXFP8:-1} +# TURBO-FREE FP8: primus_turbo is only an import-shim in this gfx1250 container, +# so the turbo FP8 path (PrimusTurboQuantConfig / primus_turbo_fp8_autocast) +# can't run. Force the TE-native fp8_autocast branch (fp8_utils.py honors this). +export PRIMUS_FP8_DISABLE_TURBO=${PRIMUS_FP8_DISABLE_TURBO:-1} +export PRECISION_TYPE=${PRECISION_TYPE:-FP8} +if [ "$PRECISION_TYPE" = "FP8" ]; then + export FP8=${FP8:-e4m3} + export FP8_RECIPE=${FP8_RECIPE:-tensorwise} + # GUARD: mxfp8 (paper ue8m0) diverges for V4 on this build. Refuse it unless + # explicitly forced, matching upstream run_deepseek_v4_pro_muon.sh. + if [ "$FP8_RECIPE" = "mxfp8" ] && [ "${MXFP8_I_KNOW_ITS_BROKEN:-0}" != "1" ]; then + echo "[FATAL] FP8_RECIPE=mxfp8 diverges for V4 on this build (TE K%128 assert +" >&2 + echo " e8m0 depth-amplified instability). Use FP8_RECIPE=tensorwise," >&2 + echo " or set MXFP8_I_KNOW_ITS_BROKEN=1 to force it anyway." >&2 + exit 1 + fi +else + export FP8=${FP8:-null} + export FP8_RECIPE=${FP8_RECIPE:-null} +fi + +if [ "$TURBO_USE_GROUPED_MLP" = "True" ]; then + export PRIMUS_BIAS_SWIGLU_FUSION=True +fi + +if [ ! -d "$PRIMUS_PATH/third_party/Megatron-LM" ] || \ + [ -z "$(ls -A "$PRIMUS_PATH/third_party/Megatron-LM" 2>/dev/null)" ]; then + echo "[ERROR] third_party/Megatron-LM missing/empty -> run: git submodule update --init --recursive" >&2 + exit 1 +fi + +echo "[pro] model=$PRIMUS_MODEL layers=$PRIMUS_TOTAL_LAYERS experts=$PRIMUS_NUM_EXPERTS seq=$PRIMUS_SEQ_LENGTH optimizer=$OPTIMIZER compress_ratios=$PRIMUS_COMPRESS_RATIOS" + +# V4-Pro single-GPU overrides (trailing args -> run_pretrain.sh -> primus cli +# train pretrain --config $EXP ...). Mirrors run_deepseek_v4_pro_muon.sh's CLI +# set, minus the DeepEP wiring, scaled to one GPU / minimum layers. +# overlap_grad_reduce/param_gather stay OFF: upstream enabled them for multi- +# node DP scaling (needs the distributed optimizer + the indexer-param freeze), +# but at single-GPU DP=1 they are no-ops, and Muon requires them off anyway. +PROXY_OVERRIDES="\ + --backend_path $PRIMUS_PATH/third_party/Megatron-LM \ + --train_iters $TRAIN_ITERS \ + --lr_warmup_iters 0 \ + --lr_decay_iters $TRAIN_ITERS \ + --num_layers $PRIMUS_TOTAL_LAYERS \ + --compress_ratios $PRIMUS_COMPRESS_RATIOS \ + --micro_batch_size $MBS \ + --global_batch_size $GBS \ + --lr $LR \ + --min_lr $MIN_LR \ + --adam_eps $ADAM_EPS \ + --moe_aux_loss_coeff $MOE_AUX_LOSS_COEFF \ + --seq_length $PRIMUS_SEQ_LENGTH \ + --max_position_embeddings $PRIMUS_MAX_POSITION_EMBEDDINGS \ + --rope_type rope \ + --tensor_model_parallel_size $PRIMUS_TP \ + --pipeline_model_parallel_size $PRIMUS_PP \ + --expert_model_parallel_size $PRIMUS_EP \ + --num_experts $PRIMUS_NUM_EXPERTS \ + --moe_router_topk $PRIMUS_MOE_TOPK \ + --moe_router_enable_expert_bias $PRIMUS_MOE_ENABLE_EXPERT_BIAS \ + --moe_ffn_hidden_size $PRIMUS_MOE_FFN_HIDDEN_SIZE \ + --index_topk $PRIMUS_INDEX_TOPK \ + --v4_grouped_experts_support_clamped_swiglu True \ + --mtp_num_layers $MTP_NUM_LAYERS \ + --mock_data True \ + --moe_router_force_load_balancing True \ + --log_avg_skip_iterations 3 \ + --optimizer $OPTIMIZER \ + --muon_momentum $MUON_MOMENTUM \ + --muon_extra_scale_factor $MUON_EXTRA_SCALE_FACTOR \ + --use_distributed_optimizer $USE_DISTRIBUTED_OPTIMIZER \ + --use_precision_aware_optimizer $USE_PRECISION_AWARE_OPTIMIZER \ + $OPT_DTYPE_ARGS \ + --enable_primus_turbo $ENABLE_PRIMUS_TURBO \ + --use_turbo_attention $USE_TURBO_ATTENTION \ + --use_turbo_deepep $USE_TURBO_DEEPEP \ + --use_turbo_grouped_gemm $TURBO_USE_GROUPED_MLP \ + --use_turbo_gemm $USE_TURBO_PARALLEL_LINEAR \ + --use_v4_attention_backend $USE_V4_ATTENTION_BACKEND \ + --use_v4_csa_attention_backend $USE_V4_CSA_ATTENTION_BACKEND \ + --use_v4_compiled_sinkhorn $USE_V4_COMPILED_SINKHORN \ + --moe_use_legacy_grouped_gemm False \ + --moe_permute_fusion $MOE_PERMUTE_FUSION \ + --fp8 $FP8 \ + --fp8_recipe $FP8_RECIPE \ + --recompute_num_layers 0 \ + --recompute_granularity full \ + --recompute_method block \ + --gradient_accumulation_fusion False \ + --overlap_grad_reduce False \ + --overlap_param_gather False \ + --disable_last_saving True \ + --disable_wandb True \ + --disable_tensorboard $DISABLE_TENSORBOARD \ + --profile $PROFILE \ + --use_pytorch_profiler $PROFILE \ + --profile_step_start $PROFILE_STEP_START \ + --profile_step_end $PROFILE_STEP_END \ + --bias_swiglu_fusion $PRIMUS_BIAS_SWIGLU_FUSION \ + --torch_profiler_use_gzip True" + +ENV_ARGS=() +for v in DOCKER_IMAGE NVTE_FUSED_ATTN NVTE_FUSED_ATTN_CK NVTE_FUSED_ATTN_AOTRITON \ + PRIMUS_TURBO_GEMM_BACKEND PRIMUS_TURBO_GROUPED_GEMM_BACKEND TURBO_WHEEL_DIR FLYDSL_PKG_DIR \ + NVTE_FLASH_ATTN NVTE_USE_CK_GEMM NVTE_ROCM_ENABLE_MXFP8 PRIMUS_FP8_DISABLE_TURBO PYTHONPATH HSA_ENABLE_SDMA HSA_NO_SCRATCH_RECLAIM \ + TORCH_COMPILE_DISABLE TORCHINDUCTOR_COMPILE_THREADS TRITON_CACHE_DIR \ + HSA_SIGNAL_ABORT_TIMEOUT HSA_ENABLE_INTERRUPT \ + HIP_LAUNCH_BLOCKING AMD_SERIALIZE_KERNEL AMD_SERIALIZE_COPY \ + AMD_LOG_LEVEL AMD_LOG_MASK MASTER_PORT TORCH_NCCL_HIGH_PRIORITY \ + NCCL_IB_DISABLE NCCL_P2P_DISABLE NCCL_IB_HCA NCCL_SOCKET_IFNAME \ + GLOO_SOCKET_IFNAME RCCL_DISABLE_AMDSMI NCCL_AMDSMI_DISABLE USING_AINIC \ + GPUS_PER_NODE NNODES PYTHONUNBUFFERED TE_DIR TE_WHEEL_DIR PRIMUS_MODEL \ + PRIMUS_SEQ_LENGTH PRIMUS_MAX_POSITION_EMBEDDINGS \ + PRIMUS_TEAM PRIMUS_USER PRIMUS_EXP_NAME \ + PRIMUS_STACK_GROUPED_WEIGHT_TRITON PRIMUS_ROPE_TRITON \ + PRIMUS_SINKHORN_TRITON PRIMUS_HC_TRITON PRIMUS_INDEXER_TRITON \ + PRIMUS_INDEXER_TRITON_FULL PRIMUS_V4_ROUTER_TRITON \ + PRIMUS_TURBO_FUSE_GROUPED_WGRAD PRIMUS_TURBO_FUSE_WGRAD_DEBUG \ + PRIMUS_MUON_BATCHED_NS PRIMUS_COMPRESS_ROPE_CACHE PRIMUS_COMPRESS_POOL_TRITON; do + ENV_ARGS+=("--env" "$v") +done +[[ -n "${HIP_VISIBLE_DEVICES:-}" ]] && ENV_ARGS+=("--env" "HIP_VISIBLE_DEVICES") +# EXTRA_CLI: extra trailing --flag value overrides appended after PROXY_OVERRIDES +# (argparse last-wins), for dimension bisects etc. +[[ -n "${EXTRA_CLI:-}" ]] && ENV_ARGS+=("--env" "EXTRA_CLI") + +# Persistent Triton compile cache. The --rm container makes TRITON_CACHE_DIR +# ephemeral, so every run recompiles ALL kernels from scratch (the slow CPU-bound +# LLVM step that dominates iteration 1, esp. under this node's MCE storm). Mount a +# host dir so compiled kernels (hsaco) are reused across runs -> iter-1 of every +# later run with the same shapes skips the cold compile. Triton keys cache entries +# by kernel-source + arch + constexpr hash, so a wheel/arch/shape change auto- +# invalidates (safe to keep warm). Disable with PRIMUS_TRITON_CACHE_DIR="". +export PRIMUS_TRITON_CACHE_DIR=${PRIMUS_TRITON_CACHE_DIR:-$PRIMUS_PATH/.triton_cache_shared} +if [ -n "$PRIMUS_TRITON_CACHE_DIR" ]; then + mkdir -p "$PRIMUS_TRITON_CACHE_DIR" + export TRITON_CACHE_DIR="$PRIMUS_TRITON_CACHE_DIR" + echo "[triton] persistent compile cache: $TRITON_CACHE_DIR ($(find "$TRITON_CACHE_DIR" -maxdepth 1 -type d 2>/dev/null | wc -l) entries)" +fi + +VOLUME_ARGS=(-v "$PRIMUS_PATH":"$PRIMUS_PATH" -v "$DATA_PATH":"$DATA_PATH") +[[ -d "$TE_WHEEL_DIR" ]] && VOLUME_ARGS+=(-v "$TE_WHEEL_DIR":"$TE_WHEEL_DIR") +[[ -d "$TE_DIR" ]] && VOLUME_ARGS+=(-v "$TE_DIR":"$TE_DIR") +[[ -n "${TURBO_WHEEL_DIR:-}" && -d "$TURBO_WHEEL_DIR" ]] && VOLUME_ARGS+=(-v "$TURBO_WHEEL_DIR":"$TURBO_WHEEL_DIR") +[[ -n "${FLYDSL_PKG_DIR:-}" && -d "$FLYDSL_PKG_DIR/flydsl" ]] && VOLUME_ARGS+=(-v "$FLYDSL_PKG_DIR":"$FLYDSL_PKG_DIR") +[[ -n "${TRITON_CACHE_DIR:-}" ]] && VOLUME_ARGS+=(-v "$TRITON_CACHE_DIR":"$TRITON_CACHE_DIR") +# Opt-in tuned hipBLASLt: mount the built library at the same path and pass the +# loader env into the container (only when enabled, to keep stock runs untouched). +if [ "$PRIMUS_TUNED_HIPBLASLT" = "1" ]; then + VOLUME_ARGS+=(-v "$HBL_TUNED_RELEASE":"$HBL_TUNED_RELEASE") + ENV_ARGS+=("--env" "PRIMUS_TUNED_HIPBLASLT" "--env" "HIPBLASLT_TENSILE_LIBPATH" \ + "--env" "HBL_TUNED_RELEASE") +fi +# Container-side loader injection for the tuned lib (prepends to the image's paths). +HBL_PRELOAD_PREFIX="" +if [ "$PRIMUS_TUNED_HIPBLASLT" = "1" ]; then + HBL_PRELOAD_PREFIX="export LD_LIBRARY_PATH=\"\$HBL_TUNED_RELEASE/library:\${LD_LIBRARY_PATH:-}\" && export LD_PRELOAD=\"\$HBL_TUNED_RELEASE/library/libhipblaslt.so.1\${LD_PRELOAD:+:\$LD_PRELOAD}\" && echo \"[hipblaslt] container LD_PRELOAD=\$LD_PRELOAD\" && " +fi + +TE_INSTALL_PREFIX="\ + if ls ${TE_WHEEL_DIR}/transformer_engine-*.whl >/dev/null 2>&1; then \ + echo '[TE] installing prebuilt wheel from ${TE_WHEEL_DIR}' && \ + pip install --quiet --force-reinstall --no-deps ${TE_WHEEL_DIR}/transformer_engine-*.whl && \ + pip install --quiet einops nvdlfw-inspect onnxscript onnx pydantic importlib-metadata packaging transformers pybind11; \ + else \ + echo '[TE] WARNING: no TE wheel found at ${TE_WHEEL_DIR}; run will likely fail'; \ + fi && \ + echo '[deps] installing Primus requirements' && \ + pip install --quiet -r requirements.txt && \ + if [ -n \"${TURBO_WHEEL_DIR:-}\" ] && ls ${TURBO_WHEEL_DIR}/primus_turbo-*.whl >/dev/null 2>&1; then \ + echo '[turbo] installing real primus_turbo wheel from ${TURBO_WHEEL_DIR}' && \ + pip install --quiet --force-reinstall --no-deps ${TURBO_WHEEL_DIR}/primus_turbo-*.whl && \ + python -c 'import primus_turbo, primus_turbo.pytorch as _; print(\"[turbo] primus_turbo\", primus_turbo.__version__, \"imported OK\")'; \ + fi && " + +docker run --rm \ + "${ENV_ARGS[@]}" \ + --ipc=host --network=host \ + --device=/dev/kfd --device=/dev/dri \ + --cap-add=SYS_PTRACE --cap-add=CAP_SYS_ADMIN \ + --security-opt seccomp=unconfined --group-add video \ + --privileged \ + --name primus-v4-pro-muon-1gpu \ + "${VOLUME_ARGS[@]}" \ + "$DOCKER_IMAGE" /bin/bash -c "\ + set -e && cd $PRIMUS_PATH && \ + ${HBL_PRELOAD_PREFIX}\ + ${TE_INSTALL_PREFIX}\ + echo '==================== V4-PRO + MUON 1-GPU PROXY (gfx1250, BF16, eager, no profiler) ====================' && \ + EXP=$EXP PRIMUS_MODEL=$PRIMUS_MODEL GPUS_PER_NODE=1 NNODES=1 bash examples/run_pretrain.sh \ + ${PROXY_OVERRIDES} ${EXTRA_CLI:-}" \ + 2>&1 | tee "$LOG" diff --git a/examples/deepseek-v4/run_dsv4_projection_1gpu.sh b/examples/deepseek-v4/run_dsv4_projection_1gpu.sh new file mode 100755 index 000000000..d3e6b98d3 --- /dev/null +++ b/examples/deepseek-v4/run_dsv4_projection_1gpu.sh @@ -0,0 +1,158 @@ +#!/bin/bash +############################################################################### +# Primus PROJECTION (memory / performance) for DeepSeek-V4 on one gfx1250 GPU. +# +# Sibling of run_deepseek_v4_pro_muon_1gpu.sh. Same validated gfx1250 docker +# recipe and the same required workarounds, but instead of pretraining it runs +# the Primus projection tool (docs/projection.md): benchmark a couple of layers +# on this single GPU and analytically project memory + training performance to +# a multi-node target cluster. +# +# Usage: +# ./run_dsv4_projection_1gpu.sh # performance, pro, ->8 nodes +# MODE=memory ./run_dsv4_projection_1gpu.sh # memory projection only +# PRIMUS_MODEL=deepseek_v4_flash ./run_dsv4_projection_1gpu.sh # flash model +# TARGET_NODES=16 ./run_dsv4_projection_1gpu.sh # project to 16 nodes +# PROFILING_MODE=simulate GPU_ARCH=mi355x MODE=performance ./run_dsv4_projection_1gpu.sh # CPU-only +############################################################################### +set -eo pipefail + +# Repo root: this script lives under examples/deepseek-v4/, so resolve two levels +# up. All paths below (TE_WHEEL_DIR, third_party, VOLUME_ARGS, cd) are repo-root-relative. +SCRIPT_DIR=$(realpath -m "$(dirname "$0")/../..") +export DOCKER_IMAGE=${DOCKER_IMAGE:-registry-sc-harbor.amd.com/framework/therock-npi@sha256:feba897e2a32a2465b8b296ed2662b2ad6136b5f1cf6f6c2716a3674aafc30f3} +export TE_WHEEL_DIR=${TE_WHEEL_DIR:-$(realpath -m "$SCRIPT_DIR/../../mi450/dist/feba897")} + +# ---------- What to project ------------------------------------------------- +export MODE=${MODE:-performance} # memory | performance +export PRIMUS_MODEL=${PRIMUS_MODEL:-deepseek_v4_pro} # deepseek_v4_pro | deepseek_v4_flash +export EXP=${EXP:-examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml} +export BENCHMARK_GPUS=${BENCHMARK_GPUS:-1} # benchmark on this many GPUs (1 here) +# This box has ONE physical GPU. We invoke `primus-cli direct --single` +# (ONE python3 process, NOT torchrun) — +# the same trick the validated run_dsv3_projection script uses. In --single mode +# torchrun is not used, so GPUS_PER_NODE no longer drives --nproc_per_node; it +# serves ONLY as the TARGET node size (8 GPUs/node), while the projection spawns +# its own nproc=1 benchmark subprocess pinned to the single physical GPU. This +# gives correct intra-/inter-node comm modeling for the real 8-GPU-node cluster. +export GPUS_PER_NODE=8 # TARGET node size (8 nodes x 8 = 64 GPUs) +export TARGET_NODES=${TARGET_NODES:-8} # production TP1*PP8*EP8 = 64 GPUs = 8 nodes +export PROFILING_MODE=${PROFILING_MODE:-benchmark} # benchmark | simulate | both +export GPU_ARCH=${GPU_ARCH:-} # e.g. mi355x for --profiling-mode simulate + +# ---------- Required gfx1250 workarounds (see the pretrain launcher) --------- +export HSA_NO_SCRATCH_RECLAIM=1 +# gfx1250 MES async-queue hang workaround — matched EXACTLY to the training +# launcher (run_deepseek_v4_pro_muon_1gpu.sh): AMD_SERIALIZE_COPY=3 alone, which +# was bisected sufficient for the V4-Pro proxy (KERNEL serialize + LAUNCH_BLOCKING +# found unnecessary, default off). If the projection still hangs with this, the +# cause is the full-model build (all 61 layers on one rank), not these knobs. +export AMD_SERIALIZE_COPY=${AMD_SERIALIZE_COPY:-3} +export AMD_SERIALIZE_KERNEL=${AMD_SERIALIZE_KERNEL:-0} +export HIP_LAUNCH_BLOCKING=${HIP_LAUNCH_BLOCKING:-0} +export HSA_ENABLE_SDMA=${HSA_ENABLE_SDMA:-0} # flaky SDMA completion-signal workaround +# Turbo-free TE-native FP8 (tensorwise / Float8CurrentScaling), matching the +# training launcher. Without this the model uses TE DelayedScaling, which asserts +# against the V4 attention's save_original_input. fp8_utils.py reads this env. +export PRIMUS_FP8_DISABLE_TURBO=${PRIMUS_FP8_DISABLE_TURBO:-1} +export NVTE_ROCM_ENABLE_MXFP8=${NVTE_ROCM_ENABLE_MXFP8:-1} +# RCCL all_reduce(AVG) hangs even at world_size=1 -> sitecustomize rewrites AVG->SUM/ws. +# Also put the vendored Emerging-Optimizers on PYTHONPATH so the muon optimizer +# (emerging_optimizers.*) imports — required for OPTIMIZER=muon. +export PYTHONPATH_IN="$SCRIPT_DIR/rccl_avg_workaround:$SCRIPT_DIR/third_party/Emerging-Optimizers" + +LOG=${LOG:-dsv4-projection-${MODE}.log} + +# Config overrides appended as trailing CLI key/value pairs. V4 uses its OWN +# attention (multi_latent_attention=false + yarn via dual_rope), but Megatron's +# stock validate_args rejects rope_type=yarn unless MLA is on. The projection +# benchmarks a stock-Megatron layer (not the V4 custom attention), so force +# rope_type=rope at the Megatron-arg level — V4 applies yarn internally and +# rope-vs-yarn is a cheap elementwise op (negligible for timing). +# (moe_router_score_function: V4 uses sqrtsoftplus, but Megatron's stock +# validate requires sigmoid for expert-bias aux-loss-free routing; the score +# function is a pointwise on router logits and does not change GEMM timing.) +# (moe_token_dispatcher_type: V4 uses the turbo "flex" dispatcher, which asserts +# TPxEP>1; the single-GPU benchmark runs at EP=1. Force "alltoall" — dispatcher +# type only affects MoE *communication* (modeled analytically), not expert-GEMM +# compute, so benchmarked layer time is unchanged.) +# (enable_primus_turbo / use_turbo_deepep: a Primus patch [moe_dispatcher_patches] +# force-replaces the dispatcher with the turbo DeepEP "flex" one when BOTH are +# true, and flex asserts TPxEP>1 (can't benchmark on 1 GPU). gfx1250 runs +# turbo-free anyway (training disables it), so force them off -> standard +# alltoall dispatcher, single-GPU-benchmarkable.) +# (gradient_accumulation_fusion: needs APEX fused_weight_gradient_mlp_cuda, not +# in this container; off here exactly as in the training launcher.) +# (optimizer: use the yaml default adam — the projection only benchmarks layer +# fwd/bwd, so the optimizer choice doesn't affect timing, and adam avoids muon's +# "Emerging Optimizers" package dependency that isn't in this container. The +# trainer.py adam/muon get_*_optimizer call sites were patched to match the +# bundled Megatron signature [dropped the removed no_wd_decay_cond/scale_lr_cond/ +# lr_mult positionals that collided with use_gloo_process_groups].) +# (tokenizer NullTokenizer: the benchmark builds a mock dataset; Megatron's +# MockGPTDataset JSON-serializes the tokenizer via .unique_identifiers, which the +# DeepSeekV4 HuggingFace tokenizer lacks -> crash. NullTokenizer has it, needs no +# HF download, and preserves vocab_size (129280) so embedding/LM-head GEMM dims +# are unchanged. Layer compute is tokenizer-independent.) +# (use_v4_triton_attention/csa: enable the fused flash-style V4 attention kernels +# instead of eager attention so the benchmarked attention time is representative +# of the real training config — eager materializes [B,H,S,S] and hugely inflates +# attention at seq 4096. Verified working on gfx1250.) +# V4-specific flags mirrored from the training launcher (now that the V4 builder +# is used, the real DeepseekV4MoE/HybridLayer is built and needs these): clamped +# SwiGLU support on the grouped backend, turbo off, legacy/permute-fusion off. +EXTRA_OVERRIDES=${EXTRA_OVERRIDES:---rope_type rope --moe_router_score_function sigmoid --moe_token_dispatcher_type alltoall --enable_primus_turbo false --use_turbo_deepep false --use_turbo_grouped_gemm false --use_turbo_gemm false --use_v4_compiled_sinkhorn false --moe_use_legacy_grouped_gemm false --moe_permute_fusion false --gradient_accumulation_fusion false --mtp_num_layers 0 --tokenizer_type NullTokenizer --use_v4_triton_attention true --use_v4_triton_csa_attention true} + +# ---------- Build the projection CLI args ----------------------------------- +PROJ_ARGS="projection $MODE --config $EXP" +if [ "$MODE" = "performance" ]; then + PROJ_ARGS="$PROJ_ARGS --benchmark-gpus $BENCHMARK_GPUS --target-nodes $TARGET_NODES --profiling-mode $PROFILING_MODE" + [ -n "$GPU_ARCH" ] && PROJ_ARGS="$PROJ_ARGS --gpu-arch $GPU_ARCH" +fi +PROJ_ARGS="$PROJ_ARGS $EXTRA_OVERRIDES" + +# TE wheel install prefix (same as pretrain launcher) +TE_INSTALL="true" +if ls "${TE_WHEEL_DIR}"/transformer_engine-*.whl >/dev/null 2>&1; then + TE_INSTALL="pip install --quiet --force-reinstall --no-deps ${TE_WHEEL_DIR}/transformer_engine-*.whl && \ + pip install --quiet einops nvdlfw-inspect onnxscript onnx pydantic importlib-metadata packaging transformers pybind11" +fi +# simulate mode (CPU-only, no model instantiation) needs the Origami GEMM model. +ORIGAMI_INSTALL="true" +if [ "$PROFILING_MODE" = "simulate" ] || [ "$PROFILING_MODE" = "both" ]; then + ORIGAMI_INSTALL="pip install --quiet 'git+https://github.com/ROCm/rocm-libraries.git#subdirectory=shared/origami/python' || echo '[warn] origami install failed'" +fi + +VOLUME_ARGS=(-v "$SCRIPT_DIR":"$SCRIPT_DIR") +[[ -d "$TE_WHEEL_DIR" ]] && VOLUME_ARGS+=(-v "$TE_WHEEL_DIR":"$TE_WHEEL_DIR") + +echo "[projection] mode=$MODE model=$PRIMUS_MODEL target_nodes=$TARGET_NODES profiling_mode=$PROFILING_MODE config=$EXP" + +# Same docker invocation as run_deepseek_v4_pro_muon_1gpu.sh (validated gfx1250 recipe). +docker run --rm \ + --ipc=host --network=host \ + --device=/dev/kfd --device=/dev/dri \ + --cap-add=SYS_PTRACE --cap-add=CAP_SYS_ADMIN \ + --security-opt seccomp=unconfined --group-add video \ + --privileged \ + --name primus-v4-projection \ + -e NNODES=1 -e GPUS_PER_NODE="$GPUS_PER_NODE" \ + -e MASTER_ADDR=localhost -e MASTER_PORT=1234 \ + -e GLOO_SOCKET_IFNAME=lo -e NCCL_SOCKET_IFNAME=lo \ + -e NCCL_IB_DISABLE=1 -e NCCL_P2P_DISABLE=1 \ + -e HSA_NO_SCRATCH_RECLAIM="$HSA_NO_SCRATCH_RECLAIM" \ + -e AMD_SERIALIZE_COPY="$AMD_SERIALIZE_COPY" -e AMD_SERIALIZE_KERNEL="$AMD_SERIALIZE_KERNEL" \ + -e HIP_LAUNCH_BLOCKING="$HIP_LAUNCH_BLOCKING" -e HSA_ENABLE_SDMA="$HSA_ENABLE_SDMA" \ + -e PRIMUS_FP8_DISABLE_TURBO="$PRIMUS_FP8_DISABLE_TURBO" -e NVTE_ROCM_ENABLE_MXFP8="$NVTE_ROCM_ENABLE_MXFP8" \ + -e PRIMUS_PROJ_MAX_LAYERS="${PRIMUS_PROJ_MAX_LAYERS:-1}" -e PRIMUS_PROJ_COMPRESS_RATIOS="${PRIMUS_PROJ_COMPRESS_RATIOS:-}" \ + -e PRIMUS_MODEL="$PRIMUS_MODEL" -e PYTHONUNBUFFERED=1 \ + "${VOLUME_ARGS[@]}" \ + "$DOCKER_IMAGE" /bin/bash -c "\ + set -e && cd $SCRIPT_DIR && \ + export PYTHONPATH=$PYTHONPATH_IN:\${PYTHONPATH:-} && \ + ${TE_INSTALL} && \ + ${ORIGAMI_INSTALL} && \ + pip install --quiet -r requirements.txt && \ + echo '==================== DSV4 PROJECTION ($MODE) ====================' && \ + bash runner/primus-cli direct --single -- $PROJ_ARGS" \ + 2>&1 | tee "$LOG" diff --git a/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml new file mode 100644 index 000000000..277c129f7 --- /dev/null +++ b/examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml @@ -0,0 +1,138 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:deepseek_v4_flash-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +# DeepSeek-V4 Flash BF16 pretraining config tuned for MI355X. +# This is a smoke / scaffold config — values (especially parallelism) will be +# revised once the V4 builder + hybrid attention land in Phase 3-5. + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + model: ${PRIMUS_MODEL:deepseek_v4_flash}.yaml + overrides: + # log + wandb_project: "Primus_DeepSeek_V4_Pretrain" + stderr_sink_level: DEBUG + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + + # hyper parameters + train_iters: 50 + micro_batch_size: 1 + global_batch_size: 256 + seq_length: ${PRIMUS_SEQ_LENGTH:4096} + max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:4096} + lr: 1.0e-5 + min_lr: 0.0 + lr_warmup_iters: 2 + lr_decay_iters: null + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: true + init_method_std: 0.008 + norm_epsilon: 1.0e-6 + + # parallel — scaffold defaults; will be retuned in Phase 6. + tensor_model_parallel_size: ${PRIMUS_TP:1} + pipeline_model_parallel_size: ${PRIMUS_PP:8} + expert_model_parallel_size: ${PRIMUS_EP:8} + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: true + + # data + mock_data: true + train_data_path: ${PRIMUS_TOKENIZED_DATA_PATH:null} + valid_data_path: null + test_data_path: null + + # ---------- DeepSeek-V4 specific ---------- + hybrid_attention_enabled: true + attn_sink: true + hc_use_sinkhorn: true + mtp_use_separate_hc_head: true + moe_router_score_function: sqrtsoftplus + swiglu_limit: 10.0 + + # ---------- Optimizer ---------- + # NOTE: Phase 7 will add Muon for the latent / hidden projections. + # Until then we run plain BF16 AdamW (precision-aware). + use_precision_aware_optimizer: true + main_grads_dtype: bf16 + exp_avg_dtype: bf16 + exp_avg_sq_dtype: bf16 + + # rope fusion + enable_experimental: true + apply_rope_fusion: false # V4 uses partial RoPE on a 512-dim head + + # recompute + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 + + # ckpt + finetune: false + auto_continue_train: false + load: null + no_load_optim: null + no_load_rng: null + save: null + save_interval: 20000 + no_save_optim: null + no_save_rng: null + disable_last_saving: true + ckpt_format: torch + eval_iters: 0 + + # turbo / deepep — keep off by default; enable via env vars when + # benchmarking the Turbo path (see run_deepseek_v4.sh). Plan-3 P22 + # routes the dense (compress_ratio == 0) attention layers through + # PrimusTurboAttention when ``use_turbo_attention=true``; the V4 + # builder auto-derives ``use_sink_attention`` / + # ``sink_sliding_window`` from ``attn_sink`` / ``attn_sliding_window`` + # so the YAML stays free of Turbo-internal knobs. + enable_primus_turbo: ${PRIMUS_ENABLE_TURBO:true} + # use_turbo_attention stays OFF: the dense (cr=0) path must run the + # triton_v2 sparse-MLA backend below, not PrimusTurboAttention (which + # would take dispatch precedence). + use_turbo_attention: ${PRIMUS_USE_TURBO_ATTENTION:false} + use_turbo_grouped_gemm: true + use_turbo_rms_norm: true + use_turbo_deepep: ${PRIMUS_USE_TURBO_DEEPEP:true} + moe_shared_expert_overlap: false + moe_router_dtype: fp32 + + # deepep tuning (64 or 80 for ep8, 32 for ep16-64 is best practice) + turbo_deepep_num_cu: 80 + turbo_deepep_use_comm_stream: false + + # sync-free moe support (stage 1-2; 0 = off). stage 2 = best perf. + turbo_sync_free_moe_stage: 1 + + # V4 attention backend selection (unified string selectors; default triton_v2). + # use_v4_attention_backend (dense cr=0 / HCA cr=128): eager|triton_v1|triton_v2|gluon + # use_v4_csa_attention_backend (CSA cr=4): eager|triton_v0|triton_v1|triton_v2|gluon|flydsl_v0 + use_v4_attention_backend: ${PRIMUS_USE_V4_ATTENTION_BACKEND:triton_v2} + use_v4_csa_attention_backend: ${PRIMUS_USE_V4_CSA_ATTENTION_BACKEND:triton_v2} + + # FP8 (E4M3) Indexer QK path (CSA selector); BF16 index-score preserved. + use_v4_fp8_indexer: ${PRIMUS_USE_V4_FP8_INDEXER:false} + + # plan-5 P29 (RESCOPED): wrap sinkhorn_normalize with a cached + # torch.compile(fullgraph=True, dynamic=False) build. Default + # false until G32 (parity) + G33b (trace) flip it on. + use_v4_compiled_sinkhorn: ${PRIMUS_USE_V4_COMPILED_SINKHORN:false} + + moe_use_fused_router_with_aux_score: true + moe_permute_fusion: true + + # Cross entropy flags + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true diff --git a/examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml b/examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml new file mode 100644 index 000000000..bc5026810 --- /dev/null +++ b/examples/megatron/configs/MI355X/deepseek_v4_flash-FP8-pretrain.yaml @@ -0,0 +1,166 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:deepseek_v4_flash-fp8-pretrain} +workspace: ${PRIMUS_WORKSPACE:./output} + +# DeepSeek-V4 Flash FP8 pretraining config tuned for MI355X. +# Derived from deepseek_v4_flash-BF16-pretrain.yaml; the ONLY substantive +# delta is the FP8 training block at the bottom. +# +# Precision recipe (paper §4.x quantization / techblog §9.6): +# Paper = "FP4 + FP8 Mixed": MoE experts + CSA-Indexer QK in FP4 (MXFP4), +# EVERYTHING ELSE in FP8, all with the **ue8m0** microscaling scale format. +# ue8m0 = OCP Microscaling (MX) block scale → on this stack that is +# `fp8_recipe: mxfp8` → TE MXFP8BlockScaling → Primus-Turbo MX_BLOCKWISE +# with scale_dtype=E8M0 (fp8_utils.py:148), native on MI355X/CDNA4. +# +# Integration gap vs the paper (NOT config — Primus V4 TODO, techblog item 10 +# "Phase 2 FP4/FP8 Mixed"): the FP4 expert / FP4-Indexer path is not yet wired +# in V4, so experts run at FP8 here (FP8 everywhere) rather than FP4. This is +# the closest supported step toward the paper recipe. FP8 is highly outlier- +# sensitive, which is why the paper pairs it with clamped SwiGLU (swiglu_limit, +# set below) — keep that on whenever FP8 is enabled. + +modules: + pre_trainer: + framework: megatron + config: pre_trainer.yaml + + model: ${PRIMUS_MODEL:deepseek_v4_flash}.yaml + overrides: + # log + wandb_project: "Primus_DeepSeek_V4_Pretrain" + stderr_sink_level: DEBUG + log_avg_skip_iterations: 2 + log_avg_reset_interval: 50 + + # hyper parameters + train_iters: 50 + micro_batch_size: 1 + global_batch_size: 256 + seq_length: ${PRIMUS_SEQ_LENGTH:4096} + max_position_embeddings: ${PRIMUS_MAX_POSITION_EMBEDDINGS:4096} + lr: 1.0e-5 + min_lr: 0.0 + lr_warmup_iters: 2 + lr_decay_iters: null + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + eod_mask_loss: true + init_method_std: 0.008 + norm_epsilon: 1.0e-6 + + # parallel — scaffold defaults; will be retuned in Phase 6. + tensor_model_parallel_size: ${PRIMUS_TP:1} + pipeline_model_parallel_size: ${PRIMUS_PP:8} + expert_model_parallel_size: ${PRIMUS_EP:8} + overlap_grad_reduce: true + overlap_param_gather: true + gradient_accumulation_fusion: true + + # data + mock_data: true + train_data_path: ${PRIMUS_TOKENIZED_DATA_PATH:null} + valid_data_path: null + test_data_path: null + + # ---------- DeepSeek-V4 specific ---------- + hybrid_attention_enabled: true + attn_sink: true + hc_use_sinkhorn: true + mtp_use_separate_hc_head: true + moe_router_score_function: sqrtsoftplus + swiglu_limit: 10.0 + + # ---------- Optimizer ---------- + # NOTE: Phase 7 will add Muon for the latent / hidden projections. + # Until then we run plain BF16 AdamW (precision-aware). + use_precision_aware_optimizer: true + main_grads_dtype: bf16 + exp_avg_dtype: bf16 + exp_avg_sq_dtype: bf16 + + # rope fusion + enable_experimental: true + apply_rope_fusion: false # V4 uses partial RoPE on a 512-dim head + + # recompute + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 + + # ckpt + finetune: false + auto_continue_train: false + load: null + no_load_optim: null + no_load_rng: null + save: null + save_interval: 20000 + no_save_optim: null + no_save_rng: null + disable_last_saving: true + ckpt_format: torch + eval_iters: 0 + + # turbo / deepep — keep off by default; enable via env vars when + # benchmarking the Turbo path (see run_deepseek_v4.sh). Plan-3 P22 + # routes the dense (compress_ratio == 0) attention layers through + # PrimusTurboAttention when ``use_turbo_attention=true``; the V4 + # builder auto-derives ``use_sink_attention`` / + # ``sink_sliding_window`` from ``attn_sink`` / ``attn_sliding_window`` + # so the YAML stays free of Turbo-internal knobs. + enable_primus_turbo: ${PRIMUS_ENABLE_TURBO:true} + # use_turbo_attention stays OFF: the dense (cr=0) path must run the + # triton_v2 sparse-MLA backend below, not PrimusTurboAttention (which + # would take dispatch precedence). + use_turbo_attention: ${PRIMUS_USE_TURBO_ATTENTION:false} + use_turbo_grouped_gemm: true + use_turbo_rms_norm: true + use_turbo_deepep: ${PRIMUS_USE_TURBO_DEEPEP:true} + moe_shared_expert_overlap: false + moe_router_dtype: fp32 + + # deepep tuning (64 or 80 for ep8, 32 for ep16-64 is best practice) + turbo_deepep_num_cu: 80 + turbo_deepep_use_comm_stream: false + + # sync-free moe support (stage 1-2; 0 = off). stage 2 = best perf. + turbo_sync_free_moe_stage: 1 + + # V4 attention backend selection (unified string selectors; default triton_v2). + # use_v4_attention_backend (dense cr=0 / HCA cr=128): eager|triton_v1|triton_v2|gluon + # use_v4_csa_attention_backend (CSA cr=4): eager|triton_v0|triton_v1|triton_v2|gluon|flydsl_v0 + use_v4_attention_backend: ${PRIMUS_USE_V4_ATTENTION_BACKEND:triton_v2} + use_v4_csa_attention_backend: ${PRIMUS_USE_V4_CSA_ATTENTION_BACKEND:triton_v2} + + # plan-5 P29 (RESCOPED): wrap sinkhorn_normalize with a cached + # torch.compile(fullgraph=True, dynamic=False) build. Default + # false until G32 (parity) + G33b (trace) flip it on. + use_v4_compiled_sinkhorn: ${PRIMUS_USE_V4_COMPILED_SINKHORN:false} + + moe_use_fused_router_with_aux_score: true + moe_permute_fusion: true + + # Cross entropy flags + cross_entropy_fusion_impl: "te" + cross_entropy_loss_fusion: true + + # ---------- FP8 training ---------- + # fp8 = number format (E4M3); fp8_recipe = scaling strategy. + # The paper uses ue8m0 microscaling (= mxfp8, 1x32 block / E8M0 scale), but + # mxfp8 is NOT runnable on this gfx950 build: turbo grouped-GEMM has no MX + # path, and TE's ROCm MXFP8 GEMM requires K%128==0 (V4 has a K=224 proj). + # So we use `tensorwise` (per-tensor scale) — the working recipe. This gives + # the paper's fp8 *layout* (all weight GEMMs fp8) with a non-ue8m0 scale. + # Override via FP8 / FP8_RECIPE env knobs (CLI wins); FP8=null => BF16. + fp8: e4m3 + fp8_recipe: tensorwise + # Route the attention/dense linear projections through PrimusTurboLinear so + # they quantize to fp8 too (not just the MoE experts) — matches the paper's + # "fp8 on all weight GEMMs". Without this the projections stay bf16. + use_turbo_gemm: true + + moe_router_padding_for_quantization: true diff --git a/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml b/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml index 7a929a43e..474ccfbea 100644 --- a/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml +++ b/examples/mlperf/llama2_70b/configs/MI355X/llama2_70b_lora_mlperf_posttrain.yaml @@ -55,6 +55,6 @@ modules: enable_primus_turbo: false use_turbo_attention: false use_turbo_rms_norm: false - use_turbo_parallel_linear: false + use_turbo_gemm: false check_for_nan_in_loss: false diff --git a/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml b/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml index cdc2efe8e..5679cac48 100644 --- a/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml +++ b/examples/mlperf/llama3.1_8b/configs/MI355X/llama3.1_8B-pretrain-FP4.yaml @@ -107,7 +107,7 @@ modules: # --- Primus Turbo Config --- enable_primus_turbo: false use_turbo_attention: false - use_turbo_parallel_linear: false # can't use together with delayed recipe - use_turbo_grouped_mlp: false + use_turbo_gemm: false # can't use together with delayed recipe + use_turbo_grouped_gemm: false moe_use_fused_router_with_aux_score: false enable_turbo_attention_float8 : false diff --git a/examples/run_pretrain.sh b/examples/run_pretrain.sh index 03af31528..d16d6baaa 100755 --- a/examples/run_pretrain.sh +++ b/examples/run_pretrain.sh @@ -203,7 +203,16 @@ LOG_INFO "GLOO_SOCKET_IFNAME: $GLOO_SOCKET_IFNAME" LOG_INFO "" # ----------------- AMD-specific GPU optimizations ----------------- -export HSA_ENABLE_SDMA=1 + +# Enable system DMA engine (SDMA) on AMD GPUs for better IO throughput. +# ${:-} guarded so a caller can set HSA_ENABLE_SDMA=0 to route copies through +# blit/compute kernels instead — workaround for hosts where an SDMA H2D copy +# intermittently never signals completion (hangs in BusyWaitSignal). +export HSA_ENABLE_SDMA=${HSA_ENABLE_SDMA:-1} + +# Prevent scratch memory from being reclaimed to stabilize large memory usage patterns (e.g., KV cache, MoE experts) +# NOTE: Must disable scratch reclaim to avoid MoE training crash on AMD GPUs +# Setting this to 0 prevents core dumps when using Mixture-of-Experts (MoE) models export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-0} export RCCL_MSCCL_ENABLE=0 export RCCL_MSCCLPP_ENABLE=0 diff --git a/primus/backends/megatron/core/extensions/_triton/__init__.py b/primus/backends/megatron/core/extensions/_triton/__init__.py new file mode 100644 index 000000000..aa31d698f --- /dev/null +++ b/primus/backends/megatron/core/extensions/_triton/__init__.py @@ -0,0 +1,14 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton kernels for the Primus Megatron extensions package. + +Currently contains: + +* :mod:`stack_grouped_weight` — plan-6 P34's fused + ``torch.stack + transpose(1, 2) + contiguous`` for the per-expert weight + tensors of :class:`PrimusTurboGroupedMLP`. +""" diff --git a/primus/backends/megatron/core/extensions/_triton/multi_tensor_add.py b/primus/backends/megatron/core/extensions/_triton/multi_tensor_add.py new file mode 100644 index 000000000..2d6224962 --- /dev/null +++ b/primus/backends/megatron/core/extensions/_triton/multi_tensor_add.py @@ -0,0 +1,207 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-7 P45 — multi-tensor BF16 add Triton kernel (prototype). + +Replaces the ~743 separate ``vec_elem`` launches that +dominate the P40 trace's optimizer-step residual (170.99 ms / iter, +32.7 % of step) with a single Triton kernel that processes a +multi-tensor batch in one launch. + +This is a **prototype** designed to demonstrate the perf headroom +of multi-tensor fusion against ``torch._foreach_add_`` (the +PyTorch reference path) and the upstream Apex / TE +``multi_tensor_apply`` (the V4-Flash production path). Production +integration with the Apex / TE optimizer call sites is plan-8 +scope — replacing those call sites bit-exactly requires +coordinating with the master-param remainder accumulation logic +that Apex / TE owns. + +Gating: import-time only; no env knob (this module is currently +a microbench / unit-test target only). +""" + +from __future__ import annotations + +from typing import List, Tuple + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _multi_tensor_add_per_tensor_kernel( + OUT_PTR, + A_PTR, + B_PTR, + N, + SCALE, + BLOCK_SIZE: tl.constexpr, +): + """Single-tensor variant. Each grid program handles ``BLOCK_SIZE`` + elements of one tensor. + + Caller is responsible for issuing the per-tensor launch grid + and serialising the loops, but the launches are CONCURRENT on + the GPU (Triton stream stacking). + """ + + pid = tl.program_id(0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offs < N + a = tl.load(A_PTR + offs, mask=mask, other=0.0).to(tl.float32) + b = tl.load(B_PTR + offs, mask=mask, other=0.0).to(tl.float32) + out = a + SCALE * b + tl.store(OUT_PTR + offs, out, mask=mask) + + +def multi_tensor_add_triton_per_tensor( + out_list: List[torch.Tensor], + a_list: List[torch.Tensor], + b_list: List[torch.Tensor], + scale: float = 1.0, + block_size: int = 8192, +) -> None: + """Compute ``out_i = a_i + scale * b_i`` for each ``i`` in place + (per-tensor variant). + + All tensors must be contiguous + on CUDA / HIP, same dtype, same + shape per ``(out_i, a_i, b_i)`` triple. + """ + assert len(out_list) == len(a_list) == len(b_list) + for out, a, b in zip(out_list, a_list, b_list): + assert out.shape == a.shape == b.shape, "shape mismatch" + n = out.numel() + grid = (triton.cdiv(n, block_size),) + _multi_tensor_add_per_tensor_kernel[grid]( + out, + a, + b, + n, + scale, + BLOCK_SIZE=block_size, + ) + + +@triton.jit +def _multi_tensor_add_packed_kernel( + OUT_PTR_BUF, # int64 [N_TENSORS] pointer array + A_PTR_BUF, # int64 [N_TENSORS] + B_PTR_BUF, # int64 [N_TENSORS] + SIZE_BUF, # int32 [N_TENSORS] -- element counts + PID_TO_TID_BUF, # int32 [N_PROGRAMS] -- program-id -> tensor-id + PID_TO_LOCAL_BUF, # int32 [N_PROGRAMS] -- program-id -> local block index + SCALE, + BLOCK_SIZE: tl.constexpr, +): + """Single-kernel multi-tensor add (CPU-side dispatch table variant). + + One program per ``BLOCK_SIZE`` chunk of exactly one tensor. + The CPU side builds a sorted dispatch table mapping each + ``program_id`` to ``(tensor_idx, local_block_idx)`` so a single + grid launch absorbs all N_TENSORS tensors without per-element + tensor lookups. + + This avoids the cross-tensor-boundary aliasing that a naive + "concatenated range" partition would create. + """ + + pid = tl.program_id(0) + tensor_idx = tl.load(PID_TO_TID_BUF + pid) + local_block = tl.load(PID_TO_LOCAL_BUF + pid) + + offs = local_block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + size_i = tl.load(SIZE_BUF + tensor_idx) + mask = offs < size_i + + out_ptr = tl.load(OUT_PTR_BUF + tensor_idx).to(tl.pointer_type(tl.bfloat16)) + a_ptr = tl.load(A_PTR_BUF + tensor_idx).to(tl.pointer_type(tl.bfloat16)) + b_ptr = tl.load(B_PTR_BUF + tensor_idx).to(tl.pointer_type(tl.bfloat16)) + + a = tl.load(a_ptr + offs, mask=mask, other=0.0).to(tl.float32) + b = tl.load(b_ptr + offs, mask=mask, other=0.0).to(tl.float32) + out = a + SCALE * b + tl.store(out_ptr + offs, out.to(tl.bfloat16), mask=mask) + + +def _build_multi_tensor_dispatch_table( + out_list: List[torch.Tensor], + a_list: List[torch.Tensor], + b_list: List[torch.Tensor], + block_size: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the GPU-side dispatch tables for the packed kernel. + + Returns + ------- + out_ptrs, a_ptrs, b_ptrs : int64 [N_TENSORS] + sizes : int32 [N_TENSORS] + pid_to_tid : int32 [N_PROGRAMS] -- block-id -> tensor-id + pid_to_local : int32 [N_PROGRAMS] -- block-id -> local-block-index + within the tensor it belongs to. + """ + len(out_list) + device = out_list[0].device + out_ptrs = torch.tensor([t.data_ptr() for t in out_list], dtype=torch.int64, device=device) + a_ptrs = torch.tensor([t.data_ptr() for t in a_list], dtype=torch.int64, device=device) + b_ptrs = torch.tensor([t.data_ptr() for t in b_list], dtype=torch.int64, device=device) + sizes = torch.tensor([t.numel() for t in out_list], dtype=torch.int32, device=device) + # Build pid -> (tid, local) on the host then copy. + pid_to_tid: List[int] = [] + pid_to_local: List[int] = [] + for t_idx, t in enumerate(out_list): + n_chunks = (t.numel() + block_size - 1) // block_size + for local in range(n_chunks): + pid_to_tid.append(t_idx) + pid_to_local.append(local) + pid_to_tid_t = torch.tensor(pid_to_tid, dtype=torch.int32, device=device) + pid_to_local_t = torch.tensor(pid_to_local, dtype=torch.int32, device=device) + return out_ptrs, a_ptrs, b_ptrs, sizes, pid_to_tid_t, pid_to_local_t + + +def multi_tensor_add_triton_packed( + out_list: List[torch.Tensor], + a_list: List[torch.Tensor], + b_list: List[torch.Tensor], + scale: float = 1.0, + block_size: int = 8192, +) -> None: + """Compute ``out_i = a_i + scale * b_i`` for each ``i`` in place + (packed multi-tensor variant; ONE kernel launch total). + + All tensors must be bf16, contiguous, on the same CUDA / HIP + device, with matching ``(out_i, a_i, b_i)`` shapes. + """ + assert len(out_list) == len(a_list) == len(b_list) > 0 + for out, a, b in zip(out_list, a_list, b_list): + assert out.dtype == torch.bfloat16 + assert a.dtype == torch.bfloat16 + assert b.dtype == torch.bfloat16 + assert out.is_contiguous() and a.is_contiguous() and b.is_contiguous() + assert out.shape == a.shape == b.shape + + out_ptrs, a_ptrs, b_ptrs, sizes, pid_to_tid, pid_to_local = _build_multi_tensor_dispatch_table( + out_list, a_list, b_list, block_size + ) + + grid = (pid_to_tid.numel(),) + _multi_tensor_add_packed_kernel[grid]( + out_ptrs, + a_ptrs, + b_ptrs, + sizes, + pid_to_tid, + pid_to_local, + scale, + BLOCK_SIZE=block_size, + ) + + +__all__ = [ + "multi_tensor_add_triton_per_tensor", + "multi_tensor_add_triton_packed", +] diff --git a/primus/backends/megatron/core/extensions/_triton/stack_grouped_weight.py b/primus/backends/megatron/core/extensions/_triton/stack_grouped_weight.py new file mode 100644 index 000000000..53bdea962 --- /dev/null +++ b/primus/backends/megatron/core/extensions/_triton/stack_grouped_weight.py @@ -0,0 +1,413 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton-fused ``torch.stack + transpose(1, 2) + contiguous`` for per-expert +GroupedMLP weights (plan-6 P34). + +The eager implementation in :class:`PrimusTurboGroupedMLP._stack_grouped_linear_weight`: + +.. code-block:: python + + weights = [getattr(module, f"weight{i}") for i in range(E)] + return torch.stack(weights, dim=0).transpose(1, 2).contiguous() + +does **two** full passes over the per-expert weight data: + +* ``torch.stack`` allocates ``[E, K, N]`` and issues ``E`` per-expert + ``copy_`` calls (one full pass aggregate); +* ``.contiguous()`` after ``.transpose(1, 2)`` allocates a second + ``[E, N, K]`` buffer and writes the transposed copy (a second full + pass). + +At V4-Flash EP=8 widths (``E=32``, fc1: ``K=4096, N=4096``; fc2: +``K=4096, N=2048``, bf16) the P32 final trace attributes +``hipMemcpyWithStream`` **289.6 ms / 32 calls** to this op chain +(``2 stack ops per layer × 8 layers × 2 (FWD + BWD VJP) = 32``). At ~9 ms +per call writing 512 MiB the effective bandwidth is only ~57 GB/s — far +below the MI355X HBM peak — because each call serializes E small +allocations plus a separate transpose copy. + +This module collapses the two passes into one Triton kernel that: + +1. Indexes per-expert weight tensors via an ``int64`` pointer tensor + (``weight_ptrs[e] = weights[e].data_ptr()``) — Triton's + ``tl.load(...).to(tl.pointer_type(...))`` idiom (same pattern as the + upstream grouped-GEMM tutorial). +2. Per program processes a ``[BLOCK_K, BLOCK_N]`` tile of one expert, + doing a tile-level transpose: reads ``weight[e][k, n]`` (row-major + ``[K, N]``) and writes ``out[e][n, k]`` (row-major ``[E, N, K]``). +3. BWD is the inverse — reads a ``[BLOCK_N, BLOCK_K]`` tile of ``dout + [e, n, k]`` and writes ``dweight[e][k, n]``. + +The two layouts form a bijection so **no atomics are needed in either +direction**. + +Nomenclature note (matches plan-6 P34 design doc): + +* ``K`` = ``N_out`` (``weight.shape[0]`` for ``nn.Linear`` — output features) +* ``N`` = ``N_in`` (``weight.shape[1]`` for ``nn.Linear`` — input features) + +The eager output is ``[E, N, K]`` (after ``.transpose(1, 2)`` on the +stacked ``[E, K, N]``); the Triton path produces the same layout. + +Gating: routed through :class:`PrimusTurboGroupedMLP._stack_grouped_linear_weight` +when ``PRIMUS_STACK_GROUPED_WEIGHT_TRITON != "0"`` (default-on). Set to +``"0"`` to fall back to the eager ``torch.stack + transpose + contiguous`` +chain (kept in tree for A/B testing and as the reference path for the +G37 unit tests). +""" + +from __future__ import annotations + +import os +from typing import List, Tuple + +import torch +import triton +import triton.language as tl + +# --------------------------------------------------------------------------- +# Triton dtype mapping +# --------------------------------------------------------------------------- + +_TORCH_TO_TL_DTYPE = { + torch.float64: tl.float64, + torch.float32: tl.float32, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, +} + + +def _triton_dtype(t: torch.dtype): + try: + return _TORCH_TO_TL_DTYPE[t] + except KeyError as exc: + raise TypeError( + f"stack_grouped_weight: unsupported dtype {t}; " f"expected one of {list(_TORCH_TO_TL_DTYPE)}" + ) from exc + + +# --------------------------------------------------------------------------- +# Triton kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _stack_grouped_weight_fwd_kernel( + WEIGHT_PTRS, # [E] int64 — tl.load() yields each expert's data_ptr() + OUT, # [E, N, K] contiguous output, row-major (strides [N*K, K, 1]) + E, + K, + N, + BLOCK_K: tl.constexpr, + BLOCK_N: tl.constexpr, + DTYPE: tl.constexpr, +): + """Per-expert ``[K, N] -> [N, K]`` tile-transpose, fused across experts. + + Each program writes one ``[BLOCK_K, BLOCK_N]`` tile of one expert's + output. Grid: ``(E, ceil(K / BLOCK_K), ceil(N / BLOCK_N))``. + + Read : ``weight[expert][k, n]`` from per-expert pointer, stride ``[N, 1]``. + Write : ``out[expert][n, k]``, stride ``[E*N*K -> N*K, K, 1]``. + + Both load and store carry bounds masks so non-multiple-of-BLOCK shapes + are supported. + """ + pid_e = tl.program_id(0) + pid_k = tl.program_id(1) + pid_n = tl.program_id(2) + + src_ptr = tl.load(WEIGHT_PTRS + pid_e).to(tl.pointer_type(DTYPE)) + + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask_k = offs_k < K + mask_n = offs_n < N + + src_offsets = offs_k[:, None] * N + offs_n[None, :] + tile = tl.load( + src_ptr + src_offsets, + mask=mask_k[:, None] & mask_n[None, :], + other=0, + ) + + dst_offsets = pid_e * (N * K) + offs_n[None, :] * K + offs_k[:, None] + tl.store( + OUT + dst_offsets, + tile, + mask=mask_k[:, None] & mask_n[None, :], + ) + + +@triton.jit +def _stack_grouped_weight_bwd_kernel( + DWEIGHT_PTRS, # [E] int64 — tl.load() yields each expert's grad data_ptr() + DOUT, # [E, N, K] grad tensor (contiguous, same layout as FWD OUT) + E, + K, + N, + BLOCK_K: tl.constexpr, + BLOCK_N: tl.constexpr, + DTYPE: tl.constexpr, +): + """Inverse of the FWD: reads ``dout[expert][n, k]`` and writes + ``dweight[expert][k, n]``. + + Grid mirrors the FWD; the kernel is a pure bijection memcpy so no + atomics needed, and the BLOCK tile is the same shape. + """ + pid_e = tl.program_id(0) + pid_k = tl.program_id(1) + pid_n = tl.program_id(2) + + dst_ptr = tl.load(DWEIGHT_PTRS + pid_e).to(tl.pointer_type(DTYPE)) + + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask_k = offs_k < K + mask_n = offs_n < N + + src_offsets = pid_e * (N * K) + offs_n[None, :] * K + offs_k[:, None] + tile = tl.load( + DOUT + src_offsets, + mask=mask_k[:, None] & mask_n[None, :], + other=0, + ) + + dst_offsets = offs_k[:, None] * N + offs_n[None, :] + tl.store( + dst_ptr + dst_offsets, + tile, + mask=mask_k[:, None] & mask_n[None, :], + ) + + +# --------------------------------------------------------------------------- +# Block-size autotune +# +# The kernel is bandwidth-bound; only a small handful of block sizes are +# worth scanning. ``BLOCK_K = BLOCK_N = 64`` is the default — it fits +# comfortably in LDS at bf16 (64 * 64 * 2 = 8 KiB / program) and matches +# the typical tile size MI355X HBM controllers like for transpose copies. +# Larger blocks ((128, 64), (64, 128), (128, 128)) widen the per-program +# tile and reduce launch count when E is small; smaller (32, 32) is +# kept as a safety floor for the fast-tier shapes. +# --------------------------------------------------------------------------- + +_BLOCK_CANDIDATES: Tuple[Tuple[int, int], ...] = ( + (32, 32), + (64, 64), + (128, 64), + (64, 128), +) + + +def _pick_block(K: int, N: int) -> Tuple[int, int]: + """Pick a ``(BLOCK_K, BLOCK_N)`` tile that divides reasonably into K, N. + + Conservative heuristic — we do not run a full autotune sweep at module + import to keep cold start cheap. The selected tile is always one of + :data:`_BLOCK_CANDIDATES`, picked as the largest block that does not + leave more than half the program's elements masked off on the small + dimension (i.e. the last program's tile-fill is at least 50 %). + """ + best: Tuple[int, int] = (64, 64) + best_score = -1.0 + for bk, bn in _BLOCK_CANDIDATES: + # Fraction of useful work in the last program along each axis; + # exact-multiple tiles get 1.0; partial tiles get the fraction. + k_full = (K // bk) * bk + k_tail = K - k_full + k_fill = 1.0 if k_tail == 0 else max(k_tail / bk, 0.5) + n_full = (N // bn) * bn + n_tail = N - n_full + n_fill = 1.0 if n_tail == 0 else max(n_tail / bn, 0.5) + # Score = tile area × tile-fill product → prefer bigger tiles when + # fill is high. + score = (bk * bn) * k_fill * n_fill + if score > best_score: + best_score = score + best = (bk, bn) + return best + + +# --------------------------------------------------------------------------- +# autograd.Function entry point +# --------------------------------------------------------------------------- + + +class StackGroupedWeightFn(torch.autograd.Function): + """Fused ``torch.stack(weights).transpose(1, 2).contiguous()`` with + in-kernel ``[K, N] -> [N, K]`` transpose, fused across all experts. + + Inputs (variadic): + ``*weights``: ``E`` per-expert tensors, each ``[K, N]``, all same + ``dtype`` and ``device``, all contiguous (Megatron's parameter + allocator always returns contiguous; a defensive assertion is + kept in :meth:`forward` regardless). + + Output: + ``[E, N, K]`` contiguous tensor — bit-identical to the eager + ``torch.stack(weights, dim=0).transpose(1, 2).contiguous()`` chain + (the operation is a pure layout transform so there is no fp + rounding to worry about). + + BWD returns one ``[K, N]`` grad tensor per input weight; PyTorch's + autograd then writes each into ``weights[i].grad``. + """ + + @staticmethod + def forward(ctx, *weights: torch.Tensor) -> torch.Tensor: + if not weights: + raise ValueError("StackGroupedWeightFn requires at least one weight tensor") + + first = weights[0] + if first.ndim != 2: + raise ValueError( + f"StackGroupedWeightFn: each weight must be 2D, got " f"weight0.shape={tuple(first.shape)}" + ) + K, N = int(first.shape[0]), int(first.shape[1]) + dtype = first.dtype + device = first.device + + for i, w in enumerate(weights): + if w.shape != first.shape: + raise ValueError( + f"StackGroupedWeightFn: weight{i}.shape={tuple(w.shape)} " + f"differs from weight0.shape={tuple(first.shape)}" + ) + if w.dtype is not dtype: + raise TypeError( + f"StackGroupedWeightFn: weight{i}.dtype={w.dtype} " f"differs from weight0.dtype={dtype}" + ) + if w.device != device: + raise RuntimeError( + f"StackGroupedWeightFn: weight{i}.device={w.device} " + f"differs from weight0.device={device}" + ) + if not w.is_contiguous(): + raise ValueError( + f"StackGroupedWeightFn: weight{i} must be contiguous; " + "Megatron's parameter allocator always returns contiguous " + "tensors so a non-contiguous input here indicates a bug " + "upstream" + ) + + E = len(weights) + + weight_ptrs = torch.tensor([w.data_ptr() for w in weights], dtype=torch.int64, device=device) + + out = torch.empty(E, N, K, dtype=dtype, device=device) + + block_k, block_n = _pick_block(K, N) + grid = (E, triton.cdiv(K, block_k), triton.cdiv(N, block_n)) + _stack_grouped_weight_fwd_kernel[grid]( + weight_ptrs, + out, + E, + K, + N, + BLOCK_K=block_k, + BLOCK_N=block_n, + DTYPE=_triton_dtype(dtype), + ) + + ctx.E = E + ctx.K = K + ctx.N = N + ctx.dtype = dtype + ctx.device = device + ctx.block_k = block_k + ctx.block_n = block_n + return out + + @staticmethod + def backward(ctx, dout: torch.Tensor): # type: ignore[override] + E = ctx.E + K = ctx.K + N = ctx.N + dtype = ctx.dtype + device = ctx.device + + if not dout.is_contiguous(): + dout = dout.contiguous() + if tuple(dout.shape) != (E, N, K): + raise ValueError( + f"StackGroupedWeightFn backward: dout.shape={tuple(dout.shape)} " + f"!= expected (E={E}, N={N}, K={K})" + ) + if dout.dtype is not dtype: + dout = dout.to(dtype) + + dweights: List[torch.Tensor] = [torch.empty(K, N, dtype=dtype, device=device) for _ in range(E)] + dweight_ptrs = torch.tensor([dw.data_ptr() for dw in dweights], dtype=torch.int64, device=device) + + block_k, block_n = ctx.block_k, ctx.block_n + grid = (E, triton.cdiv(K, block_k), triton.cdiv(N, block_n)) + _stack_grouped_weight_bwd_kernel[grid]( + dweight_ptrs, + dout, + E, + K, + N, + BLOCK_K=block_k, + BLOCK_N=block_n, + DTYPE=_triton_dtype(dtype), + ) + + return tuple(dweights) + + +# --------------------------------------------------------------------------- +# Public Python entry points +# --------------------------------------------------------------------------- + + +_ENV_FLAG = "PRIMUS_STACK_GROUPED_WEIGHT_TRITON" + + +def is_triton_path_enabled() -> bool: + """Returns ``True`` when the Triton path is active. + + Default-on; set ``PRIMUS_STACK_GROUPED_WEIGHT_TRITON=0`` to fall back + to the eager ``torch.stack + transpose + contiguous`` chain. Treated + as a soft env (not a model-config flag) so an operator can A/B the + Triton path on a live training job without re-launching with a + different YAML. + """ + return os.environ.get(_ENV_FLAG, "1") != "0" + + +def eager_stack_grouped_weight(weights: List[torch.Tensor]) -> torch.Tensor: + """Reference implementation — the exact eager chain the Triton path + replaces. Kept exported for unit tests and for the env-flag-off path + in :class:`PrimusTurboGroupedMLP`. + """ + return torch.stack(weights, dim=0).transpose(1, 2).contiguous() + + +def stack_grouped_weight(weights: List[torch.Tensor]) -> torch.Tensor: + """Dispatch entry point used by :class:`PrimusTurboGroupedMLP`. + + Routes through the Triton path when :func:`is_triton_path_enabled` + returns True; otherwise calls :func:`eager_stack_grouped_weight`. + The Triton path uses an :class:`autograd.Function` so BWD scatters + the gradient back to each ``weights[i].grad`` via the inverse + transpose kernel; the eager path inherits PyTorch's default VJP + chain for ``torch.stack + transpose + contiguous``. + """ + if is_triton_path_enabled(): + return StackGroupedWeightFn.apply(*weights) + return eager_stack_grouped_weight(weights) + + +__all__ = [ + "StackGroupedWeightFn", + "eager_stack_grouped_weight", + "is_triton_path_enabled", + "stack_grouped_weight", +] diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index 51b4f699a..536cd22e7 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -3,7 +3,9 @@ # # See LICENSE for license information. ############################################################################### +import contextlib import gc +import os from contextlib import contextmanager, nullcontext from typing import Callable, Iterable, List, Optional, Tuple, Union @@ -85,9 +87,56 @@ from primus.core.pipeline_parallel.handler.offload_handler import OFFLOAD_BUFFER +try: + import triton + import triton.language as tl + + _HAVE_TRITON = True +except (ImportError, ModuleNotFoundError): + _HAVE_TRITON = False + _dummy_wgrads = {} +if _HAVE_TRITON: + + @triton.jit + def _inplace_add_kernel(dst_ptr, src_ptr, n_elements, BLOCK: tl.constexpr): + """In-place ``dst += src`` over a flat buffer, accumulating in fp32. + + int64 offsets so a single launch covers tensors with > 2**31 elements + (e.g. the consolidated grouped-expert ``main_grad`` of [E, N, K]), + instead of Torch's ``add_`` which tiles into multiple ~528M-element + ``vectorized_elementwise_kernel`` launches. + """ + pid = tl.program_id(axis=0).to(tl.int64) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + d = tl.load(dst_ptr + offs, mask=mask, other=0.0).to(tl.float32) + s = tl.load(src_ptr + offs, mask=mask, other=0.0).to(tl.float32) + tl.store(dst_ptr + offs, (d + s).to(dst_ptr.dtype.element_ty), mask=mask) + + +def _triton_inplace_add_(dst: torch.Tensor, src: torch.Tensor) -> torch.Tensor: + """``dst.add_(src)`` via a single Triton launch (fp32 accumulate). + + Falls back to Torch's ``add_`` when Triton is unavailable or the layout is + unsupported (non-contiguous / shape mismatch). The write is in-place on + ``dst``'s storage, so ``dst`` must be contiguous. + """ + if not _HAVE_TRITON or not dst.is_cuda or not dst.is_contiguous() or dst.numel() != src.numel(): + return dst.add_(src) + + dst_flat = dst.view(-1) + # reshape (not view) so a non-contiguous grad is materialized contiguously. + src_flat = src.reshape(-1) + n_elements = dst_flat.numel() + BLOCK = 8192 + grid = (triton.cdiv(n_elements, BLOCK),) + _inplace_add_kernel[grid](dst_flat, src_flat, n_elements, BLOCK=BLOCK) + return dst + + def _get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tensor: """Returns a dummy tensor of given shape. @@ -110,6 +159,25 @@ def _get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tenso return _dummy_wgrads[key].detach() +class _MainGradShim: + """Per-expert handle for primus_turbo's ``fused_grouped_wgrad`` over a + *consolidated* ``[E, N, K]`` grouped-expert weight (OPT-1). + + ``PrimusTurboGroupedLinear`` keeps one consolidated weight with a single + ``main_grad`` block, but ``fused_grouped_wgrad`` / ``_expert_main_grad_view`` + expect a list of per-expert handles, each exposing a 2-D ``main_grad`` and a + ``grad_added_to_main_grad`` flag. These shims point at the contiguous 2-D + slices ``main_grad[i]`` of the consolidated block, so the grouped GEMM + backward accumulates each expert's wgrad straight into the right slice. + """ + + __slots__ = ("main_grad", "grad_added_to_main_grad") + + def __init__(self, main_grad_slice: torch.Tensor) -> None: + self.main_grad = main_grad_slice + self.grad_added_to_main_grad = False + + def _bridge_weight_grad( x: torch.Tensor, weight: torch.nn.Parameter, weight_buffer: PrimusTurboQuantizedTensorPair ): @@ -139,8 +207,33 @@ def backward(ctx, grad_x, grad_quantized_weight, grad_quantized_weight_trans): weight, "grad_added_to_main_grad" ), "weight.grad_added_to_main_grad don't have grad_added_to_main_grad attribute." - weight.main_grad.add_(grad_quantized_weight) - weight.grad_added_to_main_grad = True + # NOTE: Set weight.grad_added_to_main_grad to True to avoid adding the + # quantized weight gradient to main_grad twice. + if grad_quantized_weight is None: + # OPT-1 fused path: the grouped GEMM backward already accumulated the + # expert wgrad straight into main_grad (under fused_grouped_wgrad) and + # returned grad_b=None, so there is nothing to add here -- just flag it. + weight.grad_added_to_main_grad = True + else: + # `is_gfx1250` only exists in newer primus_turbo builds (it gates a + # gfx1250-specific elementwise-add workaround). Older / feature-branch + # primus_turbo that predate it (e.g. the flydsl sparse-MLA attention branch) + # don't define it; treat a missing symbol as False so those builds still work + # on non-gfx1250 archs (gfx942 / gfx950) instead of raising ImportError. + try: + from primus_turbo.pytorch.core.utils import is_gfx1250 + + _use_triton_inplace_add = is_gfx1250() + except ImportError: + _use_triton_inplace_add = False + + if _use_triton_inplace_add: + # NOTE: The bandwith of torch's elementwise add kernel has issue. Use triton to temporary workaround for gfx1250. + _triton_inplace_add_(weight.main_grad, grad_quantized_weight) + else: + weight.main_grad.add_(grad_quantized_weight) + + weight.grad_added_to_main_grad = True return grad_x, _get_dummy_wgrad(list(weight.shape), weight.dtype), None, None @@ -1000,12 +1093,14 @@ def forward_internal( assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4." if is_first_microbatch: - need_cache_colwise = not self.disable_parameter_transpose_cache ( self.quantized_weight_buffer, self.quantized_weight_t_buffer, ) = _maybe_create_quantized_weight_buffers( - weight, float4_e2m1fn_x2, quant_config, need_cache_colwise=need_cache_colwise + weight, + float4_e2m1fn_x2, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, ) x, quantized_weight = _bridge_weight_grad( @@ -1702,7 +1797,8 @@ def forward_internal( """Forward step of the legacy PrimusTurbo grouped-gemm MLP.""" weights = self.weights # NOTE: keep x and m_splits on the same device - m_splits = m_splits.to(x.device) + if m_splits.device != x.device: + m_splits = m_splits.to(x.device) if PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() @@ -1729,13 +1825,65 @@ def forward_internal( ), ) - out = primus_turbo_torch.ops.grouped_gemm_fp8( - x, - quantized_weights, - m_splits, - trans_b=True, - config=quant_config.data(), + # OPT-1 (opt-in, single-GPU): accumulate the expert wgrad straight into + # main_grad in the grouped GEMM backward (beta=1 ACCUMULATE) instead of + # GEMM-wgrad -> _WeightGradBridge.main_grad.add_(). Per-expert shims point + # at the consolidated [E,N,K] main_grad slices; the grouped backward + # accumulates into them and returns grad_b=None, then the bridge backward + # just flags grad_added. ONLY safe with no gradient all-reduce / + # reduce-scatter (TP=1 / DP=1 / EP=1), since grad_b=None skips the reduce. + # Needs a turbo wheel carrying fused_grouped_wgrad (else falls back). + _wgrad_ctx = contextlib.nullcontext() + _fwg_dbg = ( + os.environ.get("PRIMUS_TURBO_FUSE_WGRAD_DEBUG") == "1" + and getattr(type(self), "_fwg_logn", 0) < 10 ) + _fwg_flag = os.environ.get("PRIMUS_TURBO_FUSE_GROUPED_WGRAD", "0") == "1" + _mg = getattr(weights, "main_grad", None) + if _fwg_flag and _mg is not None and _mg.dim() == 3: + try: + from primus_turbo.pytorch.ops.grouped_gemm_fp8 import ( + _expert_main_grad_view, + fused_grouped_wgrad, + ) + + _shims = [_MainGradShim(_mg[i]) for i in range(_mg.shape[0])] + if _fwg_dbg: + import sys + + _v = _expert_main_grad_view(_shims) + print( + f"[OPT-1] gate PASS shape={tuple(_mg.shape)} contig={_mg.is_contiguous()} " + f"stride={_mg.stride()} view={'OK' if _v is not None else 'REJECTED'}", + file=sys.stderr, + flush=True, + ) + type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1 + _wgrad_ctx = fused_grouped_wgrad(_shims) + except ImportError as _e: + if _fwg_dbg: + import sys + + print(f"[OPT-1] ImportError: {_e}", file=sys.stderr, flush=True) + type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1 + elif _fwg_dbg: + import sys + + print( + f"[OPT-1] gate FAIL flag={_fwg_flag} main_grad={'None' if _mg is None else f'dim={_mg.dim()}'}", + file=sys.stderr, + flush=True, + ) + type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1 + + with _wgrad_ctx: + out = primus_turbo_torch.ops.grouped_gemm_fp8( + x, + quantized_weights, + m_splits, + trans_b=True, + config=quant_config.data(), + ) elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled(): assert False, "FP4 is not supported in PrimusTurboGroupedLinear" else: 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 f78172ba8..48aedcd43 100644 --- a/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py +++ b/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py @@ -5,6 +5,7 @@ ############################################################################### import warnings +from types import SimpleNamespace from typing import Optional from megatron.core.extensions.transformer_engine import ( @@ -102,11 +103,25 @@ def __init__( return PrimusLegacyGroupedMLP +def _build_default_primus_args() -> SimpleNamespace: + """Fallback args for environments without initialized Primus globals.""" + return SimpleNamespace( + enable_primus_turbo=False, + use_turbo_gemm=False, + use_turbo_attention=False, + use_turbo_grouped_gemm=False, + moe_use_legacy_grouped_gemm=False, + ) + + class PrimusTurboSpecProvider(BackendSpecProvider): """A protocol for providing the submodules used in Spec building.""" def __init__(self, fallback_to_eager_attn: bool = False): - self.cfg = get_primus_args() + try: + self.cfg = get_primus_args() + except AssertionError: + self.cfg = _build_default_primus_args() self.fallback_to_eager_attn = fallback_to_eager_attn def linear(self) -> type: @@ -133,6 +148,43 @@ def row_parallel_linear(self) -> type: else TERowParallelLinear ) + def column_parallel_linear_with_gather_output(self) -> type: + """Non-TE column-parallel linear that supports ``gather_output=True``. + + TE / Turbo column-parallel wrappers explicitly raise + ``ValueError("Transformer Engine linear layers do not support + gather_output = True")`` (see + ``third_party/Megatron-LM/megatron/core/extensions/transformer_engine.py:747`` + and ``:972``). Callers that need a column-parallel layer + whose output dim is gathered back to full width across TP + ranks (so downstream math stays TP-agnostic) must use the + upstream Megatron-native :class:`ColumnParallelLinear`. + + Plan-3 P21: V4's ``linear_q_up_proj`` is the canonical + consumer — it shards ``q_lora_rank -> n_heads * head_dim`` + across TP and gathers the heads at forward time. + """ + return ColumnParallelLinear + + def row_parallel_linear_with_scatter_input(self) -> type: + """Non-TE row-parallel linear that supports ``input_is_parallel=False``. + + TE / Turbo row-parallel wrappers explicitly raise + ``ValueError("Transformer Engine linear layers do not support + input_is_parallel = False")`` (see + ``third_party/Megatron-LM/megatron/core/extensions/transformer_engine.py:1081``). + Callers that hand the layer a full-width (non-sharded) input + and want it scattered internally + the output all-reduced + must use the upstream Megatron-native :class:`RowParallelLinear`. + + Plan-3 P21: V4's grouped-O ``linear_o_b`` is the canonical + consumer — after the inner ``[..., n_per_group, o_lora_rank] + -> [..., o_groups * o_lora_rank]`` reshape, the input is + full-width across TP and the row-parallel layer's + weight-sharding + reduce is what produces the correct sum. + """ + return RowParallelLinear + def fuse_layernorm_and_linear(self) -> bool: """TE backend chooses a single module for layernorm and linear""" return True @@ -174,8 +226,7 @@ def grouped_mlp_modules( # Megatron callers only pass ``moe_use_grouped_gemm`` here, so when Primus # args do not expose the legacy switch we must match upstream TESpecProvider # and prefer TEGroupedMLP by default. - # let it raise an error if cfg does not have moe_use_legacy_grouped_gemm - moe_use_legacy_grouped_gemm = self.cfg.moe_use_legacy_grouped_gemm + moe_use_legacy_grouped_gemm = getattr(self.cfg, "moe_use_legacy_grouped_gemm", False) use_turbo_grouped_gemm = self.cfg.use_turbo_grouped_gemm assert not ( @@ -224,3 +275,157 @@ def grouped_mlp_modules( def activation_func(self) -> type: """Which module to use for activation function""" return TEActivationOp + + +class DeepSeekV4SpecProvider(PrimusTurboSpecProvider): + """DeepSeek-V4 provider rooted on PrimusTurboSpecProvider.""" + + def __init__(self, config=None): + super().__init__() + self.config = config + + def v4_norm_module(self): + """Norm module used by V4 specs (block / layer / final).""" + return self.layer_norm(rms_norm=True) + + def v4_q_layernorm(self) -> type: + """Norm module for V4's `q_norm` (RMSNorm on `q_lora_rank`). + + Same as MLA's `q_layernorm`; we use the for_qk path so the TE + version selection picks Apex / FusedLayerNorm on older TE. + """ + return self.layer_norm(rms_norm=True, for_qk=True) + + def v4_kv_layernorm(self) -> type: + """Norm module for V4's `kv_norm` (RMSNorm on `head_dim`). + + Single-latent KV: V4 normalizes the `wkv` output (one shared head) + BEFORE broadcasting to all query heads. + """ + return self.layer_norm(rms_norm=True, for_qk=True) + + def v4_mlp_activation_func(self) -> Optional[type]: + """Activation-func selection for V4 MLP / shared-expert specs. + + Plan-2 P18 (D2 audit): the parent ``activation_func()`` returns + the TE module **type** (``TEActivationOp``); but Megatron's + ``MLP.__init__`` only consumes ``submodules.activation_func`` when + ``config.use_te_activation_func == True`` (otherwise it falls + back to the callable in ``config.activation_func``). + + Returning the TE class unconditionally caused a silent + contract mismatch in V4 yamls (which keep + ``use_te_activation_func: false`` by default — V4 wants the + clamped-SwiGLU eager path so the activation-clamp gets + applied). + + Behavior: + + * ``config.use_te_activation_func`` is True → return the TE + activation class (instantiated by Megatron MLP at build). + * Otherwise → return ``None`` so the spec leaves the + ``MLPSubmodules.activation_func`` slot empty and Megatron + MLP uses ``config.activation_func`` (V4's clamped SwiGLU). + + This keeps the spec self-consistent: if the V4 yaml opts into + the TE path, the spec carries the TE class; otherwise the + spec carries ``None`` instead of a class that would be + silently ignored. + """ + cfg = getattr(self, "config", None) + if cfg is not None and bool(getattr(cfg, "use_te_activation_func", False)): + return self.activation_func() + return None + + def v4_grouped_mlp_modules( + self, moe_use_grouped_gemm: bool, moe_use_legacy_grouped_gemm: Optional[bool] = None + ): + """Grouped-MLP module selection for V4 MoE expert path.""" + return self.grouped_mlp_modules( + moe_use_grouped_gemm=moe_use_grouped_gemm, + moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm, + ) + + # ---- V4 spec factories (plan-2 P14 §5/§6) ------------------------- + + def v4_grouped_mlp_spec( + self, + *, + swiglu_limit: float, + moe_use_grouped_gemm: bool = True, + moe_use_legacy_grouped_gemm: Optional[bool] = None, + ): + """Return a ready-to-use ``ModuleSpec`` for V4 grouped MoE experts. + + The V4 pre-multiplication clamp itself is applied through + ``config.activation_func_clamp_value`` (Megatron's MLP eager + ``glu()`` already clamps gate (max=alpha) and up (+/- alpha) + before SiLU + multiply, which is bit-equal to the HF reference + ``Expert.forward`` math). This spec only commits to the right + grouped module + the column / row-parallel linears; the runtime + config carries the clamp value. + + Args: + swiglu_limit: V4 ``alpha`` value (the released ``DeepSeek-V4-Flash`` + checkpoint uses ``7.0``). Recorded on the returned spec + so the caller can assert it lines up with the runtime + config; not consumed by the grouped-MLP module directly. + moe_use_grouped_gemm: prefer the TE / Turbo grouped-gemm + module over the local SequentialMLP fallback (default + True in production). + moe_use_legacy_grouped_gemm: optional override for the + legacy code path; when ``None`` the provider reads the + Primus arg of the same name. + + Returns: + A ``ModuleSpec`` whose ``module`` is the grouped MoE module + and whose ``submodules`` carry the linear modules. Caller + wires this into :class:`DeepseekV4MoESubmodules.grouped_experts`. + + Notes: + * Plan-2 P14 §5 calls for "downgrade to local experts with + explicit warning" when the grouped backend cannot apply + the clamp; that downgrade lives in + :class:`DeepseekV4MoE` (which already builds + :class:`ClampedSwiGLUMLP` local experts when + ``pg_collection is None`` or the backend declares no + clamp support). + """ + del swiglu_limit # documented but not consumed by the grouped MLP itself + from megatron.core.transformer.spec_utils import ModuleSpec + + module, submodules = self.v4_grouped_mlp_modules( + moe_use_grouped_gemm=moe_use_grouped_gemm, + moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm, + ) + if submodules is None: + return ModuleSpec(module=module) + return ModuleSpec(module=module, submodules=submodules) + + def v4_router_spec(self, *, learned: bool = True): + """Return a ``ModuleSpec`` for the V4 hash / learned router. + + Args: + learned: when True (default) returns the learned router + spec (``layer_idx >= num_hash_layers``); when False + returns the hash-router spec (``layer_idx < + num_hash_layers``). + + Returns: + A bare-module ``ModuleSpec`` suitable for + :class:`DeepseekV4MoESubmodules.{learned_router, hash_router}`. + Both routers are ``nn.Module`` standalones (not + ``TopKRouter`` subclasses) so they instantiate cleanly on + CPU; aux-loss / z-loss / RouterReplay inheritance is + tracked as a P19 follow-up. + """ + from megatron.core.transformer.spec_utils import ModuleSpec + + from primus.backends.megatron.core.transformer.moe.v4_hash_router import ( + DeepseekV4HashRouter, + ) + from primus.backends.megatron.core.transformer.moe.v4_topk_router import ( + DeepseekV4LearnedRouter, + ) + + return ModuleSpec(module=(DeepseekV4LearnedRouter if learned else DeepseekV4HashRouter)) diff --git a/primus/backends/megatron/core/fp8_utils.py b/primus/backends/megatron/core/fp8_utils.py index 18076c578..66e4d2200 100644 --- a/primus/backends/megatron/core/fp8_utils.py +++ b/primus/backends/megatron/core/fp8_utils.py @@ -43,6 +43,16 @@ # Primus-Turbo not importable (not installed, or a transitive dep is broken). pass +# gfx1250 / turbo-free FP8: on builds where primus_turbo is only an import shim +# (not a real install), the turbo FP8 path (PrimusTurboQuantConfig + +# primus_turbo_fp8_autocast) cannot run. Set PRIMUS_FP8_DISABLE_TURBO=1 to force +# the TE-native FP8 branch (transformer_engine.pytorch.fp8_autocast), which needs +# no turbo. Opt-in; the production turbo path is unaffected when unset. +import os as _os + +if _os.environ.get("PRIMUS_FP8_DISABLE_TURBO", "0") == "1": + HAVE_TURBO = False + SCALING_BLOCK_SIZE = 128 MXFP8_SCALING_BLOCK_SIZE = 32 diff --git a/primus/backends/megatron/core/fusions/fused_bias_swiglu.py b/primus/backends/megatron/core/fusions/fused_bias_swiglu.py new file mode 100644 index 000000000..19a689c7f --- /dev/null +++ b/primus/backends/megatron/core/fusions/fused_bias_swiglu.py @@ -0,0 +1,342 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Primus override for clamped weighted SwiGLU MoE activation fusion. + +Fuses clamp(gate), clamp(up), SiLU, multiply, and router-prob weighting into +a single ``@jit_fuser`` forward/backward pair instead of separate ``chunk`` / +``clamp`` / ``clamp`` / ``cat`` ATen ops followed by ``WeightedSwiGLUFunction``. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from megatron.core.fusions.fused_bias_swiglu import WeightedSwiGLUFunction + +# Column tile size for the elementwise kernels / reduction inner loop. +_BLOCK_N = 1024 + + +@triton.jit +def _clamped_weighted_swiglu_fwd_kernel( + y_ptr, # [M, 2*half] row-major: gate = [:, :half], up = [:, half:] + w_ptr, # [M] per-token weights (only read if HAS_WEIGHTS) + out_ptr, # [M, half] row-major output + half, + K, # = 2 * half (row stride of y) + stride_w, + clamp_value, + HAS_WEIGHTS: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask = cols < half + + y_row = y_ptr + pid_m * K + gate = tl.load(y_row + cols, mask=mask, other=0.0).to(tl.float32) + up = tl.load(y_row + half + cols, mask=mask, other=0.0).to(tl.float32) + + gate_c = tl.minimum(gate, clamp_value) + up_c = tl.minimum(tl.maximum(up, -clamp_value), clamp_value) + out = (gate_c * tl.sigmoid(gate_c)) * up_c + if HAS_WEIGHTS: + out = out * tl.load(w_ptr + pid_m * stride_w).to(tl.float32) + + tl.store(out_ptr + pid_m * half + cols, out, mask=mask) + + +@triton.jit +def _clamped_swiglu_bwd_kernel( + g_ptr, # [M, half] grad w.r.t. output + y_ptr, # [M, 2*half] pre-clamp input + w_ptr, # [M] per-token weights (only read if HAS_WEIGHTS) + dy_ptr, # [M, 2*half] grad w.r.t. input + half, + K, + stride_gm, + stride_w, + clamp_value, + HAS_WEIGHTS: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask = cols < half + + g = tl.load(g_ptr + pid_m * stride_gm + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_WEIGHTS: + g = g * tl.load(w_ptr + pid_m * stride_w).to(tl.float32) + + y_row = y_ptr + pid_m * K + gate = tl.load(y_row + cols, mask=mask, other=0.0).to(tl.float32) + up = tl.load(y_row + half + cols, mask=mask, other=0.0).to(tl.float32) + + gate_c = tl.minimum(gate, clamp_value) + up_c = tl.minimum(tl.maximum(up, -clamp_value), clamp_value) + sig = tl.sigmoid(gate_c) + silu = gate_c * sig + + # d/d gate_c [SiLU(gate_c)] = sig * (1 + gate_c * (1 - sig)) + dy_glu_c = g * sig * (1.0 + gate_c * (1.0 - sig)) * up_c + dy_linear_c = g * silu + + # Straight-through the clamp: gradient passes only inside the (un)clamped region. + gate_keep = (gate <= clamp_value).to(tl.float32) + up_keep = ((up >= -clamp_value) & (up <= clamp_value)).to(tl.float32) + + dy_row = dy_ptr + pid_m * K + tl.store(dy_row + cols, dy_glu_c * gate_keep, mask=mask) + tl.store(dy_row + half + cols, dy_linear_c * up_keep, mask=mask) + + +@triton.jit +def _clamped_swiglu_weights_grad_kernel( + g_ptr, # [M, half] grad w.r.t. output + y_ptr, # [M, 2*half] pre-clamp input + wg_ptr, # [M] per-token weights grad + half, + K, + stride_gm, + clamp_value, + NUM_TILES: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + y_row = y_ptr + pid_m * K + g_row = g_ptr + pid_m * stride_gm + + acc = tl.zeros([BLOCK_N], dtype=tl.float32) + for t in tl.static_range(NUM_TILES): + cols = t * BLOCK_N + tl.arange(0, BLOCK_N) + mask = cols < half + g = tl.load(g_row + cols, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(y_row + cols, mask=mask, other=0.0).to(tl.float32) + up = tl.load(y_row + half + cols, mask=mask, other=0.0).to(tl.float32) + gate_c = tl.minimum(gate, clamp_value) + up_c = tl.minimum(tl.maximum(up, -clamp_value), clamp_value) + activation = (gate_c * tl.sigmoid(gate_c)) * up_c + acc += activation * g # masked lanes contribute 0 (g == 0) + + tl.store(wg_ptr + pid_m, tl.sum(acc, axis=0)) + + +def clamped_weighted_swiglu(y: torch.Tensor, weights: torch.Tensor, clamp_value: float) -> torch.Tensor: + """Clamped SwiGLU with per-token weights in one fused Triton kernel. + + Semantics match the eager path: + gate_c = clamp(gate, max=alpha) + up_c = clamp(up, min=-alpha, max=alpha) + out = SiLU(gate_c) * up_c * weights + """ + lead = y.shape[:-1] + K = y.shape[-1] + half = K // 2 + y2 = y.reshape(-1, K).contiguous() + M = y2.shape[0] + w = weights.reshape(M).contiguous() + + out = torch.empty((M, half), dtype=y.dtype, device=y.device) + grid = (M, triton.cdiv(half, _BLOCK_N)) + _clamped_weighted_swiglu_fwd_kernel[grid]( + y2, + w, + out, + half, + K, + w.stride(0), + float(clamp_value), + HAS_WEIGHTS=True, + BLOCK_N=_BLOCK_N, + ) + return out.reshape(*lead, half) + + +def clamped_swiglu(y: torch.Tensor, clamp_value: float) -> torch.Tensor: + """Clamped SwiGLU without per-token weights in one fused Triton kernel. + + Semantics match the eager path used by ``mlp.MLP`` shared experts: + gate_c = clamp(gate, max=alpha) + up_c = clamp(up, min=-alpha, max=alpha) + out = SiLU(gate_c) * up_c + """ + lead = y.shape[:-1] + K = y.shape[-1] + half = K // 2 + y2 = y.reshape(-1, K).contiguous() + M = y2.shape[0] + + out = torch.empty((M, half), dtype=y.dtype, device=y.device) + grid = (M, triton.cdiv(half, _BLOCK_N)) + # w_ptr is unused when HAS_WEIGHTS=False; pass y2 as a valid placeholder. + _clamped_weighted_swiglu_fwd_kernel[grid]( + y2, + y2, + out, + half, + K, + 0, + float(clamp_value), + HAS_WEIGHTS=False, + BLOCK_N=_BLOCK_N, + ) + return out.reshape(*lead, half) + + +def clamped_swiglu_back(g: torch.Tensor, y: torch.Tensor, clamp_value: float) -> torch.Tensor: + """Backward for clamped SwiGLU w.r.t. the pre-clamp concatenated input.""" + lead = y.shape[:-1] + K = y.shape[-1] + half = K // 2 + y2 = y.reshape(-1, K).contiguous() + g2 = g.reshape(-1, half).contiguous() + M = y2.shape[0] + + dy = torch.empty((M, K), dtype=g.dtype, device=g.device) + grid = (M, triton.cdiv(half, _BLOCK_N)) + _clamped_swiglu_bwd_kernel[grid]( + g2, + y2, + g2, + dy, + half, + K, + g2.stride(0), + 0, + float(clamp_value), + HAS_WEIGHTS=False, + BLOCK_N=_BLOCK_N, + ) + return dy.reshape(*lead, K) + + +def clamped_weighted_swiglu_back(g: torch.Tensor, y: torch.Tensor, weights: torch.Tensor, clamp_value: float): + """Backward for :func:`clamped_weighted_swiglu`.""" + input_dtype = y.dtype + w_dtype = weights.dtype + lead = y.shape[:-1] + K = y.shape[-1] + half = K // 2 + y2 = y.reshape(-1, K).contiguous() + g2 = g.reshape(-1, half).contiguous() + M = y2.shape[0] + w = weights.reshape(M).contiguous() + + # input grad: reuse the clamped-swiglu backward kernel with grad pre-scaled by weights. + input_grad = torch.empty((M, K), dtype=input_dtype, device=y.device) + grid = (M, triton.cdiv(half, _BLOCK_N)) + _clamped_swiglu_bwd_kernel[grid]( + g2, + y2, + w, + input_grad, + half, + K, + g2.stride(0), + w.stride(0), + float(clamp_value), + HAS_WEIGHTS=True, + BLOCK_N=_BLOCK_N, + ) + + # weights grad: sum over the feature dim of SiLU(gate_c) * up_c * g. + weights_grad = torch.empty((M,), dtype=w_dtype, device=y.device) + _clamped_swiglu_weights_grad_kernel[(M,)]( + g2, + y2, + weights_grad, + half, + K, + g2.stride(0), + float(clamp_value), + NUM_TILES=triton.cdiv(half, _BLOCK_N), + BLOCK_N=_BLOCK_N, + ) + + return input_grad.reshape(*lead, K), weights_grad.reshape(*weights.shape) + + +class ClampedWeightedSwiGLUFunction(torch.autograd.Function): + """Autograd wrapper for clamped token-wise weighted SwiGLU.""" + + @staticmethod + def forward(ctx, input: torch.Tensor, weights: torch.Tensor, fp8_input_store: bool, clamp_value: float): + input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input + ctx.save_for_backward(input_for_backward, weights) + ctx.ori_input_dtype = input.dtype + ctx.fp8_input_store = fp8_input_store + ctx.clamp_value = float(clamp_value) + return clamped_weighted_swiglu(input, weights, ctx.clamp_value) + + @staticmethod + def backward(ctx, grad_output): + input, weights = ctx.saved_tensors + input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input + input_grad, weights_grad = clamped_weighted_swiglu_back(grad_output, input, weights, ctx.clamp_value) + return input_grad, weights_grad, None, None + + +def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False, clamp_value=None): + """Token-wise-weighted SwiGLU fusion with optional clamped gate/up.""" + ori_shape = input.shape + assert len(ori_shape) in [2, 3] + input = input.view(-1, ori_shape[-1]) + if bias is not None: + raise NotImplementedError("Bias is not supported for weighted swiglu fusion") + + if clamp_value is not None: + output = ClampedWeightedSwiGLUFunction.apply(input, weights, fp8_input_store, float(clamp_value)) + else: + output = WeightedSwiGLUFunction.apply(input, weights, fp8_input_store) + + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) + + +class ClampedSwiGLUFunction(torch.autograd.Function): + """Autograd wrapper for clamped (non-weighted) SwiGLU.""" + + @staticmethod + def forward(ctx, input: torch.Tensor, fp8_input_store: bool, clamp_value: float): + input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input + ctx.save_for_backward(input_for_backward) + ctx.ori_input_dtype = input.dtype + ctx.fp8_input_store = fp8_input_store + ctx.clamp_value = float(clamp_value) + return clamped_swiglu(input, ctx.clamp_value) + + @staticmethod + def backward(ctx, grad_output): + (input,) = ctx.saved_tensors + input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input + input_grad = clamped_swiglu_back(grad_output, input, ctx.clamp_value) + return input_grad, None, None + + +def swiglu_impl(input, bias, fp8_input_store=False, clamp_value=None): + """Non-weighted SwiGLU fusion with optional clamped gate/up. + + Mirrors ``megatron.core.fusions.fused_bias_swiglu.bias_swiglu_impl`` but + adds the DeepSeek-V4 pre-multiplication clamp when ``clamp_value`` is set. + Used by the shared-expert MLP path, which has no per-token weights. + """ + ori_shape = input.shape + assert len(ori_shape) in [2, 3] + input = input.view(-1, ori_shape[-1]) + if bias is not None: + raise NotImplementedError("Bias is not supported for clamped swiglu fusion") + + if clamp_value is not None: + output = ClampedSwiGLUFunction.apply(input, fp8_input_store, float(clamp_value)) + else: + from megatron.core.fusions.fused_bias_swiglu import SwiGLUFunction + + output = SwiGLUFunction.apply(input, fp8_input_store, False) + + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) diff --git a/primus/backends/megatron/core/fusions/fused_pad_routing_map.py b/primus/backends/megatron/core/fusions/fused_pad_routing_map.py new file mode 100644 index 000000000..c61c58d58 --- /dev/null +++ b/primus/backends/megatron/core/fusions/fused_pad_routing_map.py @@ -0,0 +1,114 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Primus override for the fused routing-map padding used by MoE quantization. + +Rewrites Megatron's ``fused_pad_routing_map`` as a self-contained Triton kernel +that operates directly on the native ``[num_tokens, num_experts]`` layout, so no +``transpose`` / ``contiguous`` / intermediate copy is needed. It is intentionally +*not* wrapped with ``@jit_fuser`` (``torch.compile``): the kernel is already the +fused region, and wrapping user-written Triton kernels with ``torch.compile`` +triggers a functionalization failure on some torch/triton combos. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import torch +from megatron.core.utils import null_decorator +from packaging import version + +try: + import triton + import triton.language as tl + + if version.parse(triton.__version__) < version.parse("3.4.0") and not torch.cuda.is_available(): + HAVE_TRITON = False + else: + HAVE_TRITON = tl.constexpr(version.parse(triton.__version__) >= version.parse("2.0.0")) +except ImportError: + HAVE_TRITON = False + +if not HAVE_TRITON: + triton = MagicMock() + triton.jit = null_decorator + triton.autotune = null_decorator + triton.heuristics = null_decorator + tl = MagicMock() + + +@triton.jit +def _pad_routing_map_kernel( + routing_map_ptr, # *Pointer* to [num_tokens, num_experts] row-major routing map + output_ptr, # *Pointer* to [num_tokens, num_experts] row-major output + num_tokens, + num_experts, + pad_multiple: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + # Each program instance handles one expert, i.e. one *column* of the + # [num_tokens, num_experts] routing map. This lets us operate directly on the + # native (row-major) layout, without transposing/copying to [num_experts, num_tokens]. + expert_idx = tl.program_id(axis=0) + + # Token indices for this block + token_indices = tl.arange(0, BLOCK_SIZE) + token_mask = token_indices < num_tokens + + # Column-strided access: element (t, expert_idx) lives at t * num_experts + expert_idx. + offsets = token_indices * num_experts + expert_idx + + # Load this expert's column; out-of-bounds tokens read as 0 and are masked on store. + row = tl.load(routing_map_ptr + offsets, mask=token_mask, other=0).to(tl.int32) + + # 1. Number of tokens currently routed to this expert. + num_ones = tl.sum(row, axis=0) + + # 2. How many zeros must be flipped to 1 to reach the next multiple of pad_multiple. + remainder = num_ones % pad_multiple + num_to_pad = tl.where(remainder != 0, pad_multiple - remainder, 0) + + # 3. 1-based cumulative rank of each zero within the column. + is_zero = row == 0 + zero_ranks = tl.cumsum(is_zero.to(tl.int32), axis=0) + + # 4. Flip only the first `num_to_pad` zeros to 1. + mask_to_flip = (zero_ranks <= num_to_pad) & is_zero + output_row = tl.where(mask_to_flip, 1, row) + + # 5. Store back in the same native layout, masking out-of-bounds tokens. + tl.store(output_ptr + offsets, output_row, mask=token_mask) + + +def fused_pad_routing_map(routing_map: torch.Tensor, pad_multiple: int) -> torch.Tensor: + """Fused version of pad_routing_map. + Args: + routing_map (torch.Tensor): A boolean or integer tensor of shape [num_tokens, + num_experts] indicating which tokens are routed to which experts. + pad_multiple (int): The multiple to pad each expert's token count to. + + Returns: + torch.Tensor: The padded routing map of shape [num_tokens, num_experts]. + """ + num_tokens, num_experts = routing_map.shape + if num_tokens == 0: + return routing_map + + # Operate directly on the native [num_tokens, num_experts] layout: the kernel reads + # each expert's column with a stride of num_experts, so no transpose/copy is needed. + routing_map = routing_map.contiguous() + output_map = torch.empty_like(routing_map, dtype=torch.int32) + + # One program instance per expert (column). + grid = (num_experts,) + BLOCK_SIZE = triton.next_power_of_2(num_tokens) + + _pad_routing_map_kernel[grid]( + routing_map, output_map, num_tokens, num_experts, pad_multiple, BLOCK_SIZE=BLOCK_SIZE + ) + + return output_map # [num_tokens, num_experts] diff --git a/primus/backends/megatron/core/models/deepseek_v4/__init__.py b/primus/backends/megatron/core/models/deepseek_v4/__init__.py new file mode 100644 index 000000000..4ffc1f8a4 --- /dev/null +++ b/primus/backends/megatron/core/models/deepseek_v4/__init__.py @@ -0,0 +1,52 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 model package. + +Plan-2 P17 surface (post-cleanup): + + DeepseekV4Model # top-level model (LanguageModule) + DeepseekV4TransformerBlock # decoder block (TransformerBlock subclass) + DeepseekV4HybridLayer # decoder layer (TransformerLayer subclass) + get_v4_mtp_block_spec # spec helper for upstream MultiTokenPredictionBlock + deepseek_v4_builder # builder used by model_provider + model_provider # Megatron pretrain() entry point + +Retired in plan-2 P17: + + DeepseekV4MTPBlock # legacy primus-owned MTP head; replaced + # by the spec-based path above. The + # ``v4_use_custom_mtp_block`` config + # toggle is also gone. +""" + +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block import ( + DeepseekV4HybridLayer, + DeepseekV4HybridLayerSubmodules, + DeepseekV4TransformerBlock, + DeepseekV4TransformerBlockSubmodules, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import ( + deepseek_v4_builder, + model_provider, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_model import ( + DeepseekV4Model, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_mtp_specs import ( + get_v4_mtp_block_spec, +) + +__all__ = [ + "DeepseekV4Model", + "DeepseekV4TransformerBlock", + "DeepseekV4TransformerBlockSubmodules", + "DeepseekV4HybridLayer", + "DeepseekV4HybridLayerSubmodules", + "get_v4_mtp_block_spec", + "deepseek_v4_builder", + "model_provider", +] diff --git a/primus/backends/megatron/core/models/deepseek_v4/build_context.py b/primus/backends/megatron/core/models/deepseek_v4/build_context.py new file mode 100644 index 000000000..1f132a049 --- /dev/null +++ b/primus/backends/megatron/core/models/deepseek_v4/build_context.py @@ -0,0 +1,102 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-2 P18 — DeepSeek-V4 build context / provider singleton. + +This module hosts the helpers that make sure a single +:class:`DeepSeekV4SpecProvider` is constructed per builder call and +threaded down to every spec helper. Without this, the provider was +re-instantiated inside ``_build_projection`` (block.py), +``DeepseekV4TransformerBlock.__init__`` (block.py), the layer-spec +factory (``deepseek_v4_layer_specs.py``), and the MTP spec helper +(``deepseek_v4_mtp_specs.py``) — every call paid the +``BackendSpecProvider`` setup cost and there was no single place to +audit which provider was actually wiring the V4 modules. + +The implementation deliberately avoids module-level globals: the +provider is cached **on the config object** under a private attribute +name. This keeps each ``DeepSeekV4TransformerConfig`` instance +self-contained (different configs get different providers) and avoids +leaking state across unit tests. + +Usage: + +.. code-block:: python + + from .build_context import resolve_v4_provider + + def some_builder(*, config): + provider = resolve_v4_provider(config) + ... +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: # avoid eager torch import in this lightweight module + from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import ( + DeepSeekV4SpecProvider, + ) + +_PROVIDER_ATTR = "_v4_spec_provider_singleton" + + +def resolve_v4_provider(config) -> "DeepSeekV4SpecProvider": + """Return a cached :class:`DeepSeekV4SpecProvider` for ``config``. + + The provider is cached on the ``config`` object itself (under a + private attribute name) so: + + * repeated calls during a single builder invocation reuse the + same provider instance, + * different configs always get fresh providers, + * the cache is naturally garbage-collected when the config is + released, + * the helper is fully thread-safe-by-construction (each builder + thread holds its own config). + + Args: + config: a :class:`DeepSeekV4TransformerConfig` instance. The + helper accesses ``config.__dict__`` directly so it works + on dataclasses without additional setattr support. + + Returns: + The cached :class:`DeepSeekV4SpecProvider`. The first call + constructs it; later calls reuse it. + """ + cached = getattr(config, _PROVIDER_ATTR, None) + if cached is not None: + return cached + + # Lazy import: this module must be lightweight enough to import + # from the dataclass module without cyclic risks. + from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import ( + DeepSeekV4SpecProvider, + ) + + provider = DeepSeekV4SpecProvider(config=config) + try: + setattr(config, _PROVIDER_ATTR, provider) + except Exception: + # Some MagicMock-style configs in unit tests reject setattr; that + # is OK — we just don't cache and pay the cost on each call. + pass + return provider + + +def reset_v4_provider_cache(config) -> None: + """Drop the cached provider on ``config``. Intended for unit tests + that need to force a re-build (e.g. after monkey-patching the + provider class).""" + if hasattr(config, _PROVIDER_ATTR): + try: + delattr(config, _PROVIDER_ATTR) + except AttributeError: + pass + + +__all__ = ["resolve_v4_provider", "reset_v4_provider_cache"] diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py new file mode 100644 index 000000000..b06eac1ea --- /dev/null +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py @@ -0,0 +1,1190 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +DeepSeek-V4 transformer block (multi-stream HC + per-layer attention dispatch). + +Reference: +* techblog §1 ("Hybrid Attention") — per-layer attention selected by + ``compress_ratios[layer_id]``. +* techblog §2 ("mHC: Manifold-Constrained Hyper-Connections") — ``hc_mult`` + parallel hidden streams, mixed via per-layer ``HyperMixer`` and final + ``HyperHead`` collapse. + +Plan-2 P15 (this commit) — Megatron parent-class integration: + +* :class:`DeepseekV4HybridLayer` now subclasses + :class:`megatron.core.transformer.transformer_layer.TransformerLayer` + (via ``MegatronModule`` ``__init__`` bypass) and reuses upstream + submodule names (``input_layernorm`` / ``self_attention`` / + ``pre_mlp_layernorm`` / ``mlp``) plus V4-specific ``attn_hc`` / ``ffn_hc`` + hooks. The submodules dataclass extends + :class:`TransformerLayerSubmodules` so Megatron spec lifecycle code that + inspects the layer's ``submodules_config`` works without bespoke V4 + branches. +* :class:`DeepseekV4TransformerBlock` now subclasses + :class:`megatron.core.transformer.transformer_block.TransformerBlock` + for type identity / sharded-state-dict integration. ``HyperHead`` is now + built **only on the post_process stage** (it was previously built on + every PP rank, which wasted memory and risked correctness drift). +* PP K-stream packing: ``[B, S, K, D] <-> [S*K, B, D]`` lift / lower + helpers (:func:`_lift_streams_in` / :func:`_lower_streams_out`) carry + the K dimension across PP P2P boundaries by folding it into the + sequence axis. The first PP stage lifts the embedded ``[S, B, D]`` to K + streams; intermediate stages preserve K (``[S*K, B, D]`` send/recv); + the final stage collapses with ``HyperHead`` to ``[B, S, D]`` and + transposes back to ``[S, B, D]`` for the output layer. +* ``token_ids`` is now a real forward kwarg threaded through + ``DeepseekV4Model.forward -> DeepseekV4TransformerBlock.forward -> + DeepseekV4HybridLayer.forward -> DeepseekV4MoE.forward -> + DeepseekV4HashRouter.forward``. The legacy ``decoder._v4_token_ids`` + attribute stash is gone. +* ``position_ids`` is now consumed from the caller (forward kwarg) when + provided; the legacy ``arange(S)`` shortcut is kept only as a fallback + for callers that omit it (e.g. tiny CPU smokes). + +Forward shape contract: +* Input ``hidden_states`` is ``[S, B, D]`` on the first PP stage and + ``[S*K, B, D]`` on subsequent stages (where K = ``hc_mult``). The block + reshapes to the K-stream form for HC math and packs back before + returning to the next PP stage. +* Output is ``[S*K, B, D]`` on non-final PP stages and ``[S, B, D]`` on + the final (post_process) stage after HyperHead collapse. +""" + +from __future__ import annotations + +import ast +import logging +from contextlib import nullcontext +from dataclasses import dataclass +from typing import Callable, List, Optional, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from megatron.core.transformer.mlp import MLPSubmodules +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.moe.shared_experts import SharedExpertMLP +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_block import TransformerBlock +from megatron.core.transformer.transformer_layer import ( + TransformerLayer, + TransformerLayerSubmodules, +) +from megatron.core.utils import make_viewless_tensor + +from primus.backends.megatron.core.models.deepseek_v4.build_context import ( + resolve_v4_provider, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, +) +from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( + DeepseekV4Attention, +) +from primus.backends.megatron.core.transformer.dual_rope import DualRoPE +from primus.backends.megatron.core.transformer.hyper_connection import ( + HyperHead, + HyperMixer, +) +from primus.backends.megatron.core.transformer.local_rmsnorm import LocalRMSNorm +from primus.backends.megatron.core.transformer.moe.v4_hash_router import ( + DeepseekV4HashRouter, +) +from primus.backends.megatron.core.transformer.moe.v4_moe import ( + DeepseekV4MoE, + DeepseekV4MoESubmodules, +) +from primus.backends.megatron.core.transformer.moe.v4_topk_router import ( + DeepseekV4LearnedRouter, +) + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Provider-aware projection helpers +# --------------------------------------------------------------------------- + + +def _default_init_method(_weight: torch.Tensor) -> None: + return None + + +def _build_projection( + in_features: int, + out_features: int, + *, + config: DeepSeekV4TransformerConfig, +) -> nn.Module: + """Build a duplicated linear via the V4 provider. + + When ``config`` is ``None`` (CPU unit tests that build the dense + SwiGLU MLP without a provider) we instantiate a plain + :class:`nn.Linear` with the same shape. Otherwise we delegate to + ``provider.linear()`` and let any constructor failure bubble up — + Plan-3 P21 retired the ``try/except/return nn.Linear`` fallback + because it produced an unsharded layer that masked real provider + bugs at TP=1 and would diverge at TP>1. + """ + if config is None: + return nn.Linear(in_features, out_features, bias=False) + + provider = resolve_v4_provider(config) + linear_module_cls = provider.linear() + init_method: Callable = config.init_method or _default_init_method + return linear_module_cls( + input_size=in_features, + output_size=out_features, + parallel_mode="duplicated", + config=config, + init_method=init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + tp_comm_buffer_name=None, + is_expert=False, + ) + + +def _projection_forward(proj: nn.Module, x: torch.Tensor) -> torch.Tensor: + out = proj(x) + if isinstance(out, tuple): + return out[0] + return out + + +# --------------------------------------------------------------------------- +# Pieces used by every layer +# --------------------------------------------------------------------------- + + +def _parse_int_sequence(value, *, field_name: str) -> Optional[List[int]]: + """Parse config-provided sequence fields into ``List[int]``. + + YAML values may arrive as: + - actual list/tuple: ``[0, 4, 128, ...]`` + - stringified list: ``"[0, 4, 128, ...]"`` + """ + if value is None: + return None + + parsed = value + if isinstance(parsed, str): + try: + parsed = ast.literal_eval(parsed) + except (SyntaxError, ValueError) as exc: + raise ValueError(f"{field_name} must be a list-like value, got {value!r}") from exc + + if isinstance(parsed, torch.Tensor): + parsed = parsed.tolist() + + if not isinstance(parsed, (list, tuple)): + raise TypeError(f"{field_name} must be list/tuple, got {type(parsed).__name__}") + + out: List[int] = [] + for i, item in enumerate(parsed): + try: + out.append(int(item)) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field_name}[{i}]={item!r} is not int-castable") from exc + return out + + +def _normalize_compress_ratios( + compress_ratios, + *, + num_layers: int, + mtp_num_layers: int, +) -> List[int]: + """Normalize ``compress_ratios`` to exactly ``num_layers`` entries.""" + ratios = _parse_int_sequence(compress_ratios, field_name="compress_ratios") + if ratios is None: + return [0] * num_layers + + if len(ratios) == num_layers: + return ratios + + # Common DeepSeek layout: decoder ratios + mtp ratios in one list. + if len(ratios) == num_layers + mtp_num_layers: + logger.warning( + "compress_ratios has decoder+MTP length (%s), truncating to decoder num_layers (%s).", + len(ratios), + num_layers, + ) + return ratios[:num_layers] + + if len(ratios) > num_layers: + logger.warning( + "compress_ratios length (%s) > num_layers (%s), truncating.", + len(ratios), + num_layers, + ) + return ratios[:num_layers] + + # len(ratios) < num_layers: extend with last ratio (or 0 if empty). + pad_value = ratios[-1] if ratios else 0 + logger.warning( + "compress_ratios length (%s) < num_layers (%s), padding with %s.", + len(ratios), + num_layers, + pad_value, + ) + return ratios + [pad_value] * (num_layers - len(ratios)) + + +# Plan-2 P17 dedup: the duplicated ``_RMSNorm`` here was retired. +# Block-level RMSNorm fallbacks now use the canonical +# :class:`LocalRMSNorm` from +# ``primus.backends.megatron.core.transformer.local_rmsnorm`` (imported +# at the top of this file). Spec-driven norms still come from +# ``DeepSeekV4SpecProvider.v4_norm_module()``. + + +class _DenseSwiGLUMLP(nn.Module): + """Plain dense SwiGLU FFN with V4 pre-multiplication clamp. + + Used for non-MoE layers (or as a fallback when ``num_routed_experts`` + is 0). V4-Flash has a tiny number of dense head/tail layers; the bulk + of layers are MoE (see :class:`DeepseekV4MoE`). + + The activation matches V4's released ``Expert.forward``: + ``SiLU(clamp(gate, max=alpha)) * clamp(up, +/- alpha)`` with + ``alpha = config.swiglu_limit`` (``0`` disables clamping). + """ + + def __init__( + self, + config: DeepSeekV4TransformerConfig, + *, + pg_collection=None, + ) -> None: + del pg_collection + super().__init__() + hidden_size = int(config.hidden_size) + ffn_hidden_size = int(config.ffn_hidden_size) + self.w_gate = _build_projection( + hidden_size, + ffn_hidden_size, + config=config, + ) + self.w_up = _build_projection( + hidden_size, + ffn_hidden_size, + config=config, + ) + self.w_down = _build_projection( + ffn_hidden_size, + hidden_size, + config=config, + ) + self.swiglu_limit = float(getattr(config, "swiglu_limit", 0.0) or 0.0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate = _projection_forward(self.w_gate, x) + up = _projection_forward(self.w_up, x) + if self.swiglu_limit > 0.0: + gate = gate.clamp(max=self.swiglu_limit) + up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) + return _projection_forward(self.w_down, F.silu(gate) * up) + + +@dataclass +class DeepseekV4HybridLayerSubmodules(TransformerLayerSubmodules): + """Spec tree for one DeepSeek-V4 hybrid layer. + + Plan-2 P15: extends :class:`TransformerLayerSubmodules` so the V4 + layer slots into Megatron's spec lifecycle without bespoke field + plumbing. The four core fields use upstream-canonical names: + + * ``input_layernorm`` (was ``attn_norm``) + * ``self_attention`` (was ``attention``) + * ``pre_mlp_layernorm`` (was ``ffn_norm``) + * ``mlp`` (was ``ffn``) + + V4 adds two HC mixer hooks that the hybrid forward consumes; both + are optional (``None`` when ``hc_mult == 1`` because there is only + one stream to mix). + + The cross-attention / BDA fields inherited from the parent stay at + their defaults (V4 has no cross-attention and no BDA wrapper — the + HC residual path replaces both). + """ + + attn_hc: Optional[Union[ModuleSpec, type]] = None + ffn_hc: Optional[Union[ModuleSpec, type]] = None + + +@dataclass +class DeepseekV4TransformerBlockSubmodules: + """Spec tree for the DeepSeek-V4 decoder block. + + Plan-2 P15: ``hyper_head`` is built **only** on the post_process + stage (the final HC collapse `[B, S, K, D] -> [B, S, D]`). On + non-final PP stages the block keeps the K-stream form intact and + relies on the lift / lower helpers for P2P shape compatibility. + """ + + layer_specs: Optional[List[ModuleSpec]] = None + hyper_head: Optional[Union[ModuleSpec, type]] = None + final_layernorm: Optional[Union[ModuleSpec, type]] = None + + +# --------------------------------------------------------------------------- +# K-stream <-> sequence-axis packing helpers (PP P2P shape carrier) +# --------------------------------------------------------------------------- + + +def _lift_streams_in( + hidden_states: torch.Tensor, + *, + pre_process: bool, + hc_mult: int, +) -> torch.Tensor: + """Reshape the block's input to ``[B, S, K, D]`` for HC math. + + Args: + hidden_states: incoming tensor. + * On the first PP stage (``pre_process=True``): ``[S, B, D]`` + (Megatron's sequence-first convention) — we expand to K + streams. + * On subsequent PP stages (``pre_process=False``): + ``[S*K, B, D]`` (K folded into the sequence axis by the + previous stage's :func:`_lower_streams_out`) — we unfold. + pre_process: ``True`` on the first PP stage. + hc_mult: number of HC streams ``K``. + + Returns: + ``[B, S, K, D]`` if ``hc_mult > 1`` else ``[B, S, D]``. + """ + if hc_mult <= 1: + # Single-stream: just transpose [S, B, D] -> [B, S, D]. + return hidden_states.transpose(0, 1).contiguous() + + if pre_process: + # [S, B, D] -> [B, S, D] -> [B, S, K, D] (broadcast across K). + x = hidden_states.transpose(0, 1).contiguous() + B, S, D = x.shape + return x.unsqueeze(2).expand(B, S, hc_mult, D).contiguous() + + # Non-first stage: [S*K, B, D] -> [B, S, K, D]. + SK, B, D = hidden_states.shape + if SK % hc_mult != 0: + raise ValueError( + f"PP boundary tensor first-dim {SK} not divisible by hc_mult={hc_mult}; " + "previous stage did not pack K via _lower_streams_out." + ) + S = SK // hc_mult + # [S*K, B, D] -> [S, K, B, D] -> [B, S, K, D] + x = hidden_states.view(S, hc_mult, B, D).permute(2, 0, 1, 3).contiguous() + return x + + +def _lower_streams_out( + x: torch.Tensor, + *, + post_process: bool, + hc_mult: int, +) -> torch.Tensor: + """Reshape the block's output back to a P2P-compatible 3D tensor. + + Args: + x: ``[B, S, K, D]`` (multi-stream) or ``[B, S, D]`` (single). + On the final stage callers pass the post-HyperHead + ``[B, S, D]``; on non-final stages they pass the + still-multi-stream ``[B, S, K, D]``. + post_process: ``True`` on the final PP stage. + hc_mult: number of HC streams ``K``. + + Returns: + ``[S, B, D]`` on the final stage (matching Megatron's + sequence-first output convention) or ``[S*K, B, D]`` on + non-final stages (K folded into the sequence axis so PP P2P + kernels see a 3D tensor of the expected shape). + """ + if hc_mult <= 1: + return x.transpose(0, 1).contiguous() + + if post_process: + # x is [B, S, D] (already collapsed by HyperHead on this stage). + if x.dim() != 3: + raise ValueError( + f"_lower_streams_out: post_process expects [B, S, D] after HyperHead, " + f"got shape {tuple(x.shape)}." + ) + return x.transpose(0, 1).contiguous() + + # Non-final stage: pack [B, S, K, D] -> [S*K, B, D]. + if x.dim() != 4: + raise ValueError( + f"_lower_streams_out: non-final stage expects [B, S, K, D], " f"got shape {tuple(x.shape)}." + ) + B, S, K, D = x.shape + if K != hc_mult: + raise ValueError(f"_lower_streams_out: K dim of input ({K}) does not match hc_mult ({hc_mult}).") + # [B, S, K, D] -> [S, K, B, D] -> [S*K, B, D] + return x.permute(1, 2, 0, 3).contiguous().view(S * K, B, D) + + +# --------------------------------------------------------------------------- +# Per-layer attention factory +# --------------------------------------------------------------------------- + + +def _build_attention( + *, + compress_ratio: int, + rope: DualRoPE, + config: Optional[DeepSeekV4TransformerConfig] = None, +): + """No-spec fallback used when the layer is built without an + ``attention`` :class:`ModuleSpec`. + + Plan-2 P13: all three V4 layer types (``compress_ratio in {0, 4, 128}``) + construct through the single faithful :class:`DeepseekV4Attention` + class, which carries its own compressor / indexer (built locally + when no spec is provided). Production paths use + :func:`get_deepseek_v4_runtime_decoder_spec` which provides full + spec submodules; this fallback is for configs / unit tests that + construct a layer without an explicit attention spec. + """ + return DeepseekV4Attention( + config=config, + rope=rope, + compress_ratio=int(compress_ratio), + ) + + +# --------------------------------------------------------------------------- +# A single V4 block layer (attention sub-block + FFN sub-block, both wrapped +# by HyperMixer for the K-stream residual) +# --------------------------------------------------------------------------- + + +class DeepseekV4HybridLayer(TransformerLayer): + """One layer of the V4 decoder. + + Holds (using upstream-canonical submodule names): + + * :attr:`input_layernorm` — pre-attention RMSNorm. + * :attr:`self_attention` — V4 attention (Dense / HCA / CSA, picked + from ``compress_ratio``). + * :attr:`pre_mlp_layernorm` — pre-MLP RMSNorm. + * :attr:`mlp` — V4 MoE (when ``num_moe_experts > 0``) + or :class:`_DenseSwiGLUMLP` fallback for non-MoE layers. + * :attr:`attn_hc` / :attr:`ffn_hc` — :class:`HyperMixer` instances + per sub-block when ``hc_mult > 1`` (else ``None`` and the residual + collapses to a vanilla ``x + sub(x)`` add). + + Plan-2 P15 inherits from :class:`TransformerLayer` for type identity + (so Megatron's ``isinstance(layer, TransformerLayer)`` checks work + and ``BaseTransformerLayer`` -derived utilities apply). The parent + ``__init__`` is *not* called because V4's submodule set differs from + upstream (no cross-attention, no BDA, V4-specific attention + signature) — we initialize via :class:`MegatronModule` directly. + """ + + def __init__( + self, + config: DeepSeekV4TransformerConfig, + *, + layer_idx: int, + compress_ratio: int, + rope: Optional[DualRoPE] = None, + pg_collection=None, + submodules: Optional[DeepseekV4HybridLayerSubmodules] = None, + layer_number: Optional[int] = None, + is_mtp_layer: bool = False, + vp_stage: Optional[int] = None, + **_mtp_build_kwargs, + ) -> None: + # Bypass TransformerLayer.__init__ (it expects the upstream + # submodule contract — cross-attn, BDA, and a self_attention + # signature that takes layer_number / cp_comm_type). We build + # the V4 attribute set directly via MegatronModule. + # + # ``is_mtp_layer`` / ``vp_stage`` / ``**_mtp_build_kwargs`` are + # accepted (and ignored) so this layer can be built by the + # upstream :class:`MultiTokenPredictionLayer`, whose + # ``build_module(self.submodules.mtp_model_layer, config=..., + # vp_stage=..., layer_number=..., is_mtp_layer=True)`` call passes + # a richer kwarg set than the main decoder block does. + del is_mtp_layer, vp_stage, _mtp_build_kwargs + MegatronModule.__init__(self, config=config) + + # ``rope`` is normally injected by ``DeepseekV4TransformerBlock`` + # (one shared DualRoPE for the whole stack). When this layer is + # built standalone — e.g. as the inner layer of a V4 MTP depth via + # the upstream MTP ``build_module`` path, which does NOT thread + # ``rope`` — fall back to constructing a private DualRoPE from the + # config. DualRoPE holds only RoPE cos/sin caches (no trainable + # params), so a per-MTP-depth instance is numerically a no-op vs. + # sharing the trunk's. + if rope is None: + rope = DualRoPE( + rotary_dim=int(config.qk_pos_emb_head_dim), + rope_theta=float(config.rotary_base), + compress_rope_theta=float(config.compress_rope_theta), + yarn_factor=float(config.rotary_scaling_factor), + yarn_beta_fast=32.0, + yarn_beta_slow=1.0, + original_max_position_embeddings=int(config.original_max_position_embeddings), + ) + # Register the privately-built RoPE as a real submodule so its + # (non-persistent) cos/sin ``inv_freq`` buffers follow the layer + # onto the GPU under ``module.to(device)`` / ``.cuda()``. The main + # decoder path instead receives the block-owned ``self.rope`` (a + # registered submodule of the block) and keeps it in the + # attention's ``self._rope`` *list* to avoid double-registration; + # here there is no other owner, so this layer must register it. + self._fallback_rope = rope + + self.layer_idx = int(layer_idx) + self.compress_ratio = int(compress_ratio) + self.hc_mult = int(config.hc_mult) + # 1-based layer_number for Megatron's sharded_state_dict + recompute. + self.layer_number = int(layer_number) if layer_number is not None else (self.layer_idx + 1) + + hidden_size = int(config.hidden_size) + norm_eps = float(config.norm_epsilon) + hc_eps = float(config.hc_eps) + hc_sinkhorn_iters = int(config.hc_sinkhorn_iters) + # Plan-5 P29 (RESCOPED): forward the V4 config flag to HyperMixer + # so its sinkhorn_normalize dispatches to the torch.compile fast + # path when the run script flips USE_V4_COMPILED_SINKHORN=True. + # Default is False — see deepseek_v4_transformer_config.py for + # the dispatch contract. + use_v4_compiled_sinkhorn = bool(getattr(config, "use_v4_compiled_sinkhorn", False)) + + use_spec_submodules = submodules is not None + submodules = submodules or DeepseekV4HybridLayerSubmodules() + # Cache for the parent (TransformerLayer) sharded_state_dict / recompute paths. + self.submodules_config = submodules + + if use_spec_submodules and submodules.input_layernorm is not None: + self.input_layernorm = build_module( + submodules.input_layernorm, + config=config, + hidden_size=hidden_size, + eps=norm_eps, + ) + else: + self.input_layernorm = LocalRMSNorm(hidden_size, eps=norm_eps) + + if use_spec_submodules and submodules.self_attention is not None: + self.self_attention = build_module( + submodules.self_attention, + config=config, + rope=rope, + ) + else: + self.self_attention = _build_attention( + compress_ratio=self.compress_ratio, + rope=rope, + config=config, + ) + + if use_spec_submodules and submodules.pre_mlp_layernorm is not None: + self.pre_mlp_layernorm = build_module( + submodules.pre_mlp_layernorm, + config=config, + hidden_size=hidden_size, + eps=norm_eps, + ) + else: + self.pre_mlp_layernorm = LocalRMSNorm(hidden_size, eps=norm_eps) + + self.is_moe = int(config.num_moe_experts) > 0 + if use_spec_submodules and submodules.mlp is not None: + self.mlp = build_module( + submodules.mlp, + config=config, + pg_collection=pg_collection, + ) + self.is_moe = isinstance(self.mlp, DeepseekV4MoE) + elif self.is_moe: + moe_use_grouped_gemm = bool(config.moe_grouped_gemm) + moe_use_legacy_grouped_gemm = bool(config.moe_use_legacy_grouped_gemm) + + provider = resolve_v4_provider(config) + grouped_mlp_module, grouped_mlp_submodules = provider.v4_grouped_mlp_modules( + moe_use_grouped_gemm=moe_use_grouped_gemm, + moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm, + ) + assert ( + grouped_mlp_module is not None + ), "DeepSeek-V4 grouped MLP module must be provided by DeepSeekV4SpecProvider." + + shared_expert_spec = ModuleSpec( + module=SharedExpertMLP, + submodules=MLPSubmodules( + linear_fc1=provider.column_parallel_linear(), + linear_fc2=provider.row_parallel_linear(), + # P18 D2: pull from the V4-aware helper so the slot + # is None when ``use_te_activation_func`` is False + # (the default in V4 yamls; clamped-SwiGLU lives on + # the eager path via ``config.activation_func``). + activation_func=provider.v4_mlp_activation_func(), + ), + ) + moe_submodules = DeepseekV4MoESubmodules( + hash_router=ModuleSpec(module=DeepseekV4HashRouter), + learned_router=ModuleSpec(module=DeepseekV4LearnedRouter), + grouped_experts=ModuleSpec( + module=grouped_mlp_module, + submodules=grouped_mlp_submodules, + ), + shared_expert=shared_expert_spec, + ) + + self.mlp = DeepseekV4MoE( + config=config, + layer_idx=self.layer_idx, + pg_collection=pg_collection, + submodules=moe_submodules, + ) + else: + self.mlp = _DenseSwiGLUMLP( + config=config, + ) + + if self.hc_mult > 1: + if use_spec_submodules and submodules.attn_hc is not None: + self.attn_hc = build_module( + submodules.attn_hc, + hidden_size=hidden_size, + hc_mult=self.hc_mult, + eps=hc_eps, + sinkhorn_iters=hc_sinkhorn_iters, + use_compiled_sinkhorn=use_v4_compiled_sinkhorn, + ) + else: + self.attn_hc = HyperMixer( + hidden_size=hidden_size, + hc_mult=self.hc_mult, + eps=hc_eps, + sinkhorn_iters=hc_sinkhorn_iters, + use_compiled_sinkhorn=use_v4_compiled_sinkhorn, + ) + if use_spec_submodules and submodules.ffn_hc is not None: + self.ffn_hc = build_module( + submodules.ffn_hc, + hidden_size=hidden_size, + hc_mult=self.hc_mult, + eps=hc_eps, + sinkhorn_iters=hc_sinkhorn_iters, + use_compiled_sinkhorn=use_v4_compiled_sinkhorn, + ) + else: + self.ffn_hc = HyperMixer( + hidden_size=hidden_size, + hc_mult=self.hc_mult, + eps=hc_eps, + sinkhorn_iters=hc_sinkhorn_iters, + use_compiled_sinkhorn=use_v4_compiled_sinkhorn, + ) + else: + self.attn_hc = None + self.ffn_hc = None + + # ------------------------------------------------------------------ + + def _hc_apply(self, mixer: Optional[HyperMixer], x: torch.Tensor, sub_block, *args): + """Run a sub-block under HC. + + ``x`` shape: ``[B, S, K, D]`` if ``hc_mult > 1``, else ``[B, S, D]``. + ``sub_block`` is ``Callable[[Tensor, *Any], Tensor]`` whose first + positional arg is the (collapsed) hidden in ``[B, S, D]``. + """ + if mixer is None: + # Single-stream: classic residual; x already has shape [B, S, D]. + out = sub_block(x, *args) + return x + out + + pre, post, comb = mixer.compute_weights(x) # [..., K], [..., K], [..., K, K] + collapsed = HyperMixer.collapse(x, pre) # [B, S, D] + out = sub_block(collapsed, *args) # sub-block first positional = collapsed + return HyperMixer.expand(x, out, post, comb) # [B, S, K, D] + + # ------------------------------------------------------------------ + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + *, + position_ids: Optional[torch.Tensor] = None, + token_ids: Optional[torch.Tensor] = None, + **kwargs, + ): + """Run one V4 layer. + + Args: + hidden_states: ``[B, S, K, D]`` (multi-stream) or + ``[B, S, D]`` (single-stream). The block calls us with + the K-stream form when ``hc_mult > 1`` and with the + collapsed form otherwise. + attention_mask: ignored — V4 manages its own SWA / sink mask + inside :class:`DeepseekV4Attention`. Accepted for + upstream :class:`TransformerLayer` API compatibility. + position_ids: ``[B, S]`` or ``[S]``. Forwarded to attention. + token_ids: ``[B, S]`` integer tensor; required when this is + a hash-routed MoE layer + (``layer_idx < num_hash_layers``). Ignored for non-MoE / + non-hash layers. + **kwargs: ignored — accepted so this layer can be placed + inside upstream :class:`MultiTokenPredictionLayer`, + which forwards a richer kwargs set (rotary buffers, + inference params, etc.). + + Returns: + ``(hidden_states, context)`` where ``context`` is always + ``None``. The upstream :class:`TransformerLayer` returns a + tuple of this shape (cross-attention context pass-through); + V4 has no cross-attention, so the second element is always + ``None``. Returning a tuple keeps V4 layers compatible with + :class:`MultiTokenPredictionLayer._proj_and_transformer_layer` + which unpacks ``hidden_states, _ = self.mtp_model_layer(...)``. + """ + del attention_mask, kwargs + + if position_ids is None: + # Tiny CPU smokes / unit tests may omit position_ids; fall + # back to the seq-only arange. The block always provides + # them in production. + S = hidden_states.shape[1] + position_ids = torch.arange(S, device=hidden_states.device) + + # Attention sub-block. The collapse passes a [B, S, D] hidden, then + # the attention runs and returns [B, S, D]; HC expand writes back. + def _attn_sub(collapsed: torch.Tensor) -> torch.Tensor: + return self.self_attention(self.input_layernorm(collapsed), position_ids) + + x = self._hc_apply(self.attn_hc, hidden_states, _attn_sub) + + # MLP / MoE sub-block. MoE needs token_ids when the layer is + # hash-routed; plain SwiGLU ignores it. + if self.is_moe: + + def _ffn_sub(collapsed: torch.Tensor) -> torch.Tensor: + return self.mlp(self.pre_mlp_layernorm(collapsed), token_ids=token_ids) + + else: + + def _ffn_sub(collapsed: torch.Tensor) -> torch.Tensor: + return self.mlp(self.pre_mlp_layernorm(collapsed)) + + x = self._hc_apply(self.ffn_hc, x, _ffn_sub) + return x, None + + +# --------------------------------------------------------------------------- +# Top-level V4 transformer block +# --------------------------------------------------------------------------- + + +class DeepseekV4TransformerBlock(TransformerBlock): + """Multi-stream HC decoder for DeepSeek-V4. + + Plan-2 P15 subclasses Megatron's + :class:`megatron.core.transformer.transformer_block.TransformerBlock` + for type identity (so any upstream ``isinstance(block, + TransformerBlock)`` checks light up) and to inherit its + sharded-state-dict / debug surface. The parent ``__init__`` is + bypassed because V4's submodule contract differs (HyperHead on + post-process only, K-stream lift / lower at PP boundaries, no + upstream-style layer-norm impl override). We initialize via + :class:`MegatronModule` directly. + + PP K-stream packing: between stages we send ``[S*K, B, D]`` (K folded + into the sequence axis), letting standard 3D PP P2P kernels carry + the multi-stream tensor unchanged. The first stage lifts ``[S, B, D]`` + to ``[B, S, K, D]``; the final stage collapses with HyperHead and + transposes back to ``[S, B, D]``. + + Phase 6 update: + - respects PP / VP layer partitioning by constructing only local layers + for this pipeline rank; + - supports ``set_input_tensor`` so non-first PP stages consume P2P input. + """ + + def __init__( + self, + config: DeepSeekV4TransformerConfig, + spec=None, + post_layer_norm: bool = True, + pre_process: bool = True, + post_process: bool = True, + pg_collection=None, + vp_stage=None, + submodules: Optional[DeepseekV4TransformerBlockSubmodules] = None, + ) -> None: + # Bypass TransformerBlock.__init__: it requires a real + # pg_collection (or pulls one from parallel_state) and runs + # upstream-specific layer construction. V4's lift / lower path + # plus the spec provider give us equivalent functionality with + # CPU instantiability. + MegatronModule.__init__(self, config=config) + # Save arguments matching the parent's interface for compatibility. + self.spec = spec + self.submodules = submodules + self.post_layer_norm = post_layer_norm + self.pre_process = pre_process + self.post_process = post_process + self.vp_stage = vp_stage + self.pg_collection = pg_collection + # Required by pipeline schedules (same contract as TransformerBlock). + self.input_tensor = None + logger.info( + "[DeepSeek-V4] decoder block initialized (pre_process=%s post_process=%s).", + pre_process, + post_process, + ) + + # ---- shape / model fields ---- + hidden_size = config.hidden_size + rotary_dim = config.qk_pos_emb_head_dim + num_layers = config.num_layers + norm_eps = config.norm_epsilon + + # ---- V4-specific fields ---- + hc_mult = config.hc_mult + hc_eps = config.hc_eps + config.hc_sinkhorn_iters + compress_ratios = _normalize_compress_ratios( + config.compress_ratios, + num_layers=num_layers, + mtp_num_layers=int(config.mtp_num_layers), + ) + self.compress_ratios: List[int] = compress_ratios + + rope_theta = config.rotary_base + compress_rope_theta = config.compress_rope_theta + yarn_factor = config.rotary_scaling_factor + original_max_pos = config.original_max_position_embeddings + + self.num_hash_layers = int(config.num_hash_layers) + + # ---- shared dual-RoPE for the whole stack ---- + self.rope = DualRoPE( + rotary_dim=rotary_dim, + rope_theta=rope_theta, + compress_rope_theta=compress_rope_theta, + yarn_factor=yarn_factor, + yarn_beta_fast=32.0, + yarn_beta_slow=1.0, + original_max_position_embeddings=original_max_pos, + ) + + # ---- stage-local layer specs (always provided by runtime spec) ---- + provided_layer_specs = submodules.layer_specs if submodules is not None else None + assert provided_layer_specs, "DeepSeek-V4 requires non-empty submodules.layer_specs." + self.layers = nn.ModuleList() + self.global_layer_indices = [] + for local_idx, layer_spec in enumerate(provided_layer_specs): + layer = build_module( + layer_spec, + config=config, + pg_collection=pg_collection, + rope=self.rope, + ) + self.layers.append(layer) + self.global_layer_indices.append(int(getattr(layer, "layer_idx", local_idx))) + self.layer_offset = self.global_layer_indices[0] if self.global_layer_indices else 0 + self.hc_mult = hc_mult + + # Final HC collapse: built only on the post_process stage. + # Earlier PP stages forward the K-stream form via + # ``_lower_streams_out`` (no HyperHead per stage). + if hc_mult > 1 and self.post_process: + if submodules is not None and submodules.hyper_head is not None: + self.hyper_head = build_module( + submodules.hyper_head, + hidden_size=hidden_size, + hc_mult=hc_mult, + eps=hc_eps, + ) + else: + self.hyper_head = HyperHead(hidden_size=hidden_size, hc_mult=hc_mult, eps=hc_eps) + else: + self.hyper_head = None + + # Final RMSNorm placement follows Megatron semantics: + # - no MTP: on post_process stage + # - with MTP: on the stage containing decoder's final layer + if self._has_final_layernorm_in_this_stage(total_decoder_layers=num_layers): + if submodules is not None and submodules.final_layernorm is not None: + self.final_layernorm = build_module( + submodules.final_layernorm, + config=self.config, + hidden_size=hidden_size, + eps=norm_eps, + ) + else: + self.final_layernorm = LocalRMSNorm(hidden_size, eps=norm_eps) + else: + self.final_layernorm = None + + # ------------------------------------------------------------------ + + def _has_final_layernorm_in_this_stage(self, *, total_decoder_layers: int) -> bool: + if not self.post_layer_norm: + return False + + mtp_num_layers = self.config.mtp_num_layers + if mtp_num_layers is None: + return self.post_process + + if not self.global_layer_indices: + return False + return self.global_layer_indices[-1] == (total_decoder_layers - 1) + + def set_input_tensor(self, input_tensor: torch.Tensor): + """Pipeline-parallel hook: stash tensor from previous PP stage.""" + self.input_tensor = input_tensor + + @property + def num_layers_per_pipeline_rank(self) -> int: + """Compatibility shim used by upstream debug / recompute code.""" + return len(self.layers) + + # ------------------------------------------------------------------ + + def _recompute_local_layer_indices(self) -> Optional[set]: + """Local indices of layers to activation-checkpoint on this stage. + + Returns ``None`` when no recompute should be applied (the caller + then runs the plain forward loop). Mirrors Megatron semantics: + + * only active when ``recompute_granularity == 'full'`` and training; + * ``recompute_method == 'block'`` -> first ``recompute_num_layers`` + layers of this PP stage; + * ``recompute_method == 'uniform'`` -> every layer on this stage + (chunk size ``recompute_num_layers``; for full-layer checkpointing + the chunk boundary does not change which layers are recomputed); + * Primus ``recompute_layer_ids`` (explicit GLOBAL indices) takes + precedence when set. + """ + if not self.training: + return None + + cfg = self.config + if getattr(cfg, "recompute_granularity", None) != "full": + return None + + n_local = len(self.layers) + if n_local == 0: + return None + + recompute_layer_ids = getattr(cfg, "recompute_layer_ids", None) + if recompute_layer_ids: + wanted = {int(i) for i in recompute_layer_ids} + local = { + local_idx + for local_idx, global_idx in enumerate(self.global_layer_indices) + if int(global_idx) in wanted + } + return local or None + + num = int(getattr(cfg, "recompute_num_layers", 0) or 0) + if num <= 0: + return None + + method = getattr(cfg, "recompute_method", None) or "block" + if method == "block": + return set(range(min(num, n_local))) + if method == "uniform": + return set(range(n_local)) + raise ValueError(f"Invalid recompute_method for DeepSeek-V4: {method!r}") + + def _layer_fp8_context(self, global_idx: int): + """FP8 autocast for a single V4 layer (no-op when fp8 is off). + + DeepseekV4TransformerBlock.forward overrides Megatron's + TransformerBlock.forward with a custom layer loop, so the per-layer + ``get_fp8_context`` wrapping that Megatron's stock forward uses to + activate fp8 — and, via Primus's ``fp8_patches``, Primus-Turbo fp8 — + is not inherited and must be re-applied here. Without it the + ``primus_turbo_fp8_autocast`` context is never entered, so + ``PRIMUS_TURBO_FP8_ENABLED`` stays False and every GEMM silently runs + bf16 regardless of ``--fp8`` / ``--fp8_recipe``. + + ``get_fp8_context`` is resolved as a module attribute at call time so + the ``before_train`` fp8 patch (which rebinds + ``megatron.core.fp8_utils.get_fp8_context``) is honored. ``global_idx`` + is the 0-based global layer index (matches the function's ``layer_no`` + contract / ``is_first_last_bf16_layer`` gating). + """ + # FP4 (mxfp4) path: enter the Primus-Turbo fp4 autocast so is_turbo_fp4_enabled() + # is set for the layer (drives the grouped-MLP MXFP4 expert path and the + # fp4 turbo linears). fp4 and fp8 are mutually exclusive (global recipe). + if getattr(self.config, "fp4", None): + from megatron.core import fp4_utils + + return fp4_utils.get_fp4_context(self.config, global_idx) + if not self.config.fp8: + return nullcontext() + from megatron.core import fp8_utils + + return fp8_utils.get_fp8_context(self.config, global_idx) + + def _forward_layer_checkpointed(self, layer, x, position_ids, token_ids, global_idx): + """Run one V4 layer under activation checkpointing. + + Only the hidden-state tensor ``x`` is passed as the checkpointed + arg (``args[0]``); ``position_ids`` / ``token_ids`` are integer + tensors that need no grad and are captured by closure. The closure + returns just the hidden tensor (V4 layers return ``(hidden, None)``) + so the checkpoint backward sees a single grad-requiring output. + + The fp8 context is entered *inside* the checkpointed closure so it is + re-established on recompute (backward), keeping the recomputed forward + in fp8 to match the original pass. + """ + from megatron.core import tensor_parallel + + def _run(hidden): + with self._layer_fp8_context(global_idx): + out, _ = layer(hidden, position_ids=position_ids, token_ids=token_ids) + return out + + return tensor_parallel.checkpoint( + _run, + self.config.distribute_saved_activations, + x, + ) + + # ------------------------------------------------------------------ + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + inference_context=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + rotary_pos_cos_sin=None, + packed_seq_params=None, + sequence_len_offset=None, + position_ids: Optional[torch.Tensor] = None, + token_ids: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """Run the V4 decoder. + + Shape contract: + + * On the first PP stage (``pre_process=True``): + ``hidden_states`` arrives as ``[S, B, D]`` (sequence-first + embedding output). The lift helper expands it to + ``[B, S, K, D]`` for HC math. + * On subsequent PP stages (``pre_process=False``): the previous + stage packed K into the sequence axis, so + ``hidden_states`` arrives as ``[S*K, B, D]``. The lift helper + unfolds back to ``[B, S, K, D]``. + * On the final stage (``post_process=True``): HyperHead + collapses ``[B, S, K, D] -> [B, S, D]``, then the lower helper + transposes to the sequence-first ``[S, B, D]`` output. + * On non-final stages: the lower helper packs + ``[B, S, K, D] -> [S*K, B, D]`` so PP P2P kernels see a 3D + tensor of the expected rank. + + ``attention_mask`` and the various ``rotary_pos_*`` kwargs are + ignored — V4 manages its own dual-RoPE and SWA mask internally. + + ``position_ids`` is the caller-provided token-position tensor. + When omitted (e.g. unit tests), we fall back to ``arange(S)``; + production callers (:class:`DeepseekV4Model.forward`) always pass + it explicitly. + + ``token_ids`` (``[B, S]`` long tensor) is required only when one + or more layers on this stage run hash routing + (``layer_idx < num_hash_layers``). The legacy + ``decoder._v4_token_ids`` attribute stash has been removed; the + model forwards ``input_ids`` here directly. + """ + del ( + inference_context, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + rotary_pos_cos_sin, + packed_seq_params, + sequence_len_offset, + attention_mask, + kwargs, + ) + + if not self.pre_process: + hidden_states = self.input_tensor if self.input_tensor is not None else hidden_states + if hidden_states is None: + raise ValueError("DeepseekV4TransformerBlock.forward received no hidden_states tensor") + + needs_hash_token_ids = self.num_hash_layers > 0 and any( + layer_idx < self.num_hash_layers for layer_idx in self.global_layer_indices + ) + if needs_hash_token_ids and token_ids is None: + raise ValueError( + "token_ids is required on this PP stage because it owns hash-routed MoE layers " + f"(global layer idx < num_hash_layers={self.num_hash_layers})." + ) + + # Lift incoming P2P tensor to the K-stream form. + x = _lift_streams_in( + hidden_states, + pre_process=self.pre_process, + hc_mult=self.hc_mult, + ) + # x is [B, S, K, D] when hc_mult > 1, else [B, S, D]. + seq_len = x.shape[1] + + if position_ids is None: + position_ids = torch.arange(seq_len, device=x.device) + + # Run the layers. Each V4 layer returns ``(hidden, None)`` for + # upstream-tuple compatibility (see DeepseekV4HybridLayer.forward). + recompute_local = self._recompute_local_layer_indices() + for local_idx, layer in enumerate(self.layers): + global_idx = self.global_layer_indices[local_idx] + if recompute_local is not None and local_idx in recompute_local: + x = self._forward_layer_checkpointed(layer, x, position_ids, token_ids, global_idx) + else: + with self._layer_fp8_context(global_idx): + x, _ = layer( + x, + position_ids=position_ids, + token_ids=token_ids, + ) + + # Final HC collapse on post_process stage; non-final stages + # forward the multi-stream form through PP P2P. + if self.post_process and self.hc_mult > 1 and self.hyper_head is not None: + x = self.hyper_head(x) # [B, S, D] + + if self.final_layernorm is not None: + x = self.final_layernorm(x) + + if not self.pre_process and len(self.layers) == 0 and self.final_layernorm is None: + x = x.clone() + + # Lower to a P2P-compatible 3D tensor. + out = _lower_streams_out( + x, + post_process=self.post_process, + hc_mult=self.hc_mult, + ) + return make_viewless_tensor(inp=out, requires_grad=out.requires_grad, keep_graph=True) + + +__all__ = [ + "DeepseekV4HybridLayerSubmodules", + "DeepseekV4TransformerBlockSubmodules", + "DeepseekV4HybridLayer", + "DeepseekV4TransformerBlock", + "_lift_streams_in", + "_lower_streams_out", +] diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_builders.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_builders.py new file mode 100644 index 000000000..6041f755a --- /dev/null +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_builders.py @@ -0,0 +1,301 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 model builder + model_provider entry points. + +In the upstream Megatron-LM convention there are two pieces: + +* ``model_provider``: a thin wrapper that handles ``args.record_memory_history``, + ModelOpt etc. and delegates to a ``model_builder`` callable. Defined once + in ``Megatron-LM/model_provider.py``. +* ``_builders.py``: contains ``_builder(args, ...)`` -- + the actual model-class instantiation logic. + +For DeepSeek-V4 we keep both in a single primus-owned module so the dispatch +in ``primus/core/utils/import_utils.py`` doesn't have to chase symbols across +``third_party/Megatron-LM``. + +Phase 8 contract: +- Resolve V4 runtime decoder spec externally in builder. +- Pass runtime spec as ``transformer_layer_spec`` into ``DeepseekV4Model``. +- DeepseekV4Model is rooted at ``LanguageModule`` and has no GPT placeholder + spec dependency. +""" + +from typing import Optional + +from megatron.core.transformer.spec_utils import import_module +from megatron.training import get_args, print_rank_0 +from megatron.training.arguments import core_transformer_config_from_args + +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import ( + get_deepseek_v4_runtime_decoder_spec, + is_v4_turbo_deepep_active, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_model import ( + DeepseekV4Model, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, +) + + +def _resolve_runtime_decoder_spec( + args, + config: DeepSeekV4TransformerConfig, + vp_stage, +): + """Resolve effective runtime decoder spec for DeepSeek-V4 decoder path.""" + if args.spec is not None: + return import_module(args.spec) + return get_deepseek_v4_runtime_decoder_spec(config=config, vp_stage=vp_stage) + + +def _maybe_plumb_v4_sink_attention_args(args) -> None: + """Plan-3 P22: derive Turbo sink-attention args from V4 config. + + :class:`PrimusTurboAttention` reads ``use_sink_attention``, + ``sink_sliding_window`` and ``sink_window_even_layers_only`` directly + from the global ``args`` namespace at module-init time. For V4 we + can derive all three from the V4 attention configuration: + + * ``use_sink_attention`` follows ``args.attn_sink`` (V4's per-head + learnable softmax sink — set to ``true`` in the V4-Flash recipe). + * ``sink_sliding_window`` is derived from ``args.attn_sliding_window``, + *capped to seq_length*. V4's dense layers use SWA with window 128 + in the released checkpoint; the released aiter Triton flash-attn + backend (``aiter/ops/triton/attention/mha.py``) does not yet + accept ``window_size != (-1, -1)`` and raises ``ValueError: + Sliding Window is not supported yet in the Triton Backend``. + We therefore zero out the window when it equals or exceeds + ``seq_length`` (a window that covers the full causal triangle is + mathematically equivalent to no window). When the window is + strictly shorter than the sequence length we *still* zero it out + and emit a warning — Turbo can't honor it, and falling back to + eager-Python every step would defeat the perf goal. Long-context + V4-Flash (seq > 128) needs an aiter SWA upgrade before Turbo can + claim full V4 fidelity; tracked as a P22 follow-up. + * ``sink_window_even_layers_only`` is hard-set to ``False`` because + V4 applies SWA on every dense layer (unlike gpt-oss, which only + windows even-numbered layers). + + The plumbing only fires when all of (a) Turbo is enabled, (b) + ``use_turbo_attention=True``, (c) the user has not set the sink + fields explicitly. This keeps the V4-Flash YAML free of + Turbo-internal knobs while still producing a correct attention + forward when the user flips ``use_turbo_attention=true``. + """ + if not getattr(args, "enable_primus_turbo", False): + return + if not getattr(args, "use_turbo_attention", False): + return + if not getattr(args, "attn_sink", False): + # Without the V4 sink, Turbo would not honor SWA either; defer + # to ``DeepseekV4Attention``'s eager-Python fallback for this + # configuration (it asserts ``_use_core_attention=False`` when + # ``attn_sink`` is off and ``attn_sliding_window > 0``). + return + + if getattr(args, "use_sink_attention", None) in (None, False): + args.use_sink_attention = True + print_rank_0( + "[Primus:DeepSeek-V4][P22] derived args.use_sink_attention=True " + "from args.attn_sink=True (Turbo flash-attn sink path)." + ) + + if getattr(args, "sink_sliding_window", None) in (None, 0): + attn_sw = int(getattr(args, "attn_sliding_window", 0) or 0) + seq_len = int(getattr(args, "seq_length", 0) or 0) + if attn_sw <= 0: + args.sink_sliding_window = 0 + elif seq_len > 0 and attn_sw >= seq_len: + # Window covers the full causal triangle — equivalent to no + # window. Drop it so Turbo's flash kernel runs without + # window_size (avoids the aiter Triton SWA gap). + args.sink_sliding_window = 0 + print_rank_0( + f"[Primus:DeepSeek-V4][P22] attn_sliding_window={attn_sw} " + f">= seq_length={seq_len}; dropping window for Turbo " + "(mathematically equivalent to full causal attention)." + ) + else: + args.sink_sliding_window = 0 + print_rank_0( + f"[Primus:DeepSeek-V4][P22] WARNING: " + f"attn_sliding_window={attn_sw} < seq_length={seq_len} " + "but aiter Triton flash-attn does not support sliding " + "window yet (raises ValueError). Setting " + "sink_sliding_window=0 — V4 dense layers will attend " + "to *all* causal-prior tokens instead of the windowed " + "subset. This deviates from V4-Flash math; for " + "checkpoint-fidelity training, keep " + "use_turbo_attention=False until aiter adds SWA support." + ) + + # gpt-oss windows only even layers; V4 windows all dense layers. + args.sink_window_even_layers_only = False + + +def _maybe_plumb_v4_turbo_deepep_args(args) -> None: + """Plan-3 P23: derive Turbo DeepEP MoE args from V4 config. + + When :func:`is_v4_turbo_deepep_active` returns True, the V4 MoE + layers will be built with :class:`PrimusTurboDeepEPTokenDispatcher` + (resolved by ``_pick_v4_dispatcher_cls`` in + ``deepseek_v4_layer_specs.py``). Two ``args`` fields control + downstream behaviour: + + * ``moe_enable_deepep`` — asserted by + :class:`PrimusTurboDeepEPTokenDispatcher.__init__`; without it + the Turbo dispatcher refuses to build. + * ``moe_token_dispatcher_type`` — read by + :class:`DeepseekV4TransformerConfig` to populate the V4 config's + ``moe_token_dispatcher_type``. V4 spec build's + ``_pick_v4_dispatcher_cls`` then sees ``"flex"`` and pairs the + Turbo class with the right dispatcher-type label so + ``DeepseekV4MoE._resolve_dispatcher_type_from_spec`` reports + ``"flex"`` (instead of falling through to the + ``"unsupported dispatcher module"`` warning + ``"alltoall"`` + fallback). + + The ``before_train`` Turbo MoE patch + (``primus.backends.megatron.patches.turbo.moe_dispatcher_patches``) + sets the same two fields, but it fires AFTER ``deepseek_v4_builder`` + has already built ``config`` from ``args``, so V4 needs to plumb + these values BEFORE ``core_transformer_config_from_args`` runs. + + The plumbing is a no-op when: + + * Turbo is not enabled / DeepEP not requested / TP > 1, OR + * the user has already set ``args.moe_token_dispatcher_type`` to + something other than ``"alltoall"`` (we only override the V4 + base.yaml default, never an explicit user opt-in like + ``"allgather"``). + """ + if not is_v4_turbo_deepep_active(args): + return + + if not bool(getattr(args, "moe_enable_deepep", False)): + args.moe_enable_deepep = True + print_rank_0( + "[Primus:DeepSeek-V4][P23] derived args.moe_enable_deepep=True " + "from args.use_turbo_deepep=True (Turbo DeepEP dispatcher path)." + ) + + current = str(getattr(args, "moe_token_dispatcher_type", "") or "").lower() + # Only override when the V4 base default ("alltoall") is in effect + # or the user already opted into "flex". Respect explicit + # "allgather" — that is a different parallelism scheme (tp+ep + # all-gather/scatter), not a deepep flavour. + if current in ("", "alltoall", "flex"): + if current != "flex": + args.moe_token_dispatcher_type = "flex" + print_rank_0( + f"[Primus:DeepSeek-V4][P23] override " + f"args.moe_token_dispatcher_type='{current or 'unset'}' -> " + f"'flex' (Turbo DeepEP dispatcher path)." + ) + else: + print_rank_0( + "[Primus:DeepSeek-V4][P23] WARNING: " + f"args.use_turbo_deepep=True but moe_token_dispatcher_type=" + f"'{current}' (not 'flex' / 'alltoall'). Keeping the user's " + "dispatcher type; the Turbo DeepEP path will not engage for V4 " + "MoE layers." + ) + + +def deepseek_v4_builder( + args, + pre_process, + post_process, + vp_stage=None, + config: Optional[DeepSeekV4TransformerConfig] = None, + pg_collection=None, +): + """Build a DeepSeek-V4 model. + + Phase 8: build from a DeepSeek runtime spec tree only. + """ + print_rank_0("[Primus:DeepSeek-V4] building DeepseekV4Model...") + + # Plan-3 P22: plumb V4 attn_sink + attn_sliding_window into the + # Turbo flash-attn sink-attention args before any V4 attention + # module is constructed (PrimusTurboAttention reads these in + # ``__init__`` from get_args()). + _maybe_plumb_v4_sink_attention_args(args) + # Plan-3 P23: plumb args.moe_enable_deepep + moe_token_dispatcher_type + # for the Turbo DeepEP MoE dispatcher path BEFORE config construction + # so the resulting ``config`` carries the right dispatcher type into + # ``_pick_v4_dispatcher_cls``. + _maybe_plumb_v4_turbo_deepep_args(args) + + if config is None: + config = core_transformer_config_from_args( + args, + config_class=DeepSeekV4TransformerConfig, + ) + + assert not args.use_legacy_models, "DeepSeek-V4 requires use_legacy_models=False (Mcore-only)." + + runtime_decoder_spec = _resolve_runtime_decoder_spec(args, config, vp_stage) + + model = DeepseekV4Model( + config=config, + transformer_layer_spec=runtime_decoder_spec, + vocab_size=args.padded_vocab_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=True, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type=args.position_embedding_type, + rotary_percent=args.rotary_percent, + rotary_base=args.rotary_base, + rope_scaling=args.use_rope_scaling, + pg_collection=pg_collection, + vp_stage=vp_stage, + ) + return model + + +def model_provider( + model_builder=None, + pre_process: bool = True, + post_process: bool = True, + vp_stage: Optional[int] = None, + config: Optional[DeepSeekV4TransformerConfig] = None, + pg_collection=None, +): + """``model_provider`` entry point used by Megatron's ``pretrain()``. + + ``MegatronPretrainTrainer`` will pass ``deepseek_v4_builder`` as the + first arg via ``functools.partial`` so the upstream ``pretrain()`` can + call this with the standard ``(pre_process, post_process, vp_stage)`` + signature. + """ + if model_builder is None: + model_builder = deepseek_v4_builder + + args = get_args() + if args.record_memory_history: + import torch + + torch.cuda.memory._record_memory_history( + True, + trace_alloc_max_entries=100000, + trace_alloc_record_context=True, + ) + + return model_builder( + args, + pre_process, + post_process, + vp_stage, + config=config, + pg_collection=pg_collection, + ) diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py new file mode 100644 index 000000000..07a9b0013 --- /dev/null +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py @@ -0,0 +1,697 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +DeepSeek-V4 spec entry points. + +This module only defines DeepSeek-native runtime specs. +""" + +import importlib +import importlib.util +import logging +import os +from typing import List, Optional, Tuple + +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.mlp import MLPSubmodules +from megatron.core.transformer.moe.shared_experts import SharedExpertMLP +from megatron.core.transformer.moe.token_dispatcher import ( + MoEAllGatherTokenDispatcher, + MoEAlltoAllTokenDispatcher, + MoEFlexTokenDispatcher, +) +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_block import get_num_layers_to_build +from megatron.core.transformer.transformer_layer import get_transformer_layer_offset + +from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import ( + DeepSeekV4SpecProvider, +) +from primus.backends.megatron.core.models.deepseek_v4.build_context import ( + resolve_v4_provider, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block import ( + DeepseekV4HybridLayer, + DeepseekV4HybridLayerSubmodules, + DeepseekV4TransformerBlock, + DeepseekV4TransformerBlockSubmodules, + _DenseSwiGLUMLP, + _normalize_compress_ratios, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, +) +from primus.backends.megatron.core.transformer.compressor import Compressor +from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( + DeepseekV4Attention, + DeepseekV4AttentionSubmodules, +) +from primus.backends.megatron.core.transformer.hyper_connection import ( + HyperHead, + HyperMixer, +) +from primus.backends.megatron.core.transformer.indexer import Indexer +from primus.backends.megatron.core.transformer.moe.v4_hash_router import ( + DeepseekV4HashRouter, +) +from primus.backends.megatron.core.transformer.moe.v4_moe import ( + DeepseekV4MoE, + DeepseekV4MoESubmodules, +) +from primus.backends.megatron.core.transformer.moe.v4_topk_router import ( + DeepseekV4LearnedRouter, +) + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Plan-3 P23 — Turbo DeepEP dispatcher gating +# --------------------------------------------------------------------------- +# +# Both ``deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args`` (which +# fires BEFORE config construction so the V4 config inherits the right +# ``moe_token_dispatcher_type`` / ``moe_enable_deepep``) and +# ``_pick_v4_dispatcher_cls`` below need the same gating predicate. +# Centralised here so the two stay in sync; mirrors +# ``primus.backends.megatron.patches.turbo.moe_dispatcher_patches._is_turbo_deepep_enabled`` +# (the patch itself fires too late for V4 spec build, but its gating +# is the canonical one). +# +# ``deepseek_v4_builders.py`` imports this helper, which is one-way +# (builders → layer_specs); putting the helper in builders would +# create a circular import because layer_specs would need to call +# back. + +# Public name so unit tests can patch / probe it directly. +PRIMUS_TURBO_DEEPEP_DISPATCHER_NAME = "PrimusTurboDeepEPTokenDispatcher" + + +def is_v4_turbo_deepep_active(args) -> bool: + """Plan-3 P23: V4-side gate for the Turbo DeepEP MoE dispatcher. + + Returns ``True`` only when **all** of the following hold (matches + the conditions enforced by the upstream + ``megatron.turbo.moe_dispatcher`` patch): + + * ``primus_turbo`` package is importable; + * ``args.enable_primus_turbo`` is True; + * ``args.use_turbo_deepep`` is True; + * ``args.tensor_model_parallel_size == 1`` + (PrimusTurboDeepEPTokenDispatcher requires TPxEP > 1; we keep + TP == 1 here because TP > 1 paths haven't been validated against + the Turbo dispatcher and the existing patch enforces the same + gate). + """ + if importlib.util.find_spec("primus_turbo") is None: + return False + if not bool(getattr(args, "enable_primus_turbo", False)): + return False + if not bool(getattr(args, "use_turbo_deepep", False)): + return False + tp_size = int(getattr(args, "tensor_model_parallel_size", 1) or 1) + if tp_size != 1: + return False + return True + + +def _import_primus_turbo_deepep_dispatcher_cls(): + """Plan-3 P23: lazy import of PrimusTurboDeepEPTokenDispatcher. + + Returns ``None`` when the import fails (callers must fall back to + the upstream MoEFlexTokenDispatcher in that case + emit a warning). + The import is lazy because the wider primus build is allowed to + run without ``primus_turbo`` installed (e.g. CPU unit tests for + non-Turbo paths). + """ + if importlib.util.find_spec("primus_turbo") is None: + return None + try: + module = importlib.import_module("primus.backends.megatron.core.extensions.primus_turbo") + except Exception as exc: # pragma: no cover - defensive + logger.warning( + "[DeepSeek-V4][P23] Failed to import " + "primus.backends.megatron.core.extensions.primus_turbo: %s", + exc, + ) + return None + return getattr(module, PRIMUS_TURBO_DEEPEP_DISPATCHER_NAME, None) + + +def _pick_v4_dispatcher_cls( + config: DeepSeekV4TransformerConfig, + *, + args=None, +) -> Tuple[type, str]: + """Plan-3 P23: pick the MoE token-dispatcher class for V4 spec build. + + Returns ``(dispatcher_cls, dispatcher_type)``. The + ``dispatcher_type`` is the string label V4 stores in its config / + annotation (used by ``DeepseekV4MoE._resolve_dispatcher_type_from_spec``); + ``dispatcher_cls`` is the class that ``ModuleSpec`` will hand to + ``build_module``. + + Selection order: + + 1. If ``config.moe_token_dispatcher_type == "allgather"`` → + :class:`MoEAllGatherTokenDispatcher` / ``"allgather"``. This is + a different parallelism scheme (TPxEP all-gather/scatter), not + a deepep flavour; never overridden by the Turbo path. + 2. If ``config.moe_token_dispatcher_type == "flex"`` AND + :func:`is_v4_turbo_deepep_active` returns True → + :class:`PrimusTurboDeepEPTokenDispatcher` / ``"flex"``. When + the package isn't importable we log a one-shot warning and + fall back to the upstream :class:`MoEFlexTokenDispatcher`. + 3. If ``config.moe_token_dispatcher_type == "flex"`` (Turbo + inactive) → :class:`MoEFlexTokenDispatcher` / ``"flex"``. + 4. Anything else (default V4 base.yaml: ``"alltoall"``) → + :class:`MoEAlltoAllTokenDispatcher` / ``"alltoall"``. Unknown + values emit a one-shot warning and fall through to alltoall. + + The ``args=`` keyword is for unit tests; production callers pass + ``None`` and the helper consults ``megatron.training.get_args()`` + lazily. ``args=None`` with no Megatron args available is + interpreted as "Turbo not active" (i.e. take the non-turbo + branch). This keeps ``_build_ffn_spec`` callable from CPU unit + tests that haven't called ``initialize_megatron``. + """ + dispatcher_type = str(getattr(config, "moe_token_dispatcher_type", None) or "").lower() + + if dispatcher_type == "allgather": + return MoEAllGatherTokenDispatcher, "allgather" + + if dispatcher_type == "flex": + if args is None: + try: + from megatron.training import get_args as _megatron_get_args + + args = _megatron_get_args() + except Exception: + args = None + if args is not None and is_v4_turbo_deepep_active(args): + turbo_cls = _import_primus_turbo_deepep_dispatcher_cls() + if turbo_cls is not None: + logger.info( + "[DeepSeek-V4][P23] MoE dispatcher class resolved to " + "PrimusTurboDeepEPTokenDispatcher (Turbo DeepEP path)." + ) + return turbo_cls, "flex" + logger.warning( + "[DeepSeek-V4][P23] use_turbo_deepep=True but " + "PrimusTurboDeepEPTokenDispatcher is not importable; " + "falling back to MoEFlexTokenDispatcher." + ) + return MoEFlexTokenDispatcher, "flex" + + if dispatcher_type not in ("", "alltoall"): + logger.warning( + "[DeepSeek-V4] unsupported moe_token_dispatcher_type=%s; fallback to alltoall.", + dispatcher_type, + ) + return MoEAlltoAllTokenDispatcher, "alltoall" + + +def _default_init_method(_weight) -> None: + return None + + +def _v4_fp8_attn_proj(config: "DeepSeekV4TransformerConfig") -> bool: + """FP8-ify the attention projections (q-up / o-proj) that otherwise fall + back to bf16 because the TE/Turbo fp8 linear rejects gather_output / + scatter-input. Only safe at TP=1, where gather/scatter are no-ops — so + we keep the bf16 gather/scatter native path for TP>1. Opt-in via + PRIMUS_V4_FP8_ATTN_PROJ=1. + """ + return ( + os.environ.get("PRIMUS_V4_FP8_ATTN_PROJ", "0") == "1" + and getattr(config, "tensor_model_parallel_size", 1) == 1 + ) + + +def _build_linear_projection_spec( + *, + config: DeepSeekV4TransformerConfig, + provider: DeepSeekV4SpecProvider, + in_features: int, + out_features: int, +) -> ModuleSpec: + """Default projection spec — a duplicated TE linear (no TP sharding). + + Used for projections that V4's grouped-low-rank O does not natively + shard along TP (``linear_q_down_proj``, ``linear_kv``, ``linear_o_a``). + Keep these duplicated for now; full TP sharding of the grouped O + projection is tracked in P14. + """ + return ModuleSpec( + module=provider.linear(), + params={ + "input_size": in_features, + "output_size": out_features, + "parallel_mode": "duplicated", + "config": config, + "init_method": config.init_method or _default_init_method, + "bias": False, + "skip_bias_add": False, + "skip_weight_param_allocation": False, + "tp_comm_buffer_name": None, + "is_expert": False, + }, + ) + + +def _build_column_parallel_spec( + *, + config: DeepSeekV4TransformerConfig, + provider: DeepSeekV4SpecProvider, + in_features: int, + out_features: int, + gather_output: bool = True, +) -> ModuleSpec: + """Column-parallel projection spec. + + With ``gather_output=True`` the output dim is gathered back to full + width across TP ranks, so downstream attention math (which assumes + ``H * head_dim`` per rank) does not need to know about TP at all. + Memory of the projection's weight matrix is sharded across TP ranks + even at ``gather_output=True``. + + Plan-3 P21: TE / Turbo column-parallel wrappers explicitly reject + ``gather_output=True`` (see + ``third_party/Megatron-LM/megatron/core/extensions/transformer_engine.py:747/972``). + When the caller asks for the gather variant we route to the + upstream Megatron-native :class:`ColumnParallelLinear` via + ``provider.column_parallel_linear_with_gather_output()``; the + standard TE path stays for ``gather_output=False``. + + Plan-2 P13 follow-up: this is used for ``linear_q_up_proj``. The + gather-then-shard variant for full sharded heads is tracked in P14 + once the grouped-O TP plan lands. + """ + # FP8 attention projections (paper recipe): the TE/Turbo fp8 column linear + # rejects gather_output=True, so q-up normally falls back to bf16 native. + # But at TP=1 the gather is a no-op, so we can route q-up through the fp8 + # turbo linear (gather_output=False ≡ True) to capture it in mxfp8. + # Gated by PRIMUS_V4_FP8_ATTN_PROJ=1 and only when TP==1. Default off. + if gather_output and _v4_fp8_attn_proj(config): + module_cls = provider.column_parallel_linear() + gather_output = False + elif gather_output: + module_cls = provider.column_parallel_linear_with_gather_output() + else: + module_cls = provider.column_parallel_linear() + return ModuleSpec( + module=module_cls, + params={ + "input_size": in_features, + "output_size": out_features, + "config": config, + "init_method": config.init_method or _default_init_method, + "gather_output": gather_output, + "bias": False, + "skip_bias_add": False, + "skip_weight_param_allocation": False, + "tp_comm_buffer_name": None, + "is_expert": False, + }, + ) + + +def _build_row_parallel_spec( + *, + config: DeepSeekV4TransformerConfig, + provider: DeepSeekV4SpecProvider, + in_features: int, + out_features: int, + input_is_parallel: bool = False, +) -> ModuleSpec: + """Row-parallel projection spec. + + With ``input_is_parallel=False`` the linear scatters the input across + TP ranks internally and all-reduces the output, so the caller can + pass a full-width input tensor and get a full-width output tensor. + Weight memory is sharded across TP ranks. Used for ``linear_o_b`` + and the flat-O fallback ``linear_proj``. + + Plan-3 P21: TE / Turbo row-parallel wrappers explicitly reject + ``input_is_parallel=False`` (see + ``third_party/Megatron-LM/megatron/core/extensions/transformer_engine.py:1081``). + When the caller asks for scatter-input we route to the upstream + Megatron-native :class:`RowParallelLinear` via + ``provider.row_parallel_linear_with_scatter_input()``; the + standard TE path stays for ``input_is_parallel=True``. + """ + # FP8 attention projections: the TE/Turbo fp8 row linear rejects + # input_is_parallel=False, so o-proj normally falls back to bf16 native. + # At TP=1 the scatter is a no-op, so route through the fp8 turbo linear + # (input_is_parallel=True ≡ False). Gated by PRIMUS_V4_FP8_ATTN_PROJ + TP==1. + if not input_is_parallel and _v4_fp8_attn_proj(config): + module_cls = provider.row_parallel_linear() + input_is_parallel = True + elif not input_is_parallel: + module_cls = provider.row_parallel_linear_with_scatter_input() + else: + module_cls = provider.row_parallel_linear() + return ModuleSpec( + module=module_cls, + params={ + "input_size": in_features, + "output_size": out_features, + "config": config, + "init_method": config.init_method or _default_init_method, + "input_is_parallel": input_is_parallel, + "bias": False, + "skip_bias_add": False, + "tp_comm_buffer_name": None, + "is_expert": False, + }, + ) + + +def _build_v4_attention_submodules( + *, + config: DeepSeekV4TransformerConfig, + provider: DeepSeekV4SpecProvider, + compress_ratio: int, +) -> DeepseekV4AttentionSubmodules: + """V4-canonical submodules for :class:`DeepseekV4Attention`. + + Field names match the released V4-Flash checkpoint layout (and MLA's + canonical names where they overlap): + + * ``linear_q_down_proj`` : ``hidden -> q_lora_rank`` (= ``wq_a``) + * ``q_layernorm`` : RMSNorm(``q_lora_rank``) (= ``q_norm``) + * ``linear_q_up_proj`` : ``q_lora_rank -> n_heads * head_dim`` (= ``wq_b``) + — built as **column-parallel** so the projection's weight is + sharded across TP at ``tp > 1``. ``gather_output=True`` keeps + downstream math TP-agnostic. + * ``linear_kv`` : ``hidden -> head_dim`` (= ``wkv``, + single-latent KV) + * ``kv_layernorm`` : RMSNorm(``head_dim``) (= ``kv_norm``) + * ``linear_o_a`` : grouped low-rank O down-proj (duplicated; + grouped-O TP plan is P14). + * ``linear_o_b`` : grouped low-rank O up-proj (-> ``hidden``) + — built as **row-parallel** so its weight is sharded across TP. + * ``linear_proj`` : flat-O fallback (``o_lora_rank == 0``, + e.g. unit tests) — also row-parallel. + * ``compressor`` : :class:`Compressor` (compressed branches) + * ``indexer`` : :class:`Indexer` (CSA branch only) + + The per-head learnable softmax sink lives directly on the attention + module as ``self.attn_sink: nn.Parameter`` — there is no separate + submodule slot (Plan-3 P21 dropped the ``attn_sink`` field; the + inline softmax-with-sink path in ``_attention_forward`` is canonical). + """ + hidden_size = int(config.hidden_size) + num_heads = int(config.num_attention_heads) + head_dim = int(config.kv_channels or (hidden_size // num_heads)) + q_lora_rank = int(config.q_lora_rank or 0) + o_groups = max(int(getattr(config, "o_groups", 1) or 1), 1) + o_lora_rank = int(getattr(config, "o_lora_rank", 0) or 0) + + if q_lora_rank <= 0: + raise ValueError( + "DeepSeek-V4 requires q_lora_rank > 0; the released checkpoint " + "always low-rank-projects Q via wq_a / wq_b." + ) + + q_out = num_heads * head_dim + submods = DeepseekV4AttentionSubmodules( + linear_q_down_proj=_build_linear_projection_spec( + config=config, + provider=provider, + in_features=hidden_size, + out_features=q_lora_rank, + ), + linear_q_up_proj=_build_column_parallel_spec( + config=config, + provider=provider, + in_features=q_lora_rank, + out_features=q_out, + ), + linear_kv=_build_linear_projection_spec( + config=config, + provider=provider, + in_features=hidden_size, + out_features=head_dim, # single-latent: K = V = wkv(hidden) + ), + q_layernorm=ModuleSpec(module=provider.v4_q_layernorm()), + kv_layernorm=ModuleSpec(module=provider.v4_kv_layernorm()), + ) + + if o_lora_rank > 0: + n_per_group = q_out // o_groups + submods.linear_o_a = _build_linear_projection_spec( + config=config, + provider=provider, + in_features=n_per_group, + out_features=o_groups * o_lora_rank, + ) + submods.linear_o_b = _build_row_parallel_spec( + config=config, + provider=provider, + in_features=o_groups * o_lora_rank, + out_features=hidden_size, + ) + else: + submods.linear_proj = _build_row_parallel_spec( + config=config, + provider=provider, + in_features=q_out, + out_features=hidden_size, + ) + + if compress_ratio > 0: + submods.compressor = ModuleSpec(module=Compressor) + if compress_ratio == 4: + submods.indexer = ModuleSpec(module=Indexer) + else: + # Plan-3 P22: dense layers route their softmax-and-attend through + # provider.core_attention() (PrimusTurboAttention when + # ``use_turbo_attention=True``, TEDotProductAttention otherwise). + # HCA + CSA layers do not get this slot — see + # ``DeepseekV4AttentionSubmodules`` docstring for why. + submods.core_attention = ModuleSpec(module=provider.core_attention()) + + return submods + + +def _build_norm_spec( + *, + config: DeepSeekV4TransformerConfig, + provider: DeepSeekV4SpecProvider, +): + del config + norm_module = provider.v4_norm_module() + assert norm_module is not None, "DeepSeek-V4 norm module must be provided by DeepSeekV4SpecProvider." + return ModuleSpec(module=norm_module) + + +def _build_attention_spec( + *, + compress_ratio: int, + config: DeepSeekV4TransformerConfig, + provider: DeepSeekV4SpecProvider, +) -> ModuleSpec: + """Plan-2 P13 attention spec — single :class:`DeepseekV4Attention` + class for all three V4 layer types (dense / HCA / CSA), dispatched + inside the class on ``compress_ratio``. + + Plan-2 P16: ``attn_mask_type=AttnMaskType.causal`` is declared on the + spec params for upstream :class:`MultiTokenPredictionLayer` + compatibility. The value is functionally inert for V4 (the V4 + attention forward manages its own SWA / sink mask internally) but + the upstream MTP layer's pre-build validator requires the field to + be one of ``{padding, causal, no_mask, padding_causal}`` when the + inner layer's submodules are :class:`TransformerLayerSubmodules` — + which they are, since :class:`DeepseekV4HybridLayerSubmodules` now + extends that dataclass. + """ + return ModuleSpec( + module=DeepseekV4Attention, + params={ + "compress_ratio": int(compress_ratio), + "attn_mask_type": AttnMaskType.causal, + }, + submodules=_build_v4_attention_submodules( + config=config, + provider=provider, + compress_ratio=int(compress_ratio), + ), + ) + + +def _build_ffn_spec( + *, + config: DeepSeekV4TransformerConfig, + provider: DeepSeekV4SpecProvider, + layer_idx: int, +) -> ModuleSpec: + num_routed_experts = int(config.num_moe_experts) + moe_use_grouped_gemm = bool(config.moe_grouped_gemm) + moe_use_legacy_grouped_gemm = bool(config.moe_use_legacy_grouped_gemm) + grouped_mlp_module, grouped_mlp_submodules = provider.v4_grouped_mlp_modules( + moe_use_grouped_gemm=moe_use_grouped_gemm, + moe_use_legacy_grouped_gemm=moe_use_legacy_grouped_gemm, + ) + # Plan-3 P23: V4-side dispatcher selection. Avoids the module-attr + # timing race where ``before_train`` patches the upstream + # ``MoEFlexTokenDispatcher`` symbol AFTER V4 spec build has captured + # it; by resolving the class locally we always get the right one. + dispatcher_cls, dispatcher_type = _pick_v4_dispatcher_cls(config) + + assert ( + grouped_mlp_module is not None + ), "DeepSeek-V4 grouped MLP module must be provided by DeepSeekV4SpecProvider." + + grouped_experts_spec = ModuleSpec( + module=grouped_mlp_module, + submodules=grouped_mlp_submodules, + ) + + shared_expert_submodules = MLPSubmodules( + linear_fc1=provider.column_parallel_linear(), + linear_fc2=provider.row_parallel_linear(), + # P18 D2: V4-aware activation_func selection — see + # ``DeepSeekV4SpecProvider.v4_mlp_activation_func`` for why + # this is None on the eager (clamped-SwiGLU) path. + activation_func=provider.v4_mlp_activation_func(), + ) + shared_expert_spec = ModuleSpec( + module=SharedExpertMLP, + submodules=shared_expert_submodules, + ) + + moe_submodules = DeepseekV4MoESubmodules( + hash_router=ModuleSpec(module=DeepseekV4HashRouter), + learned_router=ModuleSpec(module=DeepseekV4LearnedRouter), + token_dispatcher=ModuleSpec(module=dispatcher_cls), + grouped_experts=grouped_experts_spec, + shared_expert=shared_expert_spec, + ) + + if num_routed_experts > 0: + return ModuleSpec( + module=DeepseekV4MoE, + params={ + "layer_idx": layer_idx, + }, + submodules=moe_submodules, + ) + return ModuleSpec(module=_DenseSwiGLUMLP) + + +def _build_hybrid_layer_spec( + config: DeepSeekV4TransformerConfig, + *, + provider: DeepSeekV4SpecProvider, + layer_idx: int, + compress_ratio: int, +) -> ModuleSpec: + hc_mult = int(config.hc_mult) + + layer_submodules = DeepseekV4HybridLayerSubmodules( + input_layernorm=_build_norm_spec(config=config, provider=provider), + self_attention=_build_attention_spec( + compress_ratio=compress_ratio, + config=config, + provider=provider, + ), + pre_mlp_layernorm=_build_norm_spec(config=config, provider=provider), + mlp=_build_ffn_spec( + config=config, + provider=provider, + layer_idx=layer_idx, + ), + attn_hc=ModuleSpec(module=HyperMixer) if hc_mult > 1 else None, + ffn_hc=ModuleSpec(module=HyperMixer) if hc_mult > 1 else None, + ) + + return ModuleSpec( + module=DeepseekV4HybridLayer, + params={ + "layer_idx": layer_idx, + "compress_ratio": int(compress_ratio), + }, + submodules=layer_submodules, + ) + + +def _build_stage_hybrid_layer_specs( + config: DeepSeekV4TransformerConfig, + *, + provider: DeepSeekV4SpecProvider, + vp_stage: Optional[int], +) -> List[ModuleSpec]: + """Build the current stage's decoder layer specs. + + DeepSeek-V4 runtime always materializes a concrete stage-local + `layer_specs` list for `DeepseekV4TransformerBlock`. + """ + num_layers = int(config.num_layers) + mtp_num_layers = int(config.mtp_num_layers) + compress_ratios = _normalize_compress_ratios( + config.compress_ratios, + num_layers=num_layers, + mtp_num_layers=mtp_num_layers, + ) + + try: + local_layer_count = int(get_num_layers_to_build(config, vp_stage=vp_stage)) + layer_offset = int(get_transformer_layer_offset(config, vp_stage=vp_stage)) + except Exception: + local_layer_count = num_layers + layer_offset = 0 + + local_start = max(0, layer_offset) + local_end = min(num_layers, local_start + max(0, local_layer_count)) + local_layer_indices = range(local_start, local_end) + + return [ + _build_hybrid_layer_spec( + config, + provider=provider, + layer_idx=layer_idx, + compress_ratio=int(compress_ratios[layer_idx]), + ) + for layer_idx in local_layer_indices + ] + + +def get_deepseek_v4_runtime_decoder_spec( + config: DeepSeekV4TransformerConfig, + *, + vp_stage: Optional[int] = None, + pp_rank: Optional[int] = None, +) -> ModuleSpec: + """Return the effective V4 runtime decoder spec tree. + + The returned block submodules always include a non-empty `layer_specs`. + """ + del pp_rank + + provider = resolve_v4_provider(config) + logger.info("[DeepSeek-V4] resolve spec provider=%s", type(provider).__name__) + + hc_mult = int(config.hc_mult) + stage_layer_specs = _build_stage_hybrid_layer_specs( + config, + provider=provider, + vp_stage=vp_stage, + ) + assert stage_layer_specs, "DeepSeek-V4 requires non-empty stage layer specs." + block_submodules = DeepseekV4TransformerBlockSubmodules( + layer_specs=stage_layer_specs, + hyper_head=ModuleSpec(module=HyperHead) if hc_mult > 1 else None, + # DeepseekV4TransformerBlock decides whether this stage owns final norm. + final_layernorm=_build_norm_spec(config=config, provider=provider), + ) + return ModuleSpec(module=DeepseekV4TransformerBlock, submodules=block_submodules) + + +__all__ = [ + "get_deepseek_v4_runtime_decoder_spec", +] diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py new file mode 100644 index 000000000..8964bb5c2 --- /dev/null +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py @@ -0,0 +1,323 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +DeepSeek-V4 top-level model. + +This model intentionally subclasses :class:`LanguageModule` (not GPTModel) +so DeepSeek-V4 no longer depends on GPT's internal TransformerBlock +construction path. +""" + +from typing import Literal, Optional, Union + +from megatron.core import tensor_parallel +from megatron.core.models.common.embeddings.language_model_embedding import ( + LanguageModelEmbedding, +) +from megatron.core.models.common.language_module.language_module import LanguageModule +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.enums import ModelType +from megatron.core.transformer.multi_token_prediction import ( + MultiTokenPredictionBlock, + mtp_on_this_rank, + process_mtp_loss, +) +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from torch import Tensor + +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_mtp_specs import ( + get_v4_mtp_block_spec, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, +) + + +class DeepseekV4Model(LanguageModule): + """DeepSeek-V4 language model rooted on LanguageModule.""" + + def __init__( + self, + config: DeepSeekV4TransformerConfig, + transformer_layer_spec: Union[ModuleSpec, type], + vocab_size: int, + max_sequence_length: int, + pre_process: bool = True, + post_process: bool = True, + fp16_lm_cross_entropy: bool = False, + parallel_output: bool = True, + share_embeddings_and_output_weights: bool = False, + position_embedding_type: Literal[ + "learned_absolute", + "rope", + "mrope", + "yarn", + "none", + ] = "none", + rotary_percent: float = 1.0, + rotary_base: int = 10000, + rope_scaling: bool = False, + scatter_embedding_sequence_parallel: bool = True, + pg_collection: Optional[ProcessGroupCollection] = None, + vp_stage: Optional[int] = None, + **_kwargs, + ) -> None: + del rotary_percent, rotary_base, rope_scaling + super().__init__(config=config, pg_collection=pg_collection) + + self.transformer_layer_spec = transformer_layer_spec + self.vocab_size = vocab_size + self.max_sequence_length = max_sequence_length + self.pre_process = pre_process + self.post_process = post_process + self.fp16_lm_cross_entropy = fp16_lm_cross_entropy + self.parallel_output = parallel_output + self.share_embeddings_and_output_weights = share_embeddings_and_output_weights + self.vp_stage = vp_stage + self.model_type = ModelType.encoder_or_decoder + + if hasattr(self.config, "position_embedding_type"): + self.position_embedding_type = self.config.position_embedding_type + else: + self.position_embedding_type = position_embedding_type + + # Compute ``mtp_process`` (and build the MTP block spec) *before* the + # embedding so we can mirror upstream GPTModel: when MTP layers live on + # a non-pre_process PP stage, that stage still needs a (tied) copy of + # the input embedding. Without this, ``setup_embeddings_and_output_layer`` + # -> ``shared_embedding_or_output_weight`` asserts because ``self.embedding`` + # was never created on the MTP stage. + # + # Plan-2 P16/P17: V4 wires multi-token prediction exclusively via + # the spec-based upstream :class:`MultiTokenPredictionBlock`, + # built from :func:`get_v4_mtp_block_spec`. The legacy + # primus-owned ``DeepseekV4MTPBlock`` (gated by + # ``v4_use_custom_mtp_block`` in plan-2 P16) was retired in plan-2 + # P17; only the spec-based path remains. + mtp_num_layers = int(getattr(self.config, "mtp_num_layers", 0) or 0) + self.mtp_process = False + self.mtp_block_spec = None + if mtp_num_layers > 0: + self.mtp_block_spec = get_v4_mtp_block_spec( + self.config, + transformer_layer_spec=transformer_layer_spec, + vp_stage=vp_stage, + ) + # ``mtp_on_this_rank`` reads ``parallel_state`` and + # :class:`MultiTokenPredictionBlock` walks ``pg_collection.cp``; + # both require a real distributed init. On CPU smokes (no + # ``torch.distributed``) we leave ``self.mtp`` as ``None`` and + # surface the spec via ``self.mtp_block_spec`` so callers can + # still inspect the MTP wiring (the spec helper itself is fully + # CPU-testable). + try: + self.mtp_process = mtp_on_this_rank(self.config, ignore_virtual=False, vp_stage=vp_stage) + except (AssertionError, RuntimeError, AttributeError): + self.mtp_process = False + + # The embedding is needed on pre_process stages and, when MTP is + # enabled, also on MTP-process stages (which keep a tied copy). + if self.pre_process or self.mtp_process: + self.embedding = LanguageModelEmbedding( + config=self.config, + vocab_size=self.vocab_size, + max_sequence_length=self.max_sequence_length, + position_embedding_type=self.position_embedding_type, + scatter_to_sequence_parallel=scatter_embedding_sequence_parallel, + tp_group=self.pg_collection.tp, + ) + + self.decoder = build_module( + transformer_layer_spec, + config=self.config, + pre_process=self.pre_process, + post_process=self.post_process, + pg_collection=self.pg_collection, + vp_stage=vp_stage, + ) + + # NOTE: The cross-PP broadcast of ``input_ids`` for V4 hash-routed + # MoE layers (needed when middle PP stages own a hash-routed layer + # but Megatron's ``pretrain_gpt.get_batch`` returns ``None`` tokens + # for non-(first|last) PP stages) lives in a Primus patch on + # ``pretrain_gpt.get_batch``: + # ``primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py`` + # That hook runs once per ``forward_step`` (i.e. once per + # (chunk, microbatch) for both 1F1B and interleaved-1F1B), which is + # VPP-safe. An earlier in-``forward`` broadcast deadlocked the + # interleaved schedule because PP rank 0's broadcast wait blocked + # before its ``send_forward``, while PP rank 1 was simultaneously + # waiting for that ``send_forward`` in ``recv_forward``. + + # ----- MTP block --------------------------------------------------- + # NOTE: do NOT pre-assign ``self.mtp = None``. Megatron's + # ``set_current_microbatch`` (third_party/Megatron-LM/megatron/core/ + # transformer/cuda_graphs.py) probes ``hasattr(model, 'mtp')`` and + # unconditionally iterates ``model.mtp.layers``. We only create the + # attribute when MTP is actually live (matching upstream GPTModel). + if self.mtp_process: + self.mtp = MultiTokenPredictionBlock( + config=self.config, + spec=self.mtp_block_spec, + vp_stage=vp_stage, + pg_collection=self.pg_collection, + ) + + if self.post_process: + if getattr(self.config, "defer_embedding_wgrad_compute", False): + self.embedding_activation_buffer = [] + self.grad_output_buffer = [] + else: + self.embedding_activation_buffer = None + self.grad_output_buffer = None + + self.output_layer = tensor_parallel.ColumnParallelLinear( + self.config.hidden_size, + self.vocab_size, + config=self.config, + init_method=( + self.config.embedding_init_method + if getattr(self.config, "use_mup", False) and not self.share_embeddings_and_output_weights + else self.config.init_method + ), + bias=False, + skip_bias_add=False, + gather_output=not self.parallel_output, + skip_weight_param_allocation=self.pre_process and self.share_embeddings_and_output_weights, + embedding_activation_buffer=self.embedding_activation_buffer, + grad_output_buffer=self.grad_output_buffer, + tp_group=self.pg_collection.tp, + ) + + if self.pre_process or self.post_process or self.mtp_process: + self.setup_embeddings_and_output_layer() + + def set_input_tensor(self, input_tensor: Tensor) -> None: + """Pipeline-parallel hook to set decoder input tensor.""" + if not isinstance(input_tensor, list): + input_tensor = [input_tensor] + assert len(input_tensor) == 1, "input_tensor should only be length 1 for decoder-only models" + self.decoder.set_input_tensor(input_tensor[0]) + + def forward( + self, + input_ids: Optional[Tensor], + position_ids: Optional[Tensor], + attention_mask: Optional[Tensor], + decoder_input: Optional[Tensor] = None, + labels: Optional[Tensor] = None, + loss_mask: Optional[Tensor] = None, + runtime_gather_output: Optional[bool] = None, + packed_seq_params=None, + **kwargs, + ): + """Forward pass for DeepSeek-V4. + + Plan-2 P15: ``input_ids`` are passed to the decoder as the + ``token_ids`` forward kwarg (replacing the ``decoder._v4_token_ids`` + attribute stash). Hash-routed MoE layers consume them directly via + the standard kwargs propagation chain + ``model.forward -> decoder.forward -> layer.forward -> mlp.forward + -> hash_router.forward``. + + Plan-2 P16: when ``mtp_num_layers > 0`` and the spec-based MTP + path is enabled (default), this method runs the upstream + :class:`MultiTokenPredictionBlock` on the post_process stage and + feeds its concatenated output through :func:`process_mtp_loss`, + which adds the auxiliary MTP loss term to the main LM loss. + """ + if decoder_input is None: + if self.pre_process: + if input_ids is None: + raise ValueError("input_ids must be provided when pre_process=True.") + if position_ids is None: + batch, seq = input_ids.shape + position_ids = ( + input_ids.new_arange(seq, dtype=input_ids.dtype).unsqueeze(0).expand(batch, -1) + ) + decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids) + else: + decoder_input = None + + # ``input_ids`` arrives on every PP stage that owns hash-routed MoE + # layers because the Primus patch + # ``primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py`` + # broadcasts the source-of-truth tokens from PP rank 0 inside + # ``pretrain_gpt.get_batch`` before this ``forward`` runs. So + # ``input_ids`` is non-``None`` here on every middle PP stage even + # though Megatron's data loader would otherwise feed it ``None``. + + hidden_states = self.decoder( + hidden_states=decoder_input, + attention_mask=attention_mask, + position_ids=position_ids, + token_ids=input_ids, + packed_seq_params=packed_seq_params, + **kwargs, + ) + + # Run the spec-based MTP block on stages that own MTP layers. + # Mirrors GPTModel's mtp_in_postprocess gating. + if self.mtp_process and getattr(self, "mtp", None) is not None: + hidden_states = self.mtp( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + attention_mask=attention_mask, + packed_seq_params=packed_seq_params, + # MTP needs the (tied) word-embedding module to embed the + # shifted input_ids. On a dedicated MTP PP stage pre_process is + # False, but the stage still owns a tied embedding copy created + # in __init__ (gated on ``pre_process or mtp_process``), so pass + # whichever embedding module exists on this rank. + embedding=getattr(self, "embedding", None), + ) + + if not self.post_process: + return hidden_states + + output_weight = None + if self.share_embeddings_and_output_weights: + output_weight = self.shared_embedding_or_output_weight() + + # Plan-2 P16: when MTP is on, ``hidden_states`` arrives as the + # concatenation of the main-decoder hidden state plus + # ``mtp_num_layers`` shifted MTP hidden states (along the + # sequence axis). :func:`process_mtp_loss` splits the chunks, + # computes the per-depth MTP loss, and returns the main hidden + # state for the standard LM-head path below. + mtp_num_layers = int(getattr(self.config, "mtp_num_layers", 0) or 0) + if mtp_num_layers > 0 and getattr(self, "mtp", None) is not None: + cp_group = getattr(self.pg_collection, "cp", None) + hidden_states = process_mtp_loss( + hidden_states=hidden_states, + labels=labels, + loss_mask=loss_mask, + output_layer=self.output_layer, + output_weight=output_weight, + runtime_gather_output=runtime_gather_output, + is_training=self.training, + compute_language_model_loss=self.compute_language_model_loss, + config=self.config, + cp_group=cp_group, + packed_seq_params=packed_seq_params, + scale_logits_fn=self._scale_logits if getattr(self.config, "use_mup", False) else None, + ) + + logits, _ = self.output_layer( + hidden_states, + weight=output_weight, + runtime_gather_output=runtime_gather_output, + ) + logits = self._scale_logits(logits) + + if labels is None: + return logits.transpose(0, 1).contiguous() + return self.compute_language_model_loss(labels, logits) + + +__all__ = ["DeepseekV4Model"] diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_layer.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_layer.py new file mode 100644 index 000000000..8c0aa7e82 --- /dev/null +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_layer.py @@ -0,0 +1,189 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 multi-token-prediction (MTP) layer. + +V4's MTP differs from the V3 / upstream :class:`MultiTokenPredictionLayer` +in two ways that the upstream layer cannot express on its own: + +1. **Multi-stream (mHC) inner layer.** The MTP inner transformer layer is a + :class:`DeepseekV4HybridLayer`, which — when ``hc_mult > 1`` — operates on + the K-stream form ``[B, S, K, D]`` (see ``deepseek_v4_block``). The upstream + MTP layer feeds its inner layer a single-stream ``[S, B, D]`` tensor, so we + must lift to K streams before the inner layer and collapse back after it. + +2. **Per-depth ``hc_head_fn``.** The released V4 checkpoint gives each MTP + depth its *own* small :class:`HyperHead` (``hc_head_fn``) to collapse the K + streams — it does **not** reuse the main trunk's HyperHead (techblog + §8 / DeepSeek-V4 report). This is gated by + ``config.mtp_use_separate_hc_head``. + +The class subclasses upstream :class:`MultiTokenPredictionLayer` so it keeps +the shared embedding-roll / ``enorm`` / ``hnorm`` / ``eh_proj`` / +final-layernorm machinery and plugs straight into +:class:`MultiTokenPredictionBlock` (which builds each depth via +``build_module(layer_spec, ...)`` and therefore honours our ``module`` slot). +Only the inner-layer call site (``_proj_and_transformer_layer``) is overridden +to insert the lift / collapse, and ``forward`` is wrapped to capture +``position_ids`` (V4 attention derives its dual-RoPE from absolute positions, +and the upstream layer does not thread ``position_ids`` into +``_proj_and_transformer_layer``). +""" + +from __future__ import annotations + +from contextlib import nullcontext +from typing import Optional + +import torch +from megatron.core import tensor_parallel +from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionLayer +from megatron.core.transformer.spec_utils import build_module + +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block import ( + _lift_streams_in, + _lower_streams_out, +) +from primus.backends.megatron.core.transformer.hyper_connection import HyperHead + +try: # get_fp8_context lives in megatron.core.fp8_utils across recent versions + from megatron.core.fp8_utils import get_fp8_context +except Exception: # pragma: no cover - defensive: keep BF16 path working + + def get_fp8_context(config, *args, **kwargs): # type: ignore[misc] + return nullcontext() + + +class DeepseekV4MTPLayer(MultiTokenPredictionLayer): + """One V4 MTP depth (mHC-aware, per-depth HyperHead).""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + # Upstream MultiTokenPredictionLayer builds ``self.mtp_model_layer`` + # WITHOUT threading ``pg_collection`` (see + # ``multi_token_prediction.py``: ``build_module(self.submodules. + # mtp_model_layer, config=..., vp_stage=..., layer_number=..., + # is_mtp_layer=True)``). For a V4 :class:`DeepseekV4HybridLayer` whose + # MLP is a :class:`DeepseekV4MoE`, ``pg_collection=None`` selects the + # *local-experts* (non-expert-parallel) path, which (a) instantiates + # ALL routed experts on every rank and (b) breaks the DDP grad-bucket + # invariant (``len(per_param_grad_ready_counts) == len(params)`` in + # ``param_and_grad_buffer.reset``) because the local path's per-expert + # modules do not match the EP dispatcher's grad-ready bookkeeping that + # the main decoder relies on. Rebuild the inner layer WITH the + # ``pg_collection`` the MTP block already holds so the MTP MoE uses the + # exact same expert-parallel dispatcher path as the main decoder. Free + # the throwaway first build before rebuilding to bound init memory. + pg_collection = kwargs.get("pg_collection", None) + vp_stage = kwargs.get("vp_stage", None) + if pg_collection is not None and getattr(self.submodules, "mtp_model_layer", None) is not None: + self.mtp_model_layer = None # drop the pg_collection-less build + self.mtp_model_layer = build_module( + self.submodules.mtp_model_layer, + config=self.config, + vp_stage=vp_stage, + layer_number=self.layer_number, + is_mtp_layer=True, + pg_collection=pg_collection, + ) + + self.hc_mult = int(getattr(self.config, "hc_mult", 1) or 1) + use_separate_head = bool(getattr(self.config, "mtp_use_separate_hc_head", True)) + + self.mtp_hyper_head: Optional[HyperHead] = None + if self.hc_mult > 1: + if not use_separate_head: + # A single-stream MTP inner layer (no per-depth head) would + # require the inner DeepseekV4HybridLayer to be built with + # hc_mult=1, but it inherits the model's hc_mult and emits the + # K-stream form. Fail loud instead of silently shape-crashing. + raise NotImplementedError( + "DeepseekV4MTPLayer with hc_mult>1 requires " + "config.mtp_use_separate_hc_head=True (per-depth HyperHead); " + "single-stream MTP-with-mHC is not wired." + ) + self.mtp_hyper_head = HyperHead( + hidden_size=int(self.config.hidden_size), + hc_mult=self.hc_mult, + eps=float(getattr(self.config, "hc_eps", 1.0e-6)), + ) + + # Stash for ``_proj_and_transformer_layer`` (set in ``forward``). + self._mtp_position_ids: Optional[torch.Tensor] = None + + # ------------------------------------------------------------------ + + def forward(self, *args, **kwargs): + """Capture ``position_ids`` then defer to the upstream forward. + + The model calls this with keyword args + (``input_ids=``, ``position_ids=``, ``hidden_states=`` ...), so we + read ``position_ids`` from ``kwargs``. V4 attention consumes absolute + positions; the MTP token/label shift is handled by the upstream + embedding roll, so we thread the *unrolled* ``position_ids`` into the + inner hybrid layer. + """ + if "position_ids" in kwargs: + self._mtp_position_ids = kwargs["position_ids"] + elif len(args) >= 2: + self._mtp_position_ids = args[1] + return super().forward(*args, **kwargs) + + # ------------------------------------------------------------------ + + def _proj_and_transformer_layer( + self, + hidden_states: torch.Tensor, + decoder_input: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + **_kwargs, + ) -> torch.Tensor: + """eh_proj -> (lift -> V4 hybrid layer -> per-depth HyperHead) -> norm. + + Mirrors the upstream method's fp8 / rng context handling but inserts + the K-stream lift / collapse around the inner + :class:`DeepseekV4HybridLayer` so the mHC math matches the main + decoder block exactly (same ``_lift_streams_in`` / ``_lower_streams_out`` + helpers). + """ + if self.config.sequence_parallel: + rng_context = tensor_parallel.get_cuda_rng_tracker().fork() + else: + rng_context = nullcontext() + + if self.config.fp8: + fp8_context = get_fp8_context(self.config) + transformer_layer_fp8_context = get_fp8_context(self.config) + else: + fp8_context = nullcontext() + transformer_layer_fp8_context = nullcontext() + + with rng_context: + with fp8_context: + # [S, B, D] single-stream after eh_proj. + hidden_states = self._concat_embeddings(hidden_states, decoder_input) + + with transformer_layer_fp8_context: + # Lift to the K-stream form the V4 hybrid layer expects. + x = _lift_streams_in( + hidden_states, pre_process=True, hc_mult=self.hc_mult + ) # [B, S, K, D] (or [B, S, D] when hc_mult == 1) + x, _ = self.mtp_model_layer( + x, + position_ids=self._mtp_position_ids, + token_ids=None, + ) + # Per-depth collapse K streams -> single stream. + if self.hc_mult > 1 and self.mtp_hyper_head is not None: + x = self.mtp_hyper_head(x) # [B, S, D] + hidden_states = _lower_streams_out(x, post_process=True, hc_mult=self.hc_mult) # [S, B, D] + + hidden_states = self._postprocess(hidden_states) + return hidden_states + + +__all__ = ["DeepseekV4MTPLayer"] diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_specs.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_specs.py new file mode 100644 index 000000000..b0ae7e583 --- /dev/null +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_mtp_specs.py @@ -0,0 +1,201 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 MTP (Multi-Token Prediction) block spec. + +Plan-2 P16 wires V4 onto Megatron's upstream +:class:`MultiTokenPredictionBlock` (multi-token prediction with the +classic eh_proj + per-depth transformer layer + final layernorm +recipe). The V4 specialization is: + +* Each MTP-depth's inner layer is a :class:`DeepseekV4HybridLayer` (with + HC, hash routing, and clamped-SwiGLU all wired through the same spec + tree as the main decoder). +* The two pre-projection norms (``enorm`` over the embedding, + ``hnorm`` over the prior hidden state) and the post-MTP final + layernorm reuse the V4 RMSNorm provider. +* The eh_proj is a column-parallel linear (also from the provider). + +What this file does *not* include: + +* The MTP block forward path itself — that's owned by upstream + :class:`MultiTokenPredictionBlock` once we hand it the spec. +* The ``HyperHead`` / loss-aware shifting / RouterReplay snapshot + matching (those live inside :func:`process_mtp_loss` and the V4 + layer's ``forward``). + +Reference: techblog §7 ("MTP V4 head") and +``DeepSeek-V4-Flash/inference/model.py:MTPBlock``. + +This is the **only** MTP path in plan-2 (P17 retired the legacy +primus-owned :class:`DeepseekV4MTPBlock` and the +``v4_use_custom_mtp_block`` config flag). Set ``mtp_num_layers > 0`` in +the config to enable MTP; the model wires this helper into +:class:`MultiTokenPredictionBlock` automatically. +""" + +from __future__ import annotations + +from typing import Optional + +from megatron.core.transformer.multi_token_prediction import ( + MultiTokenPredictionBlock, + MultiTokenPredictionBlockSubmodules, + MultiTokenPredictionLayerSubmodules, +) +from megatron.core.transformer.spec_utils import ModuleSpec + +from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import ( + DeepSeekV4SpecProvider, +) +from primus.backends.megatron.core.models.deepseek_v4.build_context import ( + resolve_v4_provider, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block import ( + DeepseekV4TransformerBlock, + DeepseekV4TransformerBlockSubmodules, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_mtp_layer import ( + DeepseekV4MTPLayer, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, +) + + +def _extract_v4_inner_layer_spec(transformer_layer_spec: ModuleSpec) -> ModuleSpec: + """Return a single :class:`DeepseekV4HybridLayer` spec for an MTP depth. + + :class:`DeepseekV4Model` resolves the decoder as a + :class:`DeepseekV4TransformerBlock` ``ModuleSpec`` whose + ``submodules.layer_specs`` is the stage-local list of hybrid-layer specs. + Upstream :class:`MultiTokenPredictionLayer` needs a *single* + :class:`~megatron.core.transformer.transformer_layer.TransformerLayer`-style + spec as its inner layer (its ``__init__`` validates + ``mtp_model_layer.submodules`` against ``TransformerLayerSubmodules``), + so when handed the block spec we extract the last hybrid layer spec + (mirroring upstream GPT's ``spec.layer_specs[-1]`` convention; the last + decoder layer is ``cr=0`` dense in both V4-Flash and V4-Pro). When already + handed a layer spec (e.g. CPU unit tests) we thread it through unchanged. + """ + submods = getattr(transformer_layer_spec, "submodules", None) + is_block_spec = getattr( + transformer_layer_spec, "module", None + ) is DeepseekV4TransformerBlock or isinstance(submods, DeepseekV4TransformerBlockSubmodules) + if is_block_spec: + layer_specs = getattr(submods, "layer_specs", None) + if not layer_specs: + raise ValueError( + "Cannot build a V4 MTP inner layer: the decoder block spec has " + "no stage-local layer_specs to extract from." + ) + return layer_specs[-1] + return transformer_layer_spec + + +def _v4_mtp_layer_spec( + *, + config: DeepSeekV4TransformerConfig, + transformer_layer_spec: ModuleSpec, + provider: DeepSeekV4SpecProvider, +) -> ModuleSpec: + """One MTP depth's :class:`MultiTokenPredictionLayer` spec. + + Args: + config: V4 transformer config (carries ``hidden_size``, etc.). + transformer_layer_spec: the V4 hybrid-layer spec used as the + inner layer for this depth. Plan-2 wires it through + unchanged so the MTP layer reuses HC / hash routing / + clamped-SwiGLU exactly the way the main decoder does. + provider: V4 spec provider (resolves RMSNorm + column-parallel + linear modules). + """ + del config + norm_module = provider.v4_norm_module() + column_parallel = provider.column_parallel_linear() + + inner_layer_spec = _extract_v4_inner_layer_spec(transformer_layer_spec) + + return ModuleSpec( + module=DeepseekV4MTPLayer, + submodules=MultiTokenPredictionLayerSubmodules( + enorm=norm_module, + hnorm=norm_module, + eh_proj=column_parallel, + mtp_model_layer=inner_layer_spec, + layer_norm=norm_module, + ), + ) + + +def get_v4_mtp_block_spec( + config: DeepSeekV4TransformerConfig, + *, + transformer_layer_spec: ModuleSpec, + vp_stage: Optional[int] = None, +) -> ModuleSpec: + """Return a :class:`ModuleSpec` for the V4 MTP block. + + The returned spec wraps :class:`MultiTokenPredictionBlock` with one + :class:`MultiTokenPredictionLayer` spec per MTP depth (V4-Flash + uses ``mtp_num_layers=1``; larger variants may use more depths). + The block submodules carry the per-depth specs; the block's + ``__init__`` walks them and instantiates :class:`MultiTokenPredictionLayer` + instances that reuse the V4 hybrid layer for inner attention / MLP + math. + + Args: + config: V4 transformer config. Must have + ``mtp_num_layers >= 1`` (caller checks before invoking + this helper). + transformer_layer_spec: the V4 hybrid-layer spec for the main + decoder. The same spec is reused for each MTP depth so the + inner attention + MoE math matches the main decoder + exactly. Plan-2 §16 confirms V4 uses a single + (non-repeated) shape for MTP layers. + vp_stage: optional virtual-pipeline stage index. Passed through + to upstream MTP code; ignored on non-VP runs. + + Returns: + A ``ModuleSpec`` that builds a fully-wired + :class:`MultiTokenPredictionBlock` when handed to + :func:`build_module` with ``config`` + ``pg_collection`` + kwargs. + + Notes: + * V4 has no ``mtp_use_repeated_layer`` carry-over; the spec + replicates the inner layer ``mtp_num_layers`` times. (V4 + checkpoints store one set of MTP weights per depth.) + * The ``HyperHead`` per-depth collapse lives inside + :class:`DeepseekV4HybridLayer` itself (via the layer's HC + mixers); the upstream MTP block does not need V4-specific + HyperHead awareness. + """ + del vp_stage # currently only forwarded by callers; not needed here + + if int(config.mtp_num_layers or 0) < 1: + raise ValueError( + "get_v4_mtp_block_spec requires mtp_num_layers >= 1; " + f"got mtp_num_layers={config.mtp_num_layers!r}." + ) + + provider = resolve_v4_provider(config) + mtp_layer_specs = [ + _v4_mtp_layer_spec( + config=config, + transformer_layer_spec=transformer_layer_spec, + provider=provider, + ) + for _ in range(int(config.mtp_num_layers)) + ] + + return ModuleSpec( + module=MultiTokenPredictionBlock, + submodules=MultiTokenPredictionBlockSubmodules(layer_specs=mtp_layer_specs), + ) + + +__all__ = ["get_v4_mtp_block_spec"] diff --git a/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py new file mode 100644 index 000000000..f3d40315b --- /dev/null +++ b/primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_transformer_config.py @@ -0,0 +1,214 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 specific transformer config. + +This config extends Megatron's ``MLATransformerConfig`` with DeepSeek-V4 +runtime fields that are referenced by V4 modules but are not part of the +upstream ``TransformerConfig``/``MLATransformerConfig`` schema. +""" + +from __future__ import annotations + +import ast +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +from megatron.core.transformer.transformer_config import MLATransformerConfig + +# ``mtp_compress_ratios`` and ``v4_use_custom_mtp_block`` lived here in +# plan-1 / plan-2 P12-P16 as escape hatches for the legacy primus-owned +# :class:`DeepseekV4MTPBlock`. Plan-2 P17 retired that block (the MTP path +# is now exclusively the spec-based upstream +# :class:`MultiTokenPredictionBlock` route via +# :func:`get_v4_mtp_block_spec`); both fields are deliberately gone here +# and their references in YAML configs / training scripts must be removed. + + +def _normalize_compress_ratios_field( + value: Optional[Union[str, List[int], Tuple[int, ...]]], + *, + field_name: str = "compress_ratios", +) -> Optional[Tuple[int, ...]]: + """Plan-2 P18 (D4 audit): normalize ``compress_ratios`` to a tuple. + + YAML loaders deliver this field as a *string* (e.g. + ``"[0, 0, 4, 128, ...]"``) when wrapped in quotes, or as an actual + list when written without quotes. The dataclass stored both as + ``Optional[Union[str, List[int], Tuple[int, ...]]]``, and runtime + helpers (``_parse_int_sequence`` / ``_normalize_compress_ratios`` + in ``deepseek_v4_block.py``) had to ``ast.literal_eval`` the string + on every consumer path. + + With this helper running once in ``__post_init__``, every consumer + sees a single canonical type — ``tuple[int, ...]`` — and the runtime + parsing layer is reduced to a length-fitting check. + + Args: + value: raw config value (string, list, tuple, or ``None``). + field_name: only used for error messages. + + Returns: + ``None`` (when ``value`` is ``None``) or a tuple of ints. + """ + if value is None: + return None + parsed = value + if isinstance(parsed, str): + try: + parsed = ast.literal_eval(parsed) + except (SyntaxError, ValueError) as exc: + raise ValueError(f"{field_name} must be a list-like value, got {value!r}") from exc + if isinstance(parsed, (list, tuple)): + try: + return tuple(int(x) for x in parsed) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field_name} entries must be int-castable; got {parsed!r}") from exc + raise TypeError(f"{field_name} must be a list/tuple/str, got {type(parsed).__name__}") + + +@dataclass +class DeepSeekV4TransformerConfig(MLATransformerConfig): + # ---- DeepSeek-V4 hybrid attention / HC ---- + hc_mult: int = 1 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1.0e-6 + + # ---- DeepSeek-V4 MTP (multi-token prediction) ---- + # The released V4 checkpoint gives each MTP depth its OWN small + # ``hc_head_fn`` (a per-depth :class:`HyperHead`) instead of reusing the + # main trunk's HyperHead. With ``hc_mult > 1`` the MTP inner + # :class:`DeepseekV4HybridLayer` runs on the K-stream form and must be + # collapsed back to a single stream by this per-depth head before the + # MTP final-layernorm + shared LM head. Set ``False`` to fall back to a + # single-stream MTP inner layer (no per-depth head); see + # ``deepseek_v4_mtp_layer.DeepseekV4MTPLayer``. + mtp_use_separate_hc_head: bool = True + + compress_ratios: Optional[Union[str, List[int], Tuple[int, ...]]] = None + compress_rope_theta: float = 160000.0 + + # ---- DeepSeek-V4 attention extras ---- + attn_sliding_window: int = 0 + attn_sink: bool = False + index_topk: int = 0 + index_head_dim: int = 128 + index_n_heads: int = 64 + + # ---- DeepSeek-V4 FP8 Indexer (CSA selector QK path) ---- + # When True, the CSA :class:`Indexer` fake-quantizes its query / compressed + # key activations to FP8 (E4M3) before the QK scoring einsum, simulating + # the released V4 low-precision indexer QK path (the report runs the + # indexer QK in FP4/FP8 while keeping the BF16 index-score / top-k path). + # The score reduction (ReLU + per-head weight + sum + causal mask) and the + # top-k selection stay in the activation dtype (BF16). The indexer is a + # non-differentiable, frozen top-k selector, so this is an inference-side + # precision reduction (no straight-through estimator needed). Default False + # so existing BF16 runs are bit-unchanged; enable via + # ``PRIMUS_USE_V4_FP8_INDEXER`` (see run_deepseek_v4.sh / flash yaml). + use_v4_fp8_indexer: bool = False + + # ---- DeepSeek-V4 attention backend selection (unified string selectors) ---- + # ``use_v4_attention_backend`` selects the dense (cr=0) / HCA (cr=128) kernel; + # ``use_v4_csa_attention_backend`` selects the CSA (cr=4) kernel: + # dense/HCA: eager | triton_v1 | triton_v2 | gluon | gluon_v2 | flydsl_v1 | turbo + # CSA: eager | triton_v0 | triton_v1 | triton_v2 | gluon | gluon_v2 | flydsl_v0 | flydsl_v1 | turbo + # (triton_v0 = deprecated gathered; flydsl_v0 = deprecated legacy FlyDSL, + # fwd-only; ``turbo`` = Primus-Turbo native-FlyDSL sparse-MLA via the turbo + # API, primus_turbo.flydsl.attention). ``use_turbo_attention`` (when a + # ``core_attention`` module is built) still takes precedence for the dense path. + use_v4_attention_backend: str = "triton_v1" + use_v4_csa_attention_backend: str = "triton_v1" + + # ---- Per-module precision (paper recipe): routed experts in MXFP4 while the + # rest of the layer runs FP8. Decoupled from the global --fp4/--fp8 recipe + # (which are mutually exclusive): with FP8 on, the PrimusTurbo grouped MLP + # routes the expert GEMMs through the native FP4 (hipBLASLt) path. Default OFF. + moe_experts_fp4: bool = False + + # ---- DeepSeek-V4 plan-5 P29: torch.compile-fused Sinkhorn ---- + # Plan-5 P29 (RESCOPED from "small-op fusion" — see plan-5 02-phase- + # details.md). The Sinkhorn-Knopp doubly-stochastic projection inside + # ``HyperMixer.compute_weights`` (see ``hyper_connection.py``) issues + # 1 + 2 * (n_iters - 1) separate fp32 ``aten::sum`` reductions per + # call. At V4-Flash production widths the per-call shape is + # ``[B, S, K, K] = [1, 4096, 4, 4]`` and HIP's default + # ``reduce_kernel<512, 1, ...>`` runs the kernel ~250x over the + # memory-bound floor. The P28 baseline trace pinned this kernel as + # 87.3 % of step time. Setting this flag to ``True`` collapses every + # such call into a single ``torch.compile(fullgraph=True, + # dynamic=False)`` Inductor-fused Triton kernel; AOT autograd handles + # the BWD. Default ``False`` until G32 (FWD + BWD parity) and G33b + # (post-P29 trace) flip it on. + use_v4_compiled_sinkhorn: bool = False + + # ---- DeepSeek-V4 grouped low-rank output projection ---- + # Mirrors the released checkpoint's `wo_a` / `wo_b` layout. + # When ``o_lora_rank == 0`` the attention falls back to a flat O proj + # (Megatron's ``linear_proj``); set it >0 to use the grouped low-rank + # form (``linear_o_a`` + ``linear_o_b``) with ``o_groups`` groups. + o_groups: int = 8 + o_lora_rank: int = 0 + + # ---- DeepSeek-V4 MoE routing / expert extras ---- + num_hash_layers: int = 0 + hash_routing_seed: int = 0 + + moe_intermediate_size: Optional[int] = None + moe_use_legacy_grouped_gemm: bool = False + + swiglu_limit: float = 0.0 + v4_grouped_experts_support_clamped_swiglu: bool = False + + # ---- Vocab helpers used by hash router ---- + vocab_size: Optional[int] = None + padded_vocab_size: Optional[int] = None + + # ---- Compat aliases for V4 code paths ---- + norm_epsilon: Optional[float] = None + position_embedding_type: str = "none" + + def __post_init__(self) -> None: + # P18 D4: normalize compress_ratios once so downstream helpers + # always see ``tuple[int, ...]`` (or None). YAML strings like + # ``"[0, 0, 4, ...]"`` are evaluated here. + if self.compress_ratios is not None: + self.compress_ratios = _normalize_compress_ratios_field( + self.compress_ratios, field_name="compress_ratios" + ) + + # Keep V4's ``norm_epsilon`` alias consistent with MCore's + # ``layernorm_epsilon`` before parent validation runs. + if self.norm_epsilon is None: + self.norm_epsilon = float(self.layernorm_epsilon) + self.layernorm_epsilon = float(self.norm_epsilon) + + # DeepSeek naming compatibility for MoE hidden size. + if self.moe_ffn_hidden_size is None and self.moe_intermediate_size is not None: + self.moe_ffn_hidden_size = int(self.moe_intermediate_size) + + # Keep DeepSeek clamp name aligned with MCore clamp field. + clamp_from_activation = self.activation_func_clamp_value + clamp_from_swiglu = float(self.swiglu_limit) + if clamp_from_activation is None: + if clamp_from_swiglu > 0.0: + self.activation_func_clamp_value = clamp_from_swiglu + elif clamp_from_swiglu <= 0.0: + self.swiglu_limit = float(clamp_from_activation) + + # Ensure hash-router vocab lookups always have a concrete size. + if self.padded_vocab_size is None and self.vocab_size is not None: + self.padded_vocab_size = int(self.vocab_size) + if self.vocab_size is None and self.padded_vocab_size is not None: + self.vocab_size = int(self.padded_vocab_size) + + super().__post_init__() + + if self.moe_intermediate_size is None and self.moe_ffn_hidden_size is not None: + self.moe_intermediate_size = int(self.moe_ffn_hidden_size) + + +__all__ = ["DeepSeekV4TransformerConfig"] diff --git a/primus/backends/megatron/core/transformer/clamped_swiglu.py b/primus/backends/megatron/core/transformer/clamped_swiglu.py new file mode 100644 index 000000000..168986513 --- /dev/null +++ b/primus/backends/megatron/core/transformer/clamped_swiglu.py @@ -0,0 +1,192 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Clamped SwiGLU activation for DeepSeek-V4 (pre-mul layout). + +Reference: techblog §3 ("Activation: clamped SwiGLU") and the inference +reference at ``DeepSeek-V4-Flash/inference/model.py:Expert.forward``. + +DeepSeek-V4 replaces the standard ``SwiGLU(gate, up) = SiLU(gate) * up`` +with a **pre-multiplication** clamp: + + gate_c = clamp(gate, max=alpha) # one-sided (top only) + up_c = clamp(up, min=-alpha, max=alpha) # two-sided + out = SiLU(gate_c) * up_c + +This matches the released checkpoint and Megatron's +``mlp.MLP``-side clamp path (``activation_func_clamp_value``); both +clamp the gate / up *before* the activation and multiply, so the +post-multiply value is bounded by ``alpha * max(SiLU)`` and stays +well-behaved in bf16 expert summations. + +Plan-2 P14 contract: + +* :func:`clamped_swiglu_pre_mul` — split-input pointwise activation. Used + by the eager :class:`ClampedSwiGLUMLP` and as the canonical reference + for unit tests. +* :func:`clamped_swiglu_pre_mul_fused` — Megatron-fused-input + ``[..., 2I]`` gate-concat-up form. Lets a grouped-gemm expert backend + call a single function on the post-GEMM output. +* :class:`ClampedSwiGLUMLP` — eager MLP using **separate** ``w1`` (gate) + and ``w3`` (up) Linears so the parameter layout matches the + ``DeepSeek-V4-Flash`` checkpoint (``w1.weight`` / ``w3.weight``). + An optional ``fused_gate_up`` flag lets callers fuse the two GEMMs at + forward time without changing the state-dict layout — the saved / + loaded keys are always ``w1.weight`` / ``w3.weight``. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _resolve_alpha(alpha: Optional[float]) -> Optional[float]: + """Return a clamp bound or ``None`` when clamping is disabled.""" + if alpha is None: + return None + if alpha <= 0.0: + return None + return float(alpha) + + +def clamped_swiglu_pre_mul( + gate: torch.Tensor, + up: torch.Tensor, + *, + alpha: float = 7.0, +) -> torch.Tensor: + """Pre-multiplication clamped SwiGLU on split inputs. + + Args: + gate: ``[..., I]`` gate stream (output of ``w1``). + up: ``[..., I]`` up stream (output of ``w3``). + alpha: clamp bound; V4-Flash default is ``7.0``. Pass ``0`` / + ``None`` to fall back to vanilla ``SiLU(gate) * up``. + + Returns: + ``[..., I]`` activation output, ``SiLU(clamp(gate, max=alpha)) + * clamp(up, +/- alpha)``. + """ + if gate.shape != up.shape: + raise ValueError( + "clamped_swiglu_pre_mul expects matching gate / up shapes; " + f"got {tuple(gate.shape)} vs {tuple(up.shape)}." + ) + bound = _resolve_alpha(alpha) + if bound is not None: + gate_c = gate.clamp(max=bound) + up_c = up.clamp(min=-bound, max=bound) + else: + gate_c = gate + up_c = up + return F.silu(gate_c) * up_c + + +def clamped_swiglu_pre_mul_fused( + x: torch.Tensor, + *, + alpha: float = 7.0, +) -> torch.Tensor: + """Pre-multiplication clamped SwiGLU on a fused ``[..., 2I]`` input. + + The input is the Megatron-convention concatenated ``[gate | up]`` + along the last dimension (matching + :func:`megatron.core.fusions.fused_bias_swiglu.bias_swiglu` and the + eager glu path in :class:`megatron.core.transformer.mlp.MLP`). + + Args: + x: ``[..., 2 * I]`` — ``[gate | up]`` halves concatenated along + the last dim. + alpha: clamp bound; same semantics as :func:`clamped_swiglu_pre_mul`. + + Returns: + ``[..., I]`` activation output. + """ + if x.shape[-1] % 2 != 0: + raise ValueError( + "clamped_swiglu_pre_mul_fused expects a [gate | up] last dim; " + f"got shape {tuple(x.shape)} (last dim must be even)." + ) + gate, up = x.chunk(2, dim=-1) + return clamped_swiglu_pre_mul(gate, up, alpha=alpha) + + +class ClampedSwiGLUMLP(nn.Module): + """Eager SwiGLU MLP with V4's pre-multiplication clamp. + + Computes:: + + gate = w1(x) # [..., I] + up = w3(x) # [..., I] + h = SiLU(clamp(gate, max=alpha)) + * clamp(up, +/- alpha) + y = w2(h) # [..., D] + + The parameter layout (``w1`` / ``w2`` / ``w3``) mirrors the released + DeepSeek-V4-Flash ``Expert`` checkpoint exactly. The ``fused_gate_up`` + knob fuses the gate / up GEMMs at *forward time only* by stacking + ``w1.weight`` and ``w3.weight`` on the fly; the saved / loaded + ``state_dict`` keys remain ``w1.weight`` / ``w3.weight`` so released + checkpoints can be loaded without remapping. + + Notes: + * Bias is omitted by default to match V4 reference checkpoints. + * ``alpha=0`` (or ``None``) disables clamping → vanilla SwiGLU. + * This module is the canonical *eager* reference. Production + training uses Megatron's grouped-MLP path with + ``activation_func_clamp_value`` set; the math is the same. + """ + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + *, + alpha: float = 7.0, + bias: bool = False, + dtype: Optional[torch.dtype] = None, + fused_gate_up: bool = False, + ) -> None: + super().__init__() + if intermediate_size <= 0: + raise ValueError(f"intermediate_size must be > 0, got {intermediate_size}") + self.hidden_size = int(hidden_size) + self.intermediate_size = int(intermediate_size) + self.alpha = float(alpha) + self.fused_gate_up = bool(fused_gate_up) + + kw = {} if dtype is None else {"dtype": dtype} + # w1 = gate, w3 = up (V4 convention; matches DeepSeek-V4-Flash checkpoint). + self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias, **kw) + self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias, **kw) + self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias, **kw) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.fused_gate_up: + # Build the fused weight on the fly so the state_dict layout + # is unchanged (w1.weight / w3.weight). Bias is None by default. + w_gu = torch.cat([self.w1.weight, self.w3.weight], dim=0) + b_gu = None + if self.w1.bias is not None and self.w3.bias is not None: + b_gu = torch.cat([self.w1.bias, self.w3.bias], dim=0) + gate_up = F.linear(x, w_gu, b_gu) + h = clamped_swiglu_pre_mul_fused(gate_up, alpha=self.alpha) + else: + gate = self.w1(x) + up = self.w3(x) + h = clamped_swiglu_pre_mul(gate, up, alpha=self.alpha) + return self.w2(h) + + +__all__ = [ + "clamped_swiglu_pre_mul", + "clamped_swiglu_pre_mul_fused", + "ClampedSwiGLUMLP", +] diff --git a/primus/backends/megatron/core/transformer/compressor.py b/primus/backends/megatron/core/transformer/compressor.py new file mode 100644 index 000000000..c99679edd --- /dev/null +++ b/primus/backends/megatron/core/transformer/compressor.py @@ -0,0 +1,198 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +DeepSeek-V4 Compressor. + +Reference: techblog §1.3 ("Compressor: the Long-Range Compression Branch") and +the diagrams in ``deepseek-v4/develop/techblog/diagrams/csa.png`` / +``hca.png``. + +Two configurations: + +* ``ratio == 4`` (CSA branch) — overlap mode, ``coff == 2``: each compressed + token sees an effective window of ``2*ratio`` raw tokens (current window + plus the previous window's "leftover-half" channels). This smooths + boundary effects between adjacent compressed positions. +* ``ratio == 128`` (HCA branch) — non-overlap mode, ``coff == 1``: each + compressed token covers exactly ``ratio`` raw tokens. + +Compressor returns the pooled KV after a final ``kv_norm`` (RMSNorm). RoPE +at the compress branch theta is applied **outside** this module by the +caller (the dual-RoPE module produced in P4.3 + the CSA / HCA modules in +P4.4 will consume the output). +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from primus.backends.megatron.core.transformer.local_rmsnorm import LocalRMSNorm + + +class Compressor(nn.Module): + """V4 Compressor block. + + Args: + hidden_size: input feature dim ``D``. + head_dim: output channel dim per compressed position. + ratio: compression ratio ``m``. Must be a divisor of the runtime + sequence length (``S % ratio == 0``). + overlap: whether to use the overlap-stitched mode. If ``None``, + defaults to ``ratio == 4`` (the V4 convention). + rmsnorm_eps: RMSNorm stability eps. + + Shapes: + Forward input ``hidden``: ``[B, S, D]``. + Forward output ``pooled``: ``[B, S // ratio, head_dim]``. + """ + + def __init__( + self, + *, + hidden_size: int, + head_dim: int, + ratio: int, + overlap: Optional[bool] = None, + rmsnorm_eps: float = 1e-6, + ) -> None: + super().__init__() + if ratio < 1: + raise ValueError(f"ratio must be >= 1, got {ratio}") + + self.hidden_size = hidden_size + self.head_dim = head_dim + self.ratio = ratio + self.overlap = bool(ratio == 4 if overlap is None else overlap) + # coff is the projection multiplier — overlap mode needs 2x the + # channels because half goes to the "current window" half and half to + # the "previous window" half. + self.coff = 2 if self.overlap else 1 + + proj_out = self.coff * head_dim + self._proj_out = proj_out + # Fuse the kv + gate projections into ONE [hidden -> 2*proj_out] GEMM + # (default-on): ~1.5x on the projection and one launch instead of two. + # PRIMUS_COMPRESS_FUSE_PROJ=0 restores the two separate linears. + self._fuse_proj = os.environ.get("PRIMUS_COMPRESS_FUSE_PROJ", "1") != "0" + if self._fuse_proj: + self.wkv_gate = nn.Linear(hidden_size, 2 * proj_out, bias=False) + else: + self.wkv = nn.Linear(hidden_size, proj_out, bias=False) + self.wgate = nn.Linear(hidden_size, proj_out, bias=False) + + # Learnable absolute position embedding (APE) added on top of the + # softmax score. After overlap, the effective window length is + # ``2*ratio`` slots of size ``head_dim``; in non-overlap mode it's + # ``ratio`` slots of size ``head_dim``. + ape_len = 2 * ratio if self.overlap else ratio + self.ape = nn.Parameter(torch.zeros(ape_len, head_dim)) + nn.init.normal_(self.ape, std=0.02) + + self.kv_norm = LocalRMSNorm(head_dim, eps=rmsnorm_eps) + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Bridge checkpoints across the fused/unfused projection layouts. + + Old checkpoints store ``wkv.weight`` + ``wgate.weight``; the fused path + wants ``wkv_gate.weight`` = ``cat([wkv, wgate])`` (and vice-versa). Remap + in-place so either layout loads under either runtime setting. + """ + wkv_k, wgate_k, fused_k = prefix + "wkv.weight", prefix + "wgate.weight", prefix + "wkv_gate.weight" + if self._fuse_proj and wkv_k in state_dict and fused_k not in state_dict: + state_dict[fused_k] = torch.cat([state_dict.pop(wkv_k), state_dict.pop(wgate_k)], dim=0) + elif (not self._fuse_proj) and fused_k in state_dict and wkv_k not in state_dict: + w = state_dict.pop(fused_k) + state_dict[wkv_k], state_dict[wgate_k] = w[: self._proj_out], w[self._proj_out :] + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + # ------------------------------------------------------------------ + # internals + # ------------------------------------------------------------------ + + def _reshape_into_windows(self, t: torch.Tensor) -> torch.Tensor: + """``[B, S, coff*head_dim]`` → ``[B, N, ratio, coff*head_dim]``, + where ``N = S // ratio``. + """ + B, S, C = t.shape + assert S % self.ratio == 0, f"Compressor: sequence length {S} not divisible by ratio {self.ratio}" + N = S // self.ratio + return t.reshape(B, N, self.ratio, C) + + def _overlap_transform(self, t: torch.Tensor) -> torch.Tensor: + """``[B, N, ratio, 2*head_dim]`` → ``[B, N, 2*ratio, head_dim]``. + + For window ``i``, the augmented sequence is + ``[half_a[i], half_b[i-1]]`` concatenated along the per-window + axis. Window 0's "previous half" is filled with zeros (causal + padding). + """ + # Split channels. + half_a, half_b = torch.chunk(t, 2, dim=-1) # each [B, N, ratio, head_dim] + # Roll along the window dim so half_b[i] becomes "previous-window's b" of i+1. + half_b_prev = torch.cat( + [torch.zeros_like(half_b[:, :1]), half_b[:, :-1]], + dim=1, + ) + # Concat along the per-window token axis. + return torch.cat([half_a, half_b_prev], dim=2) # [B, N, 2*ratio, head_dim] + + # ------------------------------------------------------------------ + # public API + # ------------------------------------------------------------------ + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + """Pool ``hidden[B, S, D]`` to ``[B, S/ratio, head_dim]``.""" + if self._fuse_proj: + kv_proj, score_proj = self.wkv_gate(hidden).split(self._proj_out, dim=-1) + else: + kv_proj = self.wkv(hidden) # [B, S, coff*head_dim] + score_proj = self.wgate(hidden) # [B, S, coff*head_dim] + + kv = self._reshape_into_windows(kv_proj) # [B, N, ratio, coff*head_dim] + score = self._reshape_into_windows(score_proj) # [B, N, ratio, coff*head_dim] + + if self.overlap: + kv = self._overlap_transform(kv) # [B, N, 2*ratio, head_dim] + score = self._overlap_transform(score) # [B, N, 2*ratio, head_dim] + # else: kv / score already at [B, N, ratio, head_dim] + + # Per-window-softmax pool: APE bias + softmax over the window axis (dim=2) + # + weighted sum -- each compressed token is a softmax-weighted average of + # its window members. The forward burst (add + cast + softmax + cast + mul + # + reduce) is fused into one Triton launch on CUDA fp16/bf16/fp32 inputs; + # PRIMUS_COMPRESS_POOL_TRITON=0 (or non-CUDA / unsupported dtype) falls back + # to eager. + if ( + os.environ.get("PRIMUS_COMPRESS_POOL_TRITON", "1") != "0" + and kv.is_cuda + and kv.dtype + in ( + torch.float16, + torch.bfloat16, + torch.float32, + ) + ): + from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.compressor_pool import ( + fused_softmax_weighted_pool, + ) + + pooled = fused_softmax_weighted_pool(kv, score, self.ape) # [B, N, head_dim] + else: + score = score + self.ape # [B, N, win, head_dim] + weights = F.softmax(score.float(), dim=2).to(kv.dtype) + pooled = (kv * weights).sum(dim=2) # [B, N, head_dim] + + pooled = self.kv_norm(pooled) + return pooled + + +__all__ = ["Compressor"] diff --git a/primus/backends/megatron/core/transformer/deepseek_v4_attention.py b/primus/backends/megatron/core/transformer/deepseek_v4_attention.py new file mode 100644 index 000000000..543f7fbd9 --- /dev/null +++ b/primus/backends/megatron/core/transformer/deepseek_v4_attention.py @@ -0,0 +1,1652 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +DeepSeek-V4 attention. + +Plan-2 P13 — *faithful* attention rooted on Megatron's +``MLASelfAttention``. The released ``DeepSeek-V4-Flash`` checkpoint is +reproduced for **all three** layer types (``compress_ratio in {0, 4, 128}``) +inside a single attention class: + +* Single-latent KV: a single ``linear_kv`` projection ``hidden -> head_dim`` + produces both K and V, broadcast across all query heads. +* Per-head ``q_rms``: a parameter-less RMS normalization on ``head_dim`` + applied AFTER ``linear_q_up_proj`` and BEFORE partial RoPE — matches + the ``inference/model.py`` reference exactly. +* Grouped low-rank O projection: ``linear_o_a`` / ``linear_o_b`` (when + ``config.o_lora_rank > 0``) replace the standard flat ``linear_proj``. +* Learnable per-head ``attn_sink``: an extra "virtual key" column with + zero value, joined into the softmax. Drops the column after softmax + so the value-weighted sum is unaffected; the head can still spend mass + on the sink as a "no attention" fallback. +* Compressed branches (``compress_ratio > 0``) fold their compressor + (and indexer for CSA) in as :class:`Compressor` / :class:`Indexer` + spec submodules; the dense local SWA branch and the compressed branch + are softmax-joined together so the attention sink is shared across + both paths. +* Field names mirror MLA's canonical layout (``linear_q_down_proj``, + ``linear_q_up_proj``, ``q_layernorm``, ``kv_layernorm``) plus the V4 + extras (``linear_kv``, ``linear_o_a``, ``linear_o_b``, ``compressor``, + ``indexer``) so the state-dict adapter (P17) can map the released + safetensors keys + (``layers.{i}.attn.{wq_a,wq_b,wkv,q_norm,kv_norm,wo_a,wo_b,attn_sink, + compressor.*,indexer.*}``) in one straightforward table. The + per-head learnable softmax sink lives directly on the attention module + as ``self.attn_sink: nn.Parameter`` (no submodule slot — Plan-3 P21 + dropped the dead ``attn_sink`` field; the inline softmax-with-sink + path in :meth:`_attention_forward` is canonical). + +Forward signature: + +.. code-block:: python + + out = attn( + hidden, # [B, S, D] + position_ids, # [B, S] or [S] + ) + # out: [B, S, D] +""" + +from __future__ import annotations + +import atexit +import collections +import logging +import math +import os +import statistics +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.multi_latent_attention import MLASelfAttention +from megatron.core.transformer.spec_utils import ModuleSpec, build_module + +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, +) +from primus.backends.megatron.core.transformer.compressor import Compressor +from primus.backends.megatron.core.transformer.dual_rope import ( + DualRoPE, + apply_interleaved_partial_rope, +) +from primus.backends.megatron.core.transformer.indexer import Indexer +from primus.backends.megatron.core.transformer.local_rmsnorm import LocalRMSNorm +from primus.backends.megatron.core.transformer.sliding_window_kv import ( + sliding_window_causal_mask, +) + +# All attention backend entries come from the kernels package __init__ (the +# single entry point): eager, triton v1/v2. Naming: v4_attention_ +# (dense/HCA) and v4_csa_attention_ (CSA). The gluon backend is NOT imported +# here — it hard-depends on triton.experimental.gluon (gfx950 only) and is loaded +# lazily via load_gluon_attention_backends() only when a layer selects it. +from primus.backends.megatron.core.transformer.v4_attention_kernels import ( + eager_v4_attention, + eager_v4_csa_attention, + load_flydsl_attention_backends, + load_gluon_attention_backends, + load_gluon_v2_attention_backends, + load_gluon_v3_attention_backends, + load_turbo_attention_backends, + v4_attention_v1, + v4_attention_v2, + v4_csa_attention_v0, + v4_csa_attention_v1, + v4_csa_attention_v2, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rmsnorm import ( + fused_rms_norm, +) + +_SUPPORTED_COMPRESS_RATIOS = (0, 4, 128) + +logger = logging.getLogger(__name__) + + +def _require_gfx950() -> None: + """Assert the current device is gfx950 / CDNA4 before using the gluon backend. + + The gluon sparse-MLA kernels are hand-tuned for gfx950 (MI350/MI355X); running + them on any other arch is unsupported. Called only when a layer selects + ``use_v4_attention_backend`` / ``use_v4_csa_attention_backend = 'gluon'``. + """ + if not torch.cuda.is_available(): + raise RuntimeError( + "use_v4_attention_backend / use_v4_csa_attention_backend = 'gluon' requires a " + "CUDA/HIP gfx950 (CDNA4) device, but no accelerator is available. Select " + "eager | triton_v1 | triton_v2 instead." + ) + arch = str(getattr(torch.cuda.get_device_properties(0), "gcnArchName", "")) + if "gfx950" not in arch: + raise RuntimeError( + "use_v4_attention_backend / use_v4_csa_attention_backend = 'gluon' targets gfx950 / " + f"CDNA4 (MI350/MI355X); got device arch {arch!r}. Select eager | triton_v1 | triton_v2 " + "instead, or run on gfx950." + ) + + +# --------------------------------------------------------------------------- +# P32 diagnostic: collect in-context cuda.Event timings of v4_attention_v1 +# --------------------------------------------------------------------------- + + +class _DeepseekV4AttentionDiag: + """Accumulator for ``PRIMUS_V4_DIAG_TIME=1`` per-call timings.""" + + _per_mode: dict[str, list[float]] = collections.defaultdict(list) + _registered: bool = False + shape_logged: dict[str, bool] = {} + + @classmethod + def record(cls, *, mode: str, ms: float, swa: int) -> None: + cls._per_mode[mode].append(ms) + if not cls._registered: + cls._registered = True + atexit.register(cls.dump) + + @classmethod + def dump(cls) -> None: + if not cls._per_mode: + return + rank = os.environ.get("RANK", "0") + try: + local_rank = int(rank) + except (TypeError, ValueError): + local_rank = 0 + if local_rank != 0: + return + print("[PRIMUS_V4_DIAG_TIME] v4_attention_v1 inline cuda.Event timings:", flush=True) + for mode, vs in cls._per_mode.items(): + if not vs: + continue + # Drop first 3 to skip warmup. + stable = vs[3:] if len(vs) > 3 else vs + print( + f" mode={mode:<6s} n={len(vs):4d} " + f"all_med={statistics.median(vs):7.3f}ms " + f"warm_med={statistics.median(stable):7.3f}ms " + f"warm_min={min(stable):7.3f}ms " + f"warm_max={max(stable):7.3f}ms", + flush=True, + ) + + +# --------------------------------------------------------------------------- +# Spec submodules — V4 (plan-2 / MLA-canonical) +# --------------------------------------------------------------------------- + + +@dataclass +class DeepseekV4AttentionSubmodules: + """Spec submodules for the plan-2 :class:`DeepseekV4Attention`. + + The names follow MLA's canonical layout where they overlap (so that + Megatron's standard tensor-parallel / sequence-parallel / TE machinery + can apply unchanged), plus V4-specific extras for the single-latent KV + and grouped low-rank O. + + Provider-built shapes: + + * ``linear_q_down_proj`` : ``hidden -> q_lora_rank`` (= ``wq_a``) + * ``q_layernorm`` : RMSNorm on ``q_lora_rank`` (= ``q_norm``) + * ``linear_q_up_proj`` : ``q_lora_rank -> n_heads * head_dim`` (= ``wq_b``) + * ``linear_kv`` : ``hidden -> head_dim`` (= ``wkv``, + single latent — broadcast to all heads) + * ``kv_layernorm`` : RMSNorm on ``head_dim`` (= ``kv_norm``) + * ``linear_o_a`` : ``(n_heads * head_dim / o_groups) -> o_groups * o_lora_rank`` + * ``linear_o_b`` : ``o_groups * o_lora_rank -> hidden`` + * ``compressor`` : :class:`Compressor` (compress_ratio > 0 only) + * ``indexer`` : :class:`Indexer` (compress_ratio == 4 only) + + When the spec provider supplies ``linear_proj`` (instead of grouped + ``linear_o_a`` / ``linear_o_b``) the attention falls back to MLA's + standard flat output projection — useful for unit tests and the + ``o_lora_rank == 0`` fast-path config. + + Plan-3 P21: there is no ``attn_sink`` submodule slot. The per-head + learnable sink is :class:`torch.nn.Parameter` ``self.attn_sink`` + on the attention module itself (key ``layers.{i}.attn.attn_sink`` + in the released checkpoint), and the softmax-with-sink combine is + inlined in :meth:`DeepseekV4Attention._attention_forward`. + + Plan-3 P22: ``core_attention`` is the Turbo / TE flash-attention + kernel. Only the dense layer kind (``compress_ratio == 0``) emits + a spec for this slot — HCA / CSA cannot use a stock flash-attn + kernel (HCA needs a joint sink across two key streams which would + require an LSE-returning flash kernel; CSA needs per-query top-K + indexed keys which is not a flash pattern). When the dense path + receives a ``core_attention`` it bypasses the eager-Python softmax + and runs through ``provider.core_attention()`` instead. When + ``provider.core_attention()`` returns + :class:`PrimusTurboAttention` (i.e. ``use_turbo_attention=True``) + and V4's ``attn_sink`` is on, the attention module aliases + ``core_attention.sinks`` to ``self.attn_sink`` so the released + checkpoint key path is preserved. + """ + + linear_q_down_proj: Optional[Union[ModuleSpec, type]] = None + linear_q_up_proj: Optional[Union[ModuleSpec, type]] = None + linear_kv: Optional[Union[ModuleSpec, type]] = None + linear_o_a: Optional[Union[ModuleSpec, type]] = None + linear_o_b: Optional[Union[ModuleSpec, type]] = None + linear_proj: Optional[Union[ModuleSpec, type]] = None # fallback flat O + q_layernorm: Optional[Union[ModuleSpec, type]] = None + kv_layernorm: Optional[Union[ModuleSpec, type]] = None + compressor: Optional[Union[ModuleSpec, type]] = None + indexer: Optional[Union[ModuleSpec, type]] = None + # Plan-3 P22: dense (compress_ratio == 0) layers only. + core_attention: Optional[Union[ModuleSpec, type]] = None + + +# --------------------------------------------------------------------------- +# Build helpers +# --------------------------------------------------------------------------- + + +def _build_projection( + submodule: Optional[Union[ModuleSpec, type]], + *, + in_features: int, + out_features: int, +) -> nn.Module: + """Build a linear projection from a spec submodule. + + When the spec is ``None`` (CPU unit tests that exercise the + forward pass without a TP group) we instantiate a plain + :class:`nn.Linear` with the same shape. When a spec is supplied + we delegate to :func:`build_module` and let any constructor + failure bubble up — Plan-3 P21 retired the ``try/except/return + nn.Linear`` fallback because it produced an unsharded model + (vanilla ``nn.Linear`` instead of column / row parallel shards) + that silently masked spec bugs at TP=1 and would diverge at TP>1. + """ + if submodule is None: + return nn.Linear(in_features, out_features, bias=False) + return build_module(submodule) + + +def _projection_forward(proj: nn.Module, x: torch.Tensor) -> torch.Tensor: + """Run a projection and unwrap Megatron's ``(out, bias)`` tuple.""" + out = proj(x) + if isinstance(out, tuple): + return out[0] + return out + + +def _v4_o_a_fp8_enabled(config) -> bool: + """Whether to run the grouped-O ``o_a`` down-projection in MXFP8. + + ``o_a`` is a per-group (batched) matmul done as a manual einsum on + ``linear_o_a.weight``, so it bypasses the fp8 linear path and stays bf16. + When PRIMUS_V4_FP8_ATTN_PROJ is set (and TP=1, where the surrounding + projections are already routed to fp8) and the layer is in turbo-fp8, run + it as per-group fp8 GEMMs instead. Default off. + """ + if os.environ.get("PRIMUS_V4_FP8_ATTN_PROJ", "0") != "1": + return False + if getattr(config, "tensor_model_parallel_size", 1) != 1: + return False + try: + from primus.backends.megatron.core.extensions.primus_turbo import ( + PrimusTurboLowPrecisionGlobalStateManager as _M, + ) + + return _M.is_turbo_fp8_enabled() + except Exception: + return False + + +def _fp8_grouped_o_a(attn_g: torch.Tensor, wo_a_w: torch.Tensor) -> torch.Tensor: + """Fused MXFP8 grouped-O ``o_a`` down-projection (replaces the bf16 einsum). + + ``attn_g`` [B,S,G,d], ``wo_a_w`` [G,r,d] -> [B,S,G,r]. The G groups are an + independent batched matmul ``[B*S,d] @ [d,r]`` per group; we run them as a + SINGLE ``grouped_gemm_fp8`` (one fused Triton launch) instead of a per-group + ``gemm_fp8`` loop (G launches + G× quant). Stack tokens group-major into + ``[G*B*S, d]`` with ``group_lens=[B*S]*G``; weight ``[G,d,r]`` (trans_b=False). + d=(H*head_dim)/G and B*S are multiples of 32, so the MX block scale is clean. + """ + import primus_turbo.pytorch as pt + + from primus.backends.megatron.core.extensions.primus_turbo import ( + PrimusTurboLowPrecisionGlobalStateManager as _M, + ) + + cfg = _M.get_turbo_quant_config().data() + B, S, G, d = attn_g.shape + r = wo_a_w.shape[1] + a = attn_g.permute(2, 0, 1, 3).reshape(G * B * S, d).contiguous() # [G*BS, d], group-major + b = wo_a_w.transpose(1, 2).contiguous() # [G, d, r] (K=d, N=r, trans_b=False) + group_lens = torch.full((G,), B * S, dtype=torch.int64, device=a.device) + out = pt.ops.grouped_gemm_fp8(a, b, group_lens, trans_b=False, config=cfg) # [G*BS, r] + return out.reshape(G, B, S, r).permute(1, 2, 0, 3) # [B, S, G, r] + + +def _coerce_optional_bool_flag(value: object, *, field_name: str) -> bool: + """Coerce a possibly-stringified yaml flag to a clean ``bool``. + + Yaml interpolation like ``${PRIMUS_FOO:false}`` resolves to the + STRING ``"false"`` when the env var is unset, and the naive + ``bool("false") is True`` would silently flip a default-off knob + to on. Accept the common string spellings explicitly and treat + everything else as truthy/falsy via the normal ``bool(...)`` rule. + + Plan-8 P57 close-out 2 added this helper for the new + ``use_v4_tilelang_*`` flags; existing flags + (``use_v4_triton_*`` / ``use_v4_compiled_sinkhorn``) avoid the + issue because the V4 run scripts always pass them via + ``-- "False"`` and the override parser coerces to a Python + ``False`` before the config ever sees a string. + """ + if isinstance(value, bool): + return value + if isinstance(value, str): + lowered = value.strip().lower() + if lowered in ("0", "false", "no", "off", ""): + return False + if lowered in ("1", "true", "yes", "on"): + return True + raise ValueError( + f"Unrecognised string value for boolean config flag " + f"{field_name!r}: {value!r}; expected one of " + "'true' / 'false' / '1' / '0' / 'yes' / 'no' / 'on' / 'off'." + ) + return bool(value) + + +def _per_head_rms_norm(x: torch.Tensor, *, eps: float) -> torch.Tensor: + """Parameter-less per-head RMSNorm. + + Mirrors the released ``inference/model.py`` reference: + + .. code-block:: python + + q_rms = torch.rsqrt(q.float().square().mean(-1, keepdim=True) + eps) + q = (q.float() * q_rms).to(q.dtype) + + There is no learnable ``gamma`` — the per-head scale is "absorbed" + into the surrounding ``linear_q_up_proj`` weights at training time. + The check confirmed the released checkpoint has no separate + ``q_rms.weight`` parameter. + + Small-kernel-fusion (2026-07-03): the eager chain (bf16->fp32 cast + + square + mean + rsqrt + mul + fp32->bf16 cast, ~6 kernels / call × + 8 attention layers) is collapsed into one Triton FWD + one BWD kernel + via :func:`fused_rms_norm` (parameter-less, ``out_dtype = in_dtype``). + Gated by ``PRIMUS_RMSNORM_TRITON`` (default on); the dispatcher falls + back to the bit-identical eager body on CPU / when the knob is off. + """ + return fused_rms_norm(x, None, eps=eps, mid_cast=False, out_dtype=x.dtype) + + +def _build_local_rms_norm(dim: int, *, eps: float) -> nn.Module: + """Tiny CPU-friendly RMSNorm used as a fallback when no spec is given. + + Plan-2 P17 retired the closure-built ``_RMSNorm`` helper here; the + canonical implementation lives in + :class:`primus.backends.megatron.core.transformer.local_rmsnorm.LocalRMSNorm` + so the same code path is shared with + :mod:`...deepseek_v4_block` and :mod:`...compressor`. + """ + return LocalRMSNorm(dim=dim, eps=eps) + + +# --------------------------------------------------------------------------- +# DeepseekV4Attention (faithful, MLA-rooted, dense + CSA + HCA) +# --------------------------------------------------------------------------- + + +class DeepseekV4Attention(MLASelfAttention): + """V4 attention faithful to the released ``DeepSeek-V4-Flash`` checkpoint. + + Subclasses :class:`MLASelfAttention` for type identity (so downstream + Megatron isinstance checks treat V4 attention as an MLA variant) but + overrides ``__init__`` and ``forward`` because V4's parameter layout + differs from MLA's compressed-KV form: + + * V4 has **no** ``linear_kv_down_proj`` / ``linear_kv_up_proj`` — the + KV is single-latent (``wkv``) and shared as both K and V. + * V4's ``linear_proj`` is replaced by grouped low-rank + ``linear_o_a`` / ``linear_o_b`` (when ``config.o_lora_rank > 0``). + * V4 adds a per-head parameter-less ``q_rms`` and a learnable + ``attn_sink``. + * V4 layers come in three flavours selected by ``compress_ratio``: + + * ``0`` — dense / SWA over local KV. + * ``128`` — HCA: local SWA *plus* a fully-visible compressed pool + (Compressor in non-overlap mode). + * ``4`` — CSA: local SWA *plus* a per-query top-K selection over + a compressed pool (Compressor in overlap mode + Indexer). + + Because the parent's ``__init__`` builds modules we don't want, we + skip the MLA / Attention init chain and call ``nn.Module.__init__`` + directly. V4-shape modules are built from the spec submodules. + + **Plan-4 P27 — kernel dispatch precedence.** + + The softmax-and-attend kernel each layer fires through is selected + in :meth:`forward` / :meth:`_csa_forward` based on three independent + config flags (``use_turbo_attention``, ``use_v4_triton_attention``, + ``use_v4_triton_csa_attention``). The layer-kind-specific + precedence is: + + .. code-block:: text + + compress_ratio == 0 (dense / SWA, single key axis): + use_turbo_attention > use_v4_triton_attention > eager + (-> self.core_attention) (-> v4_attention_v1) (-> _attention_forward) + + compress_ratio == 128 (HCA: local SWA + full compressed pool): + use_v4_triton_attention > eager + (-> v4_attention_v1, (-> _attention_forward + HCA path with joint with cat([local, pool]) + [local | pool] mask) additive mask) + + ``use_turbo_attention`` does NOT route HCA — Turbo's + flash-attn returns no LSE so the joint local+pool softmax + cannot be decomposed into two flash calls. + + compress_ratio == 4 (CSA: local SWA + per-query top-K gather): + use_v4_triton_csa_attention > eager + (-> v4_csa_attention_v0) (-> _csa_forward eager) + + Neither ``use_turbo_attention`` nor + ``use_v4_triton_attention`` applies to CSA — the per-query + top-K gather (``gathered = pool[..., topk_idxs, :]``) is + sparse-per-row indexed attention with no flash-attn + equivalent. + + Auto-disable rules (init-side, fail-loud): + + * ``use_v4_triton_attention=True`` + ``compress_ratio == 4`` → + auto-disabled (CSA layers must opt in via the separate flag). + * ``use_v4_triton_csa_attention=True`` + ``compress_ratio != 4`` → + auto-disabled (the dense / HCA flag is ``use_v4_triton_attention``). + + On rank 0 each ``__init__`` emits one ``INFO`` log line through + :meth:`_log_kernel_choice` summarising the dispatch outcome for + the layer (e.g. ``[V4-attn] Layer 17: cr=128, kernel = v4_attention_v1 + (Triton, HCA path)``) so smoke / training logs unambiguously show + which kernel each layer is firing through. + """ + + def __init__( + self, + config: DeepSeekV4TransformerConfig, + *, + rope: DualRoPE, + compress_ratio: int = 0, + submodules: Optional[DeepseekV4AttentionSubmodules] = None, + layer_number: Optional[int] = None, + pg_collection=None, + attn_mask_type=None, + **kwargs, + ) -> None: + # We deliberately bypass the MLA / Attention parent __init__ chain + # because V4's KV layout differs from MLA's compressed-KV form. + # The class still subclasses MLASelfAttention for type identity so + # that ``isinstance(layer.self_attention, MLASelfAttention)`` keeps + # working in the Megatron stack. + # + # Plan-2 P16: ``attn_mask_type`` is accepted (and ignored) so the + # attention spec can declare a value that satisfies upstream + # :class:`MultiTokenPredictionLayer`'s pre-build validator; V4 + # manages its own SWA / sink mask internally. ``**kwargs`` swallows + # any forward-compatible kwargs upstream may add (e.g. + # ``cp_comm_type``) so the spec lifecycle keeps working. + del attn_mask_type, kwargs + nn.Module.__init__(self) + + if compress_ratio not in _SUPPORTED_COMPRESS_RATIOS: + raise ValueError( + f"DeepseekV4Attention supports compress_ratio in " + f"{_SUPPORTED_COMPRESS_RATIOS} (got {compress_ratio})." + ) + + hidden_size = int(config.hidden_size) + num_heads = int(config.num_attention_heads) + head_dim = int(config.kv_channels) + rotary_dim = int(config.qk_pos_emb_head_dim) + attn_sliding_window = int(config.attn_sliding_window) + attn_sink_enabled = bool(config.attn_sink) + attn_dropout = float(config.attention_dropout) + norm_eps = float(getattr(config, "norm_epsilon", None) or config.layernorm_epsilon) + q_lora_rank = int(config.q_lora_rank or 0) + o_groups = int(getattr(config, "o_groups", 1)) + o_lora_rank = int(getattr(config, "o_lora_rank", 0)) + + if q_lora_rank <= 0: + # V4 always uses a Q LoRA path; drop the no-LoRA branch to + # keep the math aligned with the checkpoint. + raise ValueError( + "DeepseekV4Attention requires config.q_lora_rank > 0; " + "V4 always low-rank-projects Q via wq_a / wq_b." + ) + + if num_heads * head_dim % max(o_groups, 1) != 0: + raise ValueError( + f"num_heads * head_dim ({num_heads * head_dim}) must be divisible " + f"by o_groups ({o_groups})" + ) + + self.config = config + self.compress_ratio = int(compress_ratio) + self.layer_number = int(layer_number) if layer_number is not None else 0 + self.pg_collection = pg_collection + + # ---- shape fields (read by helpers in this class) ---- + self.hidden_size = hidden_size + self.num_heads = num_heads + self.num_attention_heads_per_partition = num_heads + self.num_query_groups_per_partition = 1 # single-latent KV + self.head_dim = head_dim + self.rotary_dim = rotary_dim + self.q_head_dim = head_dim # MLA convention; here qk_head_dim + qk_pos_emb_head_dim == head_dim + self.attn_sliding_window = attn_sliding_window + self.attn_dropout = attn_dropout + self.q_lora_rank = q_lora_rank + self.o_groups = max(o_groups, 1) + self.o_lora_rank = o_lora_rank + self.norm_eps = norm_eps + + # Shared dual-RoPE (held by reference; not registered to avoid + # double-counting parameters across attention layers). + self._rope = [rope] + + submodules = submodules or DeepseekV4AttentionSubmodules() + self._submodules = submodules + + # ---- Q branch: hidden -> q_lora_rank -> n_heads * head_dim ---- + self.linear_q_down_proj = _build_projection( + submodules.linear_q_down_proj, + in_features=hidden_size, + out_features=q_lora_rank, + ) + if submodules.q_layernorm is None: + self.q_layernorm = _build_local_rms_norm(q_lora_rank, eps=norm_eps) + else: + self.q_layernorm = build_module( + submodules.q_layernorm, + hidden_size=q_lora_rank, + config=config, + eps=norm_eps, + ) + self.linear_q_up_proj = _build_projection( + submodules.linear_q_up_proj, + in_features=q_lora_rank, + out_features=num_heads * head_dim, + ) + + # ---- KV branch: single-latent ``wkv`` ---- + self.linear_kv = _build_projection( + submodules.linear_kv, + in_features=hidden_size, + out_features=head_dim, + ) + if submodules.kv_layernorm is None: + self.kv_layernorm = _build_local_rms_norm(head_dim, eps=norm_eps) + else: + self.kv_layernorm = build_module( + submodules.kv_layernorm, + hidden_size=head_dim, + config=config, + eps=norm_eps, + ) + + # ---- O projection ---- + # Two paths: + # - Grouped low-rank (V4 release): linear_o_a + linear_o_b + # - Flat MLA-style: linear_proj (used when o_lora_rank == 0) + if o_lora_rank > 0: + n_per_group = num_heads * head_dim // self.o_groups + self.linear_o_a = _build_projection( + submodules.linear_o_a, + in_features=n_per_group, + out_features=self.o_groups * o_lora_rank, + ) + self.linear_o_b = _build_projection( + submodules.linear_o_b, + in_features=self.o_groups * o_lora_rank, + out_features=hidden_size, + ) + self.linear_proj = None + else: + self.linear_o_a = None + self.linear_o_b = None + self.linear_proj = _build_projection( + submodules.linear_proj, + in_features=num_heads * head_dim, + out_features=hidden_size, + ) + + # ---- attention sink ---- + # The released checkpoint stores ``attn_sink`` as a [num_heads] + # learnable parameter directly on the attention module (key + # ``layers.{i}.attn.attn_sink`` — no wrapping submodule). + # We register it as ``self.attn_sink`` so the state-dict key + # matches the released checkpoint exactly; the softmax-with-sink + # combine is inlined in :meth:`_attention_forward`. + # + # Plan-3 P21 retired the optional ``self.attn_sink_module`` + # build branch (and the ``submodules.attn_sink`` slot) — the + # branch was never exercised in the forward path and its + # ``try/except`` masked AttentionSink build failures. A future + # TE-fused sink primitive can land as a new spec field once it + # actually replaces the inline path. + if attn_sink_enabled: + self.attn_sink = nn.Parameter(torch.zeros(num_heads)) + else: + self.register_parameter("attn_sink", None) + + # ---- compressor / indexer (compressed branches only) ---- + self.compressor: Optional[nn.Module] = None + self.indexer: Optional[nn.Module] = None + if self.compress_ratio > 0: + self.compressor = self._build_compressor(submodules.compressor) + if self.compress_ratio == 4: + self.indexer = self._build_indexer(submodules.indexer) + # The Indexer is a non-differentiable top-K *selector*: the only + # consumed output is ``topk_idxs`` (argTopK indices); its scores + # are discarded (``topk_idxs, _ = self.indexer(...)`` in forward) + # and this model has no indexer auxiliary/distillation loss, so + # none of the Indexer's params can ever receive a gradient. + # Leaving them trainable inserts permanently-dead params into the + # distributed-optimizer grad buckets, which both wastes grad / + # optimizer state + cross-node grad-sync bandwidth AND trips + # Megatron's overlap_grad_reduce invariant (every bucket param + # must fire its grad-ready backward hook) -- the latter is what + # forced overlap_grad_reduce/param_gather OFF and crippled + # cross-node DP scaling. Freeze them so Megatron excludes them + # from the grad buckets entirely. Set PRIMUS_V4_INDEXER_TRAINABLE=1 + # to re-enable (e.g. once an indexer aux loss is added). + if os.environ.get("PRIMUS_V4_INDEXER_TRAINABLE", "0") != "1": + for _indexer_param in self.indexer.parameters(): + _indexer_param.requires_grad_(False) + + # ---- core attention (Turbo / TE flash) — dense layers only ---- + # Plan-3 P22: when the spec emits a ``core_attention`` submodule + # (only on dense ``compress_ratio == 0`` layers), build it now and + # use it as the softmax-and-attend kernel instead of the + # eager-Python ``_attention_forward``. HCA + CSA always run + # eager-Python because their joint softmax / per-query top-K + # gather can't be expressed as a stock flash-attn call (see + # comments in ``forward`` / ``_csa_forward``). + # + # The constant ``softmax_scale`` is precomputed via + # ``_attention_scale()`` (the YaRN ``m_scale`` is a layer-static + # constant set at RoPE init time, so this matches the eager-path + # scale exactly for ``compress_ratio == 0``). + # Plan-4 P25: in-tree Primus Triton kernel for cr ∈ {0, 128}. + # Read the config flag once at __init__ so ``forward`` only does + # a cheap attribute load. Precedence in ``forward`` is + # ``use_turbo_attention > use_v4_triton_attention > eager``. + # ---- attention backend selection (unified string selectors) ---- + # ``use_v4_attention_backend`` selects the dense (cr=0) / HCA (cr=128) + # kernel; ``use_v4_csa_attention_backend`` selects the CSA (cr=4) kernel. + # ``use_turbo_attention`` (built as ``core_attention`` below) still takes + # precedence for the dense path when it can be built. + _ATTN_BACKENDS = ( + "eager", + "triton_v1", + "triton_v2", + "gluon", + "gluon_v2", + "gluon_v3", + "flydsl_v1", + "turbo", + ) + _CSA_BACKENDS = ( + "eager", + "triton_v0", + "triton_v1", + "triton_v2", + "gluon", + "gluon_v2", + "gluon_v3", + "flydsl_v0", + "flydsl_v1", + "turbo", + ) + self._attn_backend: str = str(getattr(config, "use_v4_attention_backend", "triton_v1") or "triton_v1") + self._csa_backend: str = str( + getattr(config, "use_v4_csa_attention_backend", "triton_v1") or "triton_v1" + ) + if self._attn_backend not in _ATTN_BACKENDS: + raise ValueError( + f"use_v4_attention_backend={self._attn_backend!r} is not a valid dense/HCA backend; " + f"expected one of {_ATTN_BACKENDS}" + ) + if self._csa_backend not in _CSA_BACKENDS: + raise ValueError( + f"use_v4_csa_attention_backend={self._csa_backend!r} is not a valid CSA backend; " + f"expected one of {_CSA_BACKENDS}" + ) + + # gluon is a gfx950/CDNA4-only backend with a hard triton.experimental.gluon + # dependency. Load it (and validate the arch) ONLY when a layer actually + # selects it, so non-gluon backends never pay the import and never crash on + # unsupported hardware / Triton builds. The loader raises a clear ImportError + # if the dependency is missing; ``_require_gfx950`` raises if the arch is wrong. + self._v4_attention_gluon = None + self._v4_csa_attention_gluon = None + if "gluon" in (self._attn_backend, self._csa_backend): + _require_gfx950() + self._v4_attention_gluon, self._v4_csa_attention_gluon = load_gluon_attention_backends() + + # gluon_v2 (2nd-gen gluon fwd+bwd) — same gfx950-only lazy-load contract as gluon. + self._v4_attention_gluon_v2 = None + self._v4_csa_attention_gluon_v2 = None + if "gluon_v2" in (self._attn_backend, self._csa_backend): + _require_gfx950() + self._v4_attention_gluon_v2, self._v4_csa_attention_gluon_v2 = load_gluon_v2_attention_backends() + + # gluon_v3 (3rd-gen optimized gluon fwd+bwd) — same gfx950-only lazy-load contract. + self._v4_attention_gluon_v3 = None + self._v4_csa_attention_gluon_v3 = None + if "gluon_v3" in (self._attn_backend, self._csa_backend): + _require_gfx950() + self._v4_attention_gluon_v3, self._v4_csa_attention_gluon_v3 = load_gluon_v3_attention_backends() + + # flydsl_v1 (native FlyDSL MFMA) is likewise a gfx950/CDNA4-only backend with + # a hard `flydsl` pip dependency; load + arch-validate only when selected. + self._v4_attention_flydsl = None + self._v4_csa_attention_flydsl = None + if "flydsl_v1" in (self._attn_backend, self._csa_backend): + _require_gfx950() + self._v4_attention_flydsl, self._v4_csa_attention_flydsl = load_flydsl_attention_backends() + + # turbo (Primus-Turbo native-FlyDSL sparse-MLA via the turbo API) — same gfx950-only + # lazy-load contract; hard-depends on the installed primus_turbo (flydsl attention) + flydsl. + self._v4_attention_turbo = None + self._v4_csa_attention_turbo = None + if "turbo" in (self._attn_backend, self._csa_backend): + _require_gfx950() + self._v4_attention_turbo, self._v4_csa_attention_turbo = load_turbo_attention_backends() + + self.core_attention: Optional[nn.Module] = None + self._use_core_attention: bool = False + if submodules.core_attention is not None and self.compress_ratio == 0: + softmax_scale = self._attention_scale() + try: + self.core_attention = build_module( + submodules.core_attention, + config=config, + layer_number=self.layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type="self", + softmax_scale=softmax_scale, + k_channels=head_dim, + v_channels=head_dim, + cp_comm_type="p2p", + pg_collection=self.pg_collection, + ) + except TypeError: + # Some core-attention classes (e.g. local CPU stubs in + # unit tests) don't accept the full TE / Turbo kwarg set. + # Retry with the minimal kwargs Megatron ships everywhere. + self.core_attention = build_module( + submodules.core_attention, + config=config, + layer_number=self.layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type="self", + softmax_scale=softmax_scale, + ) + + # Sink alias: when V4's per-head learnable sink is on AND the + # core-attention class supports learned sinks (Turbo only — + # the TE class does not), tie ``core_attention.sinks`` to + # ``self.attn_sink`` so the released-checkpoint key + # ``layers.{i}.attn.attn_sink`` keeps loading. TE classes + # that don't expose ``use_sink_attention`` get ``False`` here + # and we fall back to eager-Python so the inline + # softmax-with-sink path still produces the right math. + core_use_sink = bool(getattr(self.core_attention, "use_sink_attention", False)) + if attn_sink_enabled and core_use_sink: + # Cast V4's sink parameter to match the dtype Turbo allocated + # so the alias doesn't break dtype contracts. The eager + # path always promotes to float32 inside the softmax, so + # casting the parameter dtype is safe. + turbo_sinks = getattr(self.core_attention, "sinks", None) + if turbo_sinks is not None and turbo_sinks.dtype != self.attn_sink.dtype: + self.attn_sink.data = self.attn_sink.data.to(turbo_sinks.dtype) + self.core_attention.sinks = self.attn_sink + self._use_core_attention = True + elif not attn_sink_enabled and self.core_attention is not None: + # No-sink V4 still uses core_attention (e.g. unit tests, + # ablations). SWA is honored by Turbo only when sinks are + # on, so we accept this only when SWA is off too. + if self.attn_sliding_window <= 0: + self._use_core_attention = True + + # Plan-4 P27: surface the dispatch outcome in the training log + # so smoke / debug logs unambiguously show which kernel each + # layer is firing through (precedence is documented in the + # class docstring). Rank-0 only — every rank's own + # per-rank-file already captures the right entries. + self._log_kernel_choice() + + # ------------------------------------------------------------------ + # construction helpers (compressed branches) + # ------------------------------------------------------------------ + + def _log_kernel_choice(self) -> None: + """Emit one ``INFO`` log line summarising this layer's kernel choice. + + Plan-4 P27. Resolves the precedence outcome captured by + :meth:`forward` / :meth:`_csa_forward` (see class docstring) and + logs it once at ``__init__`` time so smoke / training logs + unambiguously show which kernel is firing for each layer. + + Rank-0 only when distributed; in single-process unit tests the + log fires unconditionally so ``caplog.at_level(logging.INFO)`` + captures it. We cannot use Megatron's ``print_rank_0`` here + because this module is also imported in CPU-only unit tests + where Megatron's parallel-state isn't initialised. + """ + try: + dist_initialized = torch.distributed.is_available() and torch.distributed.is_initialized() + except Exception: + dist_initialized = False + if dist_initialized and torch.distributed.get_rank() != 0: + return + + if self.compress_ratio == 0: + if self._use_core_attention: + kernel = "core_attention (Turbo / TE flash)" + else: + kernel = f"dense attention backend = {self._attn_backend}" + elif self.compress_ratio == 128: + kernel = f"HCA attention backend = {self._attn_backend}" + elif self.compress_ratio == 4: + kernel = f"CSA attention backend = {self._csa_backend}" + else: + # Defensive: __init__ already raises ValueError for unsupported + # compress_ratio, so this branch should be unreachable. + kernel = f"" + + logger.info( + "[V4-attn] Layer %s: cr=%s, kernel = %s", + self.layer_number, + self.compress_ratio, + kernel, + ) + + def _build_compressor(self, spec: Optional[Union[ModuleSpec, type]]) -> nn.Module: + """Build the V4 :class:`Compressor` for compressed branches. + + Plan-1 conventions (kept under V4): ``ratio=4`` → overlap mode + (CSA), ``ratio=128`` → non-overlap mode (HCA). The released + checkpoint hard-codes ``coff=2`` for overlap (CSA) and ``coff=1`` + for non-overlap (HCA); :class:`Compressor` enforces this through + its own ``overlap`` argument. + + When the spec is ``None`` (CPU unit tests, ``DeepseekV4Attention`` + constructed without a layer spec) we instantiate the local + :class:`Compressor` directly. Otherwise we delegate to + :func:`build_module` and let any constructor failure bubble up — + Plan-3 P21 retired the ``try/except/local Compressor`` fallback + because the spec passes the same :class:`Compressor` class and + the fallback handler was dead code that masked real spec bugs. + """ + kwargs = dict( + hidden_size=self.hidden_size, + head_dim=self.head_dim, + ratio=self.compress_ratio, + overlap=(self.compress_ratio == 4), + ) + if spec is None: + return Compressor(**kwargs) + return build_module(spec, **kwargs) + + def _build_indexer(self, spec: Optional[Union[ModuleSpec, type]]) -> nn.Module: + """Build the V4 :class:`Indexer` for the CSA branch. + + See :meth:`_build_compressor` for the spec-vs-fallback contract. + Plan-3 P21 retired the ``try/except/local Indexer`` fallback for + the same reason. + """ + index_topk = int(self.config.index_topk) + index_head_dim = int(self.config.index_head_dim) + index_n_heads = int(self.config.index_n_heads) + kwargs = dict( + hidden_size=self.hidden_size, + index_head_dim=index_head_dim, + index_n_heads=index_n_heads, + index_topk=index_topk, + compress_ratio=self.compress_ratio, + use_fp8_qk=bool(getattr(self.config, "use_v4_fp8_indexer", False)), + ) + if spec is None: + return Indexer(**kwargs) + return build_module(spec, **kwargs) + + # ------------------------------------------------------------------ + # internals + # ------------------------------------------------------------------ + + @property + def rope(self) -> DualRoPE: + return self._rope[0] + + def _attention_scale(self) -> float: + base = 1.0 / math.sqrt(self.head_dim) + rope_scale = self.rope.attn_scale(compress_ratio=self.compress_ratio) + return base * rope_scale + + def _apply_q(self, hidden: torch.Tensor) -> torch.Tensor: + """``[B, S, D]`` → ``[B, S, H, head_dim]`` (Q after q_norm + q_rms).""" + q_compressed = _projection_forward(self.linear_q_down_proj, hidden) + q_compressed = self.q_layernorm(q_compressed) + q = _projection_forward(self.linear_q_up_proj, q_compressed) + B, S, _ = q.shape + q = q.view(B, S, self.num_heads, self.head_dim) + # Per-head parameter-less RMS (matches `inference/model.py`). + q = _per_head_rms_norm(q, eps=self.norm_eps) + return q + + def _apply_kv(self, hidden: torch.Tensor) -> torch.Tensor: + """``[B, S, D]`` → ``[B, S, 1, head_dim]`` (single-latent K = V).""" + kv = _projection_forward(self.linear_kv, hidden) + kv = self.kv_layernorm(kv) + B, S, _ = kv.shape + return kv.view(B, S, 1, self.head_dim) + + def _apply_rope_q_k(self, q: torch.Tensor, k: torch.Tensor, position_ids: torch.Tensor): + """Apply partial RoPE (last ``rotary_dim`` channels) to Q and K + using the LAYER's compress_ratio (so CSA/HCA use the compress base + + YaRN; dense uses the main base).""" + q = self.rope.apply_rope(q, position_ids=position_ids, compress_ratio=self.compress_ratio) + k = self.rope.apply_rope(k, position_ids=position_ids, compress_ratio=self.compress_ratio) + return q, k + + def _local_mask(self, S: int, *, device, dtype) -> torch.Tensor: + """Mask for the local (SWA or full causal) branch. + + ``attn_sliding_window > 0`` enables sliding-window; ``0`` (the + default for unit tests / configs without SWA) gives full causal. + """ + window = self.attn_sliding_window if self.attn_sliding_window > 0 else 0 + if window > 0: + return sliding_window_causal_mask(S, window, device=device, dtype=dtype) + return sliding_window_causal_mask(S, S, device=device, dtype=dtype) + + def _append_sink_softmax(self, logits: torch.Tensor) -> torch.Tensor: + """Numerically-stable softmax with optional virtual-sink column. + + ``logits`` shape is ``[B, H, ..., Sk]`` — the head axis is at + ``dim=1``. Returns probabilities on the *real* keys (sink column + dropped) of the same shape as ``logits``. + """ + if self.attn_sink is None: + logits = logits - logits.amax(dim=-1, keepdim=True).detach() + return logits.softmax(dim=-1) + + # Build a sink column that broadcasts over all dims except the + # head axis (dim=1) and the key axis (dim=-1). + ndim = logits.dim() + view_shape = [1] * ndim + view_shape[1] = self.num_heads + view_shape[-1] = 1 + target_shape = list(logits.shape[:-1]) + [1] + sink_col = self.attn_sink.float().view(*view_shape).expand(*target_shape) + logits_aug = torch.cat([logits, sink_col], dim=-1) + logits_aug = logits_aug - logits_aug.amax(dim=-1, keepdim=True).detach() + probs = logits_aug.softmax(dim=-1) + return probs[..., :-1] + + def _attention_forward( + self, + q: torch.Tensor, # [B, H, Sq, head_dim] + k: torch.Tensor, # [B, H, Sk, head_dim] + v: torch.Tensor, # [B, H, Sk, head_dim] + attn_mask: torch.Tensor, # [Sq, Sk] additive (broadcasts over B,H) + ) -> torch.Tensor: + """Eager scaled-dot-product attention with optional attn_sink + for the dense / HCA paths (single key axis). + + Plan-4 P24: math lives in + :func:`primus...v4_attention_kernels._eager.reference.eager_v4_attention` + so the dense / HCA path, the plan-4 Triton kernel (P25), and the + plan-4 unit-test harness share one definition. The caller has + already pre-built the ``[Sq, Sk]`` additive mask (SWA-causal + for dense, ``cat([local_mask, hca_mask])`` for HCA) so we pass + ``swa_window=0`` and let the reference op use the supplied + ``additive_mask`` directly — bit-identical to the pre-P24 + inline implementation. + """ + return eager_v4_attention( + q, + k, + v, + sink=self.attn_sink, + swa_window=0, + additive_mask=attn_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + ) + + def _attention_forward_via_v4_triton( + self, + q: torch.Tensor, # [B, H, Sq, head_dim] + k: torch.Tensor, # [B, H, Sk, head_dim] + v: torch.Tensor, # [B, H, Sk, head_dim] + attn_mask: Optional[torch.Tensor], # [Sq, Sk] additive (broadcasts over B, H) + *, + swa_window: int = 0, + hca_local_seqlen: int = 0, + ) -> torch.Tensor: + """Run the dense / HCA softmax-and-attend through the plan-4 + in-tree :func:`v4_attention_v1` Triton kernel. + + Numerically equivalent to :meth:`_attention_forward` (same eager + ``q @ k^T * scale + mask + sink → softmax → @ v`` math) but + executes in a single fused kernel that re-materialises ``P`` + from the saved LSE during the BWD instead of storing the + ``[Sq, Sk]`` ``P`` tensor — important at full V4-Flash dims + (``S=4096`` ⇒ ``P`` is 32 MiB / microbatch). + + Plan-5 P30 flips the dense path to ``attn_mask=None`` + + ``swa_window > 0`` so the kernel can skip K tiles that are + guaranteed outside the sliding window. HCA uses the same pruning + for its local prefix by passing a pool-only mask plus + ``hca_local_seqlen``; the kernel then runs local SWA and pool + visibility as two loops under one joint softmax. + """ + # Plan-5 P32: opt-in microbench-vs-proxy timing harness, gated + # by ``PRIMUS_V4_DIAG_TIME=1``. Adds a synchronous cuda.Event + # span around the kernel call and dumps per-mode median/min/max + # at process exit (rank 0 only). Used to root-cause the dual-RoPE + # bf16 -> fp32 upcast bug that made every V4 attention kernel + # run 1.8-7x slower in the proxy than in the standalone bench; + # left in-tree for future microbench-vs-proxy regressions. + if os.environ.get("PRIMUS_V4_DIAG_TIME", "0") == "1": + mode = "hca" if hca_local_seqlen > 0 else "dense" + if not _DeepseekV4AttentionDiag.shape_logged.get(mode, False): + _DeepseekV4AttentionDiag.shape_logged[mode] = True + print( + f"[PRIMUS_V4_DIAG_TIME] mode={mode} " + f"q={tuple(q.shape)}/{q.dtype}/contig={q.is_contiguous()} " + f"k={tuple(k.shape)}/{k.dtype}/contig={k.is_contiguous()} " + f"v={tuple(v.shape)}/{v.dtype}/contig={v.is_contiguous()} " + f"swa={swa_window} hca_local={hca_local_seqlen}", + flush=True, + ) + torch.cuda.synchronize() + ev_s = torch.cuda.Event(enable_timing=True) + ev_e = torch.cuda.Event(enable_timing=True) + ev_s.record() + out = v4_attention_v1( + q, + k, + v, + sink=self.attn_sink, + swa_window=int(swa_window) if (attn_mask is None or hca_local_seqlen > 0) else 0, + additive_mask=attn_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + hca_local_seqlen=int(hca_local_seqlen), + ) + ev_e.record() + torch.cuda.synchronize() + _DeepseekV4AttentionDiag.record(mode=mode, ms=ev_s.elapsed_time(ev_e), swa=swa_window) + return out + return v4_attention_v1( + q, + k, + v, + sink=self.attn_sink, + swa_window=int(swa_window) if (attn_mask is None or hca_local_seqlen > 0) else 0, + additive_mask=attn_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + hca_local_seqlen=int(hca_local_seqlen), + ) + + def _attention_forward_via_core( + self, + q: torch.Tensor, # [B, S, H, head_dim] (post-RoPE) + kv: torch.Tensor, # [B, S, 1, head_dim] (post-RoPE, single-latent) + ) -> torch.Tensor: + """Run the dense (compress_ratio == 0) softmax-and-attend through + ``self.core_attention`` (Turbo flash-attn / TE flash-attn). + + Plan-3 P22. Avoids materialising the eager + ``[B, H, S, S] fp32`` logits tensor — at full V4-Flash dims + (``H=64, S=4096, hc_mult=4``) that's 16 GiB / microbatch, and + the dominant activation cost. + + Inputs use V4's local-frame layout (Q has all H heads, KV is + single-latent with 1 head). We forward as Turbo's required + ``qkv_format="sbhd"`` and let the underlying flash kernel + broadcast the 1-head KV across H query heads (MQA). Causal + masking + (optional) sliding window are honored by the kernel + directly — the eager ``local_mask`` is not used here. + + Returns ``[B, H, S, head_dim]`` to match the contract of + :meth:`_attention_forward`. + """ + B, S, H, Dh = q.shape + # [B, S, H, D] -> [S, B, H, D] (qkv_format="sbhd"). + q_sbhd = q.transpose(0, 1).contiguous() + kv_sbhd = kv.transpose(0, 1).contiguous() # [S, B, 1, D] + + # Turbo / TE flash-attn forward. ``attention_mask=None`` is + # legal for causal+SWA (the kernel builds the mask internally + # from ``attn_mask_type`` + the layer's ``window_size``). + out = self.core_attention( + q_sbhd, + kv_sbhd, + kv_sbhd, + None, + attn_mask_type=AttnMaskType.causal, + ) # -> [S, B, H * head_dim] + + # [S, B, H*D] -> [B, S, H, D] -> [B, H, S, D]. + out = out.view(S, B, H, Dh).permute(1, 2, 0, 3).contiguous() + return out + + def _grouped_o_projection(self, attn: torch.Tensor) -> torch.Tensor: + """Apply the V4 grouped low-rank O projection. + + Input ``attn`` shape: ``[B, S, H, head_dim]``. + Output shape: ``[B, S, hidden_size]``. + + Math (from ``inference/model.py``): + + .. code-block:: python + + # attn : [B, S, G, (H*head_dim)/G] + # wo_a.weight : [G * o_lora_rank, (H*head_dim)/G] + wo_a_w = self.linear_o_a.weight.view(G, o_lora_rank, -1) + o = einsum("bsgd,grd->bsgr", attn, wo_a_w) + o = self.linear_o_b(o.flatten(2)) + + We use the Linear's stored ``weight`` directly so the per-group + einsum semantics are exact. (Megatron's parallel linears expose + ``.weight`` after ``build_module``.) + """ + B, S, H, Dh = attn.shape + G = self.o_groups + attn_g = attn.reshape(B, S, G, (H * Dh) // G) # [B, S, G, H*Dh/G] + + wo_a = self.linear_o_a + weight = wo_a.weight if hasattr(wo_a, "weight") else None + if weight is None: + # Fall back to a dense linear apply (Megatron parallel linears + # without a directly accessible weight attribute). + o = _projection_forward(wo_a, attn_g.reshape(B, S, -1)) + o = o.view(B, S, G * self.o_lora_rank) + else: + wo_a_w = weight.view(G, self.o_lora_rank, (H * Dh) // G) + if _v4_o_a_fp8_enabled(self.config): + o = _fp8_grouped_o_a(attn_g, wo_a_w) # per-group MXFP8 + else: + o = torch.einsum("bsgd,grd->bsgr", attn_g, wo_a_w) + o = o.flatten(2) + return _projection_forward(self.linear_o_b, o) + + def _flat_o_projection(self, attn: torch.Tensor) -> torch.Tensor: + """MLA-style flat output projection (``o_lora_rank == 0`` fast path).""" + B, S, H, Dh = attn.shape + return _projection_forward(self.linear_proj, attn.reshape(B, S, H * Dh)) + + # ------------------------------------------------------------------ + # compressed branches (HCA / CSA) + # ------------------------------------------------------------------ + + def _build_compressed_pool(self, hidden: torch.Tensor) -> torch.Tensor: + """Run the compressor + compress-base partial RoPE. + + Returns ``[B, P, head_dim]`` where ``P = S // compress_ratio``. + """ + device = hidden.device + pooled = self.compressor(hidden) # [B, P, head_dim] + B, P = pooled.shape[0], pooled.shape[1] + + # Compress-base partial RoPE on compressed indices [0..P). Positions are + # the deterministic arange(P), so use the cached table instead of + # rebuilding arange -> outer -> cos/sin every forward. + cos, sin = self.rope.compress_rope.forward_arange(P, device) + cos = cos[..., : self.rotary_dim // 2] + sin = sin[..., : self.rotary_dim // 2] + cos = cos.unsqueeze(0).expand(B, -1, -1) + sin = sin.unsqueeze(0).expand(B, -1, -1) + pool_kv = pooled.unsqueeze(2) # [B, P, 1, head_dim] + pool_kv = apply_interleaved_partial_rope(pool_kv, cos, sin, rotary_dim=self.rotary_dim) + return pool_kv.squeeze(2) # [B, P, head_dim] + + def _hca_extra_kv( + self, + hidden: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build the HCA (compress_ratio == 128) compressed branch. + + Returns ``(extra_k_bh, extra_v_bh, extra_mask)`` where the + compressed pool is broadcast across H heads (single-latent + compressor output) and the additive mask is shape ``[S, P]`` + (broadcasts over B, H). + + Per the techblog: pool position ``s`` covers raw tokens + ``[s*ratio, (s+1)*ratio)``; query at raw token ``t`` may attend + to ``s`` iff ``(s+1)*ratio - 1 <= t``. + """ + B, S, _ = hidden.shape + device, dtype = hidden.device, hidden.dtype + pool = self._build_compressed_pool(hidden) # [B, P, head_dim] + P = pool.shape[1] + + # Broadcast pool across all H query-heads: [B, P, head_dim] -> [B, P, H, head_dim]. + pool_h = pool.unsqueeze(2).expand(B, P, self.num_heads, self.head_dim) + # Move heads dim to dim=1: [B, H, P, head_dim]. + pool_bh = pool_h.transpose(1, 2) + + extra_mask = self._hca_extra_mask_cached(S, P, device, dtype) + return pool_bh, pool_bh, extra_mask # K = V = compressed pool + + def _hca_extra_mask_cached(self, S: int, P: int, device, dtype): + """HCA additive causal mask ``[S, P]``, cached (data-independent). + + Pool slot ``s`` is visible to query ``t`` iff ``(s+1)*ratio - 1 <= t``; + the mask depends only on ``(S, P, compress_ratio, dtype)`` — all fixed + per run — so build it once instead of rebuilding arange + where every + compressed-layer forward. Bit-identical. PRIMUS_COMPRESS_MASK_CACHE=0 + forces the eager rebuild. + """ + if os.environ.get("PRIMUS_COMPRESS_MASK_CACHE", "1") == "0": + t = torch.arange(S, device=device).unsqueeze(1) + s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * self.compress_ratio - 1 + return torch.where(s_end <= t, 0.0, float("-inf")).to(dtype) + cache = getattr(self, "_hca_mask_cache", None) + if cache is None: + cache = self._hca_mask_cache = {} + key = (S, P, device, dtype) + m = cache.get(key) + if m is None: + t = torch.arange(S, device=device).unsqueeze(1) # [S, 1] + s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * self.compress_ratio - 1 # [1, P] + m = torch.where(s_end <= t, 0.0, float("-inf")).to(dtype) + cache[key] = m + return m + + def _csa_forward( + self, + hidden: torch.Tensor, + q_bh: torch.Tensor, # [B, H, S, head_dim] + k_local_bh: torch.Tensor, # [B, H, S, head_dim] + v_local_bh: torch.Tensor, # [B, H, S, head_dim] + local_mask: torch.Tensor, # [S, S] — built by caller; unused here, see below + ) -> torch.Tensor: + """CSA (compress_ratio == 4) joint local-SWA + sparse-compressed attention. + + The compressor produces a per-batch pool ``[B, P, head_dim]``, + the indexer picks ``index_topk`` pool positions per query, and + the attention runs softmax JOINTLY over ``[local_keys, sparse_keys]`` + so the optional ``attn_sink`` is shared across both branches. + + Plan-4 P24: the compressor / indexer / per-query top-K gather + stay here (they are V4-specific side-paths that the kernel does + not own); the joint-softmax math is delegated to + :func:`primus...v4_attention_kernels._eager.reference.eager_v4_csa_attention` + so the CSA path, the plan-4 CSA Triton kernel (P26), and the + plan-4 unit-test harness share one definition. ``local_mask`` is + retained in the signature for back-compat but unused — the + reference op rebuilds the local SWA mask deterministically from + ``swa_window`` (same call to + :func:`sliding_window_causal_mask` as :meth:`_local_mask` makes, + so the result is bit-identical). + + Plan-5 P31: when ``use_v4_triton_csa_attention=True`` the sparse + top-K pool gather moves into the Triton kernel. The eager fallback + still materialises ``gathered`` here so it remains the reference + implementation and keeps the old P26 API covered by unit tests. + """ + del local_mask # see docstring + B, H, S, Dh = q_bh.shape + dtype = hidden.dtype + + # 1) Compressed pool with compress-base RoPE. + pool = self._build_compressed_pool(hidden) # [B, P, head_dim] + P = pool.shape[1] + + # 2) Indexer top-K per query. + topk_idxs, _ = self.indexer(hidden) # [B, S, K] + # Dispatch on ``use_v4_csa_attention_backend``. gluon / triton_v2 / + # triton_v1 consume (pool, topk) directly; eager / triton_v0 / flydsl_v0 + # use the per-query gathered [B, S, K, Dh] representation. + be = self._csa_backend + if be == "gluon": + return self._v4_csa_attention_gluon( + q_bh, + k_local_bh, + v_local_bh, + pool, + topk_idxs=topk_idxs, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + ) + if be == "gluon_v2": + return self._v4_csa_attention_gluon_v2( + q_bh, + k_local_bh, + v_local_bh, + pool, + topk_idxs=topk_idxs, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + ) + if be == "gluon_v3": + return self._v4_csa_attention_gluon_v3( + q_bh, + k_local_bh, + v_local_bh, + pool, + topk_idxs=topk_idxs, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + ) + if be == "turbo": + return self._v4_csa_attention_turbo( + q_bh, + k_local_bh, + v_local_bh, + pool, + topk_idxs=topk_idxs, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + ) + if be == "triton_v2": + return v4_csa_attention_v2( + q_bh, + k_local_bh, + v_local_bh, + pool, + topk_idxs=topk_idxs, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + ) + if be == "flydsl_v1": + return self._v4_csa_attention_flydsl( + q_bh, + k_local_bh, + v_local_bh, + pool, + topk_idxs=topk_idxs, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + ) + if be == "triton_v1": + return v4_csa_attention_v1( + q_bh, + k_local_bh, + v_local_bh, + pool, + topk_idxs=topk_idxs, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + ) + + # eager / triton_v0 / flydsl_v0: build the per-query gathered slices. + K = topk_idxs.shape[-1] + valid = topk_idxs >= 0 # [B, S, K] + safe_idx = topk_idxs.clamp(min=0) + idx_expand = safe_idx.unsqueeze(-1).expand(B, S, K, Dh) + pool_expand = pool.unsqueeze(1).expand(B, S, P, Dh) + gathered = torch.gather(pool_expand, dim=2, index=idx_expand) + gathered = gathered * valid.unsqueeze(-1).to(gathered.dtype) + sparse_mask = torch.where(valid, 0.0, float("-inf")).to(dtype) # [B, S, K] + + if be in ("triton_v0", "flydsl_v0"): + return v4_csa_attention_v0( + q_bh, + k_local_bh, + v_local_bh, + gathered, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + sparse_mask=sparse_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + use_flydsl=(be == "flydsl_v0"), + ) + return eager_v4_csa_attention( + q_bh, + k_local_bh, + v_local_bh, + gathered, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + sparse_mask=sparse_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + ) + + # ------------------------------------------------------------------ + # public forward + # ------------------------------------------------------------------ + + def _attention_backend_forward(self, q_bh, k, v, *, additive_mask, hca_local_seqlen, S, device, dtype): + """Dense (cr=0) / HCA (cr=128) dispatch on ``use_v4_attention_backend``.""" + be = self._attn_backend + if be == "gluon": + return self._v4_attention_gluon( + q_bh, + k, + v, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + additive_mask=additive_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + hca_local_seqlen=hca_local_seqlen, + ) + if be == "gluon_v2": + return self._v4_attention_gluon_v2( + q_bh, + k, + v, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + additive_mask=additive_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + hca_local_seqlen=hca_local_seqlen, + ) + if be == "gluon_v3": + return self._v4_attention_gluon_v3( + q_bh, + k, + v, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + additive_mask=additive_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + hca_local_seqlen=hca_local_seqlen, + ) + if be == "turbo": + return self._v4_attention_turbo( + q_bh, + k, + v, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + additive_mask=additive_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + hca_local_seqlen=hca_local_seqlen, + ) + if be == "triton_v2": + return v4_attention_v2( + q_bh, + k, + v, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + additive_mask=additive_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + hca_local_seqlen=hca_local_seqlen, + ) + if be == "flydsl_v1": + return self._v4_attention_flydsl( + q_bh, + k, + v, + sink=self.attn_sink, + swa_window=int(self.attn_sliding_window), + additive_mask=additive_mask, + attn_dropout=self.attn_dropout, + training=self.training, + scale=self._attention_scale(), + hca_local_seqlen=hca_local_seqlen, + ) + if be == "triton_v1": + return self._attention_forward_via_v4_triton( + q_bh, + k, + v, + additive_mask, + swa_window=int(self.attn_sliding_window), + hca_local_seqlen=hca_local_seqlen, + ) + # eager + local_mask = self._local_mask(S, device=device, dtype=dtype) + mask = local_mask if additive_mask is None else torch.cat([local_mask, additive_mask], dim=-1) + return self._attention_forward(q_bh, k, v, mask) + + def forward( + self, + hidden: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + """``[B, S, D] -> [B, S, D]``. + + Dispatches on ``self.compress_ratio``: + + * ``0`` — dense / SWA over local KV (single key axis). + * ``128`` — HCA: concat compressed pool to local KV, joint softmax. + * ``4`` — CSA: per-query top-K from compressed pool, joint softmax. + """ + B, S, _ = hidden.shape + device, dtype = hidden.device, hidden.dtype + + q = self._apply_q(hidden) # [B, S, H, head_dim] + kv = self._apply_kv(hidden) # [B, S, 1, head_dim] + + # Partial RoPE on Q and K. K is post-RoPE; V uses the SAME tensor + # (V4's single-latent design: K and V share the rope-applied kv). + q, kv = self._apply_rope_q_k(q, kv, position_ids) + + if self.compress_ratio == 0 and self._use_core_attention: + # Plan-3 P22: dense layer, Turbo / TE flash path. Causal + SWA + # are handled inside the kernel; the eager ``local_mask`` is + # not consulted here. KV is forwarded as ``[S, B, 1, D]`` + # and broadcast across H query heads via MQA. + out_bh = self._attention_forward_via_core(q, kv) + out = out_bh.transpose(1, 2).contiguous() # [B, S, H, head_dim] + out = out.to(dtype=dtype) + if self.linear_o_a is not None: + return self._grouped_o_projection(out) + return self._flat_o_projection(out) + + # Broadcast K / V across the H query-head axis. + k_h = kv.expand(B, S, self.num_heads, self.head_dim) + v_h = kv.expand(B, S, self.num_heads, self.head_dim) + + # Move heads dim before sequence: [B, S, H, head_dim] -> [B, H, S, head_dim] + q_bh = q.transpose(1, 2) + k_local_bh = k_h.transpose(1, 2) + v_local_bh = v_h.transpose(1, 2) + + if self.compress_ratio == 0: + out_bh = self._attention_backend_forward( + q_bh, + k_local_bh, + v_local_bh, + additive_mask=None, + hca_local_seqlen=0, + S=S, + device=device, + dtype=dtype, + ) + elif self.compress_ratio == 128: + # HCA: the local SWA branch and the compressed-pool branch share ONE + # softmax with ONE sink column; concatenate the pool to the local + # keys and pass the pool-only additive mask. + extra_k_bh, extra_v_bh, extra_mask = self._hca_extra_kv(hidden) + k_full = torch.cat([k_local_bh, extra_k_bh], dim=2) # along Sk + v_full = torch.cat([v_local_bh, extra_v_bh], dim=2) + out_bh = self._attention_backend_forward( + q_bh, + k_full, + v_full, + additive_mask=extra_mask, + hca_local_seqlen=S, + S=S, + device=device, + dtype=dtype, + ) + elif self.compress_ratio == 4: + # CSA cannot use ``core_attention``: the per-query top-K + # gather (``gathered = pool[..., topk_idxs, :]``, shape + # ``[B, H, S, K, head_dim]``) is sparse-per-row indexed + # attention — there is no flash-attn kernel that reads a + # different per-query subset of keys from a pool. Stays on + # eager-Python under plan-3 (a custom kernel is required). + local_mask = self._local_mask(S, device=device, dtype=dtype) + out_bh = self._csa_forward(hidden, q_bh, k_local_bh, v_local_bh, local_mask) + else: + # Guarded by __init__; included for static-analysis completeness. + raise ValueError(f"Unsupported compress_ratio {self.compress_ratio}") + + out = out_bh.transpose(1, 2).contiguous() # [B, S, H, head_dim] + out = out.to(dtype=dtype) + + if self.linear_o_a is not None: + return self._grouped_o_projection(out) + return self._flat_o_projection(out) + + +__all__ = [ + "DeepseekV4Attention", + "DeepseekV4AttentionSubmodules", +] diff --git a/primus/backends/megatron/core/transformer/dual_rope.py b/primus/backends/megatron/core/transformer/dual_rope.py new file mode 100644 index 000000000..4744b889d --- /dev/null +++ b/primus/backends/megatron/core/transformer/dual_rope.py @@ -0,0 +1,362 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Dual-RoPE for DeepSeek-V4. + +Reference: techblog §6 ("RoPE: dual base + YaRN details"). + +V4 layers fall into two RoPE regimes, decided per-layer by +``compress_ratios[i]``: + +* ``compress_ratio == 0`` (dense / SWA layers): use the **main** RoPE base + (``rope_theta = 10000``), **no YaRN**. +* ``compress_ratio != 0`` (CSA / HCA layers): use the **compress** RoPE + base (``compress_rope_theta = 160000``) **with YaRN scaling** + (``factor=16, beta_fast=32, beta_slow=1, original_max_position_embeddings=65536``). + +Two important corrections over the HF Llama-style RoPE that V4 inherits +from the released weights: + +1. **Interleaved RoPE** (pairs ``(2k, 2k+1)``), **not** rotate-half pairs + ``(d, d+rd/2)``. NeMo's port did this; HF PR 45616 originally did not. +2. **Partial RoPE**: only the last ``qk_pos_emb_head_dim`` (= 64) channels of + each head are rotated. The first ``head_dim - 64`` channels stay nope. + +This module exposes: + +* :class:`DualRoPE` — owns two ``RoPECache`` instances (main + compress). + ``forward(layer_compress_ratio, x, position_ids, *, key=False)`` applies + the right partial-RoPE. +* :class:`RoPECache` — single-base RoPE, with optional YaRN scaling. +""" + +from __future__ import annotations + +import math +import os + +import torch +import torch.nn as nn + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial import ( + RoPEInterleavedPartialFn, + apply_rope_from_positions, +) + +# --------------------------------------------------------------------------- +# YaRN scaling +# --------------------------------------------------------------------------- + + +def _yarn_freq_scaling( + inv_freq: torch.Tensor, + *, + factor: float, + beta_fast: float, + beta_slow: float, + original_max_position_embeddings: int, +) -> torch.Tensor: + """Apply YaRN frequency-band scaling to ``inv_freq``. + + Each frequency component ``inv_freq[i] = 1 / theta**(2i/D)``. YaRN groups + the frequency band by wavelength ``2π / inv_freq[i]``: + + * High-freq band (wavelength < ``original_seq / beta_fast``): keep as-is. + * Low-freq band (wavelength > ``original_seq / beta_slow``): scale + ``inv_freq[i]`` by ``1/factor`` (i.e. lengthen wavelength). + * Mid-freq band: linear interpolation between the two. + + Standard YaRN reference; see e.g. NeMo's ``YarnRotaryPositionEmbedding``. + """ + if factor == 1.0 or original_max_position_embeddings <= 0: + return inv_freq + + # Wavelengths corresponding to each inv_freq component. + wavelens = 2.0 * math.pi / inv_freq # same shape as inv_freq + + low_thresh = original_max_position_embeddings / beta_fast + high_thresh = original_max_position_embeddings / beta_slow + + # Linear blend factor: 0 at high-freq side (no scaling), 1 at low-freq side (full /factor scaling). + smooth = ((wavelens - low_thresh) / max(high_thresh - low_thresh, 1e-12)).clamp(min=0.0, max=1.0) + + # Scaled inv_freq. + inv_freq_scaled = inv_freq / factor + return inv_freq * (1.0 - smooth) + inv_freq_scaled * smooth + + +def _yarn_attn_scale(factor: float) -> float: + """The standard YaRN attention magnitude scale ``m_scale``. + + ``m_scale = 0.1 * log(factor) + 1`` — used to scale the attention + softmax temperature when YaRN is on. + """ + if factor <= 1.0: + return 1.0 + return 0.1 * math.log(factor) + 1.0 + + +# --------------------------------------------------------------------------- +# RoPE cache (single base) +# --------------------------------------------------------------------------- + + +class RoPECache(nn.Module): + """Builds ``cos`` / ``sin`` tables on demand for a single RoPE base. + + Stores ``inv_freq`` as a (non-trainable) buffer so it follows the model's + device / dtype movements. ``cos`` / ``sin`` are computed lazily on + ``forward``. + """ + + def __init__( + self, + rotary_dim: int, + *, + theta: float, + yarn_factor: float = 1.0, + yarn_beta_fast: float = 32.0, + yarn_beta_slow: float = 1.0, + original_max_position_embeddings: int = 0, + ) -> None: + super().__init__() + if rotary_dim % 2 != 0: + raise ValueError(f"rotary_dim must be even, got {rotary_dim}") + + self.rotary_dim = rotary_dim + self.theta = float(theta) + self.yarn_factor = float(yarn_factor) + self.yarn_beta_fast = float(yarn_beta_fast) + self.yarn_beta_slow = float(yarn_beta_slow) + self.original_max_position_embeddings = int(original_max_position_embeddings) + + # inv_freq[i] = 1 / theta**(2i/rotary_dim), i in [0, rotary_dim/2) + i = torch.arange(0, rotary_dim, 2, dtype=torch.float32) + inv_freq = 1.0 / (theta ** (i / rotary_dim)) + inv_freq = _yarn_freq_scaling( + inv_freq, + factor=self.yarn_factor, + beta_fast=self.yarn_beta_fast, + beta_slow=self.yarn_beta_slow, + original_max_position_embeddings=self.original_max_position_embeddings, + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + # YaRN's m_scale; exposed for the caller to multiply attention scale. + self.attn_scale: float = _yarn_attn_scale(self.yarn_factor) + + # Memo of (cos, sin) tables keyed by (n, device) for arange(n) positions; + # see forward_arange. Not a buffer (recomputable, device-keyed). + self._arange_cache: dict = {} + + def forward(self, position_ids: torch.Tensor) -> tuple: + """Return (cos, sin) for the given positions. + + ``position_ids``: any integer/float tensor (typically ``[B, S]`` + or ``[S]``). + + Output ``(cos, sin)`` shape is ``position_ids.shape + (rotary_dim/2,)``; + the caller broadcasts over the head/batch axes. + """ + # outer product positions × inv_freq + freqs = position_ids.float().unsqueeze(-1) * self.inv_freq # [..., rotary_dim/2] + cos = freqs.cos() + sin = freqs.sin() + return cos, sin + + def forward_arange(self, n: int, device) -> tuple: + """``(cos, sin)`` for positions ``torch.arange(n)`` -- cached. + + Equivalent to ``self.forward(torch.arange(n, device=device))``, but + memoised by ``(n, device)``. The compressed-branch RoPE is always + evaluated at the deterministic positions ``arange(P)`` + (``P = S // compress_ratio``, fixed per run), so the table is identical + every forward; caching it skips the ``arange -> outer-product -> + cos/sin`` recompute each step (and per compressed layer). Set + ``PRIMUS_COMPRESS_ROPE_CACHE=0`` to disable the cache. + """ + if os.environ.get("PRIMUS_COMPRESS_ROPE_CACHE", "1") == "0": + return self.forward(torch.arange(n, device=device)) + key = (int(n), str(device)) + hit = self._arange_cache.get(key) + if hit is None: + hit = self.forward(torch.arange(n, device=device)) + self._arange_cache[key] = hit + return hit + + +# --------------------------------------------------------------------------- +# Partial interleaved RoPE application +# --------------------------------------------------------------------------- + + +def apply_interleaved_partial_rope( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + *, + rotary_dim: int, +) -> torch.Tensor: + """Apply RoPE to the **last** ``rotary_dim`` channels of ``x`` using the + **interleaved** pairing convention: pairs ``(2k, 2k+1)``. + + Args: + x: ``[..., head_dim]``. Typical layouts are ``[B, S, H, head_dim]`` + or ``[S, B, H, head_dim]``. + cos, sin: must broadcast to ``x[..., :rotary_dim//2]`` after the + internal pair-reshape (i.e. shape ``[..., rotary_dim/2]`` matching + x's leading dims with a singleton heads axis inserted as needed). + The simplest contract: pass shape ``[..., rotary_dim/2]`` where + the leading dims are exactly the ``position_ids`` shape; this + function will insert a singleton "heads" dim at position ``-2`` + so it broadcasts against the ``H`` axis of ``x``. + rotary_dim: the partial RoPE size; must be ``<= x.shape[-1]`` and + even. + + Returns: + ``x`` with the last ``rotary_dim`` channels rotated; first + ``head_dim - rotary_dim`` channels untouched. + """ + head_dim = x.shape[-1] + if rotary_dim > head_dim or rotary_dim % 2 != 0: + raise ValueError(f"rotary_dim must be even and <= head_dim ({head_dim}), got {rotary_dim}") + if rotary_dim == 0: + return x + + # Plan-6 P35: route through the fused Triton kernel when on CUDA / HIP + # and the env knob is not "0". The Triton path collapses the 9-op + # eager chain below (slice / reshape / four broadcast muls / stack / + # reshape / cat) into one kernel that does a single contiguous write + # with the rotation baked in -- removing the `CatArrayBatchedCopy_contig` + # bucket (~10 ms / 24 calls in the plan-5 P32 final trace) and the + # share of `elementwise_kernel_manual_unroll` that comes from the + # broadcast muls. Eager body kept in tree as the `PRIMUS_ROPE_TRITON=0` + # fallback and as the G38 unit-test reference. + if x.is_cuda and os.environ.get("PRIMUS_ROPE_TRITON", "1") != "0": + return RoPEInterleavedPartialFn.apply(x, cos, sin, rotary_dim) + + orig_dtype = x.dtype + nope = head_dim - rotary_dim + x_nope = x[..., :nope] + x_rope = x[..., nope:] + + # interleaved pairs: reshape last dim to (rotary_dim/2, 2) + x_pairs = x_rope.reshape(*x_rope.shape[:-1], rotary_dim // 2, 2) + even = x_pairs[..., 0] + odd = x_pairs[..., 1] + + # Always insert a singleton "heads" axis at position -2 so cos/sin + # broadcast across H. cos starts as ``position_ids.shape + [rd/2]``; + # after unsqueeze it becomes ``position_ids.shape + [1, rd/2]`` which + # broadcasts naturally against ``even`` of shape ``[..., H, rd/2]``. + # Cast cos/sin to x's dtype so the rotation does not promote bf16 + # inputs to fp32 -- otherwise the entire Q/K leaving RoPE doubles + # in HBM size and the downstream attention kernel runs 6-7x slower + # (see plan-5 P32 microbench-vs-proxy gap diagnostic). Math is done + # in bf16 because the freqs come from + # ``position_ids.float() * inv_freq`` which is already a + # single-precision rounding of the integer position; casting the + # final cos/sin to bf16 keeps numerics indistinguishable from the + # bf16-only reference flash-attn path used at every other site. + cos = cos.unsqueeze(-2).to(orig_dtype) + sin = sin.unsqueeze(-2).to(orig_dtype) + + rot_even = even * cos - odd * sin + rot_odd = even * sin + odd * cos + + rotated = torch.stack([rot_even, rot_odd], dim=-1).reshape(*x_rope.shape[:-1], rotary_dim) + return torch.cat([x_nope, rotated], dim=-1) + + +# --------------------------------------------------------------------------- +# DualRoPE — main + compress bases together +# --------------------------------------------------------------------------- + + +class DualRoPE(nn.Module): + """Holds the two RoPE caches V4 needs and routes per-layer applications. + + Args: + rotary_dim: partial-RoPE dim ``qk_pos_emb_head_dim`` (V4 = 64). + rope_theta: base for dense / SWA layers. + compress_rope_theta: base for CSA / HCA layers (longer base). + yarn_factor / yarn_beta_fast / yarn_beta_slow / + original_max_position_embeddings: YaRN config; applied **only** + to the compress base. Set ``yarn_factor=1.0`` to disable. + """ + + def __init__( + self, + *, + rotary_dim: int, + rope_theta: float, + compress_rope_theta: float, + yarn_factor: float = 1.0, + yarn_beta_fast: float = 32.0, + yarn_beta_slow: float = 1.0, + original_max_position_embeddings: int = 0, + ) -> None: + super().__init__() + self.rotary_dim = rotary_dim + + self.main_rope = RoPECache( + rotary_dim=rotary_dim, + theta=rope_theta, + ) + self.compress_rope = RoPECache( + rotary_dim=rotary_dim, + theta=compress_rope_theta, + yarn_factor=yarn_factor, + yarn_beta_fast=yarn_beta_fast, + yarn_beta_slow=yarn_beta_slow, + original_max_position_embeddings=original_max_position_embeddings, + ) + + def get_rope(self, *, compress_ratio: int) -> RoPECache: + """Pick the right cache for a layer. + + ``compress_ratio == 0`` → main (dense / SWA). Anything else → + compress (CSA / HCA). + """ + return self.main_rope if compress_ratio == 0 else self.compress_rope + + def apply_rope( + self, + x: torch.Tensor, + *, + position_ids: torch.Tensor, + compress_ratio: int, + ) -> torch.Tensor: + """Convenience: pick the right rope and apply partial RoPE. + + ``position_ids`` shape ``[B, S]`` or ``[S]``. cos/sin broadcast over + the heads dim of ``x``. + + Small-kernel-fusion (2026-07-03): route through + :func:`apply_rope_from_positions`, which generates cos/sin *inside* + the RoPE Triton kernel from ``(position_ids, inv_freq)`` — removing + the separate ``position_ids.float() * inv_freq`` / ``cos`` / ``sin`` + launches, the ``.to(x.dtype)`` cast, and the cos/sin HBM tensors that + the eager ``rope(position_ids)`` path materialised (and recomputed + identically for Q and K). YaRN is already baked into ``inv_freq``. + Gated by ``PRIMUS_ROPE_TRITON`` (falls back to the eager cos/sin path + on CPU / when off). + """ + rope = self.get_rope(compress_ratio=compress_ratio) + return apply_rope_from_positions(x, position_ids, rope.inv_freq, rotary_dim=self.rotary_dim) + + # Convenience accessors for callers who need the YaRN m_scale (e.g. to + # adjust attention softmax scale on compressed layers). + def attn_scale(self, *, compress_ratio: int) -> float: + return self.get_rope(compress_ratio=compress_ratio).attn_scale + + +__all__ = [ + "RoPECache", + "DualRoPE", + "apply_interleaved_partial_rope", +] diff --git a/primus/backends/megatron/core/transformer/experts.py b/primus/backends/megatron/core/transformer/experts.py index 847014b4d..694db2137 100644 --- a/primus/backends/megatron/core/transformer/experts.py +++ b/primus/backends/megatron/core/transformer/experts.py @@ -5,6 +5,12 @@ import torch import torch.nn.functional as F from megatron.core import tensor_parallel +from megatron.core.activations import squared_relu +from megatron.core.fusions.fused_bias_geglu import ( + quick_gelu, + weighted_bias_quick_geglu_impl, +) +from megatron.core.fusions.fused_weighted_squared_relu import weighted_squared_relu_impl from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( FineGrainedActivationOffloadingInterface as off_interface, ) @@ -41,6 +47,64 @@ def __init__( 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 + def bias_act_func(self, intermediate_parallel, bias_parallel, permuted_probs): + """ + Applies bias and activation function to the output of linear_fc1. + """ + if self.config.use_te_activation_func: + if bias_parallel is not None: + intermediate_parallel = intermediate_parallel + bias_parallel + intermediate_parallel = self.activation_func(intermediate_parallel) + if permuted_probs is not None: + original_dtype = intermediate_parallel.dtype + intermediate_parallel = intermediate_parallel * permuted_probs + intermediate_parallel = intermediate_parallel.to(original_dtype) + elif self.config.bias_activation_fusion: + if self.activation_func == F.silu and self.config.gated_linear_unit: + from primus.backends.megatron.core.fusions.fused_bias_swiglu import ( + weighted_bias_swiglu_impl, + ) + + # dtype is handled inside the fused kernel + intermediate_parallel = weighted_bias_swiglu_impl( + intermediate_parallel, + bias_parallel, + permuted_probs, + self.config.activation_func_fp8_input_store, + self.config.activation_func_clamp_value, + ) + elif self.activation_func == quick_gelu and self.config.gated_linear_unit: + intermediate_parallel = weighted_bias_quick_geglu_impl( + intermediate_parallel, + bias_parallel, + permuted_probs, + self.config.activation_func_fp8_input_store, + self.config.glu_linear_offset, + self.config.activation_func_clamp_value, + ) + else: + raise ValueError("Only support fusion of swiglu and quick_gelu in TEGroupedMLP.") + elif self.activation_func == squared_relu and self.config.use_fused_weighted_squared_relu: + assert bias_parallel is None, "Bias is not supported with fused weighted squared relu." + intermediate_parallel = weighted_squared_relu_impl(intermediate_parallel, permuted_probs) + else: + if self.config.gated_linear_unit: + + def glu(x): + x_glu, x_linear = torch.chunk(x, 2, dim=-1) + if (val := self.config.activation_func_clamp_value) is not None: + x_glu = x_glu.clamp(min=None, max=val) + x_linear = x_linear.clamp(min=-val, max=val) + return self.config.activation_func(x_glu) * (x_linear + self.config.glu_linear_offset) + + intermediate_parallel = glu(intermediate_parallel) + else: + intermediate_parallel = self.activation_func(intermediate_parallel) + original_dtype = intermediate_parallel.dtype + intermediate_parallel = intermediate_parallel * permuted_probs + intermediate_parallel = intermediate_parallel.to(original_dtype) + return intermediate_parallel + def bias_act_func_with_mask( self, intermediate_parallel: torch.Tensor, diff --git a/primus/backends/megatron/core/transformer/hyper_connection.py b/primus/backends/megatron/core/transformer/hyper_connection.py new file mode 100644 index 000000000..d3020e3f0 --- /dev/null +++ b/primus/backends/megatron/core/transformer/hyper_connection.py @@ -0,0 +1,468 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Manifold-Constrained Hyper-Connections (mHC) for DeepSeek-V4. + +Reference: deepseek-v4/develop/techblog/01-deepseek-v4-architecture-deep-dive.md +section 2 ("mHC: Manifold-Constrained Hyper-Connections") and section 9.3. + +Three modules live here: + +* :class:`HyperMixer` — per-layer mixer used twice per block (once before / + after the attention sub-block, once for the FFN sub-block). Produces + ``(pre, post, comb)`` triplets and exposes ``collapse`` / ``expand`` + helpers for the surrounding block to drive the K parallel hidden streams. + +* :class:`HyperHead` — final collapse used once at the end of the main trunk + (and once *per MTP layer*, with its own copy). Sigmoid-weighted sum, no + Sinkhorn. + +* :func:`sinkhorn_normalize` — alternating row/column normalization, kept in + fp32 for stability (per RedNote slide 9 and NeMo port pitfall #3). + +Phase 4 contract: +* Plain ``nn.Linear`` for the projection ``fn``. TP-friendly variants + (``ColumnParallelLinear``) come in Phase 6 once the rest of the V4 path + is correct end-to-end on a single device. +* All HC parameters (``fn.weight``, ``scale``, ``base``) live in fp32. + The block is responsible for passing in/out tensors in whatever activation + dtype it uses; the module up-casts internally to fp32 around the Sinkhorn + region and casts back at the end. + +Plan-5 P29 (RESCOPED — see ``deepseek-v4/develop/plan-5/02-phase-details.md``) +adds a ``torch.compile`` fast path for :func:`sinkhorn_normalize`. The eager +loop issues 1 + 2*(n_iters - 1) separate fp32 ``aten::sum`` reductions per +call; at V4-Flash production widths each reduction runs at ~250x over the +memory-bound floor because HIP's default ``reduce_kernel<512, 1, ...>`` is +sized for huge reductions and our ``[1, 4096, 4, 4] -> [1, 4096, 4, 1]`` +shape has only 4 elements per output. The compiled path collapses every +sum / divide / broadcast into one Inductor-fused Triton kernel +(``fullgraph=True, dynamic=False``). The compiled callable is cached +module-globally on ``(n_iters, eps, dtype)`` so each combination compiles +exactly once per process. +""" + +from __future__ import annotations + +import math +from typing import Callable, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_collapse import ( + hc_collapse_triton, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_expand import ( + hc_expand_triton, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_expand import ( + is_triton_kernel_supported as _hc_expand_supported, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_expand import ( + is_triton_path_enabled as _hc_expand_enabled, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue import ( + hc_glue_compute_tail_triton, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue import ( + is_triton_kernel_supported as _hc_glue_supported, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue import ( + is_triton_path_enabled as _hc_glue_enabled, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rmsnorm import ( + fused_rms_norm, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.sinkhorn import ( + SinkhornNormalizeFn, + is_triton_kernel_supported, + is_triton_path_enabled, +) + +# --------------------------------------------------------------------------- +# Plan-5 P29 fast path: torch.compile-fused Sinkhorn-Knopp. +# +# Cache keyed on ``(n_iters, eps, dtype)``. Each freshly-decorated callable +# uses ``@torch.compile(fullgraph=True, dynamic=True)`` so a SINGLE compiled +# artefact handles every input shape variation Inductor generates code that +# accepts shape as a runtime argument. The dynamic path is the right +# trade-off for V4 because: +# +# * In production each rank sees exactly one Sinkhorn input shape +# (``[B, S, K, K] = [1, 4096, 4, 4]`` at V4-Flash) so ``dynamic=True`` +# pays no specialization cost vs ``dynamic=False``. +# * Multi-shape harnesses (unit tests, MTP heads with different leading +# dims) all hit the same compiled artefact, so we never thrash Dynamo's +# recompile_limit. ``dynamic=False`` would force one Dynamo +# specialization per shape AND ALL CLOSURES SHARE THE SAME CODE OBJECT +# (closures from the same factory function inherit the factory's code +# object), which means Dynamo's cache_size_limit (default 8) is shared +# across every cached callable. After ~8 distinct ``(shape, stride, +# requires_grad)`` combinations the limit is hit and ``fullgraph=True`` +# raises ``FailOnRecompileLimitHit`` even though our cache is fine — +# the cache key just is not what Dynamo's internal cache is keyed on. +# +# ``in_dtype`` still participates in the cache key because the trailing +# cast back to the input activation dtype changes the Inductor IR; running +# bf16 and fp32 callers through the same compiled artefact would force a +# miscompile or a redundant recompile. +# --------------------------------------------------------------------------- + +_SinkhornCacheKey = tuple[int, float, torch.dtype] +_compiled_sinkhorn_cache: dict[_SinkhornCacheKey, Callable[[torch.Tensor], torch.Tensor]] = {} + + +def _build_compiled_sinkhorn( + n_iters: int, + eps: float, + in_dtype: torch.dtype, +) -> Callable[[torch.Tensor], torch.Tensor]: + """Build (and torch.compile) one Sinkhorn-Knopp implementation specialised + on ``(n_iters, eps, in_dtype)`` but generic over input shape. + + The body is the same algorithm as the eager :func:`sinkhorn_normalize` + path but written so Inductor sees one straight-line graph of fp32 + sums / divides / broadcasts (Python ``for`` is unrolled at compile time + because ``n_iters`` is a closure-captured Python int). + """ + + @torch.compile(fullgraph=True, dynamic=True) + def _impl(logits: torch.Tensor) -> torch.Tensor: + m = logits.float() + m = m / (m.sum(dim=-2, keepdim=True) + eps) + for _ in range(max(n_iters - 1, 0)): + m = m / (m.sum(dim=-1, keepdim=True) + eps) + m = m / (m.sum(dim=-2, keepdim=True) + eps) + return m.to(in_dtype) + + return _impl + + +def _get_compiled_sinkhorn( + n_iters: int, + eps: float, + in_dtype: torch.dtype, +) -> Callable[[torch.Tensor], torch.Tensor]: + """Cache-or-build accessor; returns the compiled callable for the given + Sinkhorn signature. Shape is NOT part of the key — see module-level + cache notes (we use ``dynamic=True`` so one artefact handles every + shape). + """ + key: _SinkhornCacheKey = (int(n_iters), float(eps), in_dtype) + fn = _compiled_sinkhorn_cache.get(key) + if fn is None: + fn = _build_compiled_sinkhorn(n_iters, eps, in_dtype) + _compiled_sinkhorn_cache[key] = fn + return fn + + +def sinkhorn_normalize( + logits: torch.Tensor, + *, + n_iters: int = 20, + eps: float = 1e-6, + use_compiled: bool = False, + use_triton: bool = False, +) -> torch.Tensor: + """Project a non-negative ``[..., K, K]`` matrix onto the doubly-stochastic + manifold via the Sinkhorn-Knopp algorithm. + + The algorithm itself is the alternating row / column ``L1`` normalization: + + .. code-block:: + + for _ in range(n_iters): + M /= M.sum(dim=-1, keepdim=True) # rows -> 1 + M /= M.sum(dim=-2, keepdim=True) # cols -> 1 + + To match the released V4 reference (cf. ``compute_weights`` in the + techblog) we emit the **first** column-normalization once before the + standard alternating loop, so that the final iteration ends on a column + normalization — this matches the (pre, post, comb) projection convention + that the surrounding block consumes. + + Stability: + * the input must already be non-negative (typically ``softmax`` output + plus an ``eps`` floor) + * runs in fp32 regardless of the input dtype, then casts back + + Args: + logits: ``[..., K, K]`` non-negative. + n_iters: total Sinkhorn iterations (= 1 priming column step + + ``n_iters - 1`` row/col cycles). + eps: numerical floor to prevent divide-by-zero. + use_compiled: when ``True``, dispatch to the cached + ``torch.compile(fullgraph=True, dynamic=False)`` build of the + same loop (plan-5 P29 RESCOPED). The compiled callable + collapses the 1 + 2*(n_iters - 1) fp32 ``aten::sum`` launches + (and their divides / broadcasts) into one Inductor-fused Triton + kernel; AOT autograd handles the BWD. Numerical contract is + unchanged (algorithm is identical; only the kernel boundary + moves). Default ``False`` until the V4 config flag + ``use_v4_compiled_sinkhorn`` is flipped on. + use_triton: plan-6 P36 hand-rolled Triton FWD/BWD path. When + ``True`` (or when the ``PRIMUS_SINKHORN_TRITON`` env knob is + not ``"0"`` and the input is supported), dispatch to + :class:`primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.sinkhorn.SinkhornNormalizeFn`. + The Triton path runs the full 1 + 2*(n_iters - 1) normalize + trajectory in registers per row of the leading axis and emits + exactly **one** FWD kernel + **one** BWD kernel -- no Dynamo + bookkeeping per call. Routing precedence at the call site: + ``use_triton or PRIMUS_SINKHORN_TRITON=1 > use_compiled > + eager``. Defaults to ``False``; production turns it on via + ``PRIMUS_SINKHORN_TRITON=1`` (default ON in the proxy + launcher and in :func:`is_triton_path_enabled`). + + Returns: + Approximately doubly-stochastic ``[..., K, K]`` (same dtype as input). + """ + # Plan-6 P36: routing precedence is Triton > compiled > eager. Env + # knob is default-ON; explicit `use_triton=True` forces the path + # even when the env is off (used by the G39 tests). The Triton path + # gracefully falls back when the input shape / device is unsupported + # (`is_triton_kernel_supported`), so callers never have to special- + # case K out-of-range or CPU inputs. + if (use_triton or is_triton_path_enabled()) and is_triton_kernel_supported(logits): + return SinkhornNormalizeFn.apply(logits, n_iters, eps) + if use_compiled: + return _get_compiled_sinkhorn(n_iters, eps, logits.dtype)(logits) + in_dtype = logits.dtype + m = logits.float() + m = m / (m.sum(dim=-2, keepdim=True) + eps) + for _ in range(max(n_iters - 1, 0)): + m = m / (m.sum(dim=-1, keepdim=True) + eps) + m = m / (m.sum(dim=-2, keepdim=True) + eps) + return m.to(in_dtype) + + +class HyperMixer(nn.Module): + """Per-layer mHC mixer. + + Maintains the three projection scales / biases (``pre``, ``post``, + ``comb``) and a single packed ``Linear`` ``fn`` that produces + ``[..., (2 + K) * K]`` from ``[..., K * D]``. + + Shapes (B and S are arbitrary; the mixer is agnostic to which is which): + + * ``compute_weights(x)`` → ``pre [..., K]``, ``post [..., K]``, + ``comb [..., K, K]`` + * ``collapse(x, pre)`` → ``[..., D]`` + * ``expand(x, out, post, comb)`` → ``[..., K, D]`` + + Args: + hidden_size: per-stream feature dim ``D``. + hc_mult: number of parallel streams ``K``. + eps: floor for ``sigmoid(...) + eps`` and for Sinkhorn. + sinkhorn_iters: Sinkhorn iteration count (``20`` matches V4 release). + use_compiled_sinkhorn: plan-5 P29 (RESCOPED) fast path. When + ``True`` the call to :func:`sinkhorn_normalize` inside + :meth:`compute_weights` dispatches to the + ``torch.compile(fullgraph=True, dynamic=False)`` build of the + same algorithm, kept in :data:`_compiled_sinkhorn_cache`. + Defaults to ``False`` so existing callers (and existing + checkpoints) are bit-equivalent. The V4 block reads + ``config.use_v4_compiled_sinkhorn`` and forwards it here. + """ + + def __init__( + self, + *, + hidden_size: int, + hc_mult: int, + eps: float = 1e-6, + sinkhorn_iters: int = 20, + use_compiled_sinkhorn: bool = False, + ) -> None: + super().__init__() + if hc_mult < 1: + raise ValueError(f"hc_mult must be >= 1, got {hc_mult}") + + self.hidden_size = hidden_size + self.hc_mult = hc_mult + self.eps = eps + self.sinkhorn_iters = sinkhorn_iters + self.use_compiled_sinkhorn = bool(use_compiled_sinkhorn) + + out_dim = (2 + hc_mult) * hc_mult + # All HC params kept in fp32; see techblog §2.2 pitfall #3. + self.fn = nn.Linear(hc_mult * hidden_size, out_dim, bias=False, dtype=torch.float32) + # Three independent scale scalars: one each for pre / post / comb. + self.scale = nn.Parameter(torch.ones(3, dtype=torch.float32)) + # Bias terms — same partition as the linear output. + self.base = nn.Parameter(torch.zeros(out_dim, dtype=torch.float32)) + + self.reset_parameters() + + def reset_parameters(self) -> None: + # NormalInitGain-style init keeps the post-rms-scale logits well-behaved. + nn.init.normal_(self.fn.weight, std=1.0 / math.sqrt(self.hc_mult * self.hidden_size)) + nn.init.zeros_(self.base) + nn.init.ones_(self.scale) + + # ---- internals ------------------------------------------------------- + + def _packed_logits(self, x: torch.Tensor) -> torch.Tensor: + """Pack streams along the last dim, RMS-normalize, project to ``out_dim``. + + ``x``: ``[..., K, D]`` → ``flat`` ``[..., K*D]`` → ``logits`` ``[..., out_dim]``. + Done in fp32 for stability (HC parameters are fp32 anyway). + + Small-kernel-fusion (2026-07-03): the parameter-less RMS over the + packed ``K*D`` axis (bf16->fp32 cast + pow + mean + rsqrt + mul) is + fused into one Triton FWD + one BWD kernel producing the fp32 + normalized tensor directly for ``F.linear``. Gated by + ``PRIMUS_RMSNORM_TRITON`` (default on). + """ + flat = x.flatten(-2) + normed = fused_rms_norm(flat, None, eps=self.eps, mid_cast=False, out_dtype=torch.float32) + logits = F.linear(normed, self.fn.weight.to(dtype=normed.dtype)) + return logits + + # ---- public API ------------------------------------------------------ + + def compute_weights(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return ``(pre, post, comb)`` for the K parallel streams ``x``. + + ``x``: ``[..., K, D]`` (typically ``[B, S, K, D]`` or ``[S, B, K, D]``). + """ + K = self.hc_mult + logits = self._packed_logits(x) # [..., (2+K)*K], fp32 + + out_dtype = x.dtype + # Plan-6 P37: fuse the elemwise tail (slice + scale + base + + # sigmoid/softmax + eps) into one Triton kernel when the env + # knob is on and the shape is supported. Falls back to the + # eager chain for the unsupported case (CPU input, K not in + # the supported set, etc.). + if _hc_glue_enabled() and _hc_glue_supported(logits, K): + pre, post, comb = hc_glue_compute_tail_triton( + logits, + self.scale, + self.base, + K=K, + eps=self.eps, + out_dtype=torch.float32, + ) + else: + pre_logit = logits[..., :K] * self.scale[0] + self.base[:K] + post_logit = logits[..., K : 2 * K] * self.scale[1] + self.base[K : 2 * K] + comb_logit = logits[..., 2 * K :].view(*logits.shape[:-1], K, K) * self.scale[2] + self.base[ + 2 * K : + ].view(K, K) + pre = torch.sigmoid(pre_logit) + self.eps # (eps, 1+eps] + post = 2.0 * torch.sigmoid(post_logit) # (0, 2) no eps + comb = torch.softmax(comb_logit, dim=-1) + self.eps + + comb = sinkhorn_normalize( + comb, + n_iters=self.sinkhorn_iters, + eps=self.eps, + use_compiled=self.use_compiled_sinkhorn, + ) + + # Cast back to the activation dtype to match downstream block compute. + return pre.to(out_dtype), post.to(out_dtype), comb.to(out_dtype) + + @staticmethod + def collapse(x: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: + """Collapse K streams into 1 via ``pre`` weights. + + ``x``: ``[..., K, D]``; ``pre``: ``[..., K]``; + returns ``[..., D]``. + + Small-kernel-fusion (2026-07-03): the eager + ``(pre.unsqueeze(-1) * x).sum(-2)`` (broadcast-mul into a full + ``[..., K, D]`` temporary + reduce = 2 kernels + ``K*D`` extra HBM + traffic) is fused into one Triton FWD + one BWD kernel. Symmetric to + the already-shipped ``expand`` fusion. Gated by + ``PRIMUS_HC_COLLAPSE_TRITON`` (default on); eager fallback on + CPU / unsupported K. + """ + return hc_collapse_triton(x, pre) + + @staticmethod + def expand( + x: torch.Tensor, + out: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ) -> torch.Tensor: + """Write the sub-block output ``out`` back to K streams. + + ``new_stream[h] = post[h] * out + Σ_k comb[h, k] * x[k]`` + + Shapes: + * ``x``: ``[..., K, D]`` — current K streams + * ``out``: ``[..., D]`` — sub-block (attn or FFN) output + * ``post``: ``[..., K]`` + * ``comb``: ``[..., K, K]`` + + Returns ``[..., K, D]``. + """ + # Triton-fused expand; falls back to eager for unsupported configs. + if _hc_expand_enabled() and _hc_expand_supported(x, post, comb): + return hc_expand_triton(x, out, post, comb) + # post[..., K] * out[..., D] -> [..., K, D] + write = post.unsqueeze(-1) * out.unsqueeze(-2) + # comb[..., K, K] @ x[..., K, D] -> [..., K, D] + mix = torch.matmul(comb, x) + return write + mix + + +class HyperHead(nn.Module): + """Final collapse from K streams → 1 stream. + + Used once at the end of the main trunk, and once *per* MTP layer (each + with its own copy — ``num_nextn_predict_layers`` separate heads). Plain + sigmoid-weighted sum; no Sinkhorn. + + Args: + hidden_size: per-stream feature dim ``D``. + hc_mult: number of input streams ``K``. + eps: floor for ``sigmoid(...) + eps``. + """ + + def __init__(self, *, hidden_size: int, hc_mult: int, eps: float = 1e-6) -> None: + super().__init__() + if hc_mult < 1: + raise ValueError(f"hc_mult must be >= 1, got {hc_mult}") + self.hidden_size = hidden_size + self.hc_mult = hc_mult + self.eps = eps + + self.fn = nn.Linear(hc_mult * hidden_size, hc_mult, bias=False, dtype=torch.float32) + self.scale = nn.Parameter(torch.ones((), dtype=torch.float32)) + self.base = nn.Parameter(torch.zeros(hc_mult, dtype=torch.float32)) + + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.normal_(self.fn.weight, std=1.0 / math.sqrt(self.hc_mult * self.hidden_size)) + nn.init.zeros_(self.base) + nn.init.ones_(self.scale) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """``x``: ``[..., K, D]`` → ``[..., D]``. + + Small-kernel-fusion (2026-07-03): the parameter-less RMS over the + packed ``K*D`` axis is fused into one Triton FWD + one BWD kernel + (fp32 output for ``F.linear``). Gated by ``PRIMUS_RMSNORM_TRITON``. + """ + flat = x.flatten(-2) + normed = fused_rms_norm(flat, None, eps=self.eps, mid_cast=False, out_dtype=torch.float32) + mixes = F.linear(normed, self.fn.weight.to(dtype=normed.dtype)) # [..., K] + pre = torch.sigmoid(mixes * self.scale + self.base) + self.eps + return (pre.unsqueeze(-1) * x).sum(dim=-2).to(x.dtype) + + +__all__ = [ + "HyperMixer", + "HyperHead", + "sinkhorn_normalize", +] diff --git a/primus/backends/megatron/core/transformer/indexer.py b/primus/backends/megatron/core/transformer/indexer.py new file mode 100644 index 000000000..bc06ca642 --- /dev/null +++ b/primus/backends/megatron/core/transformer/indexer.py @@ -0,0 +1,422 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +r""" +DeepSeek-V4 Indexer (sparse position selector for CSA). + +Reference: techblog §1.4 ("Indexer: CSA's Sparse Selector"). + +The Indexer is **only** used by CSA layers (``compress_ratio == 4``). For +each query position ``t`` it picks ``index_topk`` compressed-KV positions +``s`` (out of all compressed positions ``[0, P)`` where ``P = S // ratio``). + +Math (from the techblog): + +.. math:: + + q^Q_t = h_t W^{DQ},\quad q^I_{t,h} = q^Q_t W^{IUQ}_h,\quad + w^I_{t,h} = h_t W^w_h + +.. math:: + + I_{t,s} = \\sum_h w^I_{t,h}\\cdot \\mathrm{ReLU}(q^I_{t,h}\\cdot K^{IComp}_s) + +.. math:: + + \\mathrm{topk\\_idxs}_t = \\mathrm{argTopK}_s\\,I_{t,s} + +The Indexer carries its **own** mini-Compressor (``index_head_dim``, +``index_n_heads``); the ``K^{IComp}`` it produces is independent of the +main attention's compressed KV pool. It is only used to **select** top-k +positions; the actual values fetched into main attention come from the +main Compressor in the surrounding CSA layer. + +Phase 4 contract: +* Plain ``nn.Linear`` projections (TP integration in P6). +* Causal masking: positions ``s`` whose start raw-token index exceeds the + query's raw-token index get a value of ``-inf`` so they cannot be + selected. Out-of-range positions are returned as ``-1`` in the output + ``topk_idxs`` so the caller can treat them as "no key". +""" + +from __future__ import annotations + +import logging +import os +from typing import Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +logger = logging.getLogger(__name__) + + +def _is_rank0() -> bool: + try: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank() == 0 + except Exception: + pass + return True + + +from primus.backends.megatron.core.transformer.compressor import Compressor + +# E4M3 finite max magnitude (float8_e4m3fn): largest representable value. +_FP8_E4M3_MAX = 448.0 + + +def fake_quantize_fp8_e4m3(x: torch.Tensor) -> torch.Tensor: + """Per-tensor dynamic FP8 (E4M3) fake-quantization. + + Scales ``x`` so its max magnitude maps to the E4M3 finite range, rounds + through ``torch.float8_e4m3fn``, then dequantizes back to ``x.dtype``. This + simulates the precision of an FP8 QK GEMM input while keeping the matmul + itself in the activation dtype (QAT-style "simulated FP8" path). Returns + ``x`` unchanged when the platform/torch build lacks ``float8_e4m3fn`` or + when ``x`` is all-zero (degenerate scale). + """ + if not hasattr(torch, "float8_e4m3fn"): + return x + orig_dtype = x.dtype + amax = x.detach().abs().amax() + if not torch.isfinite(amax) or float(amax) <= 0.0: + return x + scale = (_FP8_E4M3_MAX / amax).to(x.dtype) + x_scaled = torch.clamp(x * scale, -_FP8_E4M3_MAX, _FP8_E4M3_MAX) + x_fp8 = x_scaled.to(torch.float8_e4m3fn) + return x_fp8.to(orig_dtype) / scale + + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import ( + indexer_score_triton, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import ( + is_triton_kernel_supported as _indexer_triton_full_supported, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import ( + is_triton_path_enabled as _indexer_triton_full_enabled, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post import ( + indexer_score_post_triton, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post import ( + is_triton_kernel_supported as _indexer_tail_triton_supported, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post import ( + is_triton_path_enabled as _indexer_tail_triton_enabled, +) + +# MXFP4 block size (E2M1 data + E8M0 per-32 block scales). +_MXFP4_BLOCK = 32 + + +def _indexer_fp4_enabled() -> bool: + """True iff PRIMUS_INDEXER_FP4 == "1" (default off): run the CSA-indexer QK in MXFP4.""" + return os.environ.get("PRIMUS_INDEXER_FP4", "0") == "1" + + +def _fp4_qk_gemm(q_i: torch.Tensor, k_icomp: torch.Tensor) -> torch.Tensor: + """Real MXFP4 indexer QK: per-batch [S*H,Hd] @ [P,Hd]^T (NT, trans_b) -> [B,S,H,P]. + + hipBLASLt FP4 needs K=Hd%128, M,N%16; force PRIMUS_TURBO_GEMM_BACKEND=FP4:HIPBLASLT. + """ + import primus_turbo.pytorch as pt + from primus_turbo.pytorch.core.low_precision import ( + Float4QuantConfig, + Format, + ScaleDtype, + ScalingGranularity, + ) + + cfg = Float4QuantConfig( + format=Format.E2M1_X2, + granularity=ScalingGranularity.MX_BLOCKWISE, + block_size=_MXFP4_BLOCK, + scale_dtype=ScaleDtype.E8M0, + ) + B, S, H, Hd = q_i.shape + P = k_icomp.shape[1] + outs = [] + for b in range(B): + a = q_i[b].reshape(S * H, Hd).contiguous() # [S*H, Hd] + bk = k_icomp[b].contiguous() # [P, Hd] + o = pt.ops.gemm_fp4(a, bk, trans_b=True, config=cfg) # [S*H, P] + outs.append(o.view(1, S, H, P)) + return torch.cat(outs, dim=0) + + +def _indexer_fp8_proj_enabled() -> bool: + """Run the indexer projections (w_dq/w_iuq/w_w) in MXFP8 (default off). + + Reuses the attention-proj flag PRIMUS_V4_FP8_ATTN_PROJ; only fires inside + turbo-fp8. The linears are duplicated (no TP shard), so fp8 is safe at any TP. + """ + if os.environ.get("PRIMUS_V4_FP8_ATTN_PROJ", "0") != "1": + return False + try: + from primus.backends.megatron.core.extensions.primus_turbo import ( + PrimusTurboLowPrecisionGlobalStateManager as _M, + ) + + return _M.is_turbo_fp8_enabled() + except Exception: + return False + + +def _fp8_linear(lin: nn.Linear, x: torch.Tensor) -> torch.Tensor: + """MXFP8 apply of an ``nn.Linear`` (weight [out,in], no bias): y = x @ Wᵀ.""" + import primus_turbo.pytorch as pt + + from primus.backends.megatron.core.extensions.primus_turbo import ( + PrimusTurboLowPrecisionGlobalStateManager as _M, + ) + + cfg = _M.get_turbo_quant_config().data() + orig = x.shape + x2 = x.reshape(-1, orig[-1]).contiguous() + out = pt.ops.gemm_fp8(x2, lin.weight, trans_b=True, config=cfg) # [*, out] + return out.reshape(*orig[:-1], out.shape[-1]) + + +class Indexer(nn.Module): + """Sparse position selector for CSA. + + Args: + hidden_size: input feature dim ``D`` (same as main attention). + index_head_dim: head dim used by the mini-Compressor and the + low-rank query projection. + index_n_heads: number of indexer "heads". + index_topk: number of compressed positions to select per query. + compress_ratio: ratio ``m`` of the mini-Compressor (matches the + main Compressor of the surrounding CSA layer; usually ``4``). + dq_rank: rank of the shared low-rank query projection ``W^{DQ}``. + Defaults to ``index_head_dim`` (the V4 reference doesn't expose + a separate setting; ``W^{IUQ}_h`` then projects from ``dq_rank`` + to ``index_head_dim``). + """ + + def __init__( + self, + *, + hidden_size: int, + index_head_dim: int, + index_n_heads: int, + index_topk: int, + compress_ratio: int = 4, + dq_rank: int = None, + use_fp8_qk: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.index_head_dim = index_head_dim + self.index_n_heads = index_n_heads + self.index_topk = index_topk + self.compress_ratio = compress_ratio + self.dq_rank = dq_rank if dq_rank is not None else index_head_dim + # FP8 (E4M3) fake-quant of the QK scoring inputs (V4 low-precision + # indexer QK path). See ``fake_quantize_fp8_e4m3`` / config flag + # ``use_v4_fp8_indexer``. + self.use_fp8_qk = bool(use_fp8_qk) + if self.use_fp8_qk and _is_rank0(): + logger.info( + "[V4-Indexer] FP8 (E4M3) QK scoring path ENABLED " + "(query/compressed-key activations fake-quantized; " + "BF16 index-score + top-k preserved)." + ) + + # W^{DQ} (hidden->dq_rank) and W^w (hidden->n_heads) both consume `hidden`, + # so fuse them into ONE GEMM (default-on); split the output. W^{IUQ} stays + # separate (it consumes q_q, sequentially). PRIMUS_INDEXER_FUSE_PROJ=0 keeps + # the two separate linears. + self._fuse_qw_proj = os.environ.get("PRIMUS_INDEXER_FUSE_PROJ", "1") != "0" + if self._fuse_qw_proj: + self.w_dq_w = nn.Linear(hidden_size, self.dq_rank + index_n_heads, bias=False) + else: + # W^{DQ}: low-rank query down-projection. + self.w_dq = nn.Linear(hidden_size, self.dq_rank, bias=False) + # W^w_h: per-head scalar weight. + self.w_w = nn.Linear(hidden_size, index_n_heads, bias=False) + # W^{IUQ}_h: per-head up-projection from dq_rank → index_head_dim. + self.w_iuq = nn.Linear(self.dq_rank, index_n_heads * index_head_dim, bias=False) + + # Mini-Compressor producing K^{IComp}. + self.indexer_compressor = Compressor( + hidden_size=hidden_size, + head_dim=index_head_dim, + ratio=compress_ratio, + ) + + def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs): + """Bridge checkpoints across the fused/unfused (w_dq, w_w) projection. + + Old checkpoints store ``w_dq.weight`` + ``w_w.weight``; the fused path wants + ``w_dq_w.weight`` = ``cat([w_dq, w_w])`` (and vice-versa). Remap in-place so + either layout loads under either runtime setting. + """ + dq_k, w_k, fused_k = prefix + "w_dq.weight", prefix + "w_w.weight", prefix + "w_dq_w.weight" + if self._fuse_qw_proj and dq_k in state_dict and fused_k not in state_dict: + state_dict[fused_k] = torch.cat([state_dict.pop(dq_k), state_dict.pop(w_k)], dim=0) + elif (not self._fuse_qw_proj) and fused_k in state_dict and dq_k not in state_dict: + w = state_dict.pop(fused_k) + state_dict[dq_k], state_dict[w_k] = w[: self.dq_rank], w[self.dq_rank :] + return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs) + + # ------------------------------------------------------------------ + + def _causal_mask( + self, + n_queries: int, + n_pool: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Return ``[n_queries, n_pool]`` mask: 0.0 if pool position ``s`` + is allowed for query ``t``, ``-inf`` otherwise. + + A compressed position ``s`` covers raw tokens ``[s*ratio, (s+1)*ratio)``; + a query at raw token ``t`` may attend to ``s`` iff its window + end ``(s+1)*ratio - 1 <= t``. + """ + # The mask depends only on (n_queries, n_pool, compress_ratio, dtype) — all + # fixed per run — so cache it instead of rebuilding arange + where every + # call. PRIMUS_INDEXER_MASK_CACHE=0 forces the eager rebuild. + use_cache = os.environ.get("PRIMUS_INDEXER_MASK_CACHE", "1") != "0" + if use_cache: + cache = getattr(self, "_causal_mask_cache", None) + if cache is None: + cache = self._causal_mask_cache = {} + key = (n_queries, n_pool, device, dtype) + cached = cache.get(key) + if cached is not None: + return cached + t_idx = torch.arange(n_queries, device=device).unsqueeze(1) # [t, 1] + s_end = (torch.arange(n_pool, device=device).unsqueeze(0) + 1) * self.compress_ratio - 1 # [1, s] + allowed = s_end <= t_idx # [t, s] bool + mask = torch.where(allowed, 0.0, float("-inf")).to(dtype) + if use_cache: + cache[key] = mask + return mask + + # ------------------------------------------------------------------ + + def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Select top-k compressed positions for each query. + + Args: + hidden: ``[B, S, D]``. + + Returns: + ``(topk_idxs, topk_scores)`` where: + * ``topk_idxs`` ``[B, S, K]`` (long): selected pool positions + in ``[0, P)`` for valid slots, ``-1`` for masked / invalid. + * ``topk_scores`` ``[B, S, K]``: the selection scores ``I_{t,s}`` + (``-inf`` for masked positions). + """ + B, S, D = hidden.shape + assert S % self.compress_ratio == 0, ( + f"Indexer: sequence length {S} not divisible by compress_ratio " f"{self.compress_ratio}" + ) + K = self.index_topk + H = self.index_n_heads + Hd = self.index_head_dim + + # 1) K^{IComp}: pool hidden via the mini-Compressor → [B, P, Hd] + k_icomp = self.indexer_compressor(hidden) # [B, P, Hd] + P = k_icomp.shape[1] + k_icomp = k_icomp.unsqueeze(2) # [B, P, 1, Hd] + + # 2) Per-head query and per-head weight. + # Indexer projections: FP8 (paper / NVIDIA backend.linear) when enabled, + # else the bf16 nn.Linear. No TP gather/scatter (duplicated linears). + # FP8 (paper / NVIDIA backend.linear) when enabled, else the bf16 nn.Linear. + proj = _fp8_linear if _indexer_fp8_proj_enabled() else (lambda lin, x: lin(x)) + if self._fuse_qw_proj: + dqw = proj(self.w_dq_w, hidden) # [B, S, dq_rank + H] in one GEMM + q_q = dqw[..., : self.dq_rank] # [B, S, dq_rank] + w_i = dqw[..., self.dq_rank :] # [B, S, H] + else: + q_q = proj(self.w_dq, hidden) # [B, S, dq_rank] + w_i = proj(self.w_w, hidden) # [B, S, H] + q_i = proj(self.w_iuq, q_q).view(B, S, H, Hd) # [B, S, H, Hd] + + # 3) Score I_{t,s} = Σ_h w_i[t,h] * ReLU(q_i[t,h] · k_icomp[s]) + # q_i [B,S,H,Hd] · k_icomp[B,P,Hd] → relu[B,S,H,P]; w_i[B,S,H,1] → sum over H + # 4) Causal mask + (effective topk capped at P). + # + # Dispatch precedence (P41 re-routing): + # PRIMUS_INDEXER_TRITON=1 → post-einsum tail fused + # (einsum stays eager / cuBLAS). + # PRIMUS_INDEXER_TRITON_FULL=1 → legacy P38 full-fuse path + # (einsum + tail in one kernel). + # else → fully eager. + k_icomp_2d = k_icomp.squeeze(2) + + # Indexer QK precision (both default OFF -> BF16 QK). FP8 (E4M3) fake- + # quantizes the operands before the normal score dispatch; FP4 (below) + # is a dedicated real-GEMM branch and takes precedence when both are set. + # The ReLU + per-head weight (``w_i``) + sum + causal mask + top-k stay + # in the activation dtype — only the QK operands are quantized. + if self.use_fp8_qk and not _indexer_fp4_enabled(): + q_i = fake_quantize_fp8_e4m3(q_i) + k_icomp_2d = fake_quantize_fp8_e4m3(k_icomp_2d) + + # Phase 5: FP4 CSA-indexer QK. Real MXFP4 GEMM for the QK product (paper + # §2.3.4/§5.2.1: "QK multiplied entirely in FP4"), then the eager + # ReLU/weight/sum tail (w_i + tail stay BF16/FP32 — only the QK is FP4). + if _indexer_fp4_enabled(): + dot = _fp4_qk_gemm(q_i, k_icomp_2d) # [B, S, H, P], real FP4 matmul + relu = F.relu(dot) + scores = (relu * w_i.unsqueeze(-1)).sum(dim=2) # [B, S, P] + mask = self._causal_mask(S, P, scores.device, scores.dtype) # [S, P] + scores = scores + mask.unsqueeze(0) # [B, S, P] + elif _indexer_triton_full_enabled() and _indexer_triton_full_supported(q_i, k_icomp_2d, w_i): + scores = indexer_score_triton( + q_i, + k_icomp_2d, + w_i, + compress_ratio=self.compress_ratio, + out_dtype=hidden.dtype, + ) + else: + dot = torch.einsum("bshd,bpd->bshp", q_i, k_icomp_2d) + if _indexer_tail_triton_enabled() and _indexer_tail_triton_supported(dot, w_i): + scores = indexer_score_post_triton( + dot, + w_i, + compress_ratio=self.compress_ratio, + out_dtype=hidden.dtype, + ) + else: + relu = F.relu(dot) + scores = (relu * w_i.unsqueeze(-1)).sum(dim=2) # [B, S, P] + mask = self._causal_mask(S, P, scores.device, scores.dtype) # [S, P] + scores = scores + mask.unsqueeze(0) # [B, S, P] + + topk_eff = min(K, P) + topk_scores, topk_idxs = scores.topk(topk_eff, dim=-1) # [B, S, topk_eff] + + # 5) Replace selections that are still -inf (i.e. fewer than K valid + # pool positions for very early queries) with sentinel ``-1`` so + # callers can drop them. + invalid = torch.isneginf(topk_scores) + topk_idxs = torch.where(invalid, torch.full_like(topk_idxs, -1), topk_idxs) + + # 6) Pad with -1 to exactly K columns if topk_eff < K (S smaller than + # K * ratio in unit tests). + if topk_eff < K: + pad_idxs = torch.full((B, S, K - topk_eff), -1, dtype=topk_idxs.dtype, device=topk_idxs.device) + pad_scores = torch.full( + (B, S, K - topk_eff), float("-inf"), dtype=topk_scores.dtype, device=topk_scores.device + ) + topk_idxs = torch.cat([topk_idxs, pad_idxs], dim=-1) + topk_scores = torch.cat([topk_scores, pad_scores], dim=-1) + + return topk_idxs, topk_scores + + +__all__ = ["Indexer"] diff --git a/primus/backends/megatron/core/transformer/local_rmsnorm.py b/primus/backends/megatron/core/transformer/local_rmsnorm.py new file mode 100644 index 000000000..6efa92dad --- /dev/null +++ b/primus/backends/megatron/core/transformer/local_rmsnorm.py @@ -0,0 +1,103 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Tiny RMSNorm fallback used by DeepSeek-V4 modules. + +Plan-2 P17 introduces this shared helper to retire three nearly +identical ``_RMSNorm`` implementations that lived in: + +* ``primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py`` +* ``primus/backends/megatron/core/transformer/deepseek_v4_attention.py`` + (a closure-built fallback for the attention's ``q_norm`` / ``kv_norm`` + no-spec path) +* ``primus/backends/megatron/core/transformer/compressor.py`` + +All three computed the same RMSNorm with a learnable per-channel +``weight`` and worked on the last dim. The goal of this module is to +expose **one** implementation so: + +* dead-code audits stay clean, +* state-dict round-trips treat the four call sites uniformly, +* CPU-only unit tests can build V4 norms without dragging in + TransformerEngine / Megatron's TE-backed norm kernel. + +The spec-driven path (``DeepSeekV4SpecProvider.v4_norm_module()``) is +unchanged — it returns Megatron's TE-backed RMSNorm or the local +fallback, depending on the active runtime mode. This file is the local +fallback only. +""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch +import torch.nn as nn + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rmsnorm import ( + fused_rms_norm, +) + + +class LocalRMSNorm(nn.Module): + """Single canonical RMSNorm fallback used across V4 modules. + + Args: + dim: hidden dimension to normalize over (last axis). + eps: numerical stability epsilon. + hidden_size: alias for ``dim``; if both are provided ``dim`` + wins. Provided for compatibility with Megatron-style norm + constructors that expect ``hidden_size=`` instead. + config: ignored (consumed for compatibility with Megatron's norm + factory signature). + + Notes: + * The internal compute happens in fp32 to match the V4 reference + (``inference/model.py`` uses fp32 RMS even when activations + are bf16); the result is cast back to the input dtype. + * The ``weight`` parameter has shape ``(dim,)`` and is initialized + to ones so a freshly built norm is the identity transform — + this matches the V4 / HF reference checkpoint layout exactly, + so state-dict round-trips work without remapping. + """ + + def __init__( + self, + dim: Optional[int] = None, + eps: float = 1e-6, + *, + hidden_size: Optional[int] = None, + config: Optional[Any] = None, + ) -> None: + del config + super().__init__() + if dim is None: + dim = hidden_size + if dim is None: + raise ValueError("LocalRMSNorm requires `dim` (or the alias `hidden_size`).") + self.weight = nn.Parameter(torch.ones(int(dim))) + self.eps = float(eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Small-kernel-fusion (2026-07-03): collapse the eager RMS chain + # (bf16->fp32 cast + pow + mean + rsqrt + mul + fp32->bf16 cast + + # weight mul) into one Triton FWD + one BWD kernel. ``mid_cast=True`` + # replicates the eager ``(x32*rstd).to(in_dtype)`` rounding BEFORE the + # weight multiply so numerics match bit-for-bit (within fp32 accum). + # ``out_dtype`` = the eager output dtype ``promote(in_dtype, weight)``. + # Gated by ``PRIMUS_RMSNORM_TRITON`` (default on); falls back to the + # eager body on CPU / when off. + out_dtype = torch.promote_types(x.dtype, self.weight.dtype) + return fused_rms_norm( + x, + self.weight, + eps=self.eps, + mid_cast=True, + out_dtype=out_dtype, + ) + + +__all__ = ["LocalRMSNorm"] diff --git a/primus/backends/megatron/core/transformer/moe/_triton/__init__.py b/primus/backends/megatron/core/transformer/moe/_triton/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/primus/backends/megatron/core/transformer/moe/_triton/v4_router_post.py b/primus/backends/megatron/core/transformer/moe/_triton/v4_router_post.py new file mode 100644 index 000000000..82a40c2fa --- /dev/null +++ b/primus/backends/megatron/core/transformer/moe/_triton/v4_router_post.py @@ -0,0 +1,548 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton-fused V4 router post-logits chain (plan-6 P39). + +Fuses the chain shared between +:class:`primus.backends.megatron.core.transformer.moe.v4_topk_router.DeepseekV4LearnedRouter` +and +:class:`primus.backends.megatron.core.transformer.moe.v4_hash_router.DeepseekV4HashRouter`: + +.. code-block:: python + + scores = score_fn(logits) # softmax / sigmoid / sqrtsoftplus + weights = scores.gather(1, indices) # [N, K] + if score_fn != softmax: + denom = weights.sum(dim=-1, keepdim=True).clamp(min=1e-12) + weights = weights / denom + weights *= topk_scaling_factor + + probs = zeros(N, E); probs.scatter_(1, indices, weights) + routing_map = zeros(N, E); routing_map.scatter_(1, indices, True) + +`indices` is provided by the host (hash router pre-computes via +``tid2eid[flat_ids]``; learned router pre-computes via +``torch.topk(sel_score, K)``). Keeping `topk` on the host lets us +reuse the well-tuned cuBLAS / hip topk kernel and avoids stuffing +two unrelated kernels into one register file. + +Gating: ``PRIMUS_V4_ROUTER_TRITON == "1"`` (**default-off**, P38 precedent). + +Microbench at V4-Flash widths (N=4096, E=256, K=8) shows the kernel +wins on the V4 production score function: +- ``sqrtsoftplus`` (V4 default): 1.56x FWD / 1.22x BWD +- ``softmax`` (non-V4): 1.00x FWD / 0.73x BWD +- ``sigmoid``: near-parity + +But the EP=8 proxy A/B (PRIMUS_V4_ROUTER_TRITON=1 vs =0, 10 iters +each) shows the microbench gain does **not** surface end-to-end: +~534 ms / iter both ways, lm_loss bit-identical iter-by-iter +(parity confirmed). The per-call savings (~0.06 ms FWD+BWD × +~16 router calls / iter ≈ 1 ms) are submerged in the ~500 ms / +step variance of the EP=8 dispatch + grouped-MLP pipeline. + +P38-style descope: ship the kernel behind a knob, default OFF. +Available for future tuning (e.g. when the broader graph compresses +and exposes the per-call savings) and for small-shape paths where +the FWD win is closer to 1.5-2x. Bit-identity makes flipping the +knob a safe operation. +""" + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +# Score function enum (matches FWD/BWD constexpr). +_SCORE_FN_SOFTMAX = 0 +_SCORE_FN_SIGMOID = 1 +_SCORE_FN_SQRTSOFTPLUS = 2 +_SCORE_FN_MAP = { + "softmax": _SCORE_FN_SOFTMAX, + "sigmoid": _SCORE_FN_SIGMOID, + "sqrtsoftplus": _SCORE_FN_SQRTSOFTPLUS, +} + + +# --------------------------------------------------------------------------- +# Triton kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_router_post_fwd_kernel( + LOGITS_PTR, # [N, E] fp32 + INDICES_PTR, # [N, BLOCK_K] int64 (gather positions; padded to BLOCK_K) + PROBS_PTR, # [N, E] OUT_DTYPE (must be zero-init'd by caller) + RMAP_PTR, # [N, E] bool (must be zero-init'd by caller) + SCORES_OUT_PTR, # [N, E] fp32 (saved-for-backward; full row of scores) + WEIGHTS_OUT_PTR, # [N, BLOCK_K] fp32 (saved-for-backward; gathered weights, post-denom, pre-scale) + DENOM_OUT_PTR, # [N] fp32 (saved-for-backward; the clamped denom) + N, + E, # runtime int (real expert count; NOT required to be a power of 2) + K, # runtime int (real topk; NOT required to be a power of 2) + SCORE_FN: tl.constexpr, + SCALE: tl.constexpr, + EPS: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_E: tl.constexpr, # next_pow2(E) — column block over the expert axis + BLOCK_K: tl.constexpr, # next_pow2(K) — column block over the topk axis + OUT_DTYPE: tl.constexpr, +): + """One program tile = ``BLOCK_N`` rows of the post-logits chain. + + Supports arbitrary (non-power-of-2) ``E`` / ``K``: the expert axis is + tiled by ``BLOCK_E = next_pow2(E)`` and the topk axis by + ``BLOCK_K = next_pow2(K)``, with masks zeroing the padded columns. + ``INDICES_PTR`` / ``WEIGHTS_OUT_PTR`` are laid out with row stride + ``BLOCK_K`` (host pads the indices tensor to ``[N, BLOCK_K]``). + """ + + pid = tl.program_id(0) + n_offs = pid * BLOCK_N + tl.arange(0, BLOCK_N) + n_mask = n_offs < N + + e_idx = tl.arange(0, BLOCK_E) + e_mask = e_idx < E + k_idx = tl.arange(0, BLOCK_K) + k_mask = k_idx < K + + # Load logits [BLOCK_N, BLOCK_E] in fp32 (padded cols masked out). + logits = tl.load( + LOGITS_PTR + n_offs[:, None] * E + e_idx[None, :], + mask=n_mask[:, None] & e_mask[None, :], + other=0.0, + ).to(tl.float32) + + # Apply score_fn (compile-time specialised). Padded expert columns must + # not leak into the softmax reduction, so mask them to -inf first. + if SCORE_FN == 0: # softmax + neg_inf = float("-inf") + logits_m = tl.where(e_mask[None, :], logits, neg_inf) + m = tl.max(logits_m, axis=1, keep_dims=True) + e = tl.where(e_mask[None, :], tl.exp(logits_m - m), 0.0) + s = tl.sum(e, axis=1, keep_dims=True) + scores = e / s + elif SCORE_FN == 1: # sigmoid + scores = tl.sigmoid(logits) + else: # sqrtsoftplus (SCORE_FN == 2) + softplus = tl.log(1.0 + tl.exp(logits)) + # Stable: for large x, log(1+exp(x)) ≈ x. + scores = tl.sqrt(softplus) + + # Zero padded expert columns so they never contribute to gather / denom. + scores = tl.where(e_mask[None, :], scores, 0.0) + + # Save full scores row for backward (real columns only). + tl.store( + SCORES_OUT_PTR + n_offs[:, None] * E + e_idx[None, :], + scores, + mask=n_mask[:, None] & e_mask[None, :], + ) + + # Load indices [BLOCK_N, BLOCK_K] (int64; padded cols read as 0, masked below). + indices = tl.load( + INDICES_PTR + n_offs[:, None] * BLOCK_K + k_idx[None, :], + mask=n_mask[:, None] & k_mask[None, :], + other=0, + ) + + # Gather weights from scores at the K indices per row entirely in + # registers (avoid store-then-reload-via-HBM hazard). Build + # weights[BLOCK_N, BLOCK_K] via a static loop over BLOCK_K, extracting + # each column from the already-loaded ``indices`` tile. + weights = tl.zeros((BLOCK_N, BLOCK_K), dtype=tl.float32) + for k_off in tl.static_range(BLOCK_K): + k_is_off = (k_idx == k_off).to(tl.float32) + idx_col = tl.sum(tl.where(k_idx[None, :] == k_off, indices, 0), axis=1) + # scores_at_idx[n] = scores[n, idx_col[n]] + e_is_idx = e_idx[None, :] == idx_col[:, None] + scores_at_idx = tl.sum(scores * e_is_idx.to(tl.float32), axis=1) + weights += scores_at_idx[:, None] * k_is_off[None, :] + + # Zero padded topk columns (k >= K). + weights = tl.where(k_mask[None, :], weights, 0.0) + + # If non-softmax: denom + normalize. + if SCORE_FN != 0: + s = tl.sum(weights, axis=1, keep_dims=True) + denom = tl.maximum(s, EPS) + weights = weights / denom + denom_scalar = tl.reshape(denom, (BLOCK_N,)) + else: + denom_scalar = tl.full((BLOCK_N,), 1.0, dtype=tl.float32) + tl.store(DENOM_OUT_PTR + n_offs, denom_scalar, mask=n_mask) + + # Scale. + weights = weights * SCALE + + # Save weights for backward. + tl.store( + WEIGHTS_OUT_PTR + n_offs[:, None] * BLOCK_K + k_idx[None, :], + weights, + mask=n_mask[:, None] & k_mask[None, :], + ) + + # Sparse scatter (probs[n, indices] = weights, rmap[n, indices] = True). + # No atomics needed because each (n, e) is written at most once + # per row (caller guarantees indices are unique within a row). Padded + # topk columns are masked out so they never write to expert 0. + tl.store( + PROBS_PTR + n_offs[:, None] * E + indices, + weights.to(OUT_DTYPE), + mask=n_mask[:, None] & k_mask[None, :], + ) + tl.store( + RMAP_PTR + n_offs[:, None] * E + indices, + tl.full((BLOCK_N, BLOCK_K), 1, dtype=tl.int1), + mask=n_mask[:, None] & k_mask[None, :], + ) + + +@triton.jit +def _v4_router_post_bwd_kernel( + DPROBS_PTR, # [N, E] OUT_DTYPE upstream grad + INDICES_PTR, # [N, BLOCK_K] int64 (padded) + SCORES_PTR, # [N, E] fp32 saved + WEIGHTS_PTR, # [N, BLOCK_K] fp32 saved (post-scaled, padded) + DENOM_PTR, # [N] fp32 saved + DLOGITS_PTR, # [N, E] fp32 OUT + N, + E, # runtime int + K, # runtime int + SCORE_FN: tl.constexpr, + SCALE: tl.constexpr, + EPS: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_E: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """VJP through the chain. + + FWD: + scores[n, e] = score_fn(logits[n, e]) + weights[n, k] = scores[n, indices[n, k]] + if not softmax: weights /= sum_k weights[n, k] + weights *= scale + probs[n, indices[n, k]] = weights + + BWD (gather only the K positions touched by indices; others get 0): + dprobs_at[n, k] = dprobs[n, indices[n, k]] + dweights[n, k] = dprobs_at[n, k] * SCALE (chain through scale) + If not softmax: + saved weights = (gathered / denom) * SCALE + d_pre_denom[n, k] = dweights[n, k] / denom + d_denom[n] = -sum_k (dweights[n, k] * weights_scaled[n, k]) / (denom * scale) + (because weights_scaled = gathered * scale / denom, so + ∂w/∂gathered = scale/denom, ∂w/∂denom = -gathered * scale / denom^2) + d_gathered[n, k] = dweights[n, k] * scale / denom + d_denom[n] * 1 + Wait, denom = sum_k gathered, so ∂denom/∂gathered_k = 1. + Therefore d_gathered[n, k] = dweights * scale/denom + d_denom (each gathered contributes to denom) + else (softmax): + d_gathered[n, k] = dweights[n, k] + Scatter d_gathered back into d_scores at positions indices[n, k] + (rest of d_scores is 0). + Finally chain through score_fn: + d_logits[n, e] = vjp_score_fn(d_scores[n, e], scores[n, e]) + """ + + pid = tl.program_id(0) + n_offs = pid * BLOCK_N + tl.arange(0, BLOCK_N) + n_mask = n_offs < N + + e_idx = tl.arange(0, BLOCK_E) + e_mask = e_idx < E + k_idx = tl.arange(0, BLOCK_K) + k_mask = k_idx < K + + # Load saved state. + scores = tl.load( + SCORES_PTR + n_offs[:, None] * E + e_idx[None, :], + mask=n_mask[:, None] & e_mask[None, :], + other=0.0, + ) + indices = tl.load( + INDICES_PTR + n_offs[:, None] * BLOCK_K + k_idx[None, :], + mask=n_mask[:, None] & k_mask[None, :], + other=0, + ) + + # Gather upstream grad at the K indices per row. + # dprobs_at = dprobs[n, indices[n, k]] shape [BLOCK_N, BLOCK_K], fp32. + # Padded topk columns are masked to 0 so they carry no gradient. + dprobs_at = tl.load( + DPROBS_PTR + n_offs[:, None] * E + indices, + mask=n_mask[:, None] & k_mask[None, :], + other=0.0, + ).to(tl.float32) + # Chain through the SCALE multiply: forward was + # weights_post = weights_pre * SCALE + # so d_weights_pre = d_weights_post * SCALE. + dweights_pre_scale = dprobs_at * SCALE + + if SCORE_FN != 0: + # Non-softmax: forward applied + # weights_pre = gathered / denom + # so d_gathered_k = d_weights_pre_k / denom + d_denom + # where d_denom = -sum_k d_weights_pre_k * gathered_k / denom^2. + # Recover gathered_k from saved state: + # weights_saved_k = (gathered_k / denom) * SCALE + # so gathered_k = weights_saved_k * denom / SCALE. + denom = tl.load(DENOM_PTR + n_offs, mask=n_mask, other=1.0) + weights_saved = tl.load( + WEIGHTS_PTR + n_offs[:, None] * BLOCK_K + k_idx[None, :], + mask=n_mask[:, None] & k_mask[None, :], + other=0.0, + ) + gathered_k = weights_saved * (denom[:, None] / SCALE) + d_denom_per_n = -tl.sum(dweights_pre_scale * gathered_k, axis=1) / (denom * denom) + d_gathered = dweights_pre_scale / denom[:, None] + d_denom_per_n[:, None] + else: + # Softmax: weights_pre is just gathered directly. + d_gathered = dweights_pre_scale + + # Zero padded topk columns so they contribute no gradient. + d_gathered = tl.where(k_mask[None, :], d_gathered, 0.0) + + # Build d_scores_full entirely in registers using an explicit static + # loop over BLOCK_K. For each k, we add d_gathered[:, k] at the position + # indices[:, k] in the E-axis using a broadcast compare. This avoids + # the round-trip-via-HBM hazard of a scatter-then-load pattern. Padded + # columns carry d_gathered == 0, so they add nothing even though their + # index reads as 0. + d_scores_full = tl.zeros((BLOCK_N, BLOCK_E), dtype=tl.float32) + for k_off in tl.static_range(BLOCK_K): + idx_col = tl.sum(tl.where(k_idx[None, :] == k_off, indices, 0), axis=1) + # Slice d_gathered[:, k_off] -> shape [BLOCK_N] + grad_k = tl.sum(d_gathered * (k_idx[None, :] == k_off).to(tl.float32), axis=1) + # Contribution: grad_k[:, None] where e_idx == idx_col[:, None] + e_is_idx = e_idx[None, :] == idx_col[:, None] + d_scores_full += tl.where(e_is_idx, grad_k[:, None], 0.0) + + if SCORE_FN == 0: # softmax: d_logits = scores * (d_scores - sum(d_scores * scores)) + dot = tl.sum(d_scores_full * scores, axis=1, keep_dims=True) + d_logits = scores * (d_scores_full - dot) + elif SCORE_FN == 1: # sigmoid + d_logits = d_scores_full * scores * (1.0 - scores) + else: # sqrtsoftplus: y = sqrt(softplus(x)); dy/dx = (1 / (2 * sqrt(softplus(x)))) * sigmoid(x) + # = sigmoid(x) / (2 * y); avoid div-by-zero by guarding small y. + # sigmoid(x) = exp(x) / (1 + exp(x)); when y is close to 0, x is very negative, + # and sigmoid(x) is also very small, so the limit is 0. + # Use: dy/dx = sigmoid(x) / (2 * scores) + # but we don't have x. Recompute: scores = y = sqrt(softplus(x)) + # so softplus(x) = y^2 = scores^2, and exp(x) = exp(scores^2) - 1? No: + # softplus(x) = log(1 + exp(x)) = y^2 -> exp(x) = exp(y^2) - 1 -> sigmoid(x) = (exp(y^2) - 1) / exp(y^2) + # = 1 - exp(-y^2). + sig_x = 1.0 - tl.exp(-scores * scores) + y_safe = tl.maximum(scores, EPS) + d_logits = d_scores_full * sig_x / (2.0 * y_safe) + + tl.store( + DLOGITS_PTR + n_offs[:, None] * E + e_idx[None, :], + d_logits, + mask=n_mask[:, None] & e_mask[None, :], + ) + + +# --------------------------------------------------------------------------- +# autograd.Function wrapper +# --------------------------------------------------------------------------- + + +class V4RouterPostFn(torch.autograd.Function): + """Autograd-aware wrapper around the FWD/BWD Triton kernels. + + Inputs: + logits [N, E] fp32 -- gate output (pre-score-fn). + indices [N, K] long -- gather positions (from host-side topk + or hash table). Must be unique per row. + score_function: "softmax" | "sigmoid" | "sqrtsoftplus". + topk_scaling_factor: float multiplier applied after denom. + out_dtype: dtype of the returned `probs` tensor. + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, + logits: torch.Tensor, + indices: torch.Tensor, + score_function: str, + topk_scaling_factor: float, + out_dtype: torch.dtype, + ): + if logits.dim() != 2: + raise ValueError(f"logits must be [N, E], got shape {tuple(logits.shape)}") + if indices.dim() != 2: + raise ValueError(f"indices must be [N, K], got shape {tuple(indices.shape)}") + if logits.shape[0] != indices.shape[0]: + raise ValueError(f"logits N={logits.shape[0]} != indices N={indices.shape[0]}") + if score_function not in _SCORE_FN_MAP: + raise ValueError( + f"Unknown score_function: {score_function!r}; expected one of " + f"{sorted(_SCORE_FN_MAP.keys())}" + ) + N, E = logits.shape + _, K = indices.shape + # Arbitrary (non-power-of-2) E / K are supported: the kernel tiles + # both axes by their next power of 2 and masks the padded columns. + block_e = triton.next_power_of_2(E) + block_k = triton.next_power_of_2(K) + + score_fn_enum = _SCORE_FN_MAP[score_function] + logits_c = logits.contiguous().to(torch.float32) + indices_c = indices.contiguous().to(torch.int64) + # Pad the indices tensor to [N, BLOCK_K] so the kernel's BLOCK_K-wide + # loads stay in-bounds; padded columns (value 0) are masked out + # everywhere via k_mask so they never gather / scatter / grad. + if block_k != K: + indices_pad = torch.zeros((N, block_k), dtype=torch.int64, device=indices_c.device) + indices_pad[:, :K] = indices_c + else: + indices_pad = indices_c + + device = logits_c.device + probs = torch.zeros((N, E), dtype=out_dtype, device=device) + routing_map = torch.zeros((N, E), dtype=torch.bool, device=device) + scores_saved = torch.empty((N, E), dtype=torch.float32, device=device) + weights_saved = torch.zeros((N, block_k), dtype=torch.float32, device=device) + denom_saved = torch.empty((N,), dtype=torch.float32, device=device) + + # BLOCK_N heuristic: per-row state ~ BLOCK_E + BLOCK_K + few scalars. + if block_e <= 64: + block_n = 64 + elif block_e <= 256: + block_n = 16 + else: + block_n = 4 + grid = (triton.cdiv(N, block_n),) + + _v4_router_post_fwd_kernel[grid]( + logits_c, + indices_pad, + probs, + routing_map, + scores_saved, + weights_saved, + denom_saved, + N, + E, + K, + SCORE_FN=score_fn_enum, + SCALE=float(topk_scaling_factor), + EPS=1e-12, + BLOCK_N=block_n, + BLOCK_E=block_e, + BLOCK_K=block_k, + OUT_DTYPE={ + torch.float32: tl.float32, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, + torch.float64: tl.float64, + }[out_dtype], + ) + + ctx.save_for_backward(indices_pad, scores_saved, weights_saved, denom_saved) + ctx.score_fn_enum = score_fn_enum + ctx.scale = float(topk_scaling_factor) + ctx.E = E + ctx.K = K + ctx.block_e = block_e + ctx.block_k = block_k + ctx.block_n = block_n + ctx.in_dtype = logits.dtype + return probs, routing_map + + @staticmethod + def backward(ctx, d_probs, d_routing_map): # type: ignore[override] + indices_pad, scores_saved, weights_saved, denom_saved = ctx.saved_tensors + E = ctx.E + K = ctx.K + block_e = ctx.block_e + block_k = ctx.block_k + block_n = ctx.block_n + score_fn_enum = ctx.score_fn_enum + scale = ctx.scale + in_dtype = ctx.in_dtype + + N = indices_pad.shape[0] + device = indices_pad.device + + d_probs = d_probs.contiguous() + # d_routing_map ignored (bool tensor, no grad flow). + # Initialise d_logits to 0 so the scatter writes only at K positions per row. + d_logits_fp32 = torch.zeros((N, E), dtype=torch.float32, device=device) + + grid = (triton.cdiv(N, block_n),) + _v4_router_post_bwd_kernel[grid]( + d_probs, + indices_pad, + scores_saved, + weights_saved, + denom_saved, + d_logits_fp32, + N, + E, + K, + SCORE_FN=score_fn_enum, + SCALE=scale, + EPS=1e-12, + BLOCK_N=block_n, + BLOCK_E=block_e, + BLOCK_K=block_k, + ) + + return d_logits_fp32.to(in_dtype), None, None, None, None + + +# --------------------------------------------------------------------------- +# Public entry points +# --------------------------------------------------------------------------- + + +def is_triton_path_enabled() -> bool: + """Return True iff ``PRIMUS_V4_ROUTER_TRITON != "0"`` (default ``"1"``). + + Plan-8 P57 close-out 2 (2026-05-15): default flipped from ``"0"`` + to ``"1"``. Microbench at V4-Flash widths is a clear positive on + V4's production `sqrtsoftplus` score function (1.56x FWD / 1.22x + BWD); EP=8 proxy A/B (P39 / P43) sat inside the proxy noise band, + so we default the microbench-positive kernel ON to keep it on the + production path. Set ``PRIMUS_V4_ROUTER_TRITON=0`` to revert to + the eager body. + """ + return os.environ.get("PRIMUS_V4_ROUTER_TRITON", "1") != "0" + + +def v4_router_post_triton( + logits: torch.Tensor, + indices: torch.Tensor, + *, + score_function: str, + topk_scaling_factor: float, + out_dtype: torch.dtype, +): + """Run the fused V4 router post-logits kernel. + + Returns ``(probs, routing_map)``. Caller pre-computes ``indices`` + (hash router via ``tid2eid[flat_ids]``; learned router via + ``torch.topk(sel_score, K).indices``). + + Arbitrary (non-power-of-2) ``E`` / ``K`` are supported — the kernel + tiles both axes by the next power of 2 and masks the padded columns — + so the only hard requirement is that the tensors live on the GPU. + """ + assert logits.is_cuda and indices.is_cuda, "v4_router_post_triton requires CUDA / HIP tensors" + return V4RouterPostFn.apply(logits, indices, score_function, topk_scaling_factor, out_dtype) + + +__all__ = [ + "V4RouterPostFn", + "v4_router_post_triton", + "is_triton_path_enabled", +] diff --git a/primus/backends/megatron/core/transformer/moe/shared_experts.py b/primus/backends/megatron/core/transformer/moe/shared_experts.py new file mode 100644 index 000000000..c4eeadf83 --- /dev/null +++ b/primus/backends/megatron/core/transformer/moe/shared_experts.py @@ -0,0 +1,91 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Primus shared-expert MLP with fused (clamped) SwiGLU. + +Megatron's stock :class:`SharedExpertMLP` reaches the fused SwiGLU path via +``mlp.MLP.forward`` only through ``bias_swiglu_impl``, which does **not** +support DeepSeek-V4's pre-multiplication clamp. To keep the clamp correct, +``v4_moe`` disables ``bias_activation_fusion`` for the shared expert, which +forces the un-fused eager ``chunk``/``clamp``/``SiLU``/``mul`` path. + +:class:`PrimusSharedExpertMLP` overrides the activation to call Primus's fused +clamped SwiGLU Triton kernel (:func:`swiglu_impl`), mirroring what +``PrimusGroupedMLP`` does for the routed experts. Both the normal ``forward`` +and the ``--moe-shared-expert-overlap`` ``linear_fc1_forward_and_act`` paths +are covered so the clamp semantics stay identical. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from megatron.core.transformer.moe.shared_experts import ( + SharedExpertMLP, + set_tensor_grad_fn_sequence_sr, +) +from megatron.core.typed_torch import apply_module +from megatron.core.utils import nvtx_range_pop, nvtx_range_push + +from primus.backends.megatron.core.fusions.fused_bias_swiglu import swiglu_impl + + +class PrimusSharedExpertMLP(SharedExpertMLP): + """Shared-expert MLP that fuses the (clamped) SwiGLU activation.""" + + def _can_fuse_swiglu(self) -> bool: + return ( + not self.config.use_te_activation_func + and self.config.gated_linear_unit + and self.activation_func == F.silu + ) + + def _fused_swiglu(self, intermediate_parallel, bias_parallel): + # dtype and the pre-mul clamp are handled inside the fused kernel. + return swiglu_impl( + intermediate_parallel, + bias_parallel, + self.config.activation_func_fp8_input_store, + self.config.activation_func_clamp_value, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Forward with fused clamped SwiGLU (non-overlap path).""" + if not self._can_fuse_swiglu(): + return super().forward(hidden_states) + + nvtx_range_push(suffix="linear_fc1") + intermediate_parallel, bias_parallel = apply_module(self.linear_fc1)(hidden_states) + nvtx_range_pop(suffix="linear_fc1") + + nvtx_range_push(suffix="activation") + intermediate_parallel = self._fused_swiglu(intermediate_parallel, bias_parallel) + nvtx_range_pop(suffix="activation") + + nvtx_range_push(suffix="linear_fc2") + output, _ = apply_module(self.linear_fc2)(intermediate_parallel) + nvtx_range_pop(suffix="linear_fc2") + + if self.use_shared_expert_gate: + logits = torch.nn.functional.linear(hidden_states, self.gate_weight) + gate_score = torch.nn.functional.sigmoid(logits) + output = output * gate_score + return output + + def linear_fc1_forward_and_act(self, overlapped_comm_output=None): + """Overlap-path FC1 + fused clamped SwiGLU activation.""" + if not self._can_fuse_swiglu(): + return super().linear_fc1_forward_and_act(overlapped_comm_output) + + assert self.config.moe_shared_expert_overlap + assert self.cached_fc1_input is not None + if overlapped_comm_output is not None: + set_tensor_grad_fn_sequence_sr(overlapped_comm_output, torch.iinfo(torch.int).max) + with torch.cuda.stream(self.stream): + # [s, b, 4 * h/p] + intermediate_parallel, bias_parallel = apply_module(self.linear_fc1)(self.cached_fc1_input) + self.cached_fc1_input = None + self.cached_fc2_input = self._fused_swiglu(intermediate_parallel, bias_parallel) diff --git a/primus/backends/megatron/core/transformer/moe/v4_hash_router.py b/primus/backends/megatron/core/transformer/moe/v4_hash_router.py new file mode 100644 index 000000000..eae8bca01 --- /dev/null +++ b/primus/backends/megatron/core/transformer/moe/v4_hash_router.py @@ -0,0 +1,270 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Hash router for DeepSeek-V4's first ``num_hash_layers`` MoE layers. + +Reference: techblog §4 ("MoE: hash routing for the first N layers") and +the inference reference at ``DeepSeek-V4-Flash/inference/model.py:Gate`` +(the ``self.hash`` branch). + +For the first ``num_hash_layers`` of V4 (V4-Flash = 3), expert +**selection** is static: each token id is permanently assigned to a +fixed set of ``moe_router_topk`` experts via a deterministic +``tid2eid`` table. Routing **weights**, however, are *not* uniform — +they come from the same learned linear gate as the non-hash layers; we +just gather the scores at the prescribed expert ids instead of running +a top-K argmax. + +Released-checkpoint contract (per HF ``Gate.__init__``): + +* ``weight`` : ``nn.Parameter`` of shape ``[num_experts, hidden_size]`` + — the learned gate. State-dict key matches the learned router so the + V4 state-dict adapter can map ``mlp.gate.weight`` uniformly. +* ``tid2eid`` : ``nn.Parameter`` of shape ``[vocab_size, topk]``, dtype + ``int32``, ``requires_grad=False``. The released checkpoint stores + it as a parameter (not a buffer) so it is preserved across + ``state_dict`` round-trips. + +Forward semantics: + + scores = score_fn(linear(hidden, weight)) # fp32 + indices = tid2eid[token_ids] # static + weights = scores.gather(1, indices) # learned weights + if score_fn != softmax: weights /= weights.sum(-1) + weights *= route_scale + +Plan-2 P14 contract: + +* :class:`DeepseekV4HashRouter` is a standalone ``nn.Module`` that + produces sparse ``(probs, routing_map)`` with the same ``[N, + num_experts]`` shape contract as :class:`DeepseekV4LearnedRouter`. +* ``score_function`` ∈ ``{"softmax", "sigmoid", "sqrtsoftplus"}`` and + ``topk_scaling_factor`` are honored identically. + +Back-compat alias ``HashRouter`` is exposed but deprecated; new +callers should use :class:`DeepseekV4HashRouter`. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import ( + is_triton_path_enabled as _v4_router_triton_enabled, +) +from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import ( + v4_router_post_triton, +) +from primus.backends.megatron.core.transformer.moe.v4_topk_router import ( + _VALID_SCORE_FUNCTIONS, + v4_score_fn, +) + + +def _build_default_tid2eid(*, vocab_size: int, num_experts: int, topk: int, seed: int) -> torch.Tensor: + """Build a deterministic ``[vocab_size, topk]`` int32 expert table. + + Each token id gets ``topk`` distinct expert ids drawn uniformly + without replacement from ``[0, num_experts)``. The seed is fixed + across all ranks so PP / TP / EP shards see identical routing. + + The table layout matches the HF reference: ``int32`` to keep on-disk + size small, indexed via long-cast at gather time. + """ + # Build on CPU with a CPU generator so the table is bit-identical across + # ranks and devices, independent of the ambient ``torch.set_default_device`` + # (the caller colocates the result with the module's weight). + gen = torch.Generator(device="cpu").manual_seed(int(seed)) + rows = [] + for _ in range(int(vocab_size)): + perm = torch.randperm(num_experts, generator=gen, device="cpu")[:topk] + rows.append(perm) + table = torch.stack(rows, dim=0).to(torch.int32) + return table + + +class DeepseekV4HashRouter(nn.Module): + """Static hash-based MoE router with a *learned* score gate. + + Args: + hidden_size: model dim ``D``; the gate is a single ``D -> + num_experts`` linear shared in shape with the learned router. + num_experts: total number of routed experts. + topk: number of experts each token is routed to. + vocab_size: tokenizer vocabulary size; controls the table length. + seed: deterministic seed for the hash table; same across all + ranks. Used only when ``tid2eid`` is built locally; if a + checkpoint provides ``tid2eid`` directly, the seed is + ignored at load time. + score_function: one of ``{"softmax", "sigmoid", "sqrtsoftplus"}``; + applied to the learned scores (matches the learned router). + topk_scaling_factor: scalar multiplier applied to the + renormalized routing weights (V3 ``moe_router_topk_scaling_factor``, + HF ``Gate.route_scale``). + dtype: dtype of the gate weight; defaults to fp32. + + Parameters: + weight: ``[num_experts, hidden_size]`` learned gate. + tid2eid: ``[vocab_size, topk]`` int32 frozen mapping. + ``requires_grad=False``; this is a parameter (matching the + HF reference checkpoint) so that ``state_dict`` round-trips + preserve it. + """ + + def __init__( + self, + *, + hidden_size: int, + num_experts: int, + topk: int, + vocab_size: int, + seed: int = 0, + score_function: str = "sqrtsoftplus", + topk_scaling_factor: float = 1.0, + dtype: Optional[torch.dtype] = None, + ) -> None: + super().__init__() + if num_experts <= 0: + raise ValueError(f"num_experts must be > 0, got {num_experts}") + if topk <= 0 or topk > num_experts: + raise ValueError(f"topk must be in [1, {num_experts}], got {topk}") + if vocab_size <= 0: + raise ValueError(f"vocab_size must be > 0, got {vocab_size}") + if score_function not in _VALID_SCORE_FUNCTIONS: + raise ValueError( + f"Unknown score_function: {score_function!r}. " + f"Expected one of {sorted(_VALID_SCORE_FUNCTIONS)}." + ) + + self.hidden_size = int(hidden_size) + self.num_experts = int(num_experts) + self.topk = int(topk) + self.vocab_size = int(vocab_size) + self.seed = int(seed) + self.score_function = str(score_function) + self.topk_scaling_factor = float(topk_scaling_factor) + + weight_dtype = dtype or torch.float32 + self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, dtype=weight_dtype)) + nn.init.normal_(self.weight, mean=0.0, std=0.02) + + # tid2eid is a non-trainable parameter (matches HF reference layout + # so checkpoint round-trips include it). int32 to keep memory small. + tid2eid_init = _build_default_tid2eid( + vocab_size=self.vocab_size, + num_experts=self.num_experts, + topk=self.topk, + seed=self.seed, + ) + # Colocate the CPU-built table with the module's weight (the ambient + # default device at construction, e.g. CUDA under set_default_device). + self.tid2eid = nn.Parameter(tid2eid_init.to(self.weight.device), requires_grad=False) + + # ------------------------------------------------------------------ + + def forward( + self, + hidden: torch.Tensor, + token_ids: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Route tokens whose expert ids are prescribed by ``tid2eid``. + + Args: + hidden: ``[B, S, D]`` (or any shape with last dim ``D``). + The learned gate is evaluated against ``hidden``; the + resulting scores supply the *weights*. + token_ids: ``[B, S]`` (or any compatible shape) integer + tensor of token ids. Provides the *indices* via + ``tid2eid``. + + Returns: + probs: ``[N, num_experts]`` float tensor, ``N = numel/D``. + Non-selected experts have probability 0; selected + experts hold the (renormalized + scaled) score. + routing_map: ``[N, num_experts]`` bool tensor; ``True`` at + ``(n, e)`` iff token ``n`` is routed to expert ``e``. + """ + if (token_ids.dtype != torch.long) and (token_ids.dtype != torch.int): + raise TypeError(f"DeepseekV4HashRouter expects integer token_ids, got {token_ids.dtype}") + flat_hidden = hidden.reshape(-1, self.hidden_size) + flat_ids = token_ids.reshape(-1).long() + if flat_hidden.shape[0] != flat_ids.shape[0]: + raise ValueError( + "DeepseekV4HashRouter: hidden and token_ids must flatten to the same length; " + f"got hidden={flat_hidden.shape[0]} vs token_ids={flat_ids.shape[0]}." + ) + if flat_ids.numel() == 0: + n_zero = 0 + device = flat_hidden.device + probs = torch.zeros(n_zero, self.num_experts, dtype=torch.float32, device=device) + routing_map = torch.zeros(n_zero, self.num_experts, dtype=torch.bool, device=device) + return probs, routing_map + # NOTE: the token-id bounds check below is intentionally disabled — the + # ``.item()`` forces a device->host sync every hash-router forward, which + # stalls the CPU (shows up as a blocking cpu_op in the trace). token_ids + # come straight from the (already-validated) input pipeline, so the check + # is redundant on the hot path. Re-enable for debugging if needed. + # if int(flat_ids.max().item()) >= self.vocab_size: + # raise ValueError( + # f"token_ids has values >= vocab_size ({self.vocab_size}); " + # f"max found = {int(flat_ids.max().item())}" + # ) + + # Learned scores (fp32) over the full expert axis. + logits = F.linear(flat_hidden.to(torch.float32), self.weight.to(torch.float32)) + + # Static expert assignment from the table — cast to long for gather. + indices = self.tid2eid[flat_ids].long() # [N, K] + + # The V4 router is GPU-only: both the Triton and eager paths run on + # CUDA/HIP tensors (there is no CPU compute path). + assert logits.is_cuda, "V4 router requires CUDA / HIP tensors" + # Plan-6 P39: route the post-logits chain through one fused Triton + # kernel when PRIMUS_V4_ROUTER_TRITON=1 (default ON); else the eager body. + if _v4_router_triton_enabled(): + probs, routing_map = v4_router_post_triton( + logits, + indices, + score_function=self.score_function, + topk_scaling_factor=self.topk_scaling_factor, + out_dtype=torch.float32, + ) + return probs, routing_map + + # Eager fallback (verbatim pre-P39 body). + scores = v4_score_fn(logits, score_function=self.score_function) # [N, E] + weights = scores.gather(1, indices) # [N, K] + + if self.score_function != "softmax": + denom = weights.sum(dim=-1, keepdim=True).clamp(min=1.0e-12) + weights = weights / denom + + if self.topk_scaling_factor != 1.0: + weights = weights * float(self.topk_scaling_factor) + + N = flat_hidden.shape[0] + device = flat_hidden.device + + probs = torch.zeros(N, self.num_experts, dtype=weights.dtype, device=device) + probs.scatter_(1, indices, weights) + + routing_map = torch.zeros(N, self.num_experts, dtype=torch.bool, device=device) + routing_map.scatter_(1, indices, True) + + return probs, routing_map + + +# Back-compat alias. New callers should use ``DeepseekV4HashRouter``. +HashRouter = DeepseekV4HashRouter + +__all__ = [ + "DeepseekV4HashRouter", + "HashRouter", +] diff --git a/primus/backends/megatron/core/transformer/moe/v4_moe.py b/primus/backends/megatron/core/transformer/moe/v4_moe.py new file mode 100644 index 000000000..0748c83f5 --- /dev/null +++ b/primus/backends/megatron/core/transformer/moe/v4_moe.py @@ -0,0 +1,633 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 Mixture-of-Experts module. + +Reference: techblog §4 ("MoE: hash routing + sqrtsoftplus + shared experts") +and ``DeepSeek-V4-Flash/inference/model.py:MoE``. + +V4's MoE block has three pieces: + +1. **Router** — either :class:`DeepseekV4HashRouter` (first + ``num_hash_layers`` layers) or :class:`DeepseekV4LearnedRouter` (the + rest). Both produce the same ``(probs, routing_map)`` shape contract: + ``[N, num_experts]``. The two routers share a learned gate weight; + only the *selection* differs (top-K argmax for the learned router, + ``tid2eid`` lookup for the hash router). Routing weights always come + from the same ``v4_score_fn(linear(hidden, weight))`` path. +2. **Routed experts** — ``num_experts`` clamped-SwiGLU MLPs. Each token + contributes to ``moe_router_topk`` of them, weighted by the router + probability. The clamp is **pre-multiplication**: + ``SiLU(clamp(gate, max=alpha)) * clamp(up, +/- alpha)``. +3. **Shared expert(s)** — always-on MLP(s) whose output is added to every + token's contribution. V4-Flash has 1 shared expert with the same + ``moe_intermediate_size`` as the routed experts. + +Plan-2 P14 contract: + +P14 phase-1 (committed in 1a8bf32e) — math + parameter-layout +faithfulness: pre-multiplication clamped SwiGLU activation, learned +router rewritten with HF-aligned scoring + bias-only-for-selection +semantics, hash router rewritten with a learnable gate weight + frozen +``tid2eid`` Parameter. + +P14 phase-2 (this commit) — structural bring-up: +* :class:`DeepseekV4MoE` now subclasses :class:`MegatronModule` (was + ``nn.Module``) so it integrates with Megatron's spec lifecycle and + shares config plumbing with the rest of the V4 stack. +* CPU-friendly local-experts path: when ``pg_collection`` is ``None`` + (or when the grouped backend does not declare clamped-SwiGLU support), + :class:`DeepseekV4MoE` builds a :class:`nn.ModuleList` of + :class:`ClampedSwiGLUMLP` routed experts plus a single + :class:`ClampedSwiGLUMLP` shared expert and runs a per-expert dispatch + loop in ``forward`` that mirrors the HF reference exactly. This makes + the MoE forward unit-testable on CPU at G5 (1L MoE forward agreement + vs HF reference within 1e-3 fp32) without requiring distributed init. +* :meth:`set_layer_number` mirrors :class:`BaseMoELayer` so this module + slots into ``TransformerLayer`` via the spec lifecycle. +* :attr:`local_expert_indices` exposed for compatibility with downstream + tooling that expects the ``BaseMoELayer`` public surface. + +Aux-loss / z-loss inheritance via :class:`TopKRouter` is left as a +follow-up: the V4 routers are standalone ``nn.Module``\\ s rather than +subclasses of Megatron's :class:`TopKRouter` (the parent registers CUDA +buffers in ``__init__`` and is impractical to instantiate on CPU). The +distributed re-validation phase (P19) will re-introduce that path +behind a TopKRouter subclass once the CUDA-buffer init is gated by a +device check upstream. +""" + +from __future__ import annotations + +import logging +from copy import copy +from dataclasses import dataclass +from typing import Optional, Union + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from megatron.core import parallel_state +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.moe.shared_experts import SharedExpertMLP +from megatron.core.transformer.moe.token_dispatcher import ( + MoEAllGatherTokenDispatcher, + MoEAlltoAllTokenDispatcher, + MoEFlexTokenDispatcher, + MoETokenDispatcher, +) +from megatron.core.transformer.spec_utils import ModuleSpec, build_module + +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, +) +from primus.backends.megatron.core.transformer.clamped_swiglu import ClampedSwiGLUMLP +from primus.backends.megatron.core.transformer.moe.shared_experts import ( + PrimusSharedExpertMLP, +) +from primus.backends.megatron.core.transformer.moe.v4_hash_router import ( + DeepseekV4HashRouter, +) +from primus.backends.megatron.core.transformer.moe.v4_topk_router import ( + DeepseekV4LearnedRouter, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class DeepseekV4MoESubmodules: + """Spec tree for V4 MoE construction.""" + + hash_router: Optional[Union[ModuleSpec, type]] = DeepseekV4HashRouter + learned_router: Optional[Union[ModuleSpec, type]] = DeepseekV4LearnedRouter + token_dispatcher: Optional[Union[ModuleSpec, type]] = MoEAlltoAllTokenDispatcher + grouped_experts: Optional[Union[ModuleSpec, type]] = None + shared_expert: Optional[Union[ModuleSpec, type]] = SharedExpertMLP + + +class DeepseekV4MoE(MegatronModule): + """V4 MoE FFN sub-block. + + Args: + config: runtime DeepSeek-V4 config. Core MoE dimensions and router + options are read directly from config. + layer_idx: 0-based decoder layer index. Used to pick router type + against ``num_hash_layers``. + pg_collection: Megatron process-group collection. When ``None`` + (CPU unit tests), the module skips the distributed dispatcher + and builds a local :class:`nn.ModuleList` of + :class:`ClampedSwiGLUMLP` routed experts plus a single + :class:`ClampedSwiGLUMLP` shared expert; ``forward`` runs a + per-expert dispatch loop matching the HF reference math. + submodules: spec tree describing routers / dispatcher / experts / + shared expert. Must be provided. + layer_number: optional 1-based layer number used by Megatron's + spec lifecycle (mirrors :class:`BaseMoELayer.set_layer_number`). + """ + + def __init__( + self, + config: DeepSeekV4TransformerConfig, + *, + layer_idx: int, + pg_collection=None, + submodules: Optional[DeepseekV4MoESubmodules] = None, + layer_number: Optional[int] = None, + ) -> None: + if config is None: + raise ValueError("DeepSeek-V4 MoE requires config.") + super().__init__(config=config) + self.pg_collection = pg_collection + self.submodules = submodules + assert self.submodules is not None, "DeepSeek-V4 MoE requires explicit submodules." + self.layer_number = layer_number + + hidden_size = int(config.hidden_size) + moe_intermediate_size = int( + config.moe_ffn_hidden_size or config.moe_intermediate_size or config.ffn_hidden_size + ) + num_routed_experts = int(config.num_moe_experts) + moe_router_topk = int(config.moe_router_topk) + use_shared_expert = config.moe_shared_expert_intermediate_size is not None + layer_num_hash_layers = int(config.num_hash_layers) + layer_hash_vocab_size = config.padded_vocab_size or config.vocab_size + layer_hash_seed = int(config.hash_routing_seed) + score_function = str(config.moe_router_score_function) + enable_expert_bias = bool(config.moe_router_enable_expert_bias) + topk_scaling_factor = float(getattr(config, "moe_router_topk_scaling_factor", 1.0) or 1.0) + clamp_alpha = float(config.swiglu_limit) + + if num_routed_experts <= 0: + raise ValueError(f"num_routed_experts must be > 0, got {num_routed_experts}") + if moe_router_topk <= 0 or moe_router_topk > num_routed_experts: + raise ValueError(f"moe_router_topk must be in [1, {num_routed_experts}], got {moe_router_topk}") + + self.hidden_size = hidden_size + self.moe_intermediate_size = moe_intermediate_size + self.num_routed_experts = num_routed_experts + self.moe_router_topk = moe_router_topk + self.use_shared_expert = use_shared_expert + self.layer_idx = int(layer_idx) + self.num_hash_layers = layer_num_hash_layers + self.use_hash_router = self.layer_idx < self.num_hash_layers + self.clamp_alpha = clamp_alpha + self.moe_token_dispatcher_type = "alltoall" + + # ---- EP placement ---- + self.ep_group = getattr(pg_collection, "ep", None) if pg_collection is not None else None + self.ep_size = 1 + self.ep_rank = 0 + if self.ep_group is None and dist.is_available() and dist.is_initialized(): + try: + self.ep_group = parallel_state.get_expert_model_parallel_group() + except Exception: + self.ep_group = None + if self.ep_group is not None and dist.is_available() and dist.is_initialized(): + self.ep_size = int(self.ep_group.size()) + self.ep_rank = int(self.ep_group.rank()) + + base = self.num_routed_experts // self.ep_size + remainder = self.num_routed_experts % self.ep_size + self.local_num_routed_experts = base + (1 if self.ep_rank < remainder else 0) + self.local_expert_start = (self.ep_rank * base) + min(self.ep_rank, remainder) + self.local_expert_end = self.local_expert_start + self.local_num_routed_experts + # BaseMoELayer-compatible public attribute. + self.local_expert_indices = list(range(self.local_expert_start, self.local_expert_end)) + + # ---- routers ---- + self.router = None + self.learned_router = None + self._build_router_modules( + hash_vocab_size=layer_hash_vocab_size, + hash_seed=layer_hash_seed, + score_function=score_function, + enable_expert_bias=enable_expert_bias, + topk_scaling_factor=topk_scaling_factor, + ) + + # ---- experts ---- + # Production path: full Megatron dispatcher + grouped-experts. + # CPU path: a local nn.ModuleList of ClampedSwiGLUMLP experts + a + # single ClampedSwiGLUMLP shared expert. The CPU path is used when + # ``pg_collection is None`` so unit tests can drive ``forward`` + # without distributed initialization. + self.token_dispatcher: Optional[MoETokenDispatcher] = None + self.grouped_experts: Optional[nn.Module] = None + self.local_experts: Optional[nn.ModuleList] = None + self.shared_expert: Optional[nn.Module] = None + + if pg_collection is None: + self.local_experts = self._build_local_experts(intermediate_size=self.moe_intermediate_size) + if self.use_shared_expert: + assert self.config.moe_shared_expert_intermediate_size is not None + self.shared_expert = ClampedSwiGLUMLP( + hidden_size=self.hidden_size, + intermediate_size=int(self.config.moe_shared_expert_intermediate_size), + alpha=self.clamp_alpha, + bias=False, + ) + else: + self.token_dispatcher = self._build_token_dispatcher() + self.grouped_experts = self._build_grouped_experts() + if self.use_shared_expert: + assert self.config.moe_shared_expert_intermediate_size is not None + self.shared_expert = self._build_shared_expert_module( + intermediate_size=int(self.config.moe_shared_expert_intermediate_size) + ) + + # ------------------------------------------------------------------ + + def set_layer_number(self, layer_number: int) -> None: + """Mirror :class:`BaseMoELayer.set_layer_number` for spec lifecycle. + + Megatron's :class:`TransformerLayer` walks every spec submodule + with a ``set_layer_number`` method to populate the 1-based layer + index. The V4 routers are intentionally standalone (CPU-clean), + but we still need to track ``layer_number`` here so future + TopKRouter-rooted upgrades plug in without spec changes. + """ + self.layer_number = layer_number + + def _build_local_experts(self, *, intermediate_size: int) -> nn.ModuleList: + """Build a local :class:`nn.ModuleList` of clamped-SwiGLU experts. + + Used when ``pg_collection is None`` (CPU unit tests). Each module + in the list mirrors a single HF reference ``Expert`` (separate + ``w1`` / ``w2`` / ``w3`` Linears + V4 pre-mul clamp). + """ + if self.local_num_routed_experts <= 0: + raise RuntimeError(f"DeepSeek-V4 MoE layer={self.layer_idx} has no local experts.") + return nn.ModuleList( + [ + ClampedSwiGLUMLP( + hidden_size=self.hidden_size, + intermediate_size=intermediate_size, + alpha=self.clamp_alpha, + bias=False, + ) + for _ in range(self.local_num_routed_experts) + ] + ) + + # ------------------------------------------------------------------ + + @staticmethod + def _resolve_dispatcher_type_from_spec(dispatcher_spec: Optional[Union[ModuleSpec, type]]) -> str: + module = dispatcher_spec.module if isinstance(dispatcher_spec, ModuleSpec) else dispatcher_spec + if module is MoEAllGatherTokenDispatcher: + return "allgather" + if module is MoEFlexTokenDispatcher: + return "flex" + if module is MoEAlltoAllTokenDispatcher or module is None: + return "alltoall" + # Plan-3 P23: PrimusTurboDeepEPTokenDispatcher is a "flex" + # variant — V4 spec build chose it explicitly when the user + # opted into Turbo DeepEP. We recognise it by class name so + # this resolver does not require ``primus_turbo`` to be + # importable on hosts that never opt in (CPU unit tests, + # Megatron-only consumers). + module_name = getattr(module, "__name__", "") + if module_name == "PrimusTurboDeepEPTokenDispatcher": + return "flex" + logger.warning( + "[DeepSeek-V4] unsupported dispatcher module=%s; fallback type to alltoall.", + module_name or str(module), + ) + return "alltoall" + + def _build_router_modules( + self, + *, + hash_vocab_size: Optional[int], + hash_seed: int, + score_function: str, + enable_expert_bias: bool, + topk_scaling_factor: float, + ) -> None: + if self.use_hash_router: + if hash_vocab_size is None or hash_vocab_size <= 0: + raise ValueError( + "hash_vocab_size must be provided (and > 0) when layer_idx < num_hash_layers" + ) + hash_router_spec = self.submodules.hash_router or DeepseekV4HashRouter + self.router = build_module( + hash_router_spec, + hidden_size=self.hidden_size, + num_experts=self.num_routed_experts, + topk=self.moe_router_topk, + vocab_size=hash_vocab_size, + seed=hash_seed, + score_function=score_function, + topk_scaling_factor=topk_scaling_factor, + ) + self.learned_router = None + return + + learned_router_spec = self.submodules.learned_router or DeepseekV4LearnedRouter + self.router = None + self.learned_router = build_module( + learned_router_spec, + hidden_size=self.hidden_size, + num_experts=self.num_routed_experts, + topk=self.moe_router_topk, + score_function=score_function, + enable_expert_bias=enable_expert_bias, + topk_scaling_factor=topk_scaling_factor, + ) + + def _build_shared_expert_module(self, *, intermediate_size: int) -> nn.Module: + shared_expert_spec = self.submodules.shared_expert + assert isinstance( + shared_expert_spec, ModuleSpec + ), "DeepSeek-V4 MoE requires shared_expert ModuleSpec in submodules." + shared_expert_module = shared_expert_spec.module + assert issubclass( + shared_expert_module, SharedExpertMLP + ), "DeepSeek-V4 shared_expert must be (a subclass of) SharedExpertMLP." + if self.config is None or self.pg_collection is None: + raise RuntimeError("DeepSeek-V4 MoE SharedExpertMLP requires config and pg_collection.") + + # Shared experts run with clamped SwiGLU. PrimusSharedExpertMLP fuses the + # clamp+SiLU+mul into a single Triton kernel (matching the routed experts), + # so we no longer need to force the un-fused eager path. + shared_cfg = copy(self.config) + shared_cfg.add_bias_linear = False + shared_cfg.gated_linear_unit = True + shared_cfg.activation_func = F.silu + shared_cfg.bias_activation_fusion = False + shared_cfg.use_te_activation_func = False + if self.clamp_alpha > 0: + shared_cfg.activation_func_clamp_value = float(self.clamp_alpha) + else: + shared_cfg.activation_func_clamp_value = None + if int(shared_cfg.moe_shared_expert_intermediate_size or 0) <= 0: + setattr( + shared_cfg, + "moe_shared_expert_intermediate_size", + int(intermediate_size), + ) + + # Build the Primus fused-SwiGLU shared expert while keeping the spec's + # submodules (linear layers / activation). The state-dict layout is + # unchanged since PrimusSharedExpertMLP only overrides the activation. + fused_shared_expert_spec = ModuleSpec( + module=PrimusSharedExpertMLP, + submodules=shared_expert_spec.submodules, + params=shared_expert_spec.params, + ) + try: + return build_module( + fused_shared_expert_spec, + config=shared_cfg, + pg_collection=self.pg_collection, + gate=bool(shared_cfg.moe_shared_expert_gate), + ) + except Exception as exc: + raise RuntimeError( + f"DeepSeek-V4 MoE shared expert build failed with PrimusSharedExpertMLP: {exc}" + ) from exc + + def _build_token_dispatcher(self) -> MoETokenDispatcher: + if self.config is None or self.pg_collection is None: + raise RuntimeError( + "DeepSeek-V4 MoE requires config and pg_collection for Megatron dispatcher path." + ) + if self.local_num_routed_experts <= 0: + raise RuntimeError( + f"DeepSeek-V4 MoE layer={self.layer_idx} has no local experts for dispatcher path." + ) + + dispatcher_spec: Union[ModuleSpec, type, None] = self.submodules.token_dispatcher + assert dispatcher_spec is not None, "DeepSeek-V4 MoE requires token_dispatcher spec in submodules." + requested_dispatcher_type = self._resolve_dispatcher_type_from_spec(dispatcher_spec) + ep_group = getattr(self.pg_collection, "ep", None) + tp_ep_group = getattr(self.pg_collection, "tp_ep", None) + if requested_dispatcher_type == "alltoall" and ep_group is None: + logger.info( + "[DeepSeek-V4] MoE layer=%s alltoall dispatcher requires EP group.", + self.layer_idx, + ) + if requested_dispatcher_type == "flex" and tp_ep_group is None: + logger.info( + "[DeepSeek-V4] MoE layer=%s flex dispatcher requires TPxEP group.", + self.layer_idx, + ) + self.moe_token_dispatcher_type = requested_dispatcher_type + + local_expert_indices = list(range(self.local_expert_start, self.local_expert_end)) + try: + dispatcher = build_module( + dispatcher_spec, + num_local_experts=self.local_num_routed_experts, + local_expert_indices=local_expert_indices, + config=self.config, + pg_collection=self.pg_collection, + ) + logger.info( + "[DeepSeek-V4] MoE layer=%s dispatcher active via %s.", + self.layer_idx, + type(dispatcher).__name__, + ) + return dispatcher + except Exception as exc: + raise RuntimeError(f"DeepSeek-V4 MoE layer={self.layer_idx} dispatcher build failed: {exc}") + + def _route( + self, + hidden: torch.Tensor, + token_ids: Optional[torch.Tensor], + ): + """Return ``(probs, routing_map)`` for the current router. + + Hash-routed layers feed both ``hidden`` (for the learned routing + weights) AND ``token_ids`` (for the static expert ids from + ``tid2eid``); learned layers only consume ``hidden``. + """ + if self.use_hash_router: + assert self.router is not None + if token_ids is None: + raise ValueError( + f"layer {self.layer_idx} uses DeepseekV4HashRouter; " + "token_ids is required (shape [B, S])." + ) + return self.router(hidden, token_ids) + assert self.learned_router is not None + return self.learned_router(hidden) + + # ------------------------------------------------------------------ + + def _build_grouped_experts(self): + grouped_experts_spec: Optional[Union[ModuleSpec, type]] = self.submodules.grouped_experts + assert ( + grouped_experts_spec is not None + ), "DeepSeek-V4 MoE requires grouped experts spec in submodules." + if self.local_num_routed_experts <= 0: + raise RuntimeError( + f"DeepSeek-V4 MoE layer={self.layer_idx} has no local experts for grouped backend." + ) + + if self.config is None or self.pg_collection is None: + raise RuntimeError("DeepSeek-V4 MoE requires config and pg_collection to build grouped experts.") + + try: + module = build_module( + grouped_experts_spec, + num_local_experts=self.local_num_routed_experts, + config=self.config, + pg_collection=self.pg_collection, + ) + if not self._grouped_backend_supports_clamped_swiglu(module): + raise RuntimeError( + "DeepSeek-V4 MoE grouped backend " + f"{type(module).__name__} does not declare clamped-SwiGLU support. " + "Set `v4_grouped_experts_support_clamped_swiglu=True` only " + "after backend parity is validated." + ) + logger.info( + "[DeepSeek-V4] MoE layer=%s provider grouped-gemm active via %s.", + self.layer_idx, + type(module).__name__, + ) + return module + except Exception as exc: + raise RuntimeError(f"DeepSeek-V4 MoE layer={self.layer_idx} grouped experts build failed: {exc}") + + def _grouped_backend_supports_clamped_swiglu(self, module: nn.Module) -> bool: + if self.clamp_alpha <= 0: + return True + if bool(getattr(module, "supports_clamped_swiglu", False)): + return True + if self.config is not None and bool(self.config.v4_grouped_experts_support_clamped_swiglu): + return True + return False + + def _dispatcher_expert_forward( + self, + permuted_hidden: torch.Tensor, + tokens_per_expert: torch.Tensor, + permuted_probs: torch.Tensor, + routing_map: torch.Tensor, + ) -> torch.Tensor: + assert self.grouped_experts is not None + try: + grouped_out = self.grouped_experts( + permuted_hidden, + tokens_per_expert, + permuted_probs, + routing_map=routing_map, + ) + except TypeError: + grouped_out = self.grouped_experts( + permuted_hidden, + tokens_per_expert, + permuted_probs, + ) + if isinstance(grouped_out, tuple): + return grouped_out[0] + return grouped_out + + def _dispatcher_forward( + self, + hidden: torch.Tensor, + probs: torch.Tensor, + routing_map: torch.Tensor, + ) -> torch.Tensor: + assert self.token_dispatcher is not None + hidden_states, probs = self.token_dispatcher.dispatch_preprocess(hidden, routing_map, probs) + hidden_states, probs = self.token_dispatcher.token_dispatch(hidden_states, probs) + expert_input, tokens_per_expert, permuted_probs = self.token_dispatcher.dispatch_postprocess( + hidden_states, probs + ) + + expert_output = self._dispatcher_expert_forward( + expert_input, + tokens_per_expert, + permuted_probs, + routing_map, + ) + + combined = self.token_dispatcher.combine_preprocess(expert_output) + combined = self.token_dispatcher.token_combine(combined) + return self.token_dispatcher.combine_postprocess(combined) + + def _local_experts_forward( + self, + hidden: torch.Tensor, + probs: torch.Tensor, + routing_map: torch.Tensor, + ) -> torch.Tensor: + """Per-expert dispatch loop matching the HF reference math. + + Drives :attr:`local_experts` directly (no token dispatcher); used + on the CPU path when ``pg_collection is None``. The math mirrors + ``DeepSeek-V4-Flash/inference/model.py:MoE.forward`` exactly: + + for i in local_experts: + idx = where(routing_map[:, i]) + out[idx] += probs[idx, i] * expert_i(hidden[idx]) + + Args: + hidden: ``[N, D]`` flattened input. + probs: ``[N, num_experts]`` sparse routing weights (already + renormalized + scaled by the router). + routing_map: ``[N, num_experts]`` bool mask, True at + ``(n, e)`` iff token ``n`` is routed to expert ``e``. + + Returns: + ``[N, D]`` routed-expert contribution (no shared expert). + """ + assert self.local_experts is not None + out = torch.zeros_like(hidden, dtype=hidden.dtype) + for local_i, global_i in enumerate(self.local_expert_indices): + mask = routing_map[:, global_i] # [N] + if not bool(mask.any()): + continue + idx = mask.nonzero(as_tuple=True)[0] # [n_i] + weight = probs[idx, global_i].unsqueeze(-1).to(hidden.dtype) # [n_i, 1] + expert = self.local_experts[local_i] + out_idx = expert(hidden[idx]) + out[idx] = out[idx] + weight * out_idx + return out + + def forward( + self, + hidden: torch.Tensor, + *, + token_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Run V4 MoE FFN. + + Args: + hidden: ``[B, S, D]`` input. + token_ids: ``[B, S]`` integer token ids, required only when + ``layer_idx < num_hash_layers``. + + Returns: + ``[B, S, D]`` output. Sum of routed-expert and shared-expert + contributions. + """ + probs, routing_map = self._route(hidden, token_ids) # [N, E], bool + + if self.local_experts is not None: + # CPU local-experts path. Reshape to flat then back; the + # router already returned [N, E] sparse outputs. + shape = hidden.shape + flat_hidden = hidden.reshape(-1, self.hidden_size) + flat_out = self._local_experts_forward(flat_hidden, probs, routing_map) + if self.shared_expert is not None: + flat_out = flat_out + self.shared_expert(flat_hidden) + return flat_out.view(*shape) + + # Production path: Megatron dispatcher + grouped experts. + out = self._dispatcher_forward(hidden, probs, routing_map) + if self.shared_expert is not None: + out = out + self.shared_expert(hidden) + return out + + +__all__ = ["DeepseekV4MoE", "DeepseekV4MoESubmodules"] diff --git a/primus/backends/megatron/core/transformer/moe/v4_topk_router.py b/primus/backends/megatron/core/transformer/moe/v4_topk_router.py new file mode 100644 index 000000000..a671de41a --- /dev/null +++ b/primus/backends/megatron/core/transformer/moe/v4_topk_router.py @@ -0,0 +1,275 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Learned Top-K MoE router for DeepSeek-V4. + +Reference: techblog §4 ("MoE: routing scoring") and the inference +reference at ``DeepSeek-V4-Flash/inference/model.py:Gate.forward``. + +For layers with ``layer_idx >= num_hash_layers`` V4 uses a learned +top-K router. The HF-released ``Gate`` module computes a single +``[D -> num_experts]`` linear and supports three score functions: + +* ``softmax`` — standard competitive normalization. +* ``sigmoid`` — independent expert scoring (V3 fallback). +* ``sqrtsoftplus`` — V4 default. ``sqrt(softplus(x))`` combines the + positive-only behavior of softplus with the sub-linear growth of + sqrt; sits between sigmoid (saturating) and softmax (competition) + and yields smoother routing gradients in long training runs. + +Optionally the router supports an **expert bias** correction +(``moe_router_enable_expert_bias`` / "noaux_tc" — V3-style auxiliary-free +balancing): a learnable per-expert bias is added to the score *only for +top-K selection*, and the returned routing weights are gathered from the +**un-biased** scores. This keeps gradient flow clean (probs flow back to +the gate weight, not the bias) while still letting the bias term steer +load balance. + +After top-K selection, the routing weights are renormalized to sum to 1 +**only when the score function is non-softmax** (matches HF; with +softmax the sum is already 1 by construction). A final scalar +``topk_scaling_factor`` ("route_scale" in the HF reference) is applied +multiplicatively. + +Plan-2 P14 contract: + +* :class:`DeepseekV4LearnedRouter` — standalone ``nn.Module`` that + produces sparse ``(probs, routing_map)`` with the same ``[N, num_experts]`` + shape contract as Megatron's :class:`TopKRouter`. The eager, + CPU-testable form is the canonical reference for G4 unit tests. +* Parameter layout: + - ``weight``: ``nn.Parameter`` of shape ``[num_experts, hidden_size]`` + (matches both Megatron's ``TopKRouter.weight`` and HF reference + ``Gate.weight``). + - ``expert_bias`` (optional): ``nn.Parameter`` of shape + ``[num_experts]`` (matches HF reference ``Gate.bias``). +* ``score_function`` ∈ ``{"softmax", "sigmoid", "sqrtsoftplus"}``. + +Phase-2 of P14 will subclass Megatron's :class:`TopKRouter` directly so +the router participates in aux-loss / z-loss / dispatcher lifecycle in +production. The standalone form here remains the reference for unit +tests and the state-dict adapter (P17). + +Back-compat alias ``V4TopKRouter`` is exposed but deprecated; new +callers should use :class:`DeepseekV4LearnedRouter`. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import ( + is_triton_path_enabled as _v4_router_triton_enabled, +) +from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import ( + v4_router_post_triton, +) + +_VALID_SCORE_FUNCTIONS = {"softmax", "sigmoid", "sqrtsoftplus"} + + +def v4_score_fn(logits: torch.Tensor, *, score_function: str) -> torch.Tensor: + """Apply a V4-supported score function to gate logits. + + Args: + logits: ``[..., num_experts]`` tensor of pre-score linear + outputs. Must be float (fp32 in the HF reference; we follow). + score_function: one of ``"softmax"``, ``"sigmoid"``, + ``"sqrtsoftplus"``. + + Returns: + Tensor of the same shape, post score-function. ``softmax`` sums + to 1 along the expert axis; the other two are pointwise. + """ + if score_function == "softmax": + return F.softmax(logits, dim=-1) + if score_function == "sigmoid": + return torch.sigmoid(logits) + if score_function == "sqrtsoftplus": + return F.softplus(logits).sqrt() + raise ValueError( + f"Unknown score_function: {score_function!r}. " f"Expected one of {sorted(_VALID_SCORE_FUNCTIONS)}." + ) + + +def _compute_route( + *, + logits: torch.Tensor, + expert_bias: Optional[torch.Tensor], + score_function: str, + topk: int, + topk_scaling_factor: float, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Shared selection / renormalization core for V4 routers. + + Returns sparse ``(probs[N, E], routing_map[N, E])``. The dense + ``(weights[N, K], indices[N, K])`` form follows the HF reference; + we expose only the sparse contract here so downstream Megatron + dispatchers consume it directly. + + NOTE: this helper assumes ``logits`` is already shaped ``[N, + num_experts]`` and in fp32. Callers are responsible for the cast. + """ + # Pre-compute topk indices on the host (heavy GPU compute that + # benefits from being its own kernel). When PRIMUS_V4_ROUTER_TRITON + # is on (and supported), route the rest of the chain through one + # fused Triton kernel; otherwise fall back to the eager body. + scores_for_selection = v4_score_fn(logits, score_function=score_function) + if expert_bias is not None: + sel_score = scores_for_selection + expert_bias.to(scores_for_selection.dtype) + else: + sel_score = scores_for_selection + indices = sel_score.topk(topk, dim=-1).indices # [N, K] + + # The V4 router is GPU-only: both the Triton and eager paths run on + # CUDA/HIP tensors (there is no CPU compute path). + assert logits.is_cuda, "V4 router requires CUDA / HIP tensors" + if _v4_router_triton_enabled(): + # The Triton kernel re-applies the score function inside (so + # it can save the full row for backward). This keeps the + # eager path's bit-equivalent behaviour for the gathered + # weights at the cost of one extra pass over [N, E] of fp32 + # logits -- negligible at V4-Flash widths. + probs, routing_map = v4_router_post_triton( + logits, + indices, + score_function=score_function, + topk_scaling_factor=topk_scaling_factor, + out_dtype=scores_for_selection.dtype, + ) + return probs, routing_map + + # Eager fallback (verbatim pre-P39 body). + original_scores = scores_for_selection + weights = original_scores.gather(1, indices) # [N, K] + + if score_function != "softmax": + denom = weights.sum(dim=-1, keepdim=True).clamp(min=1.0e-12) + weights = weights / denom + + if topk_scaling_factor != 1.0: + weights = weights * float(topk_scaling_factor) + + num_experts = logits.shape[-1] + N = logits.shape[0] + device = logits.device + + probs = torch.zeros(N, num_experts, dtype=weights.dtype, device=device) + probs.scatter_(1, indices, weights) + + routing_map = torch.zeros(N, num_experts, dtype=torch.bool, device=device) + routing_map.scatter_(1, indices, True) + + return probs, routing_map + + +class DeepseekV4LearnedRouter(nn.Module): + """Learned top-K router for DeepSeek-V4 MoE layers (l >= num_hash_layers). + + Args: + hidden_size: model dim ``D``; the gate is a single ``D -> num_experts`` + linear. + num_experts: total number of routed experts. + topk: number of experts each token is routed to. + score_function: one of ``{"softmax", "sigmoid", "sqrtsoftplus"}``. + V4 default is ``"sqrtsoftplus"``. + enable_expert_bias: if True, allocate a learnable per-expert bias + used for selection only ("noaux_tc"). Probabilities are + re-read from the un-biased score so probs gradient flows + only into ``weight``, not ``expert_bias``. + topk_scaling_factor: scalar multiplier applied to the + renormalized probs (V3-style ``moe_router_topk_scaling_factor``, + HF reference ``Gate.route_scale``). Defaults to ``1.0``. + dtype: dtype of the gate weight; defaults to fp32 (matches HF + reference; the routing math runs in fp32 regardless). + """ + + def __init__( + self, + *, + hidden_size: int, + num_experts: int, + topk: int, + score_function: str = "sqrtsoftplus", + enable_expert_bias: bool = False, + topk_scaling_factor: float = 1.0, + dtype: Optional[torch.dtype] = None, + ) -> None: + super().__init__() + if num_experts <= 0: + raise ValueError(f"num_experts must be > 0, got {num_experts}") + if topk <= 0 or topk > num_experts: + raise ValueError(f"topk must be in [1, {num_experts}], got {topk}") + if score_function not in _VALID_SCORE_FUNCTIONS: + raise ValueError( + f"Unknown score_function: {score_function!r}. " + f"Expected one of {sorted(_VALID_SCORE_FUNCTIONS)}." + ) + + self.hidden_size = int(hidden_size) + self.num_experts = int(num_experts) + self.topk = int(topk) + self.score_function = str(score_function) + self.topk_scaling_factor = float(topk_scaling_factor) + + weight_dtype = dtype or torch.float32 + # Gate weight: [num_experts, hidden_size] (matches Megatron TopKRouter + # AND the HF reference Gate.weight). State-dict key: ``weight``. + self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, dtype=weight_dtype)) + nn.init.normal_(self.weight, mean=0.0, std=0.02) + + if enable_expert_bias: + # Per-expert selection bias. State-dict key: ``expert_bias``. + self.expert_bias = nn.Parameter(torch.zeros(self.num_experts, dtype=weight_dtype)) + else: + self.register_parameter("expert_bias", None) + + # ------------------------------------------------------------------ + + def forward(self, hidden: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Route ``hidden`` to top-K experts. + + Args: + hidden: ``[B, S, D]`` (or any shape with last dim ``D``). + + Returns: + probs: ``[N, num_experts]`` float tensor, ``N = numel/D``. + Non-selected experts have probability 0; selected experts + hold the (possibly renormalized + scaled) un-biased + score. + routing_map: ``[N, num_experts]`` bool tensor; ``True`` at + ``(n, e)`` iff token ``n`` is routed to expert ``e``. + """ + flat = hidden.reshape(-1, self.hidden_size) + # Match HF reference: routing math runs in fp32 regardless of + # input dtype. + logits = F.linear(flat.to(torch.float32), self.weight.to(torch.float32)) + return _compute_route( + logits=logits, + expert_bias=self.expert_bias, + score_function=self.score_function, + topk=self.topk, + topk_scaling_factor=self.topk_scaling_factor, + ) + + +# Back-compat alias. New callers should use ``DeepseekV4LearnedRouter``. +V4TopKRouter = DeepseekV4LearnedRouter + +# Back-compat alias for the standalone score-function helper. The leading +# underscore was dropped because the helper is part of the test surface. +_score_fn = v4_score_fn + +__all__ = [ + "DeepseekV4LearnedRouter", + "V4TopKRouter", + "v4_score_fn", + "_score_fn", +] diff --git a/primus/backends/megatron/core/transformer/sliding_window_kv.py b/primus/backends/megatron/core/transformer/sliding_window_kv.py new file mode 100644 index 000000000..68a336696 --- /dev/null +++ b/primus/backends/megatron/core/transformer/sliding_window_kv.py @@ -0,0 +1,81 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Sliding-window mask helpers for DeepSeek-V4 dense / SWA attention layers. + +Reference: techblog §1 ("Hybrid Attention") — dense layers +(``compress_ratio == 0``) attend only to the last ``attn_sliding_window`` +tokens (default 128) plus an optional ``attn_sink``. HCA layers also use +SWA over the raw KV in addition to the compressed-KV pool. + +This module only generates the **mask** — the actual ``q @ k^T`` happens in +the surrounding attention class. Returning a plain ``[Sq, Sk]`` additive +mask keeps it composable with both eager and flash-style backends. +""" + +from __future__ import annotations + +import torch + + +def sliding_window_causal_mask( + seq_len: int, + window: int, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Return a ``[seq_len, seq_len]`` additive attention mask. + + Position ``j`` is **allowed** for query ``i`` iff: + + * ``j <= i`` (causal), and + * ``i - j < window`` (sliding window). + + Disallowed positions get ``-inf``; allowed positions get ``0``. + + A ``window`` of ``0`` or ``>= seq_len`` degenerates to the standard + causal mask. + """ + q = torch.arange(seq_len, device=device).unsqueeze(1) + k = torch.arange(seq_len, device=device).unsqueeze(0) + dist = q - k + if window <= 0 or window >= seq_len: + allowed = dist >= 0 + else: + allowed = (dist >= 0) & (dist < window) + return torch.where(allowed, 0.0, float("-inf")).to(dtype) + + +def sliding_window_kv_indices( + seq_len: int, + window: int, + *, + device: torch.device, +) -> torch.Tensor: + """For each query ``i`` return the ``window`` raw-KV indices it attends + to: ``[max(0, i-window+1), i]``. + + Returned tensor has shape ``[seq_len, window]`` (long); positions before + the start of the sequence are filled with ``-1`` so the caller can drop + or zero-mask them. + """ + if window <= 0: + # No sliding-window restriction: nothing to gather (caller should use full causal). + return torch.empty(seq_len, 0, dtype=torch.long, device=device) + + i = torch.arange(seq_len, device=device).unsqueeze(1) # [S, 1] + offset = torch.arange(window - 1, -1, -1, device=device).unsqueeze(0) # [1, W] (descending) + j = i - offset # [S, W] + j = torch.where(j >= 0, j, torch.full_like(j, -1)) + return j + + +__all__ = [ + "sliding_window_causal_mask", + "sliding_window_kv_indices", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/README.md b/primus/backends/megatron/core/transformer/v4_attention_kernels/README.md new file mode 100644 index 000000000..68b6ef159 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/README.md @@ -0,0 +1,102 @@ +# DeepSeek-V4 attention kernels + +This package holds every attention backend used by `DeepseekV4Attention` +(`primus/backends/megatron/core/transformer/deepseek_v4_attention.py`). +`__init__.py` is the single entry point: it maps each backend to its functional +entry, and the attention module imports everything from here. + +## Attention variants + +DeepSeek-V4 has three attention shapes, selected per layer by `compress_ratio`: + +| `compress_ratio` | Variant | What it does | +| ---------------- | --------- | ------------------------------------------------------------------- | +| `0` | dense/SWA | Causal + sliding-window attention over the local KV. | +| `128` | HCA | Hierarchical compressed attention: local KV concatenated with a compressed pool, joint softmax. | +| `4` | CSA | Compressed sparse attention: per-query top-K gather from the compressed pool, joint softmax. | + +Two config selectors pick the kernel per group: + +- **`use_v4_attention_backend`** → dense (`cr=0`) and HCA (`cr=128`) layers. +- **`use_v4_csa_attention_backend`** → CSA (`cr=4`) layers. + +Both default to `triton_v1`. `gluon` is a gfx950/CDNA4-only opt-in: it is imported +lazily and only when selected, and selecting it asserts the device is gfx950. + +## Backends + +| Selector value | Applies to | Folder | Entry point(s) | Notes | +| -------------- | -------------- | ------------------------ | ------------------------------------------- | --------------------------------------------------------------------- | +| `eager` | dense/HCA, CSA | `_eager/` | `eager_v4_attention`, `eager_v4_csa_attention` | Pure-PyTorch reference. Bit-identical baseline shared with unit tests; slow, used for correctness. | +| `triton_v0` | CSA only | `_triton_v0_deprecated/` | `v4_csa_attention_v0` | **Deprecated.** Gathered per-query CSA with scalar GEMV (~30–260× slower). Not for production. | +| `triton_v1` | dense/HCA, CSA | `_triton_v1/` | `v4_attention_v1`, `v4_csa_attention_v1` | **Production default.** Separate K/V, pool-based CSA + dense/HCA Triton kernels. | +| `triton_v2` | dense/HCA, CSA | `_triton_v2/` | `v4_attention_v2`, `v4_csa_attention_v2` | Fused single-latent sparse-MLA (K=V) using plain Triton `tl.dot` / MFMA. | +| `gluon` | dense/HCA, CSA | `_gluon_dsa/` | `v4_attention_gluon`, `v4_csa_attention_gluon` | Hand-tuned fused single-latent sparse-MLA for **gfx950 (CDNA4) only**. Lazily imported; selecting it asserts the arch. | +| `flydsl_v0` | CSA only | `_flydsl_v0_deprecated/` | routed via `v4_csa_attention_v0` (`use_flydsl=True`) | **Deprecated**, forward-only legacy FlyDSL scalar CSA. | + +Support folders (not directly selectable): + +- `_triton_common/` — shared Triton helpers (indexer, compressor, sinkhorn, RoPE, HC). +- `_flydsl_v1/` — WIP native FlyDSL MFMA sparse-MLA backend (forward kernel not yet implemented). +- `_tilelang/` — experimental TileLang path (not wired into the selectors). +- `v4_sparse_mla_adapter.py` — kernel-agnostic adapter mapping V4 tensors to the fused sparse-MLA interface (used by `triton_v2` / `gluon`). + +Valid selector values (enforced in `DeepseekV4Attention.__init__`): + +- dense/HCA: `eager | triton_v1 | triton_v2 | gluon` +- CSA: `eager | triton_v0 | triton_v1 | triton_v2 | gluon | flydsl_v0` + +> Note: when a Turbo `core_attention` module is built, `use_turbo_attention` +> still takes precedence for the dense (`cr=0`) path. + +## How to enable each backend + +Set the two selectors to any valid value. All three mechanisms below set the +same config fields. + +### 1. Config YAML + +```yaml +# primus/configs/models/megatron/deepseek_v4_*.yaml +use_v4_attention_backend: triton_v1 # dense (cr=0) + HCA (cr=128) +use_v4_csa_attention_backend: triton_v1 # CSA (cr=4) +``` + +### 2. CLI flags + +```bash +--use_v4_attention_backend triton_v2 \ +--use_v4_csa_attention_backend gluon +``` + +### 3. Environment variables (root run scripts) + +The `run_deepseek_v4*.sh` scripts read these and forward them as CLI flags: + +```bash +export USE_V4_ATTENTION_BACKEND=triton_v1 # dense + HCA +export USE_V4_CSA_ATTENTION_BACKEND=triton_v1 # CSA +``` + +### Examples + +```bash +# Default: production Triton v1 (separate K/V) everywhere +export USE_V4_ATTENTION_BACKEND=triton_v1 +export USE_V4_CSA_ATTENTION_BACKEND=triton_v1 + +# gfx950-only hand-tuned gluon backend (asserts arch when selected) +export USE_V4_ATTENTION_BACKEND=gluon +export USE_V4_CSA_ATTENTION_BACKEND=gluon + +# Fused single-latent sparse-MLA (Triton v2) for all groups +export USE_V4_ATTENTION_BACKEND=triton_v2 +export USE_V4_CSA_ATTENTION_BACKEND=triton_v2 + +# Eager reference (correctness / debugging) +export USE_V4_ATTENTION_BACKEND=eager +export USE_V4_CSA_ATTENTION_BACKEND=eager +``` + +The two selectors are independent, so you can mix backends per group, e.g. +`USE_V4_ATTENTION_BACKEND=triton_v1` with `USE_V4_CSA_ATTENTION_BACKEND=gluon`. diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/__init__.py new file mode 100644 index 000000000..018e57a48 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/__init__.py @@ -0,0 +1,219 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 attention kernels — single entry point for every backend. + +``DeepseekV4Attention`` imports all attention entries from here, so this module +is the one place that maps a backend to its functional entry. Naming: + +* dense (cr=0) / HCA (cr=128) entry: ``v4_attention_`` +* CSA (cr=4) entry: ``v4_csa_attention_`` + +Backends: + +* ``eager`` — pure-Python reference (:mod:`_eager`): ``eager_v4_attention`` / + ``eager_v4_csa_attention``. +* ``v0`` — Triton, DEPRECATED gathered CSA (:mod:`_triton_v0_deprecated`): + ``v4_csa_attention_v0`` (cr=4 only, ~30-260x slower; not used in production). +* ``v1`` — Triton production, separate K/V (:mod:`_triton_v1`): + ``v4_attention_v1`` (dense/HCA) / ``v4_csa_attention_v1`` (pool CSA). +* ``v2`` — Triton fused single-latent sparse-MLA (:mod:`_triton_v2`, + ``tl.dot`` / MFMA): ``v4_attention_v2`` / ``v4_csa_attention_v2``. +* ``gluon`` — hand-tuned gfx950 fused single-latent sparse-MLA (:mod:`_gluon_dsa`): + loaded LAZILY via :func:`load_gluon_attention_backends` (NOT imported eagerly). + +``gluon`` hard-depends on ``triton.experimental.gluon`` (gfx950 / CDNA4 only), so +importing it here unconditionally would make *any* ``import ...v4_attention_kernels`` +fail on a Triton build without gluon — even when the caller selected +``eager`` / ``triton_v1`` / ``triton_v2``. It is therefore imported on demand only +when a layer actually selects the ``gluon`` backend (see +:func:`load_gluon_attention_backends`). + +The eager references share exactly one definition with the kernels + unit tests +and keep the checkpoint-reproduction baseline bit-identical at the call sites. +""" + +from primus.backends.megatron.core.transformer.v4_attention_kernels._eager import ( + eager_v4_attention, + eager_v4_csa_attention, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v0_deprecated import ( + V4CSAAttentionFn, + v4_csa_attention_v0, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1 import ( + V4AttentionFn, + V4CSAPoolAttentionFn, + v4_attention_v1, + v4_csa_attention_v1, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_triton import ( + v4_attention_v2, + v4_csa_attention_v2, +) + + +def load_gluon_attention_backends(): + """Lazily import the gluon sparse-MLA attention entries. + + The gluon backend (:mod:`_gluon_dsa`) hard-depends on + ``triton.experimental.gluon`` (gfx950 / CDNA4 only). This helper defers that + import so selecting any other backend (``eager`` / ``triton_v1`` / + ``triton_v2``) never pays it — and never crashes on a Triton build / GPU arch + without gluon support. Call it only when a layer actually selects ``gluon``. + + Returns ``(v4_attention_gluon, v4_csa_attention_gluon)``. Raises + :class:`ImportError` with an actionable message when the gluon dependency is + unavailable. + + NOTE: the import is intentionally inline (optional, hardware-specific + dependency); it must not be hoisted to module scope. + """ + try: + from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon import ( + v4_attention_gluon, + v4_csa_attention_gluon, + ) + except ImportError as exc: + raise ImportError( + "use_v4_attention_backend / use_v4_csa_attention_backend = 'gluon' requires " + "the gluon sparse-MLA backend (triton.experimental.gluon, gfx950 / CDNA4 only), " + f"which failed to import: {exc}. Select a different backend " + "(eager | triton_v1 | triton_v2), or run on a gfx950 build with Triton gluon support." + ) from exc + return v4_attention_gluon, v4_csa_attention_gluon + + +def load_gluon_v2_attention_backends(): + """Lazily import the gluon_v2 sparse-MLA attention entries (:mod:`_gluon_v2`). + + Second-generation Gluon backend (gfx950 / CDNA4): Gluon forward (rope-skip + exp2 + + MFMA K=32) + Gluon backward (rope-skip + K=32 + single-chunk RMW). Same lazy-import + rationale as :func:`load_gluon_attention_backends` (hard ``triton.experimental.gluon`` + dependency). Returns ``(v4_attention_gluon_v2, v4_csa_attention_gluon_v2)``. + """ + try: + from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon_v2 import ( + v4_attention_gluon_v2, + v4_csa_attention_gluon_v2, + ) + except ImportError as exc: + raise ImportError( + "use_v4_attention_backend / use_v4_csa_attention_backend = 'gluon_v2' requires " + "the gluon_v2 sparse-MLA backend (triton.experimental.gluon, gfx950 / CDNA4 only), " + f"which failed to import: {exc}. Select a different backend " + "(eager | triton_v1 | triton_v2 | gluon), or run on a gfx950 build with Triton gluon support." + ) from exc + return v4_attention_gluon_v2, v4_csa_attention_gluon_v2 + + +def load_gluon_v3_attention_backends(): + """Lazily import the gluon_v3 sparse-MLA attention entries (:mod:`_gluon_v3`). + + Optimized 3rd-gen Gluon backend (gfx950 / CDNA4): Round-9 CSA formula-pack + + aiter Gluon LSE fwd route, gluon_v2/Round-2 bwd chunking. Same lazy-import + rationale as :func:`load_gluon_attention_backends`. Returns + ``(v4_attention_gluon_v3, v4_csa_attention_gluon_v3)``. + """ + try: + from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon_v3 import ( + v4_attention_gluon_v3, + v4_csa_attention_gluon_v3, + ) + except ImportError as exc: + raise ImportError( + "use_v4_attention_backend / use_v4_csa_attention_backend = 'gluon_v3' requires " + "the gluon_v3 sparse-MLA backend (triton.experimental.gluon, gfx950 / CDNA4 only), " + f"which failed to import: {exc}. Select a different backend " + "(eager | triton_v1 | triton_v2 | gluon | gluon_v2 | flydsl_v1 | turbo), " + "or run on a gfx950 build with Triton gluon support." + ) from exc + return v4_attention_gluon_v3, v4_csa_attention_gluon_v3 + + +def load_flydsl_attention_backends(): + """Lazily import the native-FlyDSL sparse-MLA attention entries. + + The flydsl_v1 backend (:mod:`_flydsl_v1`) hard-depends on the installed + ``flydsl`` pip package (gfx950 / CDNA4). This helper defers that import so + selecting any other backend never pays it — and never crashes on a build / + GPU arch without flydsl. Call it only when a layer actually selects + ``flydsl_v1``. + + Returns ``(v4_attention_flydsl, v4_csa_attention_flydsl)``. Raises + :class:`ImportError` with an actionable message when flydsl is unavailable. + """ + try: + from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_flydsl import ( + v4_attention_flydsl, + v4_csa_attention_flydsl, + ) + except ImportError as exc: + raise ImportError( + "use_v4_attention_backend / use_v4_csa_attention_backend = 'flydsl_v1' requires " + "the native FlyDSL sparse-MLA backend (the `flydsl` pip package, gfx950 / CDNA4), " + f"which failed to import: {exc}. Select a different backend " + "(eager | triton_v1 | triton_v2 | gluon), or install flydsl on a gfx950 build." + ) from exc + return v4_attention_flydsl, v4_csa_attention_flydsl + + +def load_turbo_attention_backends(): + """Lazily import the Primus-Turbo native-FlyDSL sparse-MLA attention entries. + + The ``turbo`` backend (:mod:`_turbo_flydsl`) binds to the installed + ``primus_turbo`` flydsl sparse-MLA v2 kernels (the "turbo API" integration), + which hard-depend on the installed ``primus_turbo`` (with the flydsl attention + submodule) and the ``flydsl`` pip package (gfx950 / CDNA4). Deferred so + selecting any other backend never pays that import — and never crashes on a + build / GPU arch without it. Call it only when a layer selects ``turbo``. + + Returns ``(v4_attention_turbo, v4_csa_attention_turbo)``. Raises + :class:`ImportError` with an actionable message when the dependency is missing. + """ + try: + from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_turbo_flydsl import ( + v4_attention_turbo, + v4_csa_attention_turbo, + ) + except ImportError as exc: + raise ImportError( + "use_v4_attention_backend / use_v4_csa_attention_backend = 'turbo' requires the " + "Primus-Turbo native-FlyDSL sparse-MLA backend (the installed `primus_turbo` with its " + "flydsl sparse-MLA attention, plus the `flydsl` pip package, gfx950 / CDNA4), " + f"which failed to import: {exc}. Select a different backend " + "(eager | triton_v1 | triton_v2 | gluon | gluon_v2 | flydsl_v1), or install a " + "primus_turbo build carrying primus_turbo.flydsl.attention on a gfx950 build." + ) from exc + return v4_attention_turbo, v4_csa_attention_turbo + + +__all__ = [ + # eager reference + "eager_v4_attention", + "eager_v4_csa_attention", + # triton v0 (deprecated gathered CSA) + "v4_csa_attention_v0", + "V4CSAAttentionFn", + # triton v1 (production, separate K/V) + "v4_attention_v1", + "v4_csa_attention_v1", + "V4AttentionFn", + "V4CSAPoolAttentionFn", + # triton v2 (fused single-latent sparse-MLA) + "v4_attention_v2", + "v4_csa_attention_v2", + # gluon (fused single-latent sparse-MLA, gfx950) — lazily loaded + "load_gluon_attention_backends", + # gluon_v2 (2nd-gen gluon fwd+bwd, gfx950) — lazily loaded + "load_gluon_v2_attention_backends", + # gluon_v3 (3rd-gen optimized gluon fwd+bwd, gfx950) — lazily loaded + "load_gluon_v3_attention_backends", + # flydsl_v1 (native FlyDSL fused single-latent sparse-MLA, gfx950) — lazily loaded + "load_flydsl_attention_backends", + # turbo (Primus-Turbo native-FlyDSL sparse-MLA via the turbo API, gfx950) — lazily loaded + "load_turbo_attention_backends", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/__init__.py new file mode 100644 index 000000000..e72167d5d --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/__init__.py @@ -0,0 +1,18 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Eager-Python reference ops for DeepSeek-V4 attention. + +The single source of "eager truth" shared by ``DeepseekV4Attention``, every +kernel backend (triton v0/v1/v2, gluon, flydsl_v2, tilelang) and the unit tests: + +* :func:`eager_v4_attention` — dense (cr=0) / HCA (cr=128) +* :func:`eager_v4_csa_attention` — CSA (cr=4) +""" + +from .reference import eager_v4_attention, eager_v4_csa_attention + +__all__ = ["eager_v4_attention", "eager_v4_csa_attention"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/reference.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/reference.py new file mode 100644 index 000000000..6fa6598d4 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_eager/reference.py @@ -0,0 +1,329 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Eager-Python references for the V4 attention kernels (plan-4 P24). + +The math here is **bit-identical** to what previously lived inline in +:meth:`DeepseekV4Attention._attention_forward` and +:meth:`DeepseekV4Attention._csa_forward`. Plan-4 extracts it into pure +functions so that: + +* ``DeepseekV4Attention`` itself, the plan-4 Triton kernels (P25 / P26), + and the plan-4 unit-test harness share exactly one definition; +* the function signatures match the kernel signatures + (``v4_attention_v1`` / ``v4_csa_attention_v0``) so the test harness can + plug reference ↔ candidate interchangeably; +* the compute kernel of V4 attention is decoupled from the + ``DeepseekV4Attention`` class (no ``self``-bound state) — the + per-call inputs are explicit. + +The two functions: + +* :func:`eager_v4_attention` — single-key-axis attention with optional + per-head learned softmax sink, optional sliding window, and optional + ``[Sq, Sk]`` additive bias. Covers ``compress_ratio == 0`` (dense + + SWA + sink, no bias) and ``compress_ratio == 128`` (HCA — caller + pre-concatenates the compressed pool to the local keys and supplies + the joint-softmax additive bias). +* :func:`eager_v4_csa_attention` — fused local-SWA + per-query top-K + sparse attention with shared per-head sink and joint softmax. + Covers ``compress_ratio == 4`` (CSA). The caller is responsible for + the per-query top-K gather; the function takes the gathered + ``[B, Sq, K, head_dim]`` tensor directly. + +Both functions: + +* keep every matmul / einsum on tensor cores in the input dtype (bf16 + in production); the matmul accumulator inside is fp32. The *only* + fp32 step is the softmax block (max-subtract + ``exp`` + sum + + divide), which :func:`_softmax_with_sink` upcasts internally and + returns in fp32; the caller-side ``probs.to(v.dtype)`` puts probs + back on the bf16 path before the V-matmul. This matches the + FlashAttention / Megatron de-facto layout (bf16 ``Q @ K^T``, fp32 + softmax, bf16 ``probs @ V``); +* honor ``attn_dropout`` only when ``training`` is also ``True`` (so + the eval / inference path is deterministic); +* return ``[B, H, Sq, head_dim]`` in ``v.dtype`` (or ``v_local.dtype`` + for the CSA path); +* are autograd-friendly: fwd produces a graph; ``out.sum().backward()`` + populates ``q.grad / k.grad / v.grad / sink.grad`` (and + ``gathered.grad`` for CSA). +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from primus.backends.megatron.core.transformer.sliding_window_kv import ( + sliding_window_causal_mask, +) + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _softmax_with_sink( + logits: torch.Tensor, + sink: Optional[torch.Tensor], +) -> torch.Tensor: + """Numerically stable softmax with optional per-head learned sink column. + + ``logits`` shape: ``[B, H, ..., Sk]`` — the head axis is at ``dim=1``. + When ``sink`` is given (shape ``[H]``) it is joined as a virtual key + column with notional value zero (so the V-weighted sum after + multiplying by ``v`` is unaffected by the sink slot), then dropped + after softmax. The head can still spend probability mass on the + sink as a "no attention" fallback. + + **dtype contract.** This is the *only* fp32 step in V4 attention. + ``logits`` may arrive in the model dtype (bf16 in production — + coming straight out of a tensor-core ``Q @ K^T`` matmul); the + softmax block (max-subtract + ``exp`` + sum + divide) is the + numerically sensitive part, so we **upcast to fp32 here** and + return probabilities in **fp32**. The caller is responsible for + ``probs.to(v.dtype)`` before the value matmul so the V-matmul + stays on bf16 tensor cores. + + Returns probabilities on the *real* keys of the same shape as + ``logits`` — but in **fp32** regardless of input dtype. + """ + # logits: [B, H, ..., Sk] (any dtype) -> logits_fp32: [B, H, ..., Sk] in fp32 + logits_fp32 = logits.float() + if sink is None: + logits_fp32 = logits_fp32 - logits_fp32.amax(dim=-1, keepdim=True).detach() + return logits_fp32.softmax(dim=-1) + + # sink: [H] (any dtype) -> sink_col: [B, H, ..., 1] in fp32 + ndim = logits_fp32.dim() + num_heads = sink.shape[0] + view_shape = [1] * ndim + view_shape[1] = num_heads + view_shape[-1] = 1 + target_shape = list(logits_fp32.shape[:-1]) + [1] + sink_col = sink.float().view(*view_shape).expand(*target_shape) + # logits_fp32: [B, H, ..., Sk], sink_col: [B, H, ..., 1] + # -> logits_aug: [B, H, ..., Sk+1] in fp32 + logits_aug = torch.cat([logits_fp32, sink_col], dim=-1) + logits_aug = logits_aug - logits_aug.amax(dim=-1, keepdim=True).detach() + probs = logits_aug.softmax(dim=-1) + return probs[..., :-1] + + +def _build_local_attention_mask( + seq_len: int, + swa_window: int, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Build the ``[seq_len, seq_len]`` local SWA-causal additive mask. + + Equivalent to :meth:`DeepseekV4Attention._local_mask`: + + * ``swa_window > 0`` — sliding-window causal (queries see the last + ``swa_window`` keys); + * otherwise — full causal (queries see all earlier keys). + + The two-call contract is deterministic: same ``(seq_len, + swa_window, device, dtype)`` always returns the same tensor, so + rebuilding inside this function vs. passing a pre-built mask gives + bit-identical attention output. + """ + window = swa_window if swa_window > 0 else seq_len + return sliding_window_causal_mask(seq_len, window, device=device, dtype=dtype) + + +# --------------------------------------------------------------------------- +# Public reference ops +# --------------------------------------------------------------------------- + + +def eager_v4_attention( + q: torch.Tensor, # [B, H, Sq, D] + k: torch.Tensor, # [B, H, Sk, D] + v: torch.Tensor, # [B, H, Sk, D] + *, + sink: Optional[torch.Tensor], # [H] or None + swa_window: int, + additive_mask: Optional[torch.Tensor], # [Sq, Sk] or None + attn_dropout: float, + training: bool, + scale: float, +) -> torch.Tensor: + """Eager-Python V4 dense / HCA attention. + + Math:: + + logits = (q @ k.T) * scale + mask # bf16 tensor-core matmul + probs = softmax_with_sink(logits, sink) # softmax internally upcasts to fp32 + if attn_dropout > 0 and training: probs = dropout(probs, attn_dropout) + out = probs.to(v.dtype) @ v # bf16 tensor-core matmul + + **dtype contract.** All matmuls (``Q @ K^T`` and ``probs @ V``) run + on tensor cores in the input dtype (bf16 in production); the + accumulator inside the matmul is fp32. The *only* fp32 step is the + softmax block, which :func:`_softmax_with_sink` handles internally + (input may be bf16, output is fp32). The caller-side + ``probs.to(v.dtype)`` puts probs back on the bf16 path before the + V-matmul. + + Mask resolution: + + * ``additive_mask is not None`` — used directly. ``swa_window`` is + ignored. (HCA's ``compress_ratio == 128`` caller pre-builds the + joint-softmax mask via ``cat([local_mask, hca_mask])`` and + passes it here.) + * ``additive_mask is None`` and ``swa_window > 0`` — sliding-window + causal mask is built internally. Requires ``Sq == Sk``. + * ``additive_mask is None`` and ``swa_window <= 0`` — full causal + mask is built internally. Requires ``Sq == Sk``. + + Returns ``[B, H, Sq, D]`` in ``v.dtype``. + """ + # q: [B, H, Sq, D], k: [B, H, Sk, D] -> Sq: scalar, Sk: scalar + Sq = q.shape[2] + Sk = k.shape[2] + + # mask: [Sq, Sk] in q.dtype (broadcasts over B, H at the addition site below) + if additive_mask is None: + if Sq != Sk: + raise ValueError( + "eager_v4_attention requires `additive_mask` when Sq != Sk; " + f"got Sq={Sq}, Sk={Sk}. The HCA caller must pre-concatenate " + "the compressed pool to the local keys and supply the joint " + "additive mask." + ) + # _build_local_attention_mask(Sq, swa_window) -> mask: [Sq, Sq] (== [Sq, Sk]) in q.dtype + mask = _build_local_attention_mask(Sq, swa_window, device=q.device, dtype=q.dtype) + else: + # additive_mask: [Sq, Sk] -> mask: [Sq, Sk] in caller's dtype + mask = additive_mask + + # q: [B, H, Sq, D], k.transpose(-2,-1): [B, H, D, Sk] -> matmul(...): [B, H, Sq, Sk] in q.dtype + # bf16 tensor-core matmul (fp32 accumulator inside, output bf16) + # * scale: [B, H, Sq, Sk] in q.dtype + logits = torch.matmul(q, k.transpose(-2, -1)) * scale + # logits: [B, H, Sq, Sk], mask: [Sq, Sk] -> logits: [B, H, Sq, Sk] (mask broadcasts over B, H) + logits = logits + mask + # logits: [B, H, Sq, Sk] (bf16), sink: [H] or None -> probs: [B, H, Sq, Sk] in fp32 + # softmax block internally upcasts logits + sink to fp32 (numerical contract: only step in fp32) + probs = _softmax_with_sink(logits, sink) + # probs: [B, H, Sq, Sk] (fp32) -> probs: [B, H, Sq, Sk] (fp32) + if attn_dropout > 0.0 and training: + probs = torch.nn.functional.dropout(probs, p=attn_dropout) + # probs.to(v.dtype): [B, H, Sq, Sk] in v.dtype, v: [B, H, Sk, D] -> out: [B, H, Sq, D] in v.dtype + # bf16 tensor-core matmul (fp32 accumulator inside) + return torch.matmul(probs.to(v.dtype), v) + + +def eager_v4_csa_attention( + q: torch.Tensor, # [B, H, Sq, D] + k_local: torch.Tensor, # [B, H, Sq, D] + v_local: torch.Tensor, # [B, H, Sq, D] + gathered: torch.Tensor, # [B, Sq, K, D] — pre-gathered per-query top-K from compressed pool + *, + sink: Optional[torch.Tensor], # [H] or None + swa_window: int, + sparse_mask: torch.Tensor, # [B, Sq, K] additive (broadcasts over H) + attn_dropout: float, + training: bool, + scale: float, +) -> torch.Tensor: + """Eager-Python V4 CSA fused attention (joint local SWA + sparse top-K). + + Math:: + + local_logits = (q @ k_local.T) * scale + local_mask # bf16 matmul + sparse_logits = einsum("bhsd,bhskd->bhsk", q, gathered_h) * scale + + sparse_mask.unsqueeze(1) # bf16 einsum + joint_logits = cat([local_logits, sparse_logits], dim=-1) # [B, H, Sq, Sq+K] + probs = softmax_with_sink(joint_logits, sink) # JOINT softmax in fp32 + if attn_dropout > 0 and training: probs = dropout(probs, attn_dropout) + probs_local, probs_sparse = probs[..., :Sq], probs[..., Sq:] + out = probs_local @ v_local + einsum("bhsk,bhskd->bhsd", probs_sparse, gathered_h) + # bf16 matmul / einsum + + where ``gathered_h = gathered.unsqueeze(1).expand(B, H, Sq, K, D)`` + (heads broadcast across the single compressor output). + + **dtype contract.** Same as :func:`eager_v4_attention`: every + matmul / einsum runs on tensor cores in the input dtype (bf16 in + production, fp32 accumulator inside); the *only* fp32 step is the + joint softmax inside :func:`_softmax_with_sink`. ``probs.to(v.dtype)`` + puts the per-branch probabilities back on bf16 before the V-stage. + + The local SWA mask is built internally from ``swa_window`` (see + :func:`_build_local_attention_mask`); the caller pre-builds + ``sparse_mask`` to flag indexer-dropped slots (``-inf`` for + ``topk_idx == -1``). + + Returns ``[B, H, Sq, D]`` in ``v_local.dtype``. + """ + # q: [B, H, Sq, D], gathered: [B, Sq, K, D] -> B, H, Sq, D, K: scalars + B, H, Sq, D = q.shape + K = gathered.shape[2] + + # _build_local_attention_mask(Sq, swa_window) -> local_mask: [Sq, Sq] in q.dtype + local_mask = _build_local_attention_mask(Sq, swa_window, device=q.device, dtype=q.dtype) + + # q: [B, H, Sq, D], k_local.transpose(-2,-1): [B, H, D, Sq] + # -> matmul(...): [B, H, Sq, Sq] in q.dtype (bf16 tensor core, fp32 accumulator inside) + # * scale: local_logits: [B, H, Sq, Sq] in q.dtype + local_logits = torch.matmul(q, k_local.transpose(-2, -1)) * scale + # local_logits: [B, H, Sq, Sq], local_mask: [Sq, Sq] + # -> local_logits: [B, H, Sq, Sq] in q.dtype (mask broadcasts over B, H) + local_logits = local_logits + local_mask + + # gathered: [B, Sq, K, D] -> unsqueeze(1): [B, 1, Sq, K, D] + # -> expand: gathered_h: [B, H, Sq, K, D] in gathered.dtype (view, no copy) + gathered_h = gathered.unsqueeze(1).expand(B, H, Sq, K, D) + # q: [B, H, Sq, D], gathered_h: [B, H, Sq, K, D] + # -> einsum("bhsd,bhskd->bhsk"): [B, H, Sq, K] in q.dtype (bf16 tensor core) + # * scale: sparse_logits: [B, H, Sq, K] in q.dtype + sparse_logits = torch.einsum("bhsd,bhskd->bhsk", q, gathered_h) * scale + # sparse_logits: [B, H, Sq, K], sparse_mask.unsqueeze(1): [B, 1, Sq, K] + # -> sparse_logits: [B, H, Sq, K] in q.dtype (mask broadcasts over H) + sparse_logits = sparse_logits + sparse_mask.unsqueeze(1) + + # local_logits: [B, H, Sq, Sq], sparse_logits: [B, H, Sq, K] + # -> joint_logits: [B, H, Sq, Sq+K] in q.dtype + joint_logits = torch.cat([local_logits, sparse_logits], dim=-1) + # joint_logits: [B, H, Sq, Sq+K] (bf16), sink: [H] or None -> probs: [B, H, Sq, Sq+K] in fp32 + # softmax block internally upcasts to fp32 (numerical contract: only step in fp32) + probs = _softmax_with_sink(joint_logits, sink) + + # probs: [B, H, Sq, Sq+K] (fp32) -> probs: [B, H, Sq, Sq+K] (fp32) + if attn_dropout > 0.0 and training: + probs = torch.nn.functional.dropout(probs, p=attn_dropout) + + # probs[..., :Sq]: [B, H, Sq, Sq] in fp32 -> probs_local: [B, H, Sq, Sq] in v_local.dtype + probs_local = probs[..., :Sq].to(v_local.dtype) + # probs[..., Sq:]: [B, H, Sq, K] in fp32 -> probs_sparse: [B, H, Sq, K] in v_local.dtype + probs_sparse = probs[..., Sq:].to(v_local.dtype) + + # probs_local: [B, H, Sq, Sq], v_local: [B, H, Sq, D] + # -> matmul(...): out_local: [B, H, Sq, D] in v_local.dtype (bf16 tensor core) + out_local = torch.matmul(probs_local, v_local) + # probs_sparse: [B, H, Sq, K], gathered_h.to(v_local.dtype): [B, H, Sq, K, D] + # -> einsum("bhsk,bhskd->bhsd"): out_sparse: [B, H, Sq, D] in v_local.dtype (bf16 tensor core) + out_sparse = torch.einsum( + "bhsk,bhskd->bhsd", + probs_sparse, + gathered_h.to(v_local.dtype), + ) + + # out_local: [B, H, Sq, D], out_sparse: [B, H, Sq, D] -> out: [B, H, Sq, D] in v_local.dtype + return out_local + out_sparse + + +__all__ = [ + "eager_v4_attention", + "eager_v4_csa_attention", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/__init__.py new file mode 100644 index 000000000..5d2445b7c --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/__init__.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FlyDSL attention-kernel backend for DeepSeek-V4 (gfx950 / MI355X). + +A *soft-dependency* alternate backend, structured exactly like the sibling +``_tilelang`` package: the V4 attention layer calls :func:`should_dispatch` +with the ``enabled`` flag from a config knob, and this module lazily imports +the FlyDSL kernel wrappers. If the ``flydsl`` runtime (or a kernel submodule) +is not importable, :func:`should_dispatch` returns ``False`` and the caller +transparently falls back to the in-tree Triton path -- so a container without +FlyDSL never breaks and never changes behaviour. + +Kernels live under ``_flydsl/kernels/``. Forward-only for now (inference/eval); +training autograd is a follow-up. +""" +from __future__ import annotations + +import os +import warnings +from typing import Any, Optional, Set + +__all__ = [ + "should_dispatch", + "is_flydsl_available", + "v4_attention_fwd_flydsl", + "v4_csa_attention_fwd_flydsl", +] + +# Kernel names the layer may ask for, mirroring the _tilelang registry. +_KNOWN_KERNEL_NAMES: Set[str] = { + "v4_attention_fwd", # SWA / HCA forward (MQA launcher) + "v4_csa_attention_fwd", # CSA forward + "v4_attention_bwd", # present in-package; training-autograd hook TBD + "v4_csa_attention_bwd", # present in-package; training-autograd hook TBD +} + +# Forward kernels that are actually wired (have a working adapter below). +_WIRED_FWD: Set[str] = {"v4_attention_fwd", "v4_csa_attention_fwd"} + +_PROBE_DONE: bool = False +_FLYDSL_AVAILABLE: bool = False + + +def _probe_flydsl() -> bool: + """Import the ``flydsl`` runtime once (cached); warn once on rank 0 if absent.""" + global _PROBE_DONE, _FLYDSL_AVAILABLE + if _PROBE_DONE: + return _FLYDSL_AVAILABLE + _PROBE_DONE = True + try: + # The kernel wrappers add the FlyDSL build dir to sys.path at import; + # allow an override for non-default install locations. + src = os.environ.get("PRIMUS_V4_FLYDSL_SRC", "/workspace/FlyDSL-amd") + import sys + + if src and src not in sys.path and os.path.isdir(src): + sys.path.insert(0, src) + import flydsl # noqa: F401 + + _FLYDSL_AVAILABLE = True + except Exception as exc: # ImportError, or a runtime/arch probe failure + if int(os.environ.get("RANK", "0")) == 0: + warnings.warn( + f"[v4-flydsl] FlyDSL runtime unavailable ({exc!r}); FlyDSL " + f"attention backend disabled, falling back to Triton.", + RuntimeWarning, + stacklevel=3, + ) + _FLYDSL_AVAILABLE = False + return _FLYDSL_AVAILABLE + + +def is_flydsl_available() -> bool: + """True iff the FlyDSL runtime imported successfully (cached).""" + return _probe_flydsl() + + +# --- lazy adapter cache ----------------------------------------------------- +_FWD_MQA = None # _launch_v4_attention_fwd_flydsl_mqa +_FWD_CSA = None # _launch_v4_attention_fwd_csa + + +def _load_fwd_mqa(): + global _FWD_MQA + if _FWD_MQA is None: + from .kernels.v4_attention_fwd_flydsl_mqa import ( + _launch_v4_attention_fwd_flydsl_mqa as fn, + ) + + _FWD_MQA = fn + return _FWD_MQA + + +def _load_fwd_csa(): + global _FWD_CSA + if _FWD_CSA is None: + from .kernels.v4_attention_fwd_flydsl_csa import ( + _launch_v4_attention_fwd_csa as fn, + ) + + _FWD_CSA = fn + return _FWD_CSA + + +def should_dispatch(kernel_name: str, *, enabled: bool) -> bool: + """True iff FlyDSL should handle ``kernel_name`` (else Triton fallback). + + Short-circuits when ``enabled`` is False so the off path never imports FlyDSL. + """ + if not enabled: + return False + if kernel_name not in _KNOWN_KERNEL_NAMES: + raise ValueError( + f"Unknown FlyDSL kernel {kernel_name!r}; expected one of " f"{sorted(_KNOWN_KERNEL_NAMES)}" + ) + if kernel_name not in _WIRED_FWD: + # bwd kernels are present in-package but lack a training-autograd hook + return False + if not _probe_flydsl(): + return False + try: + _load_fwd_csa() if kernel_name == "v4_csa_attention_fwd" else _load_fwd_mqa() + except Exception as exc: + if int(os.environ.get("RANK", "0")) == 0: + warnings.warn( + f"[v4-flydsl] kernel {kernel_name!r} import failed ({exc!r}); " f"falling back to Triton.", + RuntimeWarning, + stacklevel=3, + ) + return False + return True + + +# --- forward adapters (signatures match the dispatch sites) ----------------- +def v4_attention_fwd_flydsl( + q, + k, + v, + *, + sink: Optional[Any] = None, + swa_window: int = 0, + additive_mask: Optional[Any] = None, + scale: float, + hca_local_seqlen: int = 0, +): + """SWA/HCA forward via FlyDSL. Returns the attention output tensor.""" + out, _lse = _load_fwd_mqa()( + q, k, v, sink, int(swa_window), additive_mask, float(scale), int(hca_local_seqlen) + ) + return out + + +def v4_csa_attention_fwd_flydsl( + q, + k_local, + v_local, + gathered, + *, + sparse_mask, + sink: Optional[Any] = None, + swa_window: int = 0, + scale: float, + attn_dropout: float = 0.0, + training: bool = False, +): + """CSA forward via FlyDSL (the 2.79x kernel). Returns the output tensor.""" + out, _lse = _load_fwd_csa()( + q, k_local, v_local, gathered, sink, int(swa_window), sparse_mask, float(scale) + ) + return out diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_bwd_flydsl_mqa.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_bwd_flydsl_mqa.py new file mode 100644 index 000000000..f9ffcb6f3 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_bwd_flydsl_mqa.py @@ -0,0 +1,1270 @@ +"""V4 SWA attention backward FlyDSL launcher (Phase B STEP 1b, cr=0, SWA-only). + +API contract (matches the Triton ``_launch_v4_attention_bwd`` signature; the +bwd_modes harness calls this function): + + flydsl_v4_attention_bwd( + q, k, v, out, dout, lse, + *, + sink, swa_window, additive_mask, scale, hca_local_seqlen, + ) -> (dq, dk, dv, dsink) + +STEP 1b SCOPE +------------- +Current state of each compute stage: + * preprocess (D scalar) -> FlyDSL kernel (``v4_swa_bwd_preprocess_kernel``) + * dq -> FlyDSL kernel (``v4_swa_bwd_dq_kernel``) + forked from kernels/sla_bwd_dq.py + * dk / dv -> Triton kernel (``_v4_attention_bwd_dkv_kernel``) + * dsink -> FlyDSL kernel writes it (via atomic_fadd in dq) + +Env knobs: + V4_FLYDSL_BWD_FLY_PREPROCESS default 1 (1=FlyDSL, 0=Triton) + V4_FLYDSL_BWD_FLY_DQ default 1 (1=FlyDSL, 0=Triton) + V4_FLYDSL_BWD_VERBOSE default 0 (1=print provenance line) +""" + +from __future__ import annotations + +import os +import sys +import threading +from typing import Optional, Tuple + +import torch + +_FLYDSL_SRC = "/workspace/FlyDSL-amd" +if _FLYDSL_SRC not in sys.path: + sys.path.insert(0, _FLYDSL_SRC) + +os.environ.setdefault("FLYDSL_WAVES_PER_EU", "2") + +_HERE = os.path.dirname(os.path.abspath(__file__)) +if _HERE not in sys.path: + sys.path.insert(0, _HERE) + +if "/workspace/Primus" not in sys.path: + sys.path.insert(0, "/workspace/Primus") +import triton # noqa: E402 + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import ( # noqa: E402 + _v4_attention_bwd_dkv_kernel, + _v4_attention_bwd_dkv_pool_kernel, + _v4_attention_bwd_dq_kernel, + _v4_attention_bwd_preprocess_kernel, +) + +# Built lazily on first call; module-level kernel build is expensive. +_PREPROCESS_KERNEL_CACHE = {} +_PREPROCESS_KERNEL_LOCK = threading.Lock() + +_DQ_KERNEL_CACHE = {} +_DQ_KERNEL_LOCK = threading.Lock() +_DKV_KERNEL_CACHE = {} +_DKV_KERNEL_LOCK = threading.Lock() +_DQ_POOL_KERNEL_CACHE = {} +_DQ_POOL_KERNEL_LOCK = threading.Lock() +_DKV_POOL_KERNEL_CACHE = {} +_DKV_POOL_KERNEL_LOCK = threading.Lock() + + +def _get_fly_preprocess(head_dim: int, dtype_str: str, block_rows: int): + key = (head_dim, dtype_str, block_rows) + with _PREPROCESS_KERNEL_LOCK: + if key in _PREPROCESS_KERNEL_CACHE: + return _PREPROCESS_KERNEL_CACHE[key] + from v4_sla_bwd_kernel import build_v4_swa_bwd_preprocess_module + + launch = build_v4_swa_bwd_preprocess_module( + head_dim=head_dim, + dtype_str=dtype_str, + block_rows=block_rows, + ) + _PREPROCESS_KERNEL_CACHE[key] = launch + return launch + + +def _get_fly_dq(num_heads: int, head_dim: int, swa_window: int, dtype_str: str, mqa_kv: bool, has_sink: bool): + key = (num_heads, head_dim, swa_window, dtype_str, mqa_kv, has_sink) + with _DQ_KERNEL_LOCK: + if key in _DQ_KERNEL_CACHE: + return _DQ_KERNEL_CACHE[key] + from v4_sla_bwd_kernel import build_v4_swa_bwd_dq_module + + launch = build_v4_swa_bwd_dq_module( + num_heads=num_heads, + head_dim=head_dim, + swa_window=swa_window, + dtype_str=dtype_str, + mqa_kv=mqa_kv, + has_sink=has_sink, + ) + _DQ_KERNEL_CACHE[key] = launch + return launch + + +def _get_fly_dq_pool( + num_heads: int, head_dim: int, pool_size: int, hca_local_seqlen: int, dtype_str: str, mqa_kv: bool +): + key = (num_heads, head_dim, pool_size, hca_local_seqlen, dtype_str, mqa_kv) + with _DQ_POOL_KERNEL_LOCK: + if key in _DQ_POOL_KERNEL_CACHE: + return _DQ_POOL_KERNEL_CACHE[key] + from v4_hca_bwd_dq_pool_kernel import build_v4_hca_bwd_dq_pool_module + + launch = build_v4_hca_bwd_dq_pool_module( + num_heads=num_heads, + head_dim=head_dim, + pool_size=pool_size, + hca_local_seqlen=hca_local_seqlen, + dtype_str=dtype_str, + mqa_kv=mqa_kv, + ) + _DQ_POOL_KERNEL_CACHE[key] = launch + return launch + + +def _get_fly_dkv_pool( + num_heads: int, head_dim: int, pool_size: int, hca_local_seqlen: int, dtype_str: str, mqa_kv: bool +): + key = (num_heads, head_dim, pool_size, hca_local_seqlen, dtype_str, mqa_kv) + with _DKV_POOL_KERNEL_LOCK: + if key in _DKV_POOL_KERNEL_CACHE: + return _DKV_POOL_KERNEL_CACHE[key] + from v4_hca_bwd_dkv_pool_kernel import build_v4_hca_bwd_dkv_pool_module + + launch = build_v4_hca_bwd_dkv_pool_module( + num_heads=num_heads, + head_dim=head_dim, + pool_size=pool_size, + hca_local_seqlen=hca_local_seqlen, + dtype_str=dtype_str, + mqa_kv=mqa_kv, + ) + _DKV_POOL_KERNEL_CACHE[key] = launch + return launch + + +def _get_fly_dkv(num_heads: int, head_dim: int, swa_window: int, dtype_str: str, mqa_kv: bool): + key = (num_heads, head_dim, swa_window, dtype_str, mqa_kv) + with _DKV_KERNEL_LOCK: + if key in _DKV_KERNEL_CACHE: + return _DKV_KERNEL_CACHE[key] + from v4_sla_bwd_dkv_kernel import build_v4_swa_bwd_dkv_module + + launch = build_v4_swa_bwd_dkv_module( + num_heads=num_heads, + head_dim=head_dim, + swa_window=swa_window, + dtype_str=dtype_str, + mqa_kv=mqa_kv, + ) + _DKV_KERNEL_CACHE[key] = launch + return launch + + +def _run_preprocess( + out: torch.Tensor, + dout: torch.Tensor, + *, + use_flydsl: bool, + block_m: int, + block_dmodel: int, +) -> torch.Tensor: + """Compute D[b,h,m] = sum_d (out[b,h,m,d] * dout[b,h,m,d]) in fp32.""" + B, HQ, Sq, D = out.shape + d_buf = torch.empty((B, HQ, Sq), device=out.device, dtype=torch.float32) + if use_flydsl: + out_f = out.contiguous().view(-1, D) + dout_f = dout.contiguous().view(-1, D) + delta_f = d_buf.view(-1) + n_rows = out_f.shape[0] + dtype_str = "bf16" if out.dtype == torch.bfloat16 else "f16" + max_block_threads = 256 + threads_per_row = D // 8 + for br in (8, 4, 2, 1): + if n_rows % br == 0 and br * threads_per_row <= max_block_threads: + block_rows = br + break + else: + block_rows = 1 + launch = _get_fly_preprocess(D, dtype_str, block_rows) + launch(out_f, dout_f, delta_f, n_rows) + return d_buf + pre_grid = (triton.cdiv(Sq, block_m), B * HQ) + _v4_attention_bwd_preprocess_kernel[pre_grid]( + out, + dout, + d_buf, + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + Sq, + HEAD=HQ, + BLOCK_M=block_m, + BLOCK_DMODEL=block_dmodel, + num_warps=4, + num_stages=1, + ) + return d_buf + + +def _run_dq_triton( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + dsink_fp32, + sink_arg, + mask_arg, + *, + scale, + swa_window_constexpr, + has_sink, + has_add_mask, + hca_local_seqlen, + use_causal, + block_m, + block_n, + block_dmodel, + stride_ms, + stride_mn, + Sq, + Sk, + HQ, + HK, + B, + exact_tiles_m, + exact_tiles_n, +): + dq_grid = (triton.cdiv(Sq, block_m), B * HQ) + _v4_attention_bwd_dq_kernel[dq_grid]( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + dsink_fp32, + sink_arg, + mask_arg, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_fp32.stride(0), + dq_fp32.stride(1), + dq_fp32.stride(2), + dq_fp32.stride(3), + stride_ms, + stride_mn, + Sq, + Sk, + float(scale), + HEAD_Q=HQ, + HEAD_K=HK, + SWA_WINDOW=swa_window_constexpr, + HAS_SINK=has_sink, + HAS_ADD_MASK=has_add_mask, + HCA_LOCAL_SEQLEN=hca_local_seqlen, + USE_CAUSAL=use_causal, + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_DMODEL=block_dmodel, + EXACT_TILES_M=exact_tiles_m, + EXACT_TILES_N=exact_tiles_n, + num_warps=int(os.getenv("PRIMUS_V4_ATTN_BWD_DQ_NUM_WARPS", "2")), + num_stages=int(os.getenv("PRIMUS_V4_ATTN_BWD_DQ_NUM_STAGES", "1")), + ) + + +def _run_dq_flydsl( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + dsink_fp32, + sink_arg, + *, + scale, + swa_window, + has_sink, + Sq, + Sk, + HQ, + HK, + B, + D, +): + """V4 SWA bwd dQ via FlyDSL. Writes dq_fp32 (fp32) and adds into + dsink_fp32 (fp32, atomic) when has_sink. Sink_arg is a fp32 dummy + buffer when has_sink=False. + + Expects BHLD contiguous: q [B,HQ,Sq,D], k/v [B,HK,Sk,D], dout/lse/d_buf + same as Triton path. For MQA we ENFORCE HK==1 and pass k/v as-is + (the FlyDSL kernel uses stride_kh=0 via the mqa_kv codepath, which + expects K/V to be flat [B, Sk, D]-contiguous regardless of stride). + """ + mqa_kv = HK == 1 + # Sanity: V4 STEP 1 only supports MQA (HK==1). Keep this gate; if HK!=1 + # we fall back to the Triton path upstream. + assert mqa_kv, f"FlyDSL dq STEP 1b only supports MQA (HK==1); got HK={HK}" + dtype_str = "bf16" if q.dtype == torch.bfloat16 else "f16" + + launch = _get_fly_dq( + num_heads=HQ, + head_dim=D, + swa_window=int(swa_window), + dtype_str=dtype_str, + mqa_kv=True, + has_sink=has_sink, + ) + + # The FlyDSL kernel reads K/V via flat pointer ops, so they must be + # contiguous and rank-4 doesn't strictly matter. We pass the rank-4 + # tensors directly; the BHLD path computes ((b*1 + 0)*Sk + n)*D + col + # for KV when mqa_kv=True. That matches the stride pattern of a + # [B, 1, Sk, D]-contiguous tensor. + # NOTE: K and V come in as [B, 1, Sk, D] from the harness MQA leaves + # (see bwd_modes _build_swa_inputs). + assert q.is_contiguous(), "q must be contiguous" + assert k.is_contiguous(), "k must be contiguous" + assert v.is_contiguous(), "v must be contiguous" + assert dout.is_contiguous(), "dout must be contiguous" + assert lse.is_contiguous(), "lse must be contiguous" + assert d_buf.is_contiguous(), "d_buf must be contiguous" + assert dq_fp32.is_contiguous(), "dq_fp32 must be contiguous" + + launch( + q, # Q [B, HQ, Sq, D] + k, # K [B, 1, Sk, D] (MQA broadcast view) + v, # V [B, 1, Sk, D] + dout, # DOS [B, HQ, Sq, D] + lse, # LSE [B, HQ, Sq] fp32 + d_buf, # DELTAS [B, HQ, Sq] fp32 + dq_fp32, # DQ [B, HQ, Sq, D] fp32 (OUTPUT) + dsink_fp32, # DSINK [HQ] fp32 (OUTPUT, atomic) + sink_arg, # SINK [HQ] fp32 (INPUT, dummy if !has_sink) + int(B), + int(Sq), + int(Sk), + ) + + +def _run_dkv_flydsl( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + *, + scale, + swa_window, + has_sink, + Sq, + Sk, + HQ, + HK, + B, + D, +): + """V4 SWA bwd dKdV via FlyDSL. Writes dk_fp32 / dv_fp32 (fp32 buffers). + MQA only (HK==1). One program per (b, n_block), head-loop accumulator, + no atomics.""" + mqa_kv = HK == 1 + assert mqa_kv, f"FlyDSL dkv STEP 1c only supports MQA (HK==1); got HK={HK}" + print( + f"[provenance] FlyDSL dkv launcher invoked, B={B} H={HQ} Sk={Sk} D={D}", + flush=True, + ) + dtype_str = "bf16" if q.dtype == __import__("torch").bfloat16 else "f16" + launch = _get_fly_dkv( + num_heads=HQ, + head_dim=D, + swa_window=int(swa_window), + dtype_str=dtype_str, + mqa_kv=True, + ) + assert q.is_contiguous(), "q must be contiguous" + assert k.is_contiguous(), "k must be contiguous" + assert v.is_contiguous(), "v must be contiguous" + assert dout.is_contiguous(), "dout must be contiguous" + assert lse.is_contiguous(), "lse must be contiguous" + assert d_buf.is_contiguous(), "d_buf must be contiguous" + assert dk_fp32.is_contiguous(), "dk_fp32 must be contiguous" + assert dv_fp32.is_contiguous(), "dv_fp32 must be contiguous" + launch( + q, # Q [B, HQ, Sq, D] + k, # K [B, 1, Sk, D] + v, # V [B, 1, Sk, D] + dout, # DOS [B, HQ, Sq, D] + lse, # LSE [B, HQ, Sq] fp32 (RAW domain) + d_buf, # DELTAS [B, HQ, Sq] fp32 + dk_fp32, # DK [B, 1, Sk, D] fp32 (OUTPUT) + dv_fp32, # DV [B, 1, Sk, D] fp32 (OUTPUT) + int(B), + int(Sq), + int(Sk), + ) + + +def _run_dq_pool_flydsl( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + add_mask, + *, + pool_size, + hca_local_seqlen, + Sq, + Sk, + HQ, + HK, + B, + D, +): + """V4 HCA bwd dq POOL stream via FlyDSL. ACCUMULATES into dq_fp32. + + Called AFTER ``_run_dq_flydsl`` has written the LOCAL stream dq into + dq_fp32. Race-free since each program owns a unique + (b, qhid, m_block) slice and the launch is sequenced after the local + dq launch. + """ + mqa_kv = HK == 1 + assert mqa_kv, f"FlyDSL dq_pool only supports MQA (HK==1); got HK={HK}" + assert pool_size <= 64, f"pool_size must fit in one BLOCK_N=64 block; got pool_size={pool_size}" + assert ( + hca_local_seqlen % 64 == 0 + ), f"hca_local_seqlen must be multiple of BLOCK_N=64; got {hca_local_seqlen}" + assert ( + hca_local_seqlen + pool_size + ) == Sk, f"expected hca_local_seqlen+pool_size == Sk; got {hca_local_seqlen}+{pool_size}!={Sk}" + assert add_mask is not None and add_mask.shape == ( + Sq, + pool_size, + ), f"add_mask shape mismatch: expected ({Sq},{pool_size}); got {tuple(add_mask.shape) if add_mask is not None else None}" + assert add_mask.dtype == q.dtype, f"add_mask dtype must match q.dtype; got {add_mask.dtype} vs {q.dtype}" + dtype_str = "bf16" if q.dtype == torch.bfloat16 else "f16" + launch = _get_fly_dq_pool( + num_heads=HQ, + head_dim=D, + pool_size=int(pool_size), + hca_local_seqlen=int(hca_local_seqlen), + dtype_str=dtype_str, + mqa_kv=True, + ) + assert q.is_contiguous(), "q must be contiguous" + assert k.is_contiguous(), "k must be contiguous" + assert v.is_contiguous(), "v must be contiguous" + assert dout.is_contiguous(), "dout must be contiguous" + assert lse.is_contiguous(), "lse must be contiguous" + assert d_buf.is_contiguous(), "d_buf must be contiguous" + assert dq_fp32.is_contiguous(), "dq_fp32 must be contiguous" + assert add_mask.is_contiguous(), "add_mask must be contiguous" + launch( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + add_mask, + int(B), + int(Sq), + int(Sk), + ) + + +def _run_dkv_pool_flydsl( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + add_mask, + *, + pool_size, + hca_local_seqlen, + Sq, + Sk, + HQ, + HK, + B, + D, +): + """V4 HCA bwd dKdV POOL stream via FlyDSL. WRITES into the POOL slice of + dk_fp32/dv_fp32 (which are zero-initialized by the wrapper and have the + LOCAL slice already populated by ``_run_dkv_flydsl``). Race-free: + LOCAL writes [..., :hca_local_seqlen, :] and POOL writes + [..., hca_local_seqlen:hca_local_seqlen+pool_size, :] -- disjoint. + """ + mqa_kv = HK == 1 + assert mqa_kv, f"FlyDSL dkv_pool only supports MQA (HK==1); got HK={HK}" + assert pool_size <= 32, f"pool_size must fit in one BLOCK_N=32 block; got pool_size={pool_size}" + assert ( + hca_local_seqlen % 32 == 0 + ), f"hca_local_seqlen must be multiple of BLOCK_N=32; got {hca_local_seqlen}" + assert (hca_local_seqlen + pool_size) == Sk, ( + f"expected hca_local_seqlen+pool_size == Sk; got " f"{hca_local_seqlen}+{pool_size}!={Sk}" + ) + assert add_mask is not None and add_mask.shape == (Sq, pool_size), ( + f"add_mask shape mismatch: expected ({Sq},{pool_size}); got " + f"{tuple(add_mask.shape) if add_mask is not None else None}" + ) + assert add_mask.dtype == q.dtype, f"add_mask dtype must match q.dtype; got {add_mask.dtype} vs {q.dtype}" + dtype_str = "bf16" if q.dtype == torch.bfloat16 else "f16" + launch = _get_fly_dkv_pool( + num_heads=HQ, + head_dim=D, + pool_size=int(pool_size), + hca_local_seqlen=int(hca_local_seqlen), + dtype_str=dtype_str, + mqa_kv=True, + ) + assert q.is_contiguous(), "q must be contiguous" + assert k.is_contiguous(), "k must be contiguous" + assert v.is_contiguous(), "v must be contiguous" + assert dout.is_contiguous(), "dout must be contiguous" + assert lse.is_contiguous(), "lse must be contiguous" + assert d_buf.is_contiguous(), "d_buf must be contiguous" + assert dk_fp32.is_contiguous(), "dk_fp32 must be contiguous" + assert dv_fp32.is_contiguous(), "dv_fp32 must be contiguous" + assert add_mask.is_contiguous(), "add_mask must be contiguous" + launch( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + add_mask, + int(B), + int(Sq), + int(Sk), + ) + + +def _run_dkv_triton( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + mask_arg, + *, + scale, + swa_window_constexpr, + has_add_mask, + hca_local_seqlen, + use_causal, + block_m, + block_n, + block_dmodel, + stride_ms, + stride_mn, + Sq, + Sk, + HQ, + HK, + B, + exact_tiles_m, +): + dkv_block_n = block_n + dkv_n_blocks = triton.cdiv(Sk, dkv_block_n) + num_head_groups = 1 + if HQ > HK: + # MQA/HQ>=128 (V4-Pro): HG=2 on gfx950 mirrors the prod + # _launch_v4_attention_bwd default. Overridable via env. + _hg_default = "2" if (HQ >= 64 and HK == 1) else "1" + target = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_HEAD_GROUPS", _hg_default)) + while target > 1 and HQ % target != 0: + target //= 2 + num_head_groups = max(1, target) + dkv_grid = (dkv_n_blocks, B * HK, num_head_groups) + _v4_attention_bwd_dkv_kernel[dkv_grid]( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + mask_arg, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dk_fp32.stride(0), + dk_fp32.stride(1), + dk_fp32.stride(2), + dk_fp32.stride(3), + dv_fp32.stride(0), + dv_fp32.stride(1), + dv_fp32.stride(2), + dv_fp32.stride(3), + stride_ms, + stride_mn, + Sq, + Sk, + float(scale), + HEAD_Q=HQ, + HEAD_K=HK, + SWA_WINDOW=swa_window_constexpr, + HAS_ADD_MASK=has_add_mask, + HCA_LOCAL_SEQLEN=hca_local_seqlen, + USE_CAUSAL=use_causal, + BLOCK_M=block_m, + BLOCK_N=dkv_block_n, + BLOCK_DMODEL=block_dmodel, + NUM_HEAD_GROUPS=num_head_groups, + EXACT_TILES_M=exact_tiles_m, + EXACT_TILES_N=(Sk % dkv_block_n) == 0, + num_warps=int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_NUM_WARPS", "2")), + num_stages=int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_NUM_STAGES", "1")), + ) + + +# --------------------------------------------------------------------------- +# HCA pool-only dq accumulator (Triton, inline). +# +# This is a TEMPORARY Triton fallback for the POOL stream of HCA dq while +# the FlyDSL pool kernel is being authored. It is structurally equivalent +# to the pool branch of ``_v4_attention_bwd_dq_kernel`` (Triton ref) but +# written as a standalone kernel so we can call it as a pure accumulator +# (``DQ += pool_contrib``, no local loop). One program per (m_block, b*qhid). +# --------------------------------------------------------------------------- + + +import triton.language as tl # noqa: E402 + + +@triton.jit +def _hca_pool_dq_accumulator( + Q, + K, + V, + DOUT, + LSE, + D, + DQ, + ADD_MASK, + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_kb, + stride_kh, + stride_kn, + stride_kd, + stride_vb, + stride_vh, + stride_vn, + stride_vd, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dqb, + stride_dqh, + stride_dqm, + stride_dqd, + stride_ms, + stride_mn, + seqlen_q, + seqlen_k, + pool_size, + sm_scale, + HEAD_Q: tl.constexpr, + HEAD_K: tl.constexpr, + HCA_LOCAL_SEQLEN: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + bid = pid_bh // HEAD_Q + qhid = pid_bh % HEAD_Q + if HEAD_K == HEAD_Q: + khid = qhid + else: + khid = 0 + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_DMODEL) + pool_n = tl.arange(0, BLOCK_N) + offs_n = HCA_LOCAL_SEQLEN + pool_n + NEG_INF: tl.constexpr = -1.0e30 + pool_n_mask = pool_n < pool_size + + q_ptrs = ( + Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd + ) + dout_ptrs = ( + DOUT + + bid * stride_dob + + qhid * stride_doh + + offs_m[:, None] * stride_dom + + offs_d[None, :] * stride_dod + ) + lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm + dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm + + q_load_mask = offs_m[:, None] < seqlen_q + q = tl.load(q_ptrs, mask=q_load_mask, other=0.0) + dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0) + lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0) + dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0) + + k_ptrs = ( + K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd + ) + v_ptrs = ( + V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd + ) + kv_load_mask = pool_n_mask[:, None] + k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0) + v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0) + + qk = tl.dot(q, tl.trans(k)) * sm_scale + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & pool_n_mask[None, :] + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + qk = tl.where(pool_n_mask[None, :], qk, NEG_INF) + qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout, tl.trans(v)) + ds = p * (dp - dvec[:, None]) + dq_contrib = tl.dot(ds.to(k.dtype), k) * sm_scale + + dq_ptrs = ( + DQ + + bid * stride_dqb + + qhid * stride_dqh + + offs_m[:, None] * stride_dqm + + offs_d[None, :] * stride_dqd + ) + dq_prev = tl.load(dq_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0) + tl.store(dq_ptrs, dq_prev + dq_contrib, mask=offs_m[:, None] < seqlen_q) + + +def flydsl_v4_attention_bwd( + q: torch.Tensor, # [B, HQ, Sq, D] + k: torch.Tensor, # [B, HK, Sk, D] (HK == 1 for MQA, HK == HQ for MHA) + v: torch.Tensor, # [B, HK, Sk, D] + out: torch.Tensor, # [B, HQ, Sq, D] + dout: torch.Tensor, # [B, HQ, Sq, D] + lse: torch.Tensor, # [B, HQ, Sq] fp32 + *, + sink: Optional[torch.Tensor], + swa_window: int, + additive_mask: Optional[torch.Tensor], + scale: float, + hca_local_seqlen: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """V4 SWA backward launcher (STEP 1b scope: SWA-only, no HCA).""" + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise ValueError("rank-4 q/k/v required") + if dout.shape != out.shape or out.shape != q.shape: + raise ValueError( + f"shape mismatch: q={tuple(q.shape)} out={tuple(out.shape)} dout={tuple(dout.shape)}" + ) + B, HQ, Sq, D = q.shape + HK = k.shape[1] + Sk = k.shape[2] + if k.shape != (B, HK, Sk, D) or v.shape != k.shape: + raise ValueError(f"k/v shape mismatch: k={tuple(k.shape)} v={tuple(v.shape)}") + if q.dtype != torch.bfloat16: + raise NotImplementedError(f"bf16 only; got q.dtype={q.dtype}") + if D % 16 != 0: + raise NotImplementedError(f"head_dim must be multiple of 16; got D={D}") + # STEP 2 HCA: accept SWA-only (additive_mask=None, hca_local_seqlen=0) + # AND HCA (additive_mask is [Sq, P], hca_local_seqlen=Sq, Sk=Sq+P). + is_hca = (additive_mask is not None) and (int(hca_local_seqlen) > 0) + if (additive_mask is not None) != (int(hca_local_seqlen) > 0): + raise NotImplementedError( + "additive_mask and hca_local_seqlen must both be set (HCA) or both unset (SWA)" + ) + if is_hca: + if int(hca_local_seqlen) != Sq: + raise NotImplementedError( + f"HCA requires hca_local_seqlen == Sq; got {int(hca_local_seqlen)} vs Sq={Sq}" + ) + if int(swa_window) <= 0: + raise NotImplementedError("HCA requires swa_window > 0 for the local stream") + if Sk <= Sq: + raise NotImplementedError(f"HCA requires Sk>Sq; got Sk={Sk} Sq={Sq}") + pool_size = Sk - Sq + if additive_mask.shape != (Sq, pool_size): + raise NotImplementedError( + f"HCA additive_mask must be [Sq={Sq}, P={pool_size}]; got {tuple(additive_mask.shape)}" + ) + if B != 1: + raise NotImplementedError( + f"HCA path currently only supports B=1 (got B={B}); B>1 would mis-stride K/V." + ) + else: + if int(swa_window) <= 0: + raise NotImplementedError("SWA requires swa_window > 0") + if Sq != Sk: + raise NotImplementedError("SWA requires Sq == Sk") + + has_sink = sink is not None + has_add_mask = False + use_causal = False + swa_window_constexpr = int(swa_window) + + BLOCK_M = int(os.getenv("PRIMUS_V4_ATTN_BWD_BLOCK_M", "32")) + BLOCK_N = int(os.getenv("PRIMUS_V4_ATTN_BWD_BLOCK_N", "16")) + BLOCK_DMODEL = D + exact_tiles_m = (Sq % BLOCK_M) == 0 + exact_tiles_n = (Sk % BLOCK_N) == 0 + + dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dk_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=torch.float32) + dv_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=torch.float32) + if has_sink: + dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32) + sink_arg = sink.to(torch.float32) if sink.dtype != torch.float32 else sink + else: + # Dummy buffers; FlyDSL kernel still takes them but doesn't touch + # them when has_sink=False at build time. For Triton path we mirror + # the original behavior of passing q-shape sentinels. + dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32) + sink_arg = torch.zeros((HQ,), device=q.device, dtype=torch.float32) + + use_fly_preprocess = os.getenv("V4_FLYDSL_BWD_FLY_PREPROCESS", "1") == "1" + d_buf = _run_preprocess( + out, + dout, + use_flydsl=use_fly_preprocess, + block_m=BLOCK_M, + block_dmodel=BLOCK_DMODEL, + ) + + mask_arg = q + stride_ms = 0 + stride_mn = 0 + + use_fly_dq = os.getenv("V4_FLYDSL_BWD_FLY_DQ", "1") == "1" + use_fly_dkv = os.getenv("V4_FLYDSL_BWD_FLY_DKV", "1") == "1" + verbose = os.getenv("V4_FLYDSL_BWD_VERBOSE", "0") == "1" + + if verbose: + pp_tag = "flydsl" if use_fly_preprocess else "triton" + dq_tag = "flydsl" if use_fly_dq else "triton" + dkv_tag = "flydsl" if use_fly_dkv else "triton" + print( + f"[v4_bwd] preproc={pp_tag} dq={dq_tag} dkv={dkv_tag} " + f"has_sink={has_sink} Sq={Sq} Sk={Sk} HQ={HQ} HK={HK} D={D} swa={swa_window}", + flush=True, + ) + + if is_hca: + # ---- HCA mode (split-mask) ---- + # 1. LOCAL stream (n < HCA_LOCAL_SEQLEN): existing FlyDSL dq + dkv, + # with effective seq_len_k = HCA_LOCAL_SEQLEN. Kernel uses + # seq_len_k for both n-loop bound and batch base; B=1 makes the + # base zero so the reduced seq_len_k cleanly bounds the n-loop + # without mis-striding K/V. (B=1-only; gated above.) The LOCAL + # pass uses the JOINT lse (saved from fwd with both streams) and + # JOINT delta (preprocess used full out, dout). + # 2. POOL stream (n in [HCA_LOCAL_SEQLEN, Sk)): Triton pool kernels + # for this round; FlyDSL pool kernels are next-round work. + Sk_local = int(hca_local_seqlen) # == Sq + pool_size = Sk - Sk_local + sink_dq_arg = sink_arg # both fp32 [HQ] + use_fly_dq_pool = os.getenv("V4_FLYDSL_BWD_FLY_DQ_POOL", "0") == "1" + use_fly_dkv_pool = os.getenv("V4_FLYDSL_BWD_FLY_DKV_POOL", "0") == "1" + # Provenance: tag dq_pool / dkv_pool only based on knob (silent + # fallback would violate the strict-gate rule). + dq_pool_tag = "FlyDSL" if use_fly_dq_pool else "Triton" + dkv_pool_tag = "FlyDSL" if use_fly_dkv_pool else "Triton" + print( + f"[v4_bwd_hca] provenance: dq_local=FlyDSL dkv_local=FlyDSL " + f"dsink=FlyDSL dq_pool={dq_pool_tag} dkv_pool={dkv_pool_tag} " + f"B={B} HQ={HQ} HK={HK} Sq={Sq} Sk={Sk} pool={pool_size} D={D} " + f"swa={swa_window} hca_local_seqlen={Sk_local}", + flush=True, + ) + # ---- LOCAL FlyDSL dq (computes dq_local + dsink) ---- + _run_dq_flydsl( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + dsink_fp32, + sink_dq_arg, + scale=scale, + swa_window=swa_window_constexpr, + has_sink=has_sink, + Sq=Sq, + Sk=Sk_local, + HQ=HQ, + HK=HK, + B=B, + D=D, + ) + # ---- LOCAL FlyDSL dkv ---- + _run_dkv_flydsl( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + scale=scale, + swa_window=swa_window_constexpr, + has_sink=has_sink, + Sq=Sq, + Sk=Sk_local, + HQ=HQ, + HK=HK, + B=B, + D=D, + ) + # ---- POOL dq accumulator: FlyDSL (knob ON) or Triton (default) ---- + if use_fly_dq_pool: + _run_dq_pool_flydsl( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + additive_mask, + pool_size=int(pool_size), + hca_local_seqlen=Sk_local, + Sq=Sq, + Sk=Sk, + HQ=HQ, + HK=HK, + B=B, + D=D, + ) + else: + pool_block_n_dq = max(16, triton.next_power_of_2(pool_size)) + _hca_pool_dq_accumulator[(triton.cdiv(Sq, BLOCK_M), B * HQ)]( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + additive_mask, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_fp32.stride(0), + dq_fp32.stride(1), + dq_fp32.stride(2), + dq_fp32.stride(3), + additive_mask.stride(0), + additive_mask.stride(1), + Sq, + Sk, + pool_size, + float(scale), + HEAD_Q=HQ, + HEAD_K=HK, + HCA_LOCAL_SEQLEN=Sk_local, + BLOCK_M=BLOCK_M, + BLOCK_N=pool_block_n_dq, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=2, + num_stages=1, + ) + # ---- POOL dkv: FlyDSL (knob ON) or Triton (default) ---- + if use_fly_dkv_pool: + _run_dkv_pool_flydsl( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + additive_mask, + pool_size=int(pool_size), + hca_local_seqlen=Sk_local, + Sq=Sq, + Sk=Sk, + HQ=HQ, + HK=HK, + B=B, + D=D, + ) + else: + pool_block_n_dkv = max(16, triton.next_power_of_2(pool_size)) + pool_grid_m = triton.cdiv(Sq, BLOCK_M) + _v4_attention_bwd_dkv_pool_kernel[(pool_grid_m, B)]( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + additive_mask, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dk_fp32.stride(0), + dk_fp32.stride(1), + dk_fp32.stride(2), + dk_fp32.stride(3), + dv_fp32.stride(0), + dv_fp32.stride(1), + dv_fp32.stride(2), + dv_fp32.stride(3), + additive_mask.stride(0), + additive_mask.stride(1), + Sq, + Sk, + pool_size, + float(scale), + HEAD_Q=HQ, + HEAD_K=HK, + HCA_LOCAL_SEQLEN=Sk_local, + BLOCK_M=BLOCK_M, + BLOCK_N=pool_block_n_dkv, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=2, + num_stages=1, + ) + elif use_fly_dq: + # SWA-only path (STEP 1b unchanged). + sink_dq_arg = sink_arg # both are fp32 [HQ] + _run_dq_flydsl( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + dsink_fp32, + sink_dq_arg, + scale=scale, + swa_window=swa_window_constexpr, + has_sink=has_sink, + Sq=Sq, + Sk=Sk, + HQ=HQ, + HK=HK, + B=B, + D=D, + ) + else: + _run_dq_triton( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + dsink_fp32, + sink_arg, + mask_arg, + scale=scale, + swa_window_constexpr=swa_window_constexpr, + has_sink=has_sink, + has_add_mask=has_add_mask, + hca_local_seqlen=0, + use_causal=use_causal, + block_m=BLOCK_M, + block_n=BLOCK_N, + block_dmodel=BLOCK_DMODEL, + stride_ms=stride_ms, + stride_mn=stride_mn, + Sq=Sq, + Sk=Sk, + HQ=HQ, + HK=HK, + B=B, + exact_tiles_m=exact_tiles_m, + exact_tiles_n=exact_tiles_n, + ) + + if not is_hca: + if use_fly_dkv: + _run_dkv_flydsl( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + scale=scale, + swa_window=swa_window_constexpr, + has_sink=has_sink, + Sq=Sq, + Sk=Sk, + HQ=HQ, + HK=HK, + B=B, + D=D, + ) + else: + _run_dkv_triton( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + mask_arg, + scale=scale, + swa_window_constexpr=swa_window_constexpr, + has_add_mask=has_add_mask, + hca_local_seqlen=0, + use_causal=use_causal, + block_m=BLOCK_M, + block_n=BLOCK_N, + block_dmodel=BLOCK_DMODEL, + stride_ms=stride_ms, + stride_mn=stride_mn, + Sq=Sq, + Sk=Sk, + HQ=HQ, + HK=HK, + B=B, + exact_tiles_m=exact_tiles_m, + ) + + dq_out = dq_fp32.to(q.dtype) + dk_out = dk_fp32.to(k.dtype) + dv_out = dv_fp32.to(v.dtype) + dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None + return dq_out, dk_out, dv_out, dsink_out + + +__all__ = ["flydsl_v4_attention_bwd"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_csa.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_csa.py new file mode 100644 index 000000000..16eb93c0e --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_csa.py @@ -0,0 +1,219 @@ +"""V4 CSA attention forward FlyDSL launcher (Round 3 Step 2b). + +Stage A: correctness only. Per-row design forked from Triton monolithic CSA. +""" + +from __future__ import annotations + +import math +import os +import sys +import threading +from typing import Optional, Tuple + +import torch + +_FLYDSL_SRC = "/workspace/FlyDSL-amd" +if _FLYDSL_SRC not in sys.path: + sys.path.insert(0, _FLYDSL_SRC) + +os.environ.setdefault("FLYDSL_WAVES_PER_EU", "2") + +_HERE = os.path.dirname(os.path.abspath(__file__)) +if _HERE not in sys.path: + sys.path.insert(0, _HERE) +from v4_csa_fwd_kernel import build_v4_csa_fwd_module # noqa: E402 + +_KERNEL_CACHE = {} +_KERNEL_CACHE_LOCK = threading.Lock() + + +def _get_kernel( + num_heads_q, + head_dim, + swa_window, + dtype_str, + block_n, + block_k, + waves_per_eu, + has_sink, + has_sparse, + mqa_kv, + head_group=1, +): + key = ( + num_heads_q, + head_dim, + swa_window, + dtype_str, + block_n, + block_k, + waves_per_eu, + has_sink, + has_sparse, + mqa_kv, + head_group, + ) + with _KERNEL_CACHE_LOCK: + if key in _KERNEL_CACHE: + return _KERNEL_CACHE[key] + launch = build_v4_csa_fwd_module( + num_heads=num_heads_q, + head_dim=head_dim, + swa_window=int(swa_window), + dtype_str=dtype_str, + waves_per_eu=waves_per_eu, + block_n=block_n, + block_k=block_k, + has_sink=has_sink, + has_sparse=has_sparse, + mqa_kv=mqa_kv, + head_group=head_group, + ) + _KERNEL_CACHE[key] = launch + return launch + + +def _broadcast_kv_mqa(k: torch.Tensor, v: torch.Tensor, head_q: int): + if k.shape[1] == head_q: + return k, v + if k.shape[1] != 1: + raise ValueError(f"MQA expects K_H=1; got {k.shape[1]}") + k_full = k.expand(-1, head_q, -1, -1).clone(memory_format=torch.contiguous_format) + v_full = v.expand(-1, head_q, -1, -1).clone(memory_format=torch.contiguous_format) + return k_full, v_full + + +def _launch_v4_attention_fwd_csa( + q: torch.Tensor, # [B, H, Sq, D] + k_local: torch.Tensor, # [B, H_KV, Sq, D] (H_KV=1 for MQA) + v_local: torch.Tensor, # [B, H_KV, Sq, D] + gathered: torch.Tensor, # [B, Sq, K_topk, D] + *, + sink: Optional[torch.Tensor], # [H] fp32 or None + swa_window: int, + sparse_mask: torch.Tensor, # [B, Sq, K_topk] additive + scale: float, +) -> Tuple[torch.Tensor, torch.Tensor]: + """V4 CSA forward launcher. + + Returns (out, lse) where out is bf16 [B, H, Sq, D] and lse is fp32 [B, H, Sq] + in raw-qk-scaled-domain (m + ln(l), where m absorbs sm_scale). + """ + if q.dim() != 4 or k_local.dim() != 4 or v_local.dim() != 4: + raise ValueError("rank-4 q/k_local/v_local required") + if gathered.dim() != 4: + raise ValueError("rank-4 gathered [B,Sq,K_topk,D] required") + if sparse_mask.dim() != 3: + raise ValueError("rank-3 sparse_mask [B,Sq,K_topk] required") + B, HQ, Sq, D = q.shape + Bk, HK, Sk, Dk = k_local.shape + if (Bk, Sk, Dk) != (B, Sq, D) or v_local.shape != k_local.shape: + raise ValueError("k_local/v_local shape mismatch w.r.t. q") + if HK != 1 and HK != HQ: + raise ValueError(f"K_H must be 1 or {HQ}; got {HK}") + Bg, Sqg, K_topk, Dg = gathered.shape + if Bg != B or Sqg != Sq or Dg != D: + raise ValueError(f"gathered shape mismatch {tuple(gathered.shape)}") + Bm, Sqm, Km = sparse_mask.shape + if Bm != B or Sqm != Sq or Km != K_topk: + raise ValueError(f"sparse_mask shape mismatch {tuple(sparse_mask.shape)}") + if q.dtype != torch.bfloat16: + raise NotImplementedError(f"bf16 only; got {q.dtype}") + if D != 512: + raise NotImplementedError(f"head_dim=512 only; got {D}") + if int(swa_window) <= 0: + raise NotImplementedError("swa_window > 0 required") + expected_scale = 1.0 / math.sqrt(D) + if not math.isclose(float(scale), expected_scale, rel_tol=1e-4): + raise NotImplementedError(f"only scale=1/sqrt(D) supported; got {scale}") + + has_sink = sink is not None + has_sparse = int(K_topk) > 0 + mqa = (HK == 1) and (HQ > 1) + + if has_sink: + sink_fp32 = sink.float().contiguous() + if sink_fp32.shape != (HQ,): + raise ValueError(f"sink shape must be ({HQ},); got {tuple(sink_fp32.shape)}") + else: + sink_fp32 = torch.zeros((max(HQ, 1),), dtype=torch.float32, device=q.device) + + # MQA: avoid the .expand().clone() broadcast — pass [B,1,Sq,D] directly, + # kernel uses mqa_kv=True to drop head_idx from K/V indexing. + q_bhld = q.contiguous() + if mqa: + k_bhld = k_local.contiguous() + v_bhld = v_local.contiguous() + else: + k_bhld = k_local.contiguous() + v_bhld = v_local.contiguous() + o_bhld = torch.empty_like(q_bhld) + lse = torch.zeros((B, HQ, Sq), device=q.device, dtype=torch.float32) + + # Sparse mask & gathered must be contiguous fp32 / bf16. + if K_topk > 0: + gathered_c = gathered.contiguous() + # sparse_mask is bf16 (typically) -> upcast to fp32 here so kernel can + # just add it as f32. (Triton path takes the dtype of input.) + if sparse_mask.dtype != torch.float32: + sparse_mask_fp32 = sparse_mask.float().contiguous() + else: + sparse_mask_fp32 = sparse_mask.contiguous() + else: + gathered_c = torch.empty((B, Sq, 1, D), dtype=q.dtype, device=q.device) + sparse_mask_fp32 = torch.zeros((B, Sq, 1), dtype=torch.float32, device=q.device) + + block_n = int(os.environ.get("PRIMUS_V4_CSA_BLOCK_N", "8")) + block_k = int(os.environ.get("PRIMUS_V4_CSA_BLOCK_K", "16")) + waves_per_eu = int(os.environ.get("FLYDSL_WAVES_PER_EU", "2")) + # HG=2 is the banked default for the sparse path (same VGPR=250 / spill=0 + # footprint as HG=1). K_topk==0 (dense fallback) still forces HG=1 below. + head_group = int(os.environ.get("PRIMUS_V4_CSA_HEAD_GROUP", "2")) + # Round 5 fix: dense-fallback path (K_topk==0) has no payoff from head-group + # fusion AND triggered a flyc closure-cache collision when has_sparse differed + # between HG=1 dense and HG>1 sparse compiles. Force HG=1 when dense. + if not has_sparse: + head_group = 1 + if HQ % head_group != 0: + # Silently fall back to HG=1 if shape doesn't divide. + head_group = 1 + launch = _get_kernel( + HQ, + D, + int(swa_window), + "bf16", + block_n, + block_k, + waves_per_eu, + has_sink, + has_sparse, + mqa, + head_group=head_group, + ) + + # FlyDSL packs each tensor *shape* dim as int32 (jit_argument.py: + # `"i" * len(shape)`). A flat `.view(-1)` collapses a tensor to a single + # dim equal to its element count; for V4-Pro CSA `gathered` that is + # B*Sq*K_topk*D = 1*4096*1024*512 = 2**31, which overflows int32 ('i') + # and aborts the launch. The kernel only needs the aligned base pointer + # (it computes all offsets from the B / Sq / K_topk scalars), so passing + # `gathered` in its natural [B, Sq, K_topk, D] shape keeps every dim + # < 2**31 (strides are packed as int64) without changing kernel math. + launch( + q_bhld.view(-1), + k_bhld.view(-1), + v_bhld.view(-1), + gathered_c, + sparse_mask_fp32.view(-1), + sink_fp32.view(-1), + o_bhld.view(-1), + lse.view(-1), + B, + Sq, + int(K_topk), + ) + return o_bhld, lse + + +__all__ = ["_launch_v4_attention_fwd_csa"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_mqa.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_mqa.py new file mode 100644 index 000000000..70523c255 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_attention_fwd_flydsl_mqa.py @@ -0,0 +1,143 @@ +"""V4 SWA attention forward FlyDSL launcher with MQA stride-trick (Round 3 Stage C). + +Eliminates the K/V .expand().clone() broadcast that allocates H_Q copies of +K and V. Passes the un-broadcast [B, 1, Sk, D] view directly; the kernel reads +K/V with stride_kh=0 via the mqa_kv compile-time flag. +""" + +from __future__ import annotations + +import math +import os +import sys +import threading +from typing import Optional, Tuple + +import torch + +_FLYDSL_SRC = "/workspace/FlyDSL-amd" +if _FLYDSL_SRC not in sys.path: + sys.path.insert(0, _FLYDSL_SRC) + +os.environ.setdefault("FLYDSL_SLA_FWD_ENABLE_DMA", "1") +os.environ.setdefault("FLYDSL_WAVES_PER_EU", "2") + +_HERE = os.path.dirname(os.path.abspath(__file__)) +if _HERE not in sys.path: + sys.path.insert(0, _HERE) +from v4_sla_fwd_kernel import build_v4_swa_fwd_module # noqa: E402 + +_KERNEL_CACHE = {} +_KERNEL_CACHE_LOCK = threading.Lock() + + +def _get_kernel( + num_heads_q, head_dim, swa_window, dtype_str, block_m, block_n, waves_per_eu, mqa_kv, flat_work_group_size +): + key = ( + num_heads_q, + head_dim, + swa_window, + dtype_str, + block_m, + block_n, + waves_per_eu, + mqa_kv, + flat_work_group_size, + ) + with _KERNEL_CACHE_LOCK: + if key in _KERNEL_CACHE: + return _KERNEL_CACHE[key] + launch = build_v4_swa_fwd_module( + num_heads=num_heads_q, + head_dim=head_dim, + swa_window=int(swa_window), + dtype_str=dtype_str, + waves_per_eu=waves_per_eu, + block_m=block_m, + block_n=block_n, + flat_work_group_size=flat_work_group_size, + layout_bhld=True, + mqa_kv=mqa_kv, + ) + _KERNEL_CACHE[key] = launch + return launch + + +def _launch_v4_attention_fwd_flydsl_mqa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + sink: Optional[torch.Tensor], + swa_window: int, + additive_mask: Optional[torch.Tensor], + scale: float, + hca_local_seqlen: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor]: + """SWA-only launcher that AVOIDS the MQA broadcast. + + For K_H=1 (MQA), passes the [B, 1, Sk, D] tensor directly and uses + the mqa_kv kernel variant. For K_H=HQ (non-MQA), falls back to the + standard kernel. + """ + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise ValueError("rank-4 q/k/v required") + B, HQ, Sq, D = q.shape + Bk, HK, Sk, Dk = k.shape + if (Bk, Sk, Dk) != (B, Sk, D) or v.shape != k.shape: + raise ValueError(f"shape mismatch") + if HK != 1 and HK != HQ: + raise ValueError(f"K_H must be 1 or {HQ}; got {HK}") + if q.dtype != torch.bfloat16: + raise NotImplementedError(f"bf16 only") + if D != 512: + raise NotImplementedError(f"head_dim=512 only") + if sink is not None: + raise NotImplementedError("sink not yet supported in MQA wrapper") + if additive_mask is not None or int(hca_local_seqlen) != 0: + raise NotImplementedError("HCA mode not in this wrapper") + if int(swa_window) <= 0: + raise NotImplementedError("swa_window > 0 required") + if Sq != Sk: + raise NotImplementedError("SWA: Sq==Sk required") + expected_scale = 1.0 / math.sqrt(D) + if not math.isclose(float(scale), expected_scale, rel_tol=1e-4): + raise NotImplementedError(f"scale=1/sqrt(D) only") + + mqa = (HK == 1) and (HQ > 1) + + q_bhld = q.contiguous() + if mqa: + # K/V are [B, 1, Sk, D]. We pass them flat-view but use mqa_kv kernel + # variant that drops head_idx from K/V indexing. + k_bhld = k.contiguous() + v_bhld = v.contiguous() + else: + k_bhld = k.contiguous() + v_bhld = v.contiguous() + o_bhld = torch.empty_like(q_bhld) + lse = torch.zeros((B, HQ, Sq), device=q.device, dtype=torch.float32) + + block_m = int(os.environ.get("PRIMUS_V4_FLYDSL_BLOCK_M", "128")) + block_n = int(os.environ.get("PRIMUS_V4_FLYDSL_BLOCK_N", "32")) + waves_per_eu = int(os.environ.get("FLYDSL_WAVES_PER_EU", "2")) + fwgs_env = os.environ.get("PRIMUS_V4_FLYDSL_FWGS", "") + flat_work_group_size = int(fwgs_env) if fwgs_env else None + launch = _get_kernel( + HQ, D, int(swa_window), "bf16", block_m, block_n, waves_per_eu, mqa, flat_work_group_size + ) + + launch( + q_bhld.view(-1), + k_bhld.view(-1), + v_bhld.view(-1), + o_bhld.view(-1), + lse.view(-1), + B, + Sq, + ) + return o_bhld, lse + + +__all__ = ["_launch_v4_attention_fwd_flydsl_mqa"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_attention_bwd_flydsl_mqa.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_attention_bwd_flydsl_mqa.py new file mode 100644 index 000000000..373128656 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_attention_bwd_flydsl_mqa.py @@ -0,0 +1,399 @@ +"""V4 CSA attention backward FlyDSL launcher (Phase B STEP 3a + 3b). + +Env knobs: + V4_FLYDSL_CSA_BWD_FLY_DQ default 0 (1 = FlyDSL dq, 0 = Triton dq) + V4_FLYDSL_CSA_BWD_FLY_DKV default 0 (1 = FlyDSL dk_local/dv_local/ + dgathered/dsink, 0 = Triton) + V4_FLYDSL_BWD_VERBOSE default 0 + +Wiring: + knob_dq=0, knob_dkv=0 -> all Triton. + knob_dq=1, knob_dkv=0 -> FlyDSL dq-only; Triton emits all others. + knob_dq=0, knob_dkv=1 -> Triton dq; FlyDSL emits dk/dv/dgathered/dsink. + (rare path; uses the FULL FlyDSL kernel and + discards its dq.) + knob_dq=1, knob_dkv=1 -> ONE FlyDSL launch produces all 5 grads + (no Triton call needed; the cheapest path + when both are on). + +The launcher always emits a provenance line so the harness can prove the +expected backend was used. +""" + +from __future__ import annotations + +import os +import sys +import threading +from typing import Optional, Tuple + +import torch + +_FLYDSL_SRC = "/workspace/FlyDSL-amd" +if _FLYDSL_SRC not in sys.path: + sys.path.insert(0, _FLYDSL_SRC) + +os.environ.setdefault("FLYDSL_WAVES_PER_EU", "2") + +_HERE = os.path.dirname(os.path.abspath(__file__)) +if _HERE not in sys.path: + sys.path.insert(0, _HERE) +if "/workspace/Primus" not in sys.path: + sys.path.insert(0, "/workspace/Primus") + +import triton # noqa: E402 + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v0_deprecated.v4_csa_attention_bwd import ( # noqa: E402 + _launch_v4_csa_attention_bwd, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import ( # noqa: E402 + _v4_attention_bwd_preprocess_kernel, +) + +_DQ_KERNEL_CACHE = {} +_DQ_KERNEL_LOCK = threading.Lock() +_FULL_KERNEL_CACHE = {} +_FULL_KERNEL_LOCK = threading.Lock() + + +def _get_fly_csa_dq(num_heads, head_dim, swa_window, dtype_str, has_sink, has_sparse, block_n, block_k): + key = (num_heads, head_dim, swa_window, dtype_str, has_sink, has_sparse, block_n, block_k) + with _DQ_KERNEL_LOCK: + if key in _DQ_KERNEL_CACHE: + return _DQ_KERNEL_CACHE[key] + from v4_csa_bwd_dq_kernel import build_v4_csa_bwd_dq_module + + launch = build_v4_csa_bwd_dq_module( + num_heads=num_heads, + head_dim=head_dim, + swa_window=swa_window, + dtype_str=dtype_str, + has_sink=has_sink, + has_sparse=has_sparse, + block_n=block_n, + block_k=block_k, + mqa_kv=True, + ) + _DQ_KERNEL_CACHE[key] = launch + return launch + + +def _get_fly_csa_full(num_heads, head_dim, swa_window, dtype_str, has_sink, has_sparse, block_n, block_k): + key = (num_heads, head_dim, swa_window, dtype_str, has_sink, has_sparse, block_n, block_k) + with _FULL_KERNEL_LOCK: + if key in _FULL_KERNEL_CACHE: + return _FULL_KERNEL_CACHE[key] + from v4_csa_bwd_full_kernel import build_v4_csa_bwd_full_module + + launch = build_v4_csa_bwd_full_module( + num_heads=num_heads, + head_dim=head_dim, + swa_window=swa_window, + dtype_str=dtype_str, + has_sink=has_sink, + has_sparse=has_sparse, + block_n=block_n, + block_k=block_k, + mqa_kv=True, + ) + _FULL_KERNEL_CACHE[key] = launch + return launch + + +def _preprocess_deltas(out, dout, B, HQ, Sq, D): + BLOCK_M_PRE = 64 + d_buf = torch.empty((B, HQ, Sq), device=out.device, dtype=torch.float32) + pre_grid = (triton.cdiv(Sq, BLOCK_M_PRE), B * HQ) + _v4_attention_bwd_preprocess_kernel[pre_grid]( + out, + dout, + d_buf, + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + Sq, + HEAD=HQ, + BLOCK_M=BLOCK_M_PRE, + BLOCK_DMODEL=D, + num_warps=4, + num_stages=1, + ) + return d_buf + + +def flydsl_v4_csa_attention_bwd( + q: torch.Tensor, + k_local: torch.Tensor, + v_local: torch.Tensor, + gathered: torch.Tensor, + sparse_mask: torch.Tensor, + out: torch.Tensor, + dout: torch.Tensor, + lse: torch.Tensor, + *, + sink: Optional[torch.Tensor], + swa_window: int, + scale: float, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + if q.dim() != 4 or k_local.dim() != 4 or v_local.dim() != 4: + raise ValueError("rank-4 q / k_local / v_local required") + if dout.shape != out.shape or out.shape != q.shape: + raise ValueError( + f"shape mismatch: q={tuple(q.shape)} out={tuple(out.shape)} dout={tuple(dout.shape)}" + ) + B, HQ, Sq, D = q.shape + K_topk = gathered.shape[2] + has_sink = sink is not None + + use_fly_dq = os.getenv("V4_FLYDSL_CSA_BWD_FLY_DQ", "0") == "1" + use_fly_dkv = os.getenv("V4_FLYDSL_CSA_BWD_FLY_DKV", "0") == "1" + verbose = os.getenv("V4_FLYDSL_BWD_VERBOSE", "0") == "1" + + dq_tag = "FlyDSL" if use_fly_dq else "Triton" + dkv_tag = "FlyDSL" if use_fly_dkv else "Triton" + dsink_tag = ( + "FlyDSL" if (use_fly_dkv and has_sink) else ("FlyDSL" if (use_fly_dq and has_sink) else "Triton") + ) + print( + f"[v4_csa_bwd] provenance: dq={dq_tag} dk_local={dkv_tag} dv_local={dkv_tag} " + f"dgathered={dkv_tag} dsink={dsink_tag} " + f"B={B} HQ={HQ} Sq={Sq} K_topk={K_topk} D={D} swa={swa_window} " + f"has_sink={has_sink}", + flush=True, + ) + + # ============================================================ + # Path 1: nothing FlyDSL -> direct Triton pass-through. + # ============================================================ + if not use_fly_dq and not use_fly_dkv: + return _launch_v4_csa_attention_bwd( + q, + k_local, + v_local, + gathered, + sparse_mask, + out, + dout, + lse, + sink=sink, + swa_window=int(swa_window), + scale=float(scale), + ) + + # ============================================================ + # Validate the FlyDSL kernel's preconditions. + # ============================================================ + if q.dtype != torch.bfloat16: + raise NotImplementedError(f"FlyDSL CSA bwd supports bf16 only; got {q.dtype}") + if D % 64 != 0: + raise NotImplementedError(f"FlyDSL CSA bwd requires D % 64 == 0; got D={D}") + + # MQA view -- the Triton wrapper expands MQA->MHA by .expand().contiguous(), + # so head 0 has the original values. + k_mqa = k_local[:, :1, :, :].contiguous() + v_mqa = v_local[:, :1, :, :].contiguous() + + q_c = q.contiguous() + dout_c = dout.contiguous() + lse_c = lse.contiguous() + gathered_c = gathered.contiguous() + sparse_mask_c = sparse_mask.contiguous() + + d_buf = _preprocess_deltas(out, dout_c, B, HQ, Sq, D) + + if has_sink: + sink_arg = sink.to(torch.float32) if sink.dtype != torch.float32 else sink.contiguous() + else: + sink_arg = torch.zeros((HQ,), device=q.device, dtype=torch.float32) + + block_n = 32 + block_k = 32 + dtype_str = "bf16" + has_sparse = K_topk > 0 + + # ============================================================ + # Path 2: BOTH knobs on -- one FlyDSL full-kernel launch. + # ============================================================ + if use_fly_dq and use_fly_dkv: + dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + # dk_local / dv_local are stored at [B, HQ, Sq, D] (MHA-shape buffer + # matching Triton's atomic-add target). Caller is free to reduce. + dkl_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dvl_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dgathered_fp32 = torch.zeros((B, Sq, K_topk, D), device=q.device, dtype=torch.float32) + dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32) + + launch = _get_fly_csa_full( + num_heads=HQ, + head_dim=D, + swa_window=int(swa_window), + dtype_str=dtype_str, + has_sink=has_sink, + has_sparse=has_sparse, + block_n=block_n, + block_k=block_k, + ) + + if verbose: + print( + f"[v4_csa_bwd] FlyDSL FULL launch: B={B} HQ={HQ} Sq={Sq} K_topk={K_topk} " + f"D={D} swa={swa_window} has_sink={has_sink} has_sparse={has_sparse}", + flush=True, + ) + + launch( + q_c, + k_mqa, + v_mqa, + gathered_c, + sparse_mask_c, + dout_c, + lse_c, + d_buf, + sink_arg, + dq_fp32, + dkl_fp32, + dvl_fp32, + dgathered_fp32, + dsink_fp32, + int(B), + int(Sq), + int(K_topk), + ) + + dq_out = dq_fp32.to(q.dtype) + # dk_local / dv_local return shape [B, HQ, Sq, D] to match the Triton + # output (it returns the broadcast-MHA buffer; bwd_modes._run_csa_pair + # does the autograd reduction itself when the leaves were created + # via .expand().contiguous()). + dkl_out = dkl_fp32.to(k_local.dtype) + dvl_out = dvl_fp32.to(v_local.dtype) + dg_out = dgathered_fp32.to(gathered.dtype) + dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None + return dq_out, dkl_out, dvl_out, dg_out, dsink_out + + # ============================================================ + # Path 3: partial overlap with Triton. + # ============================================================ + # For partial paths we always run the Triton kernel and selectively + # override its outputs with FlyDSL values. + triton_out = _launch_v4_csa_attention_bwd( + q, + k_local, + v_local, + gathered, + sparse_mask, + out, + dout, + lse, + sink=sink, + swa_window=int(swa_window), + scale=float(scale), + ) + dq_t, dkl_t, dvl_t, dg_t, dsink_t = triton_out + + if use_fly_dq: + # Override dq (+ dsink for sink case) with FlyDSL dq-only kernel. + dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32) + launch = _get_fly_csa_dq( + num_heads=HQ, + head_dim=D, + swa_window=int(swa_window), + dtype_str=dtype_str, + has_sink=has_sink, + has_sparse=has_sparse, + block_n=block_n, + block_k=block_k, + ) + if verbose: + print( + f"[v4_csa_bwd] FlyDSL dq launch: B={B} HQ={HQ} Sq={Sq} K_topk={K_topk} " + f"D={D} swa={swa_window} has_sink={has_sink} has_sparse={has_sparse}", + flush=True, + ) + launch( + q_c, + k_mqa, + v_mqa, + gathered_c, + sparse_mask_c, + dout_c, + lse_c, + d_buf, + sink_arg, + dq_fp32, + dsink_fp32, + int(B), + int(Sq), + int(K_topk), + ) + dq_out = dq_fp32.to(q.dtype) + dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None + return dq_out, dkl_t, dvl_t, dg_t, dsink_out + + if use_fly_dkv: + # Run the FULL kernel; we keep dk/dv/dgathered/dsink from FlyDSL and + # discard its dq (Triton's dq is the reference here). + dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dkl_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dvl_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dgathered_fp32 = torch.zeros((B, Sq, K_topk, D), device=q.device, dtype=torch.float32) + dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32) + + launch = _get_fly_csa_full( + num_heads=HQ, + head_dim=D, + swa_window=int(swa_window), + dtype_str=dtype_str, + has_sink=has_sink, + has_sparse=has_sparse, + block_n=block_n, + block_k=block_k, + ) + + if verbose: + print( + f"[v4_csa_bwd] FlyDSL DKV launch (dq discarded): B={B} HQ={HQ} Sq={Sq} K_topk={K_topk} " + f"D={D} swa={swa_window} has_sink={has_sink} has_sparse={has_sparse}", + flush=True, + ) + + launch( + q_c, + k_mqa, + v_mqa, + gathered_c, + sparse_mask_c, + dout_c, + lse_c, + d_buf, + sink_arg, + dq_fp32, + dkl_fp32, + dvl_fp32, + dgathered_fp32, + dsink_fp32, + int(B), + int(Sq), + int(K_topk), + ) + + dkl_out = dkl_fp32.to(k_local.dtype) + dvl_out = dvl_fp32.to(v_local.dtype) + dg_out = dgathered_fp32.to(gathered.dtype) + dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None + return dq_t, dkl_out, dvl_out, dg_out, dsink_out + + raise RuntimeError("unreachable") + + +__all__ = ["flydsl_v4_csa_attention_bwd"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_dq_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_dq_kernel.py new file mode 100644 index 000000000..a3d72eb19 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_dq_kernel.py @@ -0,0 +1,665 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""v4_csa_bwd_dq: V4 CSA-backward dq kernel for FlyDSL (STEP 3a). + +Per-row design (mirrors v4_csa_fwd_kernel.py): + grid = (Sq, B * HQ) + BLOCK_SIZE = 64 (one wave). Each lane owns D_PER_LANE = HEAD_DIM // 64 + contiguous d values (=8 for D=512). + +Computes the FULL dq for one (b, qhid, q_row) by iterating BOTH + - LOCAL SWA stream (k_local / v_local, SWA-window-causal-mask) + - GATHERED stream (gathered / sparse_mask, top-K=K_topk per row) +The two streams share the saved JOINT LSE from CSA fwd. + +Inputs: + Q [B, HQ, Sq, D] bf16 + K_LOCAL [B, 1, Sk, D] bf16 (MQA, HK=1) + V_LOCAL [B, 1, Sk, D] bf16 + GATHERED [B, Sq, K_topk, D] bf16 + SPARSE_MASK[B, Sq, K_topk] bf16 (additive bias) + DOUT [B, HQ, Sq, D] bf16 + LSE [B, HQ, Sq] fp32 (JOINT LSE, RAW domain) + DELTAS [B, HQ, Sq] fp32 + SINK [HQ] fp32 (dummy if !has_sink) + +Outputs: + DQ [B, HQ, Sq, D] fp32 (overwritten -- one program owns the row) + DSINK [HQ] fp32 (atomic_add) +""" + +from __future__ import annotations + +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import math as math_dialect +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator +from kernels.kernels_common import dtype_to_elem_type + +KERNEL_NAME = "v4_csa_bwd_dq_kernel" +_LOG2E = math.log2(math.e) +_LLVM_GEP_DYNAMIC = -2147483648 + + +def _llvm_ptr_ty(): + return ir.Type.parse("!llvm.ptr") + + +def build_v4_csa_bwd_dq_module( + num_heads, + head_dim, + swa_window, + dtype_str="bf16", + sm_scale=None, + waves_per_eu=2, + block_n=32, + block_k=32, + has_sink=True, + has_sparse=True, + unsafe_fp_math=True, + fast_fp_math=True, + daz=True, + mqa_kv=True, +): + """Build the V4 CSA backward dq launcher (one program per (b, h, q_row)).""" + gpu_arch = get_hip_arch() + WARP_SIZE = 64 + BLOCK_SIZE = WARP_SIZE + BLOCK_N = int(block_n) + BLOCK_K = int(block_k) + NUM_HEADS = int(num_heads) + HEAD_DIM = int(head_dim) + assert HEAD_DIM % WARP_SIZE == 0, f"head_dim must be divisible by {WARP_SIZE}" + D_PER_LANE = HEAD_DIM // WARP_SIZE + assert mqa_kv, "v4_csa_bwd_dq only supports MQA (HK=1)" + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(HEAD_DIM) + + allocator = SmemAllocator( + None, + arch=gpu_arch, + global_sym_name=f"v4_csa_bwd_dq_smem_N{BLOCK_N}_K{BLOCK_K}_S{int(has_sink)}_HS{int(has_sparse)}", + ) + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def v4_csa_bwd_dq_kernel( + Q: fx.Tensor, + K_LOCAL: fx.Tensor, + V_LOCAL: fx.Tensor, + GATHERED: fx.Tensor, + SPARSE_MASK: fx.Tensor, + DOUT: fx.Tensor, + LSE: fx.Tensor, + DELTAS: fx.Tensor, + SINK: fx.Tensor, + DQ: fx.Tensor, + DSINK: fx.Tensor, + seq_len: fx.Int32, + K_topk: fx.Int32, + ): + elem_type = dtype_to_elem_type(dtype_str) + f16_ty = elem_type + f32_ty = T.f32 + fm_fast = arith.FastMathFlags.fast + + q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q) + kl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K_LOCAL) + vl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V_LOCAL) + g_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), GATHERED) + sm_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), SPARSE_MASK) + do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOUT) + dq_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DQ) + lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True) + deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True) + if has_sink: + sink_rsrc = buffer_ops.create_buffer_resource(SINK, max_size=True) + dsink_rsrc = buffer_ops.create_buffer_resource(DSINK, max_size=True) + + def _gep_load(base_ptr, elem_idx, vec_type, elem_t): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=elem_t, + noWrapFlags=0, + ) + return _llvm.LoadOp(vec_type, gep.result).result + + def _gep_store_f32(val, base_ptr, elem_idx): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=T.f32, + noWrapFlags=0, + ) + _llvm.StoreOp(val, gep.result) + + def load_f16_v(base_ptr, elem_idx, n): + vt = T.vec(n, f16_ty) + return _gep_load(base_ptr, elem_idx, vt, f16_ty) + + # ---- Thread / program IDs ---- + pid_m = arith.index_cast(T.index, gpu.block_idx.x) + pid_bh = arith.index_cast(T.index, gpu.block_idx.y) + tid = arith.index_cast(T.index, gpu.thread_idx.x) + lane = tid + + seq_len_v = arith.index_cast(T.index, seq_len) + K_topk_v = arith.index_cast(T.index, K_topk) + + bid = pid_bh // arith.index(NUM_HEADS) + qhid = pid_bh % arith.index(NUM_HEADS) + + q_active = arith.cmpi(arith.CmpIPredicate.slt, pid_m, seq_len_v) + pid_m_safe = arith.select(q_active, pid_m, arith.index(0)) + + NEG_INF_F = -1.0e30 + c_neg_inf = arith.constant(NEG_INF_F, type=f32_ty) + c_zero_f = arith.constant(0.0, type=f32_ty) + c_sm_scale = arith.constant(float(sm_scale), type=f32_ty) + width_i32 = arith.constant(WARP_SIZE, type=T.i32) + + zero_f32_vec = arith.constant_vector(0.0, T.vec(D_PER_LANE, f32_ty)) + q_row_base = ((bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe) * arith.index(HEAD_DIM) + q_lane_off = q_row_base + lane * arith.index(D_PER_LANE) + q_vec_raw = load_f16_v(q_ptr, q_lane_off, D_PER_LANE) + q_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), q_vec_raw) + q_f32 = arith.select(q_active, q_f32, zero_f32_vec) + + do_lane_off = q_lane_off + do_vec_raw = load_f16_v(do_ptr, do_lane_off, D_PER_LANE) + do_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), do_vec_raw) + do_f32 = arith.select(q_active, do_f32, zero_f32_vec) + + lse_delta_off = (bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe + lse_delta_off_i32 = arith.index_cast(T.i32, lse_delta_off) + lse_val = buffer_ops.buffer_load( + lse_rsrc, + lse_delta_off_i32, + vec_width=1, + dtype=f32_ty, + ) + delta_val = buffer_ops.buffer_load( + deltas_rsrc, + lse_delta_off_i32, + vec_width=1, + dtype=f32_ty, + ) + + if has_sink: + qhid_i32 = arith.index_cast(T.i32, qhid) + sink_h = buffer_ops.buffer_load( + sink_rsrc, + qhid_i32, + vec_width=1, + dtype=f32_ty, + ) + sink_h = rocdl.readfirstlane(f32_ty, sink_h) + sub_sh = arith.SubFOp(sink_h, lse_val, fastmath=fm_fast).result + p_sink = math_dialect.exp(sub_sh, fastmath=fm_fast) + neg_p_sink = arith.SubFOp(c_zero_f, p_sink, fastmath=fm_fast).result + dsink_contrib = arith.MulFOp( + neg_p_sink, + delta_val, + fastmath=fm_fast, + ).result + is_lane0 = arith.cmpi(arith.CmpIPredicate.eq, lane, arith.index(0)) + do_sink = arith.AndIOp(is_lane0, q_active).result + _if_sink = scf.IfOp(do_sink, [], has_else=False) + with ir.InsertionPoint(_if_sink.then_block): + _dsink_byte_off = arith.MulIOp( + qhid_i32, + arith.constant(4, type=T.i32), + ).result + _zero_i32_atom = arith.constant(0, type=T.i32) + rocdl.raw_ptr_buffer_atomic_fadd( + dsink_contrib, + dsink_rsrc, + _dsink_byte_off, + _zero_i32_atom, + _zero_i32_atom, + ) + scf.YieldOp([]) + + def warp_reduce_sum_f32(v): + cur = v + for off in [32, 16, 8, 4, 2, 1]: + xor_amt = arith.constant(off, type=T.i32) + peer = arith.ArithValue(cur).shuffle_xor(xor_amt, width_i32) + cur = arith.AddFOp(cur, peer, fastmath=fm_fast).result + return cur + + def vec_dot_f32(a_vec, b_vec): + s = c_zero_f + for i in range_constexpr(D_PER_LANE): + av = vector.extract(a_vec, static_position=[i], dynamic_position=[]) + bv = vector.extract(b_vec, static_position=[i], dynamic_position=[]) + p = arith.MulFOp(av, bv, fastmath=fm_fast).result + s = arith.AddFOp(s, p, fastmath=fm_fast).result + return s + + # ---- LOCAL SWA bounds ---- + _pid_p1 = pid_m + arith.index(1) + _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _pid_p1, seq_len_v) + n_loop_end_row = arith.select(_le_seq, _pid_p1, seq_len_v) + SWA = arith.index(int(swa_window)) + _ge_w = arith.cmpi(arith.CmpIPredicate.sge, _pid_p1, SWA) + _n_lo_raw = arith.select(_ge_w, _pid_p1 - SWA, arith.index(0)) + BN_idx = arith.index(BLOCK_N) + n_loop_start = (_n_lo_raw // BN_idx) * BN_idx + n_loop_end_blk = ((n_loop_end_row + BN_idx - arith.index(1)) // BN_idx) * BN_idx + + # Pad iter_args to >= 2 elements so scf.for_ always yields a tuple + # (the singleton path auto-unwraps to a scalar, breaking subscripting + # at D_PER_LANE==1 / D=64). + _PAD = 1 if D_PER_LANE == 1 else 0 + init_local = [] + for _ in range_constexpr(D_PER_LANE): + init_local.append(c_zero_f) + for _ in range_constexpr(_PAD): + init_local.append(c_zero_f) + + pid_m_i32 = arith.index_cast(T.i32, pid_m) + seq_len_i32 = arith.index_cast(T.i32, seq_len_v) + w_i32 = arith.constant(int(swa_window), type=T.i32) + K_topk_i32 = arith.index_cast(T.i32, K_topk_v) + + # ==== LOCAL SWA loop ==== + for n_start, inner_args, loop_results_local in scf.for_( + n_loop_start, + n_loop_end_blk, + BN_idx, + iter_args=init_local, + ): + dq_accs = [inner_args[d] for d in range_constexpr(D_PER_LANE)] + + n_start_i32 = arith.index_cast(T.i32, n_start) + + kl_f32_cache = [] + p_cache = [] + dp_cache = [] + + for n_off in range_constexpr(BLOCK_N): + kv_col_i32 = arith.AddIOp( + n_start_i32, + arith.constant(n_off, type=T.i32), + ).result + _kv_plus_w = arith.AddIOp(kv_col_i32, w_i32).result + is_swa = arith.cmpi( + arith.CmpIPredicate.sle, + _kv_plus_w, + pid_m_i32, + ) + is_causal = arith.cmpi( + arith.CmpIPredicate.sgt, + kv_col_i32, + pid_m_i32, + ) + is_oob = arith.cmpi( + arith.CmpIPredicate.sge, + kv_col_i32, + seq_len_i32, + ) + bad = arith.OrIOp( + arith.OrIOp(is_causal, is_swa).result, + is_oob, + ).result + + kv_col_idx = arith.index_cast(T.index, kv_col_i32) + kv_col_safe = arith.select(is_oob, arith.index(0), kv_col_idx) + kl_row_base = (bid * seq_len_v + kv_col_safe) * arith.index(HEAD_DIM) + kl_lane_off = kl_row_base + lane * arith.index(D_PER_LANE) + kl_vec = load_f16_v(kl_ptr, kl_lane_off, D_PER_LANE) + kl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), kl_vec) + kl_f32_cache.append(kl_f32) + + vl_vec = load_f16_v(vl_ptr, kl_lane_off, D_PER_LANE) + vl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), vl_vec) + + lane_dot_qk = vec_dot_f32(q_f32, kl_f32) + qk_full = warp_reduce_sum_f32(lane_dot_qk) + qk_scaled = arith.MulFOp( + qk_full, + c_sm_scale, + fastmath=fm_fast, + ).result + qk_masked = arith.select(bad, c_neg_inf, qk_scaled) + diff_qk = arith.SubFOp( + qk_masked, + lse_val, + fastmath=fm_fast, + ).result + p = math_dialect.exp(diff_qk, fastmath=fm_fast) + p_cache.append(p) + + lane_dot_dp = vec_dot_f32(do_f32, vl_f32) + dp_full = warp_reduce_sum_f32(lane_dot_dp) + dp_cache.append(dp_full) + + for n_off in range_constexpr(BLOCK_N): + p = p_cache[n_off] + dp = dp_cache[n_off] + diff = arith.SubFOp(dp, delta_val, fastmath=fm_fast).result + ds = arith.MulFOp(p, diff, fastmath=fm_fast).result + ds_scaled = arith.MulFOp( + ds, + c_sm_scale, + fastmath=fm_fast, + ).result + kl_f32 = kl_f32_cache[n_off] + for d_off in range_constexpr(D_PER_LANE): + klv = vector.extract( + kl_f32, + static_position=[d_off], + dynamic_position=[], + ) + contrib = arith.MulFOp( + ds_scaled, + klv, + fastmath=fm_fast, + ).result + dq_accs[d_off] = arith.AddFOp( + dq_accs[d_off], + contrib, + fastmath=fm_fast, + ).result + + _yield = list(dq_accs) + for _ in range_constexpr(_PAD): + _yield.append(c_zero_f) + yield _yield + + dq_accs = [loop_results_local[d] for d in range_constexpr(D_PER_LANE)] + + # ==== GATHERED branch ==== + if has_sparse: + init_sparse = list(dq_accs) + for _ in range_constexpr(_PAD): + init_sparse.append(c_zero_f) + for k_start, inner_args_g, loop_results_g in scf.for_( + arith.index(0), + K_topk_v, + arith.index(BLOCK_K), + iter_args=init_sparse, + ): + dq_accs_g = [inner_args_g[d] for d in range_constexpr(D_PER_LANE)] + + k_start_i32 = arith.index_cast(T.i32, k_start) + + g_f32_cache = [] + p_cache_g = [] + dp_cache_g = [] + + for k_off in range_constexpr(BLOCK_K): + k_pos_i32 = arith.AddIOp( + k_start_i32, + arith.constant(k_off, type=T.i32), + ).result + is_oob = arith.cmpi( + arith.CmpIPredicate.sge, + k_pos_i32, + K_topk_i32, + ) + k_pos_idx = arith.index_cast(T.index, k_pos_i32) + k_pos_safe = arith.select(is_oob, arith.index(0), k_pos_idx) + + g_row_base = ((bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe) * arith.index( + HEAD_DIM + ) + g_lane_off = g_row_base + lane * arith.index(D_PER_LANE) + g_vec = load_f16_v(g_ptr, g_lane_off, D_PER_LANE) + g_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), g_vec) + g_f32_cache.append(g_f32) + + sm_off = (bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe + sm_raw_v1 = _gep_load( + sm_ptr, + sm_off, + T.vec(1, elem_type), + elem_type, + ) + sm_raw = vector.extract( + sm_raw_v1, + static_position=[0], + dynamic_position=[], + ) + sm_val = arith.extf(f32_ty, sm_raw) + sm_val = arith.select(is_oob, c_zero_f, sm_val) + + lane_dot_qk = vec_dot_f32(q_f32, g_f32) + qk_full = warp_reduce_sum_f32(lane_dot_qk) + qk_scaled = arith.MulFOp( + qk_full, + c_sm_scale, + fastmath=fm_fast, + ).result + qk_biased = arith.AddFOp( + qk_scaled, + sm_val, + fastmath=fm_fast, + ).result + bad = arith.OrIOp( + is_oob, + arith.cmpi( + arith.CmpIPredicate.sge, + pid_m_i32, + seq_len_i32, + ), + ).result + qk_masked = arith.select(bad, c_neg_inf, qk_biased) + diff_qk = arith.SubFOp( + qk_masked, + lse_val, + fastmath=fm_fast, + ).result + p = math_dialect.exp(diff_qk, fastmath=fm_fast) + p_cache_g.append(p) + + lane_dot_dp = vec_dot_f32(do_f32, g_f32) + dp_full = warp_reduce_sum_f32(lane_dot_dp) + dp_cache_g.append(dp_full) + + for k_off in range_constexpr(BLOCK_K): + p = p_cache_g[k_off] + dp = dp_cache_g[k_off] + diff = arith.SubFOp(dp, delta_val, fastmath=fm_fast).result + ds = arith.MulFOp(p, diff, fastmath=fm_fast).result + ds_scaled = arith.MulFOp( + ds, + c_sm_scale, + fastmath=fm_fast, + ).result + g_f32 = g_f32_cache[k_off] + for d_off in range_constexpr(D_PER_LANE): + gv = vector.extract( + g_f32, + static_position=[d_off], + dynamic_position=[], + ) + contrib = arith.MulFOp( + ds_scaled, + gv, + fastmath=fm_fast, + ).result + dq_accs_g[d_off] = arith.AddFOp( + dq_accs_g[d_off], + contrib, + fastmath=fm_fast, + ).result + + _yield_g = list(dq_accs_g) + for _ in range_constexpr(_PAD): + _yield_g.append(c_zero_f) + yield _yield_g + + dq_accs = [loop_results_g[d] for d in range_constexpr(D_PER_LANE)] + + # ==== Store dq[d] into DQ buffer (fp32 direct store) ==== + _o_guard = scf.IfOp(q_active, [], has_else=False) + with ir.InsertionPoint(_o_guard.then_block): + dq_row_base = ((bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe) * arith.index( + HEAD_DIM + ) + dq_lane_off = dq_row_base + lane * arith.index(D_PER_LANE) + for d_off in range_constexpr(D_PER_LANE): + elem_off = dq_lane_off + arith.index(d_off) + _gep_store_f32(dq_accs[d_off], dq_ptr, elem_off) + scf.YieldOp([]) + + @flyc.jit + def launch_v4_csa_bwd_dq( + Q: fx.Tensor, + K_LOCAL: fx.Tensor, + V_LOCAL: fx.Tensor, + GATHERED: fx.Tensor, + SPARSE_MASK: fx.Tensor, + DOUT: fx.Tensor, + LSE: fx.Tensor, + DELTAS: fx.Tensor, + SINK: fx.Tensor, + DQ: fx.Tensor, + DSINK: fx.Tensor, + batch_size: fx.Int32, + seq_len: fx.Int32, + K_topk: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + bs_idx = arith.index_cast(T.index, batch_size) + sl_idx = arith.index_cast(T.index, seq_len) + grid_x = sl_idx + grid_y = bs_idx * arith.index(NUM_HEADS) + + launcher = v4_csa_bwd_dq_kernel( + Q, + K_LOCAL, + V_LOCAL, + GATHERED, + SPARSE_MASK, + DOUT, + LSE, + DELTAS, + SINK, + DQ, + DSINK, + seq_len, + K_topk, + ) + + if waves_per_eu is not None: + _wpe = int(waves_per_eu) + if _wpe >= 1: + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, + _wpe, + ) + + passthrough_entries = [] + if daz: + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("denormal-fp-math-f32"), + ir.StringAttr.get("preserve-sign,preserve-sign"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("no-nans-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("unsafe-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries) + + launcher.launch( + grid=(grid_x, grid_y, 1), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) + + compile_hints = { + "fast_fp_math": fast_fp_math, + "unsafe_fp_math": unsafe_fp_math, + } + + def _launch(*args, **kwargs): + with CompilationContext.compile_hints(compile_hints): + return launch_v4_csa_bwd_dq(*args, **kwargs) + + def _compile( + Q, + K_LOCAL, + V_LOCAL, + GATHERED, + SPARSE_MASK, + DOUT, + LSE, + DELTAS, + SINK, + DQ, + DSINK, + batch_size, + seq_len, + K_topk, + stream=None, + ): + with CompilationContext.compile_hints(compile_hints): + return flyc.compile( + launch_v4_csa_bwd_dq, + Q, + K_LOCAL, + V_LOCAL, + GATHERED, + SPARSE_MASK, + DOUT, + LSE, + DELTAS, + SINK, + DQ, + DSINK, + batch_size, + seq_len, + K_topk, + fx.Stream(stream), + ) + + _launch.compile = _compile + return _launch + + +build_v4_csa_bwd_dq_module_primary = build_v4_csa_bwd_dq_module diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_full_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_full_kernel.py new file mode 100644 index 000000000..9f702010a --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_bwd_full_kernel.py @@ -0,0 +1,664 @@ +# SPDX-License-Identifier: Apache-2.0 +"""v4_csa_bwd_full: V4 CSA backward kernel, full output set (STEP 3b). + +Emits dq + dk_local + dv_local + dgathered + dsink in one launch. +Mirrors `_v4_csa_attention_bwd_kernel` (Triton, _triton/v4_csa_attention_bwd.py) +1:1: grid=(Sq, B*HQ); each program owns one query row. dq is direct-stored; +dk_local, dv_local, dgathered, dsink are accumulated via atomic_fadd. +""" +from __future__ import annotations + +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import math as math_dialect +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import ( + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + vector, +) +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator +from kernels.kernels_common import dtype_to_elem_type + +KERNEL_NAME = "v4_csa_bwd_full_kernel" +_LLVM_GEP_DYNAMIC = -2147483648 + + +def _llvm_ptr_ty(): + return ir.Type.parse("!llvm.ptr") + + +def build_v4_csa_bwd_full_module( + num_heads, + head_dim, + swa_window, + dtype_str="bf16", + sm_scale=None, + waves_per_eu=2, + block_n=32, + block_k=32, + has_sink=True, + has_sparse=True, + unsafe_fp_math=True, + fast_fp_math=True, + daz=True, + mqa_kv=True, +): + gpu_arch = get_hip_arch() + WARP_SIZE = 64 + BLOCK_SIZE = WARP_SIZE + BLOCK_N = int(block_n) + BLOCK_K = int(block_k) + NUM_HEADS = int(num_heads) + HEAD_DIM = int(head_dim) + assert HEAD_DIM % WARP_SIZE == 0, f"head_dim must be divisible by {WARP_SIZE}" + D_PER_LANE = HEAD_DIM // WARP_SIZE + assert mqa_kv, "v4_csa_bwd_full only supports MQA (HK=1)" + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(HEAD_DIM) + + allocator = SmemAllocator( + None, + arch=gpu_arch, + global_sym_name=f"v4_csa_bwd_full_smem_N{BLOCK_N}_K{BLOCK_K}_S{int(has_sink)}_HS{int(has_sparse)}", + ) + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def v4_csa_bwd_full_kernel( + Q: fx.Tensor, + K_LOCAL: fx.Tensor, + V_LOCAL: fx.Tensor, + GATHERED: fx.Tensor, + SPARSE_MASK: fx.Tensor, + DOUT: fx.Tensor, + LSE: fx.Tensor, + DELTAS: fx.Tensor, + SINK: fx.Tensor, + DQ: fx.Tensor, + DK_LOCAL: fx.Tensor, + DV_LOCAL: fx.Tensor, + DGATHERED: fx.Tensor, + DSINK: fx.Tensor, + seq_len: fx.Int32, + K_topk: fx.Int32, + ): + elem_type = dtype_to_elem_type(dtype_str) + # FlyDSL >=0.2.2 compat: dtype_to_elem_type returns a Numeric meta + # (e.g. fx.BFloat16); the MLIR type-arg sites below (T.vec, GEPOp, + # SmemPtr, trunc_f) require an ir.Type. Coerce once here. + if hasattr(elem_type, "ir_type"): + elem_type = elem_type.ir_type + f16_ty = elem_type + f32_ty = T.f32 + fm_fast = arith.FastMathFlags.fast + + q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q) + kl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K_LOCAL) + vl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V_LOCAL) + g_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), GATHERED) + sm_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), SPARSE_MASK) + do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOUT) + dq_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DQ) + lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True) + deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True) + dkl_rsrc = buffer_ops.create_buffer_resource(DK_LOCAL, max_size=True) + dvl_rsrc = buffer_ops.create_buffer_resource(DV_LOCAL, max_size=True) + dg_rsrc = buffer_ops.create_buffer_resource(DGATHERED, max_size=True) + if const_expr(has_sink): + sink_rsrc = buffer_ops.create_buffer_resource(SINK, max_size=True) + dsink_rsrc = buffer_ops.create_buffer_resource(DSINK, max_size=True) + + def _gep_load(base_ptr, elem_idx, vec_type, elem_t): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=elem_t, + noWrapFlags=0, + ) + return _llvm.LoadOp(vec_type, gep.result).result + + def _gep_store_f32(val, base_ptr, elem_idx): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=T.f32, + noWrapFlags=0, + ) + _llvm.StoreOp(val, gep.result) + + def load_f16_v(base_ptr, elem_idx, n): + vt = T.vec(n, f16_ty) + return _gep_load(base_ptr, elem_idx, vt, f16_ty) + + pid_m = arith.index_cast(T.index, gpu.block_idx.x) + pid_bh = arith.index_cast(T.index, gpu.block_idx.y) + tid = arith.index_cast(T.index, gpu.thread_idx.x) + lane = tid + + seq_len_v = arith.index_cast(T.index, seq_len) + K_topk_v = arith.index_cast(T.index, K_topk) + + bid = pid_bh // arith.index(NUM_HEADS) + qhid = pid_bh % arith.index(NUM_HEADS) + + q_active = arith.cmpi(arith.CmpIPredicate.slt, pid_m, seq_len_v) + pid_m_safe = arith.select(q_active, pid_m, arith.index(0)) + + NEG_INF_F = -1.0e30 + c_neg_inf = arith.constant(NEG_INF_F, type=f32_ty) + c_zero_f = arith.constant(0.0, type=f32_ty) + c_sm_scale = arith.constant(float(sm_scale), type=f32_ty) + width_i32 = arith.constant(WARP_SIZE, type=T.i32) + c_four_i32 = arith.constant(4, type=T.i32) + c_zero_i32 = arith.constant(0, type=T.i32) + + zero_f32_vec = arith.constant_vector(0.0, T.vec(D_PER_LANE, f32_ty)) + q_row_base = ((bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe) * arith.index(HEAD_DIM) + q_lane_off = q_row_base + lane * arith.index(D_PER_LANE) + q_vec_raw = load_f16_v(q_ptr, q_lane_off, D_PER_LANE) + q_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), q_vec_raw) + q_f32 = arith.select(q_active, q_f32, zero_f32_vec) + + do_lane_off = q_lane_off + do_vec_raw = load_f16_v(do_ptr, do_lane_off, D_PER_LANE) + do_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), do_vec_raw) + do_f32 = arith.select(q_active, do_f32, zero_f32_vec) + + lse_delta_off = (bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe + lse_delta_off_i32 = arith.index_cast(T.i32, lse_delta_off) + lse_val = buffer_ops.buffer_load( + lse_rsrc, + lse_delta_off_i32, + vec_width=1, + dtype=f32_ty, + ) + delta_val = buffer_ops.buffer_load( + deltas_rsrc, + lse_delta_off_i32, + vec_width=1, + dtype=f32_ty, + ) + + qhid_i32 = arith.index_cast(T.i32, qhid) + bid_i32 = arith.index_cast(T.i32, bid) + lane_i32 = arith.index_cast(T.i32, lane) + pid_m_safe_i32 = arith.index_cast(T.i32, pid_m_safe) + head_dim_i32 = arith.constant(HEAD_DIM, type=T.i32) + d_per_lane_i32 = arith.constant(D_PER_LANE, type=T.i32) + num_heads_i32 = arith.constant(NUM_HEADS, type=T.i32) + + if const_expr(has_sink): + sink_h = buffer_ops.buffer_load( + sink_rsrc, + qhid_i32, + vec_width=1, + dtype=f32_ty, + ) + sink_h = rocdl.readfirstlane(f32_ty, sink_h) + sub_sh = arith.SubFOp(sink_h, lse_val, fastmath=fm_fast).result + p_sink = math_dialect.exp(sub_sh, fastmath=fm_fast) + neg_p_sink = arith.SubFOp(c_zero_f, p_sink, fastmath=fm_fast).result + dsink_contrib = arith.MulFOp( + neg_p_sink, + delta_val, + fastmath=fm_fast, + ).result + is_lane0 = arith.cmpi(arith.CmpIPredicate.eq, lane, arith.index(0)) + do_sink = arith.AndIOp(is_lane0, q_active).result + _if_sink = scf.IfOp(do_sink, [], has_else=False) + with ir.InsertionPoint(_if_sink.then_block): + _dsink_byte_off = arith.MulIOp(qhid_i32, c_four_i32).result + rocdl.raw_ptr_buffer_atomic_fadd( + dsink_contrib, + dsink_rsrc, + _dsink_byte_off, + c_zero_i32, + c_zero_i32, + ) + scf.YieldOp([]) + + def warp_reduce_sum_f32(v): + cur = v + for off in [32, 16, 8, 4, 2, 1]: + xor_amt = arith.constant(off, type=T.i32) + peer = arith.ArithValue(cur).shuffle_xor(xor_amt, width_i32) + cur = arith.AddFOp(cur, peer, fastmath=fm_fast).result + return cur + + def vec_dot_f32(a_vec, b_vec): + s = c_zero_f + for i in range_constexpr(D_PER_LANE): + av = vector.extract(a_vec, static_position=[i], dynamic_position=[]) + bv = vector.extract(b_vec, static_position=[i], dynamic_position=[]) + p = arith.MulFOp(av, bv, fastmath=fm_fast).result + s = arith.AddFOp(s, p, fastmath=fm_fast).result + return s + + _pid_p1 = pid_m + arith.index(1) + _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _pid_p1, seq_len_v) + n_loop_end_row = arith.select(_le_seq, _pid_p1, seq_len_v) + SWA = arith.index(int(swa_window)) + _ge_w = arith.cmpi(arith.CmpIPredicate.sge, _pid_p1, SWA) + _n_lo_raw = arith.select(_ge_w, _pid_p1 - SWA, arith.index(0)) + BN_idx = arith.index(BLOCK_N) + n_loop_start = (_n_lo_raw // BN_idx) * BN_idx + n_loop_end_blk = ((n_loop_end_row + BN_idx - arith.index(1)) // BN_idx) * BN_idx + + _PAD = 1 if D_PER_LANE == 1 else 0 + init_local = [] + for _ in range_constexpr(D_PER_LANE): + init_local.append(c_zero_f) + for _ in range_constexpr(_PAD): + init_local.append(c_zero_f) + + pid_m_i32 = arith.index_cast(T.i32, pid_m) + seq_len_i32 = arith.index_cast(T.i32, seq_len_v) + w_i32 = arith.constant(int(swa_window), type=T.i32) + K_topk_i32 = arith.index_cast(T.i32, K_topk_v) + + # ==== LOCAL SWA loop ==== + for n_start, inner_args, loop_results_local in scf.for_( + n_loop_start, + n_loop_end_blk, + BN_idx, + iter_args=init_local, + ): + dq_accs = [inner_args[d] for d in range_constexpr(D_PER_LANE)] + n_start_i32 = arith.index_cast(T.i32, n_start) + + kl_f32_cache = [] + p_cache = [] + dp_cache = [] + kv_col_i32_cache = [] + + for n_off in range_constexpr(BLOCK_N): + kv_col_i32 = arith.AddIOp( + n_start_i32, + arith.constant(n_off, type=T.i32), + ).result + kv_col_i32_cache.append(kv_col_i32) + _kv_plus_w = arith.AddIOp(kv_col_i32, w_i32).result + is_swa = arith.cmpi(arith.CmpIPredicate.sle, _kv_plus_w, pid_m_i32) + is_causal = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_i32, pid_m_i32) + is_oob = arith.cmpi(arith.CmpIPredicate.sge, kv_col_i32, seq_len_i32) + bad = arith.OrIOp( + arith.OrIOp(is_causal, is_swa).result, + is_oob, + ).result + + kv_col_idx = arith.index_cast(T.index, kv_col_i32) + kv_col_safe = arith.select(is_oob, arith.index(0), kv_col_idx) + kl_row_base = (bid * seq_len_v + kv_col_safe) * arith.index(HEAD_DIM) + kl_lane_off = kl_row_base + lane * arith.index(D_PER_LANE) + kl_vec = load_f16_v(kl_ptr, kl_lane_off, D_PER_LANE) + kl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), kl_vec) + kl_f32_cache.append(kl_f32) + + vl_vec = load_f16_v(vl_ptr, kl_lane_off, D_PER_LANE) + vl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), vl_vec) + + lane_dot_qk = vec_dot_f32(q_f32, kl_f32) + qk_full = warp_reduce_sum_f32(lane_dot_qk) + qk_scaled = arith.MulFOp(qk_full, c_sm_scale, fastmath=fm_fast).result + qk_masked = arith.select(bad, c_neg_inf, qk_scaled) + diff_qk = arith.SubFOp(qk_masked, lse_val, fastmath=fm_fast).result + p = math_dialect.exp(diff_qk, fastmath=fm_fast) + p_cache.append(p) + + lane_dot_dp = vec_dot_f32(do_f32, vl_f32) + dp_full = warp_reduce_sum_f32(lane_dot_dp) + dp_cache.append(dp_full) + + for n_off in range_constexpr(BLOCK_N): + p = p_cache[n_off] + dp = dp_cache[n_off] + diff = arith.SubFOp(dp, delta_val, fastmath=fm_fast).result + ds = arith.MulFOp(p, diff, fastmath=fm_fast).result + ds_scaled = arith.MulFOp(ds, c_sm_scale, fastmath=fm_fast).result + kl_f32 = kl_f32_cache[n_off] + + # dq accumulator + for d_off in range_constexpr(D_PER_LANE): + klv = vector.extract(kl_f32, static_position=[d_off], dynamic_position=[]) + contrib = arith.MulFOp(ds_scaled, klv, fastmath=fm_fast).result + dq_accs[d_off] = arith.AddFOp(dq_accs[d_off], contrib, fastmath=fm_fast).result + + # dk_local / dv_local atomic_add + kv_col_i32 = kv_col_i32_cache[n_off] + in_range = arith.cmpi(arith.CmpIPredicate.slt, kv_col_i32, seq_len_i32) + do_atom = arith.AndIOp(in_range, q_active).result + _if_dkv = scf.IfOp(do_atom, [], has_else=False) + with ir.InsertionPoint(_if_dkv.then_block): + _bh = arith.AddIOp( + arith.MulIOp(bid_i32, num_heads_i32).result, + qhid_i32, + ).result + _bh_n = arith.AddIOp( + arith.MulIOp(_bh, seq_len_i32).result, + kv_col_i32, + ).result + _row_d = arith.AddIOp( + arith.MulIOp(_bh_n, head_dim_i32).result, + arith.MulIOp(lane_i32, d_per_lane_i32).result, + ).result + for d_off in range_constexpr(D_PER_LANE): + elem_i32 = arith.AddIOp(_row_d, arith.constant(d_off, type=T.i32)).result + byte_off = arith.MulIOp(elem_i32, c_four_i32).result + qv = vector.extract(q_f32, static_position=[d_off], dynamic_position=[]) + dk_val = arith.MulFOp(ds_scaled, qv, fastmath=fm_fast).result + rocdl.raw_ptr_buffer_atomic_fadd( + dk_val, + dkl_rsrc, + byte_off, + c_zero_i32, + c_zero_i32, + ) + dov = vector.extract(do_f32, static_position=[d_off], dynamic_position=[]) + dv_val = arith.MulFOp(p, dov, fastmath=fm_fast).result + rocdl.raw_ptr_buffer_atomic_fadd( + dv_val, + dvl_rsrc, + byte_off, + c_zero_i32, + c_zero_i32, + ) + scf.YieldOp([]) + + _yield = list(dq_accs) + for _ in range_constexpr(_PAD): + _yield.append(c_zero_f) + yield _yield + + dq_accs = [loop_results_local[d] for d in range_constexpr(D_PER_LANE)] + + # ==== GATHERED branch ==== + if const_expr(has_sparse): + init_sparse = list(dq_accs) + for _ in range_constexpr(_PAD): + init_sparse.append(c_zero_f) + for k_start, inner_args_g, loop_results_g in scf.for_( + arith.index(0), + K_topk_v, + arith.index(BLOCK_K), + iter_args=init_sparse, + ): + dq_accs_g = [inner_args_g[d] for d in range_constexpr(D_PER_LANE)] + k_start_i32 = arith.index_cast(T.i32, k_start) + + g_f32_cache = [] + p_cache_g = [] + dp_cache_g = [] + k_pos_i32_cache = [] + + for k_off in range_constexpr(BLOCK_K): + k_pos_i32 = arith.AddIOp( + k_start_i32, + arith.constant(k_off, type=T.i32), + ).result + k_pos_i32_cache.append(k_pos_i32) + is_oob = arith.cmpi(arith.CmpIPredicate.sge, k_pos_i32, K_topk_i32) + k_pos_idx = arith.index_cast(T.index, k_pos_i32) + k_pos_safe = arith.select(is_oob, arith.index(0), k_pos_idx) + + g_row_base = ((bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe) * arith.index( + HEAD_DIM + ) + g_lane_off = g_row_base + lane * arith.index(D_PER_LANE) + g_vec = load_f16_v(g_ptr, g_lane_off, D_PER_LANE) + g_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), g_vec) + g_f32_cache.append(g_f32) + + sm_off = (bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe + sm_raw_v1 = _gep_load(sm_ptr, sm_off, T.vec(1, elem_type), elem_type) + sm_raw = vector.extract(sm_raw_v1, static_position=[0], dynamic_position=[]) + sm_val = arith.extf(f32_ty, sm_raw) + sm_val = arith.select(is_oob, c_zero_f, sm_val) + + lane_dot_qk = vec_dot_f32(q_f32, g_f32) + qk_full = warp_reduce_sum_f32(lane_dot_qk) + qk_scaled = arith.MulFOp(qk_full, c_sm_scale, fastmath=fm_fast).result + qk_biased = arith.AddFOp(qk_scaled, sm_val, fastmath=fm_fast).result + bad = arith.OrIOp( + is_oob, + arith.cmpi(arith.CmpIPredicate.sge, pid_m_i32, seq_len_i32), + ).result + qk_masked = arith.select(bad, c_neg_inf, qk_biased) + diff_qk = arith.SubFOp(qk_masked, lse_val, fastmath=fm_fast).result + p = math_dialect.exp(diff_qk, fastmath=fm_fast) + p_cache_g.append(p) + + lane_dot_dp = vec_dot_f32(do_f32, g_f32) + dp_full = warp_reduce_sum_f32(lane_dot_dp) + dp_cache_g.append(dp_full) + + for k_off in range_constexpr(BLOCK_K): + p = p_cache_g[k_off] + dp = dp_cache_g[k_off] + diff = arith.SubFOp(dp, delta_val, fastmath=fm_fast).result + ds = arith.MulFOp(p, diff, fastmath=fm_fast).result + ds_scaled = arith.MulFOp(ds, c_sm_scale, fastmath=fm_fast).result + g_f32 = g_f32_cache[k_off] + + for d_off in range_constexpr(D_PER_LANE): + gv = vector.extract(g_f32, static_position=[d_off], dynamic_position=[]) + contrib = arith.MulFOp(ds_scaled, gv, fastmath=fm_fast).result + dq_accs_g[d_off] = arith.AddFOp(dq_accs_g[d_off], contrib, fastmath=fm_fast).result + + k_pos_i32 = k_pos_i32_cache[k_off] + in_range_k = arith.cmpi(arith.CmpIPredicate.slt, k_pos_i32, K_topk_i32) + do_atom_k = arith.AndIOp(in_range_k, q_active).result + _if_dg = scf.IfOp(do_atom_k, [], has_else=False) + with ir.InsertionPoint(_if_dg.then_block): + _bm = arith.AddIOp( + arith.MulIOp(bid_i32, seq_len_i32).result, + pid_m_safe_i32, + ).result + _bm_k = arith.AddIOp( + arith.MulIOp(_bm, K_topk_i32).result, + k_pos_i32, + ).result + _row_d = arith.AddIOp( + arith.MulIOp(_bm_k, head_dim_i32).result, + arith.MulIOp(lane_i32, d_per_lane_i32).result, + ).result + for d_off in range_constexpr(D_PER_LANE): + elem_i32 = arith.AddIOp(_row_d, arith.constant(d_off, type=T.i32)).result + byte_off = arith.MulIOp(elem_i32, c_four_i32).result + qv = vector.extract(q_f32, static_position=[d_off], dynamic_position=[]) + dov = vector.extract(do_f32, static_position=[d_off], dynamic_position=[]) + t1 = arith.MulFOp(ds_scaled, qv, fastmath=fm_fast).result + t2 = arith.MulFOp(p, dov, fastmath=fm_fast).result + dg_val = arith.AddFOp(t1, t2, fastmath=fm_fast).result + rocdl.raw_ptr_buffer_atomic_fadd( + dg_val, + dg_rsrc, + byte_off, + c_zero_i32, + c_zero_i32, + ) + scf.YieldOp([]) + + _yield_g = list(dq_accs_g) + for _ in range_constexpr(_PAD): + _yield_g.append(c_zero_f) + yield _yield_g + + dq_accs = [loop_results_g[d] for d in range_constexpr(D_PER_LANE)] + + # ==== Store dq direct ==== + _o_guard = scf.IfOp(q_active, [], has_else=False) + with ir.InsertionPoint(_o_guard.then_block): + dq_row_base = ((bid * arith.index(NUM_HEADS) + qhid) * seq_len_v + pid_m_safe) * arith.index( + HEAD_DIM + ) + dq_lane_off = dq_row_base + lane * arith.index(D_PER_LANE) + for d_off in range_constexpr(D_PER_LANE): + elem_off = dq_lane_off + arith.index(d_off) + _gep_store_f32(dq_accs[d_off], dq_ptr, elem_off) + scf.YieldOp([]) + + @flyc.jit + def launch_v4_csa_bwd_full( + Q: fx.Tensor, + K_LOCAL: fx.Tensor, + V_LOCAL: fx.Tensor, + GATHERED: fx.Tensor, + SPARSE_MASK: fx.Tensor, + DOUT: fx.Tensor, + LSE: fx.Tensor, + DELTAS: fx.Tensor, + SINK: fx.Tensor, + DQ: fx.Tensor, + DK_LOCAL: fx.Tensor, + DV_LOCAL: fx.Tensor, + DGATHERED: fx.Tensor, + DSINK: fx.Tensor, + batch_size: fx.Int32, + seq_len: fx.Int32, + K_topk: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + bs_idx = arith.index_cast(T.index, batch_size) + sl_idx = arith.index_cast(T.index, seq_len) + grid_x = sl_idx + grid_y = bs_idx * arith.index(NUM_HEADS) + + launcher = v4_csa_bwd_full_kernel( + Q, + K_LOCAL, + V_LOCAL, + GATHERED, + SPARSE_MASK, + DOUT, + LSE, + DELTAS, + SINK, + DQ, + DK_LOCAL, + DV_LOCAL, + DGATHERED, + DSINK, + seq_len, + K_topk, + ) + + if waves_per_eu is not None: + _wpe = int(waves_per_eu) + if _wpe >= 1: + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe) + + passthrough_entries = [] + if daz: + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("denormal-fp-math-f32"), + ir.StringAttr.get("preserve-sign,preserve-sign"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("no-nans-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("unsafe-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries) + + launcher.launch( + grid=(grid_x, grid_y, 1), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) + + compile_hints = {"fast_fp_math": fast_fp_math, "unsafe_fp_math": unsafe_fp_math} + + def _launch(*args, **kwargs): + with CompilationContext.compile_hints(compile_hints): + return launch_v4_csa_bwd_full(*args, **kwargs) + + def _compile( + Q, + K_LOCAL, + V_LOCAL, + GATHERED, + SPARSE_MASK, + DOUT, + LSE, + DELTAS, + SINK, + DQ, + DK_LOCAL, + DV_LOCAL, + DGATHERED, + DSINK, + batch_size, + seq_len, + K_topk, + stream=None, + ): + with CompilationContext.compile_hints(compile_hints): + return flyc.compile( + launch_v4_csa_bwd_full, + Q, + K_LOCAL, + V_LOCAL, + GATHERED, + SPARSE_MASK, + DOUT, + LSE, + DELTAS, + SINK, + DQ, + DK_LOCAL, + DV_LOCAL, + DGATHERED, + DSINK, + batch_size, + seq_len, + K_topk, + fx.Stream(stream), + ) + + _launch.compile = _compile + return _launch + + +build_v4_csa_bwd_full_module_primary = build_v4_csa_bwd_full_module diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_fwd_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_fwd_kernel.py new file mode 100644 index 000000000..b225af233 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_csa_fwd_kernel.py @@ -0,0 +1,748 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""v4_csa_fwd: V4 CSA forward (FlyDSL per-row design). + +Round-3 Step 2b: ports the Triton monolithic CSA forward kernel to FlyDSL +in a 1:1 per-row design. The Triton CSA monolithic kernel uses one program +per (b, h, m), with per-key scalar dot-product. That structure ports cleanly: + + grid = (Sq, B * HQ) + BLOCK_SIZE = 64 (one wave). Lane in 0..63 handles partial D=8. + Online softmax accumulator (m, l, acc) lives in fp32 in-register and is + distributed across lanes (each lane owns D/64 = 8 elements of the + accumulator's D dimension). + +This design does NOT use MFMA -- the per-row scalar dot-product matches +Triton's monolithic kernel exactly (Triton CSA monolithic also uses +``tl.sum(k * q, axis=1)`` not ``tl.dot``). + +Layout: BHLD (Q/K_local/V_local/O all [B, H, Sq, D]). + - Gathered: [B, Sq, K_topk, D] (no H dim -- shared across heads). + - sparse_mask: [B, Sq, K_topk] (no H dim -- broadcasts over H). + - sink: [H] fp32 or None. + - LSE: [B, H, Sq] fp32 (raw-domain: m_final + ln(l_final), since m_final + already includes sm_scale via the online softmax in raw-qk*sm_scale). + +Forward computes the joint online softmax over local_SWA (block-causal, +window=swa_window, keys are k_local) + sparse (keys are gathered) + sink. + +K_topk == 0 supported (kernel skips the sparse loop). +""" + +import math +import os + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import math as math_dialect +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import ( + arith, + buffer_ops, + const_expr, + gpu, + range_constexpr, + rocdl, + vector, +) +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr +from kernels.kernels_common import dtype_to_elem_type + +KERNEL_NAME = "v4_csa_fwd_kernel" + +_LOG2E = math.log2(math.e) + +_LLVM_GEP_DYNAMIC = -2147483648 + + +def _waitcnt_lgkm_0(): + """s_waitcnt lgkmcnt(0): drain LDS/SMEM ops. vmcnt=63, expcnt=7.""" + val = 0xF | (0x7 << 4) | (0 << 8) | (0x3 << 14) + rocdl.s_waitcnt(val) + + +def _llvm_ptr_ty(): + return ir.Type.parse("!llvm.ptr") + + +def _llvm_lds_ptr_ty(): + return ir.Type.parse("!llvm.ptr<3>") + + +def build_v4_csa_fwd_module( + num_heads, + head_dim, + swa_window, + dtype_str="bf16", + sm_scale=None, + waves_per_eu=2, + block_n=32, + block_k=32, + has_sink=False, + has_sparse=True, + unsafe_fp_math=True, + fast_fp_math=True, + daz=True, + mqa_kv=False, + head_group=1, +): + """Build the V4 CSA forward per-row launcher. + + Parameters: + num_heads: int -- H_Q + head_dim: int -- D (must be divisible by 64) + swa_window: int -- SWA window (> 0) + block_n: int -- local-branch tile width (default 32) + block_k: int -- sparse-branch tile width (default 32) + has_sink: bool -- include sink epilogue + has_sparse: bool -- include sparse branch loop (K_topk > 0) + """ + gpu_arch = get_hip_arch() + WARP_SIZE = 64 + BLOCK_SIZE = WARP_SIZE # one wave per program + BLOCK_N = int(block_n) + BLOCK_K = int(block_k) + NUM_HEADS = int(num_heads) + HEAD_DIM = int(head_dim) + assert HEAD_DIM % WARP_SIZE == 0, f"head_dim must be divisible by {WARP_SIZE}" + D_PER_LANE = HEAD_DIM // WARP_SIZE # 8 for D=512 + HEAD_GROUP = int(head_group) + assert NUM_HEADS % HEAD_GROUP == 0, f"num_heads {NUM_HEADS} must be divisible by head_group {HEAD_GROUP}" + NUM_HEAD_GROUPS = NUM_HEADS // HEAD_GROUP + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(HEAD_DIM) + + # ---- LDS cache for sparse-branch gathered K-tile ---- + ENABLE_LDS_CACHE = bool(has_sparse) and (os.environ.get("PRIMUS_V4_CSA_LDS_CACHE", "0") == "1") + LDS_GATHER_TILE_ELEMS = BLOCK_K * HEAD_DIM + LDS_GATHER_TILE_BYTES = LDS_GATHER_TILE_ELEMS * 2 + + allocator = SmemAllocator( + None, + arch=gpu_arch, + global_sym_name=f"v4_csa_fwd_smem_N{BLOCK_N}_K{BLOCK_K}_C{int(ENABLE_LDS_CACHE)}_HG{HEAD_GROUP}", + ) + if ENABLE_LDS_CACHE: + lds_gather_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_gather_offset + LDS_GATHER_TILE_BYTES + else: + lds_gather_offset = 0 + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def v4_csa_fwd_kernel( + Q: fx.Tensor, + K_LOCAL: fx.Tensor, + V_LOCAL: fx.Tensor, + GATHERED: fx.Tensor, + SPARSE_MASK: fx.Tensor, + Sink: fx.Tensor, + O: fx.Tensor, + LSE: fx.Tensor, + seq_len: fx.Int32, + K_topk: fx.Int32, + ): + elem_type = dtype_to_elem_type(dtype_str) + # FlyDSL >=0.2.2 compat: dtype_to_elem_type returns a Numeric meta + # (e.g. fx.BFloat16); the MLIR type-arg sites below (T.vec, GEPOp, + # SmemPtr, trunc_f) require an ir.Type. Coerce once here. + if hasattr(elem_type, "ir_type"): + elem_type = elem_type.ir_type + T.f32 + fm_fast = arith.FastMathFlags.fast + + q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q) + kl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K_LOCAL) + vl_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V_LOCAL) + g_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), GATHERED) + sm_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), SPARSE_MASK) + o_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), O) + lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True) + if const_expr(has_sink): + sink_rsrc = buffer_ops.create_buffer_resource(Sink, max_size=True) + + f16_ty = elem_type + f32_ty = T.f32 + + # ---- Helpers ---- + def _gep_load_scalar(base_ptr, elem_idx, vec_type, elem_t): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=elem_t, + noWrapFlags=0, + ) + return _llvm.LoadOp(vec_type, gep.result).result + + def load_f16_v(base_ptr, elem_idx, n): + vt = T.vec(n, f16_ty) + return _gep_load_scalar(base_ptr, elem_idx, vt, f16_ty) + + def load_f32_scalar(base_ptr, elem_idx): + return _gep_load_scalar(base_ptr, elem_idx, f32_ty, f32_ty) + + # ---- Thread / program ---- + pid_m = arith.index_cast(T.index, gpu.block_idx.x) + pid_bh = arith.index_cast(T.index, gpu.block_idx.y) + tid = arith.index_cast(T.index, gpu.thread_idx.x) + lane = tid + + seq_len_v = arith.index_cast(T.index, seq_len) + K_topk_v = arith.index_cast(T.index, K_topk) + + bid = pid_bh // arith.index(NUM_HEAD_GROUPS) + qhid_group = pid_bh % arith.index(NUM_HEAD_GROUPS) + # qhid_base is the head index of the first head in this program's head group. + qhid_base = qhid_group * arith.index(HEAD_GROUP) + + # ---- LDS view for sparse-branch gathered cache ---- + if ENABLE_LDS_CACHE: + base_ptr = allocator.get_base() + lds_gather = SmemPtr( + base_ptr, + lds_gather_offset, + elem_type, + shape=(LDS_GATHER_TILE_ELEMS,), + ).get() + + # ---- Q row in-bounds guard ---- + q_active = arith.cmpi(arith.CmpIPredicate.slt, pid_m, seq_len_v) + pid_m_safe = arith.select(q_active, pid_m, arith.index(0)) + + # ---- Load Q for all HEAD_GROUP heads in this program ---- + zero_f32_vec = arith.constant_vector(0.0, T.vec(D_PER_LANE, f32_ty)) + q_f32_vecs = [] + for h_off in range_constexpr(HEAD_GROUP): + qhid_h = qhid_base + arith.index(h_off) + q_row_base = ((bid * arith.index(NUM_HEADS) + qhid_h) * seq_len_v + pid_m_safe) * arith.index( + HEAD_DIM + ) + q_lane_off = q_row_base + lane * arith.index(D_PER_LANE) + q_vec = load_f16_v(q_ptr, q_lane_off, D_PER_LANE) + q_f32_vec = arith.extf(T.vec(D_PER_LANE, f32_ty), q_vec) + q_f32_vec = arith.select(q_active, q_f32_vec, zero_f32_vec) + q_f32_vecs.append(q_f32_vec) + + # ---- Constants ---- + NEG_INF_F = -1.0e30 + c_neg_inf = arith.constant(NEG_INF_F, type=f32_ty) + c_zero_f = arith.constant(0.0, type=f32_ty) + c_one_f = arith.constant(1.0, type=f32_ty) + c_sm_scale_f = arith.constant(float(sm_scale), type=f32_ty) + c_log2e_f = arith.constant(_LOG2E, type=f32_ty) + + arith.index_cast(T.i32, lane) + width_i32 = arith.constant(WARP_SIZE, type=T.i32) + + def warp_reduce_sum_f32(v): + cur = v + for off in [32, 16, 8, 4, 2, 1]: + xor_amt = arith.constant(off, type=T.i32) + peer = arith.ArithValue(cur).shuffle_xor(xor_amt, width_i32) + cur = arith.AddFOp(cur, peer, fastmath=fm_fast).result + return cur + + def vec_dot_f32(a_vec, b_vec): + s = c_zero_f + for i in range_constexpr(D_PER_LANE): + av = vector.extract(a_vec, static_position=[i], dynamic_position=[]) + bv = vector.extract(b_vec, static_position=[i], dynamic_position=[]) + p = arith.MulFOp(av, bv, fastmath=fm_fast).result + s = arith.AddFOp(s, p, fastmath=fm_fast).result + return s + + # ---- Local SWA loop bounds ---- + _pid_p1 = pid_m + arith.index(1) + _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _pid_p1, seq_len_v) + n_loop_end = arith.select(_le_seq, _pid_p1, seq_len_v) + SWA = arith.index(int(swa_window)) + _ge_w = arith.cmpi(arith.CmpIPredicate.sge, _pid_p1, SWA) + _n_lo_raw = arith.select(_ge_w, _pid_p1 - SWA, arith.index(0)) + BN_idx = arith.index(BLOCK_N) + n_loop_start = (_n_lo_raw // BN_idx) * BN_idx + + # Init state: HEAD_GROUP copies of (m_i, l_i, acc[D_PER_LANE]). + # Layout: [m_0, l_0, acc_0_0..acc_0_{D-1}, m_1, l_1, acc_1_0..] + STATE_PER_HEAD = 2 + D_PER_LANE + init_args = [] + for _h in range_constexpr(HEAD_GROUP): + init_args.append(c_neg_inf) # m + init_args.append(c_zero_f) # l + for _ in range_constexpr(D_PER_LANE): + init_args.append(c_zero_f) + + # ==== LOCAL SWA loop ==== + for n_start, inner_args, loop_results_local in scf.for_( + n_loop_start, + n_loop_end, + BN_idx, + iter_args=init_args, + ): + # Unpack HEAD_GROUP states from inner_args + m_is = [inner_args[h * STATE_PER_HEAD] for h in range_constexpr(HEAD_GROUP)] + l_is = [inner_args[h * STATE_PER_HEAD + 1] for h in range_constexpr(HEAD_GROUP)] + accs = [ + [inner_args[h * STATE_PER_HEAD + 2 + d] for d in range_constexpr(D_PER_LANE)] + for h in range_constexpr(HEAD_GROUP) + ] + + n_start_i32 = arith.index_cast(T.i32, n_start) + pid_m_i32 = arith.index_cast(T.i32, pid_m) + seq_len_i32 = seq_len + if hasattr(seq_len_i32, "ir_value"): + seq_len_i32 = seq_len_i32.ir_value() + w_i32 = arith.constant(int(swa_window), type=T.i32) + + # Per-head QK values: [HEAD_GROUP][BLOCK_N] + qk_vals_per_head = [[] for _ in range_constexpr(HEAD_GROUP)] + kl_f32_cache = [] + bad_lo_cache = [] + for n_off in range_constexpr(BLOCK_N): + kv_col_i32 = arith.AddIOp( + n_start_i32, + arith.constant(n_off, type=T.i32), + ).result + _kv_plus_w_lo = arith.AddIOp(kv_col_i32, w_i32).result + is_swa_lo = arith.cmpi(arith.CmpIPredicate.sle, _kv_plus_w_lo, pid_m_i32) + is_causal_lo = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_i32, pid_m_i32) + is_oob = arith.cmpi(arith.CmpIPredicate.sge, kv_col_i32, seq_len_i32) + bad_lo = arith.OrIOp(arith.OrIOp(is_causal_lo, is_swa_lo).result, is_oob).result + bad_lo_cache.append(bad_lo) + + kv_col_idx = arith.index_cast(T.index, kv_col_i32) + kv_col_safe = arith.select(is_oob, arith.index(0), kv_col_idx) + # When mqa_kv, K is shared across heads -> load once. + # When mqa_kv=False, K differs per head, so we'd need per-head loads. + # For now this kernel requires mqa_kv when head_group > 1. + if const_expr(mqa_kv): + kl_row_base = (bid * seq_len_v + kv_col_safe) * arith.index(HEAD_DIM) + else: + # head_group > 1 with non-MQA: would need per-head load. Disallowed at setup. + kl_row_base = ( + (bid * arith.index(NUM_HEADS) + qhid_base) * seq_len_v + kv_col_safe + ) * arith.index(HEAD_DIM) + kl_lane_off = kl_row_base + lane * arith.index(D_PER_LANE) + kl_vec = load_f16_v(kl_ptr, kl_lane_off, D_PER_LANE) + kl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), kl_vec) + kl_f32_cache.append(kl_f32) + # Compute qk for each head using the shared K. + for h_off in range_constexpr(HEAD_GROUP): + lane_dot = vec_dot_f32(kl_f32, q_f32_vecs[h_off]) + qk_full = warp_reduce_sum_f32(lane_dot) + qk_scaled = arith.MulFOp(qk_full, c_sm_scale_f, fastmath=fm_fast).result + qk_masked = arith.select(bad_lo, c_neg_inf, qk_scaled) + qk_vals_per_head[h_off].append(qk_masked) + + # Per-head softmax update + new_m_is = [] + new_l_is = [] + new_accs = [] + p_vals_per_head = [] + for h_off in range_constexpr(HEAD_GROUP): + qk_vals_h = qk_vals_per_head[h_off] + m_tile = qk_vals_h[0] + for n_off in range_constexpr(BLOCK_N - 1): + m_tile = arith.MaxNumFOp(m_tile, qk_vals_h[n_off + 1], fastmath=fm_fast).result + m_new = arith.MaxNumFOp(m_is[h_off], m_tile, fastmath=fm_fast).result + + diff_m = arith.SubFOp(m_is[h_off], m_new, fastmath=fm_fast).result + diff_m_log2 = arith.MulFOp(diff_m, c_log2e_f, fastmath=fm_fast).result + alpha = arith.ArithValue(diff_m_log2).exp2(fastmath=fm_fast) + + p_vals = [] + tile_sum = c_zero_f + for n_off in range_constexpr(BLOCK_N): + d = arith.SubFOp(qk_vals_h[n_off], m_new, fastmath=fm_fast).result + dl = arith.MulFOp(d, c_log2e_f, fastmath=fm_fast).result + p = arith.ArithValue(dl).exp2(fastmath=fm_fast) + p_vals.append(p) + tile_sum = arith.AddFOp(tile_sum, p, fastmath=fm_fast).result + p_vals_per_head.append(p_vals) + + l_alpha = arith.MulFOp(l_is[h_off], alpha, fastmath=fm_fast).result + l_new = arith.AddFOp(l_alpha, tile_sum, fastmath=fm_fast).result + + acc_h = accs[h_off] + new_acc_h = [] + for d_off in range_constexpr(D_PER_LANE): + new_acc_h.append(arith.MulFOp(acc_h[d_off], alpha, fastmath=fm_fast).result) + new_m_is.append(m_new) + new_l_is.append(l_new) + new_accs.append(new_acc_h) + + # AV phase: load V once (MQA), reuse across HEAD_GROUP heads. + for n_off in range_constexpr(BLOCK_N): + kv_col_i32 = arith.AddIOp( + n_start_i32, + arith.constant(n_off, type=T.i32), + ).result + is_oob = arith.cmpi(arith.CmpIPredicate.sge, kv_col_i32, seq_len_i32) + kv_col_idx = arith.index_cast(T.index, kv_col_i32) + kv_col_safe = arith.select(is_oob, arith.index(0), kv_col_idx) + if const_expr(mqa_kv): + vl_row_base = (bid * seq_len_v + kv_col_safe) * arith.index(HEAD_DIM) + else: + vl_row_base = ( + (bid * arith.index(NUM_HEADS) + qhid_base) * seq_len_v + kv_col_safe + ) * arith.index(HEAD_DIM) + vl_lane_off = vl_row_base + lane * arith.index(D_PER_LANE) + vl_vec = load_f16_v(vl_ptr, vl_lane_off, D_PER_LANE) + vl_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), vl_vec) + for h_off in range_constexpr(HEAD_GROUP): + p_vals_h = p_vals_per_head[h_off] + new_acc_h = new_accs[h_off] + for d_off in range_constexpr(D_PER_LANE): + vv = vector.extract(vl_f32, static_position=[d_off], dynamic_position=[]) + contrib = arith.MulFOp(p_vals_h[n_off], vv, fastmath=fm_fast).result + new_acc_h[d_off] = arith.AddFOp(new_acc_h[d_off], contrib, fastmath=fm_fast).result + + # Pack yield args + yield_args = [] + for h_off in range_constexpr(HEAD_GROUP): + yield_args.append(new_m_is[h_off]) + yield_args.append(new_l_is[h_off]) + for d in range_constexpr(D_PER_LANE): + yield_args.append(new_accs[h_off][d]) + yield yield_args + + # Unpack HEAD_GROUP states from local SWA results + m_is = [loop_results_local[h * STATE_PER_HEAD] for h in range_constexpr(HEAD_GROUP)] + l_is = [loop_results_local[h * STATE_PER_HEAD + 1] for h in range_constexpr(HEAD_GROUP)] + accs = [ + [loop_results_local[h * STATE_PER_HEAD + 2 + d] for d in range_constexpr(D_PER_LANE)] + for h in range_constexpr(HEAD_GROUP) + ] + + # ==== SPARSE branch ==== + if const_expr(has_sparse): + init_sparse = [] + for h_off in range_constexpr(HEAD_GROUP): + init_sparse.append(m_is[h_off]) + init_sparse.append(l_is[h_off]) + for d in range_constexpr(D_PER_LANE): + init_sparse.append(accs[h_off][d]) + for k_start, inner_args, loop_results_sparse in scf.for_( + arith.index(0), + K_topk_v, + arith.index(BLOCK_K), + iter_args=init_sparse, + ): + m_is = [inner_args[h * STATE_PER_HEAD] for h in range_constexpr(HEAD_GROUP)] + l_is = [inner_args[h * STATE_PER_HEAD + 1] for h in range_constexpr(HEAD_GROUP)] + accs = [ + [inner_args[h * STATE_PER_HEAD + 2 + d] for d in range_constexpr(D_PER_LANE)] + for h in range_constexpr(HEAD_GROUP) + ] + + k_start_i32 = arith.index_cast(T.i32, k_start) + K_topk_i32 = K_topk + if hasattr(K_topk_i32, "ir_value"): + K_topk_i32 = K_topk_i32.ir_value() + + qk_vals_sparse_per_head = [[] for _ in range_constexpr(HEAD_GROUP)] + for k_off in range_constexpr(BLOCK_K): + k_pos_i32 = arith.AddIOp( + k_start_i32, + arith.constant(k_off, type=T.i32), + ).result + is_oob = arith.cmpi(arith.CmpIPredicate.sge, k_pos_i32, K_topk_i32) + k_pos_idx = arith.index_cast(T.index, k_pos_i32) + k_pos_safe = arith.select(is_oob, arith.index(0), k_pos_idx) + + g_row_base = ((bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe) * arith.index( + HEAD_DIM + ) + g_lane_off = g_row_base + lane * arith.index(D_PER_LANE) + g_vec = load_f16_v(g_ptr, g_lane_off, D_PER_LANE) + if ENABLE_LDS_CACHE: + lds_idx = arith.index(k_off * HEAD_DIM) + lane * arith.index(D_PER_LANE) + vector.store(g_vec, lds_gather, [lds_idx]) + g_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), g_vec) + + sm_off = (bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe + sm_val = load_f32_scalar(sm_ptr, sm_off) + # For each head, compute QK using shared g_f32. + for h_off in range_constexpr(HEAD_GROUP): + lane_dot = vec_dot_f32(g_f32, q_f32_vecs[h_off]) + qk_full = warp_reduce_sum_f32(lane_dot) + qk_scaled = arith.MulFOp(qk_full, c_sm_scale_f, fastmath=fm_fast).result + qk_biased = arith.AddFOp(qk_scaled, sm_val, fastmath=fm_fast).result + qk_masked = arith.select(is_oob, c_neg_inf, qk_biased) + qk_vals_sparse_per_head[h_off].append(qk_masked) + + new_m_is = [] + new_l_is = [] + new_accs = [] + p_vals_per_head = [] + for h_off in range_constexpr(HEAD_GROUP): + qk_vals_h = qk_vals_sparse_per_head[h_off] + m_tile = qk_vals_h[0] + for k_off in range_constexpr(BLOCK_K - 1): + m_tile = arith.MaxNumFOp(m_tile, qk_vals_h[k_off + 1], fastmath=fm_fast).result + m_new = arith.MaxNumFOp(m_is[h_off], m_tile, fastmath=fm_fast).result + + diff_m = arith.SubFOp(m_is[h_off], m_new, fastmath=fm_fast).result + diff_m_log2 = arith.MulFOp(diff_m, c_log2e_f, fastmath=fm_fast).result + alpha = arith.ArithValue(diff_m_log2).exp2(fastmath=fm_fast) + + p_vals = [] + tile_sum = c_zero_f + for k_off in range_constexpr(BLOCK_K): + d = arith.SubFOp(qk_vals_h[k_off], m_new, fastmath=fm_fast).result + dl = arith.MulFOp(d, c_log2e_f, fastmath=fm_fast).result + p = arith.ArithValue(dl).exp2(fastmath=fm_fast) + p_vals.append(p) + tile_sum = arith.AddFOp(tile_sum, p, fastmath=fm_fast).result + p_vals_per_head.append(p_vals) + + l_alpha = arith.MulFOp(l_is[h_off], alpha, fastmath=fm_fast).result + l_new = arith.AddFOp(l_alpha, tile_sum, fastmath=fm_fast).result + + acc_h = accs[h_off] + new_acc_h = [] + for d_off in range_constexpr(D_PER_LANE): + new_acc_h.append(arith.MulFOp(acc_h[d_off], alpha, fastmath=fm_fast).result) + new_m_is.append(m_new) + new_l_is.append(l_new) + new_accs.append(new_acc_h) + + # ---- AV phase: re-read gathered K-block once per k_off, reuse for all heads ---- + if ENABLE_LDS_CACHE: + _waitcnt_lgkm_0() + for k_off in range_constexpr(BLOCK_K): + if ENABLE_LDS_CACHE: + lds_idx = arith.index(k_off * HEAD_DIM) + lane * arith.index(D_PER_LANE) + g_vec = vector.load(T.vec(D_PER_LANE, f16_ty), lds_gather, [lds_idx]) + else: + k_pos_i32 = arith.AddIOp( + k_start_i32, + arith.constant(k_off, type=T.i32), + ).result + is_oob = arith.cmpi(arith.CmpIPredicate.sge, k_pos_i32, K_topk_i32) + k_pos_idx = arith.index_cast(T.index, k_pos_i32) + k_pos_safe = arith.select(is_oob, arith.index(0), k_pos_idx) + g_row_base = ((bid * seq_len_v + pid_m_safe) * K_topk_v + k_pos_safe) * arith.index( + HEAD_DIM + ) + g_lane_off = g_row_base + lane * arith.index(D_PER_LANE) + g_vec = load_f16_v(g_ptr, g_lane_off, D_PER_LANE) + g_f32 = arith.extf(T.vec(D_PER_LANE, f32_ty), g_vec) + for h_off in range_constexpr(HEAD_GROUP): + p_vals_h = p_vals_per_head[h_off] + new_acc_h = new_accs[h_off] + for d_off in range_constexpr(D_PER_LANE): + vv = vector.extract(g_f32, static_position=[d_off], dynamic_position=[]) + contrib = arith.MulFOp(p_vals_h[k_off], vv, fastmath=fm_fast).result + new_acc_h[d_off] = arith.AddFOp( + new_acc_h[d_off], contrib, fastmath=fm_fast + ).result + + # Pack yield args + yield_args = [] + for h_off in range_constexpr(HEAD_GROUP): + yield_args.append(new_m_is[h_off]) + yield_args.append(new_l_is[h_off]) + for d in range_constexpr(D_PER_LANE): + yield_args.append(new_accs[h_off][d]) + yield yield_args + + m_is = [loop_results_sparse[h * STATE_PER_HEAD] for h in range_constexpr(HEAD_GROUP)] + l_is = [loop_results_sparse[h * STATE_PER_HEAD + 1] for h in range_constexpr(HEAD_GROUP)] + accs = [ + [loop_results_sparse[h * STATE_PER_HEAD + 2 + d] for d in range_constexpr(D_PER_LANE)] + for h in range_constexpr(HEAD_GROUP) + ] + + # ==== Sink epilogue (per-head) ==== + if const_expr(has_sink): + for h_off in range_constexpr(HEAD_GROUP): + qhid_h = qhid_base + arith.index(h_off) + qhid_i32 = arith.index_cast(T.i32, qhid_h) + sink_h_val = buffer_ops.buffer_load( + sink_rsrc, + qhid_i32, + vec_width=1, + dtype=f32_ty, + ) + m_i_h = m_is[h_off] + l_i_h = l_is[h_off] + acc_h = accs[h_off] + m_new = arith.MaxNumFOp(m_i_h, sink_h_val, fastmath=fm_fast).result + d_alpha = arith.SubFOp(m_i_h, m_new, fastmath=fm_fast).result + d_alpha_log2 = arith.MulFOp(d_alpha, c_log2e_f, fastmath=fm_fast).result + alpha_sink = arith.ArithValue(d_alpha_log2).exp2(fastmath=fm_fast) + d_beta = arith.SubFOp(sink_h_val, m_new, fastmath=fm_fast).result + d_beta_log2 = arith.MulFOp(d_beta, c_log2e_f, fastmath=fm_fast).result + beta_sink = arith.ArithValue(d_beta_log2).exp2(fastmath=fm_fast) + l_alpha = arith.MulFOp(l_i_h, alpha_sink, fastmath=fm_fast).result + l_is[h_off] = arith.AddFOp(l_alpha, beta_sink, fastmath=fm_fast).result + new_acc_h = [] + for d_off in range_constexpr(D_PER_LANE): + new_acc_h.append(arith.MulFOp(acc_h[d_off], alpha_sink, fastmath=fm_fast).result) + accs[h_off] = new_acc_h + m_is[h_off] = m_new + + # ==== Final divide and store (per-head) ==== + _o_guard = scf.IfOp(q_active, [], has_else=False) + with ir.InsertionPoint(_o_guard.then_block): + for h_off in range_constexpr(HEAD_GROUP): + qhid_h = qhid_base + arith.index(h_off) + l_i_h = l_is[h_off] + m_i_h = m_is[h_off] + acc_h = accs[h_off] + inv_l = arith.DivFOp(c_one_f, l_i_h, fastmath=fm_fast).result + o_row_base = ((bid * arith.index(NUM_HEADS) + qhid_h) * seq_len_v + pid_m_safe) * arith.index( + HEAD_DIM + ) + o_lane_off = o_row_base + lane * arith.index(D_PER_LANE) + for d_off in range_constexpr(D_PER_LANE): + o_f32 = arith.MulFOp(acc_h[d_off], inv_l, fastmath=fm_fast).result + o_f16 = arith.trunc_f(elem_type, o_f32) + elem_off = o_lane_off + arith.index(d_off) + idx_i64 = arith.index_cast(T.i64, elem_off) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + o_ptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=elem_type, + noWrapFlags=0, + ) + _llvm.StoreOp(o_f16, gep.result) + + is_lane0 = arith.cmpi(arith.CmpIPredicate.eq, lane, arith.index(0)) + _lse_if = scf.IfOp(is_lane0, [], has_else=False) + with ir.InsertionPoint(_lse_if.then_block): + ln_l = math_dialect.log(l_i_h, fastmath=fm_fast) + lse_val = arith.AddFOp(m_i_h, ln_l, fastmath=fm_fast).result + lse_off = (bid * arith.index(NUM_HEADS) + qhid_h) * seq_len_v + pid_m_safe + lse_off_i32 = arith.index_cast(T.i32, lse_off) + buffer_ops.buffer_store(lse_val, lse_rsrc, lse_off_i32) + scf.YieldOp([]) + scf.YieldOp([]) + + @flyc.jit + def launch_v4_csa_fwd( + Q: fx.Tensor, + K_LOCAL: fx.Tensor, + V_LOCAL: fx.Tensor, + GATHERED: fx.Tensor, + SPARSE_MASK: fx.Tensor, + Sink: fx.Tensor, + O: fx.Tensor, + LSE: fx.Tensor, + batch_size: fx.Int32, + seq_len: fx.Int32, + K_topk: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + bs_idx = arith.index_cast(T.index, batch_size) + sl_idx = arith.index_cast(T.index, seq_len) + grid_x = sl_idx + grid_y = bs_idx * arith.index(NUM_HEAD_GROUPS) + + launcher = v4_csa_fwd_kernel( + Q, + K_LOCAL, + V_LOCAL, + GATHERED, + SPARSE_MASK, + Sink, + O, + LSE, + seq_len, + K_topk, + ) + + if waves_per_eu is not None: + _wpe = int(waves_per_eu) + if _wpe >= 1: + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe) + + passthrough_entries = [] + if daz: + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("denormal-fp-math-f32"), + ir.StringAttr.get("preserve-sign,preserve-sign"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("no-nans-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("unsafe-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries) + + launcher.launch( + grid=(grid_x, grid_y, 1), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) + + compile_hints = { + "fast_fp_math": fast_fp_math, + "unsafe_fp_math": unsafe_fp_math, + } + + def _launch(*args, **kwargs): + with CompilationContext.compile_hints(compile_hints): + return launch_v4_csa_fwd(*args, **kwargs) + + def _compile( + Q, K_LOCAL, V_LOCAL, GATHERED, SPARSE_MASK, Sink, O, LSE, batch_size, seq_len, K_topk, stream=None + ): + with CompilationContext.compile_hints(compile_hints): + return flyc.compile( + launch_v4_csa_fwd, + Q, + K_LOCAL, + V_LOCAL, + GATHERED, + SPARSE_MASK, + Sink, + O, + LSE, + batch_size, + seq_len, + K_topk, + fx.Stream(stream), + ) + + _launch.compile = _compile + + return _launch + + +build_v4_csa_fwd_module_primary = build_v4_csa_fwd_module diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dkv_pool_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dkv_pool_kernel.py new file mode 100644 index 000000000..528456c3d --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dkv_pool_kernel.py @@ -0,0 +1,926 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""v4_hca_bwd_dkv_pool: V4 HCA backward dK/dV POOL-stream kernel for FlyDSL. + +Forked from v4_sla_bwd_dkv_kernel.py (the SLA MQA dKdV head-loop accumulator +template). Computes the POOL-stream contribution to dk/dv for an HCA +(Hybrid-Causal-Additive) attention backward and stores into the POOL slice +of dk_fp32 / dv_fp32 buffers (which are zero-initialised by the wrapper). + +Key differences vs the SWA kernel: + - KV range is fixed at [HCA_LOCAL_SEQLEN, HCA_LOCAL_SEQLEN+POOL_SIZE). + POOL_SIZE is a build-time constexpr. For HCA shapes POOL_SIZE <= BLOCK_N + (BLOCK_N=32 here, POOL_SIZE in {4,32}), so each program owns the entire + pool slice for one batch. + - No SWA-window mask, no causal mask. Two element-wise predicates only: + (a) pool_n < POOL_SIZE -> NEG_INF outside the pool + (b) q_row < seq_len_q -> NEG_INF for OOB q rows + - qk + add_bias from ADD_MASK[Sq, POOL_SIZE]. Each element loaded as bf16/f16 + then cast to f32 inside the inner mask loop (matches dq_pool sibling). + - LSE / DELTAS are JOINT (saved from HCA fwd) so the same q_row LSE governs + both local and pool streams. + - Grid: (B,) -- one program per batch, owning the entire pool slice. + - Head loop is the SAME head-loop accumulator pattern: dynamic ``scf.for_`` + over HQ (NOT range_constexpr; constexpr-unroll at HQ=128 hangs MLIR). + - sm_scale applied to dK once post head-loop (same as SWA dkv, P57 cr=0). + - No atomics. Single store per (b, pool_n) slice at the end. + +LDS budget at BLOCK_N=32, BLOCK_M2=32, D=512: + K = 32 KB, V = 32 KB, DO = 32 KB, Q = 32 KB, pT = 2 KB + Total = 130 KB <= 160 KB. +Auto-fallback: if predicted > 160 KB, drop DO/Q LDS scratches and read +Q/DO directly to register packs from HBM (mirrors SWA dkv). +""" + +import math +import os + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import math as math_dialect +from flydsl._mlir.dialects import memref as _memref +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr +from kernels.kernels_common import dtype_to_elem_type + +KERNEL_NAME = "v4_hca_bwd_dkv_pool_kernel" + +_LLVM_GEP_DYNAMIC = -2147483648 + + +def _llvm_ptr_ty(): + return ir.Type.parse("!llvm.ptr") + + +def _llvm_lds_ptr_ty(): + return ir.Type.parse("!llvm.ptr<3>") + + +def build_v4_hca_bwd_dkv_pool_module( + num_heads, + head_dim, + pool_size, + hca_local_seqlen, + dtype_str="bf16", + sm_scale=None, + waves_per_eu=2, + flat_work_group_size=None, + block_n=None, + block_m2=None, + unsafe_fp_math=True, + fast_fp_math=True, + daz=True, + layout_bhld=True, + mqa_kv=True, +): + """Build the V4 HCA backward dK/dV POOL-stream launcher (MQA only).""" + gpu_arch = get_hip_arch() + + if block_n is None: + BLOCK_N = 32 + else: + BLOCK_N = int(block_n) + if block_m2 is None: + BLOCK_M2 = 32 + else: + BLOCK_M2 = int(block_m2) + WARP_SIZE = 64 + # R6d-A: MFMA 16x16x32 -> each wave covers 16 KV rows. + NUM_WAVES = max(1, BLOCK_N // 16) + ROWS_PER_WAVE = BLOCK_N // NUM_WAVES + if flat_work_group_size is None: + flat_work_group_size = NUM_WAVES * WARP_SIZE + BLOCK_SIZE = flat_work_group_size + + ENABLE_LDS_VEC16 = os.getenv("FLYDSL_SLA_FWD_ENABLE_LDS_VEC16", "1") == "1" + USE_K16 = gpu_arch.startswith("gfx950") + # R6d-A: MFMA 16x16x32 K-step = 32; A/B-frag = 8 bf16 per lane. + assert USE_K16, "R6d-A dkv_pool requires gfx950 (MFMA 16x16x32 bf16)." + K_STEP_QK = 32 + K_STEPS_QK = head_dim // K_STEP_QK + # R6d-A: each MFMA 16x16x32 tile produces a 16-wide N-chunk; D_CHUNK = 16 cols. + D_CHUNK = 16 + D_CHUNKS = head_dim // D_CHUNK + K_STEPS_PT = BLOCK_M2 // K_STEP_QK + # R6d-A: GEMM1/GEMM3 cover m_col [0..BLOCK_M2) with multiple 16-col MFMA-N tiles per ks. + assert BLOCK_M2 % 16 == 0, f"BLOCK_M2 must be a multiple of 16 (MFMA-N), got {BLOCK_M2}" + M_TILES = BLOCK_M2 // 16 + + assert BLOCK_N % NUM_WAVES == 0 + assert ROWS_PER_WAVE == 16, f"ROWS_PER_WAVE must equal 16 for MFMA 16x16x32, got {ROWS_PER_WAVE}" + assert head_dim % 32 == 0 + assert head_dim >= 64 + assert flat_work_group_size in (64, 128, 256, 512) + assert dtype_str == "bf16", "R6d-A dkv_pool currently only supports bf16." + assert BLOCK_N % 16 == 0 + assert BLOCK_M2 % K_STEP_QK == 0, ( + f"BLOCK_M2 ({BLOCK_M2}) must be a multiple of MFMA K-step " f"({K_STEP_QK}) for the dV/dK GEMMs." + ) + assert mqa_kv, "v4_hca_bwd_dkv_pool currently only supports MQA (HK=1)." + assert ( + isinstance(pool_size, int) and 0 < pool_size <= BLOCK_N + ), f"pool_size must be int in (0, {BLOCK_N}], got {pool_size!r}" + assert isinstance(hca_local_seqlen, int) and hca_local_seqlen >= 0 + assert ( + hca_local_seqlen % BLOCK_N == 0 + ), f"hca_local_seqlen must be multiple of BLOCK_N={BLOCK_N}, got {hca_local_seqlen}" + + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(head_dim) + + NUM_HEADS = num_heads + HEAD_DIM = head_dim + POOL_SIZE = int(pool_size) + HCA_LOCAL = int(hca_local_seqlen) + + K_STRIDE = HEAD_DIM + K_SWZ_ROW_MASK = (K_STRIDE // 16) - 1 + assert K_SWZ_ROW_MASK >= 0 + assert (K_SWZ_ROW_MASK & (K_SWZ_ROW_MASK + 1)) == 0 + V_STRIDE = HEAD_DIM + + VEC_WIDTH = 16 if ENABLE_LDS_VEC16 else 8 + assert HEAD_DIM % VEC_WIDTH == 0 + THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH + assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0 + ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD + + LDS_K_TILE_SIZE = BLOCK_N * K_STRIDE + LDS_V_TILE_SIZE = BLOCK_N * V_STRIDE + LDS_DO_STRIDE = HEAD_DIM + LDS_DO_ELEMS = BLOCK_M2 * LDS_DO_STRIDE + LDS_Q_STRIDE = HEAD_DIM + LDS_Q_ELEMS = BLOCK_M2 * LDS_Q_STRIDE + LDS_PT_STRIDE = BLOCK_M2 + LDS_PT_ELEMS = BLOCK_N * LDS_PT_STRIDE + + _LDS_LIMIT_BYTES = 160 * 1024 + + def _predict_lds_bytes(use_lds_for_q_do): + b = (LDS_K_TILE_SIZE + LDS_V_TILE_SIZE + LDS_PT_ELEMS) * 2 + if use_lds_for_q_do: + b += (LDS_DO_ELEMS + LDS_Q_ELEMS) * 2 + return b + + USE_LDS_FOR_Q_DO = True + _pred_full = _predict_lds_bytes(True) + if _pred_full > _LDS_LIMIT_BYTES: + USE_LDS_FOR_Q_DO = False + _pred_min = _predict_lds_bytes(False) + if _pred_min > _LDS_LIMIT_BYTES: + raise RuntimeError( + f"v4_hca_bwd_dkv_pool: minimal LDS {_pred_min}B > limit " + f"{_LDS_LIMIT_BYTES}B at D={head_dim}, BLOCK_N={BLOCK_N}, " + f"BLOCK_M2={BLOCK_M2}." + ) + import sys as _sys + + print( + f"[v4_hca_bwd_dkv_pool] auto-fallback: LDS {_pred_full}B > " + f"{_LDS_LIMIT_BYTES}B; disabling Q/DO LDS " + f"({_pred_full} -> {_pred_min})", + file=_sys.stderr, + flush=True, + ) + else: + import sys as _sys + + print( + f"[v4_hca_bwd_dkv_pool] LDS OK: predicted {_pred_full}B <= " + f"{_LDS_LIMIT_BYTES}B at D={head_dim}, BLOCK_N={BLOCK_N}, " + f"BLOCK_M2={BLOCK_M2}, USE_LDS_FOR_Q_DO=True", + file=_sys.stderr, + flush=True, + ) + + allocator = SmemAllocator( + None, + arch=gpu_arch, + global_sym_name=( + f"v4_hca_bwd_dkv_pool_smem_N{BLOCK_N}_M2_{BLOCK_M2}_P{POOL_SIZE}" + f"_L{HCA_LOCAL}_MQ{int(mqa_kv)}_QDO{int(USE_LDS_FOR_Q_DO)}" + ), + ) + lds_kv_offset = allocator._align(allocator.ptr, 16) + LDS_V_BASE = LDS_K_TILE_SIZE + LDS_KV_TOTAL_SIZE = LDS_K_TILE_SIZE + LDS_V_TILE_SIZE + allocator.ptr = lds_kv_offset + LDS_KV_TOTAL_SIZE * 2 + + lds_pt_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_pt_offset + LDS_PT_ELEMS * 2 + + if USE_LDS_FOR_Q_DO: + lds_do_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_do_offset + LDS_DO_ELEMS * 2 + lds_q_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_q_offset + LDS_Q_ELEMS * 2 + else: + lds_do_offset = None + lds_q_offset = None + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def v4_hca_bwd_dkv_pool_kernel( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + DOS: fx.Tensor, + LSE: fx.Tensor, + DELTAS: fx.Tensor, + DK: fx.Tensor, + DV: fx.Tensor, + ADD_MASK: fx.Tensor, + seq_len_q: fx.Int32, + seq_len_k: fx.Int32, + ): + elem_type = dtype_to_elem_type(dtype_str) + compute_type = T.f32 + q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q) + k_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K) + v_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V) + do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOS) + _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DK) + _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DV) + add_mask_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), ADD_MASK) + lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True) + deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True) + # R6e: dk/dv pool-slice writes are atomic_fadd (multi-program collisions). + dk_rsrc = buffer_ops.create_buffer_resource(DK, max_size=True) + dv_rsrc = buffer_ops.create_buffer_resource(DV, max_size=True) + + fm_fast = arith.FastMathFlags.fast + vxf16_type = T.vec(VEC_WIDTH, elem_type) + v8f16_type = T.vec(8, elem_type) + # R6d-A: MFMA 16x16x32 C-frag = 4 fp32 per lane. + v4f32_type = T.vec(4, compute_type) + mfma_pack_type = v8f16_type + MFMA_LANE_K = 8 + v1_elem_type = T.vec(1, elem_type) + _mfma_zero = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0) + + def mfma_acc(a, b, c): + # rocdl.mfma_f32_16x16x32_bf16 is the wrapped form: takes + # (result_type, [operands]) and returns the Value directly. + return rocdl.mfma_f32_16x16x32_bf16( + v4f32_type, + [a, b, c, _mfma_zero, _mfma_zero, _mfma_zero], + ) + + seq_len_q_v = arith.index_cast(T.index, seq_len_q) + seq_len_k_v = arith.index_cast(T.index, seq_len_k) + + base_ptr = allocator.get_base() + lds_kv = SmemPtr(base_ptr, lds_kv_offset, elem_type, shape=(LDS_KV_TOTAL_SIZE,)).get() + lds_pt = SmemPtr(base_ptr, lds_pt_offset, elem_type, shape=(LDS_PT_ELEMS,)).get() + if USE_LDS_FOR_Q_DO: + lds_do = SmemPtr(base_ptr, lds_do_offset, elem_type, shape=(LDS_DO_ELEMS,)).get() + lds_q = SmemPtr(base_ptr, lds_q_offset, elem_type, shape=(LDS_Q_ELEMS,)).get() + + block_id = arith.index_cast(T.index, gpu.block_idx.x) + tid = arith.index_cast(T.index, gpu.thread_idx.x) + wave_id = tid // WARP_SIZE + lane = tid % WARP_SIZE + # R6d-A MFMA 16x16x32 lane decomposition: + # A-frag: A[lane_mod_16, ks*32 + lane_div_16*8 + 0..7] + # B-frag: B[ks*32 + lane_div_16*8 + 0..7, lane_mod_16] + # C-frag: C[lane_div_16*4 + ii, lane_mod_16] for ii in 0..3 + lane_mod_16 = lane % 16 + lane_div_16 = lane // 16 + + wave_n_offset = wave_id * ROWS_PER_WAVE + + # R6e: grid is (B * num_m_blocks,). Decompose into (batch_idx, m_tile_idx). + # Layout: m_tile fastest -> consecutive programs share Q/dO batch. + BM2_idx_grid = arith.index(BLOCK_M2) + _one_idx_grid = arith.index(1) + num_m_blocks = (seq_len_q_v + BM2_idx_grid - _one_idx_grid) // BM2_idx_grid + m_tile_idx = block_id % num_m_blocks + batch_idx = block_id // num_m_blocks + m_start = m_tile_idx * BM2_idx_grid + # Pool slice starts at HCA_LOCAL_SEQLEN and runs for POOL_SIZE keys. + kv_start = arith.index(HCA_LOCAL) + + load_row_in_batch = tid // THREADS_PER_ROW_LOAD + load_lane_in_row = tid % THREADS_PER_ROW_LOAD + load_col_base = load_lane_in_row * VEC_WIDTH + + bh_base_tokens_kv = batch_idx * seq_len_k_v + + def bh_base_tokens_q_of(qhid_index): + return (batch_idx * NUM_HEADS + qhid_index) * seq_len_q_v + + def global_idx_q(qhid_index, token_idx, col): + return (bh_base_tokens_q_of(qhid_index) + token_idx) * arith.index(HEAD_DIM) + col + + def global_idx_kv(token_idx, col): + return (bh_base_tokens_kv + token_idx) * arith.index(HEAD_DIM) + col + + def _gep_load(base_ptr_, elem_idx, vec_type, et=elem_type): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr_, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=et, + noWrapFlags=0, + ) + return _llvm.LoadOp(vec_type, gep.result).result + + def _gep_store_f32(val, base_ptr_, elem_idx): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr_, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=T.f32, + noWrapFlags=0, + ) + _llvm.StoreOp(val, gep.result) + + def load_global_mfma_pack(base_ptr_, base_idx): + return _gep_load(base_ptr_, base_idx, mfma_pack_type) + + def load_global_f16xN(base_ptr_, base_idx): + return _gep_load(base_ptr_, base_idx, vxf16_type) + + def _k_swizzle(row_idx, col_idx): + mask = (row_idx & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + return col_idx ^ mask + + if ROWS_PER_BATCH_LOAD >= BLOCK_N: + NUM_BATCHES_KV = 1 + KV_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_N + else: + assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0 + NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD + KV_NEEDS_GUARD = False + + if ROWS_PER_BATCH_LOAD >= BLOCK_M2: + NUM_BATCHES_M = 1 + M_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_M2 + else: + assert BLOCK_M2 % ROWS_PER_BATCH_LOAD == 0 + NUM_BATCHES_M = BLOCK_M2 // ROWS_PER_BATCH_LOAD + M_NEEDS_GUARD = False + + c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type) + c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type) + c_zero_elem = arith.constant(0.0, type=elem_type) + c_zero_f = arith.constant(0.0, type=compute_type) + c_neg_inf = arith.constant(-1.0e30, type=compute_type) + c_sm_scale = arith.constant(sm_scale, type=compute_type) + # R6d-A: 4-elem fp32 acc per MFMA 16x16x32 op. + v4f32_zero = arith.constant_vector(0.0, v4f32_type) + + # ---- PROLOGUE: cooperative load K, V into LDS ---- + pool_end = kv_start + arith.index(POOL_SIZE) + + def coop_load_k(): + k_base = arith.index(0) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = kv_start + load_row_in_batch + row_offset + in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, pool_end) + row_safe = arith.select(in_bounds, row_idx, arith.index(0)) + g_idx = global_idx_kv(row_safe, load_col_base) + vec = load_global_f16xN(k_ptr, g_idx) + vec_safe = arith.select(in_bounds, vec, c_zero_vxf16) + lds_row = load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_N) + ) + _if_k = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_k.then_block): + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + vector.store(vec_safe, lds_kv, [lds_idx]) + scf.YieldOp([]) + else: + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + vector.store(vec_safe, lds_kv, [lds_idx]) + + def coop_load_v(): + v_base = arith.index(LDS_V_BASE) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = kv_start + load_row_in_batch + row_offset + in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, pool_end) + row_safe = arith.select(in_bounds, row_idx, arith.index(0)) + g_idx = global_idx_kv(row_safe, load_col_base) + vec = load_global_f16xN(v_ptr, g_idx) + vec_safe = arith.select(in_bounds, vec, c_zero_vxf16) + lds_row = load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_N) + ) + _if_v = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_v.then_block): + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = v_base + lds_row * V_STRIDE + swz_col + vector.store(vec_safe, lds_kv, [lds_idx]) + scf.YieldOp([]) + else: + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = v_base + lds_row * V_STRIDE + swz_col + vector.store(vec_safe, lds_kv, [lds_idx]) + + coop_load_k() + coop_load_v() + gpu.barrier() + + # R6d-A: K/V LDS A-frag uses lane_mod_16 (16-row MFMA) and lane_div_16*8 for col. + # Per-call swizzle mask (folds wave_n_offset bits at D=512). + def _k_idx_wave(ks): + kv_row = wave_n_offset + lane_mod_16 + col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + mask = (kv_row & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + return kv_row * arith.index(K_STRIDE) + (col ^ mask) + + def _v_idx_wave(ks): + kv_row = wave_n_offset + lane_mod_16 + col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + mask = (kv_row & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + return arith.index(LDS_V_BASE) + kv_row * arith.index(V_STRIDE) + (col ^ mask) + + # R6e: m-loop replaced by per-program m-tile. + arith.index(BLOCK_N) + arith.index(BLOCK_M2) + arith.index(1) + + outer_carry = [v4f32_zero for _ in range(D_CHUNKS)] + [v4f32_zero for _ in range(D_CHUNKS)] + + seq_len_q_i32 = arith.index_cast(T.i32, seq_len_q_v) + wave_n_off_i32 = arith.index_cast(T.i32, wave_n_offset) + lane_div_16_i32 = arith.index_cast(T.i32, lane_div_16) + pool_size_i32 = arith.constant(POOL_SIZE, type=T.i32) + + # ADD_MASK row stride in elements (= POOL_SIZE; bf16/f16 contiguous). + ADD_MASK_STRIDE_M = arith.index(POOL_SIZE) + + NUM_HEADS_idx = arith.index(NUM_HEADS) + for qhid_constexpr_idx, h_carry, h_loop_results in scf.for_( + arith.index(0), + NUM_HEADS_idx, + arith.index(1), + iter_args=outer_carry, + ): + qhid = qhid_constexpr_idx + + # R6e: per-program m-tile (no inner m-loop). Accumulators are head-loop carry. + dv_accs = [h_carry[dc] for dc in range_constexpr(D_CHUNKS)] + dk_accs = [h_carry[D_CHUNKS + dc] for dc in range_constexpr(D_CHUNKS)] + + # R6d-A: lane_mod_16 is the per-tile q-row coord; full per-tile loop builds row below. + q_row_abs = m_start + lane_mod_16 + q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, q_row_abs, seq_len_q_v) + arith.select(q_in_bounds, q_row_abs, arith.index(0)) + + if USE_LDS_FOR_Q_DO: + for batch in range_constexpr(NUM_BATCHES_M): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = m_start + load_row_in_batch + row_offset + in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, seq_len_q_v) + row_safe = arith.select(in_bounds, row_idx, arith.index(0)) + g_idx_do = global_idx_q(qhid, row_safe, load_col_base) + g_idx_q = global_idx_q(qhid, row_safe, load_col_base) + vec_do = load_global_f16xN(do_ptr, g_idx_do) + vec_q = load_global_f16xN(q_ptr, g_idx_q) + vec_do_safe = arith.select(in_bounds, vec_do, c_zero_vxf16) + vec_q_safe = arith.select(in_bounds, vec_q, c_zero_vxf16) + lds_row = load_row_in_batch + row_offset + if M_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_M2) + ) + _if_qd = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_qd.then_block): + lds_idx_do = lds_row * arith.index(LDS_DO_STRIDE) + load_col_base + lds_idx_q = lds_row * arith.index(LDS_Q_STRIDE) + load_col_base + vector.store(vec_do_safe, lds_do, [lds_idx_do]) + vector.store(vec_q_safe, lds_q, [lds_idx_q]) + scf.YieldOp([]) + else: + lds_idx_do = lds_row * arith.index(LDS_DO_STRIDE) + load_col_base + lds_idx_q = lds_row * arith.index(LDS_Q_STRIDE) + load_col_base + vector.store(vec_do_safe, lds_do, [lds_idx_do]) + vector.store(vec_q_safe, lds_q, [lds_idx_q]) + gpu.barrier() + + # R6d-A: B-frag for GEMM1 (qkT = K @ Q^T) per m-tile mt in [0, M_TILES). + # Q[mt*16 + lane_mod_16, ks*32 + lane_div_16*8 + 0..7] + q_b_packs = [[None] * K_STEPS_QK for _ in range(M_TILES)] + for mt in range_constexpr(M_TILES): + for ks in range_constexpr(K_STEPS_QK): + m_row = arith.index(mt * 16) + lane_mod_16 + d_col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + lds_idx = m_row * arith.index(LDS_Q_STRIDE) + d_col + pack = vector.load_op(mfma_pack_type, lds_q, [lds_idx]) + q_b_packs[mt][ks] = pack + + # R6d-A: B-frag for GEMM3 (dp = V @ DO^T) per m-tile mt. + do_b_packs_gemm3 = [[None] * K_STEPS_QK for _ in range(M_TILES)] + for mt in range_constexpr(M_TILES): + for ks in range_constexpr(K_STEPS_QK): + m_row = arith.index(mt * 16) + lane_mod_16 + d_col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + lds_idx = m_row * arith.index(LDS_DO_STRIDE) + d_col + pack = vector.load_op(mfma_pack_type, lds_do, [lds_idx]) + do_b_packs_gemm3[mt][ks] = pack + else: + # R6d-A fallback: HBM-direct B-frag per m-tile. row = m_start + mt*16 + lane_mod_16. + q_b_packs = [[None] * K_STEPS_QK for _ in range(M_TILES)] + do_b_packs_gemm3 = [[None] * K_STEPS_QK for _ in range(M_TILES)] + for mt in range_constexpr(M_TILES): + row_abs = m_start + arith.index(mt * 16) + lane_mod_16 + in_bnd = arith.cmpi(arith.CmpIPredicate.slt, row_abs, seq_len_q_v) + row_safe = arith.select(in_bnd, row_abs, arith.index(0)) + for ks in range_constexpr(K_STEPS_QK): + col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + g_idx = global_idx_q(qhid, row_safe, col) + q_raw = load_global_mfma_pack(q_ptr, g_idx) + q_b_packs[mt][ks] = arith.select(in_bnd, q_raw, c_zero_mfma_pack) + do_raw = load_global_mfma_pack(do_ptr, g_idx) + do_b_packs_gemm3[mt][ks] = arith.select(in_bnd, do_raw, c_zero_mfma_pack) + + # ---- GEMM1: qkT = K @ Q^T (MFMA 16x16x32, one tile per m-tile mt) ---- + s_accs = [v4f32_zero for _ in range(M_TILES)] + for ks in range_constexpr(K_STEPS_QK): + k_pack = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_wave(ks)]) + for mt in range_constexpr(M_TILES): + s_accs[mt] = mfma_acc(k_pack, q_b_packs[mt][ks], s_accs[mt]) + + # R6d-A: per m-tile lse/delta, pool-only mask, ADD_MASK bias, softmax. + # m_row for tile mt = m_start + mt*16 + lane_mod_16. + # Resulting pT_vals_per_tile[mt][ii] for ii in 0..3. + pT_vals_per_tile = [] + lse_per_tile = [] + delta_per_tile = [] + for mt in range_constexpr(M_TILES): + m_row_for_lse = m_start + arith.index(mt * 16) + lane_mod_16 + m_row_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, m_row_for_lse, seq_len_q_v) + m_row_safe = arith.select(m_row_in_bounds, m_row_for_lse, arith.index(0)) + lse_off_i32 = arith.index_cast(T.i32, bh_base_tokens_q_of(qhid) + m_row_safe) + lse_v = buffer_ops.buffer_load(lse_rsrc, lse_off_i32, vec_width=1, dtype=T.f32) + delta_v = buffer_ops.buffer_load(deltas_rsrc, lse_off_i32, vec_width=1, dtype=T.f32) + lse_per_tile.append(lse_v) + delta_per_tile.append(delta_v) + + m_row_abs_for_mask_i32 = arith.index_cast(T.i32, m_row_for_lse) + m_oob = arith.cmpi(arith.CmpIPredicate.sge, m_row_abs_for_mask_i32, seq_len_q_i32) + # ADD_MASK row base (clamp to row 0 if OOB; lane is masked). + add_mask_row_base = m_row_safe * ADD_MASK_STRIDE_M + + pT_tile = [] + for ii in range_constexpr(4): + ii_i32 = arith.constant(ii, type=T.i32) + # R6d-A: 16x16x32 C-frag lane stride is 4 (NOT 8 like 32x32x16). + pool_n_rel_i32 = arith.AddIOp( + arith.MulIOp(lane_div_16_i32, arith.constant(4, type=T.i32)).result, ii_i32 + ).result + pool_n_i32 = arith.AddIOp(wave_n_off_i32, pool_n_rel_i32).result + pool_n_oob = arith.cmpi(arith.CmpIPredicate.sge, pool_n_i32, pool_size_i32) + bad = arith.OrIOp(pool_n_oob, m_oob).result + + # Load add_bias from ADD_MASK[m_row, pool_n]. + pool_n_safe_i32 = arith.select(bad, arith.constant(0, type=T.i32), pool_n_i32) + pool_n_idx = arith.index_cast(T.index, pool_n_safe_i32) + add_elem_idx = add_mask_row_base + pool_n_idx + bias_raw_v1 = _gep_load(add_mask_ptr, add_elem_idx, T.vec(1, elem_type)) + bias_raw = vector.extract(bias_raw_v1, static_position=[0], dynamic_position=[]) + bias_f32 = arith.extf(compute_type, bias_raw) + bias_safe = arith.select(bad, c_zero_f, bias_f32) + + s_ii = vector.extract(s_accs[mt], static_position=[ii], dynamic_position=[]) + scaled = arith.MulFOp(s_ii, c_sm_scale, fastmath=fm_fast).result + scaled_plus_bias = arith.AddFOp(scaled, bias_safe, fastmath=fm_fast).result + scaled_m = arith.select(bad, c_neg_inf, scaled_plus_bias) + diff = arith.SubFOp(scaled_m, lse_v, fastmath=fm_fast).result + p = math_dialect.exp(diff, fastmath=fm_fast) + pT_tile.append(p) + pT_vals_per_tile.append(pT_tile) + + # ---- pT register -> LDS (per m-tile mt, write 4 C-frag elems/lane @ col mt*16 + lane_mod_16) ---- + for mt in range_constexpr(M_TILES): + for ii in range_constexpr(4): + kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii) + kv_row = wave_n_offset + kv_row_rel + pt_bf16 = arith.trunc_f(elem_type, pT_vals_per_tile[mt][ii]) + lds_pt_idx = kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16 + v1 = vector.from_elements(v1_elem_type, [pt_bf16]) + vector.store(v1, lds_pt, [lds_pt_idx]) + + # ---- pT A-frag for GEMM2 (dV += pT @ DO): + # A[kv_row=wave_n_offset+lane_mod_16, m_col=m_step*32+lane_div_16*8 + 0..7] + pt_a_packs = [] + for m_step in range_constexpr(K_STEPS_PT): + kv_row_a = wave_n_offset + lane_mod_16 + m_col_a = arith.index(m_step * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + lds_idx = kv_row_a * arith.index(LDS_PT_STRIDE) + m_col_a + pack = vector.load_op(mfma_pack_type, lds_pt, [lds_idx]) + pt_a_packs.append(pack) + + # ---- GEMM2: dV += pT @ DO ---- + # R6d-A GEMM2 (dV += pT @ DO): B-frag DO[m_step*32+lane_div_16*8+k, dc*16+lane_mod_16]. + if USE_LDS_FOR_Q_DO: + + def read_do_b_pack(m_step_idx, dc_idx): + d_col = arith.index(dc_idx * D_CHUNK) + lane_mod_16 + m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + vals = [] + for rk in range_constexpr(MFMA_LANE_K): + m_row = m_base + arith.index(rk) + lds_idx = m_row * arith.index(LDS_DO_STRIDE) + d_col + val = _memref.load(lds_do, [lds_idx]) + vals.append(val) + return vector.from_elements(mfma_pack_type, vals) + + else: + + def read_do_b_pack(m_step_idx, dc_idx): + d_col = arith.index(dc_idx * D_CHUNK) + lane_mod_16 + m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + vals = [] + for rk in range_constexpr(MFMA_LANE_K): + m_row_rel = m_base + arith.index(rk) + m_row_abs = m_start + m_row_rel + in_b = arith.cmpi(arith.CmpIPredicate.slt, m_row_abs, seq_len_q_v) + m_row_safe2 = arith.select(in_b, m_row_abs, arith.index(0)) + g_idx = global_idx_q(qhid, m_row_safe2, d_col) + v1 = _gep_load(do_ptr, g_idx, T.vec(1, elem_type)) + v_scalar = vector.extract(v1, static_position=[0], dynamic_position=[]) + v_safe = arith.select(in_b, v_scalar, c_zero_elem) + vals.append(v_safe) + return vector.from_elements(mfma_pack_type, vals) + + new_dv_accs = list(dv_accs) + for dc in range_constexpr(D_CHUNKS): + for pks in range_constexpr(K_STEPS_PT): + b_pack = read_do_b_pack(pks, dc) + new_dv_accs[dc] = mfma_acc(pt_a_packs[pks], b_pack, new_dv_accs[dc]) + + # ---- GEMM3: dp = V @ DO^T (MFMA 16x16x32, one tile per m-tile mt) ---- + dp_accs = [v4f32_zero for _ in range(M_TILES)] + for ks in range_constexpr(K_STEPS_QK): + v_pack = vector.load_op(mfma_pack_type, lds_kv, [_v_idx_wave(ks)]) + for mt in range_constexpr(M_TILES): + dp_accs[mt] = mfma_acc(v_pack, do_b_packs_gemm3[mt][ks], dp_accs[mt]) + + # ---- dsT = pT * (dp - delta) per m-tile ---- + dsT_vals_per_tile = [] + for mt in range_constexpr(M_TILES): + ds_tile = [] + for ii in range_constexpr(4): + dp_ii = vector.extract(dp_accs[mt], static_position=[ii], dynamic_position=[]) + diff = arith.SubFOp(dp_ii, delta_per_tile[mt], fastmath=fm_fast).result + ds_ii = arith.MulFOp(pT_vals_per_tile[mt][ii], diff, fastmath=fm_fast).result + ds_tile.append(ds_ii) + dsT_vals_per_tile.append(ds_tile) + + # ---- dsT register -> LDS (per m-tile, col = mt*16 + lane_mod_16) ---- + for mt in range_constexpr(M_TILES): + for ii in range_constexpr(4): + kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii) + kv_row = wave_n_offset + kv_row_rel + ds_bf16 = arith.trunc_f(elem_type, dsT_vals_per_tile[mt][ii]) + lds_pt_idx = kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16 + v1_ds = vector.from_elements(v1_elem_type, [ds_bf16]) + vector.store(v1_ds, lds_pt, [lds_pt_idx]) + + # ---- dsT A-frag for GEMM4 (dK += dsT @ Q): + # A[kv_row=wave_n_offset+lane_mod_16, m_col=m_step*32+lane_div_16*8 + 0..7] + ds_a_packs = [] + for m_step in range_constexpr(K_STEPS_PT): + kv_row_a = wave_n_offset + lane_mod_16 + m_col_a = arith.index(m_step * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + lds_idx = kv_row_a * arith.index(LDS_PT_STRIDE) + m_col_a + pack = vector.load_op(mfma_pack_type, lds_pt, [lds_idx]) + ds_a_packs.append(pack) + + # ---- GEMM4: dK += dsT @ Q (MFMA 16x16x32) ---- + # B-frag: Q[m_step*32+lane_div_16*8+k, dc*16+lane_mod_16]. + if USE_LDS_FOR_Q_DO: + + def read_q_b_pack(m_step_idx, dc_idx): + d_col = arith.index(dc_idx * D_CHUNK) + lane_mod_16 + m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + vals = [] + for rk in range_constexpr(MFMA_LANE_K): + m_row = m_base + arith.index(rk) + lds_idx = m_row * arith.index(LDS_Q_STRIDE) + d_col + val = _memref.load(lds_q, [lds_idx]) + vals.append(val) + return vector.from_elements(mfma_pack_type, vals) + + else: + + def read_q_b_pack(m_step_idx, dc_idx): + d_col = arith.index(dc_idx * D_CHUNK) + lane_mod_16 + m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + vals = [] + for rk in range_constexpr(MFMA_LANE_K): + m_row_rel = m_base + arith.index(rk) + m_row_abs = m_start + m_row_rel + in_b = arith.cmpi(arith.CmpIPredicate.slt, m_row_abs, seq_len_q_v) + m_row_safe2 = arith.select(in_b, m_row_abs, arith.index(0)) + g_idx = global_idx_q(qhid, m_row_safe2, d_col) + v1 = _gep_load(q_ptr, g_idx, T.vec(1, elem_type)) + v_scalar = vector.extract(v1, static_position=[0], dynamic_position=[]) + v_safe = arith.select(in_b, v_scalar, c_zero_elem) + vals.append(v_safe) + return vector.from_elements(mfma_pack_type, vals) + + new_dk_accs = list(dk_accs) + for dc in range_constexpr(D_CHUNKS): + for pks in range_constexpr(K_STEPS_PT): + q_b_pack = read_q_b_pack(pks, dc) + new_dk_accs[dc] = mfma_acc(ds_a_packs[pks], q_b_pack, new_dk_accs[dc]) + + gpu.barrier() + + # R6e: yield (dv, dk) partial directly as head-loop carry; no inner m-loop. + yield list(new_dv_accs) + list(new_dk_accs) + outer_carry = list(h_loop_results) + + # ---- Final: dK *= sm_scale, store ---- + dv_finals = [outer_carry[dc] for dc in range(D_CHUNKS)] + dk_finals = [outer_carry[D_CHUNKS + dc] for dc in range(D_CHUNKS)] + + # R6e: atomic_fadd into shared pool slice of dk/dv. + # Multiple m-tile programs collide on the same (b, pool_n, d) addresses. + # DK/DV are fp32 contiguous; byte_offset = global_idx_kv * 4. + _atom_zero_i32 = arith.constant(0, type=T.i32) + _four_i32 = arith.constant(4, type=T.i32) + for dc in range_constexpr(D_CHUNKS): + for ii in range_constexpr(4): + kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii) + kv_row_abs = kv_start + wave_n_offset + kv_row_rel + d_col_abs = arith.index(dc * D_CHUNK) + lane_mod_16 + kv_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, kv_row_abs, pool_end) + _if_kv = scf.IfOp(kv_in_bounds) + with ir.InsertionPoint(_if_kv.then_block): + dv_val = vector.extract(dv_finals[dc], static_position=[ii], dynamic_position=[]) + dk_val = vector.extract(dk_finals[dc], static_position=[ii], dynamic_position=[]) + dk_scaled = arith.MulFOp(dk_val, c_sm_scale, fastmath=fm_fast).result + g_elem_idx = global_idx_kv(kv_row_abs, d_col_abs) + g_elem_i32 = arith.index_cast(T.i32, g_elem_idx) + byte_off_i32 = arith.MulIOp(g_elem_i32, _four_i32).result + rocdl.raw_ptr_buffer_atomic_fadd( + dv_val, + dv_rsrc, + byte_off_i32, + _atom_zero_i32, + _atom_zero_i32, + ) + rocdl.raw_ptr_buffer_atomic_fadd( + dk_scaled, + dk_rsrc, + byte_off_i32, + _atom_zero_i32, + _atom_zero_i32, + ) + scf.YieldOp([]) + + @flyc.jit + def launch_v4_hca_bwd_dkv_pool( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + DOS: fx.Tensor, + LSE: fx.Tensor, + DELTAS: fx.Tensor, + DK: fx.Tensor, + DV: fx.Tensor, + ADD_MASK: fx.Tensor, + batch_size: fx.Int32, + seq_len_q: fx.Int32, + seq_len_k: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + bs_idx = arith.index_cast(T.index, batch_size) + # R6e: grid_x = B * num_m_blocks. Parallelize over m-tiles for B=1. + sq_idx_h = arith.index_cast(T.index, seq_len_q) + BM2_idx_h = arith.index(BLOCK_M2) + _one_idx_h = arith.index(1) + num_m_blocks_v = (sq_idx_h + BM2_idx_h - _one_idx_h) // BM2_idx_h + grid_x = bs_idx * num_m_blocks_v + + launcher = v4_hca_bwd_dkv_pool_kernel( + Q, + K, + V, + DOS, + LSE, + DELTAS, + DK, + DV, + ADD_MASK, + seq_len_q, + seq_len_k, + ) + + if waves_per_eu is not None: + _wpe = int(waves_per_eu) + if _wpe >= 1: + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe) + if flat_work_group_size is not None: + _fwgs = int(flat_work_group_size) + if _fwgs >= 1: + flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}") + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr + + passthrough_entries = [] + if daz: + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("denormal-fp-math-f32"), + ir.StringAttr.get("preserve-sign,preserve-sign"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("no-nans-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("unsafe-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries) + + launcher.launch( + grid=(grid_x, 1, 1), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) + + _fmha_compile_hints = { + "fast_fp_math": fast_fp_math, + "unsafe_fp_math": unsafe_fp_math, + "llvm_options": { + "enable-post-misched": False, + "lsr-drop-solution": True, + }, + } + + def _launch(*args, **kwargs): + with CompilationContext.compile_hints(_fmha_compile_hints): + return launch_v4_hca_bwd_dkv_pool(*args, **kwargs) + + def _compile(Q, K, V, DOS, LSE, DELTAS, DK, DV, ADD_MASK, batch_size, seq_len_q, seq_len_k, stream=None): + with CompilationContext.compile_hints(_fmha_compile_hints): + return flyc.compile( + launch_v4_hca_bwd_dkv_pool, + Q, + K, + V, + DOS, + LSE, + DELTAS, + DK, + DV, + ADD_MASK, + batch_size, + seq_len_q, + seq_len_k, + fx.Stream(stream), + ) + + _launch.compile = _compile + + return _launch + + +# Convenience alias. +build_v4_hca_bwd_dkv_pool_module_primary = build_v4_hca_bwd_dkv_pool_module diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dq_pool_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dq_pool_kernel.py new file mode 100644 index 000000000..f69001491 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_hca_bwd_dq_pool_kernel.py @@ -0,0 +1,1404 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""v4_hca_bwd_dq_pool: V4 HCA backward dQ POOL-stream kernel for FlyDSL. + +Forked from v4_sla_bwd_dq_kernel.py. Computes the POOL stream contribution +to dq for an HCA (Hybrid-Causal-Additive) attention backward, then +ACCUMULATES into an existing dq_fp32 buffer (which already contains the +LOCAL stream dq from the SWA dq kernel). + +Differences vs the SWA kernel: + - KV range is fixed at [HCA_LOCAL_SEQLEN, HCA_LOCAL_SEQLEN+POOL_SIZE). + POOL_SIZE is a build-time constexpr. For our shapes POOL_SIZE <= BLOCK_N + so there is exactly ONE n-block iteration. + - No SWA-window mask, no causal mask. Two element-wise predicates only: + (a) pool_n < POOL_SIZE -> NEG_INF outside the pool + (b) q_row < seq_len_q -> NEG_INF for OOB q rows + - qk + add_bias from the ADD_MASK tensor (shape [Sq, POOL_SIZE]). + Each element loaded as f32 (after cast from bf16/f16) and added to qk + before -lse and exp. + - No SINK / DSINK. Sink is handled by the LOCAL FlyDSL dq kernel; the + pool stream does not touch the sink. + - Final store ACCUMULATES into DQ (load + add + store). Race-free + because each program owns a unique (b, qhid, m_block) slice and the + pool kernel runs AFTER the SWA dq has finished writing dq for that + same slice (sequential launches from the wrapper). + - sm_scale is applied INSIDE the loop on qk (matches Triton ref); after + the n-loop (which has only one iter), sm_scale is applied ONCE MORE + on dq before accumulating into DQ (P57 cr=0 BWD). + +Layout: BHLD. Q/DOUT/DQ flat from (B, HQ, Sq, D). K/V flat from + (B, HK, Sk, D); MQA means HK=1, no head stride applied to KV. +LSE/DELTAS: (B, HQ, Sq) flat, fp32, raw-domain (JOINT lse from fwd). +ADD_MASK: (Sq, POOL_SIZE) bf16/f16; loaded as f32 inside the kernel. +DQ: (B, HQ, Sq, D) fp32; ACCUMULATED. +""" + +import math +import os + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import math as math_dialect +from flydsl._mlir.dialects import memref as _memref +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr +from kernels.kernels_common import dtype_to_elem_type + +# ---- Module-level constants ---- + +KERNEL_NAME = "v4_hca_bwd_dq_pool_kernel" + +_LOG2E = math.log2(math.e) # 1.4426950408889634 + +_LLVM_GEP_DYNAMIC = -2147483648 # LLVM kDynamicIndex sentinel (0x80000000 as signed i32) + + +def _llvm_ptr_ty(): + return ir.Type.parse("!llvm.ptr") + + +def _llvm_lds_ptr_ty(): + return ir.Type.parse("!llvm.ptr<3>") + + +_VMCNT_LO_MASK = 0xF +_LGKMCNT_EXPCNT_BASE = 0x3F70 +_VMCNT_HI_SHIFT = 14 +_VMCNT_HI_MASK = 0x3 + + +def _waitcnt_vm_n(n): + """Emit s_waitcnt vmcnt(n) only (lgkmcnt=63, expcnt=7).""" + val = (n & _VMCNT_LO_MASK) | _LGKMCNT_EXPCNT_BASE | (((n >> 4) & _VMCNT_HI_MASK) << _VMCNT_HI_SHIFT) + rocdl.s_waitcnt(val) + + +def build_v4_hca_bwd_dq_pool_module( + num_heads, + head_dim, + pool_size, + hca_local_seqlen, + dtype_str="bf16", + sm_scale=None, + waves_per_eu=2, + flat_work_group_size=None, + block_m=None, + unsafe_fp_math=True, + fast_fp_math=True, + daz=True, + layout_bhld=True, + mqa_kv=True, +): + """Build the V4 HCA backward dQ POOL-stream launcher. + + Args: + num_heads: HQ (number of Q heads). + head_dim: D (head dimension), must be % 32 == 0 and >= 64. + pool_size: int > 0; number of pool keys (<= BLOCK_N=64 in this + kernel; one n-block per program iter). Build-time constexpr. + hca_local_seqlen: int >= 0; offset into K/V at which the pool + keys start (== Sq for HCA split-mask). Build-time constexpr; + must be a multiple of BLOCK_N to keep the single-iter + n-block aligned. + sm_scale: defaults to 1/sqrt(head_dim). + mqa_kv: if True, K/V indexing drops head_idx (HK=1, stride_h=0). + layout_bhld: BHLD layout (True) or BLHD (False); V4 uses BHLD. + + Returns: launch(Q, K, V, DOUT, LSE, DELTAS, DQ_FP32, ADD_MASK, + batch_size, seq_len_q, seq_len_k, stream=None). + DQ_FP32 is ACCUMULATED into (load + add + store). + """ + gpu_arch = get_hip_arch() + + BLOCK_N = 64 + K_SUB_N = 32 + WARP_SIZE = 64 + + if block_m is not None: + BLOCK_M = block_m + else: + BLOCK_M = 128 + + if flat_work_group_size is None: + if BLOCK_M <= 128: + flat_work_group_size = 256 + else: + flat_work_group_size = 512 + NUM_WAVES = flat_work_group_size // WARP_SIZE + BLOCK_SIZE = flat_work_group_size + ROWS_PER_WAVE = BLOCK_M // NUM_WAVES + # V4 SWA: dense path. One outer iter = one BLOCK_N block. + BLOCK_N_OUT = BLOCK_N + ENABLE_PREFETCH_3BUF = os.getenv("FLYDSL_SLA_FWD_ENABLE_PREFETCH3", "0") == "1" + _has_lds_load_b128 = not gpu_arch.startswith("gfx942") + ENABLE_DMA = _has_lds_load_b128 and (os.getenv("FLYDSL_SLA_FWD_ENABLE_DMA", "1") == "1") + ENABLE_LDS_VEC16 = os.getenv("FLYDSL_SLA_FWD_ENABLE_LDS_VEC16", "1") == "1" + REDUCE_MODE = os.getenv("FLYDSL_SLA_FWD_REDUCE_MODE", "xor").strip().lower() + if REDUCE_MODE not in ("xor", "ds_bpermute"): + REDUCE_MODE = "xor" + NUM_PREFETCH_K = 3 if ENABLE_PREFETCH_3BUF else (2 if ENABLE_DMA else 1) + NUM_PREFETCH_V = 3 if ENABLE_PREFETCH_3BUF else (2 if ENABLE_DMA else 1) + (1, 2, 0, 1, 0, 1, 2, 0) if ENABLE_PREFETCH_3BUF else (0,) + + USE_HW_TR = gpu_arch.startswith("gfx950") + USE_K16 = gpu_arch.startswith("gfx950") + + # Auto-fallback: gfx950 LDS limit is 160 KB. K+V double-buffering at D=512 + # easily exceeds that. If predicted LDS exceeds budget, force DMA off + # (NUM_PREFETCH = 1) and warn. Keeps the kernel correct regardless of + # the env knob. + _LDS_LIMIT_BYTES = 160 * 1024 + + def _predicted_lds_bytes(nk, nv, dma): + # USE_HW_TR & DMA -> V_STRIDE = head_dim + # USE_HW_TR & !DMA -> V_STRIDE = head_dim + 4 + # !USE_HW_TR -> V stored transposed: LDS V = head_dim * (BLOCK_N + 2) + if USE_HW_TR: + v_str = head_dim if dma else head_dim + 4 + return (nk * BLOCK_N * head_dim + nv * BLOCK_N * v_str) * 2 + vt = BLOCK_N + 2 + return (nk * BLOCK_N * head_dim + nv * head_dim * vt) * 2 + + _pred = _predicted_lds_bytes(NUM_PREFETCH_K, NUM_PREFETCH_V, ENABLE_DMA) + if _pred > _LDS_LIMIT_BYTES and ENABLE_DMA: + # Try DMA off (single-buffered). + ENABLE_DMA = False + NUM_PREFETCH_K = 1 + NUM_PREFETCH_V = 1 + _pred2 = _predicted_lds_bytes(1, 1, False) + if _pred2 > _LDS_LIMIT_BYTES: + raise RuntimeError( + f"v4_hca_bwd_dq_pool: predicted LDS {_pred2}B > limit {_LDS_LIMIT_BYTES}B " + f"even single-buffered at D={head_dim}, BLOCK_N={BLOCK_N}" + ) + import sys as _sys + + print( + f"[v4_hca_bwd_dq_pool] LDS overflow at D={head_dim}, BLOCK_N={BLOCK_N}: " + f"auto-disabled DMA (predicted {_pred} -> {_pred2} bytes)", + file=_sys.stderr, + flush=True, + ) + K_STEP_QK = 16 if USE_K16 else 8 + K_STEPS_QK = head_dim // K_STEP_QK + D_CHUNK = 32 + D_CHUNKS = head_dim // D_CHUNK + PV_K_STEP = 16 if USE_K16 else 8 + PV_K_STEPS = K_SUB_N // PV_K_STEP # 2 steps per sub-tile (K=16) or 4 (K=8) + + assert BLOCK_M % NUM_WAVES == 0 + assert head_dim % 32 == 0, f"head_dim ({head_dim}) must be divisible by 32" + assert head_dim >= 64, f"head_dim ({head_dim}) must be >= 64" + assert flat_work_group_size in ( + 128, + 256, + 512, + ), f"flat_work_group_size must be 128, 256, or 512, got {flat_work_group_size}" + assert dtype_str in ("f16", "bf16"), "v4_hca_bwd_dq_pool only supports f16 and bf16" + assert BLOCK_N % 32 == 0 + assert BLOCK_N_OUT == BLOCK_N + assert ( + isinstance(pool_size, int) and 0 < pool_size <= BLOCK_N + ), f"pool_size must be int in (0, {BLOCK_N}], got {pool_size!r}" + assert ( + isinstance(hca_local_seqlen, int) and hca_local_seqlen >= 0 + ), f"hca_local_seqlen must be int >= 0, got {hca_local_seqlen!r}" + assert ( + hca_local_seqlen % BLOCK_N == 0 + ), f"hca_local_seqlen must be multiple of BLOCK_N={BLOCK_N}, got {hca_local_seqlen}" + + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(head_dim) + + NUM_HEADS = num_heads + HEAD_DIM = head_dim + STRIDE_TOKEN = NUM_HEADS * HEAD_DIM + + K_STRIDE = HEAD_DIM + # XOR swizzle mask must fit within the row stride. The swizzle is + # applied at 16-element granularity, so the maximum mask is + # `(K_STRIDE // 16 - 1) << 4`. For D=128 that's 7 (=0x7), giving max + # mask 112 < 128. For D=64 it's 3 (=0x3), giving max mask 48 < 64. + # SLA test only covers D=128 (mask=7); using a hardcoded `& 7` for + # D=64 wraps writes into adjacent rows -> silent corruption. + K_SWZ_ROW_MASK = (K_STRIDE // 16) - 1 + assert K_SWZ_ROW_MASK >= 0 + assert ( + K_SWZ_ROW_MASK & (K_SWZ_ROW_MASK + 1) + ) == 0, f"K_SWZ_ROW_MASK must be 2^n-1, got {K_SWZ_ROW_MASK} (K_STRIDE={K_STRIDE})" + if USE_HW_TR: + V_STRIDE = HEAD_DIM if ENABLE_DMA else HEAD_DIM + 4 + else: + VT_STRIDE = BLOCK_N + 2 + V_STRIDE = VT_STRIDE + # V swizzle: similarly bounded by V_STRIDE / 16. + V_SWZ_ROW_MASK = min(3, (V_STRIDE // 16) - 1) + assert V_SWZ_ROW_MASK >= 0 + + VEC_WIDTH = 16 if ENABLE_LDS_VEC16 else 8 + assert HEAD_DIM % VEC_WIDTH == 0 + THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH + assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0 + ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD + + if ROWS_PER_BATCH_LOAD >= BLOCK_N: + NUM_BATCHES_KV = 1 + KV_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_N + else: + assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0 + NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD + KV_NEEDS_GUARD = False + + LDS_K_TILE_SIZE = BLOCK_N * K_STRIDE + if USE_HW_TR: + LDS_V_TILE_SIZE = BLOCK_N * V_STRIDE + else: + LDS_V_TILE_SIZE = HEAD_DIM * VT_STRIDE + LDS_K_TOTAL_SIZE = NUM_PREFETCH_K * LDS_K_TILE_SIZE + LDS_V_BASE = LDS_K_TOTAL_SIZE + LDS_V_TOTAL_SIZE = NUM_PREFETCH_V * LDS_V_TILE_SIZE + LDS_KV_TOTAL_SIZE = LDS_K_TOTAL_SIZE + LDS_V_TOTAL_SIZE + + allocator = SmemAllocator( + None, + arch=gpu_arch, + global_sym_name=f"v4_hca_bwd_dq_pool_smem_M{BLOCK_M}_P{pool_size}_L{hca_local_seqlen}_MQ{int(mqa_kv)}", + ) + lds_kv_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_kv_offset + LDS_KV_TOTAL_SIZE * 2 + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def v4_hca_bwd_dq_pool_kernel( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + DOS: fx.Tensor, # grad of output (input) + LSE: fx.Tensor, # log-sum-exp from fwd (input, f32, RAW domain, JOINT) + DELTAS: fx.Tensor, # (o * do).sum(-1) preprocess (input, f32) + DQ: fx.Tensor, # grad wrt Q (output, FP32, ACCUMULATED) + ADD_MASK: fx.Tensor, # additive pool mask (input, bf16/f16, [Sq, POOL_SIZE]) + seq_len_q: fx.Int32, + seq_len_k: fx.Int32, + ): + elem_type = dtype_to_elem_type(dtype_str) + compute_type = T.f32 + q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q) + k_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K) + v_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V) + do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOS) + dq_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DQ) + add_mask_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), ADD_MASK) + # LSE / DELTAS: f32 scalar READS via buffer_load. + lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True) + deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True) + + # All FP operations use aggressive fast-math (no NaN/Inf checks, reassociation). + fm_fast = arith.FastMathFlags.fast + v4f16_type = T.vec(4, elem_type) + vxf16_type = T.vec(VEC_WIDTH, elem_type) + v8f16_type = T.vec(8, elem_type) + v16f32_type = T.vec(16, compute_type) + mfma_pack_type = v8f16_type if USE_K16 else v4f16_type + MFMA_LANE_K = 8 if USE_K16 else 4 + _mfma_zero = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0) + + def _mfma(ods_fn, a, b, c): + return ods_fn(v16f32_type, a, b, c, _mfma_zero, _mfma_zero, _mfma_zero).result + + def mfma_acc(a, b, c): + if dtype_str == "bf16": + if USE_K16: + return _mfma(rocdl.mfma_f32_32x32x16_bf16, a, b, c) + a = vector.bitcast(T.i16x4, a) + b = vector.bitcast(T.i16x4, b) + return _mfma(rocdl.mfma_f32_32x32x8bf16_1k, a, b, c) + if USE_K16: + return _mfma(rocdl.mfma_f32_32x32x16_f16, a, b, c) + return _mfma(rocdl.mfma_f32_32x32x8f16, a, b, c) + + seq_len_q_v = arith.index_cast(T.index, seq_len_q) + seq_len_k_v = arith.index_cast(T.index, seq_len_k) + + # ---- LDS view ---- + base_ptr = allocator.get_base() + lds_kv = SmemPtr( + base_ptr, + lds_kv_offset, + elem_type, + shape=(LDS_KV_TOTAL_SIZE,), + ).get() + + # ---- Thread / block indices ---- + block_id = arith.index_cast(T.index, gpu.block_idx.x) + tid = arith.index_cast(T.index, gpu.thread_idx.x) + + wave_id = tid // WARP_SIZE + lane = tid % WARP_SIZE + lane_mod_32 = lane % 32 + lane_div_32 = lane // 32 # 0/1 + + # ds_read_b64_tr_b16 lane decomposition + tr_k_group = (lane % 16) // 4 + tr_col_sub = lane % 4 + tr_col_half = (lane % 32) // 16 + + def ds_read_tr_v4f16(lds_elem_idx): + byte_offset = lds_elem_idx * 2 + lds_kv_offset + byte_i64 = arith.index_cast(T.i64, byte_offset) + ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result + return rocdl.ds_read_tr16_b64(v4f16_type, ptr).result + + wave_q_offset = wave_id * ROWS_PER_WAVE + + # ---- Decompose block_id: (batch, q_tile, head) ---- + head_idx = block_id % NUM_HEADS + batch_q_tile_id = block_id // NUM_HEADS + num_q_tiles = (seq_len_q_v + BLOCK_M - 1) // BLOCK_M + q_tile_idx = batch_q_tile_id % num_q_tiles + batch_idx = batch_q_tile_id // num_q_tiles + q_start = q_tile_idx * BLOCK_M + + # ---- V4 HCA POOL: fixed K-block range [HCA_LOCAL_SEQLEN, HCA_LOCAL_SEQLEN+POOL_SIZE) ---- + # POOL_SIZE <= BLOCK_N and HCA_LOCAL_SEQLEN % BLOCK_N == 0 (asserted + # at build time), so this is exactly ONE n-block per program. + BN = arith.index(BLOCK_N) + arith.index(BLOCK_M) + arith.index(0) + _one_idx = arith.index(1) + n_block_start = arith.index(hca_local_seqlen // BLOCK_N) + n_block_end = n_block_start + _one_idx + + # ---- Cooperative load decomposition ---- + load_row_in_batch = tid // THREADS_PER_ROW_LOAD + load_lane_in_row = tid % THREADS_PER_ROW_LOAD + load_col_base = load_lane_in_row * VEC_WIDTH + + # ---- Helper: global flat index ---- + if layout_bhld: + bh_base_tokens_q = (batch_idx * NUM_HEADS + head_idx) * seq_len_q_v + if mqa_kv: + bh_base_tokens_kv = batch_idx * seq_len_k_v + else: + bh_base_tokens_kv = (batch_idx * NUM_HEADS + head_idx) * seq_len_k_v + + def global_idx_q(token_idx, col): + return (bh_base_tokens_q + token_idx) * arith.index(HEAD_DIM) + col + + def global_idx_kv(token_idx, col): + return (bh_base_tokens_kv + token_idx) * arith.index(HEAD_DIM) + col + + else: + # BLHD path: kept for symmetry but V4 uses BHLD. + def global_idx_q(token_idx, col): + token = batch_idx * seq_len_q_v + token_idx + return token * STRIDE_TOKEN + head_idx * HEAD_DIM + col + + if mqa_kv: + + def global_idx_kv(token_idx, col): + token = batch_idx * seq_len_k_v + token_idx + return token * HEAD_DIM + col + + else: + + def global_idx_kv(token_idx, col): + token = batch_idx * seq_len_k_v + token_idx + return token * STRIDE_TOKEN + head_idx * HEAD_DIM + col + + def _gep_load(base_ptr_, elem_idx, vec_type, et=elem_type): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr_, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=et, + noWrapFlags=0, + ) + return _llvm.LoadOp(vec_type, gep.result).result + + def _gep_store_f32(val, base_ptr_, elem_idx): + """Store a single f32 value via GEP into an fp32 buffer.""" + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr_, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=T.f32, + noWrapFlags=0, + ) + _llvm.StoreOp(val, gep.result) + + def load_global_mfma_pack(base_ptr_, base_idx): + return _gep_load(base_ptr_, base_idx, mfma_pack_type) + + def load_global_f16xN(base_ptr_, base_idx): + return _gep_load(base_ptr_, base_idx, vxf16_type) + + def bf16_trunc_pack_v4(f32_vals): + _v2i32 = T.vec(2, T.i32) + _c16 = arith.constant(16, type=T.i32) + _cmask = arith.constant(0xFFFF0000, type=T.i32) + a0 = arith.ArithValue(f32_vals[0]).bitcast(T.i32) + b0 = arith.ArithValue(f32_vals[1]).bitcast(T.i32) + p0 = arith.OrIOp(arith.AndIOp(b0, _cmask).result, arith.ShRUIOp(a0, _c16).result).result + a1 = arith.ArithValue(f32_vals[2]).bitcast(T.i32) + b1 = arith.ArithValue(f32_vals[3]).bitcast(T.i32) + p1 = arith.OrIOp(arith.AndIOp(b1, _cmask).result, arith.ShRUIOp(a1, _c16).result).result + return vector.bitcast(v4f16_type, vector.from_elements(_v2i32, [p0, p1])) + + def bf16_trunc_pack_v8(f32_vals): + _v4i32 = T.vec(4, T.i32) + _c16 = arith.constant(16, type=T.i32) + _cmask = arith.constant(0xFFFF0000, type=T.i32) + pairs = [] + for j in range_constexpr(4): + a = arith.ArithValue(f32_vals[j * 2]).bitcast(T.i32) + b = arith.ArithValue(f32_vals[j * 2 + 1]).bitcast(T.i32) + p = arith.OrIOp(arith.AndIOp(b, _cmask).result, arith.ShRUIOp(a, _c16).result).result + pairs.append(p) + return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs)) + + def k_buf_base(buf_id): + if isinstance(buf_id, int): + return arith.index(buf_id * LDS_K_TILE_SIZE) + return buf_id * arith.index(LDS_K_TILE_SIZE) + + def v_buf_base(buf_id): + if isinstance(buf_id, int): + return arith.index(LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) + return arith.index(LDS_V_BASE) + buf_id * arith.index(LDS_V_TILE_SIZE) + + # ---- K XOR swizzle: col ^ ((row & K_SWZ_ROW_MASK) << 4) at 16-element granularity ---- + # K_SWZ_ROW_MASK derived from K_STRIDE to stay within the row. + def _k_swizzle(row_idx, col_idx): + mask = (row_idx & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + return col_idx ^ mask + + # ---- Cooperative K load (row-major, XOR-swizzled) ---- + def coop_load_k(tile_start, buf_id=0): + """K row-bounds-aware cooperative load. + + Pool kernel reads BLOCK_N rows starting at kv_start; for pool + sizes < BLOCK_N the top rows are past seq_len_k and would read + garbage. Gate each load with row_idx < seq_len_k_v and store + zero into the LDS slot for OOB rows. + """ + k_base = k_buf_base(buf_id) + c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid_blk = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + row_valid_seq = arith.cmpi( + arith.CmpIPredicate.ult, + row_idx, + seq_len_k_v, + ) + row_valid = arith.AndIOp(row_valid_blk, row_valid_seq).result + lds_row = load_row_in_batch + row_offset + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0)) + g_idx = global_idx_kv(g_idx_safe_row, load_col_base) + raw_vec = load_global_f16xN(k_ptr, g_idx) + vec = arith.select(row_valid, raw_vec, c_zero_vxf16) + _if_k = scf.IfOp( + arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + ) + with ir.InsertionPoint(_if_k.then_block): + vector.store(vec, lds_kv, [lds_idx]) + scf.YieldOp([]) + else: + row_valid_seq = arith.cmpi( + arith.CmpIPredicate.ult, + row_idx, + seq_len_k_v, + ) + lds_row = load_row_in_batch + row_offset + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0)) + g_idx = global_idx_kv(g_idx_safe_row, load_col_base) + raw_vec = load_global_f16xN(k_ptr, g_idx) + vec = arith.select(row_valid_seq, raw_vec, c_zero_vxf16) + vector.store(vec, lds_kv, [lds_idx]) + + # ---- Cooperative V-into-K-LDS-slot load (K-style XOR swizzle) ---- + def coop_load_v_as_k(tile_start, buf_id=0): + """V-into-K-slot with row bounds check (zero OOB rows).""" + k_base = k_buf_base(buf_id) + c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid_blk = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + row_valid_seq = arith.cmpi( + arith.CmpIPredicate.ult, + row_idx, + seq_len_k_v, + ) + row_valid = arith.AndIOp(row_valid_blk, row_valid_seq).result + lds_row = load_row_in_batch + row_offset + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0)) + g_idx = global_idx_kv(g_idx_safe_row, load_col_base) + raw_vec = load_global_f16xN(v_ptr, g_idx) + vec = arith.select(row_valid, raw_vec, c_zero_vxf16) + _if_v = scf.IfOp( + arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + ) + with ir.InsertionPoint(_if_v.then_block): + vector.store(vec, lds_kv, [lds_idx]) + scf.YieldOp([]) + else: + row_valid_seq = arith.cmpi( + arith.CmpIPredicate.ult, + row_idx, + seq_len_k_v, + ) + lds_row = load_row_in_batch + row_offset + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0)) + g_idx = global_idx_kv(g_idx_safe_row, load_col_base) + raw_vec = load_global_f16xN(v_ptr, g_idx) + vec = arith.select(row_valid_seq, raw_vec, c_zero_vxf16) + vector.store(vec, lds_kv, [lds_idx]) + + # ---- Cooperative V load (V LDS layout) ---- + def _v_store_row_major(v_base, lds_row, vec): + lds_idx = v_base + lds_row * V_STRIDE + load_col_base + vector.store(vec, lds_kv, [lds_idx]) + + _v1_type = T.vec(1, elem_type) if not USE_HW_TR else None + + def _v_store_transposed(v_base, lds_row, vec): + for _e in range_constexpr(VEC_WIDTH): + elem = vector.extract(vec, static_position=[_e], dynamic_position=[]) + vt_d = load_col_base + _e + vt_idx = v_base + vt_d * VT_STRIDE + lds_row + v1 = vector.from_elements(_v1_type, [elem]) + vector.store(v1, lds_kv, [vt_idx]) + + _v_store_to_lds = _v_store_row_major if USE_HW_TR else _v_store_transposed + + def coop_load_v(tile_start, buf_id=0): + v_base = v_buf_base(buf_id) + c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid_blk = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + row_valid_seq = arith.cmpi( + arith.CmpIPredicate.ult, + row_idx, + seq_len_k_v, + ) + row_valid = arith.AndIOp(row_valid_blk, row_valid_seq).result + g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0)) + g_idx = global_idx_kv(g_idx_safe_row, load_col_base) + raw_vec = load_global_f16xN(v_ptr, g_idx) + vec = arith.select(row_valid, raw_vec, c_zero_vxf16) + _if_v = scf.IfOp( + arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + ) + with ir.InsertionPoint(_if_v.then_block): + lds_row = load_row_in_batch + row_offset + _v_store_to_lds(v_base, lds_row, vec) + scf.YieldOp([]) + else: + row_valid_seq = arith.cmpi( + arith.CmpIPredicate.ult, + row_idx, + seq_len_k_v, + ) + g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0)) + g_idx = global_idx_kv(g_idx_safe_row, load_col_base) + raw_vec = load_global_f16xN(v_ptr, g_idx) + vec = arith.select(row_valid_seq, raw_vec, c_zero_vxf16) + lds_row = load_row_in_batch + row_offset + _v_store_to_lds(v_base, lds_row, vec) + + # ---- Cooperative K-into-V-LDS-slot load ---- + def coop_load_k_as_v(tile_start, buf_id=0): + """K-into-V-slot with row bounds check (zero OOB rows).""" + v_base = v_buf_base(buf_id) + c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid_blk = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + row_valid_seq = arith.cmpi( + arith.CmpIPredicate.ult, + row_idx, + seq_len_k_v, + ) + row_valid = arith.AndIOp(row_valid_blk, row_valid_seq).result + g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0)) + g_idx = global_idx_kv(g_idx_safe_row, load_col_base) + raw_vec = load_global_f16xN(k_ptr, g_idx) + vec = arith.select(row_valid, raw_vec, c_zero_vxf16) + _if_v = scf.IfOp( + arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + ) + with ir.InsertionPoint(_if_v.then_block): + lds_row = load_row_in_batch + row_offset + _v_store_to_lds(v_base, lds_row, vec) + scf.YieldOp([]) + else: + row_valid_seq = arith.cmpi( + arith.CmpIPredicate.ult, + row_idx, + seq_len_k_v, + ) + g_idx_safe_row = arith.select(row_valid_seq, row_idx, arith.index(0)) + g_idx = global_idx_kv(g_idx_safe_row, load_col_base) + raw_vec = load_global_f16xN(k_ptr, g_idx) + vec = arith.select(row_valid_seq, raw_vec, c_zero_vxf16) + lds_row = load_row_in_batch + row_offset + _v_store_to_lds(v_base, lds_row, vec) + + # ---- DMA loading for K (buffer_load_dwordx4 ... lds) ---- + if ENABLE_DMA: + from flydsl._mlir.dialects import llvm + + k_rsrc = buffer_ops.create_buffer_resource(K, max_size=True) + _lds_ptr_ty = _llvm_lds_ptr_ty() + DMA_BYTES = 16 + DMA_BATCH_BYTES = BLOCK_SIZE * DMA_BYTES + K_TILE_BYTES = BLOCK_N * K_STRIDE * 2 + NUM_DMA_K = K_TILE_BYTES // DMA_BATCH_BYTES + LANES_PER_K_ROW = HEAD_DIM * 2 // DMA_BYTES + ROWS_PER_DMA_BATCH = DMA_BATCH_BYTES // (HEAD_DIM * 2) + lds_kv_base_idx = _memref.extract_aligned_pointer_as_index(lds_kv) + _dma_size = arith.constant(DMA_BYTES, type=T.i32) + _dma_soff = arith.constant(0, type=T.i32) + _dma_off = arith.constant(0, type=T.i32) + _dma_aux = arith.constant(1, type=T.i32) + + def _kv_global_byte(tile_start, row_in_tile, col_byte): + if layout_bhld: + row_within = tile_start + row_in_tile + return (bh_base_tokens_kv + row_within) * arith.index(HEAD_DIM * 2) + col_byte + else: + global_row = batch_idx * seq_len_k_v + tile_start + row_in_tile + if mqa_kv: + return global_row * arith.index(HEAD_DIM * 2) + col_byte + else: + return ( + global_row * arith.index(STRIDE_TOKEN * 2) + + head_idx * arith.index(HEAD_DIM * 2) + + col_byte + ) + + def coop_dma_k(tile_start, buf_id=0): + if isinstance(buf_id, int): + k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2) + else: + k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2) + for d in range_constexpr(NUM_DMA_K): + lds_addr = ( + k_lds_byte_base + + wave_id * arith.index(WARP_SIZE * DMA_BYTES) + + arith.index(d * DMA_BATCH_BYTES) + ) + lds_i64 = arith.index_cast(T.i64, lds_addr) + lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64) + lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result + + row_in_tile = tid // LANES_PER_K_ROW + arith.index(d * ROWS_PER_DMA_BATCH) + swiz_col_f16 = (tid % LANES_PER_K_ROW) * (DMA_BYTES // 2) + xor_mask = (row_in_tile & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + unsw_col_f16 = swiz_col_f16 ^ xor_mask + col_byte = unsw_col_f16 * 2 + global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte) + voffset = arith.index_cast(T.i32, global_byte) + rocdl.raw_ptr_buffer_load_lds( + k_rsrc, + lds_ptr, + _dma_size, + voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + def _v_swizzle(row_idx, col_idx): + mask = (row_idx & arith.index(V_SWZ_ROW_MASK)) << arith.index(4) + return col_idx ^ mask + + if ENABLE_DMA: + v_rsrc = buffer_ops.create_buffer_resource(V, max_size=True) + V_TILE_BYTES = BLOCK_N * V_STRIDE * 2 + NUM_DMA_V = V_TILE_BYTES // DMA_BATCH_BYTES + LANES_PER_V_ROW = HEAD_DIM * 2 // DMA_BYTES + ROWS_PER_DMA_BATCH_V = DMA_BATCH_BYTES // (HEAD_DIM * 2) + + def coop_dma_v(tile_start, buf_id=0): + v_lds_byte_base = lds_kv_base_idx + arith.index((LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) * 2) + for d in range_constexpr(NUM_DMA_V): + lds_addr = ( + v_lds_byte_base + + wave_id * arith.index(WARP_SIZE * DMA_BYTES) + + arith.index(d * DMA_BATCH_BYTES) + ) + lds_i64 = arith.index_cast(T.i64, lds_addr) + lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64) + lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result + + row_in_tile = tid // LANES_PER_V_ROW + arith.index(d * ROWS_PER_DMA_BATCH_V) + swiz_col_f16 = (tid % LANES_PER_V_ROW) * (DMA_BYTES // 2) + xor_mask = (row_in_tile & arith.index(V_SWZ_ROW_MASK)) << arith.index(4) + unsw_col_f16 = swiz_col_f16 ^ xor_mask + col_byte = unsw_col_f16 * 2 + global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte) + voffset = arith.index_cast(T.i32, global_byte) + rocdl.raw_ptr_buffer_load_lds( + v_rsrc, + lds_ptr, + _dma_size, + voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + # ---- Bwd dQ DMA variants (cross-pointer, matching swizzle) ---- + def coop_dma_v_as_k(tile_start, buf_id=0): + if isinstance(buf_id, int): + k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2) + else: + k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2) + for d in range_constexpr(NUM_DMA_K): + lds_addr = ( + k_lds_byte_base + + wave_id * arith.index(WARP_SIZE * DMA_BYTES) + + arith.index(d * DMA_BATCH_BYTES) + ) + lds_i64 = arith.index_cast(T.i64, lds_addr) + lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64) + lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result + + row_in_tile = tid // LANES_PER_K_ROW + arith.index(d * ROWS_PER_DMA_BATCH) + swiz_col_f16 = (tid % LANES_PER_K_ROW) * (DMA_BYTES // 2) + xor_mask = (row_in_tile & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + unsw_col_f16 = swiz_col_f16 ^ xor_mask + col_byte = unsw_col_f16 * 2 + global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte) + voffset = arith.index_cast(T.i32, global_byte) + rocdl.raw_ptr_buffer_load_lds( + v_rsrc, + lds_ptr, + _dma_size, + voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + def coop_dma_k_as_v(tile_start, buf_id=0): + if isinstance(buf_id, int): + v_lds_byte_base = lds_kv_base_idx + arith.index( + (LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) * 2 + ) + else: + v_lds_byte_base = ( + lds_kv_base_idx + + arith.index(LDS_V_BASE * 2) + + buf_id * arith.index(LDS_V_TILE_SIZE * 2) + ) + for d in range_constexpr(NUM_DMA_V): + lds_addr = ( + v_lds_byte_base + + wave_id * arith.index(WARP_SIZE * DMA_BYTES) + + arith.index(d * DMA_BATCH_BYTES) + ) + lds_i64 = arith.index_cast(T.i64, lds_addr) + lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64) + lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result + + row_in_tile = tid // LANES_PER_V_ROW + arith.index(d * ROWS_PER_DMA_BATCH_V) + swiz_col_f16 = (tid % LANES_PER_V_ROW) * (DMA_BYTES // 2) + xor_mask = (row_in_tile & arith.index(V_SWZ_ROW_MASK)) << arith.index(4) + unsw_col_f16 = swiz_col_f16 ^ xor_mask + col_byte = unsw_col_f16 * 2 + global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte) + voffset = arith.index_cast(T.i32, global_byte) + rocdl.raw_ptr_buffer_load_lds( + k_rsrc, + lds_ptr, + _dma_size, + voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + # ---- Preload Q^T B-operand and DO^T B-operand packs ---- + q_row = q_start + wave_q_offset + lane_mod_32 + q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, q_row, seq_len_q_v) + q_row_safe = arith.select(q_in_bounds, q_row, arith.index(0)) + c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type) + q_b_packs = [] + do_b_packs = [] + for ks in range_constexpr(K_STEPS_QK): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + g_idx = global_idx_q(q_row_safe, col) + q_raw = load_global_mfma_pack(q_ptr, g_idx) + q_b_packs.append(arith.select(q_in_bounds, q_raw, c_zero_mfma_pack)) + do_raw = load_global_mfma_pack(do_ptr, g_idx) + do_b_packs.append(arith.select(q_in_bounds, do_raw, c_zero_mfma_pack)) + + # ---- Constants ---- + c_zero_f = arith.constant(0.0, type=compute_type) + c_neg_inf = arith.constant(-1.0e30, type=compute_type) # finite NEG_INF (matches Triton) + c_zero_v16f32 = arith.constant_vector(0.0, v16f32_type) + # V4 LSE is RAW-domain (qk*sm_scale + ln(l)). So use sm_scale, not sm_scale*log2e. + c_sm_scale = arith.constant(sm_scale, type=compute_type) + + # ---- Per-q-row scalars (LSE, delta) ---- + # bh_base_tokens_q is the Q LSE/DELTAS base too (LSE shape [B,HQ,Sq]). + lse_delta_off_i32 = arith.index_cast(T.i32, bh_base_tokens_q + q_row_safe) + lse_val = buffer_ops.buffer_load( + lse_rsrc, + lse_delta_off_i32, + vec_width=1, + dtype=T.f32, + ) + delta_val = buffer_ops.buffer_load( + deltas_rsrc, + lse_delta_off_i32, + vec_width=1, + dtype=T.f32, + ) + + # ---- (No SINK contribution: pool kernel does not touch sink.) ---- + + # ---- POOL: single n-block; double-buffered LDS optional ---- + _use_dbuf = ENABLE_DMA + + init_args = [] + for _ in range_constexpr(D_CHUNKS): + init_args.append(c_zero_v16f32) + if _use_dbuf: + init_args.append(arith.index(0)) # cur_buf_id + + # PROLOGUE: prefetch iter 0's K into BOTH slots of buf 0. + _init_kv_start = n_block_start * BN + coop_dma_k(_init_kv_start, buf_id=0) + coop_dma_k_as_v(_init_kv_start, buf_id=0) + + for block_idx, inner_iter_args, loop_results in scf.for_( + n_block_start, + n_block_end, + arith.index(1), + iter_args=init_args, + ): + dq_accs = [inner_iter_args[i] for i in range_constexpr(D_CHUNKS)] + if _use_dbuf: + cur_buf = inner_iter_args[D_CHUNKS] + next_buf = arith.index(1) - cur_buf + + # V4 SWA: block_idx is directly the K-block index. + kv_block_start = block_idx * BN + kv_start = kv_block_start + + if _use_dbuf: + rocdl.s_waitcnt(0) + gpu.barrier() + k_base = k_buf_base(cur_buf) + else: + coop_load_k(kv_start, buf_id=0) + coop_load_k_as_v(kv_start, buf_id=0) + gpu.barrier() + k_base = k_buf_base(0) + + # ==== GEMM1: s = Q @ K^T ==== + k_hi_offset = K_SUB_N * K_STRIDE + k_swz_mask = (lane_mod_32 & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + + def _k_idx_lo(ks): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + return k_base + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask) + + def _k_idx_hi(ks): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + return k_base + k_hi_offset + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask) + + s_acc_lo = c_zero_v16f32 + s_acc_hi = c_zero_v16f32 + for ks in range_constexpr(K_STEPS_QK): + k_pack_lo = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_lo(ks)]) + k_pack_hi = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_hi(ks)]) + s_acc_lo = mfma_acc(k_pack_lo, q_b_packs[ks], s_acc_lo) + s_acc_hi = mfma_acc(k_pack_hi, q_b_packs[ks], s_acc_hi) + + # ==== Compute p[r] = exp((qk + add_bias) * sm_scale_inline - LSE) ==== + # + # POOL-only mask. The s_acc registers map (per lane, per reg r) to + # the S = Q @ K^T cell at row = q_row (= lane_mod_32 + wave_q_offset + # + q_start) and column = kv_start + lane_div_32*4 + (r//4)*8 + r%4 + # (lo half), or +K_SUB_N (hi half). Translate to pool_n by + # subtracting HCA_LOCAL_SEQLEN; gate with pool_n < POOL_SIZE and + # q_row < seq_len_q. + seq_len_q_i32 = arith.index_cast(T.i32, seq_len_q_v) + q_row_i32_mask = arith.index_cast(T.i32, q_row) + lane_div_32_i32 = arith.index_cast(T.i32, lane_div_32) + lane_off_i32 = arith.MulIOp(lane_div_32_i32, arith.constant(4, type=T.i32)).result + hca_local_i32 = arith.constant(hca_local_seqlen, type=T.i32) + pool_size_i32 = arith.constant(pool_size, type=T.i32) + kv_start_i32 = arith.index_cast(T.i32, kv_start) + # pool_n base for this lane = kv_start + lane_off - hca_local_seqlen + # (kv_start = hca_local_seqlen + 0 for the single pool block). + pool_n_base_i32 = arith.SubIOp( + arith.AddIOp(kv_start_i32, lane_off_i32).result, + hca_local_i32, + ).result + + # Is this q_row out of bounds? -> entire row masked. + q_oob = arith.cmpi( + arith.CmpIPredicate.sge, + q_row_i32_mask, + seq_len_q_i32, + ) + + # --- Load add_bias from ADD_MASK[q_row, pool_n] via GEP --- + # ADD_MASK shape [Sq, POOL_SIZE] contiguous; element bytes = 2 + # (bf16 / f16). For each (lane, r), compute pool_n (lo or hi) and + # load a single bf16/f16, then cast to f32. + ADD_MASK_STRIDE_M = arith.index(pool_size) # row stride in elements + q_row_idx_for_mask = arith.select(q_oob, arith.index(0), q_row) + # Precompute add-mask row base offset in elements. + add_mask_row_base = q_row_idx_for_mask * ADD_MASK_STRIDE_M + + def _load_add_bias(pool_n_i32, bad_pred): + """Load ADD_MASK[q_row, pool_n] as f32; return 0.0 if bad.""" + # Convert pool_n to index for GEP. + pool_n_i32_safe = arith.select( + bad_pred, + arith.constant(0, type=T.i32), + pool_n_i32, + ) + pool_n_idx = arith.index_cast(T.index, pool_n_i32_safe) + elem_idx = add_mask_row_base + pool_n_idx + bias_raw_v1 = _gep_load( + add_mask_ptr, + elem_idx, + T.vec(1, elem_type), + ) + bias_raw = vector.extract( + bias_raw_v1, + static_position=[0], + dynamic_position=[], + ) + bias_f32 = arith.extf(compute_type, bias_raw) + # If bad (OOB pool_n or q_oob), zero out the bias (mask wins + # anyway via NEG_INF, but keep numerics tidy). + return arith.select(bad_pred, c_zero_f, bias_f32) + + p_vals_lo = [] + p_vals_hi = [] + for r in range_constexpr(16): + r_off_i32 = arith.constant((r % 4) + (r // 4) * 8, type=T.i32) + + # --- lo half --- + pool_n_lo_i32 = arith.AddIOp( + pool_n_base_i32, + r_off_i32, + ).result + is_pool_oob_lo = arith.cmpi( + arith.CmpIPredicate.sge, + pool_n_lo_i32, + pool_size_i32, + ) + bad_lo = arith.OrIOp(is_pool_oob_lo, q_oob).result + + s_lo_f32 = vector.extract(s_acc_lo, static_position=[r], dynamic_position=[]) + bias_lo = _load_add_bias(pool_n_lo_i32, bad_lo) + qk_plus_bias_lo = arith.AddFOp(s_lo_f32, bias_lo, fastmath=fm_fast).result + # Triton: qk = qk * sm_scale; qk = qk + add_bias. + # We deferred sm_scale to here: (qk + bias) is wrong; need + # qk*scale + bias. Re-do as: s_scaled + bias. + scaled_lo = arith.MulFOp(s_lo_f32, c_sm_scale, fastmath=fm_fast).result + scaled_plus_bias_lo = arith.AddFOp(scaled_lo, bias_lo, fastmath=fm_fast).result + scaled_lo_masked = arith.select( + bad_lo, + c_neg_inf, + scaled_plus_bias_lo, + ) + diff_lo = arith.SubFOp(scaled_lo_masked, lse_val, fastmath=fm_fast).result + p_lo = math_dialect.exp(diff_lo, fastmath=fm_fast) + p_vals_lo.append(p_lo) + + # --- hi half: pool_n_hi = pool_n_lo + K_SUB_N --- + pool_n_hi_i32 = arith.AddIOp( + pool_n_lo_i32, + arith.constant(K_SUB_N, type=T.i32), + ).result + is_pool_oob_hi = arith.cmpi( + arith.CmpIPredicate.sge, + pool_n_hi_i32, + pool_size_i32, + ) + bad_hi = arith.OrIOp(is_pool_oob_hi, q_oob).result + + s_hi_f32 = vector.extract(s_acc_hi, static_position=[r], dynamic_position=[]) + bias_hi = _load_add_bias(pool_n_hi_i32, bad_hi) + scaled_hi = arith.MulFOp(s_hi_f32, c_sm_scale, fastmath=fm_fast).result + scaled_plus_bias_hi = arith.AddFOp(scaled_hi, bias_hi, fastmath=fm_fast).result + scaled_hi_masked = arith.select( + bad_hi, + c_neg_inf, + scaled_plus_bias_hi, + ) + diff_hi = arith.SubFOp(scaled_hi_masked, lse_val, fastmath=fm_fast).result + p_hi = math_dialect.exp(diff_hi, fastmath=fm_fast) + p_vals_hi.append(p_hi) + + # ==== Overwrite K LDS slot with V for GEMM2 ==== + gpu.barrier() + if _use_dbuf: + coop_dma_v_as_k(kv_start, buf_id=cur_buf) + rocdl.s_waitcnt(0) + gpu.barrier() + elif ENABLE_DMA: + coop_dma_v_as_k(kv_start, buf_id=0) + rocdl.s_waitcnt(0) + gpu.barrier() + else: + coop_load_v_as_k(kv_start, buf_id=0) + gpu.barrier() + + # ==== GEMM2: dP = DO @ V^T ==== + dp_acc_lo = c_zero_v16f32 + dp_acc_hi = c_zero_v16f32 + for ks in range_constexpr(K_STEPS_QK): + v_pack_lo = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_lo(ks)]) + v_pack_hi = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_hi(ks)]) + dp_acc_lo = mfma_acc(v_pack_lo, do_b_packs[ks], dp_acc_lo) + dp_acc_hi = mfma_acc(v_pack_hi, do_b_packs[ks], dp_acc_hi) + + # ==== Prefetch iter+1's K DMAs (async) ==== + if _use_dbuf: + _next_block_idx = block_idx + arith.index(1) + _has_next = arith.cmpi(arith.CmpIPredicate.slt, _next_block_idx, n_block_end) + _pre_if = scf.IfOp(_has_next) + with ir.InsertionPoint(_pre_if.then_block): + _next_kv_start = _next_block_idx * BN + coop_dma_k(_next_kv_start, next_buf) + coop_dma_k_as_v(_next_kv_start, next_buf) + scf.YieldOp([]) + + # ==== Compute dS[r] = p[r] * (dp[r] - delta) ==== + ds_vals_lo = [] + ds_vals_hi = [] + for r in range_constexpr(16): + dp_lo = vector.extract(dp_acc_lo, static_position=[r], dynamic_position=[]) + dp_hi = vector.extract(dp_acc_hi, static_position=[r], dynamic_position=[]) + diff_lo = arith.SubFOp(dp_lo, delta_val, fastmath=fm_fast).result + diff_hi = arith.SubFOp(dp_hi, delta_val, fastmath=fm_fast).result + ds_lo = arith.MulFOp(p_vals_lo[r], diff_lo, fastmath=fm_fast).result + ds_hi = arith.MulFOp(p_vals_hi[r], diff_hi, fastmath=fm_fast).result + ds_vals_lo.append(ds_lo) + ds_vals_hi.append(ds_hi) + + # ==== Pack dS f32 -> mfma_pack_type (bf16/f16) ==== + if dtype_str == "bf16" and USE_K16: + ds_packs_lo = [] + ds_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + base = pks * 8 + ds_packs_lo.append(bf16_trunc_pack_v8(ds_vals_lo[base : base + 8])) + ds_packs_hi.append(bf16_trunc_pack_v8(ds_vals_hi[base : base + 8])) + elif dtype_str == "bf16": + ds_packs_lo = [] + ds_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + base = pks * 4 + ds_packs_lo.append(bf16_trunc_pack_v4(ds_vals_lo[base : base + 4])) + ds_packs_hi.append(bf16_trunc_pack_v4(ds_vals_hi[base : base + 4])) + else: + ds_f16_lo = [arith.trunc_f(elem_type, ds_vals_lo[r]) for r in range_constexpr(16)] + ds_f16_hi = [arith.trunc_f(elem_type, ds_vals_hi[r]) for r in range_constexpr(16)] + _pack_ty = v8f16_type if USE_K16 else v4f16_type + _pw = 8 if USE_K16 else 4 + ds_packs_lo = [] + ds_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + b = pks * _pw + ds_packs_lo.append(vector.from_elements(_pack_ty, [ds_f16_lo[b + i] for i in range(_pw)])) + ds_packs_hi.append(vector.from_elements(_pack_ty, [ds_f16_hi[b + i] for i in range(_pw)])) + + # ==== GEMM3: dQ += K @ dS (uses fwd's V-read schedule with K substituted) ==== + _steps = [(dc, pks) for dc in range(D_CHUNKS) for pks in range(PV_K_STEPS)] + TOTAL_PV = len(_steps) + v_base = v_buf_base(cur_buf) if _use_dbuf else v_buf_base(0) + + def _read_k_as_v_pack(step_idx): + dc, pks = _steps[step_idx] + if USE_HW_TR: + d_col = arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4 + k_row = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + tr_k_group + _d_col_eff = _v_swizzle(k_row, d_col) if ENABLE_DMA else d_col + lds_lo = v_base + k_row * V_STRIDE + _d_col_eff + lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE) + if USE_K16: + vl_a = ds_read_tr_v4f16(lds_lo) + vl_b = ds_read_tr_v4f16(lds_lo + arith.index(8 * V_STRIDE)) + vl = vector.shuffle(vl_a, vl_b, [0, 1, 2, 3, 4, 5, 6, 7]) + vh_a = ds_read_tr_v4f16(lds_hi) + vh_b = ds_read_tr_v4f16(lds_hi + arith.index(8 * V_STRIDE)) + vh = vector.shuffle(vh_a, vh_b, [0, 1, 2, 3, 4, 5, 6, 7]) + else: + vl = ds_read_tr_v4f16(lds_lo) + vh = ds_read_tr_v4f16(lds_hi) + else: + d_pos = arith.index(dc * D_CHUNK) + lane_mod_32 + kb = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + v_lo_idx = v_base + d_pos * VT_STRIDE + kb + v_hi_idx = v_lo_idx + arith.index(K_SUB_N) + vl = vector.load(v4f16_type, lds_kv, [v_lo_idx]) + vh = vector.load(v4f16_type, lds_kv, [v_hi_idx]) + return vl, vh + + k_lo_cur, k_hi_cur = _read_k_as_v_pack(0) + for si in range_constexpr(TOTAL_PV): + dc, pks = _steps[si] + if si + 1 < TOTAL_PV: + k_lo_nxt, k_hi_nxt = _read_k_as_v_pack(si + 1) + dq_accs[dc] = mfma_acc(k_lo_cur, ds_packs_lo[pks], dq_accs[dc]) + dq_accs[dc] = mfma_acc(k_hi_cur, ds_packs_hi[pks], dq_accs[dc]) + if si + 1 < TOTAL_PV: + k_lo_cur = k_lo_nxt + k_hi_cur = k_hi_nxt + + # End of iter: yield dq_accs and swapped cur_buf. + gpu.barrier() + _yield = list(dq_accs) + if _use_dbuf: + _yield.append(next_buf) + yield _yield + + # ---- Final store: dQ_fp32 += dq_acc * sm_scale (load + add + store) ---- + # Pool kernel ACCUMULATES into the existing LOCAL-stream dq buffer. + # Race-free because each program owns a unique (b, qhid, m_block) + # slice, and the launcher runs this kernel sequentially AFTER the + # SWA dq kernel has finished writing the LOCAL contribution. + dq_finals = [loop_results[dc] for dc in range_constexpr(D_CHUNKS)] + sm_scale_vec = vector.broadcast(v16f32_type, c_sm_scale) + + def _gep_load_f32(base_ptr_, elem_idx): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr_, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=T.f32, + noWrapFlags=0, + ) + return _llvm.LoadOp(T.f32, gep.result).result + + _o_guard = scf.IfOp(q_in_bounds, [], has_else=False) + with ir.InsertionPoint(_o_guard.then_block): + for dc in range_constexpr(D_CHUNKS): + dq_scaled = arith.MulFOp( + dq_finals[dc], + sm_scale_vec, + fastmath=fm_fast, + ).result + for r in range_constexpr(16): + dq_val = vector.extract( + dq_scaled, + static_position=[r], + dynamic_position=[], + ) + d_row_rel = lane_div_32 * 4 + (r // 4) * 8 + (r % 4) + d_col = arith.index(dc * D_CHUNK) + d_row_rel + dq_global = global_idx_q(q_row, d_col) + dq_prev = _gep_load_f32(dq_ptr, dq_global) + dq_sum = arith.AddFOp( + dq_val, + dq_prev, + fastmath=fm_fast, + ).result + _gep_store_f32(dq_sum, dq_ptr, dq_global) + scf.YieldOp([]) + + @flyc.jit + def launch_v4_hca_bwd_dq_pool( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + DOS: fx.Tensor, + LSE: fx.Tensor, + DELTAS: fx.Tensor, + DQ: fx.Tensor, + ADD_MASK: fx.Tensor, + batch_size: fx.Int32, + seq_len_q: fx.Int32, + seq_len_k: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + bs_idx = arith.index_cast(T.index, batch_size) + sl_idx = arith.index_cast(T.index, seq_len_q) + num_q_tiles = (sl_idx + BLOCK_M - 1) // BLOCK_M + grid_x = bs_idx * num_q_tiles * NUM_HEADS + + launcher = v4_hca_bwd_dq_pool_kernel( + Q, + K, + V, + DOS, + LSE, + DELTAS, + DQ, + ADD_MASK, + seq_len_q, + seq_len_k, + ) + + if waves_per_eu is not None: + _wpe = int(waves_per_eu) + if _wpe >= 1: + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, + _wpe, + ) + if flat_work_group_size is not None: + _fwgs = int(flat_work_group_size) + if _fwgs >= 1: + flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}") + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr + + passthrough_entries = [] + if daz: + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("denormal-fp-math-f32"), + ir.StringAttr.get("preserve-sign,preserve-sign"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("no-nans-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("unsafe-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries) + + launcher.launch( + grid=(grid_x, 1, 1), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) + + _fmha_compile_hints = { + "fast_fp_math": fast_fp_math, + "unsafe_fp_math": unsafe_fp_math, + "llvm_options": { + "enable-post-misched": False, + "lsr-drop-solution": True, + }, + } + + def _launch(*args, **kwargs): + with CompilationContext.compile_hints(_fmha_compile_hints): + return launch_v4_hca_bwd_dq_pool(*args, **kwargs) + + def _compile(Q, K, V, DOS, LSE, DELTAS, DQ, ADD_MASK, batch_size, seq_len_q, seq_len_k, stream=None): + with CompilationContext.compile_hints(_fmha_compile_hints): + return flyc.compile( + launch_v4_hca_bwd_dq_pool, + Q, + K, + V, + DOS, + LSE, + DELTAS, + DQ, + ADD_MASK, + batch_size, + seq_len_q, + seq_len_k, + fx.Stream(stream), + ) + + _launch.compile = _compile + + return _launch + + +# Convenience alias. +build_v4_hca_bwd_dq_pool_module_primary = build_v4_hca_bwd_dq_pool_module diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dkv_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dkv_kernel.py new file mode 100644 index 000000000..eda441de5 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dkv_kernel.py @@ -0,0 +1,1113 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""v4_swa_bwd_dkv: V4 SWA-causal attention backward dK/dV kernel for FlyDSL. + +Strategy: + * One workgroup per (batch, n_block). K and V loaded once into LDS + and stay resident for the program's lifetime (no LUT, no atomics). + * Inner loop iterates qhid in [0, HQ); for each head, iterates m-blocks + bounded by the SWA window. dK and dV f32 accumulators live in VGPRs. + * MQA head-loop accumulator: one dk/dv slice per (b, n_block) gets + all HQ heads accumulated -- deterministic, no atomics. + * SWA + causal mask per element. LSE is RAW-domain + (qk*sm_scale + ln(l)); p = exp(qk*sm_scale - lse). + * sm_scale applied to dK exactly once post-loop (Triton P57 cr=0). + * dK / dV written as f32 (matches launcher's dk_fp32 / dv_fp32 buffers). + +LDS budget at BLOCK_N=32, BLOCK_M2=32, D=512: + K = 32 KB, V = 32 KB, DO = 32 KB, Q = 32 KB, pT = 2 KB + Total = 130 KB <= 160 KB. +Auto-fallback: if predicted > 160 KB, drop DO/Q LDS scratches and +read Q/DO directly to register packs from HBM (mirrors STEP-1b dq). +""" + +import math +import os + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import math as math_dialect +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr +from kernels.kernels_common import dtype_to_elem_type + +KERNEL_NAME = "v4_swa_bwd_dkv_kernel" + +_LLVM_GEP_DYNAMIC = -2147483648 + + +def _llvm_ptr_ty(): + return ir.Type.parse("!llvm.ptr") + + +def _llvm_lds_ptr_ty(): + return ir.Type.parse("!llvm.ptr<3>") + + +def build_v4_swa_bwd_dkv_module( + num_heads, + head_dim, + swa_window, + dtype_str="bf16", + sm_scale=None, + waves_per_eu=2, + flat_work_group_size=None, + block_n=None, + block_m2=None, + unsafe_fp_math=True, + fast_fp_math=True, + daz=True, + layout_bhld=True, + mqa_kv=True, +): + """Build the V4 SWA backward dK/dV launcher (MQA only).""" + gpu_arch = get_hip_arch() + + if block_n is None: + # R6k-1: D-split wave remap default. BLOCK_N=16 halves per-lane + # fp32 accumulator footprint (256 -> 128 fp32) by partitioning D + # across waves instead of N. Both waves cover the same 16 KV + # rows; each wave owns half of the D-chunks of dk/dv. + BLOCK_N = 16 + else: + BLOCK_N = int(block_n) + if block_m2 is None: + BLOCK_M2 = 32 + else: + BLOCK_M2 = int(block_m2) + WARP_SIZE = 64 + # R6k-1: D-split mode. When BLOCK_N == 16 (matches MFMA-N tile), + # force NUM_WAVES=2 and partition D between the two waves instead + # of partitioning N. Each wave then carries only half the dk/dv + # accumulator footprint. For BLOCK_N > 16 we keep the legacy N-split. + D_SPLIT_WAVES = BLOCK_N == 16 + if D_SPLIT_WAVES: + NUM_WAVES = 2 + ROWS_PER_WAVE = 16 # both waves work on the same 16 rows + else: + # Legacy N-split: each wave gets BLOCK_N/NUM_WAVES rows. + NUM_WAVES = max(1, BLOCK_N // 16) + ROWS_PER_WAVE = BLOCK_N // NUM_WAVES + if flat_work_group_size is None: + flat_work_group_size = NUM_WAVES * WARP_SIZE + BLOCK_SIZE = flat_work_group_size + + ENABLE_LDS_VEC16 = os.getenv("FLYDSL_SLA_FWD_ENABLE_LDS_VEC16", "1") == "1" + USE_K16 = gpu_arch.startswith("gfx950") + # R6b: MFMA 16x16x32 K-step = 32; A/B-frag = 8 bf16 per lane. + assert USE_K16, "R6b dkv requires gfx950 (MFMA 16x16x32 bf16)." + K_STEP_QK = 32 + K_STEPS_QK = head_dim // K_STEP_QK + # R6b: each MFMA 16x16x32 tile produces a 16-wide N-chunk; D_CHUNK = 16 cols. + D_CHUNK = 16 + D_CHUNKS = head_dim // D_CHUNK + # R6k-1: per-wave D-chunk count. In D-split mode each wave owns half. + if D_SPLIT_WAVES: + assert D_CHUNKS % NUM_WAVES == 0, ( + f"D_CHUNKS ({D_CHUNKS}) must be divisible by NUM_WAVES " f"({NUM_WAVES}) for D-split mode." + ) + D_CHUNKS_LOCAL = D_CHUNKS // NUM_WAVES + else: + D_CHUNKS_LOCAL = D_CHUNKS + K_STEPS_PT = BLOCK_M2 // K_STEP_QK + # R6b: GEMM1/GEMM3 cover m_col [0..BLOCK_M2) with multiple 16-col MFMA-N tiles per ks. + assert BLOCK_M2 % 16 == 0, f"BLOCK_M2 must be a multiple of 16 (MFMA-N), got {BLOCK_M2}" + M_TILES = BLOCK_M2 // 16 + + assert BLOCK_N % NUM_WAVES == 0 + assert ROWS_PER_WAVE == 16, f"ROWS_PER_WAVE must equal 16 for MFMA 16x16x32, got {ROWS_PER_WAVE}" + assert head_dim % 32 == 0 + assert head_dim >= 64 + assert flat_work_group_size in (64, 128, 256, 512) + assert dtype_str == "bf16", "R6b dkv currently only supports bf16." + assert BLOCK_N % 16 == 0 + assert BLOCK_M2 % K_STEP_QK == 0, ( + f"BLOCK_M2 ({BLOCK_M2}) must be a multiple of MFMA K-step " f"({K_STEP_QK}) for the dV/dK GEMMs." + ) + assert isinstance(swa_window, int) and swa_window > 0 + assert mqa_kv, "v4_swa_bwd_dkv currently only supports MQA (HK=1)." + + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(head_dim) + + NUM_HEADS = num_heads + HEAD_DIM = head_dim + + K_STRIDE = HEAD_DIM + K_SWZ_ROW_MASK = (K_STRIDE // 16) - 1 + assert K_SWZ_ROW_MASK >= 0 + assert (K_SWZ_ROW_MASK & (K_SWZ_ROW_MASK + 1)) == 0 + V_STRIDE = HEAD_DIM + + VEC_WIDTH = 16 if ENABLE_LDS_VEC16 else 8 + assert HEAD_DIM % VEC_WIDTH == 0 + THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH + assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0 + ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD + + LDS_K_TILE_SIZE = BLOCK_N * K_STRIDE + LDS_V_TILE_SIZE = BLOCK_N * V_STRIDE + LDS_DO_STRIDE = HEAD_DIM + LDS_DO_ELEMS = BLOCK_M2 * LDS_DO_STRIDE + LDS_Q_STRIDE = HEAD_DIM + LDS_Q_ELEMS = BLOCK_M2 * LDS_Q_STRIDE + LDS_PT_STRIDE = BLOCK_M2 + LDS_PT_ELEMS = BLOCK_N * LDS_PT_STRIDE + + _LDS_LIMIT_BYTES = 160 * 1024 + + def _predict_lds_bytes(use_lds_for_q_do): + b = (LDS_K_TILE_SIZE + LDS_V_TILE_SIZE + LDS_PT_ELEMS) * 2 + if use_lds_for_q_do: + b += (LDS_DO_ELEMS + LDS_Q_ELEMS) * 2 + return b + + USE_LDS_FOR_Q_DO = True + _pred_full = _predict_lds_bytes(True) + if _pred_full > _LDS_LIMIT_BYTES: + USE_LDS_FOR_Q_DO = False + _pred_min = _predict_lds_bytes(False) + if _pred_min > _LDS_LIMIT_BYTES: + raise RuntimeError( + f"v4_swa_bwd_dkv: minimal LDS {_pred_min}B > limit " + f"{_LDS_LIMIT_BYTES}B at D={head_dim}, BLOCK_N={BLOCK_N}, " + f"BLOCK_M2={BLOCK_M2}." + ) + import sys as _sys + + print( + f"[v4_swa_bwd_dkv] auto-fallback: LDS {_pred_full}B > " + f"{_LDS_LIMIT_BYTES}B; disabling Q/DO LDS " + f"({_pred_full} -> {_pred_min})", + file=_sys.stderr, + flush=True, + ) + else: + import sys as _sys + + print( + f"[v4_swa_bwd_dkv] LDS OK: predicted {_pred_full}B <= " + f"{_LDS_LIMIT_BYTES}B at D={head_dim}, BLOCK_N={BLOCK_N}, " + f"BLOCK_M2={BLOCK_M2}, USE_LDS_FOR_Q_DO=True", + file=_sys.stderr, + flush=True, + ) + + allocator = SmemAllocator( + None, + arch=gpu_arch, + global_sym_name=( + f"v4_swa_bwd_dkv_smem_N{BLOCK_N}_M2_{BLOCK_M2}_W{swa_window}" + f"_MQ{int(mqa_kv)}_QDO{int(USE_LDS_FOR_Q_DO)}" + f"_DS{int(D_SPLIT_WAVES)}" + ), + ) + lds_kv_offset = allocator._align(allocator.ptr, 16) + LDS_V_BASE = LDS_K_TILE_SIZE + LDS_KV_TOTAL_SIZE = LDS_K_TILE_SIZE + LDS_V_TILE_SIZE + allocator.ptr = lds_kv_offset + LDS_KV_TOTAL_SIZE * 2 + + lds_pt_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_pt_offset + LDS_PT_ELEMS * 2 + + if USE_LDS_FOR_Q_DO: + lds_do_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_do_offset + LDS_DO_ELEMS * 2 + lds_q_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_q_offset + LDS_Q_ELEMS * 2 + else: + lds_do_offset = None + lds_q_offset = None + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def v4_swa_bwd_dkv_kernel( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + DOS: fx.Tensor, + LSE: fx.Tensor, + DELTAS: fx.Tensor, + DK: fx.Tensor, + DV: fx.Tensor, + seq_len_q: fx.Int32, + seq_len_k: fx.Int32, + ): + elem_type = dtype_to_elem_type(dtype_str) + compute_type = T.f32 + q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q) + k_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K) + v_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V) + do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOS) + dk_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DK) + dv_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DV) + lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True) + deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True) + + fm_fast = arith.FastMathFlags.fast + vxf16_type = T.vec(VEC_WIDTH, elem_type) + v8f16_type = T.vec(8, elem_type) + # R6M: tr16 returns 4 bf16 per lane (v4f16-typed). Two calls + shuffle + # form the MFMA 16x16x32 B-frag (8 bf16) via HW transpose. + v4f16_type = T.vec(4, elem_type) + # R6b: MFMA 16x16x32 C-frag = 4 fp32 per lane. + v4f32_type = T.vec(4, compute_type) + mfma_pack_type = v8f16_type + MFMA_LANE_K = 8 + v1_elem_type = T.vec(1, elem_type) + _mfma_zero = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0) + + def mfma_acc(a, b, c): + # rocdl.mfma_f32_16x16x32_bf16 is the wrapped form: takes + # (result_type, [operands]) and returns the Value directly. + return rocdl.mfma_f32_16x16x32_bf16( + v4f32_type, + [a, b, c, _mfma_zero, _mfma_zero, _mfma_zero], + ) + + seq_len_q_v = arith.index_cast(T.index, seq_len_q) + seq_len_k_v = arith.index_cast(T.index, seq_len_k) + + base_ptr = allocator.get_base() + lds_kv = SmemPtr(base_ptr, lds_kv_offset, elem_type, shape=(LDS_KV_TOTAL_SIZE,)).get() + lds_pt = SmemPtr(base_ptr, lds_pt_offset, elem_type, shape=(LDS_PT_ELEMS,)).get() + if USE_LDS_FOR_Q_DO: + lds_do = SmemPtr(base_ptr, lds_do_offset, elem_type, shape=(LDS_DO_ELEMS,)).get() + lds_q = SmemPtr(base_ptr, lds_q_offset, elem_type, shape=(LDS_Q_ELEMS,)).get() + + block_id = arith.index_cast(T.index, gpu.block_idx.x) + tid = arith.index_cast(T.index, gpu.thread_idx.x) + wave_id = tid // WARP_SIZE + lane = tid % WARP_SIZE + # R6b MFMA 16x16x32 lane decomposition: + # A-frag: A[lane_mod_16, ks*32 + lane_div_16*8 + 0..7] + # B-frag: B[ks*32 + lane_div_16*8 + 0..7, lane_mod_16] + # C-frag: C[lane_div_16*4 + ii, lane_mod_16] for ii in 0..3 + lane_mod_16 = lane % 16 + lane_div_16 = lane // 16 + # R6M: ds_read_tr16_b64 per-lane decomposition (matches proto). + tr_k_group = lane_mod_16 // arith.index(4) + tr_col_sub = lane_mod_16 % arith.index(4) + + # R6k-1: in D-split mode both waves cover the same 16 KV rows + # (wave_n_offset == 0) and instead carry disjoint D-chunks of + # dk/dv. wave_d_col_offset is the wave's D-column base offset + # (0 for wave 0, D_CHUNKS_LOCAL*D_CHUNK for wave 1). + if D_SPLIT_WAVES: + wave_n_offset = arith.index(0) + wave_d_col_offset = wave_id * arith.index(D_CHUNKS_LOCAL * D_CHUNK) + else: + wave_n_offset = wave_id * ROWS_PER_WAVE + wave_d_col_offset = arith.index(0) + + # R6S: Hoist a SINGLE per-lane base byte-pointer for DO/Q LDS + # tr16 reads. Each ds_read_tr16_b64 then uses + # `base + constexpr byte imm` so LLVM ISel folds the + # (m_step, dc, addr_0/1) deltas into the `offset:` immediate + # of the LDS instruction. Removes the per-call IntToPtrOp -> + # 32 AGPR spill chain that was throttling GEMM2/4 to 28.6% + # MFMA duty (vs Triton's 53.7%). + if USE_LDS_FOR_Q_DO: + assert LDS_DO_STRIDE == LDS_Q_STRIDE, "R6S tr16 base+imm assumes DO and Q share STRIDE" + tr_lane_base_elem = ( + lane_div_16 * arith.index(MFMA_LANE_K * LDS_DO_STRIDE) + + tr_k_group * arith.index(LDS_DO_STRIDE) + + wave_d_col_offset + + tr_col_sub * arith.index(4) + ) + tr_lane_base_byte = tr_lane_base_elem * arith.index(2) + tr_lane_base_i64 = arith.index_cast(T.i64, tr_lane_base_byte) + tr_lane_base_do_byte_i64 = arith.AddIOp( + tr_lane_base_i64, + arith.constant(lds_do_offset, type=T.i64), + ).result + tr_lane_base_q_byte_i64 = arith.AddIOp( + tr_lane_base_i64, + arith.constant(lds_q_offset, type=T.i64), + ).result + tr_lane_base_do_ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), tr_lane_base_do_byte_i64).result + tr_lane_base_q_ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), tr_lane_base_q_byte_i64).result + + # 6N-2 iter2: de-dup helper (Q/DO B-frag loads also gated). Each m-tile mt is owned by wave (mt % NUM_WAVES). + # In D-split mode the owning wave runs GEMM1/softmax/pT-write/GEMM3/dsT + # for that mt; the other wave skips. GEMM2/4 use the SHARED pT/dsT LDS. + if D_SPLIT_WAVES: + owns_mt_preds = [ + arith.cmpi(arith.CmpIPredicate.eq, wave_id, arith.index(mt % NUM_WAVES)) + for mt in range(M_TILES) + ] + else: + owns_mt_preds = None + + num_n_blocks = (seq_len_k_v + BLOCK_N - 1) // BLOCK_N + n_block_idx = block_id % num_n_blocks + batch_idx = block_id // num_n_blocks + kv_start = n_block_idx * arith.index(BLOCK_N) + + load_row_in_batch = tid // THREADS_PER_ROW_LOAD + load_lane_in_row = tid % THREADS_PER_ROW_LOAD + load_col_base = load_lane_in_row * VEC_WIDTH + + # MQA: KV stride drops head. + bh_base_tokens_kv = batch_idx * seq_len_k_v + + def bh_base_tokens_q_of(qhid_index): + return (batch_idx * NUM_HEADS + qhid_index) * seq_len_q_v + + def global_idx_q(qhid_index, token_idx, col): + return (bh_base_tokens_q_of(qhid_index) + token_idx) * arith.index(HEAD_DIM) + col + + def global_idx_kv(token_idx, col): + return (bh_base_tokens_kv + token_idx) * arith.index(HEAD_DIM) + col + + def _gep_load(base_ptr_, elem_idx, vec_type, et=elem_type): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr_, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=et, + noWrapFlags=0, + ) + return _llvm.LoadOp(vec_type, gep.result).result + + def _gep_store_f32(val, base_ptr_, elem_idx): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr_, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=T.f32, + noWrapFlags=0, + ) + _llvm.StoreOp(val, gep.result) + + def load_global_mfma_pack(base_ptr_, base_idx): + return _gep_load(base_ptr_, base_idx, mfma_pack_type) + + def load_global_f16xN(base_ptr_, base_idx): + return _gep_load(base_ptr_, base_idx, vxf16_type) + + def _k_swizzle(row_idx, col_idx): + mask = (row_idx & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + return col_idx ^ mask + + if ROWS_PER_BATCH_LOAD >= BLOCK_N: + NUM_BATCHES_KV = 1 + KV_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_N + else: + assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0 + NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD + KV_NEEDS_GUARD = False + + if ROWS_PER_BATCH_LOAD >= BLOCK_M2: + NUM_BATCHES_M = 1 + M_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_M2 + else: + assert BLOCK_M2 % ROWS_PER_BATCH_LOAD == 0 + NUM_BATCHES_M = BLOCK_M2 // ROWS_PER_BATCH_LOAD + M_NEEDS_GUARD = False + + c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type) + c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type) + c_zero_elem = arith.constant(0.0, type=elem_type) + c_neg_inf = arith.constant(-1.0e30, type=compute_type) + c_sm_scale = arith.constant(sm_scale, type=compute_type) + # R6b: 4-elem fp32 acc per MFMA 16x16x32 op. + v4f32_zero = arith.constant_vector(0.0, v4f32_type) + + def coop_load_k(tile_start): + k_base = arith.index(0) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, seq_len_k_v) + row_safe = arith.select(in_bounds, row_idx, arith.index(0)) + g_idx = global_idx_kv(row_safe, load_col_base) + vec = load_global_f16xN(k_ptr, g_idx) + vec_safe = arith.select(in_bounds, vec, c_zero_vxf16) + lds_row = load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_N) + ) + _if_k = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_k.then_block): + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + vector.store(vec_safe, lds_kv, [lds_idx]) + scf.YieldOp([]) + else: + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + vector.store(vec_safe, lds_kv, [lds_idx]) + + def coop_load_v(tile_start): + v_base = arith.index(LDS_V_BASE) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, seq_len_k_v) + row_safe = arith.select(in_bounds, row_idx, arith.index(0)) + g_idx = global_idx_kv(row_safe, load_col_base) + vec = load_global_f16xN(v_ptr, g_idx) + vec_safe = arith.select(in_bounds, vec, c_zero_vxf16) + lds_row = load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_N) + ) + _if_v = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_v.then_block): + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = v_base + lds_row * V_STRIDE + swz_col + vector.store(vec_safe, lds_kv, [lds_idx]) + scf.YieldOp([]) + else: + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = v_base + lds_row * V_STRIDE + swz_col + vector.store(vec_safe, lds_kv, [lds_idx]) + + # ---- PROLOGUE: K, V into LDS (persistent for program lifetime) ---- + coop_load_k(kv_start) + coop_load_v(kv_start) + gpu.barrier() + + # R6b: K/V LDS A-frag load uses lane_mod_16 (16-row MFMA) and lane_div_16*8 for col. + # Swizzle mask key is the per-wave KV row (wave_n_offset + lane_mod_16), but since + # NUM_WAVES <= 2 and (wave_n_offset & K_SWZ_ROW_MASK) is 0 at D=64 (and bit-aligned + # at D=512 once wave_n_offset is folded in below), we recompute the mask per call. + def _k_idx_wave(ks): + kv_row = wave_n_offset + lane_mod_16 + col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + mask = (kv_row & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + return kv_row * arith.index(K_STRIDE) + (col ^ mask) + + def _v_idx_wave(ks): + kv_row = wave_n_offset + lane_mod_16 + col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + mask = (kv_row & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + return arith.index(LDS_V_BASE) + kv_row * arith.index(V_STRIDE) + (col ^ mask) + + # ---- SWA + m-loop bounds (constant across head loop) ---- + SWA = arith.index(swa_window) + BN_idx = arith.index(BLOCK_N) + BM2_idx = arith.index(BLOCK_M2) + _one_idx = arith.index(1) + n_block_lo = kv_start + n_block_hi = kv_start + BN_idx + m_loop_start = (n_block_lo // BM2_idx) * BM2_idx + _m_end_raw = n_block_hi + SWA - _one_idx + _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _m_end_raw, seq_len_q_v) + _m_end_cl = arith.select(_le_seq, _m_end_raw, seq_len_q_v) + m_loop_end = ((_m_end_cl + BM2_idx - _one_idx) // BM2_idx) * BM2_idx + + # R6k-1: each wave only carries D_CHUNKS_LOCAL chunks of dv and dk. + outer_carry = [v4f32_zero for _ in range(D_CHUNKS_LOCAL)] + [ + v4f32_zero for _ in range(D_CHUNKS_LOCAL) + ] + + seq_len_k_i32 = arith.index_cast(T.i32, seq_len_k_v) + seq_len_q_i32 = arith.index_cast(T.i32, seq_len_q_v) + w_i32 = arith.constant(swa_window, type=T.i32) + wave_n_off_i32 = arith.index_cast(T.i32, wave_n_offset) + kv_start_i32 = arith.index_cast(T.i32, kv_start) + lane_div_16_i32 = arith.index_cast(T.i32, lane_div_16) + + NUM_HEADS_idx = arith.index(NUM_HEADS) + for qhid_constexpr_idx, h_carry, h_loop_results in scf.for_( + arith.index(0), + NUM_HEADS_idx, + arith.index(1), + iter_args=outer_carry, + ): + qhid = qhid_constexpr_idx + + m_init_args = [h_carry[i] for i in range(2 * D_CHUNKS_LOCAL)] + for m_start, m_carry, m_loop_results in scf.for_( + m_loop_start, + m_loop_end, + BM2_idx, + iter_args=m_init_args, + ): + # R6k-1: each wave's dv/dk slice has D_CHUNKS_LOCAL chunks. + dv_accs = [m_carry[dc] for dc in range_constexpr(D_CHUNKS_LOCAL)] + dk_accs = [m_carry[D_CHUNKS_LOCAL + dc] for dc in range_constexpr(D_CHUNKS_LOCAL)] + + q_row_abs = m_start + lane_mod_16 + q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, q_row_abs, seq_len_q_v) + arith.select(q_in_bounds, q_row_abs, arith.index(0)) + + # ---- Q + DO loads ---- + if USE_LDS_FOR_Q_DO: + for batch in range_constexpr(NUM_BATCHES_M): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = m_start + load_row_in_batch + row_offset + in_bounds = arith.cmpi(arith.CmpIPredicate.slt, row_idx, seq_len_q_v) + row_safe = arith.select(in_bounds, row_idx, arith.index(0)) + g_idx_do = global_idx_q(qhid, row_safe, load_col_base) + g_idx_q = global_idx_q(qhid, row_safe, load_col_base) + vec_do = load_global_f16xN(do_ptr, g_idx_do) + vec_q = load_global_f16xN(q_ptr, g_idx_q) + vec_do_safe = arith.select(in_bounds, vec_do, c_zero_vxf16) + vec_q_safe = arith.select(in_bounds, vec_q, c_zero_vxf16) + lds_row = load_row_in_batch + row_offset + if M_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, load_row_in_batch + row_offset, arith.index(BLOCK_M2) + ) + _if_qd = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_qd.then_block): + lds_idx_do = lds_row * arith.index(LDS_DO_STRIDE) + load_col_base + lds_idx_q = lds_row * arith.index(LDS_Q_STRIDE) + load_col_base + vector.store(vec_do_safe, lds_do, [lds_idx_do]) + vector.store(vec_q_safe, lds_q, [lds_idx_q]) + scf.YieldOp([]) + else: + lds_idx_do = lds_row * arith.index(LDS_DO_STRIDE) + load_col_base + lds_idx_q = lds_row * arith.index(LDS_Q_STRIDE) + load_col_base + vector.store(vec_do_safe, lds_do, [lds_idx_do]) + vector.store(vec_q_safe, lds_q, [lds_idx_q]) + gpu.barrier() + + # 6N-2 iter2: when D-split, q_b_packs / do_b_packs_gemm3 + # are loaded INSIDE the per-mt scf.IfOp gates (see + # GEMM1 / GEMM3 below) so un-owned waves do not waste + # LDS read instructions. Build trivial sentinels. + if D_SPLIT_WAVES: + q_b_packs = None # gated path: loads done in GEMM1 if-block + do_b_packs_gemm3 = None # gated path: loads done in GEMM3 if-block + else: + q_b_packs = [[None] * K_STEPS_QK for _ in range(M_TILES)] + for mt in range_constexpr(M_TILES): + for ks in range_constexpr(K_STEPS_QK): + m_row = arith.index(mt * 16) + lane_mod_16 + d_col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + lds_idx = m_row * arith.index(LDS_Q_STRIDE) + d_col + pack = vector.load_op(mfma_pack_type, lds_q, [lds_idx]) + q_b_packs[mt][ks] = pack + + do_b_packs_gemm3 = [[None] * K_STEPS_QK for _ in range(M_TILES)] + for mt in range_constexpr(M_TILES): + for ks in range_constexpr(K_STEPS_QK): + m_row = arith.index(mt * 16) + lane_mod_16 + d_col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + lds_idx = m_row * arith.index(LDS_DO_STRIDE) + d_col + pack = vector.load_op(mfma_pack_type, lds_do, [lds_idx]) + do_b_packs_gemm3[mt][ks] = pack + else: + # R6b fallback: HBM-direct B-frag per m-tile. row = m_start + mt*16 + lane_mod_16. + q_b_packs = [[None] * K_STEPS_QK for _ in range(M_TILES)] + do_b_packs_gemm3 = [[None] * K_STEPS_QK for _ in range(M_TILES)] + for mt in range_constexpr(M_TILES): + row_abs = m_start + arith.index(mt * 16) + lane_mod_16 + in_bnd = arith.cmpi(arith.CmpIPredicate.slt, row_abs, seq_len_q_v) + row_safe = arith.select(in_bnd, row_abs, arith.index(0)) + for ks in range_constexpr(K_STEPS_QK): + col = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + g_idx = global_idx_q(qhid, row_safe, col) + q_raw = load_global_mfma_pack(q_ptr, g_idx) + q_b_packs[mt][ks] = arith.select(in_bnd, q_raw, c_zero_mfma_pack) + do_raw = load_global_mfma_pack(do_ptr, g_idx) + do_b_packs_gemm3[mt][ks] = arith.select(in_bnd, do_raw, c_zero_mfma_pack) + + # 6N-2: de-dup GEMM1 + softmax + pT_write per mt-tile. + # Each mt is owned by one wave (wave_id == mt % NUM_WAVES). + # The owning wave runs everything; the other yields zeros. + # delta_per_tile must be valid in BOTH waves (used after barrier + # for the dsT compute), so we read it unconditionally. + s_accs = [None for _ in range(M_TILES)] + pT_vals_per_tile = [None for _ in range(M_TILES)] + delta_per_tile = [] + for mt in range_constexpr(M_TILES): + # delta is needed by every wave (dsT compute happens after + # GEMM3 dp_accs[mt] is restored from LDS-broadcast / read). + # In de-dup mode dsT is owned by mt-owner only, so delta + # also only needs to be live in the owning wave. But the + # buffer_load is a cheap SMEM op; we leave it unconditional + # for simplicity. Same for lse below. + m_row_for_lse = m_start + arith.index(mt * 16) + lane_mod_16 + m_row_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, m_row_for_lse, seq_len_q_v) + m_row_safe = arith.select(m_row_in_bounds, m_row_for_lse, arith.index(0)) + lse_off_i32 = arith.index_cast(T.i32, bh_base_tokens_q_of(qhid) + m_row_safe) + lse_v = buffer_ops.buffer_load(lse_rsrc, lse_off_i32, vec_width=1, dtype=T.f32) + delta_v = buffer_ops.buffer_load(deltas_rsrc, lse_off_i32, vec_width=1, dtype=T.f32) + delta_per_tile.append(delta_v) + + if D_SPLIT_WAVES: + owns = owns_mt_preds[mt] + # ---- GATED block: GEMM1[mt] + softmax + pT LDS write ---- + # Yields (s_acc, pT0, pT1, pT2, pT3). + if_g1sm = scf.IfOp( + owns, results_=[v4f32_type, T.f32, T.f32, T.f32, T.f32], has_else=True + ) + with ir.InsertionPoint(if_g1sm.then_block): + # GEMM1[mt]: accumulate over K_STEPS_QK. + # 6N-2 iter2: load Q B-frag inside the gate when LDS Q/DO. + acc = v4f32_zero + for ks in range_constexpr(K_STEPS_QK): + k_pack = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_wave(ks)]) + if USE_LDS_FOR_Q_DO: + _m_row_b = arith.index(mt * 16) + lane_mod_16 + _d_col_b = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + _lds_idx_b = _m_row_b * arith.index(LDS_Q_STRIDE) + _d_col_b + q_pack = vector.load_op(mfma_pack_type, lds_q, [_lds_idx_b]) + else: + q_pack = q_b_packs[mt][ks] + acc = mfma_acc(k_pack, q_pack, acc) + + m_row_abs_for_mask_i32 = arith.index_cast(T.i32, m_row_for_lse) + m_oob = arith.cmpi(arith.CmpIPredicate.sge, m_row_abs_for_mask_i32, seq_len_q_i32) + pT_pieces = [] + for ii in range_constexpr(4): + ii_i32 = arith.constant(ii, type=T.i32) + kv_row_rel_i32 = arith.AddIOp( + arith.MulIOp(lane_div_16_i32, arith.constant(4, type=T.i32)).result, + ii_i32, + ).result + kv_abs_i32 = arith.AddIOp( + arith.AddIOp(kv_start_i32, wave_n_off_i32).result, kv_row_rel_i32 + ).result + kv_oob = arith.cmpi(arith.CmpIPredicate.sge, kv_abs_i32, seq_len_k_i32) + is_causal = arith.cmpi( + arith.CmpIPredicate.sgt, kv_abs_i32, m_row_abs_for_mask_i32 + ) + kv_plus_w = arith.AddIOp(kv_abs_i32, w_i32).result + is_swa = arith.cmpi( + arith.CmpIPredicate.sle, kv_plus_w, m_row_abs_for_mask_i32 + ) + bad = arith.OrIOp( + arith.OrIOp(arith.OrIOp(is_causal, is_swa).result, kv_oob).result, m_oob + ).result + s_ii = vector.extract(acc, static_position=[ii], dynamic_position=[]) + scaled = arith.MulFOp(s_ii, c_sm_scale, fastmath=fm_fast).result + scaled_m = arith.select(bad, c_neg_inf, scaled) + diff = arith.SubFOp(scaled_m, lse_v, fastmath=fm_fast).result + p = math_dialect.exp(diff, fastmath=fm_fast) + pT_pieces.append(p) + # Write pT to LDS for this mt. + kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii) + kv_row = wave_n_offset + kv_row_rel + pt_bf16 = arith.trunc_f(elem_type, p) + lds_pt_idx = ( + kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16 + ) + v1_pt = vector.from_elements(v1_elem_type, [pt_bf16]) + vector.store(v1_pt, lds_pt, [lds_pt_idx]) + scf.YieldOp([acc] + pT_pieces) + with ir.InsertionPoint(if_g1sm.else_block): + _zero_f32 = arith.constant(0.0, type=T.f32) + scf.YieldOp([v4f32_zero, _zero_f32, _zero_f32, _zero_f32, _zero_f32]) + s_accs[mt] = if_g1sm.results[0] + pT_vals_per_tile[mt] = [ + if_g1sm.results[1], + if_g1sm.results[2], + if_g1sm.results[3], + if_g1sm.results[4], + ] + else: + # Legacy path: every wave does GEMM1[mt] + softmax + pT_write. + acc = v4f32_zero + for ks in range_constexpr(K_STEPS_QK): + k_pack = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_wave(ks)]) + acc = mfma_acc(k_pack, q_b_packs[mt][ks], acc) + s_accs[mt] = acc + + m_row_abs_for_mask_i32 = arith.index_cast(T.i32, m_row_for_lse) + m_oob = arith.cmpi(arith.CmpIPredicate.sge, m_row_abs_for_mask_i32, seq_len_q_i32) + pT_tile = [] + for ii in range_constexpr(4): + ii_i32 = arith.constant(ii, type=T.i32) + kv_row_rel_i32 = arith.AddIOp( + arith.MulIOp(lane_div_16_i32, arith.constant(4, type=T.i32)).result, ii_i32 + ).result + kv_abs_i32 = arith.AddIOp( + arith.AddIOp(kv_start_i32, wave_n_off_i32).result, kv_row_rel_i32 + ).result + kv_oob = arith.cmpi(arith.CmpIPredicate.sge, kv_abs_i32, seq_len_k_i32) + is_causal = arith.cmpi( + arith.CmpIPredicate.sgt, kv_abs_i32, m_row_abs_for_mask_i32 + ) + kv_plus_w = arith.AddIOp(kv_abs_i32, w_i32).result + is_swa = arith.cmpi(arith.CmpIPredicate.sle, kv_plus_w, m_row_abs_for_mask_i32) + bad = arith.OrIOp( + arith.OrIOp(arith.OrIOp(is_causal, is_swa).result, kv_oob).result, m_oob + ).result + s_ii = vector.extract(acc, static_position=[ii], dynamic_position=[]) + scaled = arith.MulFOp(s_ii, c_sm_scale, fastmath=fm_fast).result + scaled_m = arith.select(bad, c_neg_inf, scaled) + diff = arith.SubFOp(scaled_m, lse_v, fastmath=fm_fast).result + p = math_dialect.exp(diff, fastmath=fm_fast) + pT_tile.append(p) + kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii) + kv_row = wave_n_offset + kv_row_rel + pt_bf16 = arith.trunc_f(elem_type, p) + lds_pt_idx = ( + kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16 + ) + v1_pt = vector.from_elements(v1_elem_type, [pt_bf16]) + vector.store(v1_pt, lds_pt, [lds_pt_idx]) + pT_vals_per_tile[mt] = pT_tile + + # R6k-1: D-split barrier. Both waves write the SAME pT + # values (identical s_accs derived from full-D GEMM1) to + # the SAME LDS rows (wave_n_offset=0 for both waves), but + # the writer-lane != reader-lane mapping crosses wave + # boundaries: lane L in wave 0 reads cells written by + # lanes from wave 1 (and vice versa). Without a barrier + # wave 1 could observe wave 0's stale prior-iteration pT. + if D_SPLIT_WAVES: + gpu.barrier() + + # ---- pT A-frag for GEMM2 (dV += pT @ DO): + # A[kv_row=wave_n_offset+lane_mod_16, m_col=m_step*32+lane_div_16*8 + 0..7] + pt_a_packs = [] + for m_step in range_constexpr(K_STEPS_PT): + kv_row_a = wave_n_offset + lane_mod_16 + m_col_a = arith.index(m_step * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + lds_idx = kv_row_a * arith.index(LDS_PT_STRIDE) + m_col_a + pack = vector.load_op(mfma_pack_type, lds_pt, [lds_idx]) + pt_a_packs.append(pack) + + # ---- GEMM2: dV += pT @ DO ---- + # R6b GEMM2: B-frag DO[m_step*32+lane_div_16*8+k, d_col_global+lane_mod_16]. + # R6k-1: dc_idx is wave-local; d_col_global = wave_d_col_offset + dc_idx*D_CHUNK. + if USE_LDS_FOR_Q_DO: + # R6S: tr16 LDS reads via `base + constexpr byte imm` + # so LLVM ISel folds the per-call delta into the LDS + # instruction's `offset:` immediate. Replaces the + # per-call IntToPtrOp pattern that was spilling 32 + # absolute addresses across AGPRs/scratch. + def _ds_read_tr_do_imm(byte_imm): + gep = _llvm.GEPOp( + _llvm_lds_ptr_ty(), + tr_lane_base_do_ptr, + [], + rawConstantIndices=[byte_imm], + elem_type=T.i8, + noWrapFlags=0, + ) + return rocdl.ds_read_tr16_b64(v4f16_type, gep.result).result + + def read_do_b_pack(m_step_idx, dc_idx): + # All offsets are Python ints (constexpr unrolled). + base_elem = m_step_idx * K_STEP_QK * LDS_DO_STRIDE + dc_idx * D_CHUNK + byte_imm_0 = base_elem * 2 + byte_imm_1 = byte_imm_0 + 4 * LDS_DO_STRIDE * 2 + v_lo = _ds_read_tr_do_imm(byte_imm_0) + v_hi = _ds_read_tr_do_imm(byte_imm_1) + v_full = vector.shuffle(v_lo, v_hi, [0, 1, 2, 3, 4, 5, 6, 7]) + return vector.bitcast(mfma_pack_type, v_full) + + else: + + def read_do_b_pack(m_step_idx, dc_idx): + d_col = wave_d_col_offset + arith.index(dc_idx * D_CHUNK) + lane_mod_16 + m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + vals = [] + for rk in range_constexpr(MFMA_LANE_K): + m_row_rel = m_base + arith.index(rk) + m_row_abs = m_start + m_row_rel + in_b = arith.cmpi(arith.CmpIPredicate.slt, m_row_abs, seq_len_q_v) + m_row_safe2 = arith.select(in_b, m_row_abs, arith.index(0)) + g_idx = global_idx_q(qhid, m_row_safe2, d_col) + v1 = _gep_load(do_ptr, g_idx, T.vec(1, elem_type)) + v_scalar = vector.extract(v1, static_position=[0], dynamic_position=[]) + v_safe = arith.select(in_b, v_scalar, c_zero_elem) + vals.append(v_safe) + return vector.from_elements(mfma_pack_type, vals) + + # R6k-1: dc loops wave-local (D_CHUNKS_LOCAL); reader adds + # wave_d_col_offset internally. + new_dv_accs = list(dv_accs) + for dc in range_constexpr(D_CHUNKS_LOCAL): + for pks in range_constexpr(K_STEPS_PT): + b_pack = read_do_b_pack(pks, dc) + new_dv_accs[dc] = mfma_acc(pt_a_packs[pks], b_pack, new_dv_accs[dc]) + + # 6N-2: de-dup GEMM3 + dsT + dsT_write per mt-tile. + # The owning wave for mt does GEMM3 MFMA, dsT compute, and LDS write. + # The other wave skips entirely. GEMM4 readers see the full + # dsT via LDS broadcast after the trailing barrier. + if D_SPLIT_WAVES: + for mt in range_constexpr(M_TILES): + owns = owns_mt_preds[mt] + if_g3 = scf.IfOp(owns, results_=[], has_else=False) + with ir.InsertionPoint(if_g3.then_block): + # GEMM3[mt] + # 6N-2 iter2: load DO B-frag inside the gate when LDS Q/DO. + acc = v4f32_zero + for ks in range_constexpr(K_STEPS_QK): + v_pack = vector.load_op(mfma_pack_type, lds_kv, [_v_idx_wave(ks)]) + if USE_LDS_FOR_Q_DO: + _m_row_b = arith.index(mt * 16) + lane_mod_16 + _d_col_b = arith.index(ks * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + _lds_idx_b = _m_row_b * arith.index(LDS_DO_STRIDE) + _d_col_b + do_pack = vector.load_op(mfma_pack_type, lds_do, [_lds_idx_b]) + else: + do_pack = do_b_packs_gemm3[mt][ks] + acc = mfma_acc(v_pack, do_pack, acc) + # dsT + LDS write per ii. + for ii in range_constexpr(4): + dp_ii = vector.extract(acc, static_position=[ii], dynamic_position=[]) + diff = arith.SubFOp(dp_ii, delta_per_tile[mt], fastmath=fm_fast).result + ds_ii = arith.MulFOp(pT_vals_per_tile[mt][ii], diff, fastmath=fm_fast).result + kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii) + kv_row = wave_n_offset + kv_row_rel + ds_bf16 = arith.trunc_f(elem_type, ds_ii) + lds_pt_idx = ( + kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16 + ) + v1_ds = vector.from_elements(v1_elem_type, [ds_bf16]) + vector.store(v1_ds, lds_pt, [lds_pt_idx]) + scf.YieldOp([]) + else: + # Legacy path + dp_accs = [v4f32_zero for _ in range(M_TILES)] + for ks in range_constexpr(K_STEPS_QK): + v_pack = vector.load_op(mfma_pack_type, lds_kv, [_v_idx_wave(ks)]) + for mt in range_constexpr(M_TILES): + dp_accs[mt] = mfma_acc(v_pack, do_b_packs_gemm3[mt][ks], dp_accs[mt]) + for mt in range_constexpr(M_TILES): + for ii in range_constexpr(4): + dp_ii = vector.extract(dp_accs[mt], static_position=[ii], dynamic_position=[]) + diff = arith.SubFOp(dp_ii, delta_per_tile[mt], fastmath=fm_fast).result + ds_ii = arith.MulFOp(pT_vals_per_tile[mt][ii], diff, fastmath=fm_fast).result + kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii) + kv_row = wave_n_offset + kv_row_rel + ds_bf16 = arith.trunc_f(elem_type, ds_ii) + lds_pt_idx = ( + kv_row * arith.index(LDS_PT_STRIDE) + arith.index(mt * 16) + lane_mod_16 + ) + v1_ds = vector.from_elements(v1_elem_type, [ds_bf16]) + vector.store(v1_ds, lds_pt, [lds_pt_idx]) + + # R6k-1: D-split barrier (see GEMM2 pT barrier comment). + if D_SPLIT_WAVES: + gpu.barrier() + + # ---- dsT A-frag for GEMM4 (dK += dsT @ Q): + # A[kv_row=wave_n_offset+lane_mod_16, m_col=m_step*32+lane_div_16*8 + 0..7] + ds_a_packs = [] + for m_step in range_constexpr(K_STEPS_PT): + kv_row_a = wave_n_offset + lane_mod_16 + m_col_a = arith.index(m_step * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + lds_idx = kv_row_a * arith.index(LDS_PT_STRIDE) + m_col_a + pack = vector.load_op(mfma_pack_type, lds_pt, [lds_idx]) + ds_a_packs.append(pack) + + # ---- GEMM4: dK += dsT @ Q (MFMA 16x16x32) ---- + # R6k-1: Q B-frag d_col = wave_d_col_offset + dc_idx*D_CHUNK + lane_mod_16. + if USE_LDS_FOR_Q_DO: + # R6S: tr16 LDS reads via `base + constexpr byte imm` + # (same recipe as GEMM2). Uses tr_lane_base_q_ptr. + def _ds_read_tr_q_imm(byte_imm): + gep = _llvm.GEPOp( + _llvm_lds_ptr_ty(), + tr_lane_base_q_ptr, + [], + rawConstantIndices=[byte_imm], + elem_type=T.i8, + noWrapFlags=0, + ) + return rocdl.ds_read_tr16_b64(v4f16_type, gep.result).result + + def read_q_b_pack(m_step_idx, dc_idx): + base_elem = m_step_idx * K_STEP_QK * LDS_Q_STRIDE + dc_idx * D_CHUNK + byte_imm_0 = base_elem * 2 + byte_imm_1 = byte_imm_0 + 4 * LDS_Q_STRIDE * 2 + v_lo = _ds_read_tr_q_imm(byte_imm_0) + v_hi = _ds_read_tr_q_imm(byte_imm_1) + v_full = vector.shuffle(v_lo, v_hi, [0, 1, 2, 3, 4, 5, 6, 7]) + return vector.bitcast(mfma_pack_type, v_full) + + else: + + def read_q_b_pack(m_step_idx, dc_idx): + d_col = wave_d_col_offset + arith.index(dc_idx * D_CHUNK) + lane_mod_16 + m_base = arith.index(m_step_idx * K_STEP_QK) + lane_div_16 * MFMA_LANE_K + vals = [] + for rk in range_constexpr(MFMA_LANE_K): + m_row_rel = m_base + arith.index(rk) + m_row_abs = m_start + m_row_rel + in_b = arith.cmpi(arith.CmpIPredicate.slt, m_row_abs, seq_len_q_v) + m_row_safe2 = arith.select(in_b, m_row_abs, arith.index(0)) + g_idx = global_idx_q(qhid, m_row_safe2, d_col) + v1 = _gep_load(q_ptr, g_idx, T.vec(1, elem_type)) + v_scalar = vector.extract(v1, static_position=[0], dynamic_position=[]) + v_safe = arith.select(in_b, v_scalar, c_zero_elem) + vals.append(v_safe) + return vector.from_elements(mfma_pack_type, vals) + + # R6k-1: dc loops wave-local; reader adds wave_d_col_offset. + new_dk_accs = list(dk_accs) + for dc in range_constexpr(D_CHUNKS_LOCAL): + for pks in range_constexpr(K_STEPS_PT): + q_b_pack = read_q_b_pack(pks, dc) + new_dk_accs[dc] = mfma_acc(ds_a_packs[pks], q_b_pack, new_dk_accs[dc]) + + gpu.barrier() + + yield list(new_dv_accs) + list(new_dk_accs) + + yield list(m_loop_results) + outer_carry = list(h_loop_results) + + # ---- Final: dK *= sm_scale, store ---- + # R6k-1: D-split mode -> each wave only owns D_CHUNKS_LOCAL chunks + # and writes its own slice of dK/dV (no inter-wave reduction needed: + # the D-axis is disjoint per wave, and the N-axis is shared so the + # accumulators ARE complete for each wave's D-slice). + dv_finals = [outer_carry[dc] for dc in range(D_CHUNKS_LOCAL)] + dk_finals = [outer_carry[D_CHUNKS_LOCAL + dc] for dc in range(D_CHUNKS_LOCAL)] + + for dc in range_constexpr(D_CHUNKS_LOCAL): + for ii in range_constexpr(4): + kv_row_rel = lane_div_16 * arith.index(4) + arith.index(ii) + kv_row_abs = kv_start + wave_n_offset + kv_row_rel + d_col_abs = wave_d_col_offset + arith.index(dc * D_CHUNK) + lane_mod_16 + kv_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, kv_row_abs, seq_len_k_v) + _if_kv = scf.IfOp(kv_in_bounds) + with ir.InsertionPoint(_if_kv.then_block): + dv_val = vector.extract(dv_finals[dc], static_position=[ii], dynamic_position=[]) + dk_val = vector.extract(dk_finals[dc], static_position=[ii], dynamic_position=[]) + dk_scaled = arith.MulFOp(dk_val, c_sm_scale, fastmath=fm_fast).result + g_idx = global_idx_kv(kv_row_abs, d_col_abs) + _gep_store_f32(dv_val, dv_ptr, g_idx) + _gep_store_f32(dk_scaled, dk_ptr, g_idx) + scf.YieldOp([]) + + @flyc.jit + def launch_v4_swa_bwd_dkv( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + DOS: fx.Tensor, + LSE: fx.Tensor, + DELTAS: fx.Tensor, + DK: fx.Tensor, + DV: fx.Tensor, + batch_size: fx.Int32, + seq_len_q: fx.Int32, + seq_len_k: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + bs_idx = arith.index_cast(T.index, batch_size) + sk_idx = arith.index_cast(T.index, seq_len_k) + num_n_blocks = (sk_idx + BLOCK_N - 1) // BLOCK_N + grid_x = bs_idx * num_n_blocks + + launcher = v4_swa_bwd_dkv_kernel( + Q, + K, + V, + DOS, + LSE, + DELTAS, + DK, + DV, + seq_len_q, + seq_len_k, + ) + + if waves_per_eu is not None: + _wpe = int(waves_per_eu) + if _wpe >= 1: + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe) + if flat_work_group_size is not None: + _fwgs = int(flat_work_group_size) + if _fwgs >= 1: + flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}") + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr + + passthrough_entries = [] + if daz: + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("denormal-fp-math-f32"), + ir.StringAttr.get("preserve-sign,preserve-sign"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("no-nans-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("unsafe-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries) + + launcher.launch( + grid=(grid_x, 1, 1), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) + + _fmha_compile_hints = { + "fast_fp_math": fast_fp_math, + "unsafe_fp_math": unsafe_fp_math, + "llvm_options": { + "enable-post-misched": True, # R6L iter5 + "lsr-drop-solution": True, + }, + } + + def _launch(*args, **kwargs): + with CompilationContext.compile_hints(_fmha_compile_hints): + return launch_v4_swa_bwd_dkv(*args, **kwargs) + + def _compile(Q, K, V, DOS, LSE, DELTAS, DK, DV, batch_size, seq_len_q, seq_len_k, stream=None): + with CompilationContext.compile_hints(_fmha_compile_hints): + return flyc.compile( + launch_v4_swa_bwd_dkv, + Q, + K, + V, + DOS, + LSE, + DELTAS, + DK, + DV, + batch_size, + seq_len_q, + seq_len_k, + fx.Stream(stream), + ) + + _launch.compile = _compile + + return _launch + + +build_v4_swa_bwd_dkv_module_primary = build_v4_swa_bwd_dkv_module diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dq_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dq_kernel.py new file mode 100644 index 000000000..470223ae0 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_dq_kernel.py @@ -0,0 +1,1284 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""v4_swa_bwd_dq: V4 SWA-causal attention backward dQ kernel for FlyDSL. + +Forked from kernels/sla_bwd_dq.py (FlyDSL SLA backward dQ). Differences: + - Outer KV loop iterates a contiguous block range bounded by the SWA + window for the current Q tile (no LUT). + - Per-element SWA + causal mask + boundary mask (mirrors Triton ref). + - LSE is RAW-domain (Triton: lse = m + ln(l), in domain qk*sm_scale), + so p = exp(qk*sm_scale - lse) (NOT exp2). + - sm_scale is applied TWICE: once inside the loop on qk (matches + Triton: qk = qk * sm_scale), and once after the n-loop on dq + (matches Triton: dq = dq * sm_scale). The two multiplies are NOT + redundant: ds = p*(dp - dvec) carries the in-loop scale via P; + the post-loop multiply is the outer chain-rule scale on dq. + - MQA: K/V are [B, 1, Sk, D] - stride_kh = 0 - drop head_idx from + KV indexing when mqa_kv=True. + - Sink: optional SINK[HQ] fp32; dsink = -sum(p_sink * dvec) via + atomic_fadd one scalar per row-owner-lane into DSINK[qhid]. Sink + does NOT change dq (fwd already folded p_sink into lse). + - dq written as fp32 (matches the launcher's fp32 dq_fp32 buffer). + +Layout: BHLD. Q/DOUT/DQ flat from (B, HQ, Sq, D). K/V flat from + (B, HK, Sk, D); MQA means HK=1, no head stride applied to KV. +LSE/DELTAS: (B, HQ, Sq) flat, fp32, raw-domain. +SINK: (HQ,) fp32 - dummy buffer when has_sink=False. +DSINK: (HQ,) fp32 - atomic accumulator - caller must zero-init. +""" + +import math +import os + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import math as math_dialect +from flydsl._mlir.dialects import memref as _memref +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr +from kernels.kernels_common import dtype_to_elem_type + +# ---- Module-level constants ---- + +KERNEL_NAME = "v4_swa_bwd_dq_kernel" + +_LOG2E = math.log2(math.e) # 1.4426950408889634 + +_LLVM_GEP_DYNAMIC = -2147483648 # LLVM kDynamicIndex sentinel (0x80000000 as signed i32) + + +def _llvm_ptr_ty(): + return ir.Type.parse("!llvm.ptr") + + +def _llvm_lds_ptr_ty(): + return ir.Type.parse("!llvm.ptr<3>") + + +_VMCNT_LO_MASK = 0xF +_LGKMCNT_EXPCNT_BASE = 0x3F70 +_VMCNT_HI_SHIFT = 14 +_VMCNT_HI_MASK = 0x3 + + +def _waitcnt_vm_n(n): + """Emit s_waitcnt vmcnt(n) only (lgkmcnt=63, expcnt=7).""" + val = (n & _VMCNT_LO_MASK) | _LGKMCNT_EXPCNT_BASE | (((n >> 4) & _VMCNT_HI_MASK) << _VMCNT_HI_SHIFT) + rocdl.s_waitcnt(val) + + +def build_v4_swa_bwd_dq_module( + num_heads, + head_dim, + swa_window, + dtype_str="bf16", + sm_scale=None, + waves_per_eu=2, + flat_work_group_size=None, + block_m=None, + unsafe_fp_math=True, + fast_fp_math=True, + daz=True, + layout_bhld=True, + mqa_kv=True, + has_sink=True, +): + """Build the V4 SWA backward dQ launcher. + + Args: + num_heads: HQ (number of Q heads). + head_dim: D (head dimension), must be % 32 == 0 and >= 64. + swa_window: int > 0; SWA window length. + sm_scale: defaults to 1/sqrt(head_dim). + mqa_kv: if True, K/V indexing drops head_idx (HK=1, stride_h=0). + has_sink: if True, SINK/DSINK tensors are used. If False, both + are dummy buffers (the kernel still takes them but doesn't + touch them). + layout_bhld: BHLD layout (True) or BLHD (False); V4 uses BHLD. + + Returns: launch(Q, K, V, DOUT, LSE, DELTAS, DQ_FP32, DSINK, SINK, + batch_size, seq_len_q, seq_len_k, stream=None). + """ + gpu_arch = get_hip_arch() + + BLOCK_N = 64 + K_SUB_N = 32 + WARP_SIZE = 64 + + if block_m is not None: + BLOCK_M = block_m + else: + BLOCK_M = 128 + + if flat_work_group_size is None: + if BLOCK_M <= 128: + flat_work_group_size = 256 + else: + flat_work_group_size = 512 + NUM_WAVES = flat_work_group_size // WARP_SIZE + BLOCK_SIZE = flat_work_group_size + ROWS_PER_WAVE = BLOCK_M // NUM_WAVES + # V4 SWA: dense path. One outer iter = one BLOCK_N block. + BLOCK_N_OUT = BLOCK_N + ENABLE_PREFETCH_3BUF = os.getenv("FLYDSL_SLA_FWD_ENABLE_PREFETCH3", "0") == "1" + _has_lds_load_b128 = not gpu_arch.startswith("gfx942") + ENABLE_DMA = _has_lds_load_b128 and (os.getenv("FLYDSL_SLA_FWD_ENABLE_DMA", "1") == "1") + ENABLE_LDS_VEC16 = os.getenv("FLYDSL_SLA_FWD_ENABLE_LDS_VEC16", "1") == "1" + REDUCE_MODE = os.getenv("FLYDSL_SLA_FWD_REDUCE_MODE", "xor").strip().lower() + if REDUCE_MODE not in ("xor", "ds_bpermute"): + REDUCE_MODE = "xor" + NUM_PREFETCH_K = 3 if ENABLE_PREFETCH_3BUF else (2 if ENABLE_DMA else 1) + NUM_PREFETCH_V = 3 if ENABLE_PREFETCH_3BUF else (2 if ENABLE_DMA else 1) + (1, 2, 0, 1, 0, 1, 2, 0) if ENABLE_PREFETCH_3BUF else (0,) + + USE_HW_TR = gpu_arch.startswith("gfx950") + USE_K16 = gpu_arch.startswith("gfx950") + + # Auto-fallback: gfx950 LDS limit is 160 KB. K+V double-buffering at D=512 + # easily exceeds that. If predicted LDS exceeds budget, force DMA off + # (NUM_PREFETCH = 1) and warn. Keeps the kernel correct regardless of + # the env knob. + _LDS_LIMIT_BYTES = 160 * 1024 + + def _predicted_lds_bytes(nk, nv, dma): + # USE_HW_TR & DMA -> V_STRIDE = head_dim + # USE_HW_TR & !DMA -> V_STRIDE = head_dim + 4 + # !USE_HW_TR -> V stored transposed: LDS V = head_dim * (BLOCK_N + 2) + if USE_HW_TR: + v_str = head_dim if dma else head_dim + 4 + return (nk * BLOCK_N * head_dim + nv * BLOCK_N * v_str) * 2 + vt = BLOCK_N + 2 + return (nk * BLOCK_N * head_dim + nv * head_dim * vt) * 2 + + _pred = _predicted_lds_bytes(NUM_PREFETCH_K, NUM_PREFETCH_V, ENABLE_DMA) + if _pred > _LDS_LIMIT_BYTES and ENABLE_DMA: + # Try DMA off (single-buffered). + ENABLE_DMA = False + NUM_PREFETCH_K = 1 + NUM_PREFETCH_V = 1 + _pred2 = _predicted_lds_bytes(1, 1, False) + if _pred2 > _LDS_LIMIT_BYTES: + raise RuntimeError( + f"v4_swa_bwd_dq: predicted LDS {_pred2}B > limit {_LDS_LIMIT_BYTES}B " + f"even single-buffered at D={head_dim}, BLOCK_N={BLOCK_N}" + ) + import sys as _sys + + print( + f"[v4_swa_bwd_dq] LDS overflow at D={head_dim}, BLOCK_N={BLOCK_N}: " + f"auto-disabled DMA (predicted {_pred} -> {_pred2} bytes)", + file=_sys.stderr, + flush=True, + ) + K_STEP_QK = 16 if USE_K16 else 8 + K_STEPS_QK = head_dim // K_STEP_QK + D_CHUNK = 32 + D_CHUNKS = head_dim // D_CHUNK + PV_K_STEP = 16 if USE_K16 else 8 + PV_K_STEPS = K_SUB_N // PV_K_STEP # 2 steps per sub-tile (K=16) or 4 (K=8) + + assert BLOCK_M % NUM_WAVES == 0 + assert head_dim % 32 == 0, f"head_dim ({head_dim}) must be divisible by 32" + assert head_dim >= 64, f"head_dim ({head_dim}) must be >= 64" + assert flat_work_group_size in ( + 128, + 256, + 512, + ), f"flat_work_group_size must be 128, 256, or 512, got {flat_work_group_size}" + assert dtype_str in ("f16", "bf16"), "v4_swa_bwd_dq only supports f16 and bf16" + assert BLOCK_N % 32 == 0 + assert BLOCK_N_OUT == BLOCK_N + assert isinstance(swa_window, int) and swa_window > 0, f"swa_window must be int > 0, got {swa_window!r}" + + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(head_dim) + + NUM_HEADS = num_heads + HEAD_DIM = head_dim + STRIDE_TOKEN = NUM_HEADS * HEAD_DIM + + K_STRIDE = HEAD_DIM + # XOR swizzle mask must fit within the row stride. The swizzle is + # applied at 16-element granularity, so the maximum mask is + # `(K_STRIDE // 16 - 1) << 4`. For D=128 that's 7 (=0x7), giving max + # mask 112 < 128. For D=64 it's 3 (=0x3), giving max mask 48 < 64. + # SLA test only covers D=128 (mask=7); using a hardcoded `& 7` for + # D=64 wraps writes into adjacent rows -> silent corruption. + K_SWZ_ROW_MASK = (K_STRIDE // 16) - 1 + assert K_SWZ_ROW_MASK >= 0 + assert ( + K_SWZ_ROW_MASK & (K_SWZ_ROW_MASK + 1) + ) == 0, f"K_SWZ_ROW_MASK must be 2^n-1, got {K_SWZ_ROW_MASK} (K_STRIDE={K_STRIDE})" + if USE_HW_TR: + V_STRIDE = HEAD_DIM if ENABLE_DMA else HEAD_DIM + 4 + else: + VT_STRIDE = BLOCK_N + 2 + V_STRIDE = VT_STRIDE + # V swizzle: similarly bounded by V_STRIDE / 16. + V_SWZ_ROW_MASK = min(3, (V_STRIDE // 16) - 1) + assert V_SWZ_ROW_MASK >= 0 + + VEC_WIDTH = 16 if ENABLE_LDS_VEC16 else 8 + assert HEAD_DIM % VEC_WIDTH == 0 + THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH + assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0 + ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD + + if ROWS_PER_BATCH_LOAD >= BLOCK_N: + NUM_BATCHES_KV = 1 + KV_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_N + else: + assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0 + NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD + KV_NEEDS_GUARD = False + + LDS_K_TILE_SIZE = BLOCK_N * K_STRIDE + if USE_HW_TR: + LDS_V_TILE_SIZE = BLOCK_N * V_STRIDE + else: + LDS_V_TILE_SIZE = HEAD_DIM * VT_STRIDE + LDS_K_TOTAL_SIZE = NUM_PREFETCH_K * LDS_K_TILE_SIZE + LDS_V_BASE = LDS_K_TOTAL_SIZE + LDS_V_TOTAL_SIZE = NUM_PREFETCH_V * LDS_V_TILE_SIZE + LDS_KV_TOTAL_SIZE = LDS_K_TOTAL_SIZE + LDS_V_TOTAL_SIZE + + allocator = SmemAllocator( + None, + arch=gpu_arch, + global_sym_name=f"v4_swa_bwd_dq_smem_M{BLOCK_M}_W{swa_window}_S{int(has_sink)}_MQ{int(mqa_kv)}", + ) + lds_kv_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_kv_offset + LDS_KV_TOTAL_SIZE * 2 + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def v4_swa_bwd_dq_kernel( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + DOS: fx.Tensor, # grad of output (input) + LSE: fx.Tensor, # log-sum-exp from fwd (input, f32, RAW domain) + DELTAS: fx.Tensor, # (o * do).sum(-1) preprocess (input, f32) + DQ: fx.Tensor, # grad wrt Q (output, FP32) + DSINK: fx.Tensor, # grad wrt sink (output, FP32, [HQ]) + SINK: fx.Tensor, # sink param (input, FP32, [HQ]); dummy if !has_sink + seq_len_q: fx.Int32, + seq_len_k: fx.Int32, + ): + elem_type = dtype_to_elem_type(dtype_str) + compute_type = T.f32 + q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q) + k_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K) + v_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V) + do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DOS) + dq_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), DQ) + # LSE / DELTAS: f32 scalar READS via buffer_load. + lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True) + deltas_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True) + sink_rsrc = buffer_ops.create_buffer_resource(SINK, max_size=True) + # DSINK: f32 atomic_fadd via buffer_atomic_fadd. + dsink_rsrc = buffer_ops.create_buffer_resource(DSINK, max_size=True) + + # All FP operations use aggressive fast-math (no NaN/Inf checks, reassociation). + fm_fast = arith.FastMathFlags.fast + v4f16_type = T.vec(4, elem_type) + vxf16_type = T.vec(VEC_WIDTH, elem_type) + v8f16_type = T.vec(8, elem_type) + v16f32_type = T.vec(16, compute_type) + mfma_pack_type = v8f16_type if USE_K16 else v4f16_type + MFMA_LANE_K = 8 if USE_K16 else 4 + _mfma_zero = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0) + + def _mfma(ods_fn, a, b, c): + return ods_fn(v16f32_type, a, b, c, _mfma_zero, _mfma_zero, _mfma_zero).result + + def mfma_acc(a, b, c): + if dtype_str == "bf16": + if USE_K16: + return _mfma(rocdl.mfma_f32_32x32x16_bf16, a, b, c) + a = vector.bitcast(T.i16x4, a) + b = vector.bitcast(T.i16x4, b) + return _mfma(rocdl.mfma_f32_32x32x8bf16_1k, a, b, c) + if USE_K16: + return _mfma(rocdl.mfma_f32_32x32x16_f16, a, b, c) + return _mfma(rocdl.mfma_f32_32x32x8f16, a, b, c) + + seq_len_q_v = arith.index_cast(T.index, seq_len_q) + seq_len_k_v = arith.index_cast(T.index, seq_len_k) + + # ---- LDS view ---- + base_ptr = allocator.get_base() + lds_kv = SmemPtr( + base_ptr, + lds_kv_offset, + elem_type, + shape=(LDS_KV_TOTAL_SIZE,), + ).get() + + # ---- Thread / block indices ---- + block_id = arith.index_cast(T.index, gpu.block_idx.x) + tid = arith.index_cast(T.index, gpu.thread_idx.x) + + wave_id = tid // WARP_SIZE + lane = tid % WARP_SIZE + lane_mod_32 = lane % 32 + lane_div_32 = lane // 32 # 0/1 + + # ds_read_b64_tr_b16 lane decomposition + tr_k_group = (lane % 16) // 4 + tr_col_sub = lane % 4 + tr_col_half = (lane % 32) // 16 + + def ds_read_tr_v4f16(lds_elem_idx): + byte_offset = lds_elem_idx * 2 + lds_kv_offset + byte_i64 = arith.index_cast(T.i64, byte_offset) + ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result + return rocdl.ds_read_tr16_b64(v4f16_type, ptr).result + + wave_q_offset = wave_id * ROWS_PER_WAVE + + # ---- Decompose block_id: (batch, q_tile, head) ---- + head_idx = block_id % NUM_HEADS + batch_q_tile_id = block_id // NUM_HEADS + num_q_tiles = (seq_len_q_v + BLOCK_M - 1) // BLOCK_M + q_tile_idx = batch_q_tile_id % num_q_tiles + batch_idx = batch_q_tile_id // num_q_tiles + q_start = q_tile_idx * BLOCK_M + + # ---- V4 SWA: per-tile contiguous K-block range (wave-uniform) ---- + SWA = arith.index(swa_window) + BN = arith.index(BLOCK_N) + BM = arith.index(BLOCK_M) + _zero_idx = arith.index(0) + _one_idx = arith.index(1) + # n_block_start = max(0, q_start - W + 1) // BLOCK_N + # n_block_end = ceil(min(q_start + BLOCK_M, seq_len_k), BLOCK_N) + _q_plus_one = q_start + _one_idx + _ge_w = arith.cmpi(arith.CmpIPredicate.sge, _q_plus_one, SWA) + _n_start_row = arith.select(_ge_w, _q_plus_one - SWA, _zero_idx) + n_block_start = _n_start_row // BN + _n_end_row_uncl = q_start + BM + _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _n_end_row_uncl, seq_len_k_v) + n_end_row_cl = arith.select(_le_seq, _n_end_row_uncl, seq_len_k_v) + n_block_end = (n_end_row_cl + BN - _one_idx) // BN + + # ---- Cooperative load decomposition ---- + load_row_in_batch = tid // THREADS_PER_ROW_LOAD + load_lane_in_row = tid % THREADS_PER_ROW_LOAD + load_col_base = load_lane_in_row * VEC_WIDTH + + # ---- Helper: global flat index ---- + if layout_bhld: + bh_base_tokens_q = (batch_idx * NUM_HEADS + head_idx) * seq_len_q_v + if mqa_kv: + bh_base_tokens_kv = batch_idx * seq_len_k_v + else: + bh_base_tokens_kv = (batch_idx * NUM_HEADS + head_idx) * seq_len_k_v + + def global_idx_q(token_idx, col): + return (bh_base_tokens_q + token_idx) * arith.index(HEAD_DIM) + col + + def global_idx_kv(token_idx, col): + return (bh_base_tokens_kv + token_idx) * arith.index(HEAD_DIM) + col + + else: + # BLHD path: kept for symmetry but V4 uses BHLD. + def global_idx_q(token_idx, col): + token = batch_idx * seq_len_q_v + token_idx + return token * STRIDE_TOKEN + head_idx * HEAD_DIM + col + + if mqa_kv: + + def global_idx_kv(token_idx, col): + token = batch_idx * seq_len_k_v + token_idx + return token * HEAD_DIM + col + + else: + + def global_idx_kv(token_idx, col): + token = batch_idx * seq_len_k_v + token_idx + return token * STRIDE_TOKEN + head_idx * HEAD_DIM + col + + def _gep_load(base_ptr_, elem_idx, vec_type, et=elem_type): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr_, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=et, + noWrapFlags=0, + ) + return _llvm.LoadOp(vec_type, gep.result).result + + def _gep_store_f32(val, base_ptr_, elem_idx): + """Store a single f32 value via GEP into an fp32 buffer.""" + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr_, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=T.f32, + noWrapFlags=0, + ) + _llvm.StoreOp(val, gep.result) + + def load_global_mfma_pack(base_ptr_, base_idx): + return _gep_load(base_ptr_, base_idx, mfma_pack_type) + + def load_global_f16xN(base_ptr_, base_idx): + return _gep_load(base_ptr_, base_idx, vxf16_type) + + def bf16_trunc_pack_v4(f32_vals): + _v2i32 = T.vec(2, T.i32) + _c16 = arith.constant(16, type=T.i32) + _cmask = arith.constant(0xFFFF0000, type=T.i32) + a0 = arith.ArithValue(f32_vals[0]).bitcast(T.i32) + b0 = arith.ArithValue(f32_vals[1]).bitcast(T.i32) + p0 = arith.OrIOp(arith.AndIOp(b0, _cmask).result, arith.ShRUIOp(a0, _c16).result).result + a1 = arith.ArithValue(f32_vals[2]).bitcast(T.i32) + b1 = arith.ArithValue(f32_vals[3]).bitcast(T.i32) + p1 = arith.OrIOp(arith.AndIOp(b1, _cmask).result, arith.ShRUIOp(a1, _c16).result).result + return vector.bitcast(v4f16_type, vector.from_elements(_v2i32, [p0, p1])) + + def bf16_trunc_pack_v8(f32_vals): + _v4i32 = T.vec(4, T.i32) + _c16 = arith.constant(16, type=T.i32) + _cmask = arith.constant(0xFFFF0000, type=T.i32) + pairs = [] + for j in range_constexpr(4): + a = arith.ArithValue(f32_vals[j * 2]).bitcast(T.i32) + b = arith.ArithValue(f32_vals[j * 2 + 1]).bitcast(T.i32) + p = arith.OrIOp(arith.AndIOp(b, _cmask).result, arith.ShRUIOp(a, _c16).result).result + pairs.append(p) + return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs)) + + def k_buf_base(buf_id): + if isinstance(buf_id, int): + return arith.index(buf_id * LDS_K_TILE_SIZE) + return buf_id * arith.index(LDS_K_TILE_SIZE) + + def v_buf_base(buf_id): + if isinstance(buf_id, int): + return arith.index(LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) + return arith.index(LDS_V_BASE) + buf_id * arith.index(LDS_V_TILE_SIZE) + + # ---- K XOR swizzle: col ^ ((row & K_SWZ_ROW_MASK) << 4) at 16-element granularity ---- + # K_SWZ_ROW_MASK derived from K_STRIDE to stay within the row. + def _k_swizzle(row_idx, col_idx): + mask = (row_idx & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + return col_idx ^ mask + + # ---- Cooperative K load (row-major, XOR-swizzled) ---- + def coop_load_k(tile_start, buf_id=0): + k_base = k_buf_base(buf_id) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + _if_k = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_k.then_block): + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + vec = load_global_f16xN(k_ptr, g_idx) + vector.store(vec, lds_kv, [lds_idx]) + scf.YieldOp([]) + else: + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + vec = load_global_f16xN(k_ptr, g_idx) + vector.store(vec, lds_kv, [lds_idx]) + + # ---- Cooperative V-into-K-LDS-slot load (K-style XOR swizzle) ---- + def coop_load_v_as_k(tile_start, buf_id=0): + k_base = k_buf_base(buf_id) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + _if_v = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_v.then_block): + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + vec = load_global_f16xN(v_ptr, g_idx) + vector.store(vec, lds_kv, [lds_idx]) + scf.YieldOp([]) + else: + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + vec = load_global_f16xN(v_ptr, g_idx) + vector.store(vec, lds_kv, [lds_idx]) + + # ---- Cooperative V load (V LDS layout) ---- + def _v_store_row_major(v_base, lds_row, vec): + lds_idx = v_base + lds_row * V_STRIDE + load_col_base + vector.store(vec, lds_kv, [lds_idx]) + + _v1_type = T.vec(1, elem_type) if not USE_HW_TR else None + + def _v_store_transposed(v_base, lds_row, vec): + for _e in range_constexpr(VEC_WIDTH): + elem = vector.extract(vec, static_position=[_e], dynamic_position=[]) + vt_d = load_col_base + _e + vt_idx = v_base + vt_d * VT_STRIDE + lds_row + v1 = vector.from_elements(_v1_type, [elem]) + vector.store(v1, lds_kv, [vt_idx]) + + _v_store_to_lds = _v_store_row_major if USE_HW_TR else _v_store_transposed + + def coop_load_v(tile_start, buf_id=0): + v_base = v_buf_base(buf_id) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + _if_v = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_v.then_block): + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + vec = load_global_f16xN(v_ptr, g_idx) + _v_store_to_lds(v_base, lds_row, vec) + scf.YieldOp([]) + else: + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + vec = load_global_f16xN(v_ptr, g_idx) + _v_store_to_lds(v_base, lds_row, vec) + + # ---- Cooperative K-into-V-LDS-slot load ---- + def coop_load_k_as_v(tile_start, buf_id=0): + v_base = v_buf_base(buf_id) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + _if_v = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_v.then_block): + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + vec = load_global_f16xN(k_ptr, g_idx) + _v_store_to_lds(v_base, lds_row, vec) + scf.YieldOp([]) + else: + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + vec = load_global_f16xN(k_ptr, g_idx) + _v_store_to_lds(v_base, lds_row, vec) + + # ---- DMA loading for K (buffer_load_dwordx4 ... lds) ---- + if ENABLE_DMA: + from flydsl._mlir.dialects import llvm + + k_rsrc = buffer_ops.create_buffer_resource(K, max_size=True) + _lds_ptr_ty = _llvm_lds_ptr_ty() + DMA_BYTES = 16 + DMA_BATCH_BYTES = BLOCK_SIZE * DMA_BYTES + K_TILE_BYTES = BLOCK_N * K_STRIDE * 2 + NUM_DMA_K = K_TILE_BYTES // DMA_BATCH_BYTES + LANES_PER_K_ROW = HEAD_DIM * 2 // DMA_BYTES + ROWS_PER_DMA_BATCH = DMA_BATCH_BYTES // (HEAD_DIM * 2) + lds_kv_base_idx = _memref.extract_aligned_pointer_as_index(lds_kv) + _dma_size = arith.constant(DMA_BYTES, type=T.i32) + _dma_soff = arith.constant(0, type=T.i32) + _dma_off = arith.constant(0, type=T.i32) + _dma_aux = arith.constant(1, type=T.i32) + + def _kv_global_byte(tile_start, row_in_tile, col_byte): + if layout_bhld: + row_within = tile_start + row_in_tile + return (bh_base_tokens_kv + row_within) * arith.index(HEAD_DIM * 2) + col_byte + else: + global_row = batch_idx * seq_len_k_v + tile_start + row_in_tile + if mqa_kv: + return global_row * arith.index(HEAD_DIM * 2) + col_byte + else: + return ( + global_row * arith.index(STRIDE_TOKEN * 2) + + head_idx * arith.index(HEAD_DIM * 2) + + col_byte + ) + + def coop_dma_k(tile_start, buf_id=0): + if isinstance(buf_id, int): + k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2) + else: + k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2) + for d in range_constexpr(NUM_DMA_K): + lds_addr = ( + k_lds_byte_base + + wave_id * arith.index(WARP_SIZE * DMA_BYTES) + + arith.index(d * DMA_BATCH_BYTES) + ) + lds_i64 = arith.index_cast(T.i64, lds_addr) + lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64) + lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result + + row_in_tile = tid // LANES_PER_K_ROW + arith.index(d * ROWS_PER_DMA_BATCH) + swiz_col_f16 = (tid % LANES_PER_K_ROW) * (DMA_BYTES // 2) + xor_mask = (row_in_tile & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + unsw_col_f16 = swiz_col_f16 ^ xor_mask + col_byte = unsw_col_f16 * 2 + global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte) + voffset = arith.index_cast(T.i32, global_byte) + rocdl.raw_ptr_buffer_load_lds( + k_rsrc, + lds_ptr, + _dma_size, + voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + def _v_swizzle(row_idx, col_idx): + mask = (row_idx & arith.index(V_SWZ_ROW_MASK)) << arith.index(4) + return col_idx ^ mask + + if ENABLE_DMA: + v_rsrc = buffer_ops.create_buffer_resource(V, max_size=True) + V_TILE_BYTES = BLOCK_N * V_STRIDE * 2 + NUM_DMA_V = V_TILE_BYTES // DMA_BATCH_BYTES + LANES_PER_V_ROW = HEAD_DIM * 2 // DMA_BYTES + ROWS_PER_DMA_BATCH_V = DMA_BATCH_BYTES // (HEAD_DIM * 2) + + def coop_dma_v(tile_start, buf_id=0): + v_lds_byte_base = lds_kv_base_idx + arith.index((LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) * 2) + for d in range_constexpr(NUM_DMA_V): + lds_addr = ( + v_lds_byte_base + + wave_id * arith.index(WARP_SIZE * DMA_BYTES) + + arith.index(d * DMA_BATCH_BYTES) + ) + lds_i64 = arith.index_cast(T.i64, lds_addr) + lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64) + lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result + + row_in_tile = tid // LANES_PER_V_ROW + arith.index(d * ROWS_PER_DMA_BATCH_V) + swiz_col_f16 = (tid % LANES_PER_V_ROW) * (DMA_BYTES // 2) + xor_mask = (row_in_tile & arith.index(V_SWZ_ROW_MASK)) << arith.index(4) + unsw_col_f16 = swiz_col_f16 ^ xor_mask + col_byte = unsw_col_f16 * 2 + global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte) + voffset = arith.index_cast(T.i32, global_byte) + rocdl.raw_ptr_buffer_load_lds( + v_rsrc, + lds_ptr, + _dma_size, + voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + # ---- Bwd dQ DMA variants (cross-pointer, matching swizzle) ---- + def coop_dma_v_as_k(tile_start, buf_id=0): + if isinstance(buf_id, int): + k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2) + else: + k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2) + for d in range_constexpr(NUM_DMA_K): + lds_addr = ( + k_lds_byte_base + + wave_id * arith.index(WARP_SIZE * DMA_BYTES) + + arith.index(d * DMA_BATCH_BYTES) + ) + lds_i64 = arith.index_cast(T.i64, lds_addr) + lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64) + lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result + + row_in_tile = tid // LANES_PER_K_ROW + arith.index(d * ROWS_PER_DMA_BATCH) + swiz_col_f16 = (tid % LANES_PER_K_ROW) * (DMA_BYTES // 2) + xor_mask = (row_in_tile & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + unsw_col_f16 = swiz_col_f16 ^ xor_mask + col_byte = unsw_col_f16 * 2 + global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte) + voffset = arith.index_cast(T.i32, global_byte) + rocdl.raw_ptr_buffer_load_lds( + v_rsrc, + lds_ptr, + _dma_size, + voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + def coop_dma_k_as_v(tile_start, buf_id=0): + if isinstance(buf_id, int): + v_lds_byte_base = lds_kv_base_idx + arith.index( + (LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) * 2 + ) + else: + v_lds_byte_base = ( + lds_kv_base_idx + + arith.index(LDS_V_BASE * 2) + + buf_id * arith.index(LDS_V_TILE_SIZE * 2) + ) + for d in range_constexpr(NUM_DMA_V): + lds_addr = ( + v_lds_byte_base + + wave_id * arith.index(WARP_SIZE * DMA_BYTES) + + arith.index(d * DMA_BATCH_BYTES) + ) + lds_i64 = arith.index_cast(T.i64, lds_addr) + lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64) + lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result + + row_in_tile = tid // LANES_PER_V_ROW + arith.index(d * ROWS_PER_DMA_BATCH_V) + swiz_col_f16 = (tid % LANES_PER_V_ROW) * (DMA_BYTES // 2) + xor_mask = (row_in_tile & arith.index(V_SWZ_ROW_MASK)) << arith.index(4) + unsw_col_f16 = swiz_col_f16 ^ xor_mask + col_byte = unsw_col_f16 * 2 + global_byte = _kv_global_byte(tile_start, row_in_tile, col_byte) + voffset = arith.index_cast(T.i32, global_byte) + rocdl.raw_ptr_buffer_load_lds( + k_rsrc, + lds_ptr, + _dma_size, + voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + # ---- Preload Q^T B-operand and DO^T B-operand packs ---- + q_row = q_start + wave_q_offset + lane_mod_32 + q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, q_row, seq_len_q_v) + q_row_safe = arith.select(q_in_bounds, q_row, arith.index(0)) + c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type) + q_b_packs = [] + do_b_packs = [] + for ks in range_constexpr(K_STEPS_QK): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + g_idx = global_idx_q(q_row_safe, col) + q_raw = load_global_mfma_pack(q_ptr, g_idx) + q_b_packs.append(arith.select(q_in_bounds, q_raw, c_zero_mfma_pack)) + do_raw = load_global_mfma_pack(do_ptr, g_idx) + do_b_packs.append(arith.select(q_in_bounds, do_raw, c_zero_mfma_pack)) + + # ---- Constants ---- + c_zero_f = arith.constant(0.0, type=compute_type) + c_neg_inf = arith.constant(-1.0e30, type=compute_type) # finite NEG_INF (matches Triton) + c_zero_v16f32 = arith.constant_vector(0.0, v16f32_type) + # V4 LSE is RAW-domain (qk*sm_scale + ln(l)). So use sm_scale, not sm_scale*log2e. + c_sm_scale = arith.constant(sm_scale, type=compute_type) + + # ---- Per-q-row scalars (LSE, delta) ---- + # bh_base_tokens_q is the Q LSE/DELTAS base too (LSE shape [B,HQ,Sq]). + lse_delta_off_i32 = arith.index_cast(T.i32, bh_base_tokens_q + q_row_safe) + lse_val = buffer_ops.buffer_load( + lse_rsrc, + lse_delta_off_i32, + vec_width=1, + dtype=T.f32, + ) + delta_val = buffer_ops.buffer_load( + deltas_rsrc, + lse_delta_off_i32, + vec_width=1, + dtype=T.f32, + ) + + # ---- SINK contribution to DSINK ---- + # Per Triton ref: + # sink_h = SINK[qhid] (uniform across lanes of this program) + # p_sink = exp(sink_h - lse) + # dsink_contrib = sum_m -p_sink_masked * dvec_masked (over BLOCK_M rows) + # atomic_fadd(DSINK + qhid, dsink_contrib) + # We do the atomic per row-owner lane (lane_div_32==0 lane) to avoid + # the cross-wave reduction. That gives BLOCK_M atomics per program - + # slow but bulletproof. Each lane contributes -p_sink * dvec for its + # owned row only when q_in_bounds. + if has_sink: + head_idx_i32 = arith.index_cast(T.i32, head_idx) + sink_h_scalar = buffer_ops.buffer_load( + sink_rsrc, + head_idx_i32, + vec_width=1, + dtype=T.f32, + ) + sink_h_uniform = rocdl.readfirstlane(T.f32, sink_h_scalar) + sub_val = arith.SubFOp( + sink_h_uniform, + lse_val, + fastmath=fm_fast, + ).result + p_sink = math_dialect.exp(sub_val, fastmath=fm_fast) + neg_p_sink = arith.SubFOp( + c_zero_f, + p_sink, + fastmath=fm_fast, + ).result + contrib = arith.MulFOp( + neg_p_sink, + delta_val, + fastmath=fm_fast, + ).result + # Gate: only the row-owner lane (lane_div_32==0) and only when + # q_row < seq_len_q, atomically adds. + is_row_owner = arith.cmpi( + arith.CmpIPredicate.eq, + lane_div_32, + arith.index(0), + ) + do_sink_atomic = arith.AndIOp(is_row_owner, q_in_bounds).result + _if_sink = scf.IfOp(do_sink_atomic, [], has_else=False) + with ir.InsertionPoint(_if_sink.then_block): + # DSINK byte offset = head_idx * 4 (fp32). + _dsink_byte_off = arith.MulIOp( + head_idx_i32, + arith.constant(4, type=T.i32), + ).result + _zero_i32_atom = arith.constant(0, type=T.i32) + rocdl.raw_ptr_buffer_atomic_fadd( + contrib, + dsink_rsrc, + _dsink_byte_off, + _zero_i32_atom, + _zero_i32_atom, + ) + scf.YieldOp([]) + + # ---- MILESTONE: dense SWA bwd dQ with double-buffered LDS ---- + _use_dbuf = ENABLE_DMA + + init_args = [] + for _ in range_constexpr(D_CHUNKS): + init_args.append(c_zero_v16f32) + if _use_dbuf: + init_args.append(arith.index(0)) # cur_buf_id + + # PROLOGUE: prefetch iter 0's K into BOTH slots of buf 0. + _init_kv_start = n_block_start * BN + coop_dma_k(_init_kv_start, buf_id=0) + coop_dma_k_as_v(_init_kv_start, buf_id=0) + + for block_idx, inner_iter_args, loop_results in scf.for_( + n_block_start, + n_block_end, + arith.index(1), + iter_args=init_args, + ): + dq_accs = [inner_iter_args[i] for i in range_constexpr(D_CHUNKS)] + if _use_dbuf: + cur_buf = inner_iter_args[D_CHUNKS] + next_buf = arith.index(1) - cur_buf + + # V4 SWA: block_idx is directly the K-block index. + kv_block_start = block_idx * BN + kv_start = kv_block_start + + if _use_dbuf: + rocdl.s_waitcnt(0) + gpu.barrier() + k_base = k_buf_base(cur_buf) + else: + coop_load_k(kv_start, buf_id=0) + coop_load_k_as_v(kv_start, buf_id=0) + gpu.barrier() + k_base = k_buf_base(0) + + # ==== GEMM1: s = Q @ K^T ==== + k_hi_offset = K_SUB_N * K_STRIDE + k_swz_mask = (lane_mod_32 & arith.index(K_SWZ_ROW_MASK)) << arith.index(4) + + def _k_idx_lo(ks): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + return k_base + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask) + + def _k_idx_hi(ks): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + return k_base + k_hi_offset + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask) + + s_acc_lo = c_zero_v16f32 + s_acc_hi = c_zero_v16f32 + for ks in range_constexpr(K_STEPS_QK): + k_pack_lo = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_lo(ks)]) + k_pack_hi = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_hi(ks)]) + s_acc_lo = mfma_acc(k_pack_lo, q_b_packs[ks], s_acc_lo) + s_acc_hi = mfma_acc(k_pack_hi, q_b_packs[ks], s_acc_hi) + + # ==== Compute p[r] = exp(qk * sm_scale - LSE) with SWA + causal + boundary mask ==== + kv_start_i32 = arith.index_cast(T.i32, kv_start) + seq_len_k_i32 = arith.index_cast(T.i32, seq_len_k_v) + seq_len_q_i32 = arith.index_cast(T.i32, seq_len_q_v) + q_row_i32_mask = arith.index_cast(T.i32, q_row) + w_i32 = arith.constant(swa_window, type=T.i32) + lane_div_32_i32 = arith.index_cast(T.i32, lane_div_32) + lane_off_i32 = arith.MulIOp(lane_div_32_i32, arith.constant(4, type=T.i32)).result + + # Is this q_row out of bounds? In that case all mask elements = NEG_INF. + q_oob = arith.cmpi( + arith.CmpIPredicate.sge, + q_row_i32_mask, + seq_len_q_i32, + ) + + p_vals_lo = [] + p_vals_hi = [] + for r in range_constexpr(16): + r_off_i32 = arith.constant((r % 4) + (r // 4) * 8, type=T.i32) + + # lo half: col = kv_start + lane_off + r_off + kv_col_lo_i32 = arith.AddIOp( + arith.AddIOp(kv_start_i32, lane_off_i32).result, r_off_i32 + ).result + # boundary + is_oob_lo = arith.cmpi(arith.CmpIPredicate.sge, kv_col_lo_i32, seq_len_k_i32) + # causal: kv_col > q_row + is_causal_lo = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_lo_i32, q_row_i32_mask) + # SWA: kv_col + W <= q_row + kv_plus_w_lo = arith.AddIOp(kv_col_lo_i32, w_i32).result + is_swa_lo = arith.cmpi(arith.CmpIPredicate.sle, kv_plus_w_lo, q_row_i32_mask) + bad_lo = arith.OrIOp( + arith.OrIOp( + arith.OrIOp(is_causal_lo, is_swa_lo).result, + is_oob_lo, + ).result, + q_oob, + ).result + + s_lo_f32 = vector.extract(s_acc_lo, static_position=[r], dynamic_position=[]) + scaled_lo = arith.MulFOp(s_lo_f32, c_sm_scale, fastmath=fm_fast).result + scaled_lo_masked = arith.select(bad_lo, c_neg_inf, scaled_lo) + diff_lo = arith.SubFOp(scaled_lo_masked, lse_val, fastmath=fm_fast).result + p_lo = math_dialect.exp(diff_lo, fastmath=fm_fast) + p_vals_lo.append(p_lo) + + # hi half: col = lo_col + K_SUB_N + kv_col_hi_i32 = arith.AddIOp(kv_col_lo_i32, arith.constant(K_SUB_N, type=T.i32)).result + is_oob_hi = arith.cmpi(arith.CmpIPredicate.sge, kv_col_hi_i32, seq_len_k_i32) + is_causal_hi = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_hi_i32, q_row_i32_mask) + kv_plus_w_hi = arith.AddIOp(kv_col_hi_i32, w_i32).result + is_swa_hi = arith.cmpi(arith.CmpIPredicate.sle, kv_plus_w_hi, q_row_i32_mask) + bad_hi = arith.OrIOp( + arith.OrIOp( + arith.OrIOp(is_causal_hi, is_swa_hi).result, + is_oob_hi, + ).result, + q_oob, + ).result + + s_hi_f32 = vector.extract(s_acc_hi, static_position=[r], dynamic_position=[]) + scaled_hi = arith.MulFOp(s_hi_f32, c_sm_scale, fastmath=fm_fast).result + scaled_hi_masked = arith.select(bad_hi, c_neg_inf, scaled_hi) + diff_hi = arith.SubFOp(scaled_hi_masked, lse_val, fastmath=fm_fast).result + p_hi = math_dialect.exp(diff_hi, fastmath=fm_fast) + p_vals_hi.append(p_hi) + + # ==== Overwrite K LDS slot with V for GEMM2 ==== + gpu.barrier() + if _use_dbuf: + coop_dma_v_as_k(kv_start, buf_id=cur_buf) + rocdl.s_waitcnt(0) + gpu.barrier() + elif ENABLE_DMA: + coop_dma_v_as_k(kv_start, buf_id=0) + rocdl.s_waitcnt(0) + gpu.barrier() + else: + coop_load_v_as_k(kv_start, buf_id=0) + gpu.barrier() + + # ==== GEMM2: dP = DO @ V^T ==== + dp_acc_lo = c_zero_v16f32 + dp_acc_hi = c_zero_v16f32 + for ks in range_constexpr(K_STEPS_QK): + v_pack_lo = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_lo(ks)]) + v_pack_hi = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_hi(ks)]) + dp_acc_lo = mfma_acc(v_pack_lo, do_b_packs[ks], dp_acc_lo) + dp_acc_hi = mfma_acc(v_pack_hi, do_b_packs[ks], dp_acc_hi) + + # ==== Prefetch iter+1's K DMAs (async) ==== + if _use_dbuf: + _next_block_idx = block_idx + arith.index(1) + _has_next = arith.cmpi(arith.CmpIPredicate.slt, _next_block_idx, n_block_end) + _pre_if = scf.IfOp(_has_next) + with ir.InsertionPoint(_pre_if.then_block): + _next_kv_start = _next_block_idx * BN + coop_dma_k(_next_kv_start, next_buf) + coop_dma_k_as_v(_next_kv_start, next_buf) + scf.YieldOp([]) + + # ==== Compute dS[r] = p[r] * (dp[r] - delta) ==== + ds_vals_lo = [] + ds_vals_hi = [] + for r in range_constexpr(16): + dp_lo = vector.extract(dp_acc_lo, static_position=[r], dynamic_position=[]) + dp_hi = vector.extract(dp_acc_hi, static_position=[r], dynamic_position=[]) + diff_lo = arith.SubFOp(dp_lo, delta_val, fastmath=fm_fast).result + diff_hi = arith.SubFOp(dp_hi, delta_val, fastmath=fm_fast).result + ds_lo = arith.MulFOp(p_vals_lo[r], diff_lo, fastmath=fm_fast).result + ds_hi = arith.MulFOp(p_vals_hi[r], diff_hi, fastmath=fm_fast).result + ds_vals_lo.append(ds_lo) + ds_vals_hi.append(ds_hi) + + # ==== Pack dS f32 -> mfma_pack_type (bf16/f16) ==== + if dtype_str == "bf16" and USE_K16: + ds_packs_lo = [] + ds_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + base = pks * 8 + ds_packs_lo.append(bf16_trunc_pack_v8(ds_vals_lo[base : base + 8])) + ds_packs_hi.append(bf16_trunc_pack_v8(ds_vals_hi[base : base + 8])) + elif dtype_str == "bf16": + ds_packs_lo = [] + ds_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + base = pks * 4 + ds_packs_lo.append(bf16_trunc_pack_v4(ds_vals_lo[base : base + 4])) + ds_packs_hi.append(bf16_trunc_pack_v4(ds_vals_hi[base : base + 4])) + else: + ds_f16_lo = [arith.trunc_f(elem_type, ds_vals_lo[r]) for r in range_constexpr(16)] + ds_f16_hi = [arith.trunc_f(elem_type, ds_vals_hi[r]) for r in range_constexpr(16)] + _pack_ty = v8f16_type if USE_K16 else v4f16_type + _pw = 8 if USE_K16 else 4 + ds_packs_lo = [] + ds_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + b = pks * _pw + ds_packs_lo.append(vector.from_elements(_pack_ty, [ds_f16_lo[b + i] for i in range(_pw)])) + ds_packs_hi.append(vector.from_elements(_pack_ty, [ds_f16_hi[b + i] for i in range(_pw)])) + + # ==== GEMM3: dQ += K @ dS (uses fwd's V-read schedule with K substituted) ==== + _steps = [(dc, pks) for dc in range(D_CHUNKS) for pks in range(PV_K_STEPS)] + TOTAL_PV = len(_steps) + v_base = v_buf_base(cur_buf) if _use_dbuf else v_buf_base(0) + + def _read_k_as_v_pack(step_idx): + dc, pks = _steps[step_idx] + if USE_HW_TR: + d_col = arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4 + k_row = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + tr_k_group + _d_col_eff = _v_swizzle(k_row, d_col) if ENABLE_DMA else d_col + lds_lo = v_base + k_row * V_STRIDE + _d_col_eff + lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE) + if USE_K16: + vl_a = ds_read_tr_v4f16(lds_lo) + vl_b = ds_read_tr_v4f16(lds_lo + arith.index(8 * V_STRIDE)) + vl = vector.shuffle(vl_a, vl_b, [0, 1, 2, 3, 4, 5, 6, 7]) + vh_a = ds_read_tr_v4f16(lds_hi) + vh_b = ds_read_tr_v4f16(lds_hi + arith.index(8 * V_STRIDE)) + vh = vector.shuffle(vh_a, vh_b, [0, 1, 2, 3, 4, 5, 6, 7]) + else: + vl = ds_read_tr_v4f16(lds_lo) + vh = ds_read_tr_v4f16(lds_hi) + else: + d_pos = arith.index(dc * D_CHUNK) + lane_mod_32 + kb = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + v_lo_idx = v_base + d_pos * VT_STRIDE + kb + v_hi_idx = v_lo_idx + arith.index(K_SUB_N) + vl = vector.load(v4f16_type, lds_kv, [v_lo_idx]) + vh = vector.load(v4f16_type, lds_kv, [v_hi_idx]) + return vl, vh + + k_lo_cur, k_hi_cur = _read_k_as_v_pack(0) + for si in range_constexpr(TOTAL_PV): + dc, pks = _steps[si] + if si + 1 < TOTAL_PV: + k_lo_nxt, k_hi_nxt = _read_k_as_v_pack(si + 1) + dq_accs[dc] = mfma_acc(k_lo_cur, ds_packs_lo[pks], dq_accs[dc]) + dq_accs[dc] = mfma_acc(k_hi_cur, ds_packs_hi[pks], dq_accs[dc]) + if si + 1 < TOTAL_PV: + k_lo_cur = k_lo_nxt + k_hi_cur = k_hi_nxt + + # End of iter: yield dq_accs and swapped cur_buf. + gpu.barrier() + _yield = list(dq_accs) + if _use_dbuf: + _yield.append(next_buf) + yield _yield + + # ---- Final store: dQ_fp32 = dq_acc * sm_scale (NO trunc; fp32) ---- + dq_finals = [loop_results[dc] for dc in range_constexpr(D_CHUNKS)] + sm_scale_vec = vector.broadcast(v16f32_type, c_sm_scale) + + _o_guard = scf.IfOp(q_in_bounds, [], has_else=False) + with ir.InsertionPoint(_o_guard.then_block): + for dc in range_constexpr(D_CHUNKS): + dq_scaled = arith.MulFOp( + dq_finals[dc], + sm_scale_vec, + fastmath=fm_fast, + ).result + for r in range_constexpr(16): + dq_val = vector.extract( + dq_scaled, + static_position=[r], + dynamic_position=[], + ) + d_row_rel = lane_div_32 * 4 + (r // 4) * 8 + (r % 4) + d_col = arith.index(dc * D_CHUNK) + d_row_rel + dq_global = global_idx_q(q_row, d_col) + _gep_store_f32(dq_val, dq_ptr, dq_global) + scf.YieldOp([]) + + @flyc.jit + def launch_v4_swa_bwd_dq( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + DOS: fx.Tensor, + LSE: fx.Tensor, + DELTAS: fx.Tensor, + DQ: fx.Tensor, + DSINK: fx.Tensor, + SINK: fx.Tensor, + batch_size: fx.Int32, + seq_len_q: fx.Int32, + seq_len_k: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + bs_idx = arith.index_cast(T.index, batch_size) + sl_idx = arith.index_cast(T.index, seq_len_q) + num_q_tiles = (sl_idx + BLOCK_M - 1) // BLOCK_M + grid_x = bs_idx * num_q_tiles * NUM_HEADS + + launcher = v4_swa_bwd_dq_kernel( + Q, + K, + V, + DOS, + LSE, + DELTAS, + DQ, + DSINK, + SINK, + seq_len_q, + seq_len_k, + ) + + if waves_per_eu is not None: + _wpe = int(waves_per_eu) + if _wpe >= 1: + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, + _wpe, + ) + if flat_work_group_size is not None: + _fwgs = int(flat_work_group_size) + if _fwgs >= 1: + flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}") + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr + + passthrough_entries = [] + if daz: + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("denormal-fp-math-f32"), + ir.StringAttr.get("preserve-sign,preserve-sign"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("no-nans-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("unsafe-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries) + + launcher.launch( + grid=(grid_x, 1, 1), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) + + _fmha_compile_hints = { + "fast_fp_math": fast_fp_math, + "unsafe_fp_math": unsafe_fp_math, + "llvm_options": { + "enable-post-misched": False, + "lsr-drop-solution": True, + }, + } + + def _launch(*args, **kwargs): + with CompilationContext.compile_hints(_fmha_compile_hints): + return launch_v4_swa_bwd_dq(*args, **kwargs) + + def _compile(Q, K, V, DOS, LSE, DELTAS, DQ, DSINK, SINK, batch_size, seq_len_q, seq_len_k, stream=None): + with CompilationContext.compile_hints(_fmha_compile_hints): + return flyc.compile( + launch_v4_swa_bwd_dq, + Q, + K, + V, + DOS, + LSE, + DELTAS, + DQ, + DSINK, + SINK, + batch_size, + seq_len_q, + seq_len_k, + fx.Stream(stream), + ) + + _launch.compile = _compile + + return _launch + + +# Convenience alias. +build_v4_swa_bwd_dq_module_primary = build_v4_swa_bwd_dq_module diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_kernel.py new file mode 100644 index 000000000..9efb746a5 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_bwd_kernel.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""v4_sla_bwd_kernel: V4 SWA attention backward kernels (FlyDSL, STEP 1). + +Forked from /workspace/FlyDSL-amd/kernels/sla_bwd_preprocess.py. + +STEP 1 SCOPE +============ +Only the preprocess kernel is implemented in this file. dq / dkv are still +handled by Triton in the v4_attention_bwd_flydsl_mqa wrapper (see that file +for the rationale). The hooks defined here will be the landing spots for +the dq / dkv kernels in STEP 1b / STEP 1c. + +PREPROCESS KERNEL (D scalar) +--------------------------- +Computes ``delta[b, h, m] = sum_d (out[b, h, m, d] * dout[b, h, m, d])`` in +fp32 for every query row. This is the standard FA-2 pre-pass and is dense +(no SWA / sink dependency) so it is the simplest piece to lift to FlyDSL +first. + +Inputs (flat views): + OS, DOS shape (N_ROWS, D) where N_ROWS = B*HQ*Sq + DELTAS shape (N_ROWS,) fp32 + +Grid: (N_ROWS / BLOCK_ROWS, 1, 1). +Block: BLOCK_ROWS * THREADS_PER_ROW threads. + +Each row uses THREADS_PER_ROW = D // VEC_WIDTH threads. Cross-thread row +reduction uses ``shuffle_xor`` at offsets THREADS_PER_ROW/2, /4, ..., 1 +which stays inside a THREADS_PER_ROW-lane subgroup (XOR never carries out +of a power-of-2 subgroup). + +Requires: head_dim % VEC_WIDTH == 0, N_ROWS % BLOCK_ROWS == 0. +""" + +import math +import os + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import arith, buffer_ops, range_constexpr +from flydsl.expr.arith import ArithValue +from flydsl.expr.numeric import Float32 +from flydsl.expr.typing import Int32, T +from flydsl.expr.vector import ReductionOp +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from kernels.kernels_common import dtype_to_elem_type + +KERNEL_NAME_PRE = "v4_swa_bwd_preprocess_kernel" + +# Wider VEC_WIDTH than the SLA kernel (D=512 vs D=128). At D=512 with +# VEC_WIDTH=8 we would need 64 threads per row (== a full warp), which is +# fine but uses one wave per BLOCK_ROWS rows. We use VEC_WIDTH=16 so +# THREADS_PER_ROW=32 and a 256-thread block processes 8 rows = 2 waves +# (more parallel tile generators per CU at the cost of slightly wider +# loads). Either is correct. +VEC_WIDTH = 8 +WARP_SIZE = 64 + + +def build_v4_swa_bwd_preprocess_module( + head_dim, + dtype_str="bf16", + block_rows=None, +): + """Build the V4 SWA backward preprocess launcher. + + Inputs (flattened to 2D by the launcher): + OS, DOS shape (N_ROWS, D) where N_ROWS = B*H*L + DELTAS shape (N_ROWS,) f32 + + Args: + head_dim: head dimension (must be divisible by VEC_WIDTH). + dtype_str: "bf16" or "f16" for O_S / DO_S; DELTAS is always f32. + block_rows: rows processed per block. Default 8. + """ + get_hip_arch() + + if block_rows is None: + BLOCK_ROWS = 8 + else: + BLOCK_ROWS = block_rows + + assert head_dim % VEC_WIDTH == 0, f"head_dim {head_dim} must be % {VEC_WIDTH}" + assert head_dim >= VEC_WIDTH, f"head_dim {head_dim} < VEC_WIDTH {VEC_WIDTH}" + assert dtype_str in ("bf16", "f16"), f"unsupported dtype {dtype_str}" + + D = head_dim + THREADS_PER_ROW = D // VEC_WIDTH # 32 for D=512, VEC_WIDTH=16 + BLOCK_THREADS = BLOCK_ROWS * THREADS_PER_ROW + assert BLOCK_THREADS <= 1024, f"BLOCK_THREADS {BLOCK_THREADS} > 1024" + assert ( + THREADS_PER_ROW & (THREADS_PER_ROW - 1) == 0 + ), "THREADS_PER_ROW must be power of 2 for shfl_xor reduction" + + elem_bits = 16 # bf16 / f16 + + @flyc.kernel + def v4_swa_bwd_preprocess_kernel( + OS: fx.Tensor, # (N_ROWS, D) + DOS: fx.Tensor, # (N_ROWS, D) + DELTAS: fx.Tensor, # (N_ROWS,) + ): + bid = fx.block_idx.x + tid = fx.thread_idx.x + + elem_type = dtype_to_elem_type(dtype_str) + T.f32 + fm_fast = arith.FastMathFlags.fast + + # ---- Thread decomposition ---- + row_in_block_i32 = tid // Int32(THREADS_PER_ROW) + col_in_row_i32 = tid % Int32(THREADS_PER_ROW) + + global_row_i32 = bid * Int32(BLOCK_ROWS) + row_in_block_i32 + global_row_idx = ArithValue(global_row_i32).index_cast(T.index) + + # ---- Buffer-backed 2D tensors for loads ---- + OS_buf = fx.rocdl.make_buffer_tensor(OS) + DOS_buf = fx.rocdl.make_buffer_tensor(DOS) + delta_rsrc = buffer_ops.create_buffer_resource(DELTAS, max_size=True) + + row_os = fx.slice(OS_buf, (global_row_i32, None)) + row_dos = fx.slice(DOS_buf, (global_row_i32, None)) + + in_div_os = fx.logical_divide(row_os, fx.make_layout(VEC_WIDTH, 1)) + in_div_dos = fx.logical_divide(row_dos, fx.make_layout(VEC_WIDTH, 1)) + + copy_atom = fx.make_copy_atom(fx.rocdl.BufferCopy128b(), elem_bits) + vec_reg_ty = fx.MemRefType.get(elem_type, fx.LayoutType.get(VEC_WIDTH, 1), fx.AddressSpace.Register) + vec_reg_lay = fx.make_layout(VEC_WIDTH, 1) + + def _load_vec(div_tensor, col_idx): + r = fx.memref_alloca(vec_reg_ty, vec_reg_lay) + fx.copy_atom_call(copy_atom, fx.slice(div_tensor, (None, col_idx)), r) + return fx.memref_load_vec(r) + + os_vec = _load_vec(in_div_os, col_in_row_i32) + dos_vec = _load_vec(in_div_dos, col_in_row_i32) + + os_f32 = os_vec.to(Float32) + dos_f32 = dos_vec.to(Float32) + prod_f32 = os_f32 * dos_f32 + local_sum = prod_f32.reduce(ReductionOp.ADD, fastmath=fm_fast) + + width_i32 = Int32(WARP_SIZE) + val = local_sum + num_rounds = int(math.log2(THREADS_PER_ROW)) + for _sh_exp in range_constexpr(num_rounds): + off = Int32(THREADS_PER_ROW // (2 << _sh_exp)) + peer = val.shuffle_xor(off, width_i32) + val = val.addf(peer, fastmath=fm_fast) + + if col_in_row_i32 == Int32(0): + delta_off_i32 = arith.index_cast(T.i32, global_row_idx) + val_ir = val.ir_value() if hasattr(val, "ir_value") else val + buffer_ops.buffer_store(val_ir, delta_rsrc, delta_off_i32) + + @flyc.jit + def launch_v4_swa_bwd_preprocess( + OS: fx.Tensor, + DOS: fx.Tensor, + DELTAS: fx.Tensor, + n_rows: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + """Launch grid: (n_rows / BLOCK_ROWS, 1, 1). + + Expects OS, DOS to be views of shape (n_rows, D). DELTAS is (n_rows,). + """ + blocks_i32 = n_rows // Int32(BLOCK_ROWS) + blocks_idx = ArithValue(blocks_i32).index_cast(T.index) + launcher = v4_swa_bwd_preprocess_kernel(OS, DOS, DELTAS) + launcher.launch( + grid=(blocks_idx, arith.index(1), arith.index(1)), + block=(BLOCK_THREADS, 1, 1), + stream=stream, + ) + + return launch_v4_swa_bwd_preprocess + + +# Convenience alias used by the wrapper. +def build_v4_swa_bwd_preprocess(*args, **kwargs): + return build_v4_swa_bwd_preprocess_module(*args, **kwargs) + + +# --------------------------------------------------------------------------- +# STEP 1b: dq kernel re-export (forked from kernels/sla_bwd_dq.py). +# Lives in v4_sla_bwd_dq_kernel.py for file-size reasons; re-exported here +# so the wrapper can import a single module. +# --------------------------------------------------------------------------- +import os as _os +import sys as _sys + +_HERE = _os.path.dirname(_os.path.abspath(__file__)) +if _HERE not in _sys.path: + _sys.path.insert(0, _HERE) +from v4_sla_bwd_dq_kernel import ( # noqa: E402,F401 + build_v4_swa_bwd_dq_module, + build_v4_swa_bwd_dq_module_primary, +) diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_fwd_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_fwd_kernel.py new file mode 100644 index 000000000..5e331ec77 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v0_deprecated/kernels/v4_sla_fwd_kernel.py @@ -0,0 +1,1387 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""v4_swa_fwd: V4 dense sliding-window-causal attention forward (FlyDSL). + +Forked from kernels/sla_fwd.py. Differences: + - No LUT: outer KV loop iterates a contiguous block range bounded by + the SWA window for the current Q tile (depends on q_tile_idx, but + wave-uniform). + - Per-element SWA mask: kv_col > q_row OR kv_col + swa_window <= q_row + OR kv_col >= seq_len -> NEG_INF, applied before softmax max-reduce. + - LSE is fp32 raw-domain `lse = m_final*scale + ln(l_final)` to match + V4 Triton reference (m_i + tl.log(l_i)). + - A-1 scope: no sink, no additive mask, no HCA. MQA broadcast at the + Python launcher (kernel sees full MHA K/V). + +Layout: BHLD. Q/K/V/O flattened from (B, H, L, D). +LSE: (B, H, L) flat, fp32. +Grid: (B * num_q_tiles * H,), num_q_tiles = L / BLOCK_M. +""" + +import math +import os + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import math as math_dialect +from flydsl._mlir.dialects import memref as _memref +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr +from kernels.kernels_common import dtype_to_elem_type + +# ---- Module-level constants ---- + +KERNEL_NAME = "v4_swa_fwd_kernel" + +_LOG2E = math.log2(math.e) # 1.4426950408889634 + +_LLVM_GEP_DYNAMIC = -2147483648 # LLVM kDynamicIndex sentinel (0x80000000 as signed i32) + + +def _llvm_ptr_ty(): + return ir.Type.parse("!llvm.ptr") + + +def _llvm_lds_ptr_ty(): + return ir.Type.parse("!llvm.ptr<3>") + + +_VMCNT_LO_MASK = 0xF +_LGKMCNT_EXPCNT_BASE = 0x3F70 +_VMCNT_HI_SHIFT = 14 +_VMCNT_HI_MASK = 0x3 + + +def _waitcnt_vm_n(n): + """Emit s_waitcnt vmcnt(n) only (lgkmcnt=63, expcnt=7).""" + val = (n & _VMCNT_LO_MASK) | _LGKMCNT_EXPCNT_BASE | (((n >> 4) & _VMCNT_HI_MASK) << _VMCNT_HI_SHIFT) + rocdl.s_waitcnt(val) + + +def build_v4_swa_fwd_module( + num_heads, + head_dim, + swa_window, + dtype_str="bf16", + sm_scale=None, + waves_per_eu=2, + flat_work_group_size=None, + block_m=None, + block_n=None, + unsafe_fp_math=True, + fast_fp_math=True, + daz=True, + layout_bhld=True, + mqa_kv=False, +): + """Build the SLA sparse-fwd launcher. Single-variant only (no auto-dispatch). + + mqa_kv: if True, K and V are MQA-shape [B, 1, Sk, D] (stride_kh=0) and + indexing into them drops head_idx. The kernel still sees Q at [B, H, Sq, D]. + + layout_bhld: + False → Q/K/V/O in (B, L, H, D) layout (inherited from flash_attn_func, + matches torch `(batch, seq, num_heads, head_dim)` natural form). + True → Q/K/V/O in (B, H, L, D) layout (SLA's native form; avoids a + transpose at the integration boundary in `SparseLinearAttention`). + LUT and LSE layouts are unchanged either way: + LUT (B, H, M_BLOCKS, topk) flattened, dtype i32. + LSE (B, H, L) flattened, dtype f32. + """ + gpu_arch = get_hip_arch() + + if block_n is None: + BLOCK_N = 64 + else: + BLOCK_N = int(block_n) + K_SUB_N = min(BLOCK_N, 32) + N_HALVES = BLOCK_N // 32 # 1 or 2; how many K_SUB_N halves per BLOCK_N + WARP_SIZE = 64 + # SLA is always non-causal; the sparse map decides which blocks attend. + + if block_m is not None: + BLOCK_M = block_m + else: + BLOCK_M = 128 + + if flat_work_group_size is None: + if BLOCK_M <= 128: + flat_work_group_size = 256 + else: + flat_work_group_size = 512 + NUM_WAVES = flat_work_group_size // WARP_SIZE + BLOCK_SIZE = flat_work_group_size + ROWS_PER_WAVE = BLOCK_M // NUM_WAVES + # V4 SWA: always N32 path. One outer iter = one BLOCK_N block. + BLOCK_N_OUT = BLOCK_N + N_SUBTILES = 1 + ENABLE_PREFETCH_3BUF = os.getenv("FLYDSL_SLA_FWD_ENABLE_PREFETCH3", "0") == "1" + # buffer_load_dwordx4_lds (16B DMA-to-LDS) requires gfx950+; gfx94x only has dword (4B). + # For SLA we default-on DMA when hardware supports it — this is the whole point. + _has_lds_load_b128 = not gpu_arch.startswith("gfx942") + ENABLE_DMA = _has_lds_load_b128 and (os.getenv("FLYDSL_SLA_FWD_ENABLE_DMA", "1") == "1") + ENABLE_LDS_VEC16 = os.getenv("FLYDSL_SLA_FWD_ENABLE_LDS_VEC16", "1") == "1" + REDUCE_MODE = os.getenv("FLYDSL_SLA_FWD_REDUCE_MODE", "xor").strip().lower() + if REDUCE_MODE not in ("xor", "ds_bpermute"): + REDUCE_MODE = "xor" + FORCE_SINGLE_BUF_DMA = os.getenv("FLYDSL_V4_SWA_SINGLE_BUF_DMA", "0") == "1" + if ENABLE_PREFETCH_3BUF: + NUM_PREFETCH_K = 3 + elif ENABLE_DMA and not FORCE_SINGLE_BUF_DMA: + NUM_PREFETCH_K = 2 + else: + NUM_PREFETCH_K = 1 + # Lever A: double-buffer V in DMA-dbuf mode so V HBM latency overlaps + # with the QK MFMA of the next iteration (mirrors K-dbuf behavior). + if ENABLE_PREFETCH_3BUF: + NUM_PREFETCH_V = 3 + elif ENABLE_DMA and not FORCE_SINGLE_BUF_DMA: + NUM_PREFETCH_V = 2 + else: + NUM_PREFETCH_V = 1 + CK_LDS_SEQ = (1, 2, 0, 1, 0, 1, 2, 0) if ENABLE_PREFETCH_3BUF else (0,) + + # gfx950+ has ds_read_tr16_b64 (HW transpose LDS read); gfx942 needs V^T stored in LDS. + USE_HW_TR = gpu_arch.startswith("gfx950") + + # MFMA32 K-dimension: 16 on gfx950+ (CDNA4) for both GEMMs. + USE_K16 = gpu_arch.startswith("gfx950") + K_STEP_QK = 16 if USE_K16 else 8 + K_STEPS_QK = head_dim // K_STEP_QK + D_CHUNK = 32 + D_CHUNKS = head_dim // D_CHUNK + PV_K_STEP = 16 if USE_K16 else 8 + PV_K_STEPS = K_SUB_N // PV_K_STEP # 2 steps per sub-tile (K=16) or 4 (K=8) + + assert BLOCK_M % NUM_WAVES == 0 + assert head_dim % 32 == 0, f"head_dim ({head_dim}) must be divisible by 32" + assert head_dim >= 64, f"head_dim ({head_dim}) must be >= 64" + assert flat_work_group_size in ( + 128, + 256, + 512, + ), f"flat_work_group_size must be 128, 256, or 512, got {flat_work_group_size}" + assert dtype_str in ("f16", "bf16"), "sla_fwd only supports f16 and bf16" + assert BLOCK_N % 32 == 0 or BLOCK_N == 32 + assert BLOCK_N_OUT == BLOCK_N + assert N_HALVES in (1, 2) + assert isinstance(swa_window, int) and swa_window > 0, f"swa_window must be int > 0, got {swa_window!r}" + + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(head_dim) + + NUM_HEADS = num_heads + HEAD_DIM = head_dim + STRIDE_TOKEN = NUM_HEADS * HEAD_DIM + + # Bank-conflict-free LDS strides. + # K uses XOR swizzle (col ^ ((row & 7) << 4)) at 16-element granularity + # instead of padding. This enables ds_read_b128 (stride is 256B-aligned). + K_STRIDE = HEAD_DIM + if USE_HW_TR: + V_STRIDE = HEAD_DIM if ENABLE_DMA else HEAD_DIM + 4 + else: + VT_STRIDE = BLOCK_N + 2 + V_STRIDE = VT_STRIDE + + # Vectorized cooperative load constants. + VEC_WIDTH = 16 if ENABLE_LDS_VEC16 else 8 + assert HEAD_DIM % VEC_WIDTH == 0 + THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH + assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0 + ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD + + if ROWS_PER_BATCH_LOAD >= BLOCK_N: + NUM_BATCHES_KV = 1 + KV_NEEDS_GUARD = ROWS_PER_BATCH_LOAD > BLOCK_N + else: + assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0 + NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD + KV_NEEDS_GUARD = False + + # K/V circular buffers; defaults to 1/1, optional 3/3 with CK-like LDS sequence. + LDS_K_TILE_SIZE = BLOCK_N * K_STRIDE + if USE_HW_TR: + LDS_V_TILE_SIZE = BLOCK_N * V_STRIDE + else: + LDS_V_TILE_SIZE = HEAD_DIM * VT_STRIDE + LDS_K_TOTAL_SIZE = NUM_PREFETCH_K * LDS_K_TILE_SIZE + LDS_V_BASE = LDS_K_TOTAL_SIZE + LDS_V_TOTAL_SIZE = NUM_PREFETCH_V * LDS_V_TILE_SIZE + LDS_KV_TOTAL_SIZE = LDS_K_TOTAL_SIZE + LDS_V_TOTAL_SIZE + + allocator = SmemAllocator( + None, + arch=gpu_arch, + global_sym_name=f"v4_swa_fwd_smem_M{BLOCK_M}_N{BLOCK_N}_W{swa_window}", + ) + lds_kv_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_kv_offset + LDS_KV_TOTAL_SIZE * 2 + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def v4_swa_fwd_kernel( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + O: fx.Tensor, + LSE: fx.Tensor, + seq_len: fx.Int32, + ): + elem_type = dtype_to_elem_type(dtype_str) + compute_type = T.f32 + q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q) + k_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), K) + v_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), V) + o_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), O) + # LSE: f32 scalar writes via buffer_store. + lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True) + + # All FP operations use aggressive fast-math (no NaN/Inf checks, reassociation). + # The unsafe_fp_math/fast_fp_math builder params control LLVM-level attributes only. + fm_fast = arith.FastMathFlags.fast + v4f16_type = T.vec(4, elem_type) + vxf16_type = T.vec(VEC_WIDTH, elem_type) + v8f16_type = T.vec(8, elem_type) + v16f32_type = T.vec(16, compute_type) + mfma_pack_type = v8f16_type if USE_K16 else v4f16_type + MFMA_LANE_K = 8 if USE_K16 else 4 + _mfma_zero = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 0) + + def _mfma(ods_fn, a, b, c): + return ods_fn(v16f32_type, a, b, c, _mfma_zero, _mfma_zero, _mfma_zero).result + + def mfma_acc(a, b, c): + if dtype_str == "bf16": + if USE_K16: + return _mfma(rocdl.mfma_f32_32x32x16_bf16, a, b, c) + a = vector.bitcast(T.i16x4, a) + b = vector.bitcast(T.i16x4, b) + return _mfma(rocdl.mfma_f32_32x32x8bf16_1k, a, b, c) + if USE_K16: + return _mfma(rocdl.mfma_f32_32x32x16_f16, a, b, c) + return _mfma(rocdl.mfma_f32_32x32x8f16, a, b, c) + + seq_len_v = arith.index_cast(T.index, seq_len) + + # ---- LDS view ---- + base_ptr = allocator.get_base() + lds_kv = SmemPtr( + base_ptr, + lds_kv_offset, + elem_type, + shape=(LDS_KV_TOTAL_SIZE,), + ).get() + + # ---- Thread / block indices ---- + block_id = arith.index_cast(T.index, gpu.block_idx.x) + tid = arith.index_cast(T.index, gpu.thread_idx.x) + + # ---- Wave decomposition ---- + wave_id = tid // WARP_SIZE + lane = tid % WARP_SIZE + lane_mod_32 = lane % 32 + lane_div_32 = lane // 32 # 0/1 + + # ---- ds_read_b64_tr_b16 lane decomposition ---- + # Hardware does 4×4 transpose within blocks of 16 lanes. + # tr_k_group selects which of 4 K-rows within the block, + # tr_col_sub selects which 4-column sub-group within 16 columns. + tr_k_group = (lane % 16) // 4 # 0..3: K-row offset within 4-row group + tr_col_sub = lane % 4 # 0..3: 4-column sub-group + tr_col_half = (lane % 32) // 16 # 0 or 1: first/second 16-column half + + # ---- ds_read_b64_tr_b16 helper ---- + + def ds_read_tr_v4f16(lds_elem_idx): + """Read v4f16 from LDS with hardware transpose. + + Within each block of 16 lanes, the hardware performs a 4×4 + transpose across 4 groups of 4 lanes. After the transpose, + result[lane, elem_e] = Input[source_lane, lane%4] where + source_lane = e*4 + (lane%16)//4. This naturally produces + the MFMA A-operand layout when per-lane addresses point to + the correct K-row and D-column sub-group. + """ + byte_offset = lds_elem_idx * 2 + lds_kv_offset + byte_i64 = arith.index_cast(T.i64, byte_offset) + ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result + return rocdl.ds_read_tr16_b64(v4f16_type, ptr).result + + # ---- Wave offsets ---- + wave_q_offset = wave_id * ROWS_PER_WAVE + + # ---- Decompose block_id ---- + head_idx = block_id % NUM_HEADS + batch_q_tile_id = block_id // NUM_HEADS + num_q_tiles = (seq_len_v + BLOCK_M - 1) // BLOCK_M + q_tile_idx = batch_q_tile_id % num_q_tiles + batch_idx = batch_q_tile_id // num_q_tiles + q_start = q_tile_idx * BLOCK_M + + # ---- V4 SWA: per-tile contiguous K-block range (wave-uniform) ---- + # Q tile rows: [q_start, q_start + BLOCK_M). Union of visible K + # columns under SWA-causal: [q_start - W + 1, q_start + BLOCK_M). + # n_block_start = max(0, q_start - W + 1) // BLOCK_N + # n_block_end = ceil(min(q_start + BLOCK_M, seq_len), BLOCK_N) + # Per-element causal/SWA/boundary mask still applied inside the loop. + SWA = arith.index(swa_window) + BN = arith.index(BLOCK_N) + BM = arith.index(BLOCK_M) + _zero_idx = arith.index(0) + _one_idx = arith.index(1) + # q_start - W + 1 -- compute as index arithmetic, then clamp. + _q_plus_one = q_start + _one_idx + # We need max(0, _q_plus_one - SWA). To avoid negative-index issues + # we test _q_plus_one >= SWA first. + _ge_w = arith.cmpi(arith.CmpIPredicate.sge, _q_plus_one, SWA) + _n_start_row = arith.select(_ge_w, _q_plus_one - SWA, _zero_idx) + n_block_start = _n_start_row // BN + _n_end_row_uncl = q_start + BM + _le_seq = arith.cmpi(arith.CmpIPredicate.sle, _n_end_row_uncl, seq_len_v) + n_end_row_cl = arith.select(_le_seq, _n_end_row_uncl, seq_len_v) + n_block_end = (n_end_row_cl + BN - _one_idx) // BN + + # ---- Cooperative load decomposition ---- + load_row_in_batch = tid // THREADS_PER_ROW_LOAD + load_lane_in_row = tid % THREADS_PER_ROW_LOAD + load_col_base = load_lane_in_row * VEC_WIDTH + + # ---- Helper: global flat index ---- + # BLHD: stride = (L*H*D, H*D, D, 1) → (token=b*L + l) * (H*D) + h*D + col + # BHLD: stride = (H*L*D, L*D, D, 1) → ((b*H + h) * L + l) * D + col + if layout_bhld: + bh_base_tokens = (batch_idx * NUM_HEADS + head_idx) * seq_len_v + if mqa_kv: + # K/V are [B, 1, Sk, D] -- drop head_idx from KV indexing. + bh_base_tokens_kv = batch_idx * seq_len_v + else: + bh_base_tokens_kv = bh_base_tokens + + def global_idx(token_idx, col): + return (bh_base_tokens + token_idx) * arith.index(HEAD_DIM) + col + + def global_idx_kv(token_idx, col): + return (bh_base_tokens_kv + token_idx) * arith.index(HEAD_DIM) + col + + else: + + def global_idx(token_idx, col): + token = batch_idx * seq_len_v + token_idx + return token * STRIDE_TOKEN + head_idx * HEAD_DIM + col + + if mqa_kv: + # BLHD MQA: K/V token = batch_idx * seq_len_v + token_idx (no h*D added). + def global_idx_kv(token_idx, col): + token = batch_idx * seq_len_v + token_idx + return token * HEAD_DIM + col # no STRIDE_TOKEN, no head_idx + + else: + global_idx_kv = global_idx + + def _gep_load(base_ptr, elem_idx, vec_type): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=elem_type, + noWrapFlags=0, + ) + return _llvm.LoadOp(vec_type, gep.result).result + + def _gep_store(val, base_ptr, elem_idx): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + base_ptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=elem_type, + noWrapFlags=0, + ) + _llvm.StoreOp(val, gep.result) + + def load_global_f16x4(base_ptr, base_idx): + return _gep_load(base_ptr, base_idx, v4f16_type) + + def load_global_mfma_pack(base_ptr, base_idx): + return _gep_load(base_ptr, base_idx, mfma_pack_type) + + def load_global_f16xN(base_ptr, base_idx): + return _gep_load(base_ptr, base_idx, vxf16_type) + + def bf16_trunc_pack_v4(f32_vals): + """Pack 4 f32 values into v4bf16 via bitwise truncation (upper 16 bits). + ~2 fewer instructions/element vs arith.TruncFOp round-to-nearest.""" + _v2i32 = T.vec(2, T.i32) + _c16 = arith.constant(16, type=T.i32) + _cmask = arith.constant(0xFFFF0000, type=T.i32) + a0 = arith.ArithValue(f32_vals[0]).bitcast(T.i32) + b0 = arith.ArithValue(f32_vals[1]).bitcast(T.i32) + p0 = arith.OrIOp(arith.AndIOp(b0, _cmask).result, arith.ShRUIOp(a0, _c16).result).result + a1 = arith.ArithValue(f32_vals[2]).bitcast(T.i32) + b1 = arith.ArithValue(f32_vals[3]).bitcast(T.i32) + p1 = arith.OrIOp(arith.AndIOp(b1, _cmask).result, arith.ShRUIOp(a1, _c16).result).result + return vector.bitcast(v4f16_type, vector.from_elements(_v2i32, [p0, p1])) + + def bf16_trunc_pack_v8(f32_vals): + """Pack 8 f32 values into v8bf16 via bitwise truncation (upper 16 bits).""" + _v4i32 = T.vec(4, T.i32) + _c16 = arith.constant(16, type=T.i32) + _cmask = arith.constant(0xFFFF0000, type=T.i32) + pairs = [] + for j in range_constexpr(4): + a = arith.ArithValue(f32_vals[j * 2]).bitcast(T.i32) + b = arith.ArithValue(f32_vals[j * 2 + 1]).bitcast(T.i32) + p = arith.OrIOp(arith.AndIOp(b, _cmask).result, arith.ShRUIOp(a, _c16).result).result + pairs.append(p) + return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs)) + + def k_buf_base(buf_id): + if isinstance(buf_id, int): + return arith.index(buf_id * LDS_K_TILE_SIZE) + return buf_id * arith.index(LDS_K_TILE_SIZE) + + def v_buf_base(buf_id): + if isinstance(buf_id, int): + return arith.index(LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) + return arith.index(LDS_V_BASE) + buf_id * arith.index(LDS_V_TILE_SIZE) + + # ---- K XOR swizzle: col ^ ((row & 7) << 4) at 16-element granularity ---- + def _k_swizzle(row_idx, col_idx): + mask = (row_idx & arith.index(0x7)) << arith.index(4) + return col_idx ^ mask + + # ---- Cooperative K load (row-major, XOR-swizzled) ---- + def coop_load_k(tile_start, buf_id=0): + k_base = k_buf_base(buf_id) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + _if_k = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_k.then_block): + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + vec = load_global_f16xN(k_ptr, g_idx) + vector.store(vec, lds_kv, [lds_idx]) + scf.YieldOp([]) + else: + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + swz_col = _k_swizzle(lds_row, load_col_base) + lds_idx = k_base + lds_row * K_STRIDE + swz_col + vec = load_global_f16xN(k_ptr, g_idx) + vector.store(vec, lds_kv, [lds_idx]) + + # ---- Cooperative V load ---- + def _v_store_row_major(v_base, lds_row, vec): + lds_idx = v_base + lds_row * V_STRIDE + load_col_base + vector.store(vec, lds_kv, [lds_idx]) + + _v1_type = T.vec(1, elem_type) if not USE_HW_TR else None + + def _v_store_transposed(v_base, lds_row, vec): + for _e in range_constexpr(VEC_WIDTH): + elem = vector.extract(vec, static_position=[_e], dynamic_position=[]) + vt_d = load_col_base + _e + vt_idx = v_base + vt_d * VT_STRIDE + lds_row + v1 = vector.from_elements(_v1_type, [elem]) + vector.store(v1, lds_kv, [vt_idx]) + + _v_store_to_lds = _v_store_row_major if USE_HW_TR else _v_store_transposed + + def coop_load_v(tile_start, buf_id=0): + v_base = v_buf_base(buf_id) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + _if_v = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_v.then_block): + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + vec = load_global_f16xN(v_ptr, g_idx) + _v_store_to_lds(v_base, lds_row, vec) + scf.YieldOp([]) + else: + g_idx = global_idx_kv(row_idx, load_col_base) + lds_row = load_row_in_batch + row_offset + vec = load_global_f16xN(v_ptr, g_idx) + _v_store_to_lds(v_base, lds_row, vec) + + def coop_load_v_global(tile_start): + """Issue global loads for V, return vectors (non-blocking).""" + vecs = [] + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + row_idx = tile_start + load_row_in_batch + row_offset + g_idx = global_idx_kv(row_idx, load_col_base) + vecs.append(load_global_f16xN(v_ptr, g_idx)) + return vecs + + def coop_store_v_lds(vecs, buf_id=0): + """Write previously-loaded V vectors to LDS.""" + v_base = v_buf_base(buf_id) + for batch in range_constexpr(NUM_BATCHES_KV): + row_offset = batch * ROWS_PER_BATCH_LOAD + if KV_NEEDS_GUARD: + row_valid = arith.cmpi( + arith.CmpIPredicate.ult, + load_row_in_batch, + arith.index(BLOCK_N), + ) + _if_v = scf.IfOp(row_valid) + with ir.InsertionPoint(_if_v.then_block): + lds_row = load_row_in_batch + row_offset + _v_store_to_lds(v_base, lds_row, vecs[batch]) + scf.YieldOp([]) + else: + lds_row = load_row_in_batch + row_offset + _v_store_to_lds(v_base, lds_row, vecs[batch]) + + # ---- DMA loading for K (buffer_load_dwordx4 ... lds) ---- + if ENABLE_DMA: + from flydsl._mlir.dialects import llvm + + k_rsrc = buffer_ops.create_buffer_resource(K, max_size=True) + _lds_ptr_ty = _llvm_lds_ptr_ty() + DMA_BYTES = 16 # buffer_load_dwordx4 = 16 bytes per lane + DMA_BATCH_BYTES = BLOCK_SIZE * DMA_BYTES + K_TILE_BYTES = BLOCK_N * K_STRIDE * 2 + NUM_DMA_K = K_TILE_BYTES // DMA_BATCH_BYTES + LANES_PER_K_ROW = HEAD_DIM * 2 // DMA_BYTES + ROWS_PER_DMA_BATCH = DMA_BATCH_BYTES // (HEAD_DIM * 2) + lds_kv_base_idx = _memref.extract_aligned_pointer_as_index(lds_kv) + _dma_size = arith.constant(DMA_BYTES, type=T.i32) + _dma_soff = arith.constant(0, type=T.i32) + _dma_off = arith.constant(0, type=T.i32) + _dma_aux = arith.constant(1, type=T.i32) + + def coop_dma_k(tile_start, buf_id=0): + """Load K tile via DMA with XOR-swizzled global fetch.""" + if isinstance(buf_id, int): + k_lds_byte_base = lds_kv_base_idx + arith.index(buf_id * LDS_K_TILE_SIZE * 2) + else: + k_lds_byte_base = lds_kv_base_idx + buf_id * arith.index(LDS_K_TILE_SIZE * 2) + for d in range_constexpr(NUM_DMA_K): + lds_addr = ( + k_lds_byte_base + + wave_id * arith.index(WARP_SIZE * DMA_BYTES) + + arith.index(d * DMA_BATCH_BYTES) + ) + lds_i64 = arith.index_cast(T.i64, lds_addr) + lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64) + lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result + + row_in_tile = tid // LANES_PER_K_ROW + arith.index(d * ROWS_PER_DMA_BATCH) + swiz_col_f16 = (tid % LANES_PER_K_ROW) * (DMA_BYTES // 2) + xor_mask = (row_in_tile & arith.index(0x7)) << arith.index(4) + unsw_col_f16 = swiz_col_f16 ^ xor_mask + col_byte = unsw_col_f16 * 2 + # Layout-aware global byte offset. + if layout_bhld: + # BHLD: ((b*H + h)*L + (tile_start + row)) * (D*2) + col_byte + # MQA: ((b)*L + (tile_start + row)) * (D*2) + col_byte + row_within_head = tile_start + row_in_tile + global_byte = (bh_base_tokens_kv + row_within_head) * arith.index( + HEAD_DIM * 2 + ) + col_byte + else: + # BLHD: ((b*L + row) * (H*D) + h*D) * 2 + col_byte + global_row = batch_idx * seq_len_v + tile_start + row_in_tile + if mqa_kv: + global_byte = global_row * arith.index(HEAD_DIM * 2) + col_byte + else: + global_byte = ( + global_row * arith.index(STRIDE_TOKEN * 2) + + head_idx * arith.index(HEAD_DIM * 2) + + col_byte + ) + voffset = arith.index_cast(T.i32, global_byte) + + rocdl.raw_ptr_buffer_load_lds( + k_rsrc, + lds_ptr, + _dma_size, + voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + # ---- V XOR swizzle: col ^ ((row & 3) << 4) at 16-element granularity ---- + def _v_swizzle(row_idx, col_idx): + mask = (row_idx & arith.index(0x3)) << arith.index(4) + return col_idx ^ mask + + # ---- DMA loading for V (buffer_load_dwordx4 ... lds) ---- + if ENABLE_DMA: + v_rsrc = buffer_ops.create_buffer_resource(V, max_size=True) + V_TILE_BYTES = BLOCK_N * V_STRIDE * 2 + NUM_DMA_V = V_TILE_BYTES // DMA_BATCH_BYTES + LANES_PER_V_ROW = HEAD_DIM * 2 // DMA_BYTES + ROWS_PER_DMA_BATCH_V = DMA_BATCH_BYTES // (HEAD_DIM * 2) + + def coop_dma_v(tile_start, buf_id=0): + """Load V tile via DMA with XOR-swizzled global fetch.""" + if isinstance(buf_id, int): + v_lds_byte_base = lds_kv_base_idx + arith.index( + (LDS_V_BASE + buf_id * LDS_V_TILE_SIZE) * 2 + ) + else: + v_lds_byte_base = ( + lds_kv_base_idx + + arith.index(LDS_V_BASE * 2) + + buf_id * arith.index(LDS_V_TILE_SIZE * 2) + ) + for d in range_constexpr(NUM_DMA_V): + lds_addr = ( + v_lds_byte_base + + wave_id * arith.index(WARP_SIZE * DMA_BYTES) + + arith.index(d * DMA_BATCH_BYTES) + ) + lds_i64 = arith.index_cast(T.i64, lds_addr) + lds_lane0 = rocdl.readfirstlane(T.i64, lds_i64) + lds_ptr = llvm.IntToPtrOp(_lds_ptr_ty, lds_lane0).result + + row_in_tile = tid // LANES_PER_V_ROW + arith.index(d * ROWS_PER_DMA_BATCH_V) + swiz_col_f16 = (tid % LANES_PER_V_ROW) * (DMA_BYTES // 2) + xor_mask = (row_in_tile & arith.index(0x3)) << arith.index(4) + unsw_col_f16 = swiz_col_f16 ^ xor_mask + col_byte = unsw_col_f16 * 2 + # Layout-aware global byte offset (see K DMA path above). + if layout_bhld: + row_within_head = tile_start + row_in_tile + global_byte = (bh_base_tokens_kv + row_within_head) * arith.index( + HEAD_DIM * 2 + ) + col_byte + else: + global_row = batch_idx * seq_len_v + tile_start + row_in_tile + if mqa_kv: + global_byte = global_row * arith.index(HEAD_DIM * 2) + col_byte + else: + global_byte = ( + global_row * arith.index(STRIDE_TOKEN * 2) + + head_idx * arith.index(HEAD_DIM * 2) + + col_byte + ) + voffset = arith.index_cast(T.i32, global_byte) + + rocdl.raw_ptr_buffer_load_lds( + v_rsrc, + lds_ptr, + _dma_size, + voffset, + _dma_soff, + _dma_off, + _dma_aux, + ) + + # ---- Preload Q^T B-operand packs once (register-resident) ---- + # B operand uses j = lane_mod_32, k-subblock = lane_div_32*MFMA_LANE_K. + q_row = q_start + wave_q_offset + lane_mod_32 + arith.index_cast(T.i32, q_row) + q_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, q_row, seq_len_v) + q_row_safe = arith.select(q_in_bounds, q_row, arith.index(0)) + c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type) + q_b_packs = [] + for ks in range_constexpr(K_STEPS_QK): + q_col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + g_idx = global_idx(q_row_safe, q_col) + raw = load_global_mfma_pack(q_ptr, g_idx) + q_b_packs.append(arith.select(q_in_bounds, raw, c_zero_mfma_pack)) + + # ---- Constants ---- + c_neg_inf = arith.constant( + -1.0e30, type=compute_type + ) # finite -inf to keep diff_m_raw=0 in all-masked tiles + c_zero_f = arith.constant(0.0, type=compute_type) + c_one_f = arith.constant(1.0, type=compute_type) + c_sm_scale_log2e = arith.constant(sm_scale * _LOG2E, type=compute_type) + c_zero_v16f32 = arith.constant_vector(0.0, v16f32_type) + width_i32 = arith.constant(WARP_SIZE, type=T.i32) + shuf_32_i32 = arith.constant(32, type=T.i32) + c4_i32 = arith.constant(4, type=T.i32) + lane_i32 = arith.index_cast(T.i32, lane) + lane_xor_32_i32 = arith.XOrIOp(lane_i32, shuf_32_i32).result + lane_xor_32_byte = arith.MulIOp(lane_xor_32_i32, c4_i32).result + + def reduction_peer(v_f32): + if REDUCE_MODE == "ds_bpermute": + v_i32 = arith.ArithValue(v_f32).bitcast(T.i32) + peer_i32 = rocdl.ds_bpermute(T.i32, lane_xor_32_byte, v_i32) + return arith.ArithValue(peer_i32).bitcast(compute_type) + return arith.ArithValue(v_f32).shuffle_xor(shuf_32_i32, width_i32) + + # ---- SLA sparse outer loop: iterate over `topk` LUT entries ---- + # N_SUBTILES == 1 by construction: each LUT entry covers exactly one + # BLOCK_N block. The inner `kv_sub` loop collapses to one iteration, + # and the dense kernel's intra-outer "next kv_sub" prefetch path is + # unused — we prefetch the next OUTER iteration (block_idx+1) instead. + assert N_SUBTILES == 1 + + # Loop-carried: [m_old, l_old, o_acc_chunks..., (buf_id if DMA dbuf)] + _use_dma_dbuf = ENABLE_DMA and not ENABLE_PREFETCH_3BUF and NUM_PREFETCH_K >= 2 + init_args = [c_neg_inf, c_zero_f] + for _ in range_constexpr(D_CHUNKS): + init_args.append(c_zero_v16f32) + if _use_dma_dbuf: + init_args.append(arith.index(0)) + # Prefetch the first SWA-window block (K AND V — Lever A V dbuf). + _init_kv_start = n_block_start * BN + coop_dma_k(_init_kv_start, buf_id=0) + coop_dma_v(_init_kv_start, buf_id=0) + + for block_idx, inner_iter_args, loop_results in scf.for_( + n_block_start, + n_block_end, + _one_idx, + iter_args=init_args, + ): + m_running = inner_iter_args[0] + l_running = inner_iter_args[1] + o_accs = [inner_iter_args[2 + i] for i in range_constexpr(D_CHUNKS)] + _cur_buf_id = inner_iter_args[2 + D_CHUNKS] if _use_dma_dbuf else None + + # V4 SWA: block_idx is directly the K-block index. + kv_block_start = block_idx * BN + preload_k_count = 1 # N_SUBTILES == 1 + + if ENABLE_PREFETCH_3BUF: + # 3-buf prefetch for sparse: look up LUT[block_idx + pre_k] and + # fire DMA per slot. Bounded by (block_idx + pre_k) < topk. + for pre_k in range_constexpr(preload_k_count): + pre_k_slot = CK_LDS_SEQ[pre_k % len(CK_LDS_SEQ)] % NUM_PREFETCH_K + if pre_k == 0: + pre_k_start = kv_block_start + else: + _pre_idx = block_idx + arith.index(pre_k) + _pre_has = arith.cmpi(arith.CmpIPredicate.slt, _pre_idx, n_block_end) + _pre_if = scf.IfOp(_pre_has) + with ir.InsertionPoint(_pre_if.then_block): + pre_k_start = _pre_idx * BN + if ENABLE_DMA: + coop_dma_k(pre_k_start, pre_k_slot) + else: + coop_load_k(pre_k_start, pre_k_slot) + scf.YieldOp([]) + continue + if ENABLE_DMA: + coop_dma_k(pre_k_start, pre_k_slot) + else: + coop_load_k(pre_k_start, pre_k_slot) + if ENABLE_DMA: + rocdl.s_waitcnt(0) + else: + rocdl.sched_group_barrier(rocdl.mask_vmem_rd, 1, 0) + gpu.barrier() + + for kv_sub in range_constexpr(N_SUBTILES): # single iteration + kv_start = kv_block_start # sparse: kv_sub == 0 always + + if ENABLE_PREFETCH_3BUF: + k_slot = CK_LDS_SEQ[kv_sub % len(CK_LDS_SEQ)] % NUM_PREFETCH_K + elif _use_dma_dbuf: + _k_buf_id = _cur_buf_id + rocdl.s_waitcnt(0) + gpu.barrier() + _next_k_buf_id = arith.index(1) - _k_buf_id + _next_block_idx = block_idx + _one_idx + _has_next = arith.cmpi( + arith.CmpIPredicate.slt, + _next_block_idx, + n_block_end, + ) + _if_dma = scf.IfOp(_has_next) + with ir.InsertionPoint(_if_dma.then_block): + _next_kv = _next_block_idx * BN + coop_dma_k(_next_kv, _next_k_buf_id) + # Lever A: also fire next V DMA into the same buf_id. + coop_dma_v(_next_kv, _next_k_buf_id) + scf.YieldOp([]) + rocdl.sched_barrier(0) + k_base = k_buf_base(_k_buf_id) + elif ENABLE_DMA: + # Single-buf DMA: fire the DMA, then wait + barrier inline. + k_slot = 0 + coop_dma_k(kv_start, k_slot) + rocdl.s_waitcnt(0) + gpu.barrier() + else: + k_slot = 0 + coop_load_k(kv_start, k_slot) + gpu.barrier() + if not _use_dma_dbuf: + k_base = k_buf_base(k_slot) + + if not USE_HW_TR or (not ENABLE_DMA and not ENABLE_PREFETCH_3BUF): + _v_vecs_prefetch = coop_load_v_global(kv_start) + + # ==== GEMM1: bulk-read all K packs, then pipeline MFMAs ==== + k_hi_offset = K_SUB_N * K_STRIDE + # XOR swizzle: col ^ ((row & 0x7) << 4) avoids LDS bank conflicts + k_swz_mask = (lane_mod_32 & arith.index(0x7)) << arith.index(4) + + def _k_idx_lo(ks): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + return k_base + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask) + + def _k_idx_hi(ks): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + return k_base + k_hi_offset + lane_mod_32 * K_STRIDE + (col ^ k_swz_mask) + + _QK_PREFETCH_DEPTH = 2 + k_packs_lo = [None] * K_STEPS_QK + k_packs_hi = [None] * K_STEPS_QK + for p in range_constexpr(_QK_PREFETCH_DEPTH): + k_packs_lo[p] = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_lo(p)]) + if N_HALVES == 2: + k_packs_hi[p] = vector.load_op(mfma_pack_type, lds_kv, [_k_idx_hi(p)]) + + if ENABLE_DMA and not ENABLE_PREFETCH_3BUF and not _use_dma_dbuf: + # Single-buf DMA path: fire V into buf 0 inside the iter. + coop_dma_v(kv_start, 0) + rocdl.sched_barrier(0) + + s_acc_lo = c_zero_v16f32 + s_acc_hi = c_zero_v16f32 + for ks in range_constexpr(K_STEPS_QK): + s_acc_lo = mfma_acc(k_packs_lo[ks], q_b_packs[ks], s_acc_lo) + if N_HALVES == 2: + s_acc_hi = mfma_acc(k_packs_hi[ks], q_b_packs[ks], s_acc_hi) + if ks + _QK_PREFETCH_DEPTH < K_STEPS_QK: + k_packs_lo[ks + _QK_PREFETCH_DEPTH] = vector.load_op( + mfma_pack_type, lds_kv, [_k_idx_lo(ks + _QK_PREFETCH_DEPTH)] + ) + if N_HALVES == 2: + k_packs_hi[ks + _QK_PREFETCH_DEPTH] = vector.load_op( + mfma_pack_type, lds_kv, [_k_idx_hi(ks + _QK_PREFETCH_DEPTH)] + ) + + # ==== Online softmax over BLOCK_N KV positions ==== + s_raw_lo = [] + s_raw_hi = [] + for r in range_constexpr(16): + s_raw_lo.append(vector.extract(s_acc_lo, static_position=[r], dynamic_position=[])) + if N_HALVES == 2: + s_raw_hi.append(vector.extract(s_acc_hi, static_position=[r], dynamic_position=[])) + + # SWA causal + boundary mask, per element. For the MFMA + # 32x32 C-layout (BLOCK_M=32 rows, BLOCK_N=64 cols arranged + # as two 32-col halves): + # row owner: q_row = q_start + wave_q_offset + lane_mod_32 + # (preloaded above as `q_row`) + # col index in tile: lane_div_32*4 + (r//4)*8 + (r%4) + # tile column 0..31 = lo half, 32..63 = hi half + # The kernel `kv_start = block_idx * BLOCK_N` is the LEFT edge + # of the lo half (kv_col_lo). kv_col_hi = kv_col_lo + 32. + # Mask conditions (set element to NEG_INF): + # 1) causal: kv_col > q_row + # 2) SWA: kv_col + W <= q_row (window length W) + # 3) boundary: kv_col >= seq_len + # NEG_INF == -inf (sla_fwd c_neg_inf). The all-masked-tile + # case is safe because m_running=-inf and l_running=0; the + # subsequent (m, l) update is identity: + # m_new = max(-inf, -inf) = -inf + # corr = exp(0) by convention but actually NaN here -- so + # we must avoid running the masked tile. + # Mitigation: n_block_start/n_block_end above already prune + # entire tiles outside the SWA window; the only tiles we run + # have AT LEAST one in-window element per warp row. The + # per-element mask handles the remaining partial overlap. + kv_start_i32 = arith.index_cast(T.i32, kv_start) + lane_div_32_i32 = arith.index_cast(T.i32, lane_div_32) + seq_len_i32 = arith.index_cast(T.i32, seq_len_v) + q_row_i32_mask = arith.index_cast(T.i32, q_row) + w_i32 = arith.constant(swa_window, type=T.i32) + # Always mask. Tile-level gate omitted (cost negligible vs MFMA). + _bool_ty = ir.IntegerType.get_signless(1) + tile_needs_mask = arith.constant(1, type=_bool_ty) + _MASK_N_OUT = 16 * N_HALVES + _mask_if = scf.IfOp(tile_needs_mask, [T.f32] * _MASK_N_OUT, has_else=True) + with ir.InsertionPoint(_mask_if.then_block): + _m_lo = [] + _m_hi = [] + for r in range_constexpr(16): + r_off_i32 = arith.constant((r % 4) + (r // 4) * 8, type=T.i32) + lane_off_i32 = arith.MulIOp(lane_div_32_i32, arith.constant(4, type=T.i32)).result + kv_col_lo = arith.AddIOp( + arith.AddIOp(kv_start_i32, lane_off_i32).result, r_off_i32 + ).result + # Boundary + is_oob_lo = arith.cmpi(arith.CmpIPredicate.sge, kv_col_lo, seq_len_i32) + # Causal: kv_col > q_row + is_causal_lo = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_lo, q_row_i32_mask) + # SWA: kv_col + W <= q_row + kv_plus_w_lo = arith.AddIOp(kv_col_lo, w_i32).result + is_swa_lo = arith.cmpi(arith.CmpIPredicate.sle, kv_plus_w_lo, q_row_i32_mask) + bad_lo = arith.OrIOp(arith.OrIOp(is_causal_lo, is_swa_lo).result, is_oob_lo).result + _m_lo.append(arith.select(bad_lo, c_neg_inf, s_raw_lo[r])) + if N_HALVES == 2: + kv_col_hi = arith.AddIOp(kv_col_lo, arith.constant(K_SUB_N, type=T.i32)).result + is_oob_hi = arith.cmpi(arith.CmpIPredicate.sge, kv_col_hi, seq_len_i32) + is_causal_hi = arith.cmpi(arith.CmpIPredicate.sgt, kv_col_hi, q_row_i32_mask) + kv_plus_w_hi = arith.AddIOp(kv_col_hi, w_i32).result + is_swa_hi = arith.cmpi(arith.CmpIPredicate.sle, kv_plus_w_hi, q_row_i32_mask) + bad_hi = arith.OrIOp( + arith.OrIOp(is_causal_hi, is_swa_hi).result, is_oob_hi + ).result + _m_hi.append(arith.select(bad_hi, c_neg_inf, s_raw_hi[r])) + scf.YieldOp(_m_lo + _m_hi) + with ir.InsertionPoint(_mask_if.else_block): + scf.YieldOp(s_raw_lo + s_raw_hi) + s_raw_lo = [_mask_if.results[i] for i in range(16)] + if N_HALVES == 2: + s_raw_hi = [_mask_if.results[16 + i] for i in range(16)] + else: + s_raw_hi = [] + + _max_fm = {"fastmath": fm_fast} + local_max = s_raw_lo[0] + for r in range_constexpr(15): + local_max = arith.MaxNumFOp(local_max, s_raw_lo[r + 1], **_max_fm).result + if N_HALVES == 2: + for r in range_constexpr(16): + local_max = arith.MaxNumFOp(local_max, s_raw_hi[r], **_max_fm).result + peer_max = reduction_peer(local_max) + row_max = arith.MaxNumFOp(local_max, peer_max, **_max_fm).result + m_new_raw = arith.MaxNumFOp(m_running, row_max, **_max_fm).result + + diff_m_raw = arith.SubFOp(m_running, m_new_raw, fastmath=fm_fast).result + diff_m_scaled = arith.MulFOp(diff_m_raw, c_sm_scale_log2e, fastmath=fm_fast).result + corr = arith.ArithValue(diff_m_scaled).exp2(fastmath=fm_fast) + + scaled_max = arith.MulFOp(c_sm_scale_log2e, m_new_raw, fastmath=fm_fast).result + neg_scaled_max = arith.SubFOp(c_zero_f, scaled_max, fastmath=fm_fast).result + + p_vals_lo = [] + p_vals_hi = [] + local_sum = c_zero_f + for r in range_constexpr(16): + diff_lo = math_dialect.fma(s_raw_lo[r], c_sm_scale_log2e, neg_scaled_max) + p_lo = arith.ArithValue(diff_lo).exp2(fastmath=fm_fast) + p_vals_lo.append(p_lo) + local_sum = arith.AddFOp(local_sum, p_lo, fastmath=fm_fast).result + if N_HALVES == 2: + for r in range_constexpr(16): + diff_hi = math_dialect.fma(s_raw_hi[r], c_sm_scale_log2e, neg_scaled_max) + p_hi = arith.ArithValue(diff_hi).exp2(fastmath=fm_fast) + p_vals_hi.append(p_hi) + local_sum = arith.AddFOp(local_sum, p_hi, fastmath=fm_fast).result + + peer_sum = reduction_peer(local_sum) + tile_sum = arith.AddFOp(local_sum, peer_sum, fastmath=fm_fast).result + l_corr = arith.MulFOp(corr, l_running, fastmath=fm_fast).result + l_new = arith.AddFOp(l_corr, tile_sum, fastmath=fm_fast).result + + # ==== Rescale O accumulators ==== + # Lever B: defer per-dc rescale to inside the PV loop so all 16 + # corr-multiplied o_accs are not live simultaneously. This shrinks + # the register live range and (we hope) reduces VGPR spill. + corr_vec = vector.broadcast(v16f32_type, corr) + if not USE_HW_TR: + o_accs[0] = arith.MulFOp(o_accs[0], corr_vec, fastmath=fm_fast).result + # USE_HW_TR: rescale happens inline in the PV loop below. + + if ENABLE_PREFETCH_3BUF and (kv_sub + preload_k_count) < N_SUBTILES: + next_k_sub = kv_sub + preload_k_count + next_k_start = kv_block_start + next_k_sub * BLOCK_N + next_k_slot = CK_LDS_SEQ[next_k_sub % len(CK_LDS_SEQ)] % NUM_PREFETCH_K + if ENABLE_DMA: + coop_dma_k(next_k_start, next_k_slot) + else: + coop_load_k(next_k_start, next_k_slot) + + if ENABLE_PREFETCH_3BUF: + v_slot = CK_LDS_SEQ[kv_sub % len(CK_LDS_SEQ)] % NUM_PREFETCH_V + v_base = v_buf_base(v_slot) + coop_load_v(kv_start, v_slot) + rocdl.sched_group_barrier(rocdl.mask_dswr, 1, 0) + gpu.barrier() + elif _use_dma_dbuf: + # Lever A: V is in the same buf_id as K (dbuf). The + # s_waitcnt(0) + barrier issued at the K branch already + # waited for the CURRENT V tile (fired in PREVIOUS iter, + # or as part of the pre-loop prefetch for iter 0). + v_base = v_buf_base(_k_buf_id) + elif ENABLE_DMA: + v_base = v_buf_base(0) + rocdl.s_waitcnt(0) + gpu.barrier() + else: + v_slot = 0 + v_base = v_buf_base(v_slot) + _waitcnt_vm_n(0) + coop_store_v_lds(_v_vecs_prefetch, v_slot) + rocdl.sched_group_barrier(rocdl.mask_dswr, 1, 0) + gpu.barrier() + + # ==== Build P packs for lo and hi halves ==== + if dtype_str == "bf16" and not USE_K16: + p_packs_lo = [] + p_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + p_base = pks * 4 + p_packs_lo.append(bf16_trunc_pack_v4(p_vals_lo[p_base : p_base + 4])) + if N_HALVES == 2: + p_packs_hi.append(bf16_trunc_pack_v4(p_vals_hi[p_base : p_base + 4])) + elif dtype_str == "bf16" and USE_K16: + p_packs_lo = [] + p_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + p_base = pks * 8 + p_packs_lo.append(bf16_trunc_pack_v8(p_vals_lo[p_base : p_base + 8])) + if N_HALVES == 2: + p_packs_hi.append(bf16_trunc_pack_v8(p_vals_hi[p_base : p_base + 8])) + else: + p_f16_lo = [] + p_f16_hi = [] + for r in range_constexpr(16): + p_f16_lo.append(arith.trunc_f(elem_type, p_vals_lo[r])) + if N_HALVES == 2: + p_f16_hi.append(arith.trunc_f(elem_type, p_vals_hi[r])) + + if USE_K16: + p_packs_lo = [] + p_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + p_base = pks * 8 + p_packs_lo.append( + vector.from_elements( + v8f16_type, + [ + p_f16_lo[p_base + 0], + p_f16_lo[p_base + 1], + p_f16_lo[p_base + 2], + p_f16_lo[p_base + 3], + p_f16_lo[p_base + 4], + p_f16_lo[p_base + 5], + p_f16_lo[p_base + 6], + p_f16_lo[p_base + 7], + ], + ) + ) + if N_HALVES == 2: + p_packs_hi.append( + vector.from_elements( + v8f16_type, + [ + p_f16_hi[p_base + 0], + p_f16_hi[p_base + 1], + p_f16_hi[p_base + 2], + p_f16_hi[p_base + 3], + p_f16_hi[p_base + 4], + p_f16_hi[p_base + 5], + p_f16_hi[p_base + 6], + p_f16_hi[p_base + 7], + ], + ) + ) + else: + p_packs_lo = [] + p_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + p_base = pks * 4 + p_packs_lo.append( + vector.from_elements( + v4f16_type, + [ + p_f16_lo[p_base], + p_f16_lo[p_base + 1], + p_f16_lo[p_base + 2], + p_f16_lo[p_base + 3], + ], + ) + ) + if N_HALVES == 2: + p_packs_hi.append( + vector.from_elements( + v4f16_type, + [ + p_f16_hi[p_base], + p_f16_hi[p_base + 1], + p_f16_hi[p_base + 2], + p_f16_hi[p_base + 3], + ], + ) + ) + + # Build flat (dc, pks) schedule for interleaved GEMM2. + _steps = [(dc, pks) for dc in range(D_CHUNKS) for pks in range(PV_K_STEPS)] + TOTAL_PV = len(_steps) + + def _read_v_pack(step_idx): + dc, pks = _steps[step_idx] + vh = None + if USE_HW_TR: + d_col = arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4 + k_row = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + tr_k_group + _d_col_eff = _v_swizzle(k_row, d_col) if ENABLE_DMA else d_col + lds_lo = v_base + k_row * V_STRIDE + _d_col_eff + if N_HALVES == 2: + lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE) + if USE_K16: + vl_a = ds_read_tr_v4f16(lds_lo) + vl_b = ds_read_tr_v4f16(lds_lo + arith.index(8 * V_STRIDE)) + vl = vector.shuffle(vl_a, vl_b, [0, 1, 2, 3, 4, 5, 6, 7]) + if N_HALVES == 2: + vh_a = ds_read_tr_v4f16(lds_hi) + vh_b = ds_read_tr_v4f16(lds_hi + arith.index(8 * V_STRIDE)) + vh = vector.shuffle(vh_a, vh_b, [0, 1, 2, 3, 4, 5, 6, 7]) + else: + vl = ds_read_tr_v4f16(lds_lo) + if N_HALVES == 2: + vh = ds_read_tr_v4f16(lds_hi) + else: + d_pos = arith.index(dc * D_CHUNK) + lane_mod_32 + k_base = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + v_lo_idx = v_base + d_pos * VT_STRIDE + k_base + vl = vector.load(v4f16_type, lds_kv, [v_lo_idx]) + if N_HALVES == 2: + v_hi_idx = v_lo_idx + arith.index(K_SUB_N) + vh = vector.load(v4f16_type, lds_kv, [v_hi_idx]) + return vl, vh + + # Pre-read V for the first step. + v_lo_cur, v_hi_cur = _read_v_pack(0) + + # ==== GEMM2: O += V^T_lo @ P_lo (+ V^T_hi @ P_hi if N_HALVES==2) ==== + for si in range_constexpr(TOTAL_PV): + dc, pks = _steps[si] + if si + 1 < TOTAL_PV: + v_lo_nxt, v_hi_nxt = _read_v_pack(si + 1) + # Lever B: per-dc inline rescale (USE_HW_TR path). + # On the first PV pack for each dc, rescale o_accs[dc] just + # before its first accumulate; keeps live range small. + if USE_HW_TR and pks == 0: + o_accs[dc] = arith.MulFOp( + o_accs[dc], + corr_vec, + fastmath=fm_fast, + ).result + o_accs[dc] = mfma_acc(v_lo_cur, p_packs_lo[pks], o_accs[dc]) + if N_HALVES == 2: + o_accs[dc] = mfma_acc(v_hi_cur, p_packs_hi[pks], o_accs[dc]) + if not USE_HW_TR and dc == 0 and pks < D_CHUNKS - 1: + o_accs[pks + 1] = arith.MulFOp( + o_accs[pks + 1], + corr_vec, + fastmath=fm_fast, + ).result + if si + 1 < TOTAL_PV: + v_lo_cur = v_lo_nxt + v_hi_cur = v_hi_nxt + + m_running = m_new_raw + l_running = l_new + + _yield_args = [m_running, l_running] + o_accs + if _use_dma_dbuf: + if N_SUBTILES % 2 == 1: + _yield_args.append(arith.index(1) - _cur_buf_id) + else: + _yield_args.append(_cur_buf_id) + yield _yield_args + + # ---- Normalize and store O (skip OOB rows for partial Q tiles) ---- + m_final = loop_results[0] + l_final = loop_results[1] + o_finals = [loop_results[2 + dc] for dc in range_constexpr(D_CHUNKS)] + + inv_l = arith.DivFOp( + c_one_f, + l_final, + fastmath=fm_fast, + ).result + inv_l_vec = vector.broadcast(v16f32_type, inv_l) + + # V4 LSE in raw-e scaled domain to match Triton ref: + # lse = m_final * sm_scale + ln(l_final) + # (Triton stores m_i + tl.log(l_i) where m_i is in qk*sm_scale domain + # already; our m_final is raw qk so we scale here.) + c_sm_scale_f = arith.constant(float(sm_scale), type=compute_type) + scaled_m_final = arith.MulFOp( + m_final, + c_sm_scale_f, + fastmath=fm_fast, + ).result + ln_l_final = math_dialect.log(l_final, fastmath=fm_fast) + lse_val = arith.AddFOp( + scaled_m_final, + ln_l_final, + fastmath=fm_fast, + ).result + + # O and LSE stores share the Q-row in-bounds guard. Note: the MFMA + # 32x32 register layout has the four rows held by this lane at offsets + # (0, 8, 16, 24) from the base (q_row = q_start + wave_q_offset + + # lane_mod_32). The O store already uses these offsets; for LSE we + # need to pick the ROW owner (lane_div_32 == 0 && lane_mod_32 < 32) + # and write one f32 per Q row. Simpler approach: every lane with + # lane_div_32 == 0 writes its lane_mod_32 row to LSE (32 rows per + # wave, one wave per Q tile row group). Matches the way `q_row` was + # computed for the Q preload at line 592. + _o_guard = scf.IfOp(q_in_bounds, [], has_else=False) + with ir.InsertionPoint(_o_guard.then_block): + for dc in range_constexpr(D_CHUNKS): + o_norm_vec = arith.MulFOp( + o_finals[dc], + inv_l_vec, + fastmath=fm_fast, + ).result + for r in range_constexpr(16): + o_val = vector.extract( + o_norm_vec, + static_position=[r], + dynamic_position=[], + ) + o_f16 = arith.trunc_f(elem_type, o_val) + + d_row_rel = lane_div_32 * 4 + (r // 4) * 8 + (r % 4) + d_col = arith.index(dc * D_CHUNK) + d_row_rel + o_global = global_idx(q_row, d_col) + _gep_store(o_f16, o_ptr, o_global) + + # LSE store: one f32 per Q row. Wave-lane layout for 32x32 MFMA + # holds one row per lane in lane_div_32==0. Gate on lane_div_32 + # == 0 to avoid 2x duplicate stores (the lane_div_32==1 copy has + # the same q_row and the same lse_val). + _is_row_owner = arith.cmpi( + arith.CmpIPredicate.eq, + lane_div_32, + arith.index(0), + ) + _lse_if = scf.IfOp(_is_row_owner, [], has_else=False) + with ir.InsertionPoint(_lse_if.then_block): + # Flat LSE index: (batch*NUM_HEADS + head)*L + q_row + lse_off = (batch_idx * NUM_HEADS + head_idx) * seq_len_v + q_row + lse_off_i32 = arith.index_cast(T.i32, lse_off) + buffer_ops.buffer_store(lse_val, lse_rsrc, lse_off_i32) + scf.YieldOp([]) + scf.YieldOp([]) + + @flyc.jit + def launch_v4_swa_fwd( + Q: fx.Tensor, + K: fx.Tensor, + V: fx.Tensor, + O: fx.Tensor, + LSE: fx.Tensor, + batch_size: fx.Int32, + seq_len: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + bs_idx = arith.index_cast(T.index, batch_size) + sl_idx = arith.index_cast(T.index, seq_len) + num_q_tiles = (sl_idx + BLOCK_M - 1) // BLOCK_M + grid_x = bs_idx * num_q_tiles * NUM_HEADS + + launcher = v4_swa_fwd_kernel(Q, K, V, O, LSE, seq_len) + + if waves_per_eu is not None: + _wpe = int(waves_per_eu) + if _wpe >= 1: + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, + _wpe, + ) + if flat_work_group_size is not None: + _fwgs = int(flat_work_group_size) + if _fwgs >= 1: + flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}") + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr + + passthrough_entries = [] + if daz: + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("denormal-fp-math-f32"), + ir.StringAttr.get("preserve-sign,preserve-sign"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("no-nans-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("unsafe-fp-math"), + ir.StringAttr.get("true"), + ] + ) + ) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries) + + launcher.launch( + grid=(grid_x, 1, 1), + block=(BLOCK_SIZE, 1, 1), + stream=stream, + ) + + # Best MI355X FMHA numbers so far were measured with ROCm/llvm-project + # `felix/tune_fmha` at c8cf6da4367c010c7cbbb7789a9c4349e7407619. + # Other LLVM revisions can compile/run this kernel, but usually leave a + # few percent of peak throughput on the table. + _fmha_compile_hints = { + "fast_fp_math": fast_fp_math, + "unsafe_fp_math": unsafe_fp_math, + "llvm_options": { + "enable-post-misched": False, + "lsr-drop-solution": True, + }, + } + + def _launch(*args, **kwargs): + with CompilationContext.compile_hints(_fmha_compile_hints): + return launch_v4_swa_fwd(*args, **kwargs) + + def _compile(Q, K, V, O, LSE, batch_size, seq_len, stream=None): + with CompilationContext.compile_hints(_fmha_compile_hints): + return flyc.compile(launch_v4_swa_fwd, Q, K, V, O, LSE, batch_size, seq_len, fx.Stream(stream)) + + _launch.compile = _compile + + return _launch + + +build_v4_swa_fwd_module_primary = build_v4_swa_fwd_module diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/__init__.py new file mode 100644 index 000000000..39fc8c1fd --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/__init__.py @@ -0,0 +1,36 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""FlyDSL-v2 DeepSeek-V4 sparse-MLA attention backend (MFMA, gfx950 / CDNA4). + +A fresh re-implementation that mirrors the gluon backend (``_gluon_dsa``): same +fused single-latent (K == V) sparse-MLA representation and the **same public +kernel-pair API**, so it plugs straight into the kernel-agnostic V4 adapter +(:mod:`v4_sparse_mla_adapter`) with zero adapter changes: + +* ``q`` : ``[T, H, d_qk]`` bf16 (``d_qk = kv_lora_rank + rope_rank``) +* ``kv`` : ``[T, 1, d_qk]`` bf16 (single MQA latent; ``V = K[:kv_lora_rank]``) +* ``topk_indices`` : ``[T, TOPK]`` int32 (SWA window ++ sparse pool, -1 = invalid) +* ``attn_sink`` : ``[H]`` fp32 optional per-head softmax sink + +Unlike the legacy wired FlyDSL CSA path (scalarized GEMV, no MFMA), the v1 +kernels use FlyDSL MFMA (``rocdl.mfma_*`` matrix cores) over a top-k gather. + +* :func:`sparse_mla_fwd_v4_flydsl` -> ``(o, lse)`` (native FlyDSL MFMA) +* :func:`sparse_mla_bwd_v4_flydsl` -> ``(dq, dkv, d_sink)`` (native FlyDSL MFMA dQ + + shared Triton dKV intermediate/scatter-gather) + +Depends only on the installed ``flydsl`` pip package (no /workspace/FlyDSL-amd +source tree required). +""" + +from .dsa_bwd_v4_flydsl import sparse_mla_bwd_v4_flydsl +from .dsa_fwd_v4_flydsl import sparse_mla_fwd_v4_flydsl + +__all__ = [ + "sparse_mla_fwd_v4_flydsl", + "sparse_mla_bwd_v4_flydsl", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_dq_flydsl_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_dq_flydsl_kernel.py new file mode 100644 index 000000000..58362ad8f --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_dq_flydsl_kernel.py @@ -0,0 +1,592 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. + +"""dsa_fwd_v4_flydsl_kernel: DeepSeek-V4 sparse-MLA attention forward (FlyDSL MFMA). + +Native FlyDSL MFMA forward for the fused single-latent (K == V) sparse-MLA form +used by the V4 attention adapter. Same public contract as the gluon / triton_v2 +backends: + + * ``q`` : ``[T, H, D_QK]`` bf16 (D_QK = kv_lora_rank + rope; rope is a zero pad) + * ``kv`` : ``[num_kv, 1, D_QK]`` bf16 (single MQA latent; V == K[:kv_lora_rank]) + * ``topk``: ``[T, TOPK]`` int32 (SWA window ++ sparse pool; -1 = invalid) + * ``sink``: ``[H]`` fp32 (optional per-head softmax sink) + * out ``o`` : ``[T, H, kv_lora_rank]`` bf16 + * out ``lse`` : ``[T, H]`` fp32 (sink-inclusive) + +Design (adapted from the in-tree v0 SWA flash kernel v4_sla_fwd_kernel.py): + * Grid: one workgroup per query TOKEN. The MFMA "M" axis is the HEAD axis + (BLOCK_H heads, one head-group = all H), the "N" axis is the gathered key + axis (BLOCK_N per tile), the contraction "K" axis is kv_lora_rank (=512). + * Each outer tile gathers BLOCK_N latent rows kv[topk[t, tile]] into LDS + (invalid topk == -1 -> row zeroed + column masked), runs the QK MFMA, an + online (flash) softmax over the key axis, and the PV MFMA (V == the same + latent, read transposed via ds_read_tr16_b64). K == V: the gathered tile is + written to a K-LDS region (XOR swizzled, for the QK read) and a V-LDS region + (row-major, for the transposed PV read). + * Epilogue: fold the per-head sink into the denominator (V4), normalize, write + O and sink-inclusive LSE = m*scale + ln(l). + +Numerics mirror v0: raw-domain running max, exp2(scale*log2e*(s - m)) softmax, +bf16 truncation pack for P, finite NEG_INF (-1e30). One addition vs v0: a +running-max clamp (m<=-1e29 -> 0) so a fully-masked leading tile (common for the +SWA window of early tokens) contributes nothing instead of exp2(0)=1. + +gfx950 / CDNA4 only (USE_HW_TR + K16 MFMA). Non-DMA cooperative-gather path first +(correctness); DMA / double-buffer / single-latent LDS sharing are follow-ups. +""" + +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import math as math_dialect +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + +_LOG2E = math.log2(math.e) # 1.4426950408889634 +_LLVM_GEP_DYNAMIC = -2147483648 # LLVM kDynamicIndex sentinel + + +def _llvm_ptr_ty(): + return ir.Type.parse("!llvm.ptr") + + +def _llvm_lds_ptr_ty(): + return ir.Type.parse("!llvm.ptr<3>") + + +def build_dsa_bwd_dq_module( + num_heads, + kv_lora_rank, + d_qk, + topk, + dtype_str="bf16", + sm_scale=None, + has_sink=True, + block_n=64, + block_h=None, + single_latent=False, + waves_per_eu=2, + unsafe_fp_math=True, + fast_fp_math=True, + daz=True, +): + """Build the sparse-MLA forward launcher (single-variant, gfx950). + + num_heads : H (must be a multiple of 32). + kv_lora_rank : D_V = 512 (contraction of QK, output dim of PV). + d_qk : row stride of q / kv (kv_lora_rank + rope pad, e.g. 576). + topk : TOPK (padded to a multiple of block_n by the launcher). + """ + gpu_arch = get_hip_arch() + assert gpu_arch.startswith("gfx950"), "dsa_fwd_v4_flydsl targets gfx950 (CDNA4)" + assert dtype_str == "bf16", "bf16 only" + + HEAD_DIM = int(kv_lora_rank) # D_V, the MFMA contraction / output dim + D_QK = int(d_qk) # q / kv row stride (includes rope pad) + NUM_HEADS = int(num_heads) + TOPK = int(topk) + assert HEAD_DIM % 32 == 0 and HEAD_DIM >= 64 + assert NUM_HEADS % 32 == 0, f"num_heads ({NUM_HEADS}) must be a multiple of 32" + + BLOCK_N = int(block_n) + assert BLOCK_N % 32 == 0 + assert TOPK % BLOCK_N == 0, f"TOPK ({TOPK}) must be a multiple of BLOCK_N ({BLOCK_N})" + K_SUB_N = 32 + N_HALVES = BLOCK_N // 32 # halves of 32 columns per BLOCK_N + + WARP_SIZE = 64 + BLOCK_H = int(block_h) if block_h else NUM_HEADS # heads per workgroup (M tile) + assert BLOCK_H % 32 == 0 and NUM_HEADS % BLOCK_H == 0 + NUM_HEAD_GROUPS = NUM_HEADS // BLOCK_H + NUM_WAVES = BLOCK_H // 32 + BLOCK_SIZE = NUM_WAVES * WARP_SIZE + ROWS_PER_WAVE = 32 # each wave owns 32 heads (MFMA M = 32) + NUM_TILES = TOPK // BLOCK_N + + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(HEAD_DIM) + + # ---- MFMA / K-step config (gfx950 CDNA4) ---- + K_STEP_QK = 16 + K_STEPS_QK = HEAD_DIM // K_STEP_QK # 32 MFMA K-steps for the QK GEMM + D_CHUNK = 32 + D_CHUNKS = HEAD_DIM // D_CHUNK # 16 output chunks for the PV accumulator + PV_K_STEP = 16 + PV_K_STEPS = K_SUB_N // PV_K_STEP # 2 PV K-steps per 32-col sub-tile + MFMA_LANE_K = 8 + + # ---- LDS layout ---- + # SINGLE_LATENT (flash / small H, LDS-occupancy-bound): one row-major tile + # (no swizzle, +4 pad) serves both QK (K) and PV (V via ds_read_tr). Halves LDS + # (occ 1->2). DUAL (pro / large H, VGPR-bound): separate XOR-swizzled K tile + # (conflict-free QK read) + row-major V tile; smaller LDS doesn't help pro + # occupancy but the swizzle avoids QK bank conflicts. + SINGLE_LATENT = bool(single_latent) + if SINGLE_LATENT: + K_STRIDE = HEAD_DIM + 4 + V_STRIDE = HEAD_DIM + 4 + LDS_V_BASE = 0 + LDS_KV_ELEMS = BLOCK_N * (HEAD_DIM + 4) + else: + K_STRIDE = HEAD_DIM + V_STRIDE = HEAD_DIM + 4 + LDS_V_BASE = BLOCK_N * K_STRIDE + LDS_KV_ELEMS = BLOCK_N * K_STRIDE + BLOCK_N * V_STRIDE + + # ---- Cooperative gather-load decomposition ---- + VEC_WIDTH = 16 + assert HEAD_DIM % VEC_WIDTH == 0 + THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH # 32 threads per gathered row + assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0 + ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD + assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0 + NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD + + allocator = SmemAllocator( + None, + arch=gpu_arch, + global_sym_name=f"dsa_fwd_smem_H{BLOCK_H}_N{BLOCK_N}_K{TOPK}", + ) + lds_kv_offset = allocator._align(allocator.ptr, 16) + lds_valid_offset = allocator._align(lds_kv_offset + LDS_KV_ELEMS * 2, 16) # f32 region (bytes) + allocator.ptr = lds_valid_offset + BLOCK_N * 4 + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def dsa_bwd_dq_kernel( + Q: fx.Tensor, # [T, H, D_QK] bf16 flat + KV: fx.Tensor, # [num_kv, D_QK] bf16 flat (single latent) + dO: fx.Tensor, # [T, H, HEAD_DIM] bf16 flat + TopK: fx.Tensor, # [T, TOPK] int32 flat + LSE: fx.Tensor, # [T, H] fp32 flat (sink-inclusive) + Delta: fx.Tensor, # [T, H] fp32 flat (rowsum(O*dO)) + dQ: fx.Tensor, # [T, H, HEAD_DIM] bf16 flat (output; rope cols discarded) + dS: fx.Tensor, # [T, H, TOPK] bf16 flat (output for dKV kernel) + Pout: fx.Tensor, # [T, H, TOPK] bf16 flat (output for dKV kernel) + ): + elem_type = T.bf16 + compute_type = T.f32 + q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q) + kv_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), KV) + do_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), dO) + dq_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), dQ) + ds_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), dS) + pout_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Pout) + topk_rsrc = buffer_ops.create_buffer_resource(TopK, max_size=True) + lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True) + delta_rsrc = buffer_ops.create_buffer_resource(Delta, max_size=True) + + fm_fast = arith.FastMathFlags.fast + v4f16_type = T.vec(4, elem_type) + vxf16_type = T.vec(VEC_WIDTH, elem_type) + v8f16_type = T.vec(8, elem_type) + v16f32_type = T.vec(16, compute_type) + mfma_pack_type = v8f16_type + + def mfma_acc(a, b, c): + # bf16, K16 (gfx950): mfma_f32_32x32x16_bf16(result_type, [a, b, c]) + return rocdl.mfma_f32_32x32x16_bf16(v16f32_type, [a, b, c]) + + # ---- LDS view ---- + base_ptr = allocator.get_base() + lds = SmemPtr(base_ptr, lds_kv_offset, elem_type, shape=(LDS_KV_ELEMS,)).get() + lds_valid = SmemPtr(base_ptr, lds_valid_offset, compute_type, shape=(BLOCK_N,)).get() + + # ---- Thread / block indices ---- + block_id = arith.index_cast(T.index, gpu.block_idx.x) + # block_id = token * NUM_HEAD_GROUPS + hg (hg=0, hg_offset=0 when 1 group) + token = block_id // arith.index(NUM_HEAD_GROUPS) + hg_offset = (block_id % arith.index(NUM_HEAD_GROUPS)) * arith.index(BLOCK_H) + tid = arith.index_cast(T.index, gpu.thread_idx.x) + + wave_id = tid // WARP_SIZE + lane = tid % WARP_SIZE + lane_mod_32 = lane % 32 + lane_div_32 = lane // 32 # 0/1 + + # ds_read_b64_tr_b16 lane decomposition (hardware 4x4 transpose) + tr_k_group = (lane % 16) // 4 + tr_col_sub = lane % 4 + tr_col_half = (lane % 32) // 16 + + wave_h_offset = wave_id * ROWS_PER_WAVE # this wave's head-row base + + # ---- ds_read_tr helper ---- + def ds_read_tr_v4f16(lds_elem_idx): + byte_offset = lds_elem_idx * 2 + lds_kv_offset + byte_i64 = arith.index_cast(T.i64, byte_offset) + ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result + return rocdl.ds_read_tr16_b64(v4f16_type, ptr).result + + # ---- global index helpers (token-major sparse-MLA) ---- + HD_IDX = arith.index(HEAD_DIM) + DQK_IDX = arith.index(D_QK) + H_IDX = arith.index(NUM_HEADS) + TOPK_IDX = arith.index(TOPK) + + def q_global_idx(head, col): + # q[token, head, col] ; row stride = H * D_QK, head stride = D_QK + return (token * H_IDX + head) * DQK_IDX + col + + def kv_global_idx(kv_row, col): + # kv[kv_row, col] ; row stride = D_QK + return kv_row * DQK_IDX + col + + def o_global_idx(head, col): + # O[token, head, col] ; row stride = H * HEAD_DIM (no rope) + return (token * H_IDX + head) * HD_IDX + col + + def _gep_load(bptr, elem_idx, vec_type): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + bptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=elem_type, + noWrapFlags=0, + ) + return _llvm.LoadOp(vec_type, gep.result).result + + def _gep_store(val, bptr, elem_idx): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + bptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=elem_type, + noWrapFlags=0, + ) + _llvm.StoreOp(val, gep.result) + + def load_global_mfma_pack(bptr, base_idx): + return _gep_load(bptr, base_idx, mfma_pack_type) + + def load_global_f16xN(bptr, base_idx): + return _gep_load(bptr, base_idx, vxf16_type) + + def bf16_trunc_pack_v8(f32_vals): + _v4i32 = T.vec(4, T.i32) + _c16 = arith.constant(16, type=T.i32) + _cmask = arith.constant(0xFFFF0000, type=T.i32) + pairs = [] + for j in range_constexpr(4): + a = arith.ArithValue(f32_vals[j * 2]).bitcast(T.i32) + b = arith.ArithValue(f32_vals[j * 2 + 1]).bitcast(T.i32) + p = arith.OrIOp(arith.AndIOp(b, _cmask).result, arith.ShRUIOp(a, _c16).result).result + pairs.append(p) + return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs)) + + # ---- cooperative decomposition ---- + load_row_in_batch = tid // THREADS_PER_ROW_LOAD + load_lane_in_row = tid % THREADS_PER_ROW_LOAD + load_col_base = load_lane_in_row * VEC_WIDTH + + c_neg_inf = arith.constant(-1.0e30, type=compute_type) + c_zero_f = arith.constant(0.0, type=compute_type) + c_one_f = arith.constant(1.0, type=compute_type) + c_zero_v16f32 = arith.constant_vector(0.0, v16f32_type) + c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type) + c_sm_scale_log2e = arith.constant(sm_scale * _LOG2E, type=compute_type) + c_sm_scale_f = arith.constant(float(sm_scale), type=compute_type) + + width_i32 = arith.constant(WARP_SIZE, type=T.i32) + shuf_32_i32 = arith.constant(32, type=T.i32) + + def reduction_peer(v_f32): + return arith.ArithValue(v_f32).shuffle_xor(shuf_32_i32, width_i32) + + # ---- K XOR swizzle (col ^ ((row & 7) << 4)) ---- + def _k_swizzle(row_idx, col_idx): + mask = (row_idx & arith.index(0x7)) << arith.index(4) + return col_idx ^ mask + + def _swz_none(row_idx, col_idx): + return col_idx + + # build-time layout selection (no traced `if`) + _swz_k = _swz_none if SINGLE_LATENT else _k_swizzle + + def _store_row_single(lds_row, col, vec): + vector.store(vec, lds, [lds_row * K_STRIDE + col]) + + def _store_row_dual(lds_row, col, vec): + vector.store(vec, lds, [lds_row * K_STRIDE + _k_swizzle(lds_row, col)]) + vector.store(vec, lds, [arith.index(LDS_V_BASE) + lds_row * V_STRIDE + col]) + + _store_row = _store_row_single if SINGLE_LATENT else _store_row_dual + + # ---- Preload Q + dO B-operand packs and per-head lse/delta ---- + # head row = hg_offset + wave_h_offset + lane_mod_32 (MFMA M axis) + head_row = hg_offset + wave_h_offset + lane_mod_32 + head_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, head_row, H_IDX) + head_row_safe = arith.select(head_in_bounds, head_row, arith.index(0)) + c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type) + q_b_packs = [] + do_b_packs = [] + for ks in range_constexpr(K_STEPS_QK): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + q_b_packs.append( + arith.select( + head_in_bounds, + load_global_mfma_pack(q_ptr, q_global_idx(head_row_safe, col)), + c_zero_mfma_pack, + ) + ) + do_b_packs.append( + arith.select( + head_in_bounds, + load_global_mfma_pack(do_ptr, o_global_idx(head_row_safe, col)), + c_zero_mfma_pack, + ) + ) + c_log2e = arith.constant(_LOG2E, type=compute_type) + head_flat_i32 = arith.index_cast(T.i32, token * H_IDX + head_row_safe) + lse_val = buffer_ops.buffer_load(lse_rsrc, head_flat_i32, vec_width=1, dtype=T.f32) + delta_val = buffer_ops.buffer_load(delta_rsrc, head_flat_i32, vec_width=1, dtype=T.f32) + neg_lse_log2e = arith.MulFOp( + lse_val, arith.SubFOp(c_zero_f, c_log2e, fastmath=fm_fast).result, fastmath=fm_fast + ).result + + # ---- outer loop over TOPK tiles: dQ += dS @ K ---- + init_args = [] + for _ in range_constexpr(D_CHUNKS): + init_args.append(c_zero_v16f32) + + for tile_idx, inner_iter_args, loop_results in scf.for_( + arith.index(0), + arith.index(NUM_TILES), + arith.index(1), + iter_args=init_args, + ): + dq_accs = [inner_iter_args[i] for i in range_constexpr(D_CHUNKS)] + + tile_topk_start = tile_idx * arith.index(BLOCK_N) + + coop_gather_tile_dyn = tile_topk_start + # gather this tile + for batch in range_constexpr(NUM_BATCHES_KV): + lds_row = load_row_in_batch + batch * ROWS_PER_BATCH_LOAD + topk_pos = coop_gather_tile_dyn + lds_row + topk_flat = token * TOPK_IDX + topk_pos + topk_flat_i32 = arith.index_cast(T.i32, topk_flat) + idx_raw = buffer_ops.buffer_load(topk_rsrc, topk_flat_i32, vec_width=1, dtype=T.i32) + valid = arith.cmpi(arith.CmpIPredicate.sge, idx_raw, arith.constant(0, type=T.i32)) + safe_i32 = arith.select(valid, idx_raw, arith.constant(0, type=T.i32)) + kv_row = arith.index_cast(T.index, safe_i32) + g_idx = kv_global_idx(kv_row, load_col_base) + vec_raw = load_global_f16xN(kv_ptr, g_idx) + vec = arith.select(valid, vec_raw, c_zero_vxf16) + _store_row(lds_row, load_col_base, vec) + is_col0 = arith.cmpi(arith.CmpIPredicate.eq, load_col_base, arith.index(0)) + _if_c0 = scf.IfOp(is_col0) + with ir.InsertionPoint(_if_c0.then_block): + mask_add = arith.select(valid, c_zero_f, c_neg_inf) + vector.store( + vector.from_elements(T.vec(1, compute_type), [mask_add]), + lds_valid, + [lds_row], + ) + scf.YieldOp([]) + gpu.barrier() + + # ==== GEMM1: QK. bulk-read K packs, pipelined MFMA ==== + k_hi_offset = K_SUB_N * K_STRIDE + + def _k_idx_lo(ks): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + return lane_mod_32 * K_STRIDE + _swz_k(lane_mod_32, col) + + def _k_idx_hi(ks): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + return k_hi_offset + lane_mod_32 * K_STRIDE + _swz_k(lane_mod_32, col) + + _QK_PREFETCH_DEPTH = 2 + k_packs_lo = [None] * K_STEPS_QK + k_packs_hi = [None] * K_STEPS_QK + for p in range_constexpr(_QK_PREFETCH_DEPTH): + k_packs_lo[p] = vector.load_op(mfma_pack_type, lds, [_k_idx_lo(p)]) + if N_HALVES == 2: + k_packs_hi[p] = vector.load_op(mfma_pack_type, lds, [_k_idx_hi(p)]) + + s_acc_lo = c_zero_v16f32 + s_acc_hi = c_zero_v16f32 + dp_acc_lo = c_zero_v16f32 + dp_acc_hi = c_zero_v16f32 + for ks in range_constexpr(K_STEPS_QK): + s_acc_lo = mfma_acc(k_packs_lo[ks], q_b_packs[ks], s_acc_lo) + dp_acc_lo = mfma_acc(k_packs_lo[ks], do_b_packs[ks], dp_acc_lo) + if N_HALVES == 2: + s_acc_hi = mfma_acc(k_packs_hi[ks], q_b_packs[ks], s_acc_hi) + dp_acc_hi = mfma_acc(k_packs_hi[ks], do_b_packs[ks], dp_acc_hi) + if ks + _QK_PREFETCH_DEPTH < K_STEPS_QK: + k_packs_lo[ks + _QK_PREFETCH_DEPTH] = vector.load_op( + mfma_pack_type, lds, [_k_idx_lo(ks + _QK_PREFETCH_DEPTH)] + ) + if N_HALVES == 2: + k_packs_hi[ks + _QK_PREFETCH_DEPTH] = vector.load_op( + mfma_pack_type, lds, [_k_idx_hi(ks + _QK_PREFETCH_DEPTH)] + ) + + # ==== P = exp(scale*S - lse); dS = P*(dP - delta)*scale (per element) ==== + # Column mapping (v0 32x32 C-layout): tile col = lane_div_32*4 + + # (r//4)*8 + (r%4) (+ K_SUB_N for the hi half). + lane_off = lane_div_32 * arith.index(4) + + def _ds_packs(s_acc, dp_acc, half): + ds_vals = [] + for r in range_constexpr(16): + s_r = vector.extract(s_acc, static_position=[r], dynamic_position=[]) + dp_r = vector.extract(dp_acc, static_position=[r], dynamic_position=[]) + col = lane_off + arith.index((r % 4) + (r // 4) * 8 + half * K_SUB_N) + mv = vector.extract( + vector.load_op(T.vec(1, compute_type), lds_valid, [col]), + static_position=[0], + dynamic_position=[], + ) + s_m = arith.AddFOp(s_r, mv, fastmath=fm_fast).result # invalid -> -inf + # P = exp2(scale*log2e*s_m - lse*log2e) + p_arg = math_dialect.fma(s_m, c_sm_scale_log2e, neg_lse_log2e) + p_r = arith.ArithValue(p_arg).exp2(fastmath=fm_fast) + # dS = P * (dP - delta) * scale + dp_md = arith.SubFOp(dp_r, delta_val, fastmath=fm_fast).result + ds_r = arith.MulFOp( + arith.MulFOp(p_r, dp_md, fastmath=fm_fast).result, c_sm_scale_f, fastmath=fm_fast + ).result + ds_vals.append(ds_r) + # store dS and P to [T, H, TOPK] for the dKV kernel (kv_pos in [0,TOPK)) + kv_pos = tile_topk_start + col + dsp_idx = (token * H_IDX + head_row) * TOPK_IDX + kv_pos + _gep_store(arith.trunc_f(elem_type, ds_r), ds_ptr, dsp_idx) + _gep_store(arith.trunc_f(elem_type, p_r), pout_ptr, dsp_idx) + packs = [] + for pks in range_constexpr(PV_K_STEPS): + packs.append(bf16_trunc_pack_v8(ds_vals[pks * 8 : pks * 8 + 8])) + return packs + + ds_packs_lo = _ds_packs(s_acc_lo, dp_acc_lo, 0) + ds_packs_hi = _ds_packs(s_acc_hi, dp_acc_hi, 1) if N_HALVES == 2 else [] + + # ==== GEMM2: dQ += dS @ K (K == V, ds_read_tr) ==== + v_base = arith.index(LDS_V_BASE) + _steps = [(dc, pks) for dc in range(D_CHUNKS) for pks in range(PV_K_STEPS)] + TOTAL_PV = len(_steps) + + def _read_v_pack(step_idx): + dc, pks = _steps[step_idx] + d_col = arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4 + k_row = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + tr_k_group + lds_lo = v_base + k_row * V_STRIDE + d_col + vl_a = ds_read_tr_v4f16(lds_lo) + vl_b = ds_read_tr_v4f16(lds_lo + arith.index(8 * V_STRIDE)) + vl = vector.shuffle(vl_a, vl_b, [0, 1, 2, 3, 4, 5, 6, 7]) + vh = None + if N_HALVES == 2: + lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE) + vh_a = ds_read_tr_v4f16(lds_hi) + vh_b = ds_read_tr_v4f16(lds_hi + arith.index(8 * V_STRIDE)) + vh = vector.shuffle(vh_a, vh_b, [0, 1, 2, 3, 4, 5, 6, 7]) + return vl, vh + + for si in range_constexpr(TOTAL_PV): + dc, pks = _steps[si] + v_lo_cur, v_hi_cur = _read_v_pack(si) + dq_accs[dc] = mfma_acc(v_lo_cur, ds_packs_lo[pks], dq_accs[dc]) + if N_HALVES == 2: + dq_accs[dc] = mfma_acc(v_hi_cur, ds_packs_hi[pks], dq_accs[dc]) + + # protect this tile's V/K LDS reads from the next tile's gather writes + gpu.barrier() + + yield dq_accs + + # ---- epilogue: store dQ_lora [token, head, :HEAD_DIM] ---- + dq_finals = [loop_results[dc] for dc in range_constexpr(D_CHUNKS)] + _o_guard = scf.IfOp(head_in_bounds, [], has_else=False) + with ir.InsertionPoint(_o_guard.then_block): + for dc in range_constexpr(D_CHUNKS): + for r in range_constexpr(16): + dq_val = vector.extract(dq_finals[dc], static_position=[r], dynamic_position=[]) + dq_f16 = arith.trunc_f(elem_type, dq_val) + d_row_rel = lane_div_32 * 4 + (r // 4) * 8 + (r % 4) + d_col = arith.index(dc * D_CHUNK) + d_row_rel + # dQ is [T, H, d_qk] (matches q); write the :HEAD_DIM lora cols + _gep_store(dq_f16, dq_ptr, q_global_idx(head_row, d_col)) + scf.YieldOp([]) + + @flyc.jit + def launch_dsa_bwd_dq( + Q: fx.Tensor, + KV: fx.Tensor, + dO: fx.Tensor, + TopK: fx.Tensor, + LSE: fx.Tensor, + Delta: fx.Tensor, + dQ: fx.Tensor, + dS: fx.Tensor, + Pout: fx.Tensor, + total_tokens: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + grid_x = arith.index_cast(T.index, total_tokens) * arith.index(NUM_HEAD_GROUPS) + launcher = dsa_bwd_dq_kernel(Q, KV, dO, TopK, LSE, Delta, dQ, dS, Pout) + + if waves_per_eu is not None and int(waves_per_eu) >= 1: + _wpe = int(waves_per_eu) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe) + + _fwgs = int(BLOCK_SIZE) + flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}") + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr + + passthrough_entries = [] + if daz: + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("denormal-fp-math-f32"), + ir.StringAttr.get("preserve-sign,preserve-sign"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get([ir.StringAttr.get("no-nans-fp-math"), ir.StringAttr.get("true")]) + ) + passthrough_entries.append( + ir.ArrayAttr.get([ir.StringAttr.get("unsafe-fp-math"), ir.StringAttr.get("true")]) + ) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries) + + launcher.launch(grid=(grid_x, 1, 1), block=(BLOCK_SIZE, 1, 1), stream=stream) + + _compile_hints = { + "fast_fp_math": fast_fp_math, + "unsafe_fp_math": unsafe_fp_math, + "llvm_options": {"enable-post-misched": False, "lsr-drop-solution": True}, + } + + def _launch(*args, **kwargs): + with CompilationContext.compile_hints(_compile_hints): + return launch_dsa_bwd_dq(*args, **kwargs) + + return _launch diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_v4_flydsl.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_v4_flydsl.py new file mode 100644 index 000000000..cf3ff37d8 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_bwd_v4_flydsl.py @@ -0,0 +1,184 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""FlyDSL-v1 sparse-MLA backward — native FlyDSL dQ + shared dKV gather. + +``sparse_mla_bwd_v4_flydsl(q, kv, o, do, topk, lse, ...) -> (dq, dkv, d_sink)``. + +The **dQ** kernel is native FlyDSL MFMA (``dsa_bwd_dq_flydsl_kernel``): it +recomputes S = Q·Kᵀ, P = exp(scale·S − lse), dP = dO·Kᵀ, dS = P·(dP − Δ)·scale, +accumulates dQ = dS·K, and writes the per-token dS/P buffers. The dKV path +(intermediate GEMM + CSR inverted-topk scatter-reduce) reuses the shared, +proven Triton kernels — the dKV gather is a variable-length scatter-reduction +with no MFMA content, so there is nothing to gain from a FlyDSL rewrite there. + +This runs the whole top-k as a single chunk (R_CHUNK = TOPK), so dQ needs no +cross-chunk read-modify-write and the dKV intermediate is built in one pass. + +Depends only on the installed ``flydsl`` pip package (no /workspace source). +""" + +from __future__ import annotations + +import os +import sys +import threading + +import torch +import triton + +from .._gluon_dsa._dsa_bwd_gather import _build_inverted_topk_slice +from .._triton_v2.dsa_bwd_kernels import ( + _bwd_compute_dkv_intermediate, + _bwd_dkv_gather_acc, +) + +_HERE = os.path.dirname(os.path.abspath(__file__)) +if _HERE not in sys.path: + sys.path.insert(0, _HERE) + +from dsa_bwd_dq_flydsl_kernel import build_dsa_bwd_dq_module # noqa: E402 + +_DQ_CACHE = {} +_DQ_LOCK = threading.Lock() +_BLOCK_N = int(os.environ.get("PRIMUS_DSA_FLYDSL_BWD_BLOCK_N", "64")) +_BLOCK_H = int(os.environ.get("PRIMUS_DSA_FLYDSL_BWD_BLOCK_H", "64")) + + +def _get_dq_kernel(num_heads, kv_lora_rank, d_qk, topk, single_latent, scale): + block_h = min(_BLOCK_H, num_heads) + while num_heads % block_h != 0: + block_h -= 32 + key = (num_heads, kv_lora_rank, d_qk, topk, block_h, single_latent, round(float(scale), 8)) + with _DQ_LOCK: + launch = _DQ_CACHE.get(key) + if launch is None: + launch = build_dsa_bwd_dq_module( + num_heads=num_heads, + kv_lora_rank=kv_lora_rank, + d_qk=d_qk, + topk=topk, + dtype_str="bf16", + sm_scale=float(scale), + block_n=_BLOCK_N, + block_h=block_h, + single_latent=single_latent, + ) + _DQ_CACHE[key] = launch + return launch + + +def sparse_mla_bwd_v4_flydsl(q, kv, o, do, topk_indices, lse, attn_sink=None, kv_lora_rank=512, scale=None): + """DeepSeek-V4 sparse-MLA backward: native FlyDSL dQ + Triton dKV gather.""" + assert q.is_contiguous() and o.is_contiguous() and do.is_contiguous() + assert topk_indices.is_contiguous() and lse.is_contiguous() + total_tokens, num_heads, d_qk = q.shape + D = int(kv_lora_rank) + rope_rank = d_qk - D + if scale is None: + scale = 1.0 / (d_qk**0.5) + if kv.dim() == 2: + kv = kv.unsqueeze(1) + assert kv.is_contiguous() + num_kv = kv.shape[0] + assert q.dtype == torch.bfloat16 + + has_sink = attn_sink is not None + if has_sink: + assert attn_sink.dtype == torch.float32 and attn_sink.shape == (num_heads,) + + # pad topk to a BLOCK_N multiple (-1 = invalid); usually already padded upstream + topk = topk_indices.shape[1] + if topk % _BLOCK_N != 0: + pad = ((topk + _BLOCK_N - 1) // _BLOCK_N) * _BLOCK_N - topk + topk_p = torch.cat( + [topk_indices, torch.full((total_tokens, pad), -1, dtype=torch.int32, device=q.device)], dim=1 + ).contiguous() + else: + topk_p = topk_indices + TOPK = topk_p.shape[1] + + # Delta = rowsum(O * dO) (o is [T,H,D]) + delta = (o[:, :, :D].float() * do.float()).sum(-1).contiguous() + lse_c = lse.contiguous() + + # ---- native FlyDSL dQ (whole top-k) + dS/P buffers ---- + dq = torch.zeros(total_tokens, num_heads, d_qk, dtype=q.dtype, device=q.device) # rope cols stay 0 + chunk_dS = torch.empty(total_tokens, num_heads, TOPK, dtype=torch.bfloat16, device=q.device) + chunk_P = torch.empty(total_tokens, num_heads, TOPK, dtype=torch.bfloat16, device=q.device) + single_latent = num_heads <= 64 + dq_launch = _get_dq_kernel(int(num_heads), D, int(d_qk), int(TOPK), single_latent, scale) + dq_launch( + q.reshape(-1), + kv.reshape(-1), + do.reshape(-1), + topk_p.reshape(-1), + lse_c.reshape(-1), + delta.reshape(-1), + dq.reshape(-1), + chunk_dS.reshape(-1), + chunk_P.reshape(-1), + int(total_tokens), + ) + + # ---- dKV: intermediate GEMM + CSR inverted-topk scatter-reduce (Triton) ---- + # BH_DKV=32/TK_DKV=64 whole-top-k (R_CHUNK=TOPK) in one dKV-intermediate pass. + # This config's LDS stays < 160 KB in both the standalone and training Triton + # contexts (the wider 64/128 config overflows under the training runtime). + HAS_ROPE = False + BH_DKV, TK_DKV = 32, 64 + num_hg_dkv = triton.cdiv(num_heads, BH_DKV) + + interm = torch.empty(total_tokens, TOPK, d_qk, dtype=torch.bfloat16, device=q.device) + _bwd_compute_dkv_intermediate[(total_tokens,)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + num_heads, + R_CHUNK=TOPK, + TILE_K=TK_DKV, + BLOCK_H=BH_DKV, + NUM_HG=num_hg_dkv, + D_V=D, + D_ROPE=rope_rank, + HAS_ROPE=HAS_ROPE, + num_warps=4, + ) + + dkv_acc = torch.zeros(num_kv, d_qk, dtype=torch.float32, device=q.device) + inv_ptr, inv_data = _build_inverted_topk_slice(topk_p, 0, TOPK, num_kv=num_kv) + _bwd_dkv_gather_acc[(num_kv,)]( + interm, + inv_ptr, + inv_data, + dkv_acc, + interm.stride(1), + dkv_acc.stride(0), + D_V=D, + D_ROPE=rope_rank, + HAS_ROPE=HAS_ROPE, + num_warps=4, + ) + + d_sink = None + if has_sink: + d_sink = -(torch.exp(attn_sink.unsqueeze(0) - lse) * delta).sum(0) + + dkv = dkv_acc.to(kv.dtype).unsqueeze(1) + return dq, dkv, d_sink + + +__all__ = ["sparse_mla_bwd_v4_flydsl"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl.py new file mode 100644 index 000000000..8366214ce --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl.py @@ -0,0 +1,151 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""FlyDSL-v1 sparse-MLA forward (native FlyDSL MFMA) launcher. + +Public API mirrors :func:`sparse_mla_fwd_v4_gluon` / ``sparse_mla_fwd_v4_triton``: + + sparse_mla_fwd_v4_flydsl(q, kv, topk_indices, attn_sink=None, + kv_lora_rank=512, scale=None) -> (o, lse) + +The heavy lifting is the native FlyDSL MFMA kernel in +``dsa_fwd_v4_flydsl_kernel.build_dsa_fwd_module``; this module just marshals the +tensors, caches the built launcher per (H, D, TOPK, has_sink, block_n), and +launches it. See the kernel module docstring for the design. +""" + +from __future__ import annotations + +import os +import sys +import threading + +import torch + +# flydsl_v1 uses only the installed `flydsl` pip package + its own local kernel +# modules — it does NOT need the /workspace/FlyDSL-amd source tree. +_HERE = os.path.dirname(os.path.abspath(__file__)) +if _HERE not in sys.path: + sys.path.insert(0, _HERE) + +from dsa_fwd_v4_flydsl_kernel import build_dsa_fwd_module # noqa: E402 + +_KERNEL_CACHE = {} +_KERNEL_CACHE_LOCK = threading.Lock() + +_DEFAULT_BLOCK_N = int(os.environ.get("PRIMUS_DSA_FLYDSL_FWD_BLOCK_N", "0")) # 0 = shape-conditional +_DEFAULT_BLOCK_H = int(os.environ.get("PRIMUS_DSA_FLYDSL_FWD_BLOCK_H", "256")) +_DEFAULT_WPE = int(os.environ.get("PRIMUS_DSA_FLYDSL_FWD_WPE", "2")) + + +def _get_kernel( + num_heads, kv_lora_rank, d_qk, topk, has_sink, block_n, block_h, single_latent, waves_per_eu, scale +): + block_h = min(block_h, num_heads) + while num_heads % block_h != 0: + block_h -= 32 + key = ( + num_heads, + kv_lora_rank, + d_qk, + topk, + has_sink, + block_n, + block_h, + single_latent, + waves_per_eu, + round(float(scale), 8), + ) + with _KERNEL_CACHE_LOCK: + launch = _KERNEL_CACHE.get(key) + if launch is None: + launch = build_dsa_fwd_module( + num_heads=num_heads, + kv_lora_rank=kv_lora_rank, + d_qk=d_qk, + topk=topk, + dtype_str="bf16", + sm_scale=float(scale), + has_sink=has_sink, + block_n=block_n, + block_h=block_h, + single_latent=single_latent, + waves_per_eu=waves_per_eu, + ) + _KERNEL_CACHE[key] = launch + return launch + + +def sparse_mla_fwd_v4_flydsl(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None): + """DeepSeek-V4 sparse-MLA forward (native FlyDSL MFMA, gfx950).""" + assert q.is_contiguous() and topk_indices.is_contiguous() + total_tokens, num_heads, d_qk = q.shape + kv_lora_rank = int(kv_lora_rank) + if scale is None: + scale = 1.0 / (d_qk**0.5) + if kv.dim() == 2: + kv = kv.unsqueeze(1) + assert kv.is_contiguous() + assert kv.shape[0] >= total_tokens and kv.shape[-1] == d_qk + assert q.dtype == torch.bfloat16, "bf16 only" + + # Shape-conditional tile: flash (H<=64) is LDS/occupancy-bound -> small tile + # (BLOCK_N=32) doubles occupancy (occ 1->2). pro (H>=128) is VGPR-bound + # (occupancy stuck at 1 regardless of LDS) -> a smaller tile only adds + # barrier/softmax overhead, so keep the larger BLOCK_N=64 (fewer tiles). + # flash (H<=64) is LDS/occupancy-bound → single-latent (one shared tile) halves + # LDS (occ 1→2) at BLOCK_N=64 (single-latent already halves LDS, so no need for + # a smaller tile — smaller tiles only add barriers). pro (H>=128) is VGPR-bound + # → dual swizzled tile (conflict-free QK), extra LDS is free (occ stuck at 1). + single_latent = num_heads <= 64 + block_n = _DEFAULT_BLOCK_N if _DEFAULT_BLOCK_N else 64 + topk = topk_indices.shape[1] + # Pad TOPK up to a multiple of block_n with -1 (masked) if needed. + if topk % block_n != 0: + pad = ((topk + block_n - 1) // block_n) * block_n - topk + topk_indices = torch.cat( + [topk_indices, torch.full((total_tokens, pad), -1, dtype=torch.int32, device=q.device)], + dim=1, + ).contiguous() + topk = topk_indices.shape[1] + + has_sink = attn_sink is not None + if has_sink: + assert attn_sink.is_contiguous() and attn_sink.dtype == torch.float32 + assert attn_sink.shape == (num_heads,) + sink = attn_sink + else: + # kernel always folds the sink; -inf makes it a no-op (af=1, sink_e=0) + sink = torch.full((num_heads,), float("-inf"), dtype=torch.float32, device=q.device) + + o = torch.empty(total_tokens, num_heads, kv_lora_rank, dtype=q.dtype, device=q.device) + lse = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device) + + launch = _get_kernel( + int(num_heads), + kv_lora_rank, + int(d_qk), + int(topk), + has_sink, + block_n, + _DEFAULT_BLOCK_H, + single_latent, + _DEFAULT_WPE, + scale, + ) + launch( + q.reshape(-1), + kv.reshape(-1), + topk_indices.reshape(-1), + sink.reshape(-1), + o.reshape(-1), + lse.reshape(-1), + int(total_tokens), + ) + return o, lse + + +__all__ = ["sparse_mla_fwd_v4_flydsl"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl_kernel.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl_kernel.py new file mode 100644 index 000000000..4d74fafe9 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_flydsl_v1/dsa_fwd_v4_flydsl_kernel.py @@ -0,0 +1,650 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. + +"""dsa_fwd_v4_flydsl_kernel: DeepSeek-V4 sparse-MLA attention forward (FlyDSL MFMA). + +Native FlyDSL MFMA forward for the fused single-latent (K == V) sparse-MLA form +used by the V4 attention adapter. Same public contract as the gluon / triton_v2 +backends: + + * ``q`` : ``[T, H, D_QK]`` bf16 (D_QK = kv_lora_rank + rope; rope is a zero pad) + * ``kv`` : ``[num_kv, 1, D_QK]`` bf16 (single MQA latent; V == K[:kv_lora_rank]) + * ``topk``: ``[T, TOPK]`` int32 (SWA window ++ sparse pool; -1 = invalid) + * ``sink``: ``[H]`` fp32 (optional per-head softmax sink) + * out ``o`` : ``[T, H, kv_lora_rank]`` bf16 + * out ``lse`` : ``[T, H]`` fp32 (sink-inclusive) + +Design (adapted from the in-tree v0 SWA flash kernel v4_sla_fwd_kernel.py): + * Grid: one workgroup per query TOKEN. The MFMA "M" axis is the HEAD axis + (BLOCK_H heads, one head-group = all H), the "N" axis is the gathered key + axis (BLOCK_N per tile), the contraction "K" axis is kv_lora_rank (=512). + * Each outer tile gathers BLOCK_N latent rows kv[topk[t, tile]] into LDS + (invalid topk == -1 -> row zeroed + column masked), runs the QK MFMA, an + online (flash) softmax over the key axis, and the PV MFMA (V == the same + latent, read transposed via ds_read_tr16_b64). K == V: the gathered tile is + written to a K-LDS region (XOR swizzled, for the QK read) and a V-LDS region + (row-major, for the transposed PV read). + * Epilogue: fold the per-head sink into the denominator (V4), normalize, write + O and sink-inclusive LSE = m*scale + ln(l). + +Numerics mirror v0: raw-domain running max, exp2(scale*log2e*(s - m)) softmax, +bf16 truncation pack for P, finite NEG_INF (-1e30). One addition vs v0: a +running-max clamp (m<=-1e29 -> 0) so a fully-masked leading tile (common for the +SWA window of early tokens) contributes nothing instead of exp2(0)=1. + +gfx950 / CDNA4 only (USE_HW_TR + K16 MFMA). Non-DMA cooperative-gather path first +(correctness); DMA / double-buffer / single-latent LDS sharing are follow-ups. +""" + +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl._mlir.dialects import fly as _fly +from flydsl._mlir.dialects import llvm as _llvm +from flydsl._mlir.dialects import math as math_dialect +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, buffer_ops, gpu, range_constexpr, rocdl, vector +from flydsl.expr.typing import T +from flydsl.runtime.device import get_rocm_arch as get_hip_arch +from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr + +_LOG2E = math.log2(math.e) # 1.4426950408889634 +_LLVM_GEP_DYNAMIC = -2147483648 # LLVM kDynamicIndex sentinel + + +def _llvm_ptr_ty(): + return ir.Type.parse("!llvm.ptr") + + +def _llvm_lds_ptr_ty(): + return ir.Type.parse("!llvm.ptr<3>") + + +def build_dsa_fwd_module( + num_heads, + kv_lora_rank, + d_qk, + topk, + dtype_str="bf16", + sm_scale=None, + has_sink=True, + block_n=64, + block_h=None, + single_latent=False, + waves_per_eu=2, + unsafe_fp_math=True, + fast_fp_math=True, + daz=True, +): + """Build the sparse-MLA forward launcher (single-variant, gfx950). + + num_heads : H (must be a multiple of 32). + kv_lora_rank : D_V = 512 (contraction of QK, output dim of PV). + d_qk : row stride of q / kv (kv_lora_rank + rope pad, e.g. 576). + topk : TOPK (padded to a multiple of block_n by the launcher). + """ + gpu_arch = get_hip_arch() + assert gpu_arch.startswith("gfx950"), "dsa_fwd_v4_flydsl targets gfx950 (CDNA4)" + assert dtype_str == "bf16", "bf16 only" + + HEAD_DIM = int(kv_lora_rank) # D_V, the MFMA contraction / output dim + D_QK = int(d_qk) # q / kv row stride (includes rope pad) + NUM_HEADS = int(num_heads) + TOPK = int(topk) + assert HEAD_DIM % 32 == 0 and HEAD_DIM >= 64 + assert NUM_HEADS % 32 == 0, f"num_heads ({NUM_HEADS}) must be a multiple of 32" + + BLOCK_N = int(block_n) + assert BLOCK_N % 32 == 0 + assert TOPK % BLOCK_N == 0, f"TOPK ({TOPK}) must be a multiple of BLOCK_N ({BLOCK_N})" + K_SUB_N = 32 + N_HALVES = BLOCK_N // 32 # halves of 32 columns per BLOCK_N + + WARP_SIZE = 64 + BLOCK_H = int(block_h) if block_h else NUM_HEADS # heads per workgroup (M tile) + assert BLOCK_H % 32 == 0 and NUM_HEADS % BLOCK_H == 0 + NUM_HEAD_GROUPS = NUM_HEADS // BLOCK_H + NUM_WAVES = BLOCK_H // 32 + BLOCK_SIZE = NUM_WAVES * WARP_SIZE + ROWS_PER_WAVE = 32 # each wave owns 32 heads (MFMA M = 32) + NUM_TILES = TOPK // BLOCK_N + + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(HEAD_DIM) + + # ---- MFMA / K-step config (gfx950 CDNA4) ---- + K_STEP_QK = 16 + K_STEPS_QK = HEAD_DIM // K_STEP_QK # 32 MFMA K-steps for the QK GEMM + D_CHUNK = 32 + D_CHUNKS = HEAD_DIM // D_CHUNK # 16 output chunks for the PV accumulator + PV_K_STEP = 16 + PV_K_STEPS = K_SUB_N // PV_K_STEP # 2 PV K-steps per 32-col sub-tile + MFMA_LANE_K = 8 + + # ---- LDS layout ---- + # SINGLE_LATENT (flash / small H, LDS-occupancy-bound): one row-major tile + # (no swizzle, +4 pad) serves both QK (K) and PV (V via ds_read_tr). Halves LDS + # (occ 1->2). DUAL (pro / large H, VGPR-bound): separate XOR-swizzled K tile + # (conflict-free QK read) + row-major V tile; smaller LDS doesn't help pro + # occupancy but the swizzle avoids QK bank conflicts. + SINGLE_LATENT = bool(single_latent) + if SINGLE_LATENT: + K_STRIDE = HEAD_DIM + 4 + V_STRIDE = HEAD_DIM + 4 + LDS_V_BASE = 0 + LDS_KV_ELEMS = BLOCK_N * (HEAD_DIM + 4) + else: + K_STRIDE = HEAD_DIM + V_STRIDE = HEAD_DIM + 4 + LDS_V_BASE = BLOCK_N * K_STRIDE + LDS_KV_ELEMS = BLOCK_N * K_STRIDE + BLOCK_N * V_STRIDE + + # ---- Cooperative gather-load decomposition ---- + VEC_WIDTH = 16 + assert HEAD_DIM % VEC_WIDTH == 0 + THREADS_PER_ROW_LOAD = HEAD_DIM // VEC_WIDTH # 32 threads per gathered row + assert BLOCK_SIZE % THREADS_PER_ROW_LOAD == 0 + ROWS_PER_BATCH_LOAD = BLOCK_SIZE // THREADS_PER_ROW_LOAD + assert BLOCK_N % ROWS_PER_BATCH_LOAD == 0 + NUM_BATCHES_KV = BLOCK_N // ROWS_PER_BATCH_LOAD + + allocator = SmemAllocator( + None, + arch=gpu_arch, + global_sym_name=f"dsa_fwd_smem_H{BLOCK_H}_N{BLOCK_N}_K{TOPK}", + ) + lds_kv_offset = allocator._align(allocator.ptr, 16) + lds_valid_offset = allocator._align(lds_kv_offset + LDS_KV_ELEMS * 2, 16) # f32 region (bytes) + allocator.ptr = lds_valid_offset + BLOCK_N * 4 + + @flyc.kernel(known_block_size=[BLOCK_SIZE, 1, 1]) + def dsa_fwd_kernel( + Q: fx.Tensor, # [T, H, D_QK] bf16 flat + KV: fx.Tensor, # [num_kv, D_QK] bf16 flat (single latent) + TopK: fx.Tensor, # [T, TOPK] int32 flat + Sink: fx.Tensor, # [H] fp32 flat + O: fx.Tensor, # [T, H, HEAD_DIM] bf16 flat + LSE: fx.Tensor, # [T, H] fp32 flat + ): + elem_type = T.bf16 + compute_type = T.f32 + q_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), Q) + kv_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), KV) + o_ptr = _fly.extract_aligned_pointer_as_index(_llvm_ptr_ty(), O) + topk_rsrc = buffer_ops.create_buffer_resource(TopK, max_size=True) + sink_rsrc = buffer_ops.create_buffer_resource(Sink, max_size=True) + lse_rsrc = buffer_ops.create_buffer_resource(LSE, max_size=True) + + fm_fast = arith.FastMathFlags.fast + v4f16_type = T.vec(4, elem_type) + vxf16_type = T.vec(VEC_WIDTH, elem_type) + v8f16_type = T.vec(8, elem_type) + v16f32_type = T.vec(16, compute_type) + mfma_pack_type = v8f16_type + + def mfma_acc(a, b, c): + # bf16, K16 (gfx950): mfma_f32_32x32x16_bf16(result_type, [a, b, c]) + return rocdl.mfma_f32_32x32x16_bf16(v16f32_type, [a, b, c]) + + # ---- LDS view ---- + base_ptr = allocator.get_base() + lds = SmemPtr(base_ptr, lds_kv_offset, elem_type, shape=(LDS_KV_ELEMS,)).get() + lds_valid = SmemPtr(base_ptr, lds_valid_offset, compute_type, shape=(BLOCK_N,)).get() + + # ---- Thread / block indices ---- + block_id = arith.index_cast(T.index, gpu.block_idx.x) + # block_id = token * NUM_HEAD_GROUPS + hg (hg=0, hg_offset=0 when 1 group) + token = block_id // arith.index(NUM_HEAD_GROUPS) + hg_offset = (block_id % arith.index(NUM_HEAD_GROUPS)) * arith.index(BLOCK_H) + tid = arith.index_cast(T.index, gpu.thread_idx.x) + + wave_id = tid // WARP_SIZE + lane = tid % WARP_SIZE + lane_mod_32 = lane % 32 + lane_div_32 = lane // 32 # 0/1 + + # ds_read_b64_tr_b16 lane decomposition (hardware 4x4 transpose) + tr_k_group = (lane % 16) // 4 + tr_col_sub = lane % 4 + tr_col_half = (lane % 32) // 16 + + wave_h_offset = wave_id * ROWS_PER_WAVE # this wave's head-row base + + # ---- ds_read_tr helper ---- + def ds_read_tr_v4f16(lds_elem_idx): + byte_offset = lds_elem_idx * 2 + lds_kv_offset + byte_i64 = arith.index_cast(T.i64, byte_offset) + ptr = _llvm.IntToPtrOp(_llvm_lds_ptr_ty(), byte_i64).result + return rocdl.ds_read_tr16_b64(v4f16_type, ptr).result + + # ---- global index helpers (token-major sparse-MLA) ---- + HD_IDX = arith.index(HEAD_DIM) + DQK_IDX = arith.index(D_QK) + H_IDX = arith.index(NUM_HEADS) + TOPK_IDX = arith.index(TOPK) + + def q_global_idx(head, col): + # q[token, head, col] ; row stride = H * D_QK, head stride = D_QK + return (token * H_IDX + head) * DQK_IDX + col + + def kv_global_idx(kv_row, col): + # kv[kv_row, col] ; row stride = D_QK + return kv_row * DQK_IDX + col + + def o_global_idx(head, col): + # O[token, head, col] ; row stride = H * HEAD_DIM (no rope) + return (token * H_IDX + head) * HD_IDX + col + + def _gep_load(bptr, elem_idx, vec_type): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + bptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=elem_type, + noWrapFlags=0, + ) + return _llvm.LoadOp(vec_type, gep.result).result + + def _gep_store(val, bptr, elem_idx): + idx_i64 = arith.index_cast(T.i64, elem_idx) + gep = _llvm.GEPOp( + _llvm_ptr_ty(), + bptr, + [idx_i64], + rawConstantIndices=[_LLVM_GEP_DYNAMIC], + elem_type=elem_type, + noWrapFlags=0, + ) + _llvm.StoreOp(val, gep.result) + + def load_global_mfma_pack(bptr, base_idx): + return _gep_load(bptr, base_idx, mfma_pack_type) + + def load_global_f16xN(bptr, base_idx): + return _gep_load(bptr, base_idx, vxf16_type) + + def bf16_trunc_pack_v8(f32_vals): + _v4i32 = T.vec(4, T.i32) + _c16 = arith.constant(16, type=T.i32) + _cmask = arith.constant(0xFFFF0000, type=T.i32) + pairs = [] + for j in range_constexpr(4): + a = arith.ArithValue(f32_vals[j * 2]).bitcast(T.i32) + b = arith.ArithValue(f32_vals[j * 2 + 1]).bitcast(T.i32) + p = arith.OrIOp(arith.AndIOp(b, _cmask).result, arith.ShRUIOp(a, _c16).result).result + pairs.append(p) + return vector.bitcast(v8f16_type, vector.from_elements(_v4i32, pairs)) + + # ---- cooperative decomposition ---- + load_row_in_batch = tid // THREADS_PER_ROW_LOAD + load_lane_in_row = tid % THREADS_PER_ROW_LOAD + load_col_base = load_lane_in_row * VEC_WIDTH + + c_neg_inf = arith.constant(-1.0e30, type=compute_type) + c_zero_f = arith.constant(0.0, type=compute_type) + c_one_f = arith.constant(1.0, type=compute_type) + c_zero_v16f32 = arith.constant_vector(0.0, v16f32_type) + c_zero_vxf16 = arith.constant_vector(0.0, vxf16_type) + c_sm_scale_log2e = arith.constant(sm_scale * _LOG2E, type=compute_type) + c_sm_scale_f = arith.constant(float(sm_scale), type=compute_type) + + width_i32 = arith.constant(WARP_SIZE, type=T.i32) + shuf_32_i32 = arith.constant(32, type=T.i32) + + def reduction_peer(v_f32): + return arith.ArithValue(v_f32).shuffle_xor(shuf_32_i32, width_i32) + + # ---- K XOR swizzle (col ^ ((row & 7) << 4)) ---- + def _k_swizzle(row_idx, col_idx): + mask = (row_idx & arith.index(0x7)) << arith.index(4) + return col_idx ^ mask + + def _swz_none(row_idx, col_idx): + return col_idx + + # build-time layout selection (no traced `if`) + _swz_k = _swz_none if SINGLE_LATENT else _k_swizzle + + def _store_row_single(lds_row, col, vec): + vector.store(vec, lds, [lds_row * K_STRIDE + col]) + + def _store_row_dual(lds_row, col, vec): + vector.store(vec, lds, [lds_row * K_STRIDE + _k_swizzle(lds_row, col)]) + vector.store(vec, lds, [arith.index(LDS_V_BASE) + lds_row * V_STRIDE + col]) + + _store_row = _store_row_single if SINGLE_LATENT else _store_row_dual + + # ---- Preload Q B-operand packs (register-resident) ---- + # head row = hg_offset + wave_h_offset + lane_mod_32 (MFMA M axis) + head_row = hg_offset + wave_h_offset + lane_mod_32 + head_in_bounds = arith.cmpi(arith.CmpIPredicate.slt, head_row, H_IDX) + head_row_safe = arith.select(head_in_bounds, head_row, arith.index(0)) + c_zero_mfma_pack = arith.constant_vector(0.0, mfma_pack_type) + q_b_packs = [] + for ks in range_constexpr(K_STEPS_QK): + q_col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + g_idx = q_global_idx(head_row_safe, q_col) + raw = load_global_mfma_pack(q_ptr, g_idx) + q_b_packs.append(arith.select(head_in_bounds, raw, c_zero_mfma_pack)) + + # ---- outer loop over TOPK tiles ---- + init_args = [c_neg_inf, c_zero_f] + for _ in range_constexpr(D_CHUNKS): + init_args.append(c_zero_v16f32) + + for tile_idx, inner_iter_args, loop_results in scf.for_( + arith.index(0), + arith.index(NUM_TILES), + arith.index(1), + iter_args=init_args, + ): + m_running = inner_iter_args[0] + l_running = inner_iter_args[1] + o_accs = [inner_iter_args[2 + i] for i in range_constexpr(D_CHUNKS)] + + tile_topk_start = tile_idx * arith.index(BLOCK_N) + + coop_gather_tile_dyn = tile_topk_start + # gather this tile + for batch in range_constexpr(NUM_BATCHES_KV): + lds_row = load_row_in_batch + batch * ROWS_PER_BATCH_LOAD + topk_pos = coop_gather_tile_dyn + lds_row + topk_flat = token * TOPK_IDX + topk_pos + topk_flat_i32 = arith.index_cast(T.i32, topk_flat) + idx_raw = buffer_ops.buffer_load(topk_rsrc, topk_flat_i32, vec_width=1, dtype=T.i32) + valid = arith.cmpi(arith.CmpIPredicate.sge, idx_raw, arith.constant(0, type=T.i32)) + safe_i32 = arith.select(valid, idx_raw, arith.constant(0, type=T.i32)) + kv_row = arith.index_cast(T.index, safe_i32) + g_idx = kv_global_idx(kv_row, load_col_base) + vec_raw = load_global_f16xN(kv_ptr, g_idx) + vec = arith.select(valid, vec_raw, c_zero_vxf16) + _store_row(lds_row, load_col_base, vec) + is_col0 = arith.cmpi(arith.CmpIPredicate.eq, load_col_base, arith.index(0)) + _if_c0 = scf.IfOp(is_col0) + with ir.InsertionPoint(_if_c0.then_block): + mask_add = arith.select(valid, c_zero_f, c_neg_inf) + vector.store( + vector.from_elements(T.vec(1, compute_type), [mask_add]), + lds_valid, + [lds_row], + ) + scf.YieldOp([]) + gpu.barrier() + + # ==== GEMM1: QK. bulk-read K packs, pipelined MFMA ==== + k_hi_offset = K_SUB_N * K_STRIDE + + def _k_idx_lo(ks): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + return lane_mod_32 * K_STRIDE + _swz_k(lane_mod_32, col) + + def _k_idx_hi(ks): + col = arith.index(ks * K_STEP_QK) + lane_div_32 * MFMA_LANE_K + return k_hi_offset + lane_mod_32 * K_STRIDE + _swz_k(lane_mod_32, col) + + _QK_PREFETCH_DEPTH = 2 + k_packs_lo = [None] * K_STEPS_QK + k_packs_hi = [None] * K_STEPS_QK + for p in range_constexpr(_QK_PREFETCH_DEPTH): + k_packs_lo[p] = vector.load_op(mfma_pack_type, lds, [_k_idx_lo(p)]) + if N_HALVES == 2: + k_packs_hi[p] = vector.load_op(mfma_pack_type, lds, [_k_idx_hi(p)]) + + s_acc_lo = c_zero_v16f32 + s_acc_hi = c_zero_v16f32 + for ks in range_constexpr(K_STEPS_QK): + s_acc_lo = mfma_acc(k_packs_lo[ks], q_b_packs[ks], s_acc_lo) + if N_HALVES == 2: + s_acc_hi = mfma_acc(k_packs_hi[ks], q_b_packs[ks], s_acc_hi) + if ks + _QK_PREFETCH_DEPTH < K_STEPS_QK: + k_packs_lo[ks + _QK_PREFETCH_DEPTH] = vector.load_op( + mfma_pack_type, lds, [_k_idx_lo(ks + _QK_PREFETCH_DEPTH)] + ) + if N_HALVES == 2: + k_packs_hi[ks + _QK_PREFETCH_DEPTH] = vector.load_op( + mfma_pack_type, lds, [_k_idx_hi(ks + _QK_PREFETCH_DEPTH)] + ) + + # ==== Online softmax over BLOCK_N KV positions ==== + s_raw_lo = [] + s_raw_hi = [] + for r in range_constexpr(16): + s_raw_lo.append(vector.extract(s_acc_lo, static_position=[r], dynamic_position=[])) + if N_HALVES == 2: + s_raw_hi.append(vector.extract(s_acc_hi, static_position=[r], dynamic_position=[])) + + # Add validity additive mask (from LDS). Column mapping (v0 32x32 + # C-layout): tile col_lo = lane_div_32*4 + (r//4)*8 + (r%4); hi = +32. + lane_off = lane_div_32 * arith.index(4) + _m_lo = [] + _m_hi = [] + for r in range_constexpr(16): + r_off = arith.index((r % 4) + (r // 4) * 8) + col_lo = lane_off + r_off + mv_lo = vector.load_op(T.vec(1, compute_type), lds_valid, [col_lo]) + mval_lo = vector.extract(mv_lo, static_position=[0], dynamic_position=[]) + _m_lo.append(arith.AddFOp(s_raw_lo[r], mval_lo, fastmath=fm_fast).result) + if N_HALVES == 2: + col_hi = col_lo + arith.index(K_SUB_N) + mv_hi = vector.load_op(T.vec(1, compute_type), lds_valid, [col_hi]) + mval_hi = vector.extract(mv_hi, static_position=[0], dynamic_position=[]) + _m_hi.append(arith.AddFOp(s_raw_hi[r], mval_hi, fastmath=fm_fast).result) + s_raw_lo = _m_lo + s_raw_hi = _m_hi + + _max_fm = {"fastmath": fm_fast} + local_max = s_raw_lo[0] + for r in range_constexpr(15): + local_max = arith.MaxNumFOp(local_max, s_raw_lo[r + 1], **_max_fm).result + if N_HALVES == 2: + for r in range_constexpr(16): + local_max = arith.MaxNumFOp(local_max, s_raw_hi[r], **_max_fm).result + peer_max = reduction_peer(local_max) + row_max = arith.MaxNumFOp(local_max, peer_max, **_max_fm).result + m_new_raw = arith.MaxNumFOp(m_running, row_max, **_max_fm).result + # clamp: fully-masked-so-far -> 0 (placeholder; shift-invariant) + _finite = arith.cmpf( + arith.CmpFPredicate.OGT, m_new_raw, arith.constant(-1.0e29, type=compute_type) + ) + m_new_raw = arith.select(_finite, m_new_raw, c_zero_f) + + diff_m_raw = arith.SubFOp(m_running, m_new_raw, fastmath=fm_fast).result + diff_m_scaled = arith.MulFOp(diff_m_raw, c_sm_scale_log2e, fastmath=fm_fast).result + corr = arith.ArithValue(diff_m_scaled).exp2(fastmath=fm_fast) + + scaled_max = arith.MulFOp(c_sm_scale_log2e, m_new_raw, fastmath=fm_fast).result + neg_scaled_max = arith.SubFOp(c_zero_f, scaled_max, fastmath=fm_fast).result + + p_vals_lo = [] + p_vals_hi = [] + local_sum = c_zero_f + for r in range_constexpr(16): + diff_lo = math_dialect.fma(s_raw_lo[r], c_sm_scale_log2e, neg_scaled_max) + p_lo = arith.ArithValue(diff_lo).exp2(fastmath=fm_fast) + p_vals_lo.append(p_lo) + local_sum = arith.AddFOp(local_sum, p_lo, fastmath=fm_fast).result + if N_HALVES == 2: + for r in range_constexpr(16): + diff_hi = math_dialect.fma(s_raw_hi[r], c_sm_scale_log2e, neg_scaled_max) + p_hi = arith.ArithValue(diff_hi).exp2(fastmath=fm_fast) + p_vals_hi.append(p_hi) + local_sum = arith.AddFOp(local_sum, p_hi, fastmath=fm_fast).result + + peer_sum = reduction_peer(local_sum) + tile_sum = arith.AddFOp(local_sum, peer_sum, fastmath=fm_fast).result + l_corr = arith.MulFOp(corr, l_running, fastmath=fm_fast).result + l_new = arith.AddFOp(l_corr, tile_sum, fastmath=fm_fast).result + + corr_vec = vector.broadcast(v16f32_type, corr) + # online-softmax rescale of all O accumulators before this tile's PV + for dc in range_constexpr(D_CHUNKS): + o_accs[dc] = arith.MulFOp(o_accs[dc], corr_vec, fastmath=fm_fast).result + + # ==== Build P packs (bf16 truncation) ==== + p_packs_lo = [] + p_packs_hi = [] + for pks in range_constexpr(PV_K_STEPS): + p_base = pks * 8 + p_packs_lo.append(bf16_trunc_pack_v8(p_vals_lo[p_base : p_base + 8])) + if N_HALVES == 2: + p_packs_hi.append(bf16_trunc_pack_v8(p_vals_hi[p_base : p_base + 8])) + + # ==== GEMM2: PV. read V transposed (ds_read_tr), interleaved ==== + v_base = arith.index(LDS_V_BASE) + _steps = [(dc, pks) for dc in range(D_CHUNKS) for pks in range(PV_K_STEPS)] + TOTAL_PV = len(_steps) + + def _read_v_pack(step_idx): + dc, pks = _steps[step_idx] + d_col = arith.index(dc * D_CHUNK) + tr_col_half * 16 + tr_col_sub * 4 + k_row = arith.index(pks * PV_K_STEP) + lane_div_32 * 4 + tr_k_group + lds_lo = v_base + k_row * V_STRIDE + d_col + vl_a = ds_read_tr_v4f16(lds_lo) + vl_b = ds_read_tr_v4f16(lds_lo + arith.index(8 * V_STRIDE)) + vl = vector.shuffle(vl_a, vl_b, [0, 1, 2, 3, 4, 5, 6, 7]) + vh = None + if N_HALVES == 2: + lds_hi = lds_lo + arith.index(K_SUB_N * V_STRIDE) + vh_a = ds_read_tr_v4f16(lds_hi) + vh_b = ds_read_tr_v4f16(lds_hi + arith.index(8 * V_STRIDE)) + vh = vector.shuffle(vh_a, vh_b, [0, 1, 2, 3, 4, 5, 6, 7]) + return vl, vh + + for si in range_constexpr(TOTAL_PV): + dc, pks = _steps[si] + v_lo_cur, v_hi_cur = _read_v_pack(si) + o_accs[dc] = mfma_acc(v_lo_cur, p_packs_lo[pks], o_accs[dc]) + if N_HALVES == 2: + o_accs[dc] = mfma_acc(v_hi_cur, p_packs_hi[pks], o_accs[dc]) + + # protect this tile's V/K LDS reads from the next tile's gather writes + gpu.barrier() + + m_running = m_new_raw + l_running = l_new + + yield [m_running, l_running] + o_accs + + # ---- epilogue: sink fold + normalize + store ---- + m_final = loop_results[0] + l_final = loop_results[1] + o_finals = [loop_results[2 + dc] for dc in range_constexpr(D_CHUNKS)] + + # scaled-domain sink fold (matches triton_v2 / gluon): + # M = m_final * sm_scale ; l_final = sum exp(scaled_s - M) + # Always fold the per-head sink into the denominator (V4). For the + # no-sink case the launcher fills Sink with -inf, so af=1, sink_e=0 -> + # lse=M+log(l), acc_scale=1/l (identical to the no-sink formula). This + # avoids a Python `if` in the traced body (the AST rewriter does not + # propagate branch-local rebindings out of a dispatched if). + _log2e = arith.constant(_LOG2E, type=compute_type) + M_scaled = arith.MulFOp(m_final, c_sm_scale_f, fastmath=fm_fast).result + head_row_i32 = arith.index_cast(T.i32, head_row_safe) + sink_val = buffer_ops.buffer_load(sink_rsrc, head_row_i32, vec_width=1, dtype=T.f32) + m_fin = arith.MaxNumFOp(M_scaled, sink_val, fastmath=fm_fast).result + _daf = arith.MulFOp( + arith.SubFOp(M_scaled, m_fin, fastmath=fm_fast).result, _log2e, fastmath=fm_fast + ).result + af = arith.ArithValue(_daf).exp2(fastmath=fm_fast) + l_af = arith.MulFOp(l_final, af, fastmath=fm_fast).result + _dse = arith.MulFOp( + arith.SubFOp(sink_val, m_fin, fastmath=fm_fast).result, _log2e, fastmath=fm_fast + ).result + sink_e = arith.ArithValue(_dse).exp2(fastmath=fm_fast) + l_total = arith.AddFOp(l_af, sink_e, fastmath=fm_fast).result + ln_l = math_dialect.log(l_total, fastmath=fm_fast) + lse_val = arith.AddFOp(m_fin, ln_l, fastmath=fm_fast).result + inv_l = arith.DivFOp(c_one_f, l_total, fastmath=fm_fast).result + acc_scale = arith.MulFOp(af, inv_l, fastmath=fm_fast).result + + acc_scale_vec = vector.broadcast(v16f32_type, acc_scale) + + _o_guard = scf.IfOp(head_in_bounds, [], has_else=False) + with ir.InsertionPoint(_o_guard.then_block): + for dc in range_constexpr(D_CHUNKS): + o_norm_vec = arith.MulFOp(o_finals[dc], acc_scale_vec, fastmath=fm_fast).result + for r in range_constexpr(16): + o_val = vector.extract(o_norm_vec, static_position=[r], dynamic_position=[]) + o_f16 = arith.trunc_f(elem_type, o_val) + d_row_rel = lane_div_32 * 4 + (r // 4) * 8 + (r % 4) + d_col = arith.index(dc * D_CHUNK) + d_row_rel + _gep_store(o_f16, o_ptr, o_global_idx(head_row, d_col)) + + _is_row_owner = arith.cmpi(arith.CmpIPredicate.eq, lane_div_32, arith.index(0)) + _lse_if = scf.IfOp(_is_row_owner, [], has_else=False) + with ir.InsertionPoint(_lse_if.then_block): + lse_off = token * H_IDX + head_row + lse_off_i32 = arith.index_cast(T.i32, lse_off) + buffer_ops.buffer_store(lse_val, lse_rsrc, lse_off_i32) + scf.YieldOp([]) + scf.YieldOp([]) + + @flyc.jit + def launch_dsa_fwd( + Q: fx.Tensor, + KV: fx.Tensor, + TopK: fx.Tensor, + Sink: fx.Tensor, + O: fx.Tensor, + LSE: fx.Tensor, + total_tokens: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + allocator.finalized = False + ctx = CompilationContext.get_current() + with ir.InsertionPoint(ctx.gpu_module_body): + allocator.finalize() + + grid_x = arith.index_cast(T.index, total_tokens) * arith.index(NUM_HEAD_GROUPS) + launcher = dsa_fwd_kernel(Q, KV, TopK, Sink, O, LSE) + + if waves_per_eu is not None and int(waves_per_eu) >= 1: + _wpe = int(waves_per_eu) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get(T.i32, _wpe) + + _fwgs = int(BLOCK_SIZE) + flat_wg_attr = ir.StringAttr.get(f"{_fwgs},{_fwgs}") + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["rocdl.flat_work_group_size"] = flat_wg_attr + + passthrough_entries = [] + if daz: + passthrough_entries.append( + ir.ArrayAttr.get( + [ + ir.StringAttr.get("denormal-fp-math-f32"), + ir.StringAttr.get("preserve-sign,preserve-sign"), + ] + ) + ) + passthrough_entries.append( + ir.ArrayAttr.get([ir.StringAttr.get("no-nans-fp-math"), ir.StringAttr.get("true")]) + ) + passthrough_entries.append( + ir.ArrayAttr.get([ir.StringAttr.get("unsafe-fp-math"), ir.StringAttr.get("true")]) + ) + for op in ctx.gpu_module_body.operations: + if getattr(op, "OPERATION_NAME", None) == "gpu.func": + op.attributes["passthrough"] = ir.ArrayAttr.get(passthrough_entries) + + launcher.launch(grid=(grid_x, 1, 1), block=(BLOCK_SIZE, 1, 1), stream=stream) + + _compile_hints = { + "fast_fp_math": fast_fp_math, + "unsafe_fp_math": unsafe_fp_math, + "llvm_options": {"enable-post-misched": False, "lsr-drop-solution": True}, + } + + def _launch(*args, **kwargs): + with CompilationContext.compile_hints(_compile_hints): + return launch_dsa_fwd(*args, **kwargs) + + return _launch diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/__init__.py new file mode 100644 index 000000000..3d38c10ab --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/__init__.py @@ -0,0 +1,37 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Gluon (hardware-controlled Triton) DeepSeek-V4 sparse-MLA attention backend. + +Ported from ROCm/aiter PR #2922 (``aiter/ops/triton/_gluon_kernels/gfx950``) +for gfx950 / CDNA4 (MI350/MI355X). These kernels operate on the **sparse-MLA +latent** representation used by the DeepSeek V4 paper / FlashMLA: + +* ``q`` : ``[T, H, d_qk]`` with ``d_qk = kv_lora_rank (512) + rope_rank (64)`` +* ``kv`` : ``[T, 1, d_qk]`` single MQA latent (K and V share it; ``V_lora`` is + the first ``kv_lora_rank`` channels of ``K_lora``) +* ``topk_indices`` : ``[T, TOPK]`` int32 absolute KV-token indices (SWA window + + sparse top-k already concatenated by the caller; ``-1`` = invalid) +* ``attn_sink`` : ``[H]`` fp32 optional per-head learnable softmax sink + +This is a different (latent + per-token-topk) representation than the in-tree +CSA path (``v4_csa_attention_v0``: ``q / k_local / v_local / gathered / +sparse_mask``); it is exposed here as a standalone ``gluon`` backend. + +Public API mirrors aiter's ``sparse_mla_fwd_v4`` / ``sparse_mla_bwd_v4`` with +``backend="gluon"``: + +* :func:`sparse_mla_fwd_v4_gluon` -> ``(o, lse)`` +* :func:`sparse_mla_bwd_v4_gluon` -> ``(dq, dkv, d_sink)`` +""" + +from .dsa_bwd_v4_gluon import sparse_mla_bwd_v4_gluon +from .dsa_fwd_v4_gluon import sparse_mla_fwd_v4_gluon + +__all__ = [ + "sparse_mla_fwd_v4_gluon", + "sparse_mla_bwd_v4_gluon", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_gather.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_gather.py new file mode 100644 index 000000000..0626c846e --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_gather.py @@ -0,0 +1,617 @@ +import torch +import triton +import triton.language as tl + + +# ===================================================================== +# Backward method="gather" — intermediate storage + inverted topk gather +# ===================================================================== +@triton.jit +def _bwd_compute_dkv_intermediate( + Q_T_ptr, # [T, D_QK, H] bf16 (q transposed: stride_qt_t = D_QK * H) + dO_T_ptr, # [T, D_V, H] bf16 + dS_ptr, # [T, H, TOPK] bf16 + P_ptr, # [T, H, TOPK] bf16 + TopK_ptr, # [T, TOPK] int32 + Interm_ptr, # [T, TOPK, D_QK] bf16 — output, one writer per (q, topk_rank) + stride_qt_t: tl.int64, + stride_dot_t: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + stride_topk_t: tl.int64, + stride_interm_t: tl.int64, # TOPK * D_QK + stride_interm_k: tl.int64, # D_QK + num_heads: tl.int32, + TOPK: tl.constexpr, + TILE_K: tl.constexpr, + BLOCK_H: tl.constexpr, + NUM_HG: tl.constexpr, + D_V: tl.constexpr, + D_ROPE: tl.constexpr, +): + """ + Same compute as _bwd_dkv_hg_fused but writes to a private intermediate + [T, TOPK, D] bf16 instead of atomic_add to shared dKV — no atomics needed. + + Grid: (total_tokens,) -- one program per query token q. + For each tile of TOPK, stores dKV_lora/rope for that (q, tile) block. + """ + token_idx = tl.program_id(0) + + NUM_TILES: tl.constexpr = (TOPK + TILE_K - 1) // TILE_K + token_idx * stride_topk_t + offs_tile = tl.arange(0, TILE_K) + offs_v = tl.arange(0, D_V) + offs_r = tl.arange(0, D_ROPE) + + interm_base_t = token_idx * stride_interm_t # base for this query token + + for t in range(NUM_TILES): + tile_start = t * TILE_K + tile_offs = tile_start + offs_tile + valid = tile_offs < TOPK + + dKV_lora = tl.zeros([D_V, TILE_K], dtype=tl.float32) + dKV_rope = tl.zeros([D_ROPE, TILE_K], dtype=tl.float32) + + for hg in range(NUM_HG): + offs_h = hg * BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = offs_h < num_heads + + qt_base = token_idx * stride_qt_t + Q_lora_T = tl.load( + Q_T_ptr + qt_base + offs_v[:, None] * num_heads + offs_h[None, :], + mask=mask_h[None, :], + other=0.0, + ) + Q_rope_T = tl.load( + Q_T_ptr + qt_base + (D_V + offs_r[:, None]) * num_heads + offs_h[None, :], + mask=mask_h[None, :], + other=0.0, + ) + + dot_base = token_idx * stride_dot_t + dO_T = tl.load( + dO_T_ptr + dot_base + offs_v[:, None] * num_heads + offs_h[None, :], + mask=mask_h[None, :], + other=0.0, + ) + + ds_base = token_idx * stride_ds_t + dS_val = tl.load( + dS_ptr + ds_base + offs_h[:, None] * stride_ds_h + tile_offs[None, :], + mask=mask_h[:, None] & valid[None, :], + other=0.0, + ) + P_val = tl.load( + P_ptr + ds_base + offs_h[:, None] * stride_ds_h + tile_offs[None, :], + mask=mask_h[:, None] & valid[None, :], + other=0.0, + ) + + dKV_lora += tl.dot(Q_lora_T, dS_val.to(Q_lora_T.dtype)).to(tl.float32) + dKV_lora += tl.dot(dO_T, P_val.to(dO_T.dtype)).to(tl.float32) + dKV_rope += tl.dot(Q_rope_T, dS_val.to(Q_rope_T.dtype)).to(tl.float32) + + # Store to intermediate: Interm[token_idx, tile_start:tile_start+TILE_K, 0:D_V] + # Layout: [T, TOPK, D] so pointer = interm_base_t + tile_offs[None,:]*D + offs_v[:,None] + interm_lora_ptrs = Interm_ptr + interm_base_t + tile_offs[None, :] * stride_interm_k + offs_v[:, None] + tl.store(interm_lora_ptrs, dKV_lora.to(tl.bfloat16), mask=valid[None, :]) + + interm_rope_ptrs = ( + Interm_ptr + interm_base_t + tile_offs[None, :] * stride_interm_k + D_V + offs_r[:, None] + ) + tl.store(interm_rope_ptrs, dKV_rope.to(tl.bfloat16), mask=valid[None, :]) + + +@triton.jit +def _bwd_dkv_gather( + Interm_ptr, # [T, TOPK, D] bf16, flattened as [T*TOPK, D] + InvPtr_ptr, # [T+1] int32 — CSR row pointers (kv_token -> range in inv_data) + InvData_ptr, # [T*TOPK] int32 — encoded as q*TOPK+r, sorted by KV token + dKV_ptr, # [T, D] bf16 — output + stride_interm_k: tl.int64, # D_V + D_ROPE + stride_dkv_t: tl.int64, + TOPK: tl.constexpr, + D_V: tl.constexpr, + D_ROPE: tl.constexpr, +): + """ + Gather dKV from intermediate buffer using CSR-style inverted topk index. + + Grid: (total_tokens,) -- one CTA per KV token k. + Accumulates in fp32, stores bf16. No atomics. + """ + k = tl.program_id(0) + offs_v = tl.arange(0, D_V) + offs_r = tl.arange(0, D_ROPE) + + start = tl.load(InvPtr_ptr + k) + end = tl.load(InvPtr_ptr + k + 1) + + dkv_acc_lora = tl.zeros([D_V], dtype=tl.float32) + dkv_acc_rope = tl.zeros([D_ROPE], dtype=tl.float32) + + n_entries = end - start + for i in range(n_entries): + # entry = q*TOPK + r, used directly as flat index into [T*TOPK, D] intermediate + entry = tl.load(InvData_ptr + start + i).to(tl.int64) + base = entry * stride_interm_k + lora_val = tl.load(Interm_ptr + base + offs_v) + rope_val = tl.load(Interm_ptr + base + D_V + offs_r) + dkv_acc_lora += lora_val.to(tl.float32) + dkv_acc_rope += rope_val.to(tl.float32) + + dkv_base = k.to(tl.int64) * stride_dkv_t + tl.store(dKV_ptr + dkv_base + offs_v, dkv_acc_lora.to(tl.bfloat16)) + tl.store(dKV_ptr + dkv_base + D_V + offs_r, dkv_acc_rope.to(tl.bfloat16)) + + +def _build_inverted_topk(topk_indices): + """ + Build CSR-style inverted index from topk_indices [T, TOPK] int32. + + Returns: + inv_ptr: [T+1] int32 — row pointers (kv_token -> range in inv_data) + inv_data: [T*TOPK] int32 — encoded (q*TOPK+r) values, sorted by KV token + """ + T, TOPK = topk_indices.shape + device = topk_indices.device + + flat_kv = topk_indices.reshape(-1).long() # [T*TOPK] KV token indices + # argsort by KV token to get the flat indices (q*TOPK+r) in sorted order + order = torch.argsort(flat_kv, stable=True) + inv_data = order.to(torch.int32) # [T*TOPK] + + counts = torch.zeros(T, dtype=torch.int32, device=device) + counts.scatter_add_(0, flat_kv, torch.ones(T * TOPK, dtype=torch.int32, device=device)) + + inv_ptr = torch.zeros(T + 1, dtype=torch.int32, device=device) + torch.cumsum(counts, dim=0, out=inv_ptr[1:]) + + return inv_ptr, inv_data + + +def _build_inverted_topk_slice(topk_indices_slice, r_start, R_CHUNK, num_kv=None): + """ + Build CSR-style inverted index for a topk slice, excluding invalid (-1) entries. + + Args: + topk_indices_slice: [T, R_CHUNK] int32 — topk_indices[:, r_start:r_start+R_CHUNK] + May contain -1 for padding (when actual chunk < R_CHUNK at the last chunk). + r_start: int — first rank index in this slice (unused, for documentation) + R_CHUNK: int — number of ranks in this slice (constexpr width) + num_kv: int or None — number of KV tokens. When the KV buffer holds + MORE rows than query tokens (V4 [local ++ pool], num_kv = S + P), pass + it so ``inv_ptr`` has length ``num_kv + 1`` even if the last KV tokens + are referenced by no query. Defaults to ``T`` (query tokens). + + Returns: + inv_ptr: [num_kv+1] int32 — row pointers (kv_token -> range in inv_data) + inv_data: [valid_entries] int32 — flat indices q*R_CHUNK+local_r, sorted by KV token + """ + T, RC = topk_indices_slice.shape + n_kv = T if num_kv is None else int(num_kv) + flat_kv = topk_indices_slice.reshape(-1).long() # [T*R_CHUNK]; -1 marks invalid + + # Sort ALL entries by KV token. argsort returns the flat positions (q*RC+r) in + # KV-token order; invalid (-1) entries sort to the FRONT and are never referenced + # because inv_ptr[0] starts past them. This avoids the boolean-mask `nonzero` + # + per-element `scatter_add` + extra `index` gathers of the previous version + # (all slow torch ops, especially on gfx1250 where they dominated the bwd). + inv_data = torch.argsort(flat_kv, stable=True).to(torch.int32) # [T*R_CHUNK] + # bincount of (kv+1): bin 0 = #invalid, bins 1..num_kv = per-KV counts. + counts = torch.bincount(flat_kv + 1, minlength=n_kv + 1) # [num_kv+1] + inv_ptr = torch.cumsum(counts, dim=0).to(torch.int32) # [num_kv+1]; inv_ptr[0]=#invalid + + return inv_ptr, inv_data + + +# ===================================================================== +# Backward method="chunked_gather" — three kernels per chunk (no atomics) +# ===================================================================== +@triton.jit +def _bwd_chunk_dq_store_ds( + Q_ptr, # [T, H, D] bf16 + KV_ptr, # [T, 1, D] bf16 + dO_ptr, # [T, H, D_V] bf16 + TopK_ptr, # [T, TOPK] int32 + LSE_ptr, # [T, H] fp32 + Delta_ptr, # [T, H] fp32 + dQ_ptr, # [T, H, D] bf16 — read-modify-write across chunks + dS_ptr, # [T, H, R_CHUNK] bf16 — output chunk dS + P_ptr, # [T, H, R_CHUNK] bf16 — output chunk P + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_dq_t: tl.int64, + stride_dq_h: tl.int64, + stride_topk_t: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + R_START: tl.int32, + R_CHUNK: tl.constexpr, + BLOCK_H: tl.constexpr, + TILE_K: tl.constexpr, + D_V: tl.constexpr, + D_ROPE: tl.constexpr, + IS_FIRST_CHUNK: tl.constexpr, +): + """ + dQ accumulation for rank chunk [R_START, R_START+R_CHUNK), plus stores + chunk dS and P to [T, H, R_CHUNK] buffers for use by _bwd_compute_dkv_intermediate. + Grid: (total_tokens, num_hg). + """ + token_idx = tl.program_id(0) + hg_idx = tl.program_id(1) + offs_h = hg_idx * BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = offs_h < num_heads + offs_v = tl.arange(0, D_V) + offs_r = tl.arange(0, D_ROPE) + + q_base = token_idx * stride_q_t + Q_lora = tl.load( + Q_ptr + q_base + offs_h[:, None] * stride_q_h + offs_v[None, :], mask=mask_h[:, None], other=0.0 + ) + Q_rope = tl.load( + Q_ptr + q_base + offs_h[:, None] * stride_q_h + (D_V + offs_r[None, :]), + mask=mask_h[:, None], + other=0.0, + ) + do_base = token_idx * stride_do_t + dO_val = tl.load( + dO_ptr + do_base + offs_h[:, None] * stride_do_h + offs_v[None, :], mask=mask_h[:, None], other=0.0 + ) + lse = tl.load(LSE_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0) + delta = tl.load(Delta_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0) + + dq_base = token_idx * stride_dq_t + if IS_FIRST_CHUNK: + dQ_lora = tl.zeros([BLOCK_H, D_V], dtype=tl.float32) + dQ_rope = tl.zeros([BLOCK_H, D_ROPE], dtype=tl.float32) + else: + dQ_lora = tl.load( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :], + mask=mask_h[:, None], + other=0.0, + ).to(tl.float32) + dQ_rope = tl.load( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]), + mask=mask_h[:, None], + other=0.0, + ).to(tl.float32) + + NUM_TILES: tl.constexpr = (R_CHUNK + TILE_K - 1) // TILE_K + topk_base = token_idx * stride_topk_t + R_START + offs_tile = tl.arange(0, TILE_K) + ds_base = token_idx * stride_ds_t + hg_idx * BLOCK_H * stride_ds_h + + for t in range(NUM_TILES): + tile_start = t * TILE_K + tile_offs = tile_start + offs_tile + valid = tile_offs < R_CHUNK + topk_pos = tl.load(TopK_ptr + topk_base + tile_offs, mask=valid, other=-1) + valid = valid & (topk_pos != -1) + safe_pos = tl.where(valid, topk_pos, 0) + + K_lora_T = tl.load( + KV_ptr + safe_pos[None, :] * stride_kv_t + offs_v[:, None], mask=valid[None, :], other=0.0 + ) + K_rope_T = tl.load( + KV_ptr + safe_pos[None, :] * stride_kv_t + (D_V + offs_r[:, None]), mask=valid[None, :], other=0.0 + ) + + S = tl.dot(Q_lora, K_lora_T) + tl.dot(Q_rope, K_rope_T) + S = tl.where(valid[None, :] & mask_h[:, None], S * scale, float("-inf")) + P = tl.exp(S - lse[:, None]) + P = tl.where(valid[None, :] & mask_h[:, None], P, 0.0) + dP = tl.dot(dO_val, K_lora_T) + dS = P * (dP - delta[:, None]) * scale + dS = tl.where(valid[None, :] & mask_h[:, None], dS, 0.0) + + dQ_lora += tl.dot(dS.to(tl.bfloat16), tl.trans(K_lora_T)).to(tl.float32) + dQ_rope += tl.dot(dS.to(tl.bfloat16), tl.trans(K_rope_T)).to(tl.float32) + + # Store chunk dS and P for this tile — use local head offsets (0..BLOCK_H-1) + # since ds_base already encodes hg_idx*BLOCK_H*stride_ds_h + local_h = tl.arange(0, BLOCK_H) + tl.store( + dS_ptr + ds_base + local_h[:, None] * stride_ds_h + tile_offs[None, :], + dS.to(tl.bfloat16), + mask=mask_h[:, None] & valid[None, :], + ) + tl.store( + P_ptr + ds_base + local_h[:, None] * stride_ds_h + tile_offs[None, :], + P.to(tl.bfloat16), + mask=mask_h[:, None] & valid[None, :], + ) + + tl.store( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :], + dQ_lora.to(Q_lora.dtype), + mask=mask_h[:, None], + ) + tl.store( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]), + dQ_rope.to(Q_rope.dtype), + mask=mask_h[:, None], + ) + + +@triton.jit +def _bwd_chunk_dq( + Q_ptr, # [T, H, D] bf16 + KV_ptr, # [T, 1, D] bf16 + dO_ptr, # [T, H, D_V] bf16 + TopK_ptr, # [T, TOPK] int32 + LSE_ptr, # [T, H] fp32 + Delta_ptr, # [T, H] fp32 + dQ_ptr, # [T, H, D] bf16 — read-modify-write across chunks + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_dq_t: tl.int64, + stride_dq_h: tl.int64, + stride_topk_t: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + R_START: tl.int32, + R_CHUNK: tl.constexpr, + BLOCK_H: tl.constexpr, + TILE_K: tl.constexpr, + D_V: tl.constexpr, + D_ROPE: tl.constexpr, + IS_FIRST_CHUNK: tl.constexpr, +): + """ + dQ accumulation for one rank chunk [R_START, R_START+R_CHUNK). + Grid: (total_tokens, num_hg). No writes to intermediate buffer. + IS_FIRST_CHUNK=True: initialises dQ to zero (avoids a global memory read). + """ + token_idx = tl.program_id(0) + hg_idx = tl.program_id(1) + offs_h = hg_idx * BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = offs_h < num_heads + offs_v = tl.arange(0, D_V) + offs_r = tl.arange(0, D_ROPE) + + q_base = token_idx * stride_q_t + Q_lora = tl.load( + Q_ptr + q_base + offs_h[:, None] * stride_q_h + offs_v[None, :], mask=mask_h[:, None], other=0.0 + ) + Q_rope = tl.load( + Q_ptr + q_base + offs_h[:, None] * stride_q_h + (D_V + offs_r[None, :]), + mask=mask_h[:, None], + other=0.0, + ) + do_base = token_idx * stride_do_t + dO_val = tl.load( + dO_ptr + do_base + offs_h[:, None] * stride_do_h + offs_v[None, :], mask=mask_h[:, None], other=0.0 + ) + lse = tl.load(LSE_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0) + delta = tl.load(Delta_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0) + + dq_base = token_idx * stride_dq_t + if IS_FIRST_CHUNK: + dQ_lora = tl.zeros([BLOCK_H, D_V], dtype=tl.float32) + dQ_rope = tl.zeros([BLOCK_H, D_ROPE], dtype=tl.float32) + else: + dQ_lora = tl.load( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :], + mask=mask_h[:, None], + other=0.0, + ).to(tl.float32) + dQ_rope = tl.load( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]), + mask=mask_h[:, None], + other=0.0, + ).to(tl.float32) + + NUM_TILES: tl.constexpr = (R_CHUNK + TILE_K - 1) // TILE_K + topk_base = token_idx * stride_topk_t + R_START + offs_tile = tl.arange(0, TILE_K) + topk_pos = tl.load(TopK_ptr + topk_base + offs_tile, mask=offs_tile < R_CHUNK, other=-1) + topk_pos_next = topk_pos + + for t in range(NUM_TILES): + tile_start = t * TILE_K + valid = (tile_start + offs_tile) < R_CHUNK + valid = valid & (topk_pos != -1) + if t + 1 < NUM_TILES: + next_offs = (t + 1) * TILE_K + offs_tile + topk_pos_next = tl.load(TopK_ptr + topk_base + next_offs, mask=next_offs < R_CHUNK, other=-1) + safe_pos = tl.where(valid, topk_pos, 0) + + K_lora_T = tl.load( + KV_ptr + safe_pos[None, :] * stride_kv_t + offs_v[:, None], mask=valid[None, :], other=0.0 + ) + K_rope_T = tl.load( + KV_ptr + safe_pos[None, :] * stride_kv_t + (D_V + offs_r[:, None]), mask=valid[None, :], other=0.0 + ) + + S = tl.dot(Q_lora, K_lora_T) + tl.dot(Q_rope, K_rope_T) + S = tl.where(valid[None, :] & mask_h[:, None], S * scale, float("-inf")) + P = tl.exp(S - lse[:, None]) + P = tl.where(valid[None, :] & mask_h[:, None], P, 0.0) + dP = tl.dot(dO_val, K_lora_T) + dS = P * (dP - delta[:, None]) * scale + dS = tl.where(valid[None, :] & mask_h[:, None], dS, 0.0) + + dQ_lora += tl.dot(dS.to(tl.bfloat16), tl.trans(K_lora_T)).to(tl.float32) + dQ_rope += tl.dot(dS.to(tl.bfloat16), tl.trans(K_rope_T)).to(tl.float32) + + if t + 1 < NUM_TILES: + topk_pos = topk_pos_next + + tl.store( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :], + dQ_lora.to(Q_lora.dtype), + mask=mask_h[:, None], + ) + tl.store( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]), + dQ_rope.to(Q_rope.dtype), + mask=mask_h[:, None], + ) + + +@triton.jit +def _bwd_chunk_dkv_interm( + Q_T_ptr, # [T, D_QK, H] bf16 (transposed: stride = D_QK * H) + dO_T_ptr, # [T, D_V, H] bf16 + TopK_ptr, # [T, TOPK] int32 + LSE_ptr, # [T, H] fp32 + Delta_ptr, # [T, H] fp32 + KV_ptr, # [T, 1, D] bf16 + Interm_ptr, # [T, R_CHUNK, D] bf16 — output (plain store, one writer per slot) + stride_qt_t: tl.int64, + stride_dot_t: tl.int64, + stride_topk_t: tl.int64, + stride_kv_t: tl.int64, + stride_interm_t: tl.int64, # R_CHUNK * D + stride_interm_r: tl.int64, # D + scale: tl.float32, + num_heads: tl.int32, + R_START: tl.int32, + R_CHUNK: tl.constexpr, + TILE_K: tl.constexpr, + BLOCK_H: tl.constexpr, + NUM_HG: tl.constexpr, + D_V: tl.constexpr, + D_ROPE: tl.constexpr, +): + """ + dKV intermediate for one rank chunk [R_START, R_START+R_CHUNK). + Grid: (total_tokens,) — one CTA per query token, inner loop over head groups. + Recomputes S/P/dS on-the-fly. Plain stores to bf16 interm — no atomics. + """ + token_idx = tl.program_id(0) + offs_v = tl.arange(0, D_V) + offs_r = tl.arange(0, D_ROPE) + offs_tile = tl.arange(0, TILE_K) + + NUM_TILES: tl.constexpr = (R_CHUNK + TILE_K - 1) // TILE_K + topk_base = token_idx * stride_topk_t + R_START + interm_base_t = token_idx * stride_interm_t + + qt_base = token_idx * stride_qt_t + dot_base = token_idx * stride_dot_t + + for t in range(NUM_TILES): + tile_start = t * TILE_K + tile_offs = tile_start + offs_tile + valid_tile = tile_offs < R_CHUNK + + topk_pos = tl.load(TopK_ptr + topk_base + tile_start + offs_tile, mask=valid_tile, other=-1) + valid = valid_tile & (topk_pos != -1) + safe_pos = tl.where(valid, topk_pos, 0) + + K_lora = tl.load( + KV_ptr + safe_pos[:, None] * stride_kv_t + offs_v[None, :], mask=valid[:, None], other=0.0 + ) # [TILE_K, D_V] + K_rope = tl.load( + KV_ptr + safe_pos[:, None] * stride_kv_t + (D_V + offs_r[None, :]), mask=valid[:, None], other=0.0 + ) # [TILE_K, D_ROPE] + + dKV_lora = tl.zeros([TILE_K, D_V], dtype=tl.float32) + dKV_rope = tl.zeros([TILE_K, D_ROPE], dtype=tl.float32) + + for hg in range(NUM_HG): + offs_h = hg * BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = offs_h < num_heads + + Q_lora_T = tl.load( + Q_T_ptr + qt_base + offs_v[:, None] * num_heads + offs_h[None, :], + mask=mask_h[None, :], + other=0.0, + ) # [D_V, BLOCK_H] + Q_rope_T = tl.load( + Q_T_ptr + qt_base + (D_V + offs_r[:, None]) * num_heads + offs_h[None, :], + mask=mask_h[None, :], + other=0.0, + ) # [D_ROPE, BLOCK_H] + dO_T = tl.load( + dO_T_ptr + dot_base + offs_v[:, None] * num_heads + offs_h[None, :], + mask=mask_h[None, :], + other=0.0, + ) # [D_V, BLOCK_H] + lse_h = tl.load(LSE_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0) + delta_h = tl.load(Delta_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0) + + # S = K @ Q^T [TILE_K, BLOCK_H] + S = tl.dot(K_lora, Q_lora_T) + tl.dot(K_rope, Q_rope_T) + S = tl.where(valid[:, None] & mask_h[None, :], S * scale, float("-inf")) + P = tl.exp(S - lse_h[None, :]) + P = tl.where(valid[:, None] & mask_h[None, :], P, 0.0) + dP = tl.dot(K_lora, dO_T) # [TILE_K, BLOCK_H] + dS = P * (dP - delta_h[None, :]) * scale + dS = tl.where(valid[:, None] & mask_h[None, :], dS, 0.0) + + dKV_lora += tl.dot(dS.to(tl.bfloat16), tl.trans(Q_lora_T)).to(tl.float32) + dKV_lora += tl.dot(P.to(tl.bfloat16), tl.trans(dO_T)).to(tl.float32) + dKV_rope += tl.dot(dS.to(tl.bfloat16), tl.trans(Q_rope_T)).to(tl.float32) + + # Plain store to bf16 interm — one writer per (token, local_r) slot + interm_lora_ptrs = Interm_ptr + interm_base_t + tile_offs[:, None] * stride_interm_r + offs_v[None, :] + tl.store(interm_lora_ptrs, dKV_lora.to(tl.bfloat16), mask=valid[:, None]) + + interm_rope_ptrs = ( + Interm_ptr + interm_base_t + tile_offs[:, None] * stride_interm_r + (D_V + offs_r[None, :]) + ) + tl.store(interm_rope_ptrs, dKV_rope.to(tl.bfloat16), mask=valid[:, None]) + + +@triton.jit +def _bwd_dkv_gather_acc( + Interm_ptr, # [T, R_CHUNK, D] bf16 — chunk intermediate + InvPtr_ptr, # [T+1] int32 — CSR row pointers + InvData_ptr, # [T*R_CHUNK] int32 — encoded as q*R_CHUNK + local_r + dKV_acc_ptr, # [T, D] fp32 — accumulator (read-modify-write across chunks) + stride_interm_r: tl.int64, # D + stride_acc_t: tl.int64, # D + D_V: tl.constexpr, + D_ROPE: tl.constexpr, + BLOCK_K: tl.constexpr = 64, +): + """ + Gather one chunk's bf16 intermediate into the fp32 dKV accumulator. + Grid: (total_tokens,) — one CTA per KV token k, no atomics. + + Tiled over the CSR segment in BLOCK_K-row blocks: each iteration vector-gathers + BLOCK_K interm rows ([BLOCK_K, D]) and reduces with tl.sum, instead of the old + scalar `for i in range(n_entries)` one-row-at-a-time loop (~6.5x faster, gfx1250; + memory-bound 327 -> 2130 GB/s at BLOCK_K=64). RMW semantics unchanged. + """ + k = tl.program_id(0) + offs_v = tl.arange(0, D_V) + offs_r = tl.arange(0, D_ROPE) + offs_k = tl.arange(0, BLOCK_K) + + start = tl.load(InvPtr_ptr + k) + end = tl.load(InvPtr_ptr + k + 1) + + acc_base = k.to(tl.int64) * stride_acc_t + dkv_acc_lora = tl.load(dKV_acc_ptr + acc_base + offs_v).to(tl.float32) + dkv_acc_rope = tl.load(dKV_acc_ptr + acc_base + D_V + offs_r).to(tl.float32) + + n_entries = end - start + for ti in range(0, tl.cdiv(n_entries, BLOCK_K)): + e_local = ti * BLOCK_K + offs_k + valid = e_local < n_entries + entry = tl.load(InvData_ptr + start + e_local, mask=valid, other=0).to(tl.int64) + rp = entry[:, None] * stride_interm_r + rows_v = tl.load(Interm_ptr + rp + offs_v[None, :], mask=valid[:, None], other=0.0).to(tl.float32) + rows_r = tl.load(Interm_ptr + rp + D_V + offs_r[None, :], mask=valid[:, None], other=0.0).to( + tl.float32 + ) + dkv_acc_lora += tl.sum(rows_v, axis=0) + dkv_acc_rope += tl.sum(rows_r, axis=0) + + tl.store(dKV_acc_ptr + acc_base + offs_v, dkv_acc_lora) + tl.store(dKV_acc_ptr + acc_base + D_V + offs_r, dkv_acc_rope) diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_preprocess.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_preprocess.py new file mode 100644 index 000000000..6315f6ff1 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/_dsa_bwd_preprocess.py @@ -0,0 +1,50 @@ +import triton +import triton.language as tl + + +# ===================================================================== +# Backward — preprocess kernel (Delta computation) +# ===================================================================== +@triton.jit +def _sparse_mla_bwd_preprocess( + O_ptr, # [total_tokens, num_heads, D_V] + dO_ptr, # [total_tokens, num_heads, D_V] + Delta_ptr, # [total_tokens, num_heads] + stride_o_t: tl.int64, + stride_o_h: tl.int64, + num_heads: tl.int32, + D_V: tl.constexpr, + BLOCK_H: tl.constexpr, +): + """ + Delta[t, h] = sum_d(O[t, h, d] * dO[t, h, d]) + + Grid: (total_tokens, cdiv(num_heads, BLOCK_H)) + """ + token_idx = tl.program_id(0) + hg_idx = tl.program_id(1) + + offs_h = hg_idx * BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = offs_h < num_heads + offs_d = tl.arange(0, D_V) + + base = token_idx * stride_o_t + + O = tl.load( + O_ptr + base + offs_h[:, None] * stride_o_h + offs_d[None, :], + mask=mask_h[:, None], + other=0.0, + ) + dO = tl.load( + dO_ptr + base + offs_h[:, None] * stride_o_h + offs_d[None, :], + mask=mask_h[:, None], + other=0.0, + ) + + delta = tl.sum(O.to(tl.float32) * dO.to(tl.float32), axis=1) + + tl.store( + Delta_ptr + token_idx * num_heads + offs_h, + delta, + mask=mask_h, + ) diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dkv_interm.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dkv_interm.py new file mode 100644 index 000000000..0529e0764 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dkv_interm.py @@ -0,0 +1,229 @@ +""" +Gluon dKV-intermediate backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X). + +M3 port of the Triton `_bwd_compute_dkv_intermediate`. Key gluon delta vs Triton: +the Triton path materializes a transposed Q/dO in HBM (`q.transpose(1,2).contiguous()`); +this kernel loads Q/dO UNtransposed and transposes in-LDS via `ds_read_*_tr`, removing +the external transpose copy. + +Per program: 1 query token. Grid: (total_tokens,). +Per rank-tile (loop NUM_TILES = R_CHUNK / TILE_K), summed over head groups: + dKV_lora[D_V, TILE_K] = sum_hg ( Q_lora_T @ dS + dO_T @ P ) # contract over heads + dKV_rope[D_ROPE, TILE_K] = sum_hg ( Q_rope_T @ dS ) + store interm[token, rank, :D_QK] + +Q_lora_T/dO_T/Q_rope_T ([D, BLOCK_H]) are the opIdx-0 *transposed* operands -> staged in +LDS, read transposed with ds_read_tr. dS/P ([BLOCK_H, TILE_K]) are opIdx-1 natural-layout +-> register load + convert. M1 config: BLOCK_H=64, TILE_K=64, single-buffered. +""" + +import torch +import triton +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +@gluon.jit +def _sparse_mla_bwd_dkv_interm_gl_kernel( + Q_ptr, # [T, H, D_QK] bf16 (UNtransposed) + dO_ptr, # [T, H, D_V] bf16 (UNtransposed) + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + Interm_ptr, # [T, R_CHUNK, D_QK] bf16 + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + stride_interm_t: tl.int64, + stride_interm_r: tl.int64, + num_heads: tl.int32, + R_CHUNK: gl.constexpr, + TILE_K: gl.constexpr, + BLOCK_H: gl.constexpr, + NUM_HG: gl.constexpr, + D_V: gl.constexpr, + D_ROPE: gl.constexpr, +): + # ===================== constexpr layouts ===================== + # MMA output is [D_V, TILE_K] (and [D_ROPE, TILE_K]); contraction over BLOCK_H heads. + mfma: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + + # ---- Blocked layouts for HBM loads ---- + # Q/dO [BLOCK_H, D_V] : load coalesced then stage to LDS for transpose-read. + _q_tpw_k: gl.constexpr = min(64, D_V // 8) + _q_tpw_m: gl.constexpr = 64 // _q_tpw_k + blk_q: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[_q_tpw_m, _q_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_qrope: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D_ROPE] + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[4, 1], + order=[1, 0], + ) + # dS / P [BLOCK_H, TILE_K] : opIdx-1, register load + convert (no transpose). + blk_ds: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 4], + threads_per_warp=[16, 4], + warps_per_cta=[4, 1], + order=[1, 0], + ) + + # ---- Shared layouts (Q/dO/Q_rope staged for transpose read) ---- + sh_q: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0]) + sh_do: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0]) + sh_qrope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[1, 0]) + + # ---- Dot operand layouts ---- + dot_qT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_doT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_qropeT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_ds_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8) + dot_p_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8) + + token_idx = gl.program_id(axis=0) + NUM_TILES: gl.constexpr = R_CHUNK // TILE_K + + # ---- LDS for Q/dO/Q_rope (single-buffered) ---- + smem_q = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_q) + smem_do = gl.allocate_shared_memory(dO_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_do) + smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope) + + q_base = token_idx.to(tl.int64) * stride_q_t + do_base = token_idx.to(tl.int64) * stride_do_t + ds_base = token_idx.to(tl.int64) * stride_ds_t + interm_base = token_idx.to(tl.int64) * stride_interm_t + + # store offsets (mfma layout): dKV[d, col] -> interm[token, t*TILE_K+col, d] + offs_d_st = gl.arange(0, D_V, layout=gl.SliceLayout(1, mfma)) + offs_col_st = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma)) + offs_dr_st = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, mfma)) + + for t in range(NUM_TILES): + dKV_lora = gl.zeros([D_V, TILE_K], dtype=gl.float32, layout=mfma) + dKV_rope = gl.zeros([D_ROPE, TILE_K], dtype=gl.float32, layout=mfma) + + for hg in range(NUM_HG): + hg_off = hg * BLOCK_H + + # ---- stage Q/dO/Q_rope (this head group) into LDS ---- + offs_h_q = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_q)) + offs_v_q = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_q)) + mask_h_q = offs_h_q < num_heads + q_offs = q_base + offs_h_q[:, None].to(tl.int64) * stride_q_h + offs_v_q[None, :].to(tl.int64) + do_offs = do_base + offs_h_q[:, None].to(tl.int64) * stride_do_h + offs_v_q[None, :].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_q, ptr=Q_ptr, offsets=q_offs.to(tl.int32), mask=mask_h_q[:, None] + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_do, ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=mask_h_q[:, None] + ) + + offs_h_qr = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_qr = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_qr = offs_h_qr < num_heads + qr_offs = ( + q_base + + offs_h_qr[:, None].to(tl.int64) * stride_q_h + + (D_V + offs_r_qr[None, :]).to(tl.int64) + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qrope, ptr=Q_ptr, offsets=qr_offs.to(tl.int32), mask=mask_h_qr[:, None] + ) + gl.amd.cdna4.async_copy.commit_group() + + # ---- load dS / P (this tile, this head group) -> dot operands ---- + offs_h_ds = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_ds)) + offs_col_ds = t * TILE_K + gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_ds)) + mask_h_ds = offs_h_ds < num_heads + dsp_offs = ( + ds_base + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + offs_col_ds[None, :].to(tl.int64) + ) + dS_blk = gl.amd.cdna4.buffer_load( + ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + P_blk = gl.amd.cdna4.buffer_load( + ptr=P_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + dS_dot = gl.convert_layout(dS_blk, dot_ds_b) + P_dot = gl.convert_layout(P_blk, dot_p_b) + + # ---- wait + transpose-read Q/dO/Q_rope ---- + gl.amd.cdna4.async_copy.wait_group(0) + Q_T = smem_q.permute([1, 0]).load(dot_qT_a) # [D_V, BLOCK_H] + dO_T = smem_do.permute([1, 0]).load(dot_doT_a) # [D_V, BLOCK_H] + Q_rope_T = smem_qrope.permute([1, 0]).load(dot_qropeT_a) # [D_ROPE, BLOCK_H] + + dKV_lora = gl.amd.cdna4.mfma(Q_T, dS_dot, dKV_lora) + dKV_lora = gl.amd.cdna4.mfma(dO_T, P_dot, dKV_lora) + dKV_rope = gl.amd.cdna4.mfma(Q_rope_T, dS_dot, dKV_rope) + + # ---- store interm[token, t*TILE_K : +TILE_K, :] (direct from mfma layout) ---- + col = t * TILE_K + offs_col_st + interm_lora_offs = ( + interm_base + col[None, :].to(tl.int64) * stride_interm_r + offs_d_st[:, None].to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dKV_lora.to(Interm_ptr.dtype.element_ty), + ptr=Interm_ptr, + offsets=interm_lora_offs.to(tl.int32), + ) + interm_rope_offs = ( + interm_base + + col[None, :].to(tl.int64) * stride_interm_r + + (D_V + offs_dr_st[:, None]).to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dKV_rope.to(Interm_ptr.dtype.element_ty), + ptr=Interm_ptr, + offsets=interm_rope_offs.to(tl.int32), + ) + + +def sparse_mla_bwd_dkv_interm_gl(q, do, chunk_dS, chunk_P, R_CHUNK, kv_lora_rank=512, BLOCK_H=32, TILE_K=64): + """ + Gluon dKV-intermediate for one chunk. Takes UNtransposed q/do (transposes in-kernel). + + Returns interm [T, R_CHUNK, D_QK] bf16. + """ + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + assert R_CHUNK % TILE_K == 0 + num_hg = triton.cdiv(num_heads, BLOCK_H) + interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device) + + _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + num_heads, + R_CHUNK=R_CHUNK, + TILE_K=TILE_K, + BLOCK_H=BLOCK_H, + NUM_HG=num_hg, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + num_warps=4, + ) + return interm diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dq.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dq.py new file mode 100644 index 000000000..11b61b4e5 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_dq.py @@ -0,0 +1,502 @@ +""" +Gluon dQ backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X). + +M1 port of the Triton `_bwd_chunk_dq_store_ds_v4` (V4 chunked_gather dQ) onto the +gluon hardware-control structure of Leon's V3.2 forward. + +Per program: 1 query token x BLOCK_H heads x one rank chunk [R_START, R_START+R_CHUNK). +Grid: (total_tokens, cdiv(num_heads, BLOCK_H)). Dispatch loops chunks externally, +dQ is read-modify-written across chunks (IS_FIRST_CHUNK zero-inits). + +Per-tile math (TILE_K wide, looping NUM_TILES = R_CHUNK / TILE_K): + S = Q_lora @ K_lora_T + Q_rope @ K_rope_T # 2 MMAs (mfma_s), contract over D + P = exp(S*scale - lse) # lse is sink-inclusive (from fwd) + dP = dO @ K_lora_T # 1 MMA (mfma_s), reuses K_lora_T_dot + dS = P * (dP - delta) * scale + dQ_lora += dS @ K_lora # 1 MMA (mfma_acc), K_lora = K_lora_T.T view + dQ_rope += dS @ K_rope # 1 MMA (mfma_acc), K_rope = K_rope_T.T view + store dS, P chunk -> HBM (consumed by dKV-intermediate kernel) + +Differences vs Leon's fwd (the structural template): + * dO added as a second stationary [BH, D_V] operand (async-loaded with Q). + * No online softmax: single P = exp(S - lse) (lse precomputed by fwd). + * 5 MMAs/tile vs 3; K_lora read 3 ways (S, dP, dQ_lora), K_rope 2 ways. + * Per-tile dS/P stores + final dQ store (RMW across chunks) replace the O/LSE write. + * Sink: d_sink is NOT done here (handled by a torch reduction in the launcher), + so this kernel needs no atomics. lse from fwd already folds the sink in. + +M1 config: BLOCK_H=32, TILE_K=16 (matches Triton TILE_K_DQ and the 16x16x16 MFMA +k-dim). LDS at BH=32/TILE_K=16 ~= 104 KB < 160 KB (gfx950). +""" + +import torch +import triton +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +@gluon.jit +def _sparse_mla_bwd_dq_gl_kernel( + Q_ptr, # [T, H, D_QK] bf16 + KV_ptr, # [T, 1, D_QK] bf16 + dO_ptr, # [T, H, D_V] bf16 + TopK_ptr, # [T, TOPK_padded] int32 + LSE_ptr, # [T, H] fp32 (sink-inclusive) + Delta_ptr, # [T, H] fp32 + dQ_ptr, # [T, H, D_QK] bf16 (read-modify-write across chunks) + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_dq_t: tl.int64, + stride_dq_h: tl.int64, + stride_topk_t: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + R_START: tl.int32, + R_CHUNK: gl.constexpr, + BLOCK_H: gl.constexpr, + TILE_K: gl.constexpr, + D_V: gl.constexpr, + D_ROPE: gl.constexpr, + IS_FIRST_CHUNK: gl.constexpr, +): + # ===================== constexpr layouts ===================== + mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + + # ---- Blocked layouts for global loads (per Leon's fwd) ---- + _qlora_tpw_k: gl.constexpr = min(64, D_V // 8) + _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k + blk_qlora: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D_V] (Q_lora, dO, dQ_lora) + size_per_thread=[1, 8], + threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_qrope: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D_ROPE] + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[4, 1], + order=[1, 0], + ) + + _klora_tpw_m: gl.constexpr = min(64, D_V // 8) + _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m + blk_klora: gl.constexpr = gl.BlockedLayout( # [D_V, TILE_K] + size_per_thread=[8, 1], + threads_per_warp=[_klora_tpw_m, _klora_tpw_n], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_krope: gl.constexpr = gl.BlockedLayout( # [D_ROPE, TILE_K] + size_per_thread=[2, 1], + threads_per_warp=[32, 2], + warps_per_cta=[1, 4], + order=[0, 1], + ) + # ---- Shared layouts (only K is staged in LDS) ---- + sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [D_V, TILE_K], + [0, 1], + ) + sh_krope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[0, 1]) + + # ---- Dot operand layouts ---- + dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_do_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_ds_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4) + dot_klora_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4) + dot_krope_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4) + + # ===================== program ids ===================== + token_idx = gl.program_id(axis=0) + hg_idx = gl.program_id(axis=1) + hg_offset = hg_idx * BLOCK_H + + NUM_TILES: gl.constexpr = R_CHUNK // TILE_K + + # ===================== Q / dO offsets ===================== + offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_qlora = offs_h_qlora < num_heads + q_base = token_idx.to(tl.int64) * stride_q_t + q_offs_lora = ( + q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64) + ) + q_mask_lora = mask_h_qlora[:, None] + + offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_qrope = offs_h_qrope < num_heads + q_offs_rope = ( + q_base + offs_h_qrope[:, None].to(tl.int64) * stride_q_h + (D_V + offs_r_qrope[None, :]).to(tl.int64) + ) + q_mask_rope = mask_h_qrope[:, None] + + offs_h_do = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_do = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_do = offs_h_do < num_heads + do_base = token_idx.to(tl.int64) * stride_do_t + do_offs = do_base + offs_h_do[:, None].to(tl.int64) * stride_do_h + offs_v_do[None, :].to(tl.int64) + do_mask = mask_h_do[:, None] + + # ===================== load Q_lora, Q_rope, dO -> registers (no LDS staging) ===================== + # Stationary opIdx-0 operands: load HBM->VGPR (blocked, coalesced) then convert to + # the dot-operand layout once. The convert's LDS scratch is transient (freed before + # the K loop), unlike persistent staging, so only K occupies LDS during the loop. + q_lora_blk = gl.amd.cdna4.buffer_load( + ptr=Q_ptr, offsets=q_offs_lora.to(tl.int32), mask=q_mask_lora, other=0.0 + ) + q_rope_blk = gl.amd.cdna4.buffer_load( + ptr=Q_ptr, offsets=q_offs_rope.to(tl.int32), mask=q_mask_rope, other=0.0 + ) + do_blk = gl.amd.cdna4.buffer_load(ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=do_mask, other=0.0) + Q_lora_dot = gl.convert_layout(q_lora_blk, dot_qlora_a) + Q_rope_dot = gl.convert_layout(q_rope_blk, dot_qrope_a) + dO_dot = gl.convert_layout(do_blk, dot_do_a) + + # ===================== topk / KV offsets ===================== + topk_base = token_idx.to(tl.int64) * stride_topk_t + R_START + + offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora)) + offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope)) + offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s)) + offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora)) + offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope)) + + # ===================== shared mem for K loop (double-buffered) ===================== + smem_krope = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_ROPE, TILE_K], layout=sh_krope) + smem_klora = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_V, TILE_K], layout=sh_klora) + + # ===================== dQ accumulators ===================== + # Always zero-init; the read-modify-write across chunks is folded in at STORE + # time in the blocked layout (avoids a big blocked->mfma_acc convert_layout + # whose LDS scratch would overflow the K/Q/dO buffers). + dQ_lora = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc) + dQ_rope = gl.zeros([BLOCK_H, D_ROPE], dtype=gl.float32, layout=mfma_acc) + + # ===================== lse / delta (in mfma_s row-slice layout) ===================== + offs_h_s = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + mask_h_s = offs_h_s < num_heads + lse = gl.amd.cdna4.buffer_load( + ptr=LSE_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0 + ) + delta = gl.amd.cdna4.buffer_load( + ptr=Delta_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0 + ) + + # ===================== prologue: K tile 0 (group B, buffer 0) ===================== + topk_pos_klora = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + offs_tile_klora).to(tl.int32), + mask=offs_tile_klora < R_CHUNK, + other=-1, + ) + topk_pos_krope = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + offs_tile_krope).to(tl.int32), + mask=offs_tile_krope < R_CHUNK, + other=-1, + ) + topk_pos_mma = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, offsets=(topk_base + offs_tile_mma).to(tl.int32), mask=offs_tile_mma < R_CHUNK, other=-1 + ) + + valid_klora = topk_pos_klora != -1 + valid_krope = topk_pos_krope != -1 + valid_mma = topk_pos_mma != -1 + safe_klora = gl.where(valid_klora, topk_pos_klora, 0) + safe_krope = gl.where(valid_krope, topk_pos_krope, 0) + + klora_offs = safe_klora[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(0), ptr=KV_ptr, offsets=klora_offs.to(tl.int32), mask=valid_klora[None, :] + ) + krope_offs = safe_krope[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(0), ptr=KV_ptr, offsets=krope_offs.to(tl.int32), mask=valid_krope[None, :] + ) + gl.amd.cdna4.async_copy.commit_group() + + # dS / P store offsets in the mfma_s layout (store directly from the compute + # layout -> no convert_layout/LDS shuffle; less-coalesced HBM write instead). + ds_base = token_idx.to(tl.int64) * stride_ds_t + hg_idx.to(tl.int64) * BLOCK_H * stride_ds_h + offs_h_dsp = gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + offs_tile_dsp = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s)) + mask_h_dsp = (hg_offset + offs_h_dsp) < num_heads + + # ===================== main loop: prefetch t+1, compute t ===================== + cur_buf = 0 + for t in range(NUM_TILES - 1): + next_offs_klora = (t + 1) * TILE_K + offs_tile_klora + next_offs_krope = (t + 1) * TILE_K + offs_tile_krope + next_offs_mma = (t + 1) * TILE_K + offs_tile_mma + + topk_pos_klora_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_klora).to(tl.int32), + mask=next_offs_klora < R_CHUNK, + other=-1, + ) + topk_pos_krope_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_krope).to(tl.int32), + mask=next_offs_krope < R_CHUNK, + other=-1, + ) + topk_pos_mma_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_mma).to(tl.int32), + mask=next_offs_mma < R_CHUNK, + other=-1, + ) + + valid_klora_next = (next_offs_klora < R_CHUNK) & (topk_pos_klora_next != -1) + valid_krope_next = (next_offs_krope < R_CHUNK) & (topk_pos_krope_next != -1) + valid_mma_next = (next_offs_mma < R_CHUNK) & (topk_pos_mma_next != -1) + safe_klora_next = gl.where(valid_klora_next, topk_pos_klora_next, 0) + safe_krope_next = gl.where(valid_krope_next, topk_pos_krope_next, 0) + + next_buf = 1 - cur_buf + klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to( + tl.int64 + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(next_buf), + ptr=KV_ptr, + offsets=klora_offs_next.to(tl.int32), + mask=valid_klora_next[None, :], + ) + krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + ( + D_V + offs_r_krope[:, None] + ).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(next_buf), + ptr=KV_ptr, + offsets=krope_offs_next.to(tl.int32), + mask=valid_krope_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + + gl.amd.cdna4.async_copy.wait_group(1) + + # ----- read K views from current buffer ----- + klora_smem_cur = smem_klora.index(cur_buf) + K_lora_T_dot = klora_smem_cur.load(dot_klora_b) # [D_V, TILE_K] opIdx1 mfma_s + K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b) # [TILE_K, D_V] opIdx1 mfma_acc + krope_smem_cur = smem_krope.index(cur_buf) + K_rope_T_dot = krope_smem_cur.load(dot_krope_b) + K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b) + + # ----- S = Q_lora@K_lora_T + Q_rope@K_rope_T ----- + S = gl.amd.cdna4.mfma( + Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + ) + S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S) + S = S * scale + offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None] + S = gl.where(valid_mask, S, float("-inf")) + + # ----- P = exp(S - lse) ; dP = dO@K_lora_T ; dS = P*(dP-delta)*scale ----- + P = gl.exp(S - lse[:, None]) + P = gl.where(valid_mask, P, 0.0) + dP = gl.amd.cdna4.mfma( + dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + ) + dS = P * (dP - delta[:, None]) * scale + dS = gl.where(valid_mask, dS, 0.0) + + # ----- dQ_lora += dS@K_lora ; dQ_rope += dS@K_rope ----- + dS_bf = dS.to(KV_ptr.dtype.element_ty) + dS_dot = gl.convert_layout(dS_bf, dot_ds_a) + dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora) + dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope) + + # ----- store dS, P chunk ----- + col = t * TILE_K + offs_tile_dsp + dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64) + gl.amd.cdna4.buffer_store( + stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None] + ) + gl.amd.cdna4.buffer_store( + stored_value=P.to(KV_ptr.dtype.element_ty), + ptr=P_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + + # promote prefetch -> current + cur_buf = next_buf + valid_mma = valid_mma_next + + # ===================== epilogue: last tile ===================== + gl.amd.cdna4.async_copy.wait_group(0) + t = NUM_TILES - 1 + klora_smem_cur = smem_klora.index(cur_buf) + K_lora_T_dot = klora_smem_cur.load(dot_klora_b) + K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b) + krope_smem_cur = smem_krope.index(cur_buf) + K_rope_T_dot = krope_smem_cur.load(dot_krope_b) + K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b) + + S = gl.amd.cdna4.mfma( + Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + ) + S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S) + S = S * scale + offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None] + S = gl.where(valid_mask, S, float("-inf")) + + P = gl.exp(S - lse[:, None]) + P = gl.where(valid_mask, P, 0.0) + dP = gl.amd.cdna4.mfma(dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)) + dS = P * (dP - delta[:, None]) * scale + dS = gl.where(valid_mask, dS, 0.0) + + dS_bf = dS.to(KV_ptr.dtype.element_ty) + dS_dot = gl.convert_layout(dS_bf, dot_ds_a) + dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora) + dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope) + + col = t * TILE_K + offs_tile_dsp + dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64) + gl.amd.cdna4.buffer_store( + stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None] + ) + gl.amd.cdna4.buffer_store( + stored_value=P.to(KV_ptr.dtype.element_ty), + ptr=P_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + + # ===================== store dQ (lora + rope) ===================== + dq_base = token_idx.to(tl.int64) * stride_dq_t + offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_o = offs_h_o < num_heads + dq_offs_lora = dq_base + offs_h_o[:, None].to(tl.int64) * stride_dq_h + offs_v_o[None, :].to(tl.int64) + dq_lora_blk = gl.convert_layout(dQ_lora.to(dQ_ptr.dtype.element_ty), blk_qlora) + if not IS_FIRST_CHUNK: + prev_lora = gl.amd.cdna4.buffer_load( + ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None], other=0.0 + ) + dq_lora_blk = (dq_lora_blk.to(gl.float32) + prev_lora.to(gl.float32)).to(dQ_ptr.dtype.element_ty) + gl.amd.cdna4.buffer_store( + stored_value=dq_lora_blk, ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None] + ) + + offs_h_or = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_or = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_or = offs_h_or < num_heads + dq_offs_rope = ( + dq_base + offs_h_or[:, None].to(tl.int64) * stride_dq_h + (D_V + offs_r_or[None, :]).to(tl.int64) + ) + dq_rope_blk = gl.convert_layout(dQ_rope.to(dQ_ptr.dtype.element_ty), blk_qrope) + if not IS_FIRST_CHUNK: + prev_rope = gl.amd.cdna4.buffer_load( + ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None], other=0.0 + ) + dq_rope_blk = (dq_rope_blk.to(gl.float32) + prev_rope.to(gl.float32)).to(dQ_ptr.dtype.element_ty) + gl.amd.cdna4.buffer_store( + stored_value=dq_rope_blk, ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None] + ) + + +# ===================================================================== +# Launcher — runs the dQ pass only (chunk loop + RMW), returns dq, chunk dS/P. +# d_sink (if needed) is a torch reduction handled by the caller. +# ===================================================================== +def sparse_mla_bwd_dq_gl( + q, + kv, + do, + topk_indices_padded, + lse, + delta, + R_CHUNK, + topk, + kv_lora_rank=512, + scale=None, + BLOCK_H=64, + TILE_K=16, +): + """ + Gluon dQ pass. Mirrors the dQ portion of `sparse_mla_bwd_v4`'s chunk loop. + + Returns: + dq: [T, H, D_QK] bf16 (fully accumulated across chunks) + chunk_dS: [T, H, R_CHUNK] bf16 (LAST chunk's dS — for spot validation) + chunk_P: [T, H, R_CHUNK] bf16 (LAST chunk's P) + """ + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + if scale is None: + scale = 1.0 / (d_qk**0.5) + assert R_CHUNK % TILE_K == 0, "TILE_K must divide R_CHUNK" + + dq = torch.empty_like(q) + chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + + num_hg = triton.cdiv(num_heads, BLOCK_H) + grid = (total_tokens, num_hg) + + for r_start in range(0, topk, R_CHUNK): + is_first = r_start == 0 + _sparse_mla_bwd_dq_gl_kernel[grid]( + q, + kv, + do, + topk_indices_padded, + lse, + delta, + dq, + chunk_dS, + chunk_P, + q.stride(0), + q.stride(1), + kv.stride(0), + do.stride(0), + do.stride(1), + dq.stride(0), + dq.stride(1), + topk_indices_padded.stride(0), + chunk_dS.stride(0), + chunk_dS.stride(1), + scale, + num_heads, + r_start, + R_CHUNK=R_CHUNK, + BLOCK_H=BLOCK_H, + TILE_K=TILE_K, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + IS_FIRST_CHUNK=is_first, + num_warps=4, + waves_per_eu=1, + ) + return dq, chunk_dS, chunk_P diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_v4_gluon.py new file mode 100644 index 000000000..4be3b9a17 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_bwd_v4_gluon.py @@ -0,0 +1,166 @@ +""" +End-to-end gluon backward for DeepSeek V4 sparse MLA (M4). + +Wires: Triton preprocess (Delta) + gluon dQ + gluon dKV-intermediate + +Triton gather (CSR inverted-topk) + torch d_sink reduction. + +Drop-in for `sparse_mla_bwd_v4` (chunked_gather). Differences vs the Triton path: + - dQ kernel = gluon `_sparse_mla_bwd_dq_gl_kernel` (BH=64, TK=16) + - dKV-interm kernel = gluon `_sparse_mla_bwd_dkv_interm_gl_kernel` (BH=32, TK=64), + takes UNtransposed q/do -> the `q.transpose(1,2).contiguous()` copies are GONE. + - d_sink computed in torch (no in-kernel atomics). + - gather + preprocess unchanged (Triton). +""" + +import torch +import triton + +from ._dsa_bwd_gather import _build_inverted_topk_slice, _bwd_dkv_gather_acc +from ._dsa_bwd_preprocess import _sparse_mla_bwd_preprocess +from .dsa_bwd_dkv_interm import _sparse_mla_bwd_dkv_interm_gl_kernel +from .dsa_bwd_dq import _sparse_mla_bwd_dq_gl_kernel + + +def sparse_mla_bwd_v4_gluon(q, kv, o, do, topk_indices, lse, attn_sink=None, kv_lora_rank=512, scale=None): + assert q.is_contiguous() and kv.is_contiguous() and o.is_contiguous() + assert do.is_contiguous() and topk_indices.is_contiguous() and lse.is_contiguous() + + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + topk = topk_indices.shape[1] + if scale is None: + scale = 1.0 / (d_qk**0.5) + if kv.dim() == 2: + kv = kv.unsqueeze(1) + # num_kv may exceed total_tokens (V4 [local ++ compressed-pool] buffer): + # dKV is accumulated per KV-token, so size it by kv rows, not query tokens. + num_kv = kv.shape[0] + + has_sink = attn_sink is not None + if has_sink: + assert attn_sink.dtype == torch.float32 and attn_sink.shape == (num_heads,) + + # ---- preprocess: Delta = rowsum(O*dO) (Triton, unchanged) ---- + delta = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device) + BLOCK_H_PRE = triton.next_power_of_2(min(64, num_heads)) + _sparse_mla_bwd_preprocess[(total_tokens, triton.cdiv(num_heads, BLOCK_H_PRE))]( + O_ptr=o, + dO_ptr=do, + Delta_ptr=delta, + stride_o_t=o.stride(0), + stride_o_h=o.stride(1), + num_heads=num_heads, + D_V=kv_lora_rank, + BLOCK_H=BLOCK_H_PRE, + ) + + # ---- config ---- + R_CHUNK = min(256, topk) + BH_DQ, TK_DQ = 64, 16 + BH_DKV, TK_DKV = 32, 64 + num_hg_dq = triton.cdiv(num_heads, BH_DQ) + num_hg_dkv = triton.cdiv(num_heads, BH_DKV) + + dq = torch.empty_like(q) + chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + dkv_acc = torch.zeros(num_kv, d_qk, dtype=torch.float32, device=q.device) + interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device) + + # ---- pad topk to R_CHUNK multiple ---- + topk_padded_len = ((topk + R_CHUNK - 1) // R_CHUNK) * R_CHUNK + if topk_padded_len != topk: + pad = torch.full((total_tokens, topk_padded_len - topk), -1, dtype=torch.int32, device=q.device) + topk_padded = torch.cat([topk_indices, pad], dim=1).contiguous() + else: + topk_padded = topk_indices + + all_csr = [ + _build_inverted_topk_slice(topk_padded[:, rs : rs + R_CHUNK], rs, R_CHUNK, num_kv=num_kv) + for rs in range(0, topk, R_CHUNK) + ] + + for chunk_idx, r_start in enumerate(range(0, topk, R_CHUNK)): + is_first = r_start == 0 + + # gluon dQ (writes dq RMW, chunk_dS, chunk_P) + _sparse_mla_bwd_dq_gl_kernel[(total_tokens, num_hg_dq)]( + q, + kv, + do, + topk_padded, + lse, + delta, + dq, + chunk_dS, + chunk_P, + q.stride(0), + q.stride(1), + kv.stride(0), + do.stride(0), + do.stride(1), + dq.stride(0), + dq.stride(1), + topk_padded.stride(0), + chunk_dS.stride(0), + chunk_dS.stride(1), + scale, + num_heads, + r_start, + R_CHUNK=R_CHUNK, + BLOCK_H=BH_DQ, + TILE_K=TK_DQ, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + IS_FIRST_CHUNK=is_first, + num_warps=4, + waves_per_eu=1, # sweep: 1 is best (+3-7%); >=2 spills catastrophically + ) + + # gluon dKV-intermediate (untransposed q/do -> no external transpose) + _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + num_heads, + R_CHUNK=R_CHUNK, + TILE_K=TK_DKV, + BLOCK_H=BH_DKV, + NUM_HG=num_hg_dkv, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + num_warps=4, + ) + + # Triton gather (CSR inverted-topk reduce interm -> dkv_acc). + # Grid is over KV tokens (num_kv), which may exceed query tokens. + inv_ptr, inv_data = all_csr[chunk_idx] + _bwd_dkv_gather_acc[(num_kv,)]( + interm, + inv_ptr, + inv_data, + dkv_acc, + interm.stride(1), + dkv_acc.stride(0), + D_V=kv_lora_rank, + D_ROPE=rope_rank, + num_warps=4, + ) + + # ---- d_sink in torch: -sum_t exp(sink - lse) * delta ---- + d_sink = None + if has_sink: + d_sink = -(torch.exp(attn_sink.unsqueeze(0) - lse) * delta).sum(0) + + dkv_out = dkv_acc.to(kv.dtype).unsqueeze(1) + return dq, dkv_out, d_sink diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_fwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_fwd_v4_gluon.py new file mode 100644 index 000000000..5cc778e19 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_dsa/dsa_fwd_v4_gluon.py @@ -0,0 +1,625 @@ +""" +Gluon forward for DeepSeek V4 sparse MLA (gfx950 / CDNA4), with attention sink. + +Based on Leon's (leonling-ll) V3.2 gluon forward from `leonling-ll/aiter` branch +`liyang/dsa` -- which adapted our V3.2 Triton forward and added the gfx950 hardware +control (MFMA4 layouts, padded/swizzled shared, double-buffered K, async DMA pipeline, +ds_read_tr transpose, explicit dot-operand layouts); see also his DSA PR ROCm/aiter#3456. +This file adds the V4 attention-sink epilogue (sink-inclusive LSE) so it matches +`sparse_mla_fwd_v4`, and is exposed via `sparse_mla_fwd_v4(..., backend="gluon")`. + +Pipeline (Leon's): + Prologue: Q_lora + Q_rope -> shared via async DMA (group A). + K_lora tile 0 + K_rope tile 0 -> shared via async DMA (group B, double-buffered). + wait_group(1) -> Q in shared; load Q dot operands once. + Loop tile t (0..N-2): prefetch K[t+1]; wait_group(1); read K_lora/V_lora(permute)/K_rope; + S = Q_lora @ K_lora_T + Q_rope @ K_rope_T; online softmax; acc += P @ V_lora. + Epilogue: wait_group(0) -> last tile; fold sink into the denominator (V4); write O, LSE. +""" + +import torch +import triton +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +# ===================================================================== +# Utility +# ===================================================================== +def _get_lds_limit(): + """Return the per-CU LDS limit in bytes for the current GPU. + + gfx942 (MI300X): 64 KB = 65536 bytes + gfx950 (MI355X): 160 KB = 163840 bytes + """ + if torch.cuda.is_available(): + prop = torch.cuda.get_device_properties(0) + gcn_arch = getattr(prop, "gcnArchName", "") + if "gfx950" in gcn_arch: + return 163840 + return 65536 + + +_LDS_LIMIT = _get_lds_limit() + + +# ===================================================================== +# Forward — autotune configs and pruning +# ===================================================================== +def _fwd_prune_configs(configs, named_args, **kwargs): + """Prune autotune configs that would exceed per-CU LDS.""" + D_V = kwargs.get("D_V", named_args.get("D_V")) + D_ROPE = kwargs.get("D_ROPE", named_args.get("D_ROPE")) + pruned = [] + for config in configs: + config.kwargs["BLOCK_H"] + tk = config.kwargs["TILE_K"] + ns = config.num_stages + kv_lds = (D_V + D_ROPE) * tk * 2 * ns + if kv_lds <= _LDS_LIMIT: + pruned.append(config) + if not pruned: + pruned.append(configs[0]) + return pruned + + +def _get_fwd_autotune_configs(): + configs = [ + triton.Config( + {"BLOCK_H": BLOCK_H, "TILE_K": TILE_K, "waves_per_eu": WPE}, + num_warps=nw, + ) + for BLOCK_H in [16, 32, 64] + for TILE_K in [16, 32, 64, 128] + for WPE in [0, 1, 2] + for nw in [4] # num_warps must be 4 to align with kernel implementation + ] + # configs = [triton.Config({"BLOCK_H": 64, "TILE_K": 32, "waves_per_eu": 0}, num_warps=4),] + return configs + + +@triton.autotune( + configs=_get_fwd_autotune_configs(), + key=["num_heads", "TOPK", "D_V", "D_ROPE"], + prune_configs_by={"early_config_prune": _fwd_prune_configs}, +) +@gluon.jit +def _sparse_mla_fwd_gl_kernel( + Q_ptr, # [total_tokens, num_heads, D_QK] bf16 + KV_ptr, # [total_tokens, 1, D_QK] bf16 + TopK_ptr, # [total_tokens, TOPK] int32 + Sink_ptr, # [num_heads] fp32; ignored if HAS_SINK == False + O_ptr, # [total_tokens, num_heads, D_V] bf16 + LSE_ptr, # [total_tokens, num_heads] fp32 (sink-inclusive if HAS_SINK) + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_o_t: tl.int64, + stride_o_h: tl.int64, + stride_topk_t: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + TOPK: gl.constexpr, + BLOCK_H: gl.constexpr, + TILE_K: gl.constexpr, + D_V: gl.constexpr, + D_ROPE: gl.constexpr, + HAS_SINK: gl.constexpr, +): + # ---------- constexpr layouts ---------- + mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + + # Blocked layouts for global loads. + _qlora_tpw_k: gl.constexpr = min(64, D_V // 8) + _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k + blk_qlora: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_qrope: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[4, 1], + order=[1, 0], + ) + + _klora_tpw_m: gl.constexpr = min(64, D_V // 8) + _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m + blk_klora: gl.constexpr = gl.BlockedLayout( # [D_V, TILE_K] + size_per_thread=[8, 1], + threads_per_warp=[_klora_tpw_m, _klora_tpw_n], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_krope: gl.constexpr = gl.BlockedLayout( # [D_ROPE, TILE_K] = [64, 16] + size_per_thread=[2, 1], + threads_per_warp=[32, 2], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_topk: gl.constexpr = gl.BlockedLayout( # [TILE_K] int32 + size_per_thread=[1], + threads_per_warp=[64], + warps_per_cta=[4], + order=[0], + ) + blk_lse: gl.constexpr = gl.BlockedLayout( # [BLOCK_H] fp32 + size_per_thread=[1], + threads_per_warp=[64], + warps_per_cta=[4], + order=[0], + ) + + # Shared layouts. + sh_qlora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [BLOCK_H, D_V], + [1, 0], + ) + sh_qrope: gl.constexpr = gl.SwizzledSharedLayout( + vec=8, + per_phase=2, + max_phase=8, + order=[1, 0], + ) + sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [D_V, TILE_K], + [0, 1], + ) + sh_krope: gl.constexpr = gl.SwizzledSharedLayout( + vec=8, + per_phase=2, + max_phase=8, + order=[0, 1], + ) + + # Dot operand layouts + dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_p_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4) + dot_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4) + + # ---------- program ids ---------- + token_idx = gl.program_id(axis=0) + hg_idx = gl.program_id(axis=1) + hg_offset = hg_idx * BLOCK_H + + # ---------- offsets for Q ---------- + # Q_lora [BLOCK_H, D_V] + offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_qlora = offs_h_qlora < num_heads + + q_base = token_idx.to(tl.int64) * stride_q_t + q_offs_lora = ( + q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64) + ) + q_mask_lora = mask_h_qlora[:, None] + + # Q_rope [BLOCK_H, D_ROPE] + offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_qrope = offs_h_qrope < num_heads + + q_offs_rope = ( + q_base + offs_h_qrope[:, None].to(tl.int64) * stride_q_h + (D_V + offs_r_qrope[None, :]).to(tl.int64) + ) + q_mask_rope = mask_h_qrope[:, None] + + smem_qlora = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_qlora) + smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qlora, + ptr=Q_ptr, + offsets=q_offs_lora.to(tl.int32), + mask=q_mask_lora, + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qrope, + ptr=Q_ptr, + offsets=q_offs_rope.to(tl.int32), + mask=q_mask_rope, + ) + gl.amd.cdna4.async_copy.commit_group() + + # ---------- topk and KV offsets ---------- + NUM_TILES: gl.constexpr = (TOPK + TILE_K - 1) // TILE_K + topk_base = token_idx.to(tl.int64) * stride_topk_t + + # offs_tile in three layouts (sliced from each of the three loaders) + offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora)) + offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope)) + offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s)) + offs_tile_topk = gl.arange(0, TILE_K, layout=blk_topk) + + offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora)) + offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope)) + + # (removed dead `topk_pos_reg` prologue load — was never consumed) + + # ---------- shared mem allocations for the K loop ---------- + smem_krope = gl.allocate_shared_memory( + KV_ptr.dtype.element_ty, + [2, D_ROPE, TILE_K], + layout=sh_krope, + ) + smem_klora = gl.allocate_shared_memory( + KV_ptr.dtype.element_ty, + [2, D_V, TILE_K], + layout=sh_klora, + ) + + # ---------- accumulators ---------- + m_i = gl.full([BLOCK_H], float("-inf"), dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s)) + l_i = gl.full([BLOCK_H], 0.0, dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s)) + acc = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc) + + # ---------- tile-0 prefetch (prologue) ---------- + # Load K_lora and K_rope for tile 0. + topk_pos_klora = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + offs_tile_klora, + mask=offs_tile_klora < TOPK, + other=-1, + ) + topk_pos_krope = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + offs_tile_krope, + mask=offs_tile_krope < TOPK, + other=-1, + ) + topk_pos_mma = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + offs_tile_mma, + mask=offs_tile_mma < TOPK, + other=-1, + ) + + # Deep-prefetch tile-1 topk ONCE in the neutral blk_topk layout (DEDUP). Carried as a + # single register set; converted to the klora/krope/mma layouts at point of use, to + # minimize carried register pressure (the 3-layout carry caused an acc-rescale codegen + # regression -- see att_fwd_gluon_mi350/RESULTS.md). + p1_off_topk = TILE_K + offs_tile_topk + tkraw = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + p1_off_topk, + mask=p1_off_topk < TOPK, + other=-1, + ) + + valid_klora = topk_pos_klora != -1 # tile_start=0 -> offs_tile buf1, drain K[0], QK[0] -> S_prev (no softmax/PV yet). + tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora)) + tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope)) + tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s)) + valid_klora_next = ((TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1) + valid_krope_next = ((TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1) + valid_qk = ((TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1) + safe_klora_next = gl.where(valid_klora_next, tk_klora, 0) + safe_krope_next = gl.where(valid_krope_next, tk_krope, 0) + klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(1), + ptr=KV_ptr, + offsets=klora_offs_next.to(tl.int32), + mask=valid_klora_next[None, :], + ) + krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to( + tl.int64 + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(1), + ptr=KV_ptr, + offsets=krope_offs_next.to(tl.int32), + mask=valid_krope_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + tkraw = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + (2 * TILE_K + offs_tile_topk), + mask=(2 * TILE_K + offs_tile_topk) < TOPK, + other=-1, + ) + gl.amd.cdna4.async_copy.wait_group(1) + S_prev = gl.amd.cdna4.mfma( + Q_lora_dot, + smem_klora.index(0).load(dot_klora_b), + gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s), + ) + S_prev = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(0).load(dot_krope_b), S_prev) + S_prev = S_prev * scale + S_prev = gl.where(valid_mma[None, :] & mask_h_mma[:, None], S_prev, float("-inf")) + cur_buf = 1 + + for t in range(NUM_TILES - 2): + gl.amd.cdna4.async_copy.wait_group(0) # drain K[t+1] (cur_buf) before QK reads it + # 2-BUFFER EARLY-GATHER: evacuate V[t] from pv_buf into REGISTERS first, freeing that buffer, + # then gather tile t+2 into it BEFORE the QK/PV MFMAs so the DMA overlaps both (no 3rd buffer). + # V_lora_dot in regs => no read/async-write race on the recycled buffer. Costs VGPR live range. + V_lora_dot = smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b) + tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora)) + tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope)) + tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s)) + valid_klora_next = (((t + 2) * TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1) + valid_krope_next = (((t + 2) * TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1) + valid_qk_next = (((t + 2) * TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1) + safe_klora_next = gl.where(valid_klora_next, tk_klora, 0) + safe_krope_next = gl.where(valid_krope_next, tk_krope, 0) + klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to( + tl.int64 + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(1 - cur_buf), + ptr=KV_ptr, + offsets=klora_offs_next.to(tl.int32), + mask=valid_klora_next[None, :], + ) + krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + ( + D_V + offs_r_krope[:, None] + ).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(1 - cur_buf), + ptr=KV_ptr, + offsets=krope_offs_next.to(tl.int32), + mask=valid_krope_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + tkraw_n = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + ((t + 3) * TILE_K + offs_tile_topk), + mask=((t + 3) * TILE_K + offs_tile_topk) < TOPK, + other=-1, + ) + # QK tile (t+1) from cur_buf -- matrix; overlaps the gather above + softmax below + S_cur = gl.amd.cdna4.mfma( + Q_lora_dot, + smem_klora.index(cur_buf).load(dot_klora_b), + gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s), + ) + S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur) + S_cur = S_cur * scale + S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf")) + # softmax(S_prev = tile t) [VALU, overlaps QK] + m_j = gl.max(S_prev, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > float("-inf"), m_new, 0.0) + alpha = gl.exp(m_i - m_new) + P = gl.exp(S_prev - m_new[:, None]) + l_i = alpha * l_i + gl.sum(P, axis=1) + m_i = m_new + # PV tile t from registers -- matrix; overlaps the gather still in flight + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(P_dot, V_lora_dot, acc) + # promote + S_prev = S_cur + valid_qk = valid_qk_next + tkraw = tkraw_n + cur_buf = 1 - cur_buf + + # ---------- PRE-DRAIN: QK[N-1] (cur_buf) || softmax+PV[N-2] (pv_buf); no gather ---------- + gl.amd.cdna4.async_copy.wait_group(0) # drain K[N-1] (last loop gather) + S_cur = gl.amd.cdna4.mfma( + Q_lora_dot, + smem_klora.index(cur_buf).load(dot_klora_b), + gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s), + ) + S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur) + S_cur = S_cur * scale + S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf")) + m_j = gl.max(S_prev, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > float("-inf"), m_new, 0.0) + alpha = gl.exp(m_i - m_new) + P = gl.exp(S_prev - m_new[:, None]) + l_i = alpha * l_i + gl.sum(P, axis=1) + m_i = m_new + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b), acc) + S_prev = S_cur + + # ---------- DRAIN: softmax+PV[N-1] (S_prev = QK[N-1], V from cur_buf) ---------- + m_j = gl.max(S_prev, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > float("-inf"), m_new, 0.0) + alpha = gl.exp(m_i - m_new) + P = gl.exp(S_prev - m_new[:, None]) + l_new = alpha * l_i + gl.sum(P, axis=1) + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(cur_buf).permute([1, 0]).load(dot_v_b), acc) + m_i = m_new + l_i = l_new + + # ---------- epilogue: fold sink into the denominator (V4 delta) ---------- + if HAS_SINK: + offs_h_sink = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + sink = gl.amd.cdna4.buffer_load( + ptr=Sink_ptr, + offsets=offs_h_sink.to(tl.int32), + mask=offs_h_sink < num_heads, + other=float("-inf"), + ) + m_final = gl.maximum(m_i, sink) + alpha_fix = gl.exp(m_i - m_final) + l_total = l_i * alpha_fix + gl.exp(sink - m_final) + alpha_fix_acc = gl.convert_layout(alpha_fix, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_fix_acc[:, None] + l_total_acc = gl.convert_layout(l_total, gl.SliceLayout(1, mfma_acc)) + acc = acc / l_total_acc[:, None] + lse = m_final + gl.log(l_total) + else: + l_i_acc = gl.convert_layout(l_i, gl.SliceLayout(1, mfma_acc)) + acc = acc / l_i_acc[:, None] + lse = m_i + gl.log(l_i) + + # Output O[token_idx, h, v] + offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_o = offs_h_o < num_heads + o_base = token_idx.to(tl.int64) * stride_o_t + o_offs = o_base + offs_h_o[:, None].to(tl.int64) * stride_o_h + offs_v_o[None, :].to(tl.int64) + acc_bf = acc.to(O_ptr.dtype.element_ty) + acc_bf_blk = gl.convert_layout(acc_bf, blk_qlora) + gl.amd.cdna4.buffer_store( + stored_value=acc_bf_blk, + ptr=O_ptr, + offsets=o_offs.to(tl.int32), + mask=mask_h_o[:, None], + ) + + # LSE[token_idx, h] + offs_h_lse = hg_offset + gl.arange(0, BLOCK_H, layout=blk_lse) + mask_h_lse = offs_h_lse < num_heads + lse_base = token_idx * num_heads + lse_offs = lse_base + offs_h_lse + lse_blk = gl.convert_layout(lse, blk_lse) + gl.amd.cdna4.buffer_store( + stored_value=lse_blk, + ptr=LSE_ptr, + offsets=lse_offs.to(tl.int32), + mask=mask_h_lse, + ) + + +# ===================================================================== +# Launcher +# ===================================================================== +def sparse_mla_fwd_v4_gluon(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None): + """ + DeepSeek V4 sparse MLA forward (Gluon, gfx950 / CDNA4), with attention sink. + + Args: + q: [total_tokens, num_heads, d_qk] bfloat16 + kv: [total_tokens, 1, d_qk] bfloat16 (or [total_tokens, d_qk]) + topk_indices: [total_tokens, topk] int32 (SWA + sparse, -1 marks invalid) + attn_sink: [num_heads] fp32, optional per-head learnable sink logit. + When None, behaves like the V3.2 forward. + kv_lora_rank: int, default 512 + scale: float, default 1/sqrt(d_qk) + + Returns: + o: [total_tokens, num_heads, kv_lora_rank] same dtype as q + lse: [total_tokens, num_heads] float32 (sink-inclusive when attn_sink is given) + """ + assert q.is_contiguous() + assert kv.is_contiguous() + assert topk_indices.is_contiguous() + + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + topk = topk_indices.shape[1] + + if scale is None: + scale = 1.0 / (d_qk**0.5) + + if kv.dim() == 2: + kv = kv.unsqueeze(1) + # kv may hold MORE rows than there are query tokens (V4 feeds a + # [local ++ compressed-pool] buffer, so num_kv = S + P > total_tokens). + # The kernel only dereferences kv via topk indices (stride_kv_t), so any + # num_kv >= max(topk_index)+1 is valid. + assert kv.shape[0] >= total_tokens and kv.shape[-1] == d_qk + + has_sink = attn_sink is not None + if has_sink: + assert attn_sink.is_contiguous() + assert attn_sink.dtype == torch.float32 + assert attn_sink.shape == (num_heads,) + sink_ptr = attn_sink + else: + sink_ptr = torch.empty(1, dtype=torch.float32, device=q.device) # guarded by HAS_SINK + + o = torch.empty(total_tokens, num_heads, kv_lora_rank, dtype=q.dtype, device=q.device) + lse = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device) + + # Grid is autotune-aware: BLOCK_H comes from the chosen config. + grid = lambda META: (total_tokens, triton.cdiv(num_heads, META["BLOCK_H"])) + + _sparse_mla_fwd_gl_kernel[grid]( + Q_ptr=q, + KV_ptr=kv, + TopK_ptr=topk_indices, + Sink_ptr=sink_ptr, + O_ptr=o, + LSE_ptr=lse, + stride_q_t=q.stride(0), + stride_q_h=q.stride(1), + stride_kv_t=kv.stride(0), + stride_o_t=o.stride(0), + stride_o_h=o.stride(1), + stride_topk_t=topk_indices.stride(0), + scale=scale, + num_heads=num_heads, + TOPK=topk, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_SINK=has_sink, + ) + + return o, lse diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/__init__.py new file mode 100644 index 000000000..e3c606370 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/__init__.py @@ -0,0 +1,31 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Gluon DeepSeek-V4 sparse-MLA attention backend ("gluon v2"). + +Second-generation Gluon (gfx950 / CDNA4) sparse-MLA backend, using the same fused +single-latent (K == V) sparse-MLA representation and public API as the other V4 +backends. Both directions are explicit-layout Gluon kernels: + +* forward (``dsa_fwd_v4_gluon``): MFMA layouts, padded/swizzled shared, async + double-buffered pipeline, rope-skip, exp2 softmax, MFMA K=32. +* backward (``dsa_bwd_v4_gluon``): Gluon dQ + dKV-intermediate kernels (rope-skip, + MFMA K=32, single-chunk dQ RMW) + Triton Delta preprocess + CSR inverted-topk gather. + +The Gluon kernels need a Gluon-capable (recompiled) triton whose CDNA4 async_copy +accepts general offset layouts, and raise a clear install hint otherwise. + +* :func:`sparse_mla_fwd_v4_gluon_v2` -> ``(o, lse)`` +* :func:`sparse_mla_bwd_v4_gluon_v2` -> ``(dq, dkv, d_sink)`` +""" + +from .dsa_bwd_v4_gluon import sparse_mla_bwd_v4_gluon_v2 +from .dsa_fwd_v4_gluon import sparse_mla_fwd_v4_gluon_v2 + +__all__ = [ + "sparse_mla_fwd_v4_gluon_v2", + "sparse_mla_bwd_v4_gluon_v2", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dkv_interm_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dkv_interm_gluon.py new file mode 100644 index 000000000..dc2098b2a --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dkv_interm_gluon.py @@ -0,0 +1,237 @@ +""" +Gluon dKV-intermediate backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X). + +M3 port of the Triton `_bwd_compute_dkv_intermediate`. Key gluon delta vs Triton: +the Triton path materializes a transposed Q/dO in HBM (`q.transpose(1,2).contiguous()`); +this kernel loads Q/dO UNtransposed and transposes in-LDS via `ds_read_*_tr`, removing +the external transpose copy. + +Per program: 1 query token. Grid: (total_tokens,). +Per rank-tile (loop NUM_TILES = R_CHUNK / TILE_K), summed over head groups: + dKV_lora[D_V, TILE_K] = sum_hg ( Q_lora_T @ dS + dO_T @ P ) # contract over heads + dKV_rope[D_ROPE, TILE_K] = sum_hg ( Q_rope_T @ dS ) + store interm[token, rank, :D_QK] + +Q_lora_T/dO_T/Q_rope_T ([D, BLOCK_H]) are the opIdx-0 *transposed* operands -> staged in +LDS, read transposed with ds_read_tr. dS/P ([BLOCK_H, TILE_K]) are opIdx-1 natural-layout +-> register load + convert. M1 config: BLOCK_H=64, TILE_K=64, single-buffered. +""" + +import torch +import triton +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +@gluon.jit +def _sparse_mla_bwd_dkv_interm_gl_kernel( + Q_ptr, # [T, H, D_QK] bf16 (UNtransposed) + dO_ptr, # [T, H, D_V] bf16 (UNtransposed) + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + Interm_ptr, # [T, R_CHUNK, D_QK] bf16 + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + stride_interm_t: tl.int64, + stride_interm_r: tl.int64, + num_heads: tl.int32, + R_CHUNK: gl.constexpr, + TILE_K: gl.constexpr, + BLOCK_H: gl.constexpr, + NUM_HG: gl.constexpr, + D_V: gl.constexpr, + D_ROPE: gl.constexpr, + HAS_ROPE: gl.constexpr, +): + # ===================== constexpr layouts ===================== + # MMA output is [D_V, TILE_K] (and [D_ROPE, TILE_K]); contraction over BLOCK_H heads. + mfma: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + + # ---- Blocked layouts for HBM loads ---- + # Q/dO [BLOCK_H, D_V] : load coalesced then stage to LDS for transpose-read. + _q_tpw_k: gl.constexpr = min(64, D_V // 8) + _q_tpw_m: gl.constexpr = 64 // _q_tpw_k + blk_q: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[_q_tpw_m, _q_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_qrope: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D_ROPE] + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[4, 1], + order=[1, 0], + ) + # dS / P [BLOCK_H, TILE_K] : opIdx-1, register load + convert (no transpose). + blk_ds: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 4], + threads_per_warp=[16, 4], + warps_per_cta=[4, 1], + order=[1, 0], + ) + + # ---- Shared layouts (Q/dO/Q_rope staged for transpose read) ---- + sh_q: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0]) + sh_do: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0]) + sh_qrope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[1, 0]) + + # ---- Dot operand layouts ---- + dot_qT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_doT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_qropeT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_ds_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8) + dot_p_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8) + + token_idx = gl.program_id(axis=0) + NUM_TILES: gl.constexpr = R_CHUNK // TILE_K + + # ---- LDS for Q/dO/Q_rope (single-buffered) ---- + smem_q = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_q) + smem_do = gl.allocate_shared_memory(dO_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_do) + if HAS_ROPE: + smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope) + + q_base = token_idx.to(tl.int64) * stride_q_t + do_base = token_idx.to(tl.int64) * stride_do_t + ds_base = token_idx.to(tl.int64) * stride_ds_t + interm_base = token_idx.to(tl.int64) * stride_interm_t + + # store offsets (mfma layout): dKV[d, col] -> interm[token, t*TILE_K+col, d] + offs_d_st = gl.arange(0, D_V, layout=gl.SliceLayout(1, mfma)) + offs_col_st = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma)) + offs_dr_st = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, mfma)) + + for t in range(NUM_TILES): + dKV_lora = gl.zeros([D_V, TILE_K], dtype=gl.float32, layout=mfma) + if HAS_ROPE: + dKV_rope = gl.zeros([D_ROPE, TILE_K], dtype=gl.float32, layout=mfma) + + for hg in range(NUM_HG): + hg_off = hg * BLOCK_H + + # ---- stage Q/dO/Q_rope (this head group) into LDS ---- + offs_h_q = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_q)) + offs_v_q = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_q)) + mask_h_q = offs_h_q < num_heads + q_offs = q_base + offs_h_q[:, None].to(tl.int64) * stride_q_h + offs_v_q[None, :].to(tl.int64) + do_offs = do_base + offs_h_q[:, None].to(tl.int64) * stride_do_h + offs_v_q[None, :].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_q, ptr=Q_ptr, offsets=q_offs.to(tl.int32), mask=mask_h_q[:, None] + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_do, ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=mask_h_q[:, None] + ) + + if HAS_ROPE: + offs_h_qr = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_qr = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_qr = offs_h_qr < num_heads + qr_offs = ( + q_base + + offs_h_qr[:, None].to(tl.int64) * stride_q_h + + (D_V + offs_r_qr[None, :]).to(tl.int64) + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qrope, ptr=Q_ptr, offsets=qr_offs.to(tl.int32), mask=mask_h_qr[:, None] + ) + gl.amd.cdna4.async_copy.commit_group() + + # ---- load dS / P (this tile, this head group) -> dot operands ---- + offs_h_ds = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_ds)) + offs_col_ds = t * TILE_K + gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_ds)) + mask_h_ds = offs_h_ds < num_heads + dsp_offs = ( + ds_base + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + offs_col_ds[None, :].to(tl.int64) + ) + dS_blk = gl.amd.cdna4.buffer_load( + ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + P_blk = gl.amd.cdna4.buffer_load( + ptr=P_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + dS_dot = gl.convert_layout(dS_blk, dot_ds_b) + P_dot = gl.convert_layout(P_blk, dot_p_b) + + # ---- wait + transpose-read Q/dO/Q_rope ---- + gl.amd.cdna4.async_copy.wait_group(0) + Q_T = smem_q.permute([1, 0]).load(dot_qT_a) # [D_V, BLOCK_H] + dO_T = smem_do.permute([1, 0]).load(dot_doT_a) # [D_V, BLOCK_H] + + dKV_lora = gl.amd.cdna4.mfma(Q_T, dS_dot, dKV_lora) + dKV_lora = gl.amd.cdna4.mfma(dO_T, P_dot, dKV_lora) + if HAS_ROPE: + Q_rope_T = smem_qrope.permute([1, 0]).load(dot_qropeT_a) # [D_ROPE, BLOCK_H] + dKV_rope = gl.amd.cdna4.mfma(Q_rope_T, dS_dot, dKV_rope) + + # ---- store interm[token, t*TILE_K : +TILE_K, :] (direct from mfma layout) ---- + col = t * TILE_K + offs_col_st + interm_lora_offs = ( + interm_base + col[None, :].to(tl.int64) * stride_interm_r + offs_d_st[:, None].to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dKV_lora.to(Interm_ptr.dtype.element_ty), + ptr=Interm_ptr, + offsets=interm_lora_offs.to(tl.int32), + ) + # dKV_rope is provably zero for the V4 zero-rope-pad (discarded downstream by the + # gather/adapter, which use dkv[..., :D_V]); skip its compute + store when HAS_ROPE=False. + if HAS_ROPE: + interm_rope_offs = ( + interm_base + + col[None, :].to(tl.int64) * stride_interm_r + + (D_V + offs_dr_st[:, None]).to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dKV_rope.to(Interm_ptr.dtype.element_ty), + ptr=Interm_ptr, + offsets=interm_rope_offs.to(tl.int32), + ) + + +def sparse_mla_bwd_dkv_interm_gl(q, do, chunk_dS, chunk_P, R_CHUNK, kv_lora_rank=512, BLOCK_H=32, TILE_K=64): + """ + Gluon dKV-intermediate for one chunk. Takes UNtransposed q/do (transposes in-kernel). + + Returns interm [T, R_CHUNK, D_QK] bf16. + """ + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + assert R_CHUNK % TILE_K == 0 + num_hg = triton.cdiv(num_heads, BLOCK_H) + interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device) + + _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + num_heads, + R_CHUNK=R_CHUNK, + TILE_K=TILE_K, + BLOCK_H=BLOCK_H, + NUM_HG=num_hg, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + num_warps=4, + ) + return interm diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dq_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dq_gluon.py new file mode 100644 index 000000000..78e6f58e1 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_dq_gluon.py @@ -0,0 +1,523 @@ +""" +Gluon dQ backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X). + +M1 port of the Triton `_bwd_chunk_dq_store_ds_v4` (V4 chunked_gather dQ) onto the +gluon hardware-control structure of Leon's V3.2 forward. + +Per program: 1 query token x BLOCK_H heads x one rank chunk [R_START, R_START+R_CHUNK). +Grid: (total_tokens, cdiv(num_heads, BLOCK_H)). Dispatch loops chunks externally, +dQ is read-modify-written across chunks (IS_FIRST_CHUNK zero-inits). + +Per-tile math (TILE_K wide, looping NUM_TILES = R_CHUNK / TILE_K): + S = Q_lora @ K_lora_T + Q_rope @ K_rope_T # 2 MMAs (mfma_s), contract over D + P = exp(S*scale - lse) # lse is sink-inclusive (from fwd) + dP = dO @ K_lora_T # 1 MMA (mfma_s), reuses K_lora_T_dot + dS = P * (dP - delta) * scale + dQ_lora += dS @ K_lora # 1 MMA (mfma_acc), K_lora = K_lora_T.T view + dQ_rope += dS @ K_rope # 1 MMA (mfma_acc), K_rope = K_rope_T.T view + store dS, P chunk -> HBM (consumed by dKV-intermediate kernel) + +Differences vs Leon's fwd (the structural template): + * dO added as a second stationary [BH, D_V] operand (async-loaded with Q). + * No online softmax: single P = exp(S - lse) (lse precomputed by fwd). + * 5 MMAs/tile vs 3; K_lora read 3 ways (S, dP, dQ_lora), K_rope 2 ways. + * Per-tile dS/P stores + final dQ store (RMW across chunks) replace the O/LSE write. + * Sink: d_sink is NOT done here (handled by a torch reduction in the launcher), + so this kernel needs no atomics. lse from fwd already folds the sink in. + +M1 config: BLOCK_H=32, TILE_K=16 (matches Triton TILE_K_DQ and the 16x16x16 MFMA +k-dim). LDS at BH=32/TILE_K=16 ~= 104 KB < 160 KB (gfx950). +""" + +import torch +import triton +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +@gluon.jit +def _sparse_mla_bwd_dq_gl_kernel( + Q_ptr, # [T, H, D_QK] bf16 + KV_ptr, # [T, 1, D_QK] bf16 + dO_ptr, # [T, H, D_V] bf16 + TopK_ptr, # [T, TOPK_padded] int32 + LSE_ptr, # [T, H] fp32 (sink-inclusive) + Delta_ptr, # [T, H] fp32 + dQ_ptr, # [T, H, D_QK] bf16 (read-modify-write across chunks) + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_dq_t: tl.int64, + stride_dq_h: tl.int64, + stride_topk_t: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + R_START: tl.int32, + R_CHUNK: gl.constexpr, + BLOCK_H: gl.constexpr, + TILE_K: gl.constexpr, + D_V: gl.constexpr, + D_ROPE: gl.constexpr, + HAS_ROPE: gl.constexpr, + IS_FIRST_CHUNK: gl.constexpr, +): + # ===================== constexpr layouts ===================== + # mfma_s drives S = Q@K_T and dP = dO@K_T, both reducing D_V=512 -> K=32 halves their + # MFMA instruction count vs K=16 (mirrors the fwd R15 win). mfma_acc (dQ += dS@K) stays + # K=16 since its reduction dim is TILE_K (=16). + mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 32], + transposed=True, + warps_per_cta=[4, 1], + ) + mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + + # ---- Blocked layouts for global loads (per Leon's fwd) ---- + _qlora_tpw_k: gl.constexpr = min(64, D_V // 8) + _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k + blk_qlora: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D_V] (Q_lora, dO, dQ_lora) + size_per_thread=[1, 8], + threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_qrope: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D_ROPE] + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[4, 1], + order=[1, 0], + ) + + _klora_tpw_m: gl.constexpr = min(64, D_V // 8) + _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m + blk_klora: gl.constexpr = gl.BlockedLayout( # [D_V, TILE_K] + size_per_thread=[8, 1], + threads_per_warp=[_klora_tpw_m, _klora_tpw_n], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_krope: gl.constexpr = gl.BlockedLayout( # [D_ROPE, TILE_K] + size_per_thread=[2, 1], + threads_per_warp=[32, 2], + warps_per_cta=[1, 4], + order=[0, 1], + ) + # ---- Shared layouts (only K is staged in LDS) ---- + sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [D_V, TILE_K], + [0, 1], + ) + sh_krope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[0, 1]) + + # ---- Dot operand layouts ---- + dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_do_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_ds_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4) + dot_klora_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4) + dot_krope_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4) + + # ===================== program ids ===================== + token_idx = gl.program_id(axis=0) + hg_idx = gl.program_id(axis=1) + hg_offset = hg_idx * BLOCK_H + + NUM_TILES: gl.constexpr = R_CHUNK // TILE_K + + # ===================== Q / dO offsets ===================== + offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_qlora = offs_h_qlora < num_heads + q_base = token_idx.to(tl.int64) * stride_q_t + q_offs_lora = ( + q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64) + ) + q_mask_lora = mask_h_qlora[:, None] + + if HAS_ROPE: + offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_qrope = offs_h_qrope < num_heads + q_offs_rope = ( + q_base + + offs_h_qrope[:, None].to(tl.int64) * stride_q_h + + (D_V + offs_r_qrope[None, :]).to(tl.int64) + ) + q_mask_rope = mask_h_qrope[:, None] + + offs_h_do = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_do = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_do = offs_h_do < num_heads + do_base = token_idx.to(tl.int64) * stride_do_t + do_offs = do_base + offs_h_do[:, None].to(tl.int64) * stride_do_h + offs_v_do[None, :].to(tl.int64) + do_mask = mask_h_do[:, None] + + # ===================== load Q_lora, Q_rope, dO -> registers (no LDS staging) ===================== + # Stationary opIdx-0 operands: load HBM->VGPR (blocked, coalesced) then convert to + # the dot-operand layout once. The convert's LDS scratch is transient (freed before + # the K loop), unlike persistent staging, so only K occupies LDS during the loop. + q_lora_blk = gl.amd.cdna4.buffer_load( + ptr=Q_ptr, offsets=q_offs_lora.to(tl.int32), mask=q_mask_lora, other=0.0 + ) + do_blk = gl.amd.cdna4.buffer_load(ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=do_mask, other=0.0) + Q_lora_dot = gl.convert_layout(q_lora_blk, dot_qlora_a) + dO_dot = gl.convert_layout(do_blk, dot_do_a) + if HAS_ROPE: + q_rope_blk = gl.amd.cdna4.buffer_load( + ptr=Q_ptr, offsets=q_offs_rope.to(tl.int32), mask=q_mask_rope, other=0.0 + ) + Q_rope_dot = gl.convert_layout(q_rope_blk, dot_qrope_a) + + # ===================== topk / KV offsets ===================== + topk_base = token_idx.to(tl.int64) * stride_topk_t + R_START + + offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora)) + offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope)) + offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s)) + offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora)) + offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope)) + + # ===================== shared mem for K loop (double-buffered) ===================== + if HAS_ROPE: + smem_krope = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_ROPE, TILE_K], layout=sh_krope) + smem_klora = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_V, TILE_K], layout=sh_klora) + + # ===================== dQ accumulators ===================== + # Always zero-init; the read-modify-write across chunks is folded in at STORE + # time in the blocked layout (avoids a big blocked->mfma_acc convert_layout + # whose LDS scratch would overflow the K/Q/dO buffers). + dQ_lora = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc) + if HAS_ROPE: + dQ_rope = gl.zeros([BLOCK_H, D_ROPE], dtype=gl.float32, layout=mfma_acc) + + # ===================== lse / delta (in mfma_s row-slice layout) ===================== + offs_h_s = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + mask_h_s = offs_h_s < num_heads + lse = gl.amd.cdna4.buffer_load( + ptr=LSE_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0 + ) + delta = gl.amd.cdna4.buffer_load( + ptr=Delta_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0 + ) + + # ===================== prologue: K tile 0 (group B, buffer 0) ===================== + topk_pos_klora = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + offs_tile_klora).to(tl.int32), + mask=offs_tile_klora < R_CHUNK, + other=-1, + ) + topk_pos_mma = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, offsets=(topk_base + offs_tile_mma).to(tl.int32), mask=offs_tile_mma < R_CHUNK, other=-1 + ) + + valid_klora = topk_pos_klora != -1 + valid_mma = topk_pos_mma != -1 + safe_klora = gl.where(valid_klora, topk_pos_klora, 0) + + klora_offs = safe_klora[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(0), ptr=KV_ptr, offsets=klora_offs.to(tl.int32), mask=valid_klora[None, :] + ) + if HAS_ROPE: + topk_pos_krope = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + offs_tile_krope).to(tl.int32), + mask=offs_tile_krope < R_CHUNK, + other=-1, + ) + valid_krope = topk_pos_krope != -1 + safe_krope = gl.where(valid_krope, topk_pos_krope, 0) + krope_offs = safe_krope[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to( + tl.int64 + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(0), ptr=KV_ptr, offsets=krope_offs.to(tl.int32), mask=valid_krope[None, :] + ) + gl.amd.cdna4.async_copy.commit_group() + + # dS / P store offsets in the mfma_s layout (store directly from the compute + # layout -> no convert_layout/LDS shuffle; less-coalesced HBM write instead). + ds_base = token_idx.to(tl.int64) * stride_ds_t + hg_idx.to(tl.int64) * BLOCK_H * stride_ds_h + offs_h_dsp = gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + offs_tile_dsp = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s)) + mask_h_dsp = (hg_offset + offs_h_dsp) < num_heads + + # ===================== main loop: prefetch t+1, compute t ===================== + cur_buf = 0 + for t in range(NUM_TILES - 1): + next_offs_klora = (t + 1) * TILE_K + offs_tile_klora + next_offs_krope = (t + 1) * TILE_K + offs_tile_krope + next_offs_mma = (t + 1) * TILE_K + offs_tile_mma + + topk_pos_klora_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_klora).to(tl.int32), + mask=next_offs_klora < R_CHUNK, + other=-1, + ) + topk_pos_mma_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_mma).to(tl.int32), + mask=next_offs_mma < R_CHUNK, + other=-1, + ) + + valid_klora_next = (next_offs_klora < R_CHUNK) & (topk_pos_klora_next != -1) + valid_mma_next = (next_offs_mma < R_CHUNK) & (topk_pos_mma_next != -1) + safe_klora_next = gl.where(valid_klora_next, topk_pos_klora_next, 0) + + next_buf = 1 - cur_buf + klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to( + tl.int64 + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(next_buf), + ptr=KV_ptr, + offsets=klora_offs_next.to(tl.int32), + mask=valid_klora_next[None, :], + ) + if HAS_ROPE: + topk_pos_krope_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_krope).to(tl.int32), + mask=next_offs_krope < R_CHUNK, + other=-1, + ) + valid_krope_next = (next_offs_krope < R_CHUNK) & (topk_pos_krope_next != -1) + safe_krope_next = gl.where(valid_krope_next, topk_pos_krope_next, 0) + krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + ( + D_V + offs_r_krope[:, None] + ).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(next_buf), + ptr=KV_ptr, + offsets=krope_offs_next.to(tl.int32), + mask=valid_krope_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + + gl.amd.cdna4.async_copy.wait_group(1) + + # ----- read K views from current buffer ----- + klora_smem_cur = smem_klora.index(cur_buf) + K_lora_T_dot = klora_smem_cur.load(dot_klora_b) # [D_V, TILE_K] opIdx1 mfma_s + K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b) # [TILE_K, D_V] opIdx1 mfma_acc + + # ----- S = Q_lora@K_lora_T (+ Q_rope@K_rope_T when HAS_ROPE) ----- + S = gl.amd.cdna4.mfma( + Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + ) + if HAS_ROPE: + krope_smem_cur = smem_krope.index(cur_buf) + K_rope_T_dot = krope_smem_cur.load(dot_krope_b) + S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S) + S = S * scale + offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None] + S = gl.where(valid_mask, S, float("-inf")) + + # ----- P = exp(S - lse) ; dP = dO@K_lora_T ; dS = P*(dP-delta)*scale ----- + P = gl.exp(S - lse[:, None]) + P = gl.where(valid_mask, P, 0.0) + dP = gl.amd.cdna4.mfma( + dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + ) + dS = P * (dP - delta[:, None]) * scale + dS = gl.where(valid_mask, dS, 0.0) + + # ----- dQ_lora += dS@K_lora (+ dQ_rope += dS@K_rope when HAS_ROPE) ----- + dS_bf = dS.to(KV_ptr.dtype.element_ty) + dS_dot = gl.convert_layout(dS_bf, dot_ds_a) + dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora) + if HAS_ROPE: + K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b) + dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope) + + # ----- store dS, P chunk ----- + col = t * TILE_K + offs_tile_dsp + dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64) + gl.amd.cdna4.buffer_store( + stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None] + ) + gl.amd.cdna4.buffer_store( + stored_value=P.to(KV_ptr.dtype.element_ty), + ptr=P_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + + # promote prefetch -> current + cur_buf = next_buf + valid_mma = valid_mma_next + + # ===================== epilogue: last tile ===================== + gl.amd.cdna4.async_copy.wait_group(0) + t = NUM_TILES - 1 + klora_smem_cur = smem_klora.index(cur_buf) + K_lora_T_dot = klora_smem_cur.load(dot_klora_b) + K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b) + + S = gl.amd.cdna4.mfma( + Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + ) + if HAS_ROPE: + krope_smem_cur = smem_krope.index(cur_buf) + K_rope_T_dot = krope_smem_cur.load(dot_krope_b) + S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S) + S = S * scale + offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None] + S = gl.where(valid_mask, S, float("-inf")) + + P = gl.exp(S - lse[:, None]) + P = gl.where(valid_mask, P, 0.0) + dP = gl.amd.cdna4.mfma(dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)) + dS = P * (dP - delta[:, None]) * scale + dS = gl.where(valid_mask, dS, 0.0) + + dS_bf = dS.to(KV_ptr.dtype.element_ty) + dS_dot = gl.convert_layout(dS_bf, dot_ds_a) + dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora) + if HAS_ROPE: + K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b) + dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope) + + col = t * TILE_K + offs_tile_dsp + dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64) + gl.amd.cdna4.buffer_store( + stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None] + ) + gl.amd.cdna4.buffer_store( + stored_value=P.to(KV_ptr.dtype.element_ty), + ptr=P_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + + # ===================== store dQ (lora + rope) ===================== + dq_base = token_idx.to(tl.int64) * stride_dq_t + offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_o = offs_h_o < num_heads + dq_offs_lora = dq_base + offs_h_o[:, None].to(tl.int64) * stride_dq_h + offs_v_o[None, :].to(tl.int64) + dq_lora_blk = gl.convert_layout(dQ_lora.to(dQ_ptr.dtype.element_ty), blk_qlora) + if not IS_FIRST_CHUNK: + prev_lora = gl.amd.cdna4.buffer_load( + ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None], other=0.0 + ) + dq_lora_blk = (dq_lora_blk.to(gl.float32) + prev_lora.to(gl.float32)).to(dQ_ptr.dtype.element_ty) + gl.amd.cdna4.buffer_store( + stored_value=dq_lora_blk, ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None] + ) + + # dQ_rope is provably zero for the V4 zero-rope-pad (the adapter discards dq[..., D_V:]), + # so skip its accumulation + store entirely when HAS_ROPE is False. + if HAS_ROPE: + offs_h_or = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_or = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_or = offs_h_or < num_heads + dq_offs_rope = ( + dq_base + offs_h_or[:, None].to(tl.int64) * stride_dq_h + (D_V + offs_r_or[None, :]).to(tl.int64) + ) + dq_rope_blk = gl.convert_layout(dQ_rope.to(dQ_ptr.dtype.element_ty), blk_qrope) + if not IS_FIRST_CHUNK: + prev_rope = gl.amd.cdna4.buffer_load( + ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None], other=0.0 + ) + dq_rope_blk = (dq_rope_blk.to(gl.float32) + prev_rope.to(gl.float32)).to(dQ_ptr.dtype.element_ty) + gl.amd.cdna4.buffer_store( + stored_value=dq_rope_blk, ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None] + ) + + +# ===================================================================== +# Launcher — runs the dQ pass only (chunk loop + RMW), returns dq, chunk dS/P. +# d_sink (if needed) is a torch reduction handled by the caller. +# ===================================================================== +def sparse_mla_bwd_dq_gl( + q, + kv, + do, + topk_indices_padded, + lse, + delta, + R_CHUNK, + topk, + kv_lora_rank=512, + scale=None, + BLOCK_H=64, + TILE_K=16, +): + """ + Gluon dQ pass. Mirrors the dQ portion of `sparse_mla_bwd_v4`'s chunk loop. + + Returns: + dq: [T, H, D_QK] bf16 (fully accumulated across chunks) + chunk_dS: [T, H, R_CHUNK] bf16 (LAST chunk's dS — for spot validation) + chunk_P: [T, H, R_CHUNK] bf16 (LAST chunk's P) + """ + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + if scale is None: + scale = 1.0 / (d_qk**0.5) + assert R_CHUNK % TILE_K == 0, "TILE_K must divide R_CHUNK" + + dq = torch.empty_like(q) + chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + + num_hg = triton.cdiv(num_heads, BLOCK_H) + grid = (total_tokens, num_hg) + + for r_start in range(0, topk, R_CHUNK): + is_first = r_start == 0 + _sparse_mla_bwd_dq_gl_kernel[grid]( + q, + kv, + do, + topk_indices_padded, + lse, + delta, + dq, + chunk_dS, + chunk_P, + q.stride(0), + q.stride(1), + kv.stride(0), + do.stride(0), + do.stride(1), + dq.stride(0), + dq.stride(1), + topk_indices_padded.stride(0), + chunk_dS.stride(0), + chunk_dS.stride(1), + scale, + num_heads, + r_start, + R_CHUNK=R_CHUNK, + BLOCK_H=BLOCK_H, + TILE_K=TILE_K, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + IS_FIRST_CHUNK=is_first, + num_warps=4, + waves_per_eu=1, + ) + return dq, chunk_dS, chunk_P diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_v4_gluon.py new file mode 100644 index 000000000..62d6cd72a --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_bwd_v4_gluon.py @@ -0,0 +1,175 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Gluon DeepSeek-V4 sparse-MLA backward for the "gluon_v2" backend. + +Companion to the gluon_v2 forward (:func:`sparse_mla_fwd_v4_gluon_v2`). Wires the Gluon +dQ + Gluon dKV-intermediate compute kernels with a Triton Delta preprocess and the +backend-neutral CSR inverted-topk gather + torch d_sink reduction (non-atomic +chunked-gather scheme). + +The dQ / dKV-intermediate Gluon kernels apply the forward campaign's accepted techniques +to the backward: rope-skip (the V4 zero-rope-pad makes the rope gradients provably zero) ++ MFMA K=32 for the D_V=512-reduction matmuls, plus a single-chunk dQ read-modify-write +for high head counts. Beats the plain-Triton backward ~1.12x geomean over the 6 +flash/pro x cr{0,4,128} shapes (eager-UT 9/9). +""" + +import torch +import triton + +from .._gluon_dsa._dsa_bwd_gather import _build_inverted_topk_slice, _bwd_dkv_gather_acc +from .._gluon_dsa._dsa_bwd_preprocess import _sparse_mla_bwd_preprocess +from .dsa_bwd_dkv_interm_gluon import _sparse_mla_bwd_dkv_interm_gl_kernel +from .dsa_bwd_dq_gluon import _sparse_mla_bwd_dq_gl_kernel + + +def sparse_mla_bwd_v4_gluon_v2(q, kv, o, do, topk_indices, lse, attn_sink=None, kv_lora_rank=512, scale=None): + """DeepSeek-V4 sparse-MLA backward (Gluon dQ/dKV). Returns ``(dq, dkv, d_sink)``.""" + assert q.is_contiguous() and kv.is_contiguous() and o.is_contiguous() + assert do.is_contiguous() and topk_indices.is_contiguous() and lse.is_contiguous() + + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + topk = topk_indices.shape[1] + if scale is None: + scale = 1.0 / (d_qk**0.5) + if kv.dim() == 2: + kv = kv.unsqueeze(1) + num_kv = kv.shape[0] + + has_sink = attn_sink is not None + if has_sink: + assert attn_sink.dtype == torch.float32 and attn_sink.shape == (num_heads,) + + # ---- preprocess: Delta = rowsum(O*dO) (Triton, unchanged) ---- + delta = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device) + BLOCK_H_PRE = triton.next_power_of_2(min(64, num_heads)) + _sparse_mla_bwd_preprocess[(total_tokens, triton.cdiv(num_heads, BLOCK_H_PRE))]( + O_ptr=o, + dO_ptr=do, + Delta_ptr=delta, + stride_o_t=o.stride(0), + stride_o_h=o.stride(1), + num_heads=num_heads, + D_V=kv_lora_rank, + BLOCK_H=BLOCK_H_PRE, + ) + + # ---- config ---- + # R2: dQ is read-modify-written across chunks, so more chunks = more redundant dq + # reload passes + repeated CSR builds. For high head counts (H>=128) the dq RMW volume + # is large, so a single chunk over the whole topk (bounded for memory) is a big win at + # pro cr4 (mirrors the triton bwd). Low head counts keep the 256 cap. + if num_heads >= 128: + R_CHUNK = min(topk, 1536) + else: + R_CHUNK = min(256, topk) + BH_DQ, TK_DQ = 64, 16 + BH_DKV, TK_DKV = 32, 64 + num_hg_dq = triton.cdiv(num_heads, BH_DQ) + num_hg_dkv = triton.cdiv(num_heads, BH_DKV) + + dq = torch.empty_like(q) + chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + dkv_acc = torch.zeros(num_kv, d_qk, dtype=torch.float32, device=q.device) + interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device) + + # ---- pad topk to R_CHUNK multiple ---- + topk_padded_len = ((topk + R_CHUNK - 1) // R_CHUNK) * R_CHUNK + if topk_padded_len != topk: + pad = torch.full((total_tokens, topk_padded_len - topk), -1, dtype=torch.int32, device=q.device) + topk_padded = torch.cat([topk_indices, pad], dim=1).contiguous() + else: + topk_padded = topk_indices + + all_csr = [ + _build_inverted_topk_slice(topk_padded[:, rs : rs + R_CHUNK], rs, R_CHUNK, num_kv=num_kv) + for rs in range(0, topk, R_CHUNK) + ] + + for chunk_idx, r_start in enumerate(range(0, topk, R_CHUNK)): + is_first = r_start == 0 + + _sparse_mla_bwd_dq_gl_kernel[(total_tokens, num_hg_dq)]( + q, + kv, + do, + topk_padded, + lse, + delta, + dq, + chunk_dS, + chunk_P, + q.stride(0), + q.stride(1), + kv.stride(0), + do.stride(0), + do.stride(1), + dq.stride(0), + dq.stride(1), + topk_padded.stride(0), + chunk_dS.stride(0), + chunk_dS.stride(1), + scale, + num_heads, + r_start, + R_CHUNK=R_CHUNK, + BLOCK_H=BH_DQ, + TILE_K=TK_DQ, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_ROPE=False, # V4 zero-rope-pad: dQ_rope is provably zero + discarded by the adapter + IS_FIRST_CHUNK=is_first, + num_warps=4, + waves_per_eu=1, + ) + + _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + num_heads, + R_CHUNK=R_CHUNK, + TILE_K=TK_DKV, + BLOCK_H=BH_DKV, + NUM_HG=num_hg_dkv, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_ROPE=False, # V4 zero-rope-pad: dKV_rope provably zero + discarded downstream + num_warps=4, + ) + + inv_ptr, inv_data = all_csr[chunk_idx] + _bwd_dkv_gather_acc[(num_kv,)]( + interm, + inv_ptr, + inv_data, + dkv_acc, + interm.stride(1), + dkv_acc.stride(0), + D_V=kv_lora_rank, + D_ROPE=rope_rank, + num_warps=4, + ) + + d_sink = None + if has_sink: + d_sink = -(torch.exp(attn_sink.unsqueeze(0) - lse) * delta).sum(0) + + dkv_out = dkv_acc.to(kv.dtype).unsqueeze(1) + return dq, dkv_out, d_sink diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_fwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_fwd_v4_gluon.py new file mode 100644 index 000000000..e41e1a510 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v2/dsa_fwd_v4_gluon.py @@ -0,0 +1,722 @@ +""" +Gluon forward for DeepSeek V4 sparse MLA (gfx950 / CDNA4), with attention sink. + +Based on Leon's (leonling-ll) V3.2 gluon forward from `leonling-ll/aiter` branch +`liyang/dsa` -- which adapted our V3.2 Triton forward and added the gfx950 hardware +control (MFMA4 layouts, padded/swizzled shared, double-buffered K, async DMA pipeline, +ds_read_tr transpose, explicit dot-operand layouts); see also his DSA PR ROCm/aiter#3456. +This file adds the V4 attention-sink epilogue (sink-inclusive LSE) so it matches +`sparse_mla_fwd_v4`; it is the forward of the "gluon_v2" backend. + +Optimizations accepted over the base gluon forward (gfx950 campaign): + * rope-skip (HAS_ROPE=False): the V4 latent bakes RoPE in-place over the 512, so the + rope QK term is provably zero -- skip the 64-wide rope MFMA + K_rope loads entirely. + * exp2 softmax: fold log2(e) into the QK scale so the per-element exp is one hardware + exp2 (m_i/l_i in log2 units; LSE converted back to natural log for the backward). + * MFMA K=32 for the QK score matmul (instr_shape [16,16,32]): the score reduces D_V=512, + so K=32 halves the QK MFMA instruction count vs K=16 (PV/acc stays K=16, TILE_K-bound). + * register-prefetched topk index (off the KV-gather critical path) + async double-buffered + K gather overlapping the QK/softmax/PV MFMAs. + +Pipeline: + Prologue: Q -> shared (async); K tile 0 -> shared (async, double-buffered); deep-prefetch topk. + Loop tile t: gather tile t+2 (async) while computing QK[t+1] + softmax(t) + PV(t); promote. + Epilogue: drain; fold sink into the denominator (V4); write O, LSE. +""" + +import functools + +import torch +import triton +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + +# --------------------------------------------------------------------------- +# Triton capability gate: the gluon_v2 forward is a Gluon kernel (the backend's backward +# is currently plain-Triton, being migrated to Gluon). The Gluon fwd needs a Gluon-capable triton whose CDNA4 +# async_copy accepts arbitrary (DistributedLinearLayout) offsets. Released +# triton 3.7.0/3.7.1 still restricts async_copy offsets to Blocked/Slice and +# will NOT compile this path; build triton from the commit below. +# --------------------------------------------------------------------------- +_GLUON_V2_REQUIRED_COMMIT = "09500db9f0" +_GLUON_V2_INSTALL_HINT = ( + "gluon_v2 forward (Gluon) requires a Gluon-capable triton whose CDNA4 " + "async_copy accepts general offset layouts. The installed triton ({ver}) does not (released " + "3.7.0/3.7.1 restrict async_copy offsets to BlockedLayout/SliceLayout).\n" + "Build & install triton-lang/triton @ commit " + _GLUON_V2_REQUIRED_COMMIT + ":\n" + " git clone https://github.com/triton-lang/triton.git third_party/triton\n" + " cd third_party/triton && git checkout " + _GLUON_V2_REQUIRED_COMMIT + "\n" + " pip install -r python/requirements.txt\n" + " TRITON_CODEGEN_BACKENDS=amd MAX_JOBS=128 pip wheel --no-build-isolation --no-deps . -w dist\n" + " pip install --force-reinstall --no-deps dist/triton-*.whl" +) + + +@functools.lru_cache(maxsize=1) +def _gluon_available() -> bool: + """True iff triton exposes the experimental CDNA4 Gluon dialect this fwd uses.""" + try: + from triton.experimental import gluon as _gl # noqa: F401 + from triton.experimental.gluon.language import amd as _amd + + return hasattr(_amd, "cdna4") + except Exception: # noqa: BLE001 + return False + + +def _require_gluon_v2_triton() -> None: + """Fail fast with a build hint when Gluon is unavailable. The kernel compile is + ALSO guarded at the launch site: a Gluon-capable-but-incompatible triton raises a + CompilationError there, which is re-wrapped with the same install hint (so the user + always gets the commit rather than a raw layout/compile error).""" + if not _gluon_available(): + raise RuntimeError(_GLUON_V2_INSTALL_HINT.format(ver=getattr(triton, "__version__", "unknown"))) + + +def _wrap_compile_error(exc: Exception) -> RuntimeError: + return RuntimeError( + _GLUON_V2_INSTALL_HINT.format(ver=getattr(triton, "__version__", "unknown")) + + f"\n(triton failed to compile the Gluon fwd: {type(exc).__name__}: " + + (str(exc).splitlines()[0] if str(exc).strip() else "") + + ")" + ) + + +# ===================================================================== +# Utility +# ===================================================================== +def _get_lds_limit(): + """Return the per-CU LDS limit in bytes for the current GPU. + + gfx942 (MI300X): 64 KB = 65536 bytes + gfx950 (MI355X): 160 KB = 163840 bytes + """ + if torch.cuda.is_available(): + prop = torch.cuda.get_device_properties(0) + gcn_arch = getattr(prop, "gcnArchName", "") + if "gfx950" in gcn_arch: + return 163840 + return 65536 + + +_LDS_LIMIT = _get_lds_limit() + + +# ===================================================================== +# Forward — autotune configs and pruning +# ===================================================================== +def _fwd_prune_configs(configs, named_args, **kwargs): + """Prune autotune configs that would exceed per-CU LDS.""" + D_V = kwargs.get("D_V", named_args.get("D_V")) + D_ROPE = kwargs.get("D_ROPE", named_args.get("D_ROPE")) + pruned = [] + for config in configs: + config.kwargs["BLOCK_H"] + tk = config.kwargs["TILE_K"] + ns = config.num_stages + kv_lds = (D_V + D_ROPE) * tk * 2 * ns + if kv_lds <= _LDS_LIMIT: + pruned.append(config) + if not pruned: + pruned.append(configs[0]) + return pruned + + +def _get_fwd_autotune_configs(): + configs = [ + triton.Config( + {"BLOCK_H": BLOCK_H, "TILE_K": TILE_K, "waves_per_eu": WPE}, + num_warps=nw, + ) + for BLOCK_H in [16, 32, 64] + for TILE_K in [16, 32, 64, 128] + for WPE in [0, 1, 2] + for nw in [4] # num_warps must be 4 to align with kernel implementation + ] + # configs = [triton.Config({"BLOCK_H": 64, "TILE_K": 32, "waves_per_eu": 0}, num_warps=4),] + return configs + + +@triton.autotune( + configs=_get_fwd_autotune_configs(), + key=["num_heads", "TOPK", "D_V", "D_ROPE"], + prune_configs_by={"early_config_prune": _fwd_prune_configs}, +) +@gluon.jit +def _sparse_mla_fwd_gl_v2_kernel( + Q_ptr, # [total_tokens, num_heads, D_QK] bf16 + KV_ptr, # [total_tokens, 1, D_QK] bf16 + TopK_ptr, # [total_tokens, TOPK] int32 + Sink_ptr, # [num_heads] fp32; ignored if HAS_SINK == False + O_ptr, # [total_tokens, num_heads, D_V] bf16 + LSE_ptr, # [total_tokens, num_heads] fp32 (sink-inclusive if HAS_SINK) + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_o_t: tl.int64, + stride_o_h: tl.int64, + stride_topk_t: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + TOPK: gl.constexpr, + BLOCK_H: gl.constexpr, + TILE_K: gl.constexpr, + D_V: gl.constexpr, + D_ROPE: gl.constexpr, + HAS_SINK: gl.constexpr, + HAS_ROPE: gl.constexpr, +): + # ---------- constexpr layouts ---------- + # QK MFMA uses K=32 (aiter): the score matmul reduces D_V=512, so instr_shape=[16,16,32] + # halves the MFMA instruction count vs [16,16,16]. mfma_acc (PV) stays K=16 since its + # reduction dim is TILE_K (autotuned down to 16). + mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 32], + transposed=True, + warps_per_cta=[4, 1], + ) + mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + + # Blocked layouts for global loads. + _qlora_tpw_k: gl.constexpr = min(64, D_V // 8) + _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k + blk_qlora: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_qrope: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[4, 1], + order=[1, 0], + ) + + _klora_tpw_m: gl.constexpr = min(64, D_V // 8) + _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m + blk_klora: gl.constexpr = gl.BlockedLayout( # [D_V, TILE_K] + size_per_thread=[8, 1], + threads_per_warp=[_klora_tpw_m, _klora_tpw_n], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_krope: gl.constexpr = gl.BlockedLayout( # [D_ROPE, TILE_K] = [64, 16] + size_per_thread=[2, 1], + threads_per_warp=[32, 2], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_topk: gl.constexpr = gl.BlockedLayout( # [TILE_K] int32 + size_per_thread=[1], + threads_per_warp=[64], + warps_per_cta=[4], + order=[0], + ) + blk_lse: gl.constexpr = gl.BlockedLayout( # [BLOCK_H] fp32 + size_per_thread=[1], + threads_per_warp=[64], + warps_per_cta=[4], + order=[0], + ) + + # Shared layouts. + sh_qlora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [BLOCK_H, D_V], + [1, 0], + ) + sh_qrope: gl.constexpr = gl.SwizzledSharedLayout( + vec=8, + per_phase=2, + max_phase=8, + order=[1, 0], + ) + sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [D_V, TILE_K], + [0, 1], + ) + sh_krope: gl.constexpr = gl.SwizzledSharedLayout( + vec=8, + per_phase=2, + max_phase=8, + order=[0, 1], + ) + + # Dot operand layouts + dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_p_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4) + dot_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4) + + # ---------- program ids ---------- + token_idx = gl.program_id(axis=0) + hg_idx = gl.program_id(axis=1) + hg_offset = hg_idx * BLOCK_H + + # ---------- offsets for Q ---------- + # Q_lora [BLOCK_H, D_V] + offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_qlora = offs_h_qlora < num_heads + + q_base = token_idx.to(tl.int64) * stride_q_t + q_offs_lora = ( + q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64) + ) + q_mask_lora = mask_h_qlora[:, None] + + smem_qlora = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_qlora) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qlora, + ptr=Q_ptr, + offsets=q_offs_lora.to(tl.int32), + mask=q_mask_lora, + ) + # V4 zero-rope-pad: skip the rope Q load + rope MFMA entirely when HAS_ROPE is False + # (RoPE is baked in-place over the 512 latent, so the rope QK term is provably zero). + if HAS_ROPE: + # Q_rope [BLOCK_H, D_ROPE] + offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_qrope = offs_h_qrope < num_heads + q_offs_rope = ( + q_base + + offs_h_qrope[:, None].to(tl.int64) * stride_q_h + + (D_V + offs_r_qrope[None, :]).to(tl.int64) + ) + q_mask_rope = mask_h_qrope[:, None] + smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qrope, + ptr=Q_ptr, + offsets=q_offs_rope.to(tl.int32), + mask=q_mask_rope, + ) + gl.amd.cdna4.async_copy.commit_group() + + # ---------- topk and KV offsets ---------- + NUM_TILES: gl.constexpr = (TOPK + TILE_K - 1) // TILE_K + topk_base = token_idx.to(tl.int64) * stride_topk_t + + # offs_tile in three layouts (sliced from each of the three loaders) + offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora)) + offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope)) + offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s)) + offs_tile_topk = gl.arange(0, TILE_K, layout=blk_topk) + + offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora)) + offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope)) + + # (removed dead `topk_pos_reg` prologue load — was never consumed) + + # ---------- shared mem allocations for the K loop ---------- + if HAS_ROPE: + smem_krope = gl.allocate_shared_memory( + KV_ptr.dtype.element_ty, + [2, D_ROPE, TILE_K], + layout=sh_krope, + ) + smem_klora = gl.allocate_shared_memory( + KV_ptr.dtype.element_ty, + [2, D_V, TILE_K], + layout=sh_klora, + ) + + # ---------- accumulators ---------- + m_i = gl.full([BLOCK_H], float("-inf"), dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s)) + l_i = gl.full([BLOCK_H], 0.0, dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s)) + acc = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc) + # exp2 softmax (aiter technique): fold log2(e) into the QK scale so the per-element + # softmax exp becomes a single hardware exp2 (no per-element *log2e). m_i / l_i are then + # in log2 units; lse is converted back to natural log at the epilogue (the bwd needs nat-log). + scale_log2 = scale * 1.4426950408889634 + + # ---------- tile-0 prefetch (prologue) ---------- + # Load K_lora and K_rope for tile 0. + topk_pos_klora = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + offs_tile_klora, + mask=offs_tile_klora < TOPK, + other=-1, + ) + if HAS_ROPE: + topk_pos_krope = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + offs_tile_krope, + mask=offs_tile_krope < TOPK, + other=-1, + ) + topk_pos_mma = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + offs_tile_mma, + mask=offs_tile_mma < TOPK, + other=-1, + ) + + # Deep-prefetch tile-1 topk ONCE in the neutral blk_topk layout (DEDUP). Carried as a + # single register set; converted to the klora/krope/mma layouts at point of use, to + # minimize carried register pressure (the 3-layout carry caused an acc-rescale codegen + # regression -- see att_fwd_gluon_mi350/RESULTS.md). + p1_off_topk = TILE_K + offs_tile_topk + tkraw = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + p1_off_topk, + mask=p1_off_topk < TOPK, + other=-1, + ) + + valid_klora = topk_pos_klora != -1 # tile_start=0 -> offs_tile buf1, drain K[0], QK[0] -> S_prev (no softmax/PV yet). + tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora)) + tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s)) + valid_klora_next = ((TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1) + valid_qk = ((TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1) + safe_klora_next = gl.where(valid_klora_next, tk_klora, 0) + klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(1), + ptr=KV_ptr, + offsets=klora_offs_next.to(tl.int32), + mask=valid_klora_next[None, :], + ) + if HAS_ROPE: + tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope)) + valid_krope_next = ((TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1) + safe_krope_next = gl.where(valid_krope_next, tk_krope, 0) + krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + ( + D_V + offs_r_krope[:, None] + ).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(1), + ptr=KV_ptr, + offsets=krope_offs_next.to(tl.int32), + mask=valid_krope_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + tkraw = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + (2 * TILE_K + offs_tile_topk), + mask=(2 * TILE_K + offs_tile_topk) < TOPK, + other=-1, + ) + gl.amd.cdna4.async_copy.wait_group(1) + S_prev = gl.amd.cdna4.mfma( + Q_lora_dot, + smem_klora.index(0).load(dot_klora_b), + gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s), + ) + if HAS_ROPE: + S_prev = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(0).load(dot_krope_b), S_prev) + S_prev = S_prev * scale_log2 + S_prev = gl.where(valid_mma[None, :] & mask_h_mma[:, None], S_prev, float("-inf")) + cur_buf = 1 + + for t in range(NUM_TILES - 2): + gl.amd.cdna4.async_copy.wait_group(0) # drain K[t+1] (cur_buf) before QK reads it + # 2-BUFFER EARLY-GATHER: evacuate V[t] from pv_buf into REGISTERS first, freeing that buffer, + # then gather tile t+2 into it BEFORE the QK/PV MFMAs so the DMA overlaps both (no 3rd buffer). + # V_lora_dot in regs => no read/async-write race on the recycled buffer. Costs VGPR live range. + V_lora_dot = smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b) + tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora)) + tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s)) + valid_klora_next = (((t + 2) * TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1) + valid_qk_next = (((t + 2) * TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1) + safe_klora_next = gl.where(valid_klora_next, tk_klora, 0) + klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to( + tl.int64 + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(1 - cur_buf), + ptr=KV_ptr, + offsets=klora_offs_next.to(tl.int32), + mask=valid_klora_next[None, :], + ) + if HAS_ROPE: + tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope)) + valid_krope_next = (((t + 2) * TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1) + safe_krope_next = gl.where(valid_krope_next, tk_krope, 0) + krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + ( + D_V + offs_r_krope[:, None] + ).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(1 - cur_buf), + ptr=KV_ptr, + offsets=krope_offs_next.to(tl.int32), + mask=valid_krope_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + tkraw_n = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + ((t + 3) * TILE_K + offs_tile_topk), + mask=((t + 3) * TILE_K + offs_tile_topk) < TOPK, + other=-1, + ) + # QK tile (t+1) from cur_buf -- matrix; overlaps the gather above + softmax below + S_cur = gl.amd.cdna4.mfma( + Q_lora_dot, + smem_klora.index(cur_buf).load(dot_klora_b), + gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s), + ) + if HAS_ROPE: + S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur) + S_cur = S_cur * scale_log2 + S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf")) + # softmax(S_prev = tile t) [VALU, overlaps QK] + m_j = gl.max(S_prev, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > float("-inf"), m_new, 0.0) + alpha = gl.exp2(m_i - m_new) + P = gl.exp2(S_prev - m_new[:, None]) + l_i = alpha * l_i + gl.sum(P, axis=1) + m_i = m_new + # PV tile t from registers -- matrix; overlaps the gather still in flight + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(P_dot, V_lora_dot, acc) + # promote + S_prev = S_cur + valid_qk = valid_qk_next + tkraw = tkraw_n + cur_buf = 1 - cur_buf + + # ---------- PRE-DRAIN: QK[N-1] (cur_buf) || softmax+PV[N-2] (pv_buf); no gather ---------- + gl.amd.cdna4.async_copy.wait_group(0) # drain K[N-1] (last loop gather) + S_cur = gl.amd.cdna4.mfma( + Q_lora_dot, + smem_klora.index(cur_buf).load(dot_klora_b), + gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s), + ) + if HAS_ROPE: + S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur) + S_cur = S_cur * scale_log2 + S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf")) + m_j = gl.max(S_prev, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > float("-inf"), m_new, 0.0) + alpha = gl.exp2(m_i - m_new) + P = gl.exp2(S_prev - m_new[:, None]) + l_i = alpha * l_i + gl.sum(P, axis=1) + m_i = m_new + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b), acc) + S_prev = S_cur + + # ---------- DRAIN: softmax+PV[N-1] (S_prev = QK[N-1], V from cur_buf) ---------- + m_j = gl.max(S_prev, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > float("-inf"), m_new, 0.0) + alpha = gl.exp2(m_i - m_new) + P = gl.exp2(S_prev - m_new[:, None]) + l_new = alpha * l_i + gl.sum(P, axis=1) + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(cur_buf).permute([1, 0]).load(dot_v_b), acc) + m_i = m_new + l_i = l_new + + # ---------- epilogue: fold sink into the denominator (V4 delta) ---------- + if HAS_SINK: + offs_h_sink = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + sink = gl.amd.cdna4.buffer_load( + ptr=Sink_ptr, + offsets=offs_h_sink.to(tl.int32), + mask=offs_h_sink < num_heads, + other=float("-inf"), + ) + sink = sink * 1.4426950408889634 # natural-log sink -> log2 units (exp2 softmax) + m_final = gl.maximum(m_i, sink) + alpha_fix = gl.exp2(m_i - m_final) + l_total = l_i * alpha_fix + gl.exp2(sink - m_final) + alpha_fix_acc = gl.convert_layout(alpha_fix, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_fix_acc[:, None] + l_total_acc = gl.convert_layout(l_total, gl.SliceLayout(1, mfma_acc)) + acc = acc / l_total_acc[:, None] + # lse back to natural log for the backward: m_final is log2, l_total is the natural denom. + lse = m_final * 0.6931471805599453 + gl.log(l_total) + else: + l_i_acc = gl.convert_layout(l_i, gl.SliceLayout(1, mfma_acc)) + acc = acc / l_i_acc[:, None] + lse = m_i * 0.6931471805599453 + gl.log(l_i) + + # Output O[token_idx, h, v] + offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_o = offs_h_o < num_heads + o_base = token_idx.to(tl.int64) * stride_o_t + o_offs = o_base + offs_h_o[:, None].to(tl.int64) * stride_o_h + offs_v_o[None, :].to(tl.int64) + acc_bf = acc.to(O_ptr.dtype.element_ty) + acc_bf_blk = gl.convert_layout(acc_bf, blk_qlora) + gl.amd.cdna4.buffer_store( + stored_value=acc_bf_blk, + ptr=O_ptr, + offsets=o_offs.to(tl.int32), + mask=mask_h_o[:, None], + ) + + # LSE[token_idx, h] + offs_h_lse = hg_offset + gl.arange(0, BLOCK_H, layout=blk_lse) + mask_h_lse = offs_h_lse < num_heads + lse_base = token_idx * num_heads + lse_offs = lse_base + offs_h_lse + lse_blk = gl.convert_layout(lse, blk_lse) + gl.amd.cdna4.buffer_store( + stored_value=lse_blk, + ptr=LSE_ptr, + offsets=lse_offs.to(tl.int32), + mask=mask_h_lse, + ) + + +# ===================================================================== +# Launcher +# ===================================================================== +def sparse_mla_fwd_v4_gluon_v2(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None): + """ + DeepSeek V4 sparse MLA forward (Gluon, gfx950 / CDNA4), with attention sink. + + Args: + q: [total_tokens, num_heads, d_qk] bfloat16 + kv: [total_tokens, 1, d_qk] bfloat16 (or [total_tokens, d_qk]) + topk_indices: [total_tokens, topk] int32 (SWA + sparse, -1 marks invalid) + attn_sink: [num_heads] fp32, optional per-head learnable sink logit. + When None, behaves like the V3.2 forward. + kv_lora_rank: int, default 512 + scale: float, default 1/sqrt(d_qk) + + Returns: + o: [total_tokens, num_heads, kv_lora_rank] same dtype as q + lse: [total_tokens, num_heads] float32 (sink-inclusive when attn_sink is given) + """ + _require_gluon_v2_triton() + assert q.is_contiguous() + assert kv.is_contiguous() + assert topk_indices.is_contiguous() + + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + topk = topk_indices.shape[1] + + if scale is None: + scale = 1.0 / (d_qk**0.5) + + if kv.dim() == 2: + kv = kv.unsqueeze(1) + # kv may hold MORE rows than there are query tokens (V4 feeds a + # [local ++ compressed-pool] buffer, so num_kv = S + P > total_tokens). + # The kernel only dereferences kv via topk indices (stride_kv_t), so any + # num_kv >= max(topk_index)+1 is valid. + assert kv.shape[0] >= total_tokens and kv.shape[-1] == d_qk + + has_sink = attn_sink is not None + if has_sink: + assert attn_sink.is_contiguous() + assert attn_sink.dtype == torch.float32 + assert attn_sink.shape == (num_heads,) + sink_ptr = attn_sink + else: + sink_ptr = torch.empty(1, dtype=torch.float32, device=q.device) # guarded by HAS_SINK + + o = torch.empty(total_tokens, num_heads, kv_lora_rank, dtype=q.dtype, device=q.device) + lse = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device) + + # V4 single-latent form: the D_ROPE block of q/kv is a zero pad (RoPE baked in-place + # over the 512 latent), so the rope QK term is provably zero. Skip it — bit-identical + # to computing it, but avoids the wasteful 64-wide rope MFMA + K_rope loads every tile + # (this is the win triton_v2 already has that our ported gluon fwd lacked). + has_rope = False + + # Grid is autotune-aware: BLOCK_H comes from the chosen config. + grid = lambda META: (total_tokens, triton.cdiv(num_heads, META["BLOCK_H"])) + + try: + _sparse_mla_fwd_gl_v2_kernel[grid]( + Q_ptr=q, + KV_ptr=kv, + TopK_ptr=topk_indices, + Sink_ptr=sink_ptr, + O_ptr=o, + LSE_ptr=lse, + stride_q_t=q.stride(0), + stride_q_h=q.stride(1), + stride_kv_t=kv.stride(0), + stride_o_t=o.stride(0), + stride_o_h=o.stride(1), + stride_topk_t=topk_indices.stride(0), + scale=scale, + num_heads=num_heads, + TOPK=topk, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_SINK=has_sink, + HAS_ROPE=has_rope, + ) + except Exception as exc: # noqa: BLE001 - surface a build hint on Gluon compile failures + _n = type(exc).__name__.lower() + if "compil" in _n or "compil" in str(exc).lower() or "layout" in str(exc).lower(): + raise _wrap_compile_error(exc) from exc + raise + + return o, lse diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/__init__.py new file mode 100644 index 000000000..a736001ef --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/__init__.py @@ -0,0 +1,31 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Gluon DeepSeek-V4 sparse-MLA attention backend ("gluon v3"). + +Third-generation Gluon (gfx950 / CDNA4) sparse-MLA optimization campaign backend. +Round 1 intentionally starts from the stable ``gluon_v2`` implementation so every +future round can change exactly one kernel variable and be compared linearly. + +* forward (``dsa_fwd_v4_gluon``): MFMA layouts, padded/swizzled shared, async + double-buffered pipeline, rope-skip, exp2 softmax, MFMA K=32. +* backward (``dsa_bwd_v4_gluon``): Gluon dQ + dKV-intermediate kernels (rope-skip, + MFMA K=32, single-chunk dQ RMW) + Triton Delta preprocess + CSR inverted-topk gather. + +The Gluon kernels need a Gluon-capable (recompiled) triton whose CDNA4 async_copy +accepts general offset layouts, and raise a clear install hint otherwise. + +* :func:`sparse_mla_fwd_v4_gluon_v3` -> ``(o, lse)`` +* :func:`sparse_mla_bwd_v4_gluon_v3` -> ``(dq, dkv, d_sink)`` +""" + +from .dsa_bwd_v4_gluon import sparse_mla_bwd_v4_gluon_v2 as sparse_mla_bwd_v4_gluon_v3 +from .dsa_fwd_v4_gluon import sparse_mla_fwd_v4_gluon_v2 as sparse_mla_fwd_v4_gluon_v3 + +__all__ = [ + "sparse_mla_fwd_v4_gluon_v3", + "sparse_mla_bwd_v4_gluon_v3", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_lse_fwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_lse_fwd.py new file mode 100644 index 000000000..3ea978efc --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_lse_fwd.py @@ -0,0 +1,172 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from .aiter_mla_gluon import mla_gluon + + +@triton.jit +def _count_valid_topk_kernel(topk_ptr, counts_ptr, stride_t, topk: tl.constexpr, BLOCK: tl.constexpr): + row = tl.program_id(0) + offs = tl.arange(0, BLOCK) + vals = tl.load(topk_ptr + row * stride_t + offs, mask=offs < topk, other=-1) + valid = (offs < topk) & (vals >= 0) + count = tl.sum(valid.to(tl.int32), axis=0) + tl.store(counts_ptr + row, count) + + +@triton.jit +def _pack_valid_topk_kernel( + topk_ptr, + indptr_ptr, + flat_ptr, + stride_t, + topk: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + offs = tl.arange(0, BLOCK) + vals = tl.load(topk_ptr + row * stride_t + offs, mask=offs < topk, other=-1) + valid = (offs < topk) & (vals >= 0) + pos = tl.cumsum(valid.to(tl.int32), axis=0) - 1 + start = tl.load(indptr_ptr + row) + tl.store(flat_ptr + start + pos, vals.to(tl.int32), mask=valid) + + +def _v4_topk_to_ragged_gpu(topk_indices: torch.Tensor): + total_tokens, topk = topk_indices.shape + block = triton.next_power_of_2(topk) + counts = torch.empty(total_tokens, dtype=torch.int32, device=topk_indices.device) + _count_valid_topk_kernel[(total_tokens,)]( + topk_indices, + counts, + topk_indices.stride(0), + topk, + BLOCK=block, + ) + indptr = torch.empty(total_tokens + 1, dtype=torch.int32, device=topk_indices.device) + indptr[0] = 0 + torch.cumsum(counts, dim=0, out=indptr[1:]) + + # Allocate the max possible nnz to avoid a CPU sync on indptr[-1]. The Gluon + # kernel uses indptr for bounds, so unused tail elements are never read. + flat = torch.empty(total_tokens * topk, dtype=torch.int32, device=topk_indices.device) + _pack_valid_topk_kernel[(total_tokens,)]( + topk_indices, + indptr, + flat, + topk_indices.stride(0), + topk, + BLOCK=block, + ) + return flat, indptr + + +@triton.jit +def _pack_v4_csa_h64_kernel( + topk_ptr, + indptr_ptr, + flat_ptr, + stride_t, + total_tokens: tl.constexpr, + TOPK: tl.constexpr, + WINDOW: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + pool_k: tl.constexpr = TOPK - WINDOW + local_count = tl.minimum(row + 1, WINDOW) + + if row < WINDOW: + local_prefix = row * (row + 1) // 2 + else: + local_prefix = WINDOW * (WINDOW + 1) // 2 + (row - WINDOW) * WINDOW + start = row * pool_k + local_prefix + tl.store(indptr_ptr + row, start) + + offs = tl.arange(0, BLOCK) + local_mask = offs < local_count + pool_offs = offs - local_count + pool_mask = (offs >= local_count) & (pool_offs < pool_k) + src = tl.where(local_mask, WINDOW - local_count + offs, WINDOW + pool_offs) + vals = tl.load(topk_ptr + row * stride_t + src, mask=local_mask | pool_mask, other=-1) + tl.store(flat_ptr + start + offs, vals.to(tl.int32), mask=local_mask | pool_mask) + + if row == total_tokens - 1: + tl.store(indptr_ptr + total_tokens, start + local_count + pool_k) + + +def _v4_csa_h64_topk_to_ragged(topk_indices: torch.Tensor): + total_tokens, topk = topk_indices.shape + block = triton.next_power_of_2(topk) + indptr = torch.empty(total_tokens + 1, dtype=torch.int32, device=topk_indices.device) + flat = torch.empty(total_tokens * topk, dtype=torch.int32, device=topk_indices.device) + _pack_v4_csa_h64_kernel[(total_tokens,)]( + topk_indices, + indptr, + flat, + topk_indices.stride(0), + total_tokens, + topk, + WINDOW=128, + BLOCK=block, + ) + return flat, indptr + + +def sparse_mla_fwd_v4_aiter_lse(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None): + """Aiter Gluon sparse-MLA fwd with LSE, adapted to the V4 dense-topk API.""" + _, _, d_qk = q.shape + if scale is None: + scale = 1.0 / (d_qk**0.5) + + q_nope = q[..., :kv_lora_rank].contiguous() + kv2 = kv[:, 0, :] if kv.dim() == 3 else kv + kv_c = kv2[:, :kv_lora_rank].contiguous() + flat, indptr = _v4_topk_to_ragged_gpu(topk_indices) + out = torch.empty(q.shape[0], q.shape[1], kv_lora_rank, dtype=q.dtype, device=q.device) + + return mla_gluon( + q_nope, + None, + kv_c, + out, + page_table=flat, + seq_info=indptr, + sm_scale=float(scale), + has_pe=False, + min_kv_seq_len=float("inf"), + attn_sink=attn_sink, + return_lse=True, + ) + + +def sparse_mla_fwd_v4_aiter_lse_csa_formula( + q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None +): + """CSA specialization using the V4 [SWA128 + pool topk] dense-topk layout.""" + _, _, d_qk = q.shape + if scale is None: + scale = 1.0 / (d_qk**0.5) + + q_nope = q[..., :kv_lora_rank].contiguous() + kv2 = kv[:, 0, :] if kv.dim() == 3 else kv + kv_c = kv2[:, :kv_lora_rank].contiguous() + flat, indptr = _v4_csa_h64_topk_to_ragged(topk_indices) + out = torch.empty(q.shape[0], q.shape[1], kv_lora_rank, dtype=q.dtype, device=q.device) + + return mla_gluon( + q_nope, + None, + kv_c, + out, + page_table=flat, + seq_info=indptr, + sm_scale=float(scale), + has_pe=False, + min_kv_seq_len=float("inf"), + attn_sink=attn_sink, + return_lse=True, + ) diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_mla_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_mla_gluon.py new file mode 100644 index 000000000..811a74428 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/aiter_mla_gluon.py @@ -0,0 +1,999 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +# Gluon MLA decode kernel originated from FlashMLA triton kernel(https://github.com/deepseek-ai/FlashMLA/blob/main/benchmark/bench_flash_mla.py). +# Stage-1 split-KV MLA attention using explicit Gluon layouts. Three regimes: +# +# REGIME='bh64' - bf16 Q + bf16 KV, BLOCK_H=64, BLOCK_N=64, +# nhead in {64, 128}, batch_size in {64, 128, 256}, +# NUM_KV_SPLITS auto-picked to fill ~256 WGs (in {1,2,4}). +# Fast path: when NUM_KV_SPLITS==1, stage-1 writes the +# final output directly to O and stage-2 reduce is skipped. +# REGIME='bh16bn128' - bf16 Q + fp8 KV, BLOCK_H=16, BLOCK_N=128, +# nhead <= 16, batch_size=1, NUM_KV_SPLITS=256. +# 2-D (batch, split) grid. Always splits + always +# reduces. NHEAD < BLOCK_H masks OOB heads on Q load +# and O store. +# REGIME='bh16bn64' - bf16 Q + bf16 KV, BLOCK_H=16, BLOCK_N=64, +# nhead <= 16, batch_size >= 1, 2-D (batch, split) grid, +# NUM_KV_SPLITS = max(1, 256 // batch_size). Full decode +# (stage-1 + stage-2 reduce into the final O). +# NHEAD < BLOCK_H masks OOB heads on Q load and O store. +# +# The bh16 regimes support num_iter in {1, 2, ...} (no gl.assume(num_iter>=3)); +# only bh64 assumes >= 3. See epilogue-1 handling below. +# +# Wrapper dispatch: nhead in {64,128} -> bh64; nhead <= 16 routes by KV dtype +# (bf16 -> bh16bn64, fp8 -> bh16bn128). +# +# Full decode for all regimes. For NUM_KV_SPLITS>1 stage-1 writes per-split acc + +# fp32 lse; stage-2 (_mla_softmax_reducev_kernel) reduces into O. RETURN_LSE also +# returns the merged fp32 lse [B, H] (stage-2 for splits>1, else stage-1). +# +# 3-stage software pipeline (double-buffered, BLOCK_N with 2x(BLOCK_N/2) KV slices): +# AC = async_copy (global->LDS), LL = load (LDS->reg), P = page, K = K-cache, V = V-cache +# +# iter i iter i+1 iter i+2 +# ACP(page): [i+2] [i+3] [i+4] +# LLP+ACK(K): [i+1] [i+2] [i+3] +# LLK+MFMA+LLV: [i] [i+1] [i+2] +# +# Within each loop iteration (operating on buf_idx=current, async_idx=next): +# ACP -- async_copy page numbers [i+2] +# LLP, ACK -- local_load pages [i+1], async_copy K/KPE [i+1] +# LLK, MFMA0, softmax, LLV, MFMA1 -- compute on [i]: QK dot, softmax, PV dot + +import aiter.ops.triton.utils._triton.arch_info as arch_info +import torch +import triton +import triton.language as tl +from aiter.ops.triton.utils.device_info import get_num_xcds +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +# fmt: off +@gluon.jit +def _mla_gluon( + Q_nope, + Q_pe, + Kv_c_cache, + K_pe_cache, + Req_to_tokens, + B_seq_len, + O, # noqa: E741 + Attn_sink, + sm_scale, + kv_scale, + stride_q_nope_bs, + stride_q_nope_h, + stride_q_pe_bs, + stride_q_pe_h, + stride_kv_c_bs, + stride_k_pe_bs, + stride_req_to_tokens_bs, + stride_o_b, + stride_o_h, + stride_o_s, + Mid_lse, # split>1: per-split fp32 lse [B, H, NUM_KV_SPLITS] (else None) + stride_mid_lse_b, + stride_mid_lse_h, + stride_mid_lse_s, + Final_lse, # RETURN_LSE only: merged fp32 lse [B, H] (else None) + stride_final_lse_b, + stride_final_lse_h, + BLOCK_H: gl.constexpr, + BLOCK_N: gl.constexpr, + NUM_KV_SPLITS: gl.constexpr, + PAGE_SIZE: gl.constexpr, + HEAD_DIM_CKV: gl.constexpr, + HEAD_DIM_KPE: gl.constexpr, + KV_PE_OFFSET: gl.constexpr, + USE_2D_VIEW: gl.constexpr, + WITHIN_2GB: gl.constexpr, + NUM_XCDS: gl.constexpr, + NHEAD: gl.constexpr, + REGIME: gl.constexpr, + RETURN_LSE: gl.constexpr, + # --- dsv4-prefill knobs --- + HAS_PE: gl.constexpr, + HAS_ATTN_SINK: gl.constexpr, +): + # Grid mapping: bh64 uses 3-D XCD-aware multi-batch; bh16bn64 and bh16bn128 + # use 2-D (batch, split) — for batch_size=1 this is (1, NUM_KV_SPLITS). + if REGIME == 'bh64': + cur_batch = gl.program_id(0) + (gl.program_id(2) // NUM_KV_SPLITS) * NUM_XCDS + cur_head_id = gl.program_id(1) + split_kv_id = gl.program_id(2) % NUM_KV_SPLITS + else: + cur_batch = gl.program_id(0) + cur_head_id = 0 + split_kv_id = gl.program_id(1) + + # USE_2D_VIEW=True: fixed len or max padded VarLen + # Req_to_tokens = block_table[batch, max_seqlen], B_seq_len = cache_seqlens[batch] + # USE_2D_VIEW=False: flattened VarLen + # Req_to_tokens = kv_indices[total_kv], B_seq_len = kv_indptr[batch+1] + if USE_2D_VIEW: + batch_page_start = stride_req_to_tokens_bs * cur_batch + cur_batch_seq_len = gl.load(B_seq_len + cur_batch) + else: + batch_page_start = gl.load(B_seq_len + cur_batch) + cur_batch_seq_len = gl.load(B_seq_len + cur_batch + 1) - batch_page_start + + # split-KV: each program covers [split_kv_start, split_kv_end). + # OLD: ceil-based per_split. The LAST split could be empty (num_iter=0), + # which breaks the unconditional epilogue-2 consume. Kept here as commented + # reference; remove in cleanup. + # kv_len_per_split = gl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS) + # split_kv_start = kv_len_per_split * split_kv_id + # split_kv_end = gl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len) + # + # NEW: floor per_split with the last split absorbing the remainder + # (remainder = seq mod NUM_KV_SPLITS, in [0, NUM_KV_SPLITS)). Combined with + # the wrapper bound min_kv_seq_len >= NUM_KV_SPLITS this guarantees every + # split is non-empty (split_len >= floor >= 1, hence num_iter >= 1); bh64 + # additionally bounds min_kv_seq_len so num_iter >= 3 for its gl.assume. + # Trade-off: at seqs just above the wrapper minimum the last CU does up to + # ~(floor + NUM_KV_SPLITS - 1)/floor more work than the others. + kv_len_per_split = cur_batch_seq_len // NUM_KV_SPLITS + split_kv_start = kv_len_per_split * split_kv_id + split_kv_end = split_kv_start + kv_len_per_split + if split_kv_id == NUM_KV_SPLITS - 1: + split_kv_end = cur_batch_seq_len + num_iter = gl.cdiv(split_kv_end - split_kv_start, BLOCK_N) + start_n = split_kv_start + + # early return with empty kv slice to save compute + if split_kv_start >= split_kv_end: + return + + ######### layout setting begin ######### + # Q-side layouts + mfma_layout: switch by BLOCK_H. + # bh64 has BLOCK_H=64; bh16bn128 and bh16bn64 share BLOCK_H=16 (identical Q layouts + mfma orientation). + if BLOCK_H == 64: + # bh64: Q is [64, 512] / [64, 64]; warps tile M. + blocked_q_nope: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[1, 64], + warps_per_cta=[4, 1], + order=[1, 0], + ) + shared_q_nope: gl.constexpr = gl.PaddedSharedLayout( + interval_padding_pairs=[[512, 16]], + offset_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [0, 32], [0, 64], [0, 128], [0, 256], [1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [32, 0]], + cga_layout=[], + shape=[64, 512] + ) + blocked_q_pe: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((0, 1), (0, 2), (0, 4), (32, 0)), + lane_bases=((0, 8), (0, 16), (0, 32), (4, 0), (8, 0), (16, 0)), + warp_bases=((1, 0), (2, 0)), + block_bases=[], + shape=[64, 64], + ) + shared_q_pe: gl.constexpr = gl.PaddedSharedLayout( + interval_padding_pairs=[[512, 16]], + offset_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [0, 32], [4, 0], [8, 0], [16, 0], [1, 0], [2, 0], [32, 0]], + cga_layout=[], + shape=[64, 64] + ) + mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 32], + transposed=True, + warps_per_cta=[4, 1], + ) + else: + # BLOCK_H == 16: shared by bh16bn128 and bh16bn64. Q is [16, 512] / [16, 64]; warps tile K. + blocked_q_nope: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[1, 64], + warps_per_cta=[4, 1], + order=[1, 0], + ) + shared_q_nope: gl.constexpr = gl.PaddedSharedLayout( + interval_padding_pairs=[[512, 16]], + offset_bases=[[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [0, 32], [0, 64], [0, 128], [0, 256], [1, 0], [2, 0], [4, 0], [8, 0]], + cga_layout=[], + shape=[16, 512] + ) + blocked_q_pe: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((0, 1), (0, 2), (0, 4)), + lane_bases=((0, 8), (0, 16), (0, 32), (1, 0), (2, 0), (4, 0)), + warp_bases=((8, 0), (0, 0)), + block_bases=[], + shape=[16, 64], + ) + shared_q_pe: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[1, 0]) + mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 32], + transposed=True, + warps_per_cta=[1, 4], + ) + + # KV-side layouts: switch by BLOCK_N. + # bh16bn128 (BLOCK_N=128, fp8 KV) needs distinct K layouts; bh64 and bh16bn64 share BLOCK_N=64 bf16 KV. + if BLOCK_N == 128: + # bh16bn128: K is [512, 128]fp8, KPE is [64, 128]fp8. + blocked_kv: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 8), (0, 4), (0, 32), (0, 64)), + lane_bases=((16, 0), (32, 0), (64, 0), (128, 0), (256, 0), (0, 16)), + warp_bases=((0, 1), (0, 2)), + block_bases=[], + shape=[512, 128], + ) + shared_kv: gl.constexpr = gl.PaddedSharedLayout( + interval_padding_pairs=[[1024, 32], [8192, 16]], + offset_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [32, 0], [64, 0], [128, 0], [256, 0], [0, 16], [0, 1], [0, 2], [0, 8], [0, 4], [0, 32], [0, 64]], + cga_layout=[], + shape=[512, 128] + ) + blocked_kpe: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 2)), + lane_bases=((16, 0), (32, 0), (0, 4), (0, 8), (0, 16), (0, 32)), + warp_bases=((0, 64), (0, 1)), + block_bases=[], + shape=[64, 128], + ) + shared_kpe: gl.constexpr = gl.PaddedSharedLayout( + interval_padding_pairs=[[2048, 16]], + offset_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [32, 0], [0, 4], [0, 8], [0, 16], [0, 32], [0, 64], [0, 1], [0, 2]], + cga_layout=[], + shape=[64, 128] + ) + blocked_page: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((0,),), + lane_bases=((1,), (2,), (4,), (8,), (16,), (32,)), + warp_bases=((64,), (0,)), + block_bases=[], + shape=[128], + ) + blocked_kv_slice: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 8), (0, 4), (0, 32)), + lane_bases=((16, 0), (32, 0), (64, 0), (128, 0), (256, 0), (0, 16)), + warp_bases=((0, 1), (0, 2)), + block_bases=[], + shape=[512, 64], + ) + else: + # BLOCK_N == 64: shared by bh64 and bh16bn64 (both bf16 KV). + # K is [512, 64]bf16, KPE is [64, 64]bf16. + blocked_kv: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((1, 0), (2, 0), (4, 0), (0, 8), (0, 4), (0, 16), (0, 32)), + lane_bases=((8, 0), (16, 0), (32, 0), (64, 0), (128, 0), (256, 0)), + warp_bases=((0, 1), (0, 2)), + block_bases=[], + shape=[512, 64], + ) + shared_kv: gl.constexpr = gl.PaddedSharedLayout( + interval_padding_pairs=[[512, 16]], + offset_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [32, 0], [64, 0], [128, 0], [256, 0], [0, 1], [0, 2], [0, 8], [0, 4], [0, 16], [0, 32]], + cga_layout=[], + shape=[512, 64] + ) + blocked_kpe: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((1, 0), (2, 0), (4, 0), (0, 32)), + lane_bases=((8, 0), (16, 0), (32, 0), (0, 4), (0, 8), (0, 16)), + warp_bases=((0, 1), (0, 2)), + block_bases=[], + shape=[64, 64], + ) + shared_kpe: gl.constexpr = gl.PaddedSharedLayout( + interval_padding_pairs=[[512, 16]], + offset_bases=[[1, 0], [2, 0], [4, 0], [8, 0], [16, 0], [32, 0], [0, 4], [0, 8], [0, 16], [0, 1], [0, 2], [0, 32]], + cga_layout=[], + shape=[64, 64] + ) + blocked_page: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((0,),), + lane_bases=((1,), (2,), (4,), (8,), (16,), (32,)), + warp_bases=((0,), (0,)), + block_bases=[], + shape=[64], + ) + blocked_kv_slice: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((1, 0), (2, 0), (4, 0), (0, 8), (0, 4), (0, 16)), + lane_bases=((8, 0), (16, 0), (32, 0), (64, 0), (128, 0), (256, 0)), + warp_bases=((0, 1), (0, 2)), + block_bases=[], + shape=[512, 32], + ) + + # linear_v: each regime has unique warp/reg mapping (bh64 has degenerate warp_bases, + # bh16bn128 has an extra K reg base for the 128-wide K, bh16bn64 has the bh16 warp layout at 64-wide K). + if REGIME == 'bh64': + linear_v: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((0, 1), (0, 2), (0, 4), (0, 32), (16, 0), (32, 0), (64, 0), (128, 0), (256, 0)), + lane_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 8), (0, 16)), + warp_bases=((0, 0), (0, 0)), + block_bases=[], + shape=[512, 64], + ) + elif REGIME == 'bh16bn128': + linear_v: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((0, 1), (0, 2), (0, 4), (0, 32), (0, 64), (64, 0), (128, 0), (256, 0)), + lane_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 8), (0, 16)), + warp_bases=((16, 0), (32, 0)), + block_bases=[], + shape=[512, 128], + ) + else: + linear_v: gl.constexpr = gl.DistributedLinearLayout( + reg_bases=((0, 1), (0, 2), (0, 4), (0, 32), (64, 0), (128, 0), (256, 0)), + lane_bases=((1, 0), (2, 0), (4, 0), (8, 0), (0, 8), (0, 16)), + warp_bases=((16, 0), (32, 0)), + block_bases=[], + shape=[512, 64], + ) + + mfma_layout_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_layout, k_width=8) + mfma_layout_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_layout, k_width=8) + dtype = Q_nope.type.element_ty + kvtype = Kv_c_cache.type.element_ty + ######### layout setting end ######### + + buf_q_nope = gl.allocate_shared_memory(dtype, shape=[BLOCK_H, HEAD_DIM_CKV], layout=shared_q_nope) + if HAS_PE: + buf_q_pe = gl.allocate_shared_memory(dtype, shape=[BLOCK_H, HEAD_DIM_KPE], layout=shared_q_pe) + + # load q_nope + offs_d_ckv = gl.arange(0, HEAD_DIM_CKV, layout=gl.SliceLayout(0, blocked_q_nope)) + cur_head = cur_head_id * BLOCK_H + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blocked_q_nope)) + offs_q_nope = cur_batch * stride_q_nope_bs + cur_head[:, None] * stride_q_nope_h + offs_d_ckv[None, :] + ### For nhead < BLOCK_H, mask OOB heads to zero on Q load and skip OOB O stores; wasted MFMA lanes are free (memory-bound). + gl.amd.cdna4.async_copy.buffer_load_to_shared(buf_q_nope, Q_nope, offs_q_nope, mask = (cur_head < NHEAD)[:, None] if NHEAD < BLOCK_H else None) + gl.amd.cdna4.async_copy.commit_group() + + # load q_pe + if HAS_PE: + offs_d_kpe = gl.arange(0, HEAD_DIM_KPE, layout=gl.SliceLayout(0, blocked_q_pe)) + cur_head_qpe = cur_head_id * BLOCK_H + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blocked_q_pe)) + offs_q_pe = cur_batch * stride_q_pe_bs + cur_head_qpe[:, None] * stride_q_pe_h + offs_d_kpe[None, :] + gl.amd.cdna4.async_copy.buffer_load_to_shared(buf_q_pe, Q_pe, offs_q_pe, mask = (cur_head_qpe < NHEAD)[:, None] if NHEAD < BLOCK_H else None) + gl.amd.cdna4.async_copy.commit_group() + + e_max = gl.zeros([BLOCK_H], dtype=gl.float32, layout=gl.SliceLayout(1, mfma_layout)) - float("inf") + e_sum = gl.zeros([BLOCK_H], dtype=gl.float32, layout=gl.SliceLayout(1, mfma_layout)) + acc = gl.zeros([BLOCK_H, HEAD_DIM_CKV], dtype=gl.float32, layout=mfma_layout) + + # Fold KV dequant scale into the QK temperature. For fp8 KV the real + # logits are (Q @ K_fp8^T) * kv_scale * sm_scale; softmax is shift- but + # not scale-invariant, so kv_scale must affect qk (not just acc). + # For bf16 KV the wrapper passes kv_scale=1.0, so this is a no-op. + qk_scale = sm_scale * kv_scale + + ### bufs of page_number + shared_page: gl.constexpr = gl.SwizzledSharedLayout(vec=1, per_phase=1, max_phase=1, order=[0]) + bufs_page = gl.allocate_shared_memory(gl.int32, shape=[2, BLOCK_N], layout=shared_page) + gl.static_assert(PAGE_SIZE == 1) + + offs_page_raw = gl.arange(0, BLOCK_N, layout=blocked_page) + + ################ prologue + #### global load page number + offs_n_page = start_n + offs_page_raw + offs_page = batch_page_start + offs_n_page // PAGE_SIZE + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_page.index(0), Req_to_tokens, offs_page, offs_n_page < split_kv_end) + gl.amd.cdna4.async_copy.commit_group() + + start_n += BLOCK_N + #### global load page number + offs_n_page = start_n + offs_page_raw + offs_page = batch_page_start + offs_n_page // PAGE_SIZE + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_page.index(1), Req_to_tokens, offs_page, offs_n_page < split_kv_end) + gl.amd.cdna4.async_copy.commit_group() + + #### local load Q + gl.amd.cdna4.async_copy.wait_group(2) + q_nope = gl.amd.cdna4.async_copy.load_shared_relaxed(buf_q_nope, mfma_layout_a) + if HAS_PE: + q_pe = gl.amd.cdna4.async_copy.load_shared_relaxed(buf_q_pe, mfma_layout_a) + + #################### move here to work around allocate_shared_memory bug + bufs_kv = gl.allocate_shared_memory(kvtype, shape=[2, HEAD_DIM_CKV, BLOCK_N], layout=shared_kv) + if HAS_PE: + bufs_kpe = gl.allocate_shared_memory(kvtype, shape=[2, HEAD_DIM_KPE, BLOCK_N], layout=shared_kpe) + + #### global load K + # local load page number + gl.amd.cdna4.async_copy.wait_group(1) + if HAS_PE: + kv_page_number_pe = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page.index(0), gl.SliceLayout(0, blocked_kpe)) + # simplify for page_size 1 + kv_loc_pe = kv_page_number_pe + + # local load page number for slice 0 + bufs_page_0 = bufs_page.index(0).slice(0, BLOCK_N // 2, 0) + kv_page_number_0 = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page_0, gl.SliceLayout(0, blocked_kv_slice)) + kv_loc0 = kv_page_number_0 + + # global load K_nope slice 0 + offs_n_nope0 = split_kv_start + gl.arange(0, BLOCK_N // 2, layout=gl.SliceLayout(0, blocked_kv_slice)) + offs_d_ckv_10 = gl.arange(0, HEAD_DIM_CKV, layout=gl.SliceLayout(1, blocked_kv_slice)) + offs_k_c0 = kv_loc0[None, :] * stride_kv_c_bs + offs_d_ckv_10[:, None] + bufs_kv0 = bufs_kv.index(0).slice(0, BLOCK_N // 2, 1) + if WITHIN_2GB: + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv0, Kv_c_cache, offs_k_c0, mask=offs_n_nope0[None, :] < split_kv_end) + else: + gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv0, Kv_c_cache + offs_k_c0) + gl.amd.cdna4.async_copy.commit_group() + + # global load K_pe + if HAS_PE: + offs_n_pe0 = split_kv_start + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, blocked_kpe)) + offs_d_kpe_1 = gl.arange(0, HEAD_DIM_KPE, layout=gl.SliceLayout(1, blocked_kpe)) + offs_k_pe = kv_loc_pe[None, :] * stride_k_pe_bs + offs_d_kpe_1[:, None] + KV_PE_OFFSET + if WITHIN_2GB: + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kpe.index(0), K_pe_cache, offs_k_pe, mask=offs_n_pe0[None, :] < split_kv_end) + else: + gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kpe.index(0), K_pe_cache + offs_k_pe) + gl.amd.cdna4.async_copy.commit_group() + + # local load page number for slice 1 + bufs_page_1 = bufs_page.index(0).slice(BLOCK_N // 2, BLOCK_N // 2, 0) + kv_page_number_1 = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page_1, gl.SliceLayout(0, blocked_kv_slice)) + kv_loc1 = kv_page_number_1 + + # global load K_nope slice 1 + offs_n_nope1 = offs_n_nope0 + BLOCK_N // 2 + bufs_kv1 = bufs_kv.index(0).slice(BLOCK_N // 2, BLOCK_N // 2, 1) + offs_k_c1 = kv_loc1[None, :] * stride_kv_c_bs + offs_d_ckv_10[:, None] + if WITHIN_2GB: + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv1, Kv_c_cache, offs_k_c1, mask=offs_n_nope1[None, :] < split_kv_end) + else: + gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv1, Kv_c_cache + offs_k_c1) + gl.amd.cdna4.async_copy.commit_group() + + if REGIME == 'bh64': + # bh64 guarantees >= 3 iters/split; this constant-folds the + # `if num_iter >= 2` epilogue-1 guard below so its codegen is unchanged. + gl.assume(num_iter >= 3) + buf_idx = 0 + ################ loop + for i in range(num_iter - 2): + async_idx = (buf_idx + 1) % 2 + + gl.amd.cdna4.async_copy.wait_group(0) + #### global load page number + offs_n_page = start_n + BLOCK_N + offs_page_raw + offs_page = batch_page_start + offs_n_page // PAGE_SIZE + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_page.index(buf_idx), Req_to_tokens, offs_page, offs_n_page < split_kv_end) + gl.amd.cdna4.async_copy.commit_group() + + #### global load K + bufs_kv0 = bufs_kv.index(async_idx).slice(0, BLOCK_N // 2, 1) + bufs_kv1 = bufs_kv.index(async_idx).slice(BLOCK_N // 2, BLOCK_N // 2, 1) + # local load page number for slice 0 + bufs_page_0 = bufs_page.index(async_idx).slice(0, BLOCK_N // 2, 0) + kv_page_number_0 = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page_0, gl.SliceLayout(0, blocked_kv_slice)) + kv_loc0 = kv_page_number_0 + # global load K_nope slice 0 + offs_n_nope0 = start_n + gl.arange(0, BLOCK_N // 2, layout=gl.SliceLayout(0, blocked_kv_slice)) + offs_d_ckv_10 = gl.arange(0, HEAD_DIM_CKV, layout=gl.SliceLayout(1, blocked_kv_slice)) + offs_k_c0 = kv_loc0[None, :] * stride_kv_c_bs + offs_d_ckv_10[:, None] + if WITHIN_2GB: + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv0, Kv_c_cache, offs_k_c0, mask=offs_n_nope0[None, :] < split_kv_end) + else: + # No mask needed on global_load path in the loop body: all + # iterations are guaranteed in-bounds by num_iter arithmetic. + # Only the epilogue uses mask + other=0 for the last + # potentially-partial block. + gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv0, Kv_c_cache + offs_k_c0) + gl.amd.cdna4.async_copy.commit_group() + + # local load page_number_pe + if HAS_PE: + kv_page_number_pe = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page.index(async_idx), gl.SliceLayout(0, blocked_kpe)) + kv_loc_pe = kv_page_number_pe + # global load K_pe + offs_n_pe = start_n + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, blocked_kpe)) + offs_d_kpe_1 = gl.arange(0, HEAD_DIM_KPE, layout=gl.SliceLayout(1, blocked_kpe)) + offs_k_pe = kv_loc_pe[None, :] * stride_k_pe_bs + offs_d_kpe_1[:, None] + KV_PE_OFFSET + if WITHIN_2GB: + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kpe.index(async_idx), K_pe_cache, offs_k_pe, mask=offs_n_pe[None, :] < split_kv_end) + else: + # No mask needed: loop iterations are in-bounds (see KV slice 0 comment). + gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kpe.index(async_idx), K_pe_cache + offs_k_pe) + gl.amd.cdna4.async_copy.commit_group() + + #### dot, softmax, dot (part0) + k_c = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_kv.index(buf_idx), mfma_layout_b) + zeros = gl.zeros([BLOCK_H, BLOCK_N], dtype=gl.float32, layout=mfma_layout) + qk = gl.amd.cdna4.mfma(q_nope, k_c.to(dtype), zeros) + if HAS_PE: + k_pe = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_kpe.index(buf_idx), mfma_layout_b) + qk = gl.amd.cdna4.mfma(q_pe, k_pe.to(dtype), qk) + + # local load page number for slice 1 + bufs_page_1 = bufs_page.index(async_idx).slice(BLOCK_N // 2, BLOCK_N // 2, 0) + kv_page_number_1 = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page_1, gl.SliceLayout(0, blocked_kv_slice)) + kv_loc1 = kv_page_number_1 + # global load K_nope slice 1 + offs_n1 = offs_n_nope0 + BLOCK_N // 2 + offs_k_c1 = kv_loc1[None, :] * stride_kv_c_bs + offs_d_ckv_10[:, None] + if WITHIN_2GB: + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv1, Kv_c_cache, offs_k_c1, mask=offs_n1[None, :] < split_kv_end) + else: + # No mask needed: loop iterations are in-bounds (see KV slice 0 comment). + gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv1, Kv_c_cache + offs_k_c1) + gl.amd.cdna4.async_copy.commit_group() + + #### dot, softmax, dot (part1) + qk *= qk_scale + offs_n_qk = split_kv_start + i * BLOCK_N + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, mfma_layout)) + qk = gl.where(offs_n_qk[None, :] < split_kv_end, qk, float("-inf")) + n_e_max = gl.maximum(gl.max(qk, 1), e_max) + LOG2E: gl.constexpr = 1.4426950408889634 + re_scale = gl.exp2((e_max - n_e_max) * LOG2E) + p = gl.exp2((qk - n_e_max[:, None]) * LOG2E) + e_sum = e_sum * re_scale + gl.sum(p, 1) + e_max = n_e_max + p = p.to(dtype) + p = gl.convert_layout(p, mfma_layout_a) + acc *= re_scale[:, None] + v_c = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_kv.index(buf_idx), linear_v) + v_c = v_c.to(dtype) + v_c = gl.permute(v_c, [1, 0]) + v_c = gl.convert_layout(v_c, mfma_layout_b) + acc = gl.amd.cdna4.mfma(p, v_c, acc) + + start_n += BLOCK_N + buf_idx = (buf_idx + 1) % 2 + + LOG2E: gl.constexpr = 1.4426950408889634 + + ################ epilogue 1 + # Skip when num_iter < 2 (possible for bh16bn64 / bh16bn128 in either mode). + # bh64 has gl.assume(num_iter >= 3) above so the compiler folds this branch + # out there; for the bh16 regimes it stays a runtime branch. + if num_iter >= 2: + async_idx = (buf_idx + 1) % 2 + + #### global load K + # local load page number + gl.amd.cdna4.async_copy.wait_group(3 if HAS_PE else 2) + kv_page_number = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page.index(async_idx), gl.SliceLayout(0, blocked_kv)) + kv_loc = kv_page_number + if HAS_PE: + kv_page_number_pe = gl.amd.cdna4.async_copy.load_shared_relaxed(bufs_page.index(async_idx), gl.SliceLayout(0, blocked_kpe)) + kv_loc_pe = kv_page_number_pe + # global load K_nope + offs_n_nope = start_n + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, blocked_kv)) + offs_d_ckv_1 = gl.arange(0, HEAD_DIM_CKV, layout=gl.SliceLayout(1, blocked_kv)) + offs_k_c = kv_loc[None, :] * stride_kv_c_bs + offs_d_ckv_1[:, None] + if WITHIN_2GB: + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kv.index(async_idx), Kv_c_cache, offs_k_c, mask=offs_n_nope[None, :] < split_kv_end) + else: + # No mask needed: out-of-range positions are discarded by the qk score mask + gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kv.index(async_idx), Kv_c_cache + offs_k_c) + gl.amd.cdna4.async_copy.commit_group() + # global load K_pe + if HAS_PE: + offs_n_pe = start_n + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, blocked_kpe)) + offs_d_kpe_1 = gl.arange(0, HEAD_DIM_KPE, layout=gl.SliceLayout(1, blocked_kpe)) + offs_k_pe = kv_loc_pe[None, :] * stride_k_pe_bs + offs_d_kpe_1[:, None] + KV_PE_OFFSET + if WITHIN_2GB: + gl.amd.cdna4.async_copy.buffer_load_to_shared(bufs_kpe.index(async_idx), K_pe_cache, offs_k_pe, mask=offs_n_pe[None, :] < split_kv_end) + else: + gl.amd.cdna4.async_copy.global_load_to_shared(bufs_kpe.index(async_idx), K_pe_cache + offs_k_pe) + gl.amd.cdna4.async_copy.commit_group() + + # dot, softmax, dot + gl.amd.cdna4.async_copy.wait_group(2 if HAS_PE else 1) + k_c = bufs_kv.index(buf_idx).load(layout=mfma_layout_b) + zeros = gl.zeros([BLOCK_H, BLOCK_N], dtype=gl.float32, layout=mfma_layout) + qk = gl.amd.cdna4.mfma(q_nope, k_c.to(dtype), zeros) + + if HAS_PE: + k_pe = bufs_kpe.index(buf_idx).load(layout=mfma_layout_b) + qk = gl.amd.cdna4.mfma(q_pe, k_pe.to(dtype), qk) + qk *= qk_scale + offs_n_qk = split_kv_start + (num_iter - 2) * BLOCK_N + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, mfma_layout)) + qk = gl.where(offs_n_qk[None, :] < split_kv_end, qk, float("-inf")) + n_e_max = gl.maximum(gl.max(qk, 1), e_max) + re_scale = gl.exp2((e_max - n_e_max) * LOG2E) + p = gl.exp2((qk - n_e_max[:, None]) * LOG2E) + e_sum = e_sum * re_scale + gl.sum(p, 1) + e_max = n_e_max + p = p.to(dtype) + p = gl.convert_layout(p, mfma_layout_a) + acc *= re_scale[:, None] + v_c = bufs_kv.index(buf_idx).load(layout=linear_v) + v_c = v_c.to(dtype) + v_c = gl.permute(v_c, [1, 0]) + v_c = gl.convert_layout(v_c, mfma_layout_b) + acc = gl.amd.cdna4.mfma(p, v_c, acc) + + start_n += BLOCK_N + buf_idx = (buf_idx + 1) % 2 + + ################ epilogue 2 + #### dot, softmax, dot + gl.amd.cdna4.async_copy.wait_group(0) + k_c = bufs_kv.index(buf_idx).load(layout=mfma_layout_b) + zeros = gl.zeros([BLOCK_H, BLOCK_N], dtype=gl.float32, layout=mfma_layout) + qk = gl.amd.cdna4.mfma(q_nope, k_c.to(dtype), zeros) + + if HAS_PE: + k_pe = bufs_kpe.index(buf_idx).load(layout=mfma_layout_b) + qk = gl.amd.cdna4.mfma(q_pe, k_pe.to(dtype), qk) + qk *= qk_scale + offs_n_qk = split_kv_start + (num_iter - 1) * BLOCK_N + gl.arange(0, BLOCK_N, layout=gl.SliceLayout(0, mfma_layout)) + qk = gl.where(offs_n_qk[None, :] < split_kv_end, qk, float("-inf")) + n_e_max = gl.maximum(gl.max(qk, 1), e_max) + re_scale = gl.exp2((e_max - n_e_max) * LOG2E) + p = gl.exp2((qk - n_e_max[:, None]) * LOG2E) + e_sum = e_sum * re_scale + gl.sum(p, 1) + e_max = n_e_max + p = p.to(dtype) + p = gl.convert_layout(p, mfma_layout_a) + acc *= re_scale[:, None] + v_c = bufs_kv.index(buf_idx).load(layout=linear_v) + v_c = v_c.to(dtype) + v_c = gl.permute(v_c, [1, 0]) + v_c = gl.convert_layout(v_c, mfma_layout_b) + acc = gl.amd.cdna4.mfma(p, v_c, acc) + + cur_head_o = cur_head_id * BLOCK_H + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_layout)) + offs_d_ckv_o = gl.arange(0, HEAD_DIM_CKV, layout=gl.SliceLayout(0, mfma_layout)) + offs_o = cur_batch * stride_o_b + cur_head_o[:, None] * stride_o_h + split_kv_id * stride_o_s + offs_d_ckv_o[None, :] + + if HAS_ATTN_SINK: + # Fold the optional per-head sink into the softmax denom (no V contribution). + # e_max/e_sum are natural-log units (the *LOG2E is inside exp2), so is sink. + if NHEAD < BLOCK_H: + sink = gl.load(Attn_sink + cur_head_o, mask=cur_head_o < NHEAD, other=float("-inf")).to(gl.float32) + else: + sink = gl.load(Attn_sink + cur_head_o).to(gl.float32) + n_e_max = gl.maximum(e_max, sink) + re_scale = gl.exp2((e_max - n_e_max) * LOG2E) + acc *= re_scale[:, None] + e_sum = e_sum * re_scale + gl.exp2((sink - n_e_max) * LOG2E) + e_max = n_e_max + + acc *= kv_scale + rcp = 1.0 / e_sum + stored_value = (acc * rcp[:, None]).to(dtype) + if NHEAD < BLOCK_H: + gl.amd.cdna4.buffer_store(stored_value, ptr=O, offsets=offs_o, mask=(cur_head_o < NHEAD)[:, None]) + else: + gl.amd.cdna4.buffer_store(stored_value, ptr=O, offsets=offs_o) + + ### store lse + blocked_lse: gl.constexpr = gl.BlockedLayout(size_per_thread=[1], threads_per_warp=[64], warps_per_cta=[4], order=[0]) + cur_head_lse = cur_head_id * BLOCK_H + gl.arange(0, BLOCK_H, layout=blocked_lse) + if RETURN_LSE and NUM_KV_SPLITS == 1: + # split==1: single split is the whole sequence, so its lse is the final lse. + offs_final_lse = cur_batch * stride_final_lse_b + cur_head_lse * stride_final_lse_h + lse = e_max + gl.log(e_sum) + lse = gl.convert_layout(lse, blocked_lse) + if NHEAD < BLOCK_H: + gl.amd.cdna4.buffer_store(lse, ptr=Final_lse, offsets=offs_final_lse, mask=(cur_head_lse < NHEAD)) + else: + gl.amd.cdna4.buffer_store(lse, ptr=Final_lse, offsets=offs_final_lse) + elif NUM_KV_SPLITS > 1: + # per-split lse for stage-2 reduce. + offs_mid_lse = cur_batch * stride_mid_lse_b + cur_head_lse * stride_mid_lse_h + split_kv_id * stride_mid_lse_s + lse = e_max + gl.log(e_sum) + lse = gl.convert_layout(lse, blocked_lse) + if NHEAD < BLOCK_H: + gl.amd.cdna4.buffer_store(lse, ptr=Mid_lse, offsets=offs_mid_lse, mask=(cur_head_lse < NHEAD)) + else: + gl.amd.cdna4.buffer_store(lse, ptr=Mid_lse, offsets=offs_mid_lse) +# fmt: on + + +# fmt: off +@triton.jit +def _mla_softmax_reducev_kernel( + Logits, + Mid_lse, + O, # noqa: E741 + Final_lse, + B_seq_len, # same seq_info as the decode kernel to derive empty kv splits + stride_l_b, + stride_l_h, + stride_l_s, + stride_ml_b, + stride_ml_h, + stride_ml_s, + stride_o_b, + stride_o_h, + stride_fl_b, + stride_fl_h, + NUM_KV_SPLITS: tl.constexpr, + HEAD_DIM_CKV: tl.constexpr, + HAS_FINAL_LSE: tl.constexpr, + USE_2D_VIEW: tl.constexpr, +): + cur_batch = tl.program_id(0) + cur_head = tl.program_id(1) + + # Recompute this batch's seq len exactly as the decode kernel did, so we can + # rederive which splits are empty. Stage-1 early-returns on empty splits + # (num_iter == 0) and writes nothing, so their logits_buf / mid_lse slots hold + # raw, uninitialised memory - they cannot be loaded or reduced. + if USE_2D_VIEW: + cur_batch_seq_len = tl.load(B_seq_len + cur_batch) + else: + batch_page_start = tl.load(B_seq_len + cur_batch) + cur_batch_seq_len = tl.load(B_seq_len + cur_batch + 1) - batch_page_start + kv_len_per_split = cur_batch_seq_len // NUM_KV_SPLITS + + offs_d_ckv = tl.arange(0, HEAD_DIM_CKV) + offs_l = cur_batch * stride_l_b + cur_head * stride_l_h + offs_d_ckv + offs_ml = cur_batch * stride_ml_b + cur_head * stride_ml_h + + e_sum = 0.0 + e_max = -float("inf") + acc = tl.zeros([HEAD_DIM_CKV], dtype=tl.float32) + + LOOP_START = NUM_KV_SPLITS - 1 if kv_len_per_split == 0 else 0 + for split_kv_id in range(LOOP_START, NUM_KV_SPLITS): + logits = tl.load(Logits + offs_l + split_kv_id * stride_l_s) + logits_1 = tl.load(Mid_lse + offs_ml + split_kv_id * stride_ml_s) + + n_e_max = tl.maximum(logits_1, e_max) + old_scale = tl.where(e_max == -float("inf"), 0.0, tl.exp(e_max - n_e_max)) + acc *= old_scale + exp_logic = tl.where(logits_1 == -float("inf"), 0.0, tl.exp(logits_1 - n_e_max)) + acc += exp_logic * logits + + e_sum = e_sum * old_scale + exp_logic + e_max = n_e_max + + out = acc / e_sum if e_sum > 0.0 else tl.zeros([HEAD_DIM_CKV], dtype=tl.float32) + tl.store( + O + cur_batch * stride_o_b + cur_head * stride_o_h + offs_d_ckv, + out, + ) + if HAS_FINAL_LSE: + tl.store( + Final_lse + cur_batch * stride_fl_b + cur_head * stride_fl_h, + e_max + tl.log(e_sum), + ) +# fmt: on + + +def mla_gluon( + q_nope, # [batch, nhead, kv_lora_rank] + q_pe, # [batch, nhead, qk_rope_head_dim] + # Shared: kv_c=[N, kv_lora_rank+qk_rope_head_dim], k_pe=None, kv_pe_offset=kv_lora_rank + # Split: kv_c=[N, kv_lora_rank], k_pe=[N,qk_rope_head_dim], kv_pe_offset=0 + kv_c, + # final output [batch, nhead, kv_lora_rank]. + o, + page_table, # 2D: block_table [batch, max_seqlen] | 1D: kv_indices [total_kv] + seq_info, # 2D: cache_seqlens [batch] | 1D: kv_indptr [batch+1] + sm_scale, + k_pe=None, + kv_pe_offset=512, + use_2d_view=True, + kv_scale=1.0, + min_kv_seq_len=1, + return_lse=False, + has_pe=True, + attn_sink=None, # [nhead] fp32 per-head sink bias, None means no sink +): + """Unified Gluon MLA entry (gfx950 / CDNA4) — decode and DeepSeek V4 sparse prefill. + + `mla_gluon` supports the full decode (stage-1 + stage-2 reduce, or the stage-1-only + fast path when NUM_KV_SPLITS==1) and writes the final attention into the + caller's `o` ([batch, nhead, kv_lora_rank]). + + return_lse=False (default): returns (o, None). + + return_lse=True: additionally returns the merged log-sum-exp, a separate + fp32 tensor [batch, nhead] + + DSv4 Sparse prefill packs NoPE and RoPE in to one contiguous row (448+64). + To run DSv4 prefill, it requires has_pe=False, prepares valid Q / K in q_nope / kv_c, + and attn_sink, q_pe / k_pe are unused placeholders. + """ + if k_pe is None: + k_pe = kv_c + + batch_size, nhead, head_dim_ckv = q_nope.shape + # Decode carries a real q_pe [.., 64]; + # DSV4 prefill (HAS_PE=False) has no PE, so q_pe may be None and RoPE head_dim is the fixed 64. + head_dim_kpe = q_pe.shape[-1] if has_pe else 64 + if not has_pe: + q_pe = q_nope + k_pe = kv_c + kv_pe_offset = 0 + use_2d_view = False + + assert arch_info.get_arch() == "gfx950", f"mla_gluon requires gfx950 (CDNA4), got {arch_info.get_arch()}" + assert head_dim_ckv == 512, f"mla_gluon requires head_dim_ckv=512, got {head_dim_ckv}" + assert head_dim_kpe == 64, f"mla_gluon requires head_dim_kpe=64, got {head_dim_kpe}" + + # attn sink: decode never sets one; prefill may pass a per-head [H] bias. + has_attn_sink = attn_sink is not None + if attn_sink is None: + attn_sink = torch.empty(1, device=o.device, dtype=torch.float32) # dummy ptr + + # Pick regime by (nhead, kv dtype). + if nhead in (64, 128): + REGIME = "bh64" + elif 1 <= nhead <= 16: + if kv_c.dtype == torch.bfloat16: + REGIME = "bh16bn64" + elif kv_c.dtype == torch.float8_e4m3fn: # gfx950 fp8 (e4m3fn, not e4m3fnuz) + REGIME = "bh16bn128" + else: + raise AssertionError( + f"mla_gluon[bh16*] requires kv_c.dtype in (bfloat16, float8_e4m3fn), got {kv_c.dtype}" + ) + else: + raise AssertionError( + f"mla_gluon requires nhead <= 16 [bh16bn128/bh16bn64] or nhead in (64,128) [bh64], got {nhead}" + ) + + PAGE_SIZE = 1 + + if REGIME == "bh64": + BLOCK_H, BLOCK_N = 64, 64 + NUM_XCDS = get_num_xcds() + # Auto-pick NUM_KV_SPLITS so the launch fills ~256 workgroups (one wave on + # MI350). For the supported (batch, nhead) matrix the result is in {1, 2, 4}. + base_grid = NUM_XCDS * triton.cdiv(nhead, BLOCK_H) * (batch_size // NUM_XCDS) + NUM_KV_SPLITS = max(1, triton.next_power_of_2(triton.cdiv(256, base_grid))) + + assert batch_size % 64 == 0, f"mla_gluon[bh64] requires batch_size divisible by 64, got {batch_size}" + # gl.assume(num_iter > 3) inside the kernel requires every split to have + # > 3*BLOCK_N tokens. Smallest split (last) for batch length s is + # s - (k-1)*ceil(s/k); a sufficient bound is min_kv_seq_len > k*(3*BLOCK_N + k). + min_kv_seq_len_required = NUM_KV_SPLITS * (3 * BLOCK_N + NUM_KV_SPLITS) + assert ( + min_kv_seq_len > min_kv_seq_len_required + ), f"mla_gluon[bh64] requires min_kv_seq_len > {min_kv_seq_len_required} (NUM_KV_SPLITS={NUM_KV_SPLITS}), got {min_kv_seq_len}" + assert ( + q_nope.dtype == torch.bfloat16 and q_pe.dtype == torch.bfloat16 + ), f"q_nope/q_pe must be bf16, got {q_nope.dtype}/{q_pe.dtype}" + assert ( + kv_c.dtype == torch.bfloat16 and k_pe.dtype == torch.bfloat16 + ), f"kv_c/k_pe must be bf16, got {kv_c.dtype}/{k_pe.dtype}" + else: # bh16bn128 (fp8 KV) or bh16bn64 (bf16 KV) + BLOCK_H = 16 + BLOCK_N = 128 if REGIME == "bh16bn128" else 64 + kv_dtype = torch.float8_e4m3fn if REGIME == "bh16bn128" else torch.bfloat16 + NUM_XCDS = 1 # unused by 2-D split grid mapping + # 2-D grid (batch, split). Both bh16 regimes support num_iter in {1, 2, ...} + # (no gl.assume(num_iter >= 3) in the kernel); the only correctness need is + # that every split is non-empty (floor split size = min_kv_seq_len // + # NUM_KV_SPLITS >= 1). Each clamp below keeps NUM_KV_SPLITS <= min_kv_seq_len, + if REGIME == "bh16bn128": + assert batch_size == 1, f"mla_gluon[bh16bn128] requires batch_size=1, got {batch_size}" + NUM_KV_SPLITS = max(1, min(256 // batch_size, min_kv_seq_len)) + else: # bh16bn64 + # Fill ~256 WGs (total WGs = B * NUM_KV_SPLITS <= 256, one MI350 wave), + # but never split a sequence into more blocks than it has: bound by the + # shortest seq's block count so every split holds >= 1 block (no wasted + # partial-block MFMA). For min_kv_seq_len <= BLOCK_N this collapses to + # NUM_KV_SPLITS=1, i.e. one WG per batch computing the whole (short) seq. + NUM_KV_SPLITS = max(1, min(256 // batch_size, triton.cdiv(min_kv_seq_len, BLOCK_N))) + assert ( + q_nope.dtype == torch.bfloat16 and q_pe.dtype == torch.bfloat16 + ), f"q_nope/q_pe must be bf16, got {q_nope.dtype}/{q_pe.dtype}" + assert ( + kv_c.dtype == kv_dtype and k_pe.dtype == kv_dtype + ), f"kv_c/k_pe must be {kv_dtype}, got {kv_c.dtype}/{k_pe.dtype}" + + # buffer_load uses scalar base + 32-bit offsets, limiting addressable range. + # For KV caches > 2 GB the kernel falls back to global_load (64-bit pointers). + max_kv_bytes = kv_c.shape[0] * kv_c.stride(0) * kv_c.element_size() + within_2gb = max_kv_bytes <= 0x80000000 # 2 GB + + if NUM_KV_SPLITS == 1: + # Fast path: stage-1 writes the final attention (and lse) directly to o. + logits_buf = o.view(batch_size, nhead, NUM_KV_SPLITS, head_dim_ckv) + mid_lse = None + stride_mid_lse_b, stride_mid_lse_h, stride_mid_lse_s = 0, 0, 0 + else: + # stage-1 -> per-split (acc, lse); stage-2 reduces into o. + logits_buf = torch.empty( + (batch_size, nhead, NUM_KV_SPLITS, head_dim_ckv), + dtype=o.dtype, + device=o.device, + ) + mid_lse = torch.empty( + (batch_size, nhead, NUM_KV_SPLITS), + dtype=torch.float32, + device=o.device, + ) + stride_mid_lse_b, stride_mid_lse_h, stride_mid_lse_s = mid_lse.stride() + + if return_lse: + final_lse = torch.empty((batch_size, nhead), dtype=torch.float32, device=q_nope.device) + stride_final_lse_b, stride_final_lse_h = final_lse.stride() + else: + final_lse = None + stride_final_lse_b, stride_final_lse_h = 0, 0 + + if REGIME == "bh64": + grid = ( + NUM_XCDS, + triton.cdiv(nhead, BLOCK_H), + (batch_size // NUM_XCDS) * NUM_KV_SPLITS, + ) + else: + grid = (batch_size, NUM_KV_SPLITS) + stride_page_bs = page_table.stride(0) if use_2d_view else 0 + + _mla_gluon[grid]( + q_nope, + q_pe, + kv_c, + k_pe, + page_table, + seq_info, + logits_buf, + attn_sink, + sm_scale, + kv_scale, + q_nope.stride(0), + q_nope.stride(1), + q_pe.stride(0), + q_pe.stride(1), + kv_c.stride(-2), + k_pe.stride(-2), + stride_page_bs, + logits_buf.stride(0), + logits_buf.stride(1), + logits_buf.stride(2), + mid_lse, + stride_mid_lse_b, + stride_mid_lse_h, + stride_mid_lse_s, + final_lse, + stride_final_lse_b, + stride_final_lse_h, + BLOCK_H=BLOCK_H, + BLOCK_N=BLOCK_N, + NUM_KV_SPLITS=NUM_KV_SPLITS, + PAGE_SIZE=PAGE_SIZE, + HEAD_DIM_CKV=head_dim_ckv, + HEAD_DIM_KPE=head_dim_kpe, + KV_PE_OFFSET=kv_pe_offset, + USE_2D_VIEW=use_2d_view, + WITHIN_2GB=within_2gb, + NUM_XCDS=NUM_XCDS, + NHEAD=nhead, + REGIME=REGIME, + RETURN_LSE=return_lse, + HAS_PE=has_pe, + HAS_ATTN_SINK=has_attn_sink, + ) + + if NUM_KV_SPLITS == 1: + # Fast path: stage-1 already wrote o (and lse) directly. + return o, final_lse + + # Stage-2: reduce per-split (acc, lse) into o (and lse when return_lse). + grid_reduce = (batch_size, nhead) + _mla_softmax_reducev_kernel[grid_reduce]( + logits_buf, + mid_lse, + o, + final_lse, + seq_info, + logits_buf.stride(0), + logits_buf.stride(1), + logits_buf.stride(2), + stride_mid_lse_b, + stride_mid_lse_h, + stride_mid_lse_s, + o.stride(0), + o.stride(1), + stride_final_lse_b, + stride_final_lse_h, + NUM_KV_SPLITS=NUM_KV_SPLITS, + HEAD_DIM_CKV=head_dim_ckv, + HAS_FINAL_LSE=return_lse, + USE_2D_VIEW=use_2d_view, + num_warps=8, + ) + + return o, final_lse diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dkv_interm_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dkv_interm_gluon.py new file mode 100644 index 000000000..dc2098b2a --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dkv_interm_gluon.py @@ -0,0 +1,237 @@ +""" +Gluon dKV-intermediate backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X). + +M3 port of the Triton `_bwd_compute_dkv_intermediate`. Key gluon delta vs Triton: +the Triton path materializes a transposed Q/dO in HBM (`q.transpose(1,2).contiguous()`); +this kernel loads Q/dO UNtransposed and transposes in-LDS via `ds_read_*_tr`, removing +the external transpose copy. + +Per program: 1 query token. Grid: (total_tokens,). +Per rank-tile (loop NUM_TILES = R_CHUNK / TILE_K), summed over head groups: + dKV_lora[D_V, TILE_K] = sum_hg ( Q_lora_T @ dS + dO_T @ P ) # contract over heads + dKV_rope[D_ROPE, TILE_K] = sum_hg ( Q_rope_T @ dS ) + store interm[token, rank, :D_QK] + +Q_lora_T/dO_T/Q_rope_T ([D, BLOCK_H]) are the opIdx-0 *transposed* operands -> staged in +LDS, read transposed with ds_read_tr. dS/P ([BLOCK_H, TILE_K]) are opIdx-1 natural-layout +-> register load + convert. M1 config: BLOCK_H=64, TILE_K=64, single-buffered. +""" + +import torch +import triton +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +@gluon.jit +def _sparse_mla_bwd_dkv_interm_gl_kernel( + Q_ptr, # [T, H, D_QK] bf16 (UNtransposed) + dO_ptr, # [T, H, D_V] bf16 (UNtransposed) + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + Interm_ptr, # [T, R_CHUNK, D_QK] bf16 + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + stride_interm_t: tl.int64, + stride_interm_r: tl.int64, + num_heads: tl.int32, + R_CHUNK: gl.constexpr, + TILE_K: gl.constexpr, + BLOCK_H: gl.constexpr, + NUM_HG: gl.constexpr, + D_V: gl.constexpr, + D_ROPE: gl.constexpr, + HAS_ROPE: gl.constexpr, +): + # ===================== constexpr layouts ===================== + # MMA output is [D_V, TILE_K] (and [D_ROPE, TILE_K]); contraction over BLOCK_H heads. + mfma: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + + # ---- Blocked layouts for HBM loads ---- + # Q/dO [BLOCK_H, D_V] : load coalesced then stage to LDS for transpose-read. + _q_tpw_k: gl.constexpr = min(64, D_V // 8) + _q_tpw_m: gl.constexpr = 64 // _q_tpw_k + blk_q: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[_q_tpw_m, _q_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_qrope: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D_ROPE] + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[4, 1], + order=[1, 0], + ) + # dS / P [BLOCK_H, TILE_K] : opIdx-1, register load + convert (no transpose). + blk_ds: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 4], + threads_per_warp=[16, 4], + warps_per_cta=[4, 1], + order=[1, 0], + ) + + # ---- Shared layouts (Q/dO/Q_rope staged for transpose read) ---- + sh_q: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0]) + sh_do: gl.constexpr = gl.PaddedSharedLayout.with_identity_for([[512, 16]], [BLOCK_H, D_V], [1, 0]) + sh_qrope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[1, 0]) + + # ---- Dot operand layouts ---- + dot_qT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_doT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_qropeT_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_ds_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8) + dot_p_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8) + + token_idx = gl.program_id(axis=0) + NUM_TILES: gl.constexpr = R_CHUNK // TILE_K + + # ---- LDS for Q/dO/Q_rope (single-buffered) ---- + smem_q = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_q) + smem_do = gl.allocate_shared_memory(dO_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_do) + if HAS_ROPE: + smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope) + + q_base = token_idx.to(tl.int64) * stride_q_t + do_base = token_idx.to(tl.int64) * stride_do_t + ds_base = token_idx.to(tl.int64) * stride_ds_t + interm_base = token_idx.to(tl.int64) * stride_interm_t + + # store offsets (mfma layout): dKV[d, col] -> interm[token, t*TILE_K+col, d] + offs_d_st = gl.arange(0, D_V, layout=gl.SliceLayout(1, mfma)) + offs_col_st = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma)) + offs_dr_st = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, mfma)) + + for t in range(NUM_TILES): + dKV_lora = gl.zeros([D_V, TILE_K], dtype=gl.float32, layout=mfma) + if HAS_ROPE: + dKV_rope = gl.zeros([D_ROPE, TILE_K], dtype=gl.float32, layout=mfma) + + for hg in range(NUM_HG): + hg_off = hg * BLOCK_H + + # ---- stage Q/dO/Q_rope (this head group) into LDS ---- + offs_h_q = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_q)) + offs_v_q = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_q)) + mask_h_q = offs_h_q < num_heads + q_offs = q_base + offs_h_q[:, None].to(tl.int64) * stride_q_h + offs_v_q[None, :].to(tl.int64) + do_offs = do_base + offs_h_q[:, None].to(tl.int64) * stride_do_h + offs_v_q[None, :].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_q, ptr=Q_ptr, offsets=q_offs.to(tl.int32), mask=mask_h_q[:, None] + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_do, ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=mask_h_q[:, None] + ) + + if HAS_ROPE: + offs_h_qr = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_qr = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_qr = offs_h_qr < num_heads + qr_offs = ( + q_base + + offs_h_qr[:, None].to(tl.int64) * stride_q_h + + (D_V + offs_r_qr[None, :]).to(tl.int64) + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qrope, ptr=Q_ptr, offsets=qr_offs.to(tl.int32), mask=mask_h_qr[:, None] + ) + gl.amd.cdna4.async_copy.commit_group() + + # ---- load dS / P (this tile, this head group) -> dot operands ---- + offs_h_ds = hg_off + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_ds)) + offs_col_ds = t * TILE_K + gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_ds)) + mask_h_ds = offs_h_ds < num_heads + dsp_offs = ( + ds_base + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + offs_col_ds[None, :].to(tl.int64) + ) + dS_blk = gl.amd.cdna4.buffer_load( + ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + P_blk = gl.amd.cdna4.buffer_load( + ptr=P_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + dS_dot = gl.convert_layout(dS_blk, dot_ds_b) + P_dot = gl.convert_layout(P_blk, dot_p_b) + + # ---- wait + transpose-read Q/dO/Q_rope ---- + gl.amd.cdna4.async_copy.wait_group(0) + Q_T = smem_q.permute([1, 0]).load(dot_qT_a) # [D_V, BLOCK_H] + dO_T = smem_do.permute([1, 0]).load(dot_doT_a) # [D_V, BLOCK_H] + + dKV_lora = gl.amd.cdna4.mfma(Q_T, dS_dot, dKV_lora) + dKV_lora = gl.amd.cdna4.mfma(dO_T, P_dot, dKV_lora) + if HAS_ROPE: + Q_rope_T = smem_qrope.permute([1, 0]).load(dot_qropeT_a) # [D_ROPE, BLOCK_H] + dKV_rope = gl.amd.cdna4.mfma(Q_rope_T, dS_dot, dKV_rope) + + # ---- store interm[token, t*TILE_K : +TILE_K, :] (direct from mfma layout) ---- + col = t * TILE_K + offs_col_st + interm_lora_offs = ( + interm_base + col[None, :].to(tl.int64) * stride_interm_r + offs_d_st[:, None].to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dKV_lora.to(Interm_ptr.dtype.element_ty), + ptr=Interm_ptr, + offsets=interm_lora_offs.to(tl.int32), + ) + # dKV_rope is provably zero for the V4 zero-rope-pad (discarded downstream by the + # gather/adapter, which use dkv[..., :D_V]); skip its compute + store when HAS_ROPE=False. + if HAS_ROPE: + interm_rope_offs = ( + interm_base + + col[None, :].to(tl.int64) * stride_interm_r + + (D_V + offs_dr_st[:, None]).to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dKV_rope.to(Interm_ptr.dtype.element_ty), + ptr=Interm_ptr, + offsets=interm_rope_offs.to(tl.int32), + ) + + +def sparse_mla_bwd_dkv_interm_gl(q, do, chunk_dS, chunk_P, R_CHUNK, kv_lora_rank=512, BLOCK_H=32, TILE_K=64): + """ + Gluon dKV-intermediate for one chunk. Takes UNtransposed q/do (transposes in-kernel). + + Returns interm [T, R_CHUNK, D_QK] bf16. + """ + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + assert R_CHUNK % TILE_K == 0 + num_hg = triton.cdiv(num_heads, BLOCK_H) + interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device) + + _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + num_heads, + R_CHUNK=R_CHUNK, + TILE_K=TILE_K, + BLOCK_H=BLOCK_H, + NUM_HG=num_hg, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + num_warps=4, + ) + return interm diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dq_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dq_gluon.py new file mode 100644 index 000000000..78e6f58e1 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_dq_gluon.py @@ -0,0 +1,523 @@ +""" +Gluon dQ backward kernel for DeepSeek V4 sparse MLA (gfx950 / MI355X). + +M1 port of the Triton `_bwd_chunk_dq_store_ds_v4` (V4 chunked_gather dQ) onto the +gluon hardware-control structure of Leon's V3.2 forward. + +Per program: 1 query token x BLOCK_H heads x one rank chunk [R_START, R_START+R_CHUNK). +Grid: (total_tokens, cdiv(num_heads, BLOCK_H)). Dispatch loops chunks externally, +dQ is read-modify-written across chunks (IS_FIRST_CHUNK zero-inits). + +Per-tile math (TILE_K wide, looping NUM_TILES = R_CHUNK / TILE_K): + S = Q_lora @ K_lora_T + Q_rope @ K_rope_T # 2 MMAs (mfma_s), contract over D + P = exp(S*scale - lse) # lse is sink-inclusive (from fwd) + dP = dO @ K_lora_T # 1 MMA (mfma_s), reuses K_lora_T_dot + dS = P * (dP - delta) * scale + dQ_lora += dS @ K_lora # 1 MMA (mfma_acc), K_lora = K_lora_T.T view + dQ_rope += dS @ K_rope # 1 MMA (mfma_acc), K_rope = K_rope_T.T view + store dS, P chunk -> HBM (consumed by dKV-intermediate kernel) + +Differences vs Leon's fwd (the structural template): + * dO added as a second stationary [BH, D_V] operand (async-loaded with Q). + * No online softmax: single P = exp(S - lse) (lse precomputed by fwd). + * 5 MMAs/tile vs 3; K_lora read 3 ways (S, dP, dQ_lora), K_rope 2 ways. + * Per-tile dS/P stores + final dQ store (RMW across chunks) replace the O/LSE write. + * Sink: d_sink is NOT done here (handled by a torch reduction in the launcher), + so this kernel needs no atomics. lse from fwd already folds the sink in. + +M1 config: BLOCK_H=32, TILE_K=16 (matches Triton TILE_K_DQ and the 16x16x16 MFMA +k-dim). LDS at BH=32/TILE_K=16 ~= 104 KB < 160 KB (gfx950). +""" + +import torch +import triton +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +@gluon.jit +def _sparse_mla_bwd_dq_gl_kernel( + Q_ptr, # [T, H, D_QK] bf16 + KV_ptr, # [T, 1, D_QK] bf16 + dO_ptr, # [T, H, D_V] bf16 + TopK_ptr, # [T, TOPK_padded] int32 + LSE_ptr, # [T, H] fp32 (sink-inclusive) + Delta_ptr, # [T, H] fp32 + dQ_ptr, # [T, H, D_QK] bf16 (read-modify-write across chunks) + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_dq_t: tl.int64, + stride_dq_h: tl.int64, + stride_topk_t: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + R_START: tl.int32, + R_CHUNK: gl.constexpr, + BLOCK_H: gl.constexpr, + TILE_K: gl.constexpr, + D_V: gl.constexpr, + D_ROPE: gl.constexpr, + HAS_ROPE: gl.constexpr, + IS_FIRST_CHUNK: gl.constexpr, +): + # ===================== constexpr layouts ===================== + # mfma_s drives S = Q@K_T and dP = dO@K_T, both reducing D_V=512 -> K=32 halves their + # MFMA instruction count vs K=16 (mirrors the fwd R15 win). mfma_acc (dQ += dS@K) stays + # K=16 since its reduction dim is TILE_K (=16). + mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 32], + transposed=True, + warps_per_cta=[4, 1], + ) + mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + + # ---- Blocked layouts for global loads (per Leon's fwd) ---- + _qlora_tpw_k: gl.constexpr = min(64, D_V // 8) + _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k + blk_qlora: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D_V] (Q_lora, dO, dQ_lora) + size_per_thread=[1, 8], + threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_qrope: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D_ROPE] + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[4, 1], + order=[1, 0], + ) + + _klora_tpw_m: gl.constexpr = min(64, D_V // 8) + _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m + blk_klora: gl.constexpr = gl.BlockedLayout( # [D_V, TILE_K] + size_per_thread=[8, 1], + threads_per_warp=[_klora_tpw_m, _klora_tpw_n], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_krope: gl.constexpr = gl.BlockedLayout( # [D_ROPE, TILE_K] + size_per_thread=[2, 1], + threads_per_warp=[32, 2], + warps_per_cta=[1, 4], + order=[0, 1], + ) + # ---- Shared layouts (only K is staged in LDS) ---- + sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [D_V, TILE_K], + [0, 1], + ) + sh_krope: gl.constexpr = gl.SwizzledSharedLayout(vec=8, per_phase=2, max_phase=8, order=[0, 1]) + + # ---- Dot operand layouts ---- + dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_do_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_ds_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4) + dot_klora_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4) + dot_krope_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4) + + # ===================== program ids ===================== + token_idx = gl.program_id(axis=0) + hg_idx = gl.program_id(axis=1) + hg_offset = hg_idx * BLOCK_H + + NUM_TILES: gl.constexpr = R_CHUNK // TILE_K + + # ===================== Q / dO offsets ===================== + offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_qlora = offs_h_qlora < num_heads + q_base = token_idx.to(tl.int64) * stride_q_t + q_offs_lora = ( + q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64) + ) + q_mask_lora = mask_h_qlora[:, None] + + if HAS_ROPE: + offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_qrope = offs_h_qrope < num_heads + q_offs_rope = ( + q_base + + offs_h_qrope[:, None].to(tl.int64) * stride_q_h + + (D_V + offs_r_qrope[None, :]).to(tl.int64) + ) + q_mask_rope = mask_h_qrope[:, None] + + offs_h_do = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_do = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_do = offs_h_do < num_heads + do_base = token_idx.to(tl.int64) * stride_do_t + do_offs = do_base + offs_h_do[:, None].to(tl.int64) * stride_do_h + offs_v_do[None, :].to(tl.int64) + do_mask = mask_h_do[:, None] + + # ===================== load Q_lora, Q_rope, dO -> registers (no LDS staging) ===================== + # Stationary opIdx-0 operands: load HBM->VGPR (blocked, coalesced) then convert to + # the dot-operand layout once. The convert's LDS scratch is transient (freed before + # the K loop), unlike persistent staging, so only K occupies LDS during the loop. + q_lora_blk = gl.amd.cdna4.buffer_load( + ptr=Q_ptr, offsets=q_offs_lora.to(tl.int32), mask=q_mask_lora, other=0.0 + ) + do_blk = gl.amd.cdna4.buffer_load(ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=do_mask, other=0.0) + Q_lora_dot = gl.convert_layout(q_lora_blk, dot_qlora_a) + dO_dot = gl.convert_layout(do_blk, dot_do_a) + if HAS_ROPE: + q_rope_blk = gl.amd.cdna4.buffer_load( + ptr=Q_ptr, offsets=q_offs_rope.to(tl.int32), mask=q_mask_rope, other=0.0 + ) + Q_rope_dot = gl.convert_layout(q_rope_blk, dot_qrope_a) + + # ===================== topk / KV offsets ===================== + topk_base = token_idx.to(tl.int64) * stride_topk_t + R_START + + offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora)) + offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope)) + offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s)) + offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora)) + offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope)) + + # ===================== shared mem for K loop (double-buffered) ===================== + if HAS_ROPE: + smem_krope = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_ROPE, TILE_K], layout=sh_krope) + smem_klora = gl.allocate_shared_memory(KV_ptr.dtype.element_ty, [2, D_V, TILE_K], layout=sh_klora) + + # ===================== dQ accumulators ===================== + # Always zero-init; the read-modify-write across chunks is folded in at STORE + # time in the blocked layout (avoids a big blocked->mfma_acc convert_layout + # whose LDS scratch would overflow the K/Q/dO buffers). + dQ_lora = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc) + if HAS_ROPE: + dQ_rope = gl.zeros([BLOCK_H, D_ROPE], dtype=gl.float32, layout=mfma_acc) + + # ===================== lse / delta (in mfma_s row-slice layout) ===================== + offs_h_s = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + mask_h_s = offs_h_s < num_heads + lse = gl.amd.cdna4.buffer_load( + ptr=LSE_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0 + ) + delta = gl.amd.cdna4.buffer_load( + ptr=Delta_ptr, offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), mask=mask_h_s, other=0.0 + ) + + # ===================== prologue: K tile 0 (group B, buffer 0) ===================== + topk_pos_klora = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + offs_tile_klora).to(tl.int32), + mask=offs_tile_klora < R_CHUNK, + other=-1, + ) + topk_pos_mma = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, offsets=(topk_base + offs_tile_mma).to(tl.int32), mask=offs_tile_mma < R_CHUNK, other=-1 + ) + + valid_klora = topk_pos_klora != -1 + valid_mma = topk_pos_mma != -1 + safe_klora = gl.where(valid_klora, topk_pos_klora, 0) + + klora_offs = safe_klora[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(0), ptr=KV_ptr, offsets=klora_offs.to(tl.int32), mask=valid_klora[None, :] + ) + if HAS_ROPE: + topk_pos_krope = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + offs_tile_krope).to(tl.int32), + mask=offs_tile_krope < R_CHUNK, + other=-1, + ) + valid_krope = topk_pos_krope != -1 + safe_krope = gl.where(valid_krope, topk_pos_krope, 0) + krope_offs = safe_krope[None, :].to(tl.int64) * stride_kv_t + (D_V + offs_r_krope[:, None]).to( + tl.int64 + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(0), ptr=KV_ptr, offsets=krope_offs.to(tl.int32), mask=valid_krope[None, :] + ) + gl.amd.cdna4.async_copy.commit_group() + + # dS / P store offsets in the mfma_s layout (store directly from the compute + # layout -> no convert_layout/LDS shuffle; less-coalesced HBM write instead). + ds_base = token_idx.to(tl.int64) * stride_ds_t + hg_idx.to(tl.int64) * BLOCK_H * stride_ds_h + offs_h_dsp = gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + offs_tile_dsp = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s)) + mask_h_dsp = (hg_offset + offs_h_dsp) < num_heads + + # ===================== main loop: prefetch t+1, compute t ===================== + cur_buf = 0 + for t in range(NUM_TILES - 1): + next_offs_klora = (t + 1) * TILE_K + offs_tile_klora + next_offs_krope = (t + 1) * TILE_K + offs_tile_krope + next_offs_mma = (t + 1) * TILE_K + offs_tile_mma + + topk_pos_klora_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_klora).to(tl.int32), + mask=next_offs_klora < R_CHUNK, + other=-1, + ) + topk_pos_mma_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_mma).to(tl.int32), + mask=next_offs_mma < R_CHUNK, + other=-1, + ) + + valid_klora_next = (next_offs_klora < R_CHUNK) & (topk_pos_klora_next != -1) + valid_mma_next = (next_offs_mma < R_CHUNK) & (topk_pos_mma_next != -1) + safe_klora_next = gl.where(valid_klora_next, topk_pos_klora_next, 0) + + next_buf = 1 - cur_buf + klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to( + tl.int64 + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(next_buf), + ptr=KV_ptr, + offsets=klora_offs_next.to(tl.int32), + mask=valid_klora_next[None, :], + ) + if HAS_ROPE: + topk_pos_krope_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_krope).to(tl.int32), + mask=next_offs_krope < R_CHUNK, + other=-1, + ) + valid_krope_next = (next_offs_krope < R_CHUNK) & (topk_pos_krope_next != -1) + safe_krope_next = gl.where(valid_krope_next, topk_pos_krope_next, 0) + krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + ( + D_V + offs_r_krope[:, None] + ).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(next_buf), + ptr=KV_ptr, + offsets=krope_offs_next.to(tl.int32), + mask=valid_krope_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + + gl.amd.cdna4.async_copy.wait_group(1) + + # ----- read K views from current buffer ----- + klora_smem_cur = smem_klora.index(cur_buf) + K_lora_T_dot = klora_smem_cur.load(dot_klora_b) # [D_V, TILE_K] opIdx1 mfma_s + K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b) # [TILE_K, D_V] opIdx1 mfma_acc + + # ----- S = Q_lora@K_lora_T (+ Q_rope@K_rope_T when HAS_ROPE) ----- + S = gl.amd.cdna4.mfma( + Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + ) + if HAS_ROPE: + krope_smem_cur = smem_krope.index(cur_buf) + K_rope_T_dot = krope_smem_cur.load(dot_krope_b) + S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S) + S = S * scale + offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None] + S = gl.where(valid_mask, S, float("-inf")) + + # ----- P = exp(S - lse) ; dP = dO@K_lora_T ; dS = P*(dP-delta)*scale ----- + P = gl.exp(S - lse[:, None]) + P = gl.where(valid_mask, P, 0.0) + dP = gl.amd.cdna4.mfma( + dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + ) + dS = P * (dP - delta[:, None]) * scale + dS = gl.where(valid_mask, dS, 0.0) + + # ----- dQ_lora += dS@K_lora (+ dQ_rope += dS@K_rope when HAS_ROPE) ----- + dS_bf = dS.to(KV_ptr.dtype.element_ty) + dS_dot = gl.convert_layout(dS_bf, dot_ds_a) + dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora) + if HAS_ROPE: + K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b) + dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope) + + # ----- store dS, P chunk ----- + col = t * TILE_K + offs_tile_dsp + dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64) + gl.amd.cdna4.buffer_store( + stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None] + ) + gl.amd.cdna4.buffer_store( + stored_value=P.to(KV_ptr.dtype.element_ty), + ptr=P_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + + # promote prefetch -> current + cur_buf = next_buf + valid_mma = valid_mma_next + + # ===================== epilogue: last tile ===================== + gl.amd.cdna4.async_copy.wait_group(0) + t = NUM_TILES - 1 + klora_smem_cur = smem_klora.index(cur_buf) + K_lora_T_dot = klora_smem_cur.load(dot_klora_b) + K_lora_v_dot = klora_smem_cur.permute([1, 0]).load(dot_klora_v_b) + + S = gl.amd.cdna4.mfma( + Q_lora_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + ) + if HAS_ROPE: + krope_smem_cur = smem_krope.index(cur_buf) + K_rope_T_dot = krope_smem_cur.load(dot_krope_b) + S = gl.amd.cdna4.mfma(Q_rope_dot, K_rope_T_dot, S) + S = S * scale + offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None] + S = gl.where(valid_mask, S, float("-inf")) + + P = gl.exp(S - lse[:, None]) + P = gl.where(valid_mask, P, 0.0) + dP = gl.amd.cdna4.mfma(dO_dot, K_lora_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s)) + dS = P * (dP - delta[:, None]) * scale + dS = gl.where(valid_mask, dS, 0.0) + + dS_bf = dS.to(KV_ptr.dtype.element_ty) + dS_dot = gl.convert_layout(dS_bf, dot_ds_a) + dQ_lora = gl.amd.cdna4.mfma(dS_dot, K_lora_v_dot, dQ_lora) + if HAS_ROPE: + K_rope_v_dot = krope_smem_cur.permute([1, 0]).load(dot_krope_v_b) + dQ_rope = gl.amd.cdna4.mfma(dS_dot, K_rope_v_dot, dQ_rope) + + col = t * TILE_K + offs_tile_dsp + dsp_offs = ds_base + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + col[None, :].to(tl.int64) + gl.amd.cdna4.buffer_store( + stored_value=dS_bf, ptr=dS_ptr, offsets=dsp_offs.to(tl.int32), mask=mask_h_dsp[:, None] + ) + gl.amd.cdna4.buffer_store( + stored_value=P.to(KV_ptr.dtype.element_ty), + ptr=P_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + + # ===================== store dQ (lora + rope) ===================== + dq_base = token_idx.to(tl.int64) * stride_dq_t + offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_o = offs_h_o < num_heads + dq_offs_lora = dq_base + offs_h_o[:, None].to(tl.int64) * stride_dq_h + offs_v_o[None, :].to(tl.int64) + dq_lora_blk = gl.convert_layout(dQ_lora.to(dQ_ptr.dtype.element_ty), blk_qlora) + if not IS_FIRST_CHUNK: + prev_lora = gl.amd.cdna4.buffer_load( + ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None], other=0.0 + ) + dq_lora_blk = (dq_lora_blk.to(gl.float32) + prev_lora.to(gl.float32)).to(dQ_ptr.dtype.element_ty) + gl.amd.cdna4.buffer_store( + stored_value=dq_lora_blk, ptr=dQ_ptr, offsets=dq_offs_lora.to(tl.int32), mask=mask_h_o[:, None] + ) + + # dQ_rope is provably zero for the V4 zero-rope-pad (the adapter discards dq[..., D_V:]), + # so skip its accumulation + store entirely when HAS_ROPE is False. + if HAS_ROPE: + offs_h_or = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_or = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_or = offs_h_or < num_heads + dq_offs_rope = ( + dq_base + offs_h_or[:, None].to(tl.int64) * stride_dq_h + (D_V + offs_r_or[None, :]).to(tl.int64) + ) + dq_rope_blk = gl.convert_layout(dQ_rope.to(dQ_ptr.dtype.element_ty), blk_qrope) + if not IS_FIRST_CHUNK: + prev_rope = gl.amd.cdna4.buffer_load( + ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None], other=0.0 + ) + dq_rope_blk = (dq_rope_blk.to(gl.float32) + prev_rope.to(gl.float32)).to(dQ_ptr.dtype.element_ty) + gl.amd.cdna4.buffer_store( + stored_value=dq_rope_blk, ptr=dQ_ptr, offsets=dq_offs_rope.to(tl.int32), mask=mask_h_or[:, None] + ) + + +# ===================================================================== +# Launcher — runs the dQ pass only (chunk loop + RMW), returns dq, chunk dS/P. +# d_sink (if needed) is a torch reduction handled by the caller. +# ===================================================================== +def sparse_mla_bwd_dq_gl( + q, + kv, + do, + topk_indices_padded, + lse, + delta, + R_CHUNK, + topk, + kv_lora_rank=512, + scale=None, + BLOCK_H=64, + TILE_K=16, +): + """ + Gluon dQ pass. Mirrors the dQ portion of `sparse_mla_bwd_v4`'s chunk loop. + + Returns: + dq: [T, H, D_QK] bf16 (fully accumulated across chunks) + chunk_dS: [T, H, R_CHUNK] bf16 (LAST chunk's dS — for spot validation) + chunk_P: [T, H, R_CHUNK] bf16 (LAST chunk's P) + """ + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + if scale is None: + scale = 1.0 / (d_qk**0.5) + assert R_CHUNK % TILE_K == 0, "TILE_K must divide R_CHUNK" + + dq = torch.empty_like(q) + chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + + num_hg = triton.cdiv(num_heads, BLOCK_H) + grid = (total_tokens, num_hg) + + for r_start in range(0, topk, R_CHUNK): + is_first = r_start == 0 + _sparse_mla_bwd_dq_gl_kernel[grid]( + q, + kv, + do, + topk_indices_padded, + lse, + delta, + dq, + chunk_dS, + chunk_P, + q.stride(0), + q.stride(1), + kv.stride(0), + do.stride(0), + do.stride(1), + dq.stride(0), + dq.stride(1), + topk_indices_padded.stride(0), + chunk_dS.stride(0), + chunk_dS.stride(1), + scale, + num_heads, + r_start, + R_CHUNK=R_CHUNK, + BLOCK_H=BLOCK_H, + TILE_K=TILE_K, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + IS_FIRST_CHUNK=is_first, + num_warps=4, + waves_per_eu=1, + ) + return dq, chunk_dS, chunk_P diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_v4_gluon.py new file mode 100644 index 000000000..458adaf9c --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_bwd_v4_gluon.py @@ -0,0 +1,177 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Gluon DeepSeek-V4 sparse-MLA backward for the "gluon_v2" backend. + +Companion to the gluon_v2 forward (:func:`sparse_mla_fwd_v4_gluon_v2`). Wires the Gluon +dQ + Gluon dKV-intermediate compute kernels with a Triton Delta preprocess and the +backend-neutral CSR inverted-topk gather + torch d_sink reduction (non-atomic +chunked-gather scheme). + +The dQ / dKV-intermediate Gluon kernels apply the forward campaign's accepted techniques +to the backward: rope-skip (the V4 zero-rope-pad makes the rope gradients provably zero) ++ MFMA K=32 for the D_V=512-reduction matmuls, plus a single-chunk dQ read-modify-write +for high head counts. Beats the plain-Triton backward ~1.12x geomean over the 6 +flash/pro x cr{0,4,128} shapes (eager-UT 9/9). +""" + +import torch +import triton + +from .._gluon_dsa._dsa_bwd_gather import _build_inverted_topk_slice, _bwd_dkv_gather_acc +from .._gluon_dsa._dsa_bwd_preprocess import _sparse_mla_bwd_preprocess +from .dsa_bwd_dkv_interm_gluon import _sparse_mla_bwd_dkv_interm_gl_kernel +from .dsa_bwd_dq_gluon import _sparse_mla_bwd_dq_gl_kernel + + +def sparse_mla_bwd_v4_gluon_v2(q, kv, o, do, topk_indices, lse, attn_sink=None, kv_lora_rank=512, scale=None): + """DeepSeek-V4 sparse-MLA backward (Gluon dQ/dKV). Returns ``(dq, dkv, d_sink)``.""" + assert q.is_contiguous() and kv.is_contiguous() and o.is_contiguous() + assert do.is_contiguous() and topk_indices.is_contiguous() and lse.is_contiguous() + + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + topk = topk_indices.shape[1] + if scale is None: + scale = 1.0 / (d_qk**0.5) + if kv.dim() == 2: + kv = kv.unsqueeze(1) + num_kv = kv.shape[0] + + has_sink = attn_sink is not None + if has_sink: + assert attn_sink.dtype == torch.float32 and attn_sink.shape == (num_heads,) + + # ---- preprocess: Delta = rowsum(O*dO) (Triton, unchanged) ---- + delta = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device) + BLOCK_H_PRE = triton.next_power_of_2(min(64, num_heads)) + _sparse_mla_bwd_preprocess[(total_tokens, triton.cdiv(num_heads, BLOCK_H_PRE))]( + O_ptr=o, + dO_ptr=do, + Delta_ptr=delta, + stride_o_t=o.stride(0), + stride_o_h=o.stride(1), + num_heads=num_heads, + D_V=kv_lora_rank, + BLOCK_H=BLOCK_H_PRE, + ) + + # ---- config ---- + # R2: dQ is read-modify-written across chunks, so more chunks = more redundant dq + # reload passes + repeated CSR builds. H=64 CSA has TOPK=640, so the old 256 cap + # split it into 3 chunks; a 320 cap tests a two-chunk schedule without changing + # high-head whole-topk behavior. + if num_heads >= 128: + R_CHUNK = min(topk, 1536) + elif num_heads >= 64: + R_CHUNK = min(topk, 320) + else: + R_CHUNK = min(256, topk) + BH_DQ, TK_DQ = 64, 16 + BH_DKV, TK_DKV = 32, 64 + num_hg_dq = triton.cdiv(num_heads, BH_DQ) + num_hg_dkv = triton.cdiv(num_heads, BH_DKV) + + dq = torch.empty_like(q) + chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + dkv_acc = torch.zeros(num_kv, d_qk, dtype=torch.float32, device=q.device) + interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device) + + # ---- pad topk to R_CHUNK multiple ---- + topk_padded_len = ((topk + R_CHUNK - 1) // R_CHUNK) * R_CHUNK + if topk_padded_len != topk: + pad = torch.full((total_tokens, topk_padded_len - topk), -1, dtype=torch.int32, device=q.device) + topk_padded = torch.cat([topk_indices, pad], dim=1).contiguous() + else: + topk_padded = topk_indices + + all_csr = [ + _build_inverted_topk_slice(topk_padded[:, rs : rs + R_CHUNK], rs, R_CHUNK, num_kv=num_kv) + for rs in range(0, topk, R_CHUNK) + ] + + for chunk_idx, r_start in enumerate(range(0, topk, R_CHUNK)): + is_first = r_start == 0 + + _sparse_mla_bwd_dq_gl_kernel[(total_tokens, num_hg_dq)]( + q, + kv, + do, + topk_padded, + lse, + delta, + dq, + chunk_dS, + chunk_P, + q.stride(0), + q.stride(1), + kv.stride(0), + do.stride(0), + do.stride(1), + dq.stride(0), + dq.stride(1), + topk_padded.stride(0), + chunk_dS.stride(0), + chunk_dS.stride(1), + scale, + num_heads, + r_start, + R_CHUNK=R_CHUNK, + BLOCK_H=BH_DQ, + TILE_K=TK_DQ, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_ROPE=False, # V4 zero-rope-pad: dQ_rope is provably zero + discarded by the adapter + IS_FIRST_CHUNK=is_first, + num_warps=4, + waves_per_eu=1, + ) + + _sparse_mla_bwd_dkv_interm_gl_kernel[(total_tokens,)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + num_heads, + R_CHUNK=R_CHUNK, + TILE_K=TK_DKV, + BLOCK_H=BH_DKV, + NUM_HG=num_hg_dkv, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_ROPE=False, # V4 zero-rope-pad: dKV_rope provably zero + discarded downstream + num_warps=4, + ) + + inv_ptr, inv_data = all_csr[chunk_idx] + _bwd_dkv_gather_acc[(num_kv,)]( + interm, + inv_ptr, + inv_data, + dkv_acc, + interm.stride(1), + dkv_acc.stride(0), + D_V=kv_lora_rank, + D_ROPE=rope_rank, + num_warps=4, + ) + + d_sink = None + if has_sink: + d_sink = -(torch.exp(attn_sink.unsqueeze(0) - lse) * delta).sum(0) + + dkv_out = dkv_acc.to(kv.dtype).unsqueeze(1) + return dq, dkv_out, d_sink diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_fwd_v4_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_fwd_v4_gluon.py new file mode 100644 index 000000000..5c89b83e5 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_gluon_v3/dsa_fwd_v4_gluon.py @@ -0,0 +1,733 @@ +""" +Gluon forward for DeepSeek V4 sparse MLA (gfx950 / CDNA4), with attention sink. + +Based on Leon's (leonling-ll) V3.2 gluon forward from `leonling-ll/aiter` branch +`liyang/dsa` -- which adapted our V3.2 Triton forward and added the gfx950 hardware +control (MFMA4 layouts, padded/swizzled shared, double-buffered K, async DMA pipeline, +ds_read_tr transpose, explicit dot-operand layouts); see also his DSA PR ROCm/aiter#3456. +This file adds the V4 attention-sink epilogue (sink-inclusive LSE) so it matches +`sparse_mla_fwd_v4`; it is the forward of the "gluon_v2" backend. + +Optimizations accepted over the base gluon forward (gfx950 campaign): + * rope-skip (HAS_ROPE=False): the V4 latent bakes RoPE in-place over the 512, so the + rope QK term is provably zero -- skip the 64-wide rope MFMA + K_rope loads entirely. + * exp2 softmax: fold log2(e) into the QK scale so the per-element exp is one hardware + exp2 (m_i/l_i in log2 units; LSE converted back to natural log for the backward). + * MFMA K=32 for the QK score matmul (instr_shape [16,16,32]): the score reduces D_V=512, + so K=32 halves the QK MFMA instruction count vs K=16 (PV/acc stays K=16, TILE_K-bound). + * register-prefetched topk index (off the KV-gather critical path) + async double-buffered + K gather overlapping the QK/softmax/PV MFMAs. + +Pipeline: + Prologue: Q -> shared (async); K tile 0 -> shared (async, double-buffered); deep-prefetch topk. + Loop tile t: gather tile t+2 (async) while computing QK[t+1] + softmax(t) + PV(t); promote. + Epilogue: drain; fold sink into the denominator (V4); write O, LSE. +""" + +import functools + +import torch +import triton +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + +from .aiter_lse_fwd import sparse_mla_fwd_v4_aiter_lse_csa_formula + +# --------------------------------------------------------------------------- +# Triton capability gate: the gluon_v2 forward is a Gluon kernel (the backend's backward +# is currently plain-Triton, being migrated to Gluon). The Gluon fwd needs a Gluon-capable triton whose CDNA4 +# async_copy accepts arbitrary (DistributedLinearLayout) offsets. Released +# triton 3.7.0/3.7.1 still restricts async_copy offsets to Blocked/Slice and +# will NOT compile this path; build triton from the commit below. +# --------------------------------------------------------------------------- +_GLUON_V2_REQUIRED_COMMIT = "09500db9f0" +_GLUON_V2_INSTALL_HINT = ( + "gluon_v2 forward (Gluon) requires a Gluon-capable triton whose CDNA4 " + "async_copy accepts general offset layouts. The installed triton ({ver}) does not (released " + "3.7.0/3.7.1 restrict async_copy offsets to BlockedLayout/SliceLayout).\n" + "Build & install triton-lang/triton @ commit " + _GLUON_V2_REQUIRED_COMMIT + ":\n" + " git clone https://github.com/triton-lang/triton.git third_party/triton\n" + " cd third_party/triton && git checkout " + _GLUON_V2_REQUIRED_COMMIT + "\n" + " pip install -r python/requirements.txt\n" + " TRITON_CODEGEN_BACKENDS=amd MAX_JOBS=128 pip wheel --no-build-isolation --no-deps . -w dist\n" + " pip install --force-reinstall --no-deps dist/triton-*.whl" +) + + +@functools.lru_cache(maxsize=1) +def _gluon_available() -> bool: + """True iff triton exposes the experimental CDNA4 Gluon dialect this fwd uses.""" + try: + from triton.experimental import gluon as _gl # noqa: F401 + from triton.experimental.gluon.language import amd as _amd + + return hasattr(_amd, "cdna4") + except Exception: # noqa: BLE001 + return False + + +def _require_gluon_v2_triton() -> None: + """Fail fast with a build hint when Gluon is unavailable. The kernel compile is + ALSO guarded at the launch site: a Gluon-capable-but-incompatible triton raises a + CompilationError there, which is re-wrapped with the same install hint (so the user + always gets the commit rather than a raw layout/compile error).""" + if not _gluon_available(): + raise RuntimeError(_GLUON_V2_INSTALL_HINT.format(ver=getattr(triton, "__version__", "unknown"))) + + +def _wrap_compile_error(exc: Exception) -> RuntimeError: + return RuntimeError( + _GLUON_V2_INSTALL_HINT.format(ver=getattr(triton, "__version__", "unknown")) + + f"\n(triton failed to compile the Gluon fwd: {type(exc).__name__}: " + + (str(exc).splitlines()[0] if str(exc).strip() else "") + + ")" + ) + + +# ===================================================================== +# Utility +# ===================================================================== +def _get_lds_limit(): + """Return the per-CU LDS limit in bytes for the current GPU. + + gfx942 (MI300X): 64 KB = 65536 bytes + gfx950 (MI355X): 160 KB = 163840 bytes + """ + if torch.cuda.is_available(): + prop = torch.cuda.get_device_properties(0) + gcn_arch = getattr(prop, "gcnArchName", "") + if "gfx950" in gcn_arch: + return 163840 + return 65536 + + +_LDS_LIMIT = _get_lds_limit() + + +# ===================================================================== +# Forward — autotune configs and pruning +# ===================================================================== +def _fwd_prune_configs(configs, named_args, **kwargs): + """Prune autotune configs that would exceed per-CU LDS.""" + D_V = kwargs.get("D_V", named_args.get("D_V")) + D_ROPE = kwargs.get("D_ROPE", named_args.get("D_ROPE")) + pruned = [] + for config in configs: + config.kwargs["BLOCK_H"] + tk = config.kwargs["TILE_K"] + ns = config.num_stages + kv_lds = (D_V + D_ROPE) * tk * 2 * ns + if kv_lds <= _LDS_LIMIT: + pruned.append(config) + if not pruned: + pruned.append(configs[0]) + return pruned + + +def _get_fwd_autotune_configs(): + configs = [ + triton.Config( + {"BLOCK_H": BLOCK_H, "TILE_K": TILE_K, "waves_per_eu": WPE}, + num_warps=nw, + ) + for BLOCK_H in [16, 32, 64] + for TILE_K in [16, 32, 64, 128] + for WPE in [0, 1, 2] + for nw in [4] # num_warps must be 4 to align with kernel implementation + ] + # configs = [triton.Config({"BLOCK_H": 64, "TILE_K": 32, "waves_per_eu": 0}, num_warps=4),] + return configs + + +@triton.autotune( + configs=_get_fwd_autotune_configs(), + key=["num_heads", "TOPK", "D_V", "D_ROPE"], + prune_configs_by={"early_config_prune": _fwd_prune_configs}, +) +@gluon.jit +def _sparse_mla_fwd_gl_v2_kernel( + Q_ptr, # [total_tokens, num_heads, D_QK] bf16 + KV_ptr, # [total_tokens, 1, D_QK] bf16 + TopK_ptr, # [total_tokens, TOPK] int32 + Sink_ptr, # [num_heads] fp32; ignored if HAS_SINK == False + O_ptr, # [total_tokens, num_heads, D_V] bf16 + LSE_ptr, # [total_tokens, num_heads] fp32 (sink-inclusive if HAS_SINK) + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_o_t: tl.int64, + stride_o_h: tl.int64, + stride_topk_t: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + TOPK: gl.constexpr, + BLOCK_H: gl.constexpr, + TILE_K: gl.constexpr, + D_V: gl.constexpr, + D_ROPE: gl.constexpr, + HAS_SINK: gl.constexpr, + HAS_ROPE: gl.constexpr, +): + # ---------- constexpr layouts ---------- + # QK MFMA uses K=32 (aiter): the score matmul reduces D_V=512, so instr_shape=[16,16,32] + # halves the MFMA instruction count vs [16,16,16]. mfma_acc (PV) stays K=16 since its + # reduction dim is TILE_K (autotuned down to 16). + mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 32], + transposed=True, + warps_per_cta=[4, 1], + ) + mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + + # Blocked layouts for global loads. + _qlora_tpw_k: gl.constexpr = min(64, D_V // 8) + _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k + blk_qlora: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_qrope: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[4, 1], + order=[1, 0], + ) + + _klora_tpw_m: gl.constexpr = min(64, D_V // 8) + _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m + blk_klora: gl.constexpr = gl.BlockedLayout( # [D_V, TILE_K] + size_per_thread=[8, 1], + threads_per_warp=[_klora_tpw_m, _klora_tpw_n], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_krope: gl.constexpr = gl.BlockedLayout( # [D_ROPE, TILE_K] = [64, 16] + size_per_thread=[2, 1], + threads_per_warp=[32, 2], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_topk: gl.constexpr = gl.BlockedLayout( # [TILE_K] int32 + size_per_thread=[1], + threads_per_warp=[64], + warps_per_cta=[4], + order=[0], + ) + blk_lse: gl.constexpr = gl.BlockedLayout( # [BLOCK_H] fp32 + size_per_thread=[1], + threads_per_warp=[64], + warps_per_cta=[4], + order=[0], + ) + + # Shared layouts. + sh_qlora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [BLOCK_H, D_V], + [1, 0], + ) + sh_qrope: gl.constexpr = gl.SwizzledSharedLayout( + vec=8, + per_phase=2, + max_phase=8, + order=[1, 0], + ) + sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [D_V, TILE_K], + [0, 1], + ) + sh_krope: gl.constexpr = gl.SwizzledSharedLayout( + vec=8, + per_phase=2, + max_phase=8, + order=[0, 1], + ) + + # Dot operand layouts + dot_qlora_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_qrope_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_s, k_width=8) + dot_klora_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_krope_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_s, k_width=8) + dot_p_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma_acc, k_width=4) + dot_v_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma_acc, k_width=4) + + # ---------- program ids ---------- + token_idx = gl.program_id(axis=0) + hg_idx = gl.program_id(axis=1) + hg_offset = hg_idx * BLOCK_H + + # ---------- offsets for Q ---------- + # Q_lora [BLOCK_H, D_V] + offs_h_qlora = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_qlora = offs_h_qlora < num_heads + + q_base = token_idx.to(tl.int64) * stride_q_t + q_offs_lora = ( + q_base + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + offs_v_qlora[None, :].to(tl.int64) + ) + q_mask_lora = mask_h_qlora[:, None] + + smem_qlora = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_V], layout=sh_qlora) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qlora, + ptr=Q_ptr, + offsets=q_offs_lora.to(tl.int32), + mask=q_mask_lora, + ) + # V4 zero-rope-pad: skip the rope Q load + rope MFMA entirely when HAS_ROPE is False + # (RoPE is baked in-place over the 512 latent, so the rope QK term is provably zero). + if HAS_ROPE: + # Q_rope [BLOCK_H, D_ROPE] + offs_h_qrope = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qrope)) + offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_qrope = offs_h_qrope < num_heads + q_offs_rope = ( + q_base + + offs_h_qrope[:, None].to(tl.int64) * stride_q_h + + (D_V + offs_r_qrope[None, :]).to(tl.int64) + ) + q_mask_rope = mask_h_qrope[:, None] + smem_qrope = gl.allocate_shared_memory(Q_ptr.dtype.element_ty, [BLOCK_H, D_ROPE], layout=sh_qrope) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qrope, + ptr=Q_ptr, + offsets=q_offs_rope.to(tl.int32), + mask=q_mask_rope, + ) + gl.amd.cdna4.async_copy.commit_group() + + # ---------- topk and KV offsets ---------- + NUM_TILES: gl.constexpr = (TOPK + TILE_K - 1) // TILE_K + topk_base = token_idx.to(tl.int64) * stride_topk_t + + # offs_tile in three layouts (sliced from each of the three loaders) + offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora)) + offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope)) + offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s)) + offs_tile_topk = gl.arange(0, TILE_K, layout=blk_topk) + + offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora)) + offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope)) + + # (removed dead `topk_pos_reg` prologue load — was never consumed) + + # ---------- shared mem allocations for the K loop ---------- + if HAS_ROPE: + smem_krope = gl.allocate_shared_memory( + KV_ptr.dtype.element_ty, + [2, D_ROPE, TILE_K], + layout=sh_krope, + ) + smem_klora = gl.allocate_shared_memory( + KV_ptr.dtype.element_ty, + [2, D_V, TILE_K], + layout=sh_klora, + ) + + # ---------- accumulators ---------- + m_i = gl.full([BLOCK_H], float("-inf"), dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s)) + l_i = gl.full([BLOCK_H], 0.0, dtype=gl.float32, layout=gl.SliceLayout(1, mfma_s)) + acc = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc) + # exp2 softmax (aiter technique): fold log2(e) into the QK scale so the per-element + # softmax exp becomes a single hardware exp2 (no per-element *log2e). m_i / l_i are then + # in log2 units; lse is converted back to natural log at the epilogue (the bwd needs nat-log). + scale_log2 = scale * 1.4426950408889634 + + # ---------- tile-0 prefetch (prologue) ---------- + # Load K_lora and K_rope for tile 0. + topk_pos_klora = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + offs_tile_klora, + mask=offs_tile_klora < TOPK, + other=-1, + ) + if HAS_ROPE: + topk_pos_krope = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + offs_tile_krope, + mask=offs_tile_krope < TOPK, + other=-1, + ) + topk_pos_mma = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + offs_tile_mma, + mask=offs_tile_mma < TOPK, + other=-1, + ) + + # Deep-prefetch tile-1 topk ONCE in the neutral blk_topk layout (DEDUP). Carried as a + # single register set; converted to the klora/krope/mma layouts at point of use, to + # minimize carried register pressure (the 3-layout carry caused an acc-rescale codegen + # regression -- see att_fwd_gluon_mi350/RESULTS.md). + p1_off_topk = TILE_K + offs_tile_topk + tkraw = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + p1_off_topk, + mask=p1_off_topk < TOPK, + other=-1, + ) + + valid_klora = topk_pos_klora != -1 # tile_start=0 -> offs_tile buf1, drain K[0], QK[0] -> S_prev (no softmax/PV yet). + tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora)) + tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s)) + valid_klora_next = ((TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1) + valid_qk = ((TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1) + safe_klora_next = gl.where(valid_klora_next, tk_klora, 0) + klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(1), + ptr=KV_ptr, + offsets=klora_offs_next.to(tl.int32), + mask=valid_klora_next[None, :], + ) + if HAS_ROPE: + tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope)) + valid_krope_next = ((TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1) + safe_krope_next = gl.where(valid_krope_next, tk_krope, 0) + krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + ( + D_V + offs_r_krope[:, None] + ).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(1), + ptr=KV_ptr, + offsets=krope_offs_next.to(tl.int32), + mask=valid_krope_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + tkraw = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + (2 * TILE_K + offs_tile_topk), + mask=(2 * TILE_K + offs_tile_topk) < TOPK, + other=-1, + ) + gl.amd.cdna4.async_copy.wait_group(1) + S_prev = gl.amd.cdna4.mfma( + Q_lora_dot, + smem_klora.index(0).load(dot_klora_b), + gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s), + ) + if HAS_ROPE: + S_prev = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(0).load(dot_krope_b), S_prev) + S_prev = S_prev * scale_log2 + S_prev = gl.where(valid_mma[None, :] & mask_h_mma[:, None], S_prev, float("-inf")) + cur_buf = 1 + + for t in range(NUM_TILES - 2): + gl.amd.cdna4.async_copy.wait_group(0) # drain K[t+1] (cur_buf) before QK reads it + # 2-BUFFER EARLY-GATHER: evacuate V[t] from pv_buf into REGISTERS first, freeing that buffer, + # then gather tile t+2 into it BEFORE the QK/PV MFMAs so the DMA overlaps both (no 3rd buffer). + # V_lora_dot in regs => no read/async-write race on the recycled buffer. Costs VGPR live range. + V_lora_dot = smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b) + tk_klora = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_klora)) + tk_mma = gl.convert_layout(tkraw, gl.SliceLayout(0, mfma_s)) + valid_klora_next = (((t + 2) * TILE_K + offs_tile_klora) < TOPK) & (tk_klora != -1) + valid_qk_next = (((t + 2) * TILE_K + offs_tile_mma) < TOPK) & (tk_mma != -1) + safe_klora_next = gl.where(valid_klora_next, tk_klora, 0) + klora_offs_next = safe_klora_next[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[:, None].to( + tl.int64 + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(1 - cur_buf), + ptr=KV_ptr, + offsets=klora_offs_next.to(tl.int32), + mask=valid_klora_next[None, :], + ) + if HAS_ROPE: + tk_krope = gl.convert_layout(tkraw, gl.SliceLayout(0, blk_krope)) + valid_krope_next = (((t + 2) * TILE_K + offs_tile_krope) < TOPK) & (tk_krope != -1) + safe_krope_next = gl.where(valid_krope_next, tk_krope, 0) + krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + ( + D_V + offs_r_krope[:, None] + ).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(1 - cur_buf), + ptr=KV_ptr, + offsets=krope_offs_next.to(tl.int32), + mask=valid_krope_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + tkraw_n = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=topk_base.to(tl.int32) + ((t + 3) * TILE_K + offs_tile_topk), + mask=((t + 3) * TILE_K + offs_tile_topk) < TOPK, + other=-1, + ) + # QK tile (t+1) from cur_buf -- matrix; overlaps the gather above + softmax below + S_cur = gl.amd.cdna4.mfma( + Q_lora_dot, + smem_klora.index(cur_buf).load(dot_klora_b), + gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s), + ) + if HAS_ROPE: + S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur) + S_cur = S_cur * scale_log2 + S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf")) + # softmax(S_prev = tile t) [VALU, overlaps QK] + m_j = gl.max(S_prev, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > float("-inf"), m_new, 0.0) + alpha = gl.exp2(m_i - m_new) + P = gl.exp2(S_prev - m_new[:, None]) + l_i = alpha * l_i + gl.sum(P, axis=1) + m_i = m_new + # PV tile t from registers -- matrix; overlaps the gather still in flight + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(P_dot, V_lora_dot, acc) + # promote + S_prev = S_cur + valid_qk = valid_qk_next + tkraw = tkraw_n + cur_buf = 1 - cur_buf + + # ---------- PRE-DRAIN: QK[N-1] (cur_buf) || softmax+PV[N-2] (pv_buf); no gather ---------- + gl.amd.cdna4.async_copy.wait_group(0) # drain K[N-1] (last loop gather) + S_cur = gl.amd.cdna4.mfma( + Q_lora_dot, + smem_klora.index(cur_buf).load(dot_klora_b), + gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s), + ) + if HAS_ROPE: + S_cur = gl.amd.cdna4.mfma(Q_rope_dot, smem_krope.index(cur_buf).load(dot_krope_b), S_cur) + S_cur = S_cur * scale_log2 + S_cur = gl.where(valid_qk[None, :] & mask_h_mma[:, None], S_cur, float("-inf")) + m_j = gl.max(S_prev, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > float("-inf"), m_new, 0.0) + alpha = gl.exp2(m_i - m_new) + P = gl.exp2(S_prev - m_new[:, None]) + l_i = alpha * l_i + gl.sum(P, axis=1) + m_i = m_new + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(1 - cur_buf).permute([1, 0]).load(dot_v_b), acc) + S_prev = S_cur + + # ---------- DRAIN: softmax+PV[N-1] (S_prev = QK[N-1], V from cur_buf) ---------- + m_j = gl.max(S_prev, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > float("-inf"), m_new, 0.0) + alpha = gl.exp2(m_i - m_new) + P = gl.exp2(S_prev - m_new[:, None]) + l_new = alpha * l_i + gl.sum(P, axis=1) + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + P_dot = gl.convert_layout(P.to(Q_ptr.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(P_dot, smem_klora.index(cur_buf).permute([1, 0]).load(dot_v_b), acc) + m_i = m_new + l_i = l_new + + # ---------- epilogue: fold sink into the denominator (V4 delta) ---------- + if HAS_SINK: + offs_h_sink = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mfma_s)) + sink = gl.amd.cdna4.buffer_load( + ptr=Sink_ptr, + offsets=offs_h_sink.to(tl.int32), + mask=offs_h_sink < num_heads, + other=float("-inf"), + ) + sink = sink * 1.4426950408889634 # natural-log sink -> log2 units (exp2 softmax) + m_final = gl.maximum(m_i, sink) + alpha_fix = gl.exp2(m_i - m_final) + l_total = l_i * alpha_fix + gl.exp2(sink - m_final) + alpha_fix_acc = gl.convert_layout(alpha_fix, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_fix_acc[:, None] + l_total_acc = gl.convert_layout(l_total, gl.SliceLayout(1, mfma_acc)) + acc = acc / l_total_acc[:, None] + # lse back to natural log for the backward: m_final is log2, l_total is the natural denom. + lse = m_final * 0.6931471805599453 + gl.log(l_total) + else: + l_i_acc = gl.convert_layout(l_i, gl.SliceLayout(1, mfma_acc)) + acc = acc / l_i_acc[:, None] + lse = m_i * 0.6931471805599453 + gl.log(l_i) + + # Output O[token_idx, h, v] + offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_o = offs_h_o < num_heads + o_base = token_idx.to(tl.int64) * stride_o_t + o_offs = o_base + offs_h_o[:, None].to(tl.int64) * stride_o_h + offs_v_o[None, :].to(tl.int64) + acc_bf = acc.to(O_ptr.dtype.element_ty) + acc_bf_blk = gl.convert_layout(acc_bf, blk_qlora) + gl.amd.cdna4.buffer_store( + stored_value=acc_bf_blk, + ptr=O_ptr, + offsets=o_offs.to(tl.int32), + mask=mask_h_o[:, None], + ) + + # LSE[token_idx, h] + offs_h_lse = hg_offset + gl.arange(0, BLOCK_H, layout=blk_lse) + mask_h_lse = offs_h_lse < num_heads + lse_base = token_idx * num_heads + lse_offs = lse_base + offs_h_lse + lse_blk = gl.convert_layout(lse, blk_lse) + gl.amd.cdna4.buffer_store( + stored_value=lse_blk, + ptr=LSE_ptr, + offsets=lse_offs.to(tl.int32), + mask=mask_h_lse, + ) + + +# ===================================================================== +# Launcher +# ===================================================================== +def sparse_mla_fwd_v4_gluon_v2(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None): + """ + DeepSeek V4 sparse MLA forward (Gluon, gfx950 / CDNA4), with attention sink. + + Args: + q: [total_tokens, num_heads, d_qk] bfloat16 + kv: [total_tokens, 1, d_qk] bfloat16 (or [total_tokens, d_qk]) + topk_indices: [total_tokens, topk] int32 (SWA + sparse, -1 marks invalid) + attn_sink: [num_heads] fp32, optional per-head learnable sink logit. + When None, behaves like the V3.2 forward. + kv_lora_rank: int, default 512 + scale: float, default 1/sqrt(d_qk) + + Returns: + o: [total_tokens, num_heads, kv_lora_rank] same dtype as q + lse: [total_tokens, num_heads] float32 (sink-inclusive when attn_sink is given) + """ + _require_gluon_v2_triton() + assert q.is_contiguous() + assert kv.is_contiguous() + assert topk_indices.is_contiguous() + + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + topk = topk_indices.shape[1] + + if scale is None: + scale = 1.0 / (d_qk**0.5) + + if kv.dim() == 2: + kv = kv.unsqueeze(1) + # kv may hold MORE rows than there are query tokens (V4 feeds a + # [local ++ compressed-pool] buffer, so num_kv = S + P > total_tokens). + # The kernel only dereferences kv via topk indices (stride_kv_t), so any + # num_kv >= max(topk_index)+1 is valid. + assert kv.shape[0] >= total_tokens and kv.shape[-1] == d_qk + + has_sink = attn_sink is not None + if has_sink: + assert attn_sink.is_contiguous() + assert attn_sink.dtype == torch.float32 + assert attn_sink.shape == (num_heads,) + sink_ptr = attn_sink + else: + sink_ptr = torch.empty(1, dtype=torch.float32, device=q.device) # guarded by HAS_SINK + + if num_heads == 128 and topk == 1152 and kv_lora_rank == 512 and rope_rank == 64: + return sparse_mla_fwd_v4_aiter_lse_csa_formula( + q, kv, topk_indices, attn_sink=attn_sink, kv_lora_rank=kv_lora_rank, scale=scale + ) + if num_heads == 64 and topk == 640 and kv_lora_rank == 512 and rope_rank == 64: + return sparse_mla_fwd_v4_aiter_lse_csa_formula( + q, kv, topk_indices, attn_sink=attn_sink, kv_lora_rank=kv_lora_rank, scale=scale + ) + + o = torch.empty(total_tokens, num_heads, kv_lora_rank, dtype=q.dtype, device=q.device) + lse = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device) + + # V4 single-latent form: the D_ROPE block of q/kv is a zero pad (RoPE baked in-place + # over the 512 latent), so the rope QK term is provably zero. Skip it — bit-identical + # to computing it, but avoids the wasteful 64-wide rope MFMA + K_rope loads every tile + # (this is the win triton_v2 already has that our ported gluon fwd lacked). + has_rope = False + + # Grid is autotune-aware: BLOCK_H comes from the chosen config. + grid = lambda META: (total_tokens, triton.cdiv(num_heads, META["BLOCK_H"])) + + try: + _sparse_mla_fwd_gl_v2_kernel[grid]( + Q_ptr=q, + KV_ptr=kv, + TopK_ptr=topk_indices, + Sink_ptr=sink_ptr, + O_ptr=o, + LSE_ptr=lse, + stride_q_t=q.stride(0), + stride_q_h=q.stride(1), + stride_kv_t=kv.stride(0), + stride_o_t=o.stride(0), + stride_o_h=o.stride(1), + stride_topk_t=topk_indices.stride(0), + scale=scale, + num_heads=num_heads, + TOPK=topk, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_SINK=has_sink, + HAS_ROPE=has_rope, + ) + except Exception as exc: # noqa: BLE001 - surface a build hint on Gluon compile failures + _n = type(exc).__name__.lower() + if "compil" in _n or "compil" in str(exc).lower() or "layout" in str(exc).lower(): + raise _wrap_compile_error(exc) from exc + raise + + return o, lse diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/__init__.py new file mode 100644 index 000000000..97ad10bf1 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/__init__.py @@ -0,0 +1,388 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-8 P49 — Tilelang V4 attention dispatcher (infra-only). + +This module is the dispatch point for the plan-8 tilelang-backed V4 +attention kernels (cr ∈ {0, 4, 128}). At P49 it ships only the +**infra**: + +* The pinned tilelang version probe (one-time module-import warning + if the installed tilelang ≠ the plan-8 pin). +* The :func:`should_dispatch` predicate that takes the per-call + ``enabled`` flag (from the ``use_v4_tilelang_attention`` / + ``use_v4_tilelang_csa_attention`` config flags surfaced by + ``DeepSeekV4TransformerConfig``). Plan-8 P57 close-out 2: the + dispatcher no longer reads ``PRIMUS_V4_TILELANG_ATTN`` — callers + pass the config flag explicitly so a container without tilelang + installed can simply leave the flag ``False`` and never trigger + any tilelang import. +* The :func:`is_tilelang_kernel_available` predicate that lets each + plan-8 phase (P50..P55) register its kernel name as it lands. + Empty at P49 — every dispatcher call falls through to the + plan-4 P25 / P26 Triton path with a one-time rank-0 warning. +* Four stub entry points that raise :class:`NotImplementedError` + until the corresponding plan-8 phase lands them: + + - :func:`v4_attention_fwd_tilelang` (P50 — cr=0 / cr=128 FWD) + - :func:`v4_attention_bwd_tilelang` (P51 — cr=0 / cr=128 BWD) + - :func:`v4_csa_attention_fwd_tilelang` (P54 — cr=4 FWD) + - :func:`v4_csa_attention_bwd_tilelang` (P55 — cr=4 BWD) + +The dispatcher precedence (enforced inside the V4 attention +functional wrappers ``v4_attention_v1`` / ``v4_csa_attention_v0``) is: + +.. code-block:: text + + cr ∈ {0, 128}: + use_turbo_attention > use_v4_tilelang_attention + > use_v4_triton_attention > eager + cr == 4: + use_v4_tilelang_csa_attention > use_v4_triton_csa_attention > eager + +At P49 (default-off + empty available-kernels set) the dispatcher +emits **no** behaviour change vs the plan-7 P48 anchor. Plan-4..7 +ratchet stays green by construction. + +R6.2 — Tilelang is vendored at ``tilelang/`` and NEVER edited. This +module imports it lazily so a missing install does not raise at +module import time; it only logs a one-time rank-0 warning the +first time a caller actually asks the dispatcher to use tilelang. +""" + +from __future__ import annotations + +import warnings +from typing import Any, Set + +# --------------------------------------------------------------------------- +# Tilelang version pin (set by plan-8 P49; bump when upstream is updated) +# --------------------------------------------------------------------------- + +# Read from ``tilelang/VERSION`` at plan-8 P49 land time. Future plan-8 +# phases (P50..P55) MUST re-run the G50..G55 ratchets before bumping +# this pin. +TILELANG_VERSION_PIN: str = "0.1.9+cuda.gitbcb2da33" + + +# --------------------------------------------------------------------------- +# Dispatch-enable signal (replaces the PRIMUS_V4_TILELANG_ATTN env knob) +# --------------------------------------------------------------------------- +# +# Plan-8 P57 close-out 2 (2026-05-15): the dispatcher no longer reads an +# environment variable. Callers pass ``enabled`` to :func:`should_dispatch` +# from the V4 attention layer's config flag +# (``use_v4_tilelang_attention`` / ``use_v4_tilelang_csa_attention``), +# which is plumbed through :class:`DeepSeekV4TransformerConfig` from +# the run-time CLI args. Default-False everywhere — running on a +# container without tilelang installed leaves the flag off and the +# dispatcher never imports tilelang. + + +# --------------------------------------------------------------------------- +# Per-kernel availability registry +# --------------------------------------------------------------------------- + +# Mutable set of kernel names that plan-8 phases register as they +# land. Empty at P49; populated by P50..P55 as each kernel ships +# via ``register_available_kernel()``. Names follow the +# four-stub convention below. +_AVAILABLE_KERNELS: Set[str] = set() + +_KNOWN_KERNEL_NAMES: Set[str] = { + "v4_attention_fwd", # P50 — dense / HCA FWD + "v4_attention_bwd", # P51 — dense / HCA BWD + "v4_csa_attention_fwd", # P54 — CSA FWD + "v4_csa_attention_bwd", # P55 — CSA BWD +} + + +def is_tilelang_kernel_available(name: str) -> bool: + """Return True iff the plan-8 kernel ``name`` has landed. + + Empty at P49; each plan-8 phase (P50..P55) registers its + kernel name via :func:`register_available_kernel` once the + parity gate passes. + """ + if name not in _KNOWN_KERNEL_NAMES: + raise ValueError( + f"Unknown tilelang kernel name {name!r}; expected one of " f"{sorted(_KNOWN_KERNEL_NAMES)}" + ) + return name in _AVAILABLE_KERNELS + + +def register_available_kernel(name: str) -> None: + """Mark a plan-8 tilelang kernel as available. + + Called at module-import time by the phase that lands the + corresponding kernel. P50: ``register_available_kernel("v4_attention_fwd")`` + inside ``v4_attention_fwd_tilelang.py``. + """ + if name not in _KNOWN_KERNEL_NAMES: + raise ValueError( + f"Unknown tilelang kernel name {name!r}; expected one of " f"{sorted(_KNOWN_KERNEL_NAMES)}" + ) + _AVAILABLE_KERNELS.add(name) + + +# Map kernel name -> submodule name to lazy-import. Each submodule +# registers its kernel name + overrides the stub in this namespace on +# first import. +_KERNEL_SUBMODULES: dict[str, str] = { + "v4_attention_fwd": "v4_attention_fwd_tilelang", + "v4_attention_bwd": "v4_attention_bwd_tilelang", + "v4_csa_attention_fwd": "v4_csa_attention_fwd_tilelang", + "v4_csa_attention_bwd": "v4_csa_attention_bwd_tilelang", +} + +_LAZY_LOADED: Set[str] = set() + + +def _lazy_load(name: str) -> bool: + """Import the submodule that implements kernel ``name`` (if landed). + + Returns True iff the import succeeded. Each plan-8 kernel + submodule's import-time code: + + 1. Calls :func:`register_available_kernel` to flip + :func:`is_tilelang_kernel_available` to True. + 2. Replaces the module-level stub in this namespace via + ``setattr(, name + "_tilelang", real_fn)``. + + Idempotent — subsequent calls no-op. + """ + if name in _LAZY_LOADED: + return name in _AVAILABLE_KERNELS + _LAZY_LOADED.add(name) + sub = _KERNEL_SUBMODULES.get(name) + if sub is None: + return False + try: + from importlib import import_module + + sub_mod = import_module(__name__ + "." + sub) + except ImportError as exc: + warnings.warn( + f"[plan-8] tilelang submodule {sub!r} import failed ({exc}); " + f"kernel {name!r} stays unavailable.", + RuntimeWarning, + stacklevel=3, + ) + return False + # Override the stub: the submodule exposes a function with the same + # name as the submodule (e.g. `v4_attention_fwd_tilelang.py` + # exposes `v4_attention_fwd_tilelang`). Importing the submodule + # would otherwise set `_tilelang.v4_attention_fwd_tilelang` to + # the SUBMODULE, shadowing the stub. We restore the function + # attribute here. + real_fn = getattr(sub_mod, sub, None) + if real_fn is None: + warnings.warn( + f"[plan-8] tilelang submodule {sub!r} does not expose a " + f"function named {sub!r}; kernel {name!r} stays unavailable.", + RuntimeWarning, + stacklevel=3, + ) + return False + import sys + + sys.modules[__name__].__dict__[sub] = real_fn + return name in _AVAILABLE_KERNELS + + +# --------------------------------------------------------------------------- +# Tilelang import probe (lazy) +# --------------------------------------------------------------------------- + +_TILELANG_PROBE_DONE: bool = False +_TILELANG_AVAILABLE: bool = False + + +def _probe_tilelang() -> bool: + """Lazily import tilelang on first dispatcher call. + + Returns True iff the import succeeds at the pinned version. + Emits a one-time warning when the installed version differs from + the pin or when tilelang is not importable. + """ + global _TILELANG_PROBE_DONE, _TILELANG_AVAILABLE + if _TILELANG_PROBE_DONE: + return _TILELANG_AVAILABLE + _TILELANG_PROBE_DONE = True + try: + import tilelang # noqa: F401 + except ImportError as exc: + warnings.warn( + f"[plan-8 P49] tilelang import failed ({exc}); falling back to " + "the plan-4 P25 / P26 Triton kernels. Install the pinned " + f"tilelang {TILELANG_VERSION_PIN} to enable the plan-8 path.", + RuntimeWarning, + stacklevel=2, + ) + _TILELANG_AVAILABLE = False + return False + installed = getattr(tilelang, "__version__", "") + if installed != TILELANG_VERSION_PIN: + warnings.warn( + f"[plan-8 P49] installed tilelang {installed!r} != pinned " + f"{TILELANG_VERSION_PIN!r}; the plan-8 G50..G55 parity gates were " + "run against the pin. Proceeding anyway; please re-run the " + "ratchets if the version drifted.", + RuntimeWarning, + stacklevel=2, + ) + _TILELANG_AVAILABLE = True + return True + + +# --------------------------------------------------------------------------- +# Cache directory +# --------------------------------------------------------------------------- + + +def cache_dir() -> str: + """Tilelang autotune cache dir. + + Defaults to ``output/.tilelang_cache/v4/`` (gitignored under + ``output/``). Override via ``PRIMUS_V4_TILELANG_CACHE_DIR``. + """ + import os as _os + + override = _os.environ.get("PRIMUS_V4_TILELANG_CACHE_DIR") + if override: + return override + return _os.path.join("output", ".tilelang_cache", "v4") + + +# --------------------------------------------------------------------------- +# Stub entry points — replaced by P50 / P51 / P54 / P55 +# --------------------------------------------------------------------------- + + +def _stub_raise(phase: str, name: str) -> None: + raise NotImplementedError( + f"plan-8 {phase} tilelang kernel {name!r} has not landed yet; " + "the dispatcher should fall back to the plan-4 / plan-5 Triton " + "path via `is_tilelang_kernel_available({name!r})` returning False." + ) + + +def v4_attention_fwd_tilelang(*args: Any, **kwargs: Any): # type: ignore[no-untyped-def] + """Stub for plan-8 P50 (dense / HCA FWD).""" + _stub_raise("P50", "v4_attention_fwd") + + +def v4_attention_bwd_tilelang(*args: Any, **kwargs: Any): # type: ignore[no-untyped-def] + """Stub for plan-8 P51 (dense / HCA BWD).""" + _stub_raise("P51", "v4_attention_bwd") + + +def v4_csa_attention_fwd_tilelang(*args: Any, **kwargs: Any): # type: ignore[no-untyped-def] + """Stub for plan-8 P54 (CSA FWD).""" + _stub_raise("P54", "v4_csa_attention_fwd") + + +def v4_csa_attention_bwd_tilelang(*args: Any, **kwargs: Any): # type: ignore[no-untyped-def] + """Stub for plan-8 P55 (CSA BWD).""" + _stub_raise("P55", "v4_csa_attention_bwd") + + +# --------------------------------------------------------------------------- +# Dispatcher helpers — used by the v4_attention_v1 / v4_csa_attention_v0 wrappers +# --------------------------------------------------------------------------- + +_FALLBACK_WARNED: Set[str] = set() + + +def _maybe_warn_fallback(kernel_name: str) -> None: + """Emit a one-time rank-0 warning when a tilelang dispatch falls + back to Triton. + + Hit conditions: + + * Config flag ``use_v4_tilelang_attention`` / + ``use_v4_tilelang_csa_attention`` set but the kernel hasn't + landed yet (``is_tilelang_kernel_available`` returns False). + * Config flag set but tilelang import failed (e.g. missing + install / version drift). + + Once per kernel name per process. Banned-warning ratchet is + extended in `plan-8/03-test-strategy.md` to allow this string + (since it's the documented dispatcher fallback signal). + """ + if kernel_name in _FALLBACK_WARNED: + return + _FALLBACK_WARNED.add(kernel_name) + # Best-effort rank-0 detection — the V4 attention wrappers may + # run in unit tests without parallel-state initialised. + rank = 0 + try: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + rank = dist.get_rank() + except Exception: + pass + if rank == 0: + warnings.warn( + f"[plan-8 P57] use_v4_tilelang_*=True but kernel " + f"{kernel_name!r} is not available; falling back to the " + "plan-4 / plan-5 Triton path.", + RuntimeWarning, + stacklevel=3, + ) + + +def should_dispatch(kernel_name: str, enabled: bool = False) -> bool: + """Single dispatcher predicate used by ``v4_attention_v1`` / + ``v4_csa_attention_v0`` wrappers. + + Returns True iff all three conditions hold: + + 1. ``enabled`` is True (i.e. the caller passes the relevant + ``use_v4_tilelang_*`` config flag). + 2. The named plan-8 kernel has been registered (P50..P55 land it). + 3. The tilelang import succeeds at the pinned version (probed + lazily on the first dispatcher call). + + Otherwise emits a one-time rank-0 warning + returns False so the + caller falls through to the Triton path. ``enabled=False`` short- + circuits before any tilelang import attempt, so containers without + tilelang installed can leave the config flag off and the dispatcher + never touches tilelang. + + Plan-8 P57 close-out 2 (2026-05-15): previously gated by the + ``PRIMUS_V4_TILELANG_ATTN`` env var; now driven by the caller's + config flag so default-off runs in tilelang-free containers do + not need to set or unset any env knob. + """ + if not enabled: + return False + if not _probe_tilelang(): + _maybe_warn_fallback(kernel_name) + return False + # Lazy-import the submodule that implements ``kernel_name``. + # The submodule's import-time code flips + # ``is_tilelang_kernel_available(...)`` to True and overrides + # the stub function in this namespace. Cheap on subsequent calls. + _lazy_load(kernel_name) + if not is_tilelang_kernel_available(kernel_name): + _maybe_warn_fallback(kernel_name) + return False + return True + + +__all__ = [ + "TILELANG_VERSION_PIN", + "cache_dir", + "is_tilelang_kernel_available", + "register_available_kernel", + "should_dispatch", + "v4_attention_bwd_tilelang", + "v4_attention_fwd_tilelang", + "v4_csa_attention_bwd_tilelang", + "v4_csa_attention_fwd_tilelang", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_autograd_tilelang.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_autograd_tilelang.py new file mode 100644 index 000000000..10769790d --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_autograd_tilelang.py @@ -0,0 +1,147 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-8 P51 — autograd wrapper that ties the tilelang FWD + BWD. + +Mirrors :class:`primus.backends.megatron.core.transformer.v4_attention_kernels.v4_attention.V4AttentionFn` +(the Triton autograd Function) but routes through the tilelang +FWD / BWD kernels. The wrapper falls back to the Triton path +inside the FWD / BWD wrapper functions when the tilelang kernels +don't yet support the requested feature (e.g. additive_mask, +hca_local_seqlen > 0). + +Implementation note: no ``from __future__ import annotations`` +because tilelang's annotation eval is eager. +""" + +from typing import Optional + +import torch + +from primus.backends.megatron.core.transformer.v4_attention_kernels._tilelang.v4_attention_bwd_tilelang import ( + v4_attention_bwd_tilelang, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._tilelang.v4_attention_fwd_tilelang import ( + v4_attention_fwd_tilelang_with_lse, +) + + +class V4AttentionTilelangFn(torch.autograd.Function): + """Autograd-aware V4 dense / SWA / sink attention via tilelang. + + FWD: :func:`v4_attention_fwd_tilelang_with_lse` returns + ``(out, lse)``; we save the inputs + ``lse`` for backward. + + BWD: :func:`v4_attention_bwd_tilelang` consumes the saved + tensors + the incoming ``dO``; returns ``(dq, dk, dv, dsink)``. + For unsupported features (additive_mask, hca_local_seqlen), + both wrappers fall back to the Triton kernels so this autograd + path remains correct even when tilelang's scope is limited. + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sink: Optional[torch.Tensor], + additive_mask: Optional[torch.Tensor], + swa_window: int, + attn_dropout: float, + training: bool, + scale: float, + hca_local_seqlen: int, + ) -> torch.Tensor: + out, lse = v4_attention_fwd_tilelang_with_lse( + q, + k, + v, + sink=sink, + additive_mask=additive_mask, + swa_window=int(swa_window), + attn_dropout=float(attn_dropout), + training=bool(training), + scale=float(scale), + hca_local_seqlen=int(hca_local_seqlen), + ) + # Save tensors needed for BWD. `additive_mask` and `sink` + # may be None; save_for_backward only stores tensors, so we + # stash None-ness on ctx and conditionally save. + ctx.save_for_backward(q, k, v, out, lse) + ctx.sink = sink + ctx.additive_mask = additive_mask + ctx.swa_window = int(swa_window) + ctx.scale = float(scale) + ctx.hca_local_seqlen = int(hca_local_seqlen) + return out + + @staticmethod + def backward(ctx, d_out): # type: ignore[override] + q, k, v, out, lse = ctx.saved_tensors + dq, dk, dv, dsink = v4_attention_bwd_tilelang( + q, + k, + v, + out, + lse, + d_out, + sink=ctx.sink, + additive_mask=ctx.additive_mask, + swa_window=ctx.swa_window, + scale=ctx.scale, + hca_local_seqlen=ctx.hca_local_seqlen, + ) + # Match the input-arg count of `forward(...)`: q, k, v, sink, + # additive_mask, swa_window, attn_dropout, training, scale, + # hca_local_seqlen. Non-tensor args get None. + return ( + dq, + dk, + dv, + dsink, + None, # additive_mask + None, # swa_window + None, # attn_dropout + None, # training + None, # scale + None, # hca_local_seqlen + ) + + +def v4_attention_tilelang( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + sink: Optional[torch.Tensor] = None, + additive_mask: Optional[torch.Tensor] = None, + swa_window: int = 0, + attn_dropout: float = 0.0, + training: bool = False, + scale: Optional[float] = None, + hca_local_seqlen: int = 0, +) -> torch.Tensor: + """Convenience wrapper matching the existing `v4_attention_v1()` + functional API but routing through the tilelang autograd + Function.""" + if scale is None: + scale = 1.0 / (q.shape[-1] ** 0.5) + return V4AttentionTilelangFn.apply( + q, + k, + v, + sink, + additive_mask, + int(swa_window), + float(attn_dropout), + bool(training), + float(scale), + int(hca_local_seqlen), + ) + + +__all__ = ["V4AttentionTilelangFn", "v4_attention_tilelang"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_bwd_tilelang.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_bwd_tilelang.py new file mode 100644 index 000000000..937f8a70e --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_bwd_tilelang.py @@ -0,0 +1,427 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-8 P51 — V4 dense BWD tilelang kernels (compress_ratio == 0). + +Three-kernel pipeline matching the tilelang FlashAttention v2 BWD +reference (`tilelang/examples/amd/example_amd_flash_attn_bwd.py`): + +* :func:`_make_preprocess_kernel` — ``Delta[b, h, m] = + sum_d O[b, h, m, d] * dO[b, h, m, d]``. One tiny kernel. +* :func:`_make_bwd_kernel` — main BWD pass. One program per + ``(h, n_tile, b)``; loops over m_tile; re-materialises ``P`` + from saved ``LSE``; accumulates ``dQ`` / ``dK`` / ``dV`` + via ``tl.atomic_add``. Sink BWD reuses the dq stride. +* The sink gradient is computed in the Python wrapper as a small + reduction over the saved ``LSE`` + ``Delta`` (cheap; one ATen + launch). + +Scope at P51 (matches P50): + +* Q/K/V in BHSD layout with optional MQA broadcast (`K_H ∈ {1, H}`). +* Optional per-head sink. +* Optional sliding-window-causal mask. +* No `additive_mask`, no `hca_local_seqlen > 0` — those route + through the plan-4 P25 / plan-5 P32 final Triton BWD kernel. + +dtype contract (matches :func:`reference.eager_v4_attention` +backward): matmuls run in input dtype on tensor cores; the +accumulator + softmax recompute live in fp32. Output gradients +are in `q.dtype` / `k.dtype` / `v.dtype` (typically all bf16); +internal accumulators are fp32 buffers cast at the end. + +Implementation note: this module does NOT use ``from __future__ +import annotations`` because tilelang's ``@T.prim_func`` +decorator evaluates annotations eagerly via +``typing.get_type_hints``. +""" + +from typing import Optional, Tuple + +import tilelang +import tilelang.language as T +import torch + +from primus.backends.megatron.core.transformer.v4_attention_kernels import _tilelang +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import ( + _launch_v4_attention_bwd, +) + +_PREPROCESS_CACHE: dict[tuple, object] = {} +_BWD_CACHE: dict[tuple, object] = {} + + +# --------------------------------------------------------------------------- +# Preprocess: Delta[b, h, m] = (O * dO).sum(-1) +# --------------------------------------------------------------------------- + + +@tilelang.jit(out_idx=[2]) +def _make_preprocess_kernel( + batch: int, + heads: int, + seq_q: int, + dim: int, + dtype: "T.dtype" = T.bfloat16, + block: int = 32, +): + """Preprocess kernel: emit ``Delta [B, H, Sq]`` fp32. + + Tiles `(b, h, m_tile)`; each program reduces over `dim` in + chunks of `block`. + """ + accum_dtype = T.float32 + bhsd_shape = [batch, heads, seq_q, dim] + delta_shape = [batch, heads, seq_q] + + @T.prim_func + def main( + O: T.Tensor(bhsd_shape, dtype), + dO: T.Tensor(bhsd_shape, dtype), + Delta: T.Tensor(delta_shape, accum_dtype), + ): + with T.Kernel(batch, heads, T.ceildiv(seq_q, block)) as (bz, bx, by): + o = T.alloc_fragment([block, block], dtype) + do = T.alloc_fragment([block, block], dtype) + acc = T.alloc_fragment([block, block], accum_dtype) + delta = T.alloc_fragment([block], accum_dtype) + T.clear(acc) + for k in range(T.ceildiv(dim, block)): + T.copy( + O[bz, bx, by * block : (by + 1) * block, k * block : (k + 1) * block], + o, + ) + T.copy( + dO[bz, bx, by * block : (by + 1) * block, k * block : (k + 1) * block], + do, + ) + for i, j in T.Parallel(block, block): + acc[i, j] += o[i, j] * do[i, j] + T.reduce_sum(acc, delta, 1) + T.copy(delta, Delta[bz, bx, by * block : (by + 1) * block]) + + return main + + +# --------------------------------------------------------------------------- +# Main BWD: dQ, dK, dV via re-materialised P from saved LSE +# --------------------------------------------------------------------------- + + +# Note: dK / dV are emitted as fp32 (not the input dtype) so the +# MQA path can use `T.atomic_add` — tilelang's HIP runtime only +# supports `AtomicAddx2(float*, ...)`. The wrapper casts to +# k.dtype / v.dtype before returning. +@tilelang.jit(out_idx=[6, 7, 8]) +def _make_bwd_kernel( + batch: int, + heads_q: int, + heads_k: int, + seq_q: int, + seq_k: int, + dim: int, + has_sink: bool, + swa_window: int, + dtype: "T.dtype" = T.bfloat16, + block_M: int = 32, + block_N: int = 32, + threads: int = 64, +): + """Plan-8 P51 main BWD: dQ / dK / dV via re-materialised softmax. + + Grid: one program per `(b, h, n_tile)` (n_tile along K/V seq dim). + Each program loops over `m_tile` along Q. Pattern mirrors the + tilelang AMD example BWD with three adaptations: + + * BHSD layout instead of BSHD. + * MQA via `k_head = bx // groups`. + * SWA-aware m_tile loop bounds (skip tiles outside the window). + + Output `dQ` is in fp32 (caller casts back to q.dtype after a + `dsink` host-side reduction); `dK`, `dV` are emitted in + `k.dtype` / `v.dtype` via the kernel epilogue. + """ + groups = heads_q // heads_k + accum_dtype = T.float32 + sm_scale = 1.0 / (dim**0.5) + q_shape = [batch, heads_q, seq_q, dim] + kv_shape = [batch, heads_k, seq_k, dim] + dq_shape = [batch, heads_q, seq_q, dim] + dkv_shape = [batch, heads_k, seq_k, dim] + lse_shape = [batch, heads_q, seq_q] + delta_shape = [batch, heads_q, seq_q] + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), + K: T.Tensor(kv_shape, dtype), + V: T.Tensor(kv_shape, dtype), + dO: T.Tensor(q_shape, dtype), + LSE: T.Tensor(lse_shape, accum_dtype), + Delta: T.Tensor(delta_shape, accum_dtype), + dQ: T.Tensor(dq_shape, accum_dtype), + dK: T.Tensor(dkv_shape, accum_dtype), + dV: T.Tensor(dkv_shape, accum_dtype), + ): + with T.Kernel(heads_q, T.ceildiv(seq_k, block_M), batch, threads=threads) as ( + bx, + by, + bz, + ): + K_shared = T.alloc_shared([block_M, dim], dtype) + V_shared = T.alloc_shared([block_M, dim], dtype) + q_shared = T.alloc_shared([block_N, dim], dtype) + do_shared = T.alloc_shared([block_N, dim], dtype) + lse_shared = T.alloc_shared([block_N], accum_dtype) + delta_shared = T.alloc_shared([block_N], accum_dtype) + ds_shared = T.alloc_shared([block_M, block_N], dtype) + p_cast = T.alloc_fragment([block_M, block_N], dtype) + qkT = T.alloc_fragment([block_M, block_N], accum_dtype) + P_acc = T.alloc_fragment([block_M, block_N], accum_dtype) + dP = T.alloc_fragment([block_M, block_N], accum_dtype) + dv = T.alloc_fragment([block_M, dim], accum_dtype) + dk = T.alloc_fragment([block_M, dim], accum_dtype) + dq_tile = T.alloc_fragment([block_N, dim], accum_dtype) + + k_head = bx // groups + + T.copy(K[bz, k_head, by * block_M : (by + 1) * block_M, :], K_shared) + T.copy(V[bz, k_head, by * block_M : (by + 1) * block_M, :], V_shared) + T.clear(dv) + T.clear(dk) + + # Causal: only m_tiles that touch this n_tile contribute. + loop_st = by * block_M // block_N + loop_ed = T.ceildiv(seq_q, block_N) + + for k in T.Pipelined(loop_st, loop_ed, num_stages=1): + T.copy(Q[bz, bx, k * block_N : (k + 1) * block_N, :], q_shared) + T.clear(qkT) + T.gemm( + K_shared, + q_shared, + qkT, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(LSE[bz, bx, k * block_N : (k + 1) * block_N], lse_shared) + + for i, j in T.Parallel(block_M, block_N): + P_acc[i, j] = T.exp(qkT[i, j] * sm_scale - lse_shared[j]) + + # Causal mask + optional SWA: keep entries where the + # query token (`k * block_N + j`) attends to the key + # token (`by * block_M + i`). SWA: also require + # `q_idx - k_idx < swa_window`. + for i, j in T.Parallel(block_M, block_N): + q_idx = k * block_N + j + k_idx = by * block_M + i + valid = q_idx >= k_idx + if swa_window > 0: + valid = valid and (q_idx - k_idx < swa_window) + P_acc[i, j] = T.if_then_else(valid, P_acc[i, j], 0.0) + + T.copy(dO[bz, bx, k * block_N : (k + 1) * block_N, :], do_shared) + T.clear(dP) + T.gemm( + V_shared, + do_shared, + dP, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(P_acc, p_cast) + T.gemm(p_cast, do_shared, dv, policy=T.GemmWarpPolicy.FullRow) + + T.copy(Delta[bz, bx, k * block_N : (k + 1) * block_N], delta_shared) + for i, j in T.Parallel(block_M, block_N): + p_cast[i, j] = P_acc[i, j] * (dP[i, j] - delta_shared[j]) * sm_scale + T.gemm(p_cast, q_shared, dk, policy=T.GemmWarpPolicy.FullRow) + T.copy(p_cast, ds_shared) + + T.clear(dq_tile) + T.gemm(ds_shared, K_shared, dq_tile, transpose_A=True) + for i, j in T.Parallel(block_N, dim): + T.atomic_add(dQ[bz, bx, k * block_N + i, j], dq_tile[i, j]) + + # MQA: every query head writes to the same K/V head; + # use atomic_add to merge across query-head programs. + # `dK` / `dV` are fp32 buffers (tilelang's HIP atomic-add + # only supports float*); the wrapper casts to k.dtype / + # v.dtype after the kernel returns. + if groups > 1: + for i, j in T.Parallel(block_M, dim): + T.atomic_add(dV[bz, k_head, by * block_M + i, j], dv[i, j]) + T.atomic_add(dK[bz, k_head, by * block_M + i, j], dk[i, j]) + else: + for i, j in T.Parallel(block_M, dim): + dV[bz, k_head, by * block_M + i, j] = dv[i, j] + dK[bz, k_head, by * block_M + i, j] = dk[i, j] + + return main + + +# --------------------------------------------------------------------------- +# Wrapper API +# --------------------------------------------------------------------------- + + +def _kernel_supports( + *, + additive_mask: Optional[torch.Tensor], + hca_local_seqlen: int, +) -> bool: + if additive_mask is not None: + return False + if hca_local_seqlen > 0: + return False + return True + + +def _torch_dtype_to_tilelang(dtype: torch.dtype): + if dtype == torch.bfloat16: + return T.bfloat16 + if dtype == torch.float16: + return T.float16 + if dtype == torch.float32: + return T.float32 + raise ValueError(f"unsupported dtype {dtype}") + + +def _get_preprocess(batch, heads, seq_q, dim, dtype): + key = (batch, heads, seq_q, dim, str(dtype)) + if key in _PREPROCESS_CACHE: + return _PREPROCESS_CACHE[key] + k = _make_preprocess_kernel(batch, heads, seq_q, dim, dtype=_torch_dtype_to_tilelang(dtype)) + _PREPROCESS_CACHE[key] = k + return k + + +def _get_bwd(batch, heads_q, heads_k, seq_q, seq_k, dim, has_sink, swa_window, dtype): + key = ( + batch, + heads_q, + heads_k, + seq_q, + seq_k, + dim, + has_sink, + int(swa_window > 0) * swa_window, + str(dtype), + ) + if key in _BWD_CACHE: + return _BWD_CACHE[key] + # Same SMEM-budget heuristic as P50. + if dim >= 256: + block_M, block_N, threads = 32, 32, 64 + else: + block_M, block_N, threads = 64, 64, 128 + k = _make_bwd_kernel( + batch, + heads_q, + heads_k, + seq_q, + seq_k, + dim, + has_sink, + swa_window, + dtype=_torch_dtype_to_tilelang(dtype), + block_M=block_M, + block_N=block_N, + threads=threads, + ) + _BWD_CACHE[key] = k + return k + + +def v4_attention_bwd_tilelang( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + lse: torch.Tensor, + do: torch.Tensor, + *, + sink: Optional[torch.Tensor] = None, + additive_mask: Optional[torch.Tensor] = None, + swa_window: int = 0, + scale: Optional[float] = None, + hca_local_seqlen: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Plan-8 P51 tilelang dense / SWA / sink BWD wrapper. + + Drop-in replacement for ``_launch_v4_attention_bwd`` (plan-5 + P32 final Triton split BWD) on the dense subset of the V4 + attention contract. Falls back to Triton for `additive_mask` + / `hca_local_seqlen > 0` paths (P53 territory). + + Returns ``(dq, dk, dv, dsink)`` matching the Triton launcher + signature. + """ + if not _kernel_supports(additive_mask=additive_mask, hca_local_seqlen=hca_local_seqlen): + return _launch_v4_attention_bwd( + q, + k, + v, + out, + lse, + do, + sink=sink, + additive_mask=additive_mask, + swa_window=swa_window, + scale=scale if scale is not None else 1.0 / (q.shape[-1] ** 0.5), + hca_local_seqlen=hca_local_seqlen, + ) + + B, HQ, Sq, D = q.shape + _, HK, Sk, _ = k.shape + assert HK in (1, HQ) + dtype = q.dtype + + q_c = q.contiguous() + k_c = k.contiguous() + v_c = v.contiguous() + do_c = do.contiguous() + out_c = out.contiguous() + lse_c = lse.contiguous() + + # Preprocess: Delta = (O * dO).sum(-1) + preprocess = _get_preprocess(B, HQ, Sq, D, dtype) + delta = preprocess(out_c, do_c) + + has_sink = sink is not None + bwd = _get_bwd(B, HQ, HK, Sq, Sk, D, has_sink, int(swa_window), dtype) + dq_fp32, dk_fp32, dv_fp32 = bwd(q_c, k_c, v_c, do_c, lse_c, delta) + # All three grads are emitted in fp32 (MQA atomic-add target dtype); + # cast back to input dtypes. + dq = dq_fp32.to(dtype) + dk = dk_fp32.to(k.dtype) + dv = dv_fp32.to(v.dtype) + + # Sink BWD: computed host-side (cheap). When sink is present: + # dsink[h] = sum_b sum_m exp(sink[h] - lse[b, h, m]) * delta[b, h, m] + if has_sink: + # Re-cast delta to fp32 for the reduce (it's already fp32 from + # preprocess). + dsink = (torch.exp(sink.float().unsqueeze(0).unsqueeze(-1) - lse_c.float()) * delta).sum(dim=(0, 2)) + dsink = dsink.to(sink.dtype) + else: + dsink = None + + return dq, dk, dv, dsink + + +# --------------------------------------------------------------------------- +# Module-level registration +# --------------------------------------------------------------------------- + +_tilelang.register_available_kernel("v4_attention_bwd") +# `_tilelang._lazy_load(...)` restores the parent-module attribute +# (`_tilelang.v4_attention_bwd_tilelang`) after Python's import +# machinery would otherwise shadow the stub with this submodule. + + +__all__ = ["v4_attention_bwd_tilelang"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_fwd_tilelang.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_fwd_tilelang.py new file mode 100644 index 000000000..7340b03c9 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_tilelang/v4_attention_fwd_tilelang.py @@ -0,0 +1,480 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-8 P50 — V4 dense FWD tilelang kernel (compress_ratio == 0). + +Drop-in replacement for ``_launch_v4_attention_fwd`` (plan-4 P25 +Triton FWD) at the dense / SWA / sink subset of the V4 attention +contract. HCA (``hca_local_seqlen > 0``) and explicit +``additive_mask`` paths are deferred to P52; the wrapper falls +back to the Triton kernel when those features are requested. + +Scope at P50: + +* ``Q [B, H, Sq, D]``, ``K [B, K_H, Sk, D]``, ``V [B, K_H, Sk, D]`` + with ``K_H ∈ {1 (MQA), H (MHA)}``. +* Optional per-head sink ``[H]`` — joined as a virtual key column + at the end of the softmax (matches the plan-4 P25 / plan-2 P14 + sink contract). +* Optional sliding-window-causal mask (``swa_window > 0``) applied + in the kernel via `start = max(0, (m + 1 - window_size) // + block_N)`. +* MQA broadcast: when ``K_H == 1``, every query head reads the + shared K/V head; computed via per-program ``h // (HQ // HK)``. + +Out of scope at P50 (Triton fallback handles these): + +* ``additive_mask`` not None — needs the kernel to accept a + [Sq, Sk] additive bias. Deferred to P52 (alongside the HCA + split-mask path). +* ``hca_local_seqlen > 0`` — HCA split-mask path. Deferred to + P52. + +dtype contract (must match :func:`reference.eager_v4_attention`): + +* Q / K / V matmuls run in input dtype on tensor cores; the + matmul accumulator inside is fp32. +* Online softmax accumulator (``m_i``, ``l_i``, ``acc_o``) is fp32. +* Output is in ``v.dtype``; saved ``LSE`` is fp32 (BWD walks back + from it). + +The kernel borrows from +``tilelang/examples/attention_sink/example_mha_sink_fwd_bhsd.py`` +(sink fusion at end of softmax via ``exp2(sinks * log2e - max * +scale)``) and ``tilelang/examples/amd/example_amd_flash_attn_fwd.py`` +(MI355X-tuned block sizes + WMMA ``k_pack`` + ``T.use_swizzle``). + +Implementation note: this module deliberately does NOT use +``from __future__ import annotations`` because tilelang's +``@T.prim_func`` decorator evaluates annotations eagerly via +``typing.get_type_hints``, which fails when annotations are lazy +strings + the referenced names live in an enclosing function's +local scope (e.g. ``q_shape`` defined inside the JIT factory). +""" + +from typing import Optional, Tuple + +# Lazy import — tilelang is heavy. The wrapper checks +# ``_tilelang.should_dispatch(...)`` upstream so we only land here +# when the env knob is set + tilelang is importable. +import tilelang +import tilelang.language as T +import torch + +from primus.backends.megatron.core.transformer.v4_attention_kernels import _tilelang +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import ( + _launch_v4_attention_fwd, +) + +# Compiled-kernel cache keyed by (B, HQ, HK, Sq, Sk, D, has_sink, +# swa_window, dtype-str). Tilelang's own JIT also caches per-shape; +# this dict is a Python-side memoisation that avoids the +# ``@tilelang.jit`` recompile cost on every wrapper call. +_KERNEL_CACHE: dict[tuple, object] = {} + + +# --------------------------------------------------------------------------- +# Tilelang kernel definition +# --------------------------------------------------------------------------- + + +@tilelang.jit( + out_idx=[3, 4], + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True}, +) +def _make_v4_attention_fwd_kernel( + batch: int, + heads_q: int, + heads_k: int, + seq_q: int, + seq_k: int, + dim: int, + has_sink: bool, + swa_window: int, + dtype: "T.dtype" = T.bfloat16, + block_M: int = 64, + block_N: int = 64, + num_stages: int = 1, + threads: int = 128, +): + """Plan-8 P50 tilelang FWD kernel JIT factory. + + `has_sink`, `swa_window`, MQA-vs-MHA branching are baked into the + kernel at JIT time so the inner loop has no host-side branching + cost. All shape args + flags are kwargs to this JIT'd function; + tilelang's `@tilelang.jit` decorator caches compiled binaries per + (kwargs) tuple via its own cache. + """ + + groups = heads_q // heads_k # 1 for MHA, heads_q for MQA + accum_dtype = T.float32 + sm_scale = 1.0 / (dim**0.5) + scale = sm_scale * 1.44269504 # log2(e) — enables exp2/FFMA + + q_shape = [batch, heads_q, seq_q, dim] + kv_shape = [batch, heads_k, seq_k, dim] + sink_shape = [heads_q] + out_shape = [batch, heads_q, seq_q, dim] + lse_shape = [batch, heads_q, seq_q] + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), + K: T.Tensor(kv_shape, dtype), + V: T.Tensor(kv_shape, dtype), + Output: T.Tensor(out_shape, dtype), + Lse: T.Tensor(lse_shape, accum_dtype), + Sinks: T.Tensor(sink_shape, dtype), + ): + with T.Kernel(T.ceildiv(seq_q, block_M), heads_q, batch, threads=threads) as ( + bx, + by, + bz, + ): + Q_shared = T.alloc_shared([block_M, dim], dtype) + K_shared = T.alloc_shared([block_N, dim], dtype) + V_shared = T.alloc_shared([block_N, dim], dtype) + O_shared = T.alloc_shared([block_M, dim], dtype) + acc_s = T.alloc_fragment([block_M, block_N], accum_dtype) + acc_s_cast = T.alloc_fragment([block_M, block_N], dtype) + acc_o = T.alloc_fragment([block_M, dim], accum_dtype) + scores_max = T.alloc_fragment([block_M], accum_dtype) + scores_max_prev = T.alloc_fragment([block_M], accum_dtype) + scores_scale = T.alloc_fragment([block_M], accum_dtype) + scores_sum = T.alloc_fragment([block_M], accum_dtype) + logsum = T.alloc_fragment([block_M], accum_dtype) + sinks_reg = T.alloc_fragment([block_M], dtype) + + # MQA broadcast: kv-head index for this query head. + k_head = by // groups + + T.copy(Q[bz, by, bx * block_M : (bx + 1) * block_M, :], Q_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(accum_dtype)) + if has_sink: + for i in T.Parallel(block_M): + sinks_reg[i] = Sinks[by] + + # K-tile loop bounds. Causal: end at the right edge of + # the m-tile. SWA: start at `max(0, (m+1 - W) // BN)`. + end = T.min(T.ceildiv(seq_k, block_N), T.ceildiv((bx + 1) * block_M, block_N)) + if swa_window > 0: + start = T.max(0, (bx * block_M - swa_window) // block_N) + else: + start = 0 + + for k in T.Pipelined(start, end, num_stages=num_stages): + T.copy(K[bz, k_head, k * block_N : (k + 1) * block_N, :], K_shared) + for i, j in T.Parallel(block_M, block_N): + q_idx = bx * block_M + i + k_idx = k * block_N + j + if swa_window > 0: + acc_s[i, j] = T.if_then_else( + q_idx >= k_idx and q_idx < k_idx + swa_window, + 0, + -T.infinity(acc_s.dtype), + ) + else: + acc_s[i, j] = T.if_then_else(q_idx >= k_idx, 0, -T.infinity(acc_s.dtype)) + T.gemm( + Q_shared, + K_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(accum_dtype)) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(block_M): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + # Check-inf for SWA: a tile that is fully outside the + # window has scores_max = -inf; reset to 0 to avoid + # NaN from `exp2(-inf - -inf)`. + for i in T.Parallel(block_M): + if swa_window > 0: + scores_max[i] = T.if_then_else( + scores_max[i] == -T.infinity(accum_dtype), + 0, + scores_max[i], + ) + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Parallel(block_M, block_N): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + T.reduce_sum(acc_s, scores_sum, dim=1) + for i in T.Parallel(block_M): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + T.copy(acc_s, acc_s_cast) + + for i, j in T.Parallel(block_M, dim): + acc_o[i, j] *= scores_scale[i] + + T.copy(V[bz, k_head, k * block_N : (k + 1) * block_N, :], V_shared) + T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + # Sink as virtual key column at end-of-loop. + if has_sink: + for i in T.Parallel(block_M): + logsum[i] += T.exp2(sinks_reg[i] * 1.44269504 - scores_max[i] * scale) + + # Final normalise + cast + LSE emit. + for i, j in T.Parallel(block_M, dim): + acc_o[i, j] /= logsum[i] + T.copy(acc_o, O_shared) + T.copy(O_shared, Output[bz, by, bx * block_M : (bx + 1) * block_M, :]) + + # LSE = ln(l_i) + m_i*sm_scale (note: scores were + # accumulated with `*scale = sm_scale*log2e`, so the + # base-e LSE is `ln(l_i) + m_i*sm_scale`). + for i in T.Parallel(block_M): + logsum[i] = T.log(logsum[i]) + scores_max[i] * sm_scale + T.copy(logsum, Lse[bz, by, bx * block_M : (bx + 1) * block_M]) + + return main + + +# --------------------------------------------------------------------------- +# Wrapper API +# --------------------------------------------------------------------------- + + +def _kernel_supports( + *, + sink: Optional[torch.Tensor], + swa_window: int, + additive_mask: Optional[torch.Tensor], + hca_local_seqlen: int, +) -> bool: + """Return True iff the P50 tilelang kernel covers this call site. + + P50 ships only the dense / SWA / sink subset. Additive-mask + + HCA split-mask are P52 territory; we fall back to Triton there. + """ + if additive_mask is not None: + return False + if hca_local_seqlen > 0: + return False + return True + + +def _torch_dtype_to_tilelang(dtype: torch.dtype) -> "T.dtype": + if dtype == torch.bfloat16: + return T.bfloat16 + if dtype == torch.float16: + return T.float16 + if dtype == torch.float32: + return T.float32 + raise ValueError(f"unsupported dtype {dtype}; expected bf16 / fp16 / fp32") + + +def _get_or_compile_kernel( + *, + batch: int, + heads_q: int, + heads_k: int, + seq_q: int, + seq_k: int, + dim: int, + has_sink: bool, + swa_window: int, + dtype: torch.dtype, +): + """Compile (or fetch from tilelang's per-args cache) the FWD kernel + for the given shape envelope. + + Tilelang's `@tilelang.jit` decorator caches compiled binaries per + (kwargs) tuple via its own cache, so this wrapper is itself fast + on cache hits. + """ + key = ( + batch, + heads_q, + heads_k, + seq_q, + seq_k, + dim, + has_sink, + int(swa_window > 0) * swa_window, + str(dtype), + ) + if key in _KERNEL_CACHE: + return _KERNEL_CACHE[key] + tl_dtype = _torch_dtype_to_tilelang(dtype) + # SMEM budget at head_dim=512 on MI355X is tight (160 KiB). + # Each `T.alloc_shared([B, dim], bf16)` allocates `B * dim * 2` + # bytes; the kernel uses Q + K + V + O shared = 4 * B_tile * D + # * 2. At D=512 even 32x32 hits some tilelang-internal hidden + # allocations that overshoot. Conservative choice: shrink to + # 16x32 with threads=64 at D >= 256. + if dim >= 256: + block_M, block_N, threads = 32, 32, 64 + else: + block_M, block_N, threads = 64, 64, 128 + kernel = _make_v4_attention_fwd_kernel( + batch, + heads_q, + heads_k, + seq_q, + seq_k, + dim, + has_sink, + swa_window, + dtype=tl_dtype, + block_M=block_M, + block_N=block_N, + threads=threads, + ) + _KERNEL_CACHE[key] = kernel + return kernel + + +def v4_attention_fwd_tilelang_with_lse( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + sink: Optional[torch.Tensor] = None, + additive_mask: Optional[torch.Tensor] = None, + swa_window: int = 0, + attn_dropout: float = 0.0, + training: bool = False, + scale: Optional[float] = None, + hca_local_seqlen: int = 0, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Same dispatch logic as :func:`v4_attention_fwd_tilelang` but + also returns the saved ``LSE`` tensor when the tilelang path + runs. + + When the wrapper falls back to Triton, ``LSE`` is computed by + the Triton launcher and returned alongside ``out``; the + autograd Function in P51 saves it for backward. + """ + if attn_dropout > 0.0 and training: + return _launch_v4_attention_fwd( + q, + k, + v, + sink=sink, + swa_window=swa_window, + additive_mask=additive_mask, + scale=scale if scale is not None else 1.0 / (q.shape[-1] ** 0.5), + hca_local_seqlen=hca_local_seqlen, + ) + if not _kernel_supports( + sink=sink, + swa_window=swa_window, + additive_mask=additive_mask, + hca_local_seqlen=hca_local_seqlen, + ): + return _launch_v4_attention_fwd( + q, + k, + v, + sink=sink, + swa_window=swa_window, + additive_mask=additive_mask, + scale=scale if scale is not None else 1.0 / (q.shape[-1] ** 0.5), + hca_local_seqlen=hca_local_seqlen, + ) + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise ValueError( + "v4_attention_fwd_tilelang expects q / k / v of rank 4 " + f"(got q.dim={q.dim()}, k.dim={k.dim()}, v.dim={v.dim()})" + ) + B, HQ, Sq, D = q.shape + _, HK, Sk, _ = k.shape + assert HK in (1, HQ) + has_sink = sink is not None + kernel = _get_or_compile_kernel( + batch=B, + heads_q=HQ, + heads_k=HK, + seq_q=Sq, + seq_k=Sk, + dim=D, + has_sink=has_sink, + swa_window=int(swa_window), + dtype=q.dtype, + ) + if not has_sink: + sink_arg = torch.zeros(HQ, device=q.device, dtype=q.dtype) + else: + sink_arg = sink + q_c = q.contiguous() + k_c = k.contiguous() + v_c = v.contiguous() + sink_c = sink_arg.contiguous() + out, lse = kernel(q_c, k_c, v_c, sink_c) + return out, lse + + +def v4_attention_fwd_tilelang( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + sink: Optional[torch.Tensor] = None, + additive_mask: Optional[torch.Tensor] = None, + swa_window: int = 0, + attn_dropout: float = 0.0, + training: bool = False, + scale: Optional[float] = None, + hca_local_seqlen: int = 0, +) -> torch.Tensor: + """Plan-8 P50 tilelang dense / SWA / sink FWD wrapper. + + Drop-in replacement for the relevant subset of ``v4_attention_v1`` + (the Triton autograd path). When the call site requires features + P50 doesn't cover (`additive_mask`, `hca_local_seqlen > 0`), the + wrapper falls back to the Triton kernel via + ``_launch_v4_attention_fwd`` so production smokes never see an + error in the regression window. + + Returns ``out [B, H, Sq, D]`` in ``v.dtype``. + + ``attn_dropout`` and ``training`` are accepted for parity with + the Triton wrapper but not used — dropout on attention is off + in production V4 (plan-4 P25 design note); P50 inherits that + contract. + + Note: this entry point is the **plain** (no-autograd) wrapper. + For the autograd-aware path used by V4AttentionFn, the call + site uses :func:`v4_attention_fwd_tilelang_with_lse` directly + so it can save ``LSE`` for backward. + """ + out, _ = v4_attention_fwd_tilelang_with_lse( + q, + k, + v, + sink=sink, + additive_mask=additive_mask, + swa_window=swa_window, + attn_dropout=attn_dropout, + training=training, + scale=scale, + hca_local_seqlen=hca_local_seqlen, + ) + return out + + +# --------------------------------------------------------------------------- +# Module-level registration — flips +# `_tilelang.is_tilelang_kernel_available("v4_attention_fwd")` to True +# AND overrides the parent module's stub with the real wrapper, so +# callers can do `_tilelang.v4_attention_fwd_tilelang(...)` and get +# the real implementation after lazy-load. +# --------------------------------------------------------------------------- + + +_tilelang.register_available_kernel("v4_attention_fwd") +# Note: the parent module's `v4_attention_fwd_tilelang` attribute +# is restored by `_tilelang._lazy_load(...)` after the import_module +# call (Python's import machinery sets it to the submodule otherwise). + + +__all__ = ["v4_attention_fwd_tilelang"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/__init__.py new file mode 100644 index 000000000..7d452efde --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/__init__.py @@ -0,0 +1,20 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Shared (non-attention) Triton kernels for the DeepSeek-V4 pipeline. + +These back the surrounding V4 components rather than the attention core, and +are imported by full submodule path (no package-level re-export): + +* :mod:`indexer_score` / :mod:`indexer_score_post` — lightning indexer scores +* :mod:`compressor_pool` — compressed-pool builder +* :mod:`hc_expand` / :mod:`hc_glue` — hierarchical-compression expand / glue +* :mod:`sinkhorn` — fused Sinkhorn-Knopp normalize (FWD/BWD) +* :mod:`rope_interleaved_partial` — fused interleaved partial RoPE (FWD/BWD) + +The attention kernels live in :mod:`.._triton_v0_deprecated` / :mod:`.._triton_v1` / +:mod:`.._triton_v2`. +""" diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/compressor_pool.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/compressor_pool.py new file mode 100644 index 000000000..64d0f7234 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/compressor_pool.py @@ -0,0 +1,283 @@ +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +"""Triton-fused softmax-weighted pool for the V4 Compressor (forward burst). + +Fuses the per-window-softmax pooling tail of :class:`Compressor.forward`:: + + score = score + ape # [B, N, W, hd] + weights = softmax(score.float(), dim=2) # over the window W + pooled = (kv * weights).sum(dim=2) # [B, N, hd] + +into a single forward kernel: one program per ``(b, n)`` row computes, for each +channel ``d``, the softmax over the ``W`` window slots of ``score[..,:,d]+ape[:,d]`` +and the weighted sum of ``kv[..,:,d]`` -- collapsing the ~5-launch forward burst +(add + cast + softmax + cast + mul + reduce) into one launch per compressed layer. + +The forward runs the reduction in fp32 (more accurate than the eager bf16-weights +path, numerically equivalent within bf16 noise). The backward is the explicit +analytic gradient in eager torch (softmax jacobian + weighted-sum), so there is no +second Triton kernel to validate; the forward burst -- which is what dominates the +compressed-layer forward -- is the part fused here. + +Gating: routed through :func:`fused_softmax_weighted_pool` when +``PRIMUS_COMPRESS_POOL_TRITON != "0"`` (default-on) for CUDA float16/bfloat16/float32 +inputs; falls back to eager only for non-CUDA / unsupported dtypes. Handles BOTH +the HCA (non-overlap, window ``W = ratio``) and CSA (overlap, window ``W = 2*ratio``) +pools — the kernel softmaxes over whatever window axis ``W`` it is given (verified +fp32-exact at both W=128 and W=8). +""" +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _compressor_pool_fwd_kernel( + kv_ptr, + score_ptr, + ape_ptr, + out_ptr, + R, + HD, + stride_kv_bn, + stride_kv_r, + stride_kv_d, + stride_sc_bn, + stride_sc_r, + stride_sc_d, + stride_ape_r, + stride_ape_d, + stride_out_bn, + stride_out_d, + BLOCK_R: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_bn = tl.program_id(0) + pid_d = tl.program_id(1) + offs_r = tl.arange(0, BLOCK_R) + offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + r_mask = offs_r < R + d_mask = offs_d < HD + mask = r_mask[:, None] & d_mask[None, :] + + sc_ptrs = ( + score_ptr + pid_bn * stride_sc_bn + offs_r[:, None] * stride_sc_r + offs_d[None, :] * stride_sc_d + ) + ape_ptrs = ape_ptr + offs_r[:, None] * stride_ape_r + offs_d[None, :] * stride_ape_d + s = tl.load(sc_ptrs, mask=mask, other=0.0).to(tl.float32) + a = tl.load(ape_ptrs, mask=mask, other=0.0).to(tl.float32) + s = s + a + + # softmax over the window axis R (axis 0), masking padded rows to -inf. + s = tl.where(r_mask[:, None], s, float("-inf")) + s_max = tl.max(s, axis=0) + e = tl.exp(s - s_max[None, :]) + e = tl.where(r_mask[:, None], e, 0.0) + denom = tl.sum(e, axis=0) + w = e / denom[None, :] # [BLOCK_R, BLOCK_D] fp32 softmax weights + + kv_ptrs = kv_ptr + pid_bn * stride_kv_bn + offs_r[:, None] * stride_kv_r + offs_d[None, :] * stride_kv_d + kv = tl.load(kv_ptrs, mask=mask, other=0.0).to(tl.float32) + pooled = tl.sum(kv * w, axis=0) # [BLOCK_D] + + out_ptrs = out_ptr + pid_bn * stride_out_bn + offs_d * stride_out_d + tl.store(out_ptrs, pooled.to(out_ptr.dtype.element_ty), mask=d_mask) + + +def _pool_fwd_triton(kv: torch.Tensor, score: torch.Tensor, ape: torch.Tensor) -> torch.Tensor: + """kv/score: ``[B, N, W, hd]``; ape: ``[W, hd]`` -> pooled ``[B, N, hd]``.""" + B, N, W, HD = kv.shape + bn = B * N + kv3 = kv.reshape(bn, W, HD) + sc3 = score.reshape(bn, W, HD) + out = torch.empty((bn, HD), dtype=kv.dtype, device=kv.device) + BLOCK_R = max(16, triton.next_power_of_2(W)) + BLOCK_D = min(64, max(16, triton.next_power_of_2(HD))) + grid = (bn, triton.cdiv(HD, BLOCK_D)) + _compressor_pool_fwd_kernel[grid]( + kv3, + sc3, + ape, + out, + W, + HD, + kv3.stride(0), + kv3.stride(1), + kv3.stride(2), + sc3.stride(0), + sc3.stride(1), + sc3.stride(2), + ape.stride(0), + ape.stride(1), + out.stride(0), + out.stride(1), + BLOCK_R=BLOCK_R, + BLOCK_D=BLOCK_D, + ) + return out.reshape(B, N, HD) + + +@triton.jit +def _compressor_pool_bwd_kernel( + kv_ptr, + score_ptr, + ape_ptr, + dout_ptr, + dkv_ptr, + dscore_ptr, + R, + HD, + stride_kv_bn, + stride_kv_r, + stride_kv_d, + stride_sc_bn, + stride_sc_r, + stride_sc_d, + stride_ape_r, + stride_ape_d, + stride_do_bn, + stride_do_d, + stride_dkv_bn, + stride_dkv_r, + stride_dkv_d, + stride_dsc_bn, + stride_dsc_r, + stride_dsc_d, + BLOCK_R: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_bn = tl.program_id(0) + pid_d = tl.program_id(1) + offs_r = tl.arange(0, BLOCK_R) + offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + r_mask = offs_r < R + d_mask = offs_d < HD + mask = r_mask[:, None] & d_mask[None, :] + + # Recompute the fwd softmax weights w over the window axis R (axis 0), fp32. + sc_ptrs = ( + score_ptr + pid_bn * stride_sc_bn + offs_r[:, None] * stride_sc_r + offs_d[None, :] * stride_sc_d + ) + ape_ptrs = ape_ptr + offs_r[:, None] * stride_ape_r + offs_d[None, :] * stride_ape_d + s = tl.load(sc_ptrs, mask=mask, other=0.0).to(tl.float32) + a = tl.load(ape_ptrs, mask=mask, other=0.0).to(tl.float32) + s = s + a + s = tl.where(r_mask[:, None], s, float("-inf")) + s_max = tl.max(s, axis=0) + e = tl.exp(s - s_max[None, :]) + e = tl.where(r_mask[:, None], e, 0.0) + denom = tl.sum(e, axis=0) + w = e / denom[None, :] # [BLOCK_R, BLOCK_D] + + kv_ptrs = kv_ptr + pid_bn * stride_kv_bn + offs_r[:, None] * stride_kv_r + offs_d[None, :] * stride_kv_d + kv = tl.load(kv_ptrs, mask=mask, other=0.0).to(tl.float32) + do_ptrs = dout_ptr + pid_bn * stride_do_bn + offs_d * stride_do_d + dout = tl.load(do_ptrs, mask=d_mask, other=0.0).to(tl.float32) # [BLOCK_D] + + # pooled = sum_r w*kv -> dkv = dout*w ; dw = dout*kv ; + # softmax jacobian: dscore = w * (dw - sum_r(dw*w)) + dkv = dout[None, :] * w + dw = dout[None, :] * kv + sum_dw_w = tl.sum(dw * w, axis=0) # [BLOCK_D] + dscore = w * (dw - sum_dw_w[None, :]) + + dkv_ptrs = ( + dkv_ptr + pid_bn * stride_dkv_bn + offs_r[:, None] * stride_dkv_r + offs_d[None, :] * stride_dkv_d + ) + tl.store(dkv_ptrs, dkv.to(dkv_ptr.dtype.element_ty), mask=mask) + dsc_ptrs = ( + dscore_ptr + pid_bn * stride_dsc_bn + offs_r[:, None] * stride_dsc_r + offs_d[None, :] * stride_dsc_d + ) + tl.store(dsc_ptrs, dscore.to(dscore_ptr.dtype.element_ty), mask=mask) + + +def _pool_bwd_triton(dout, kv, score, ape): + """Fused backward of the softmax-weighted pool. Returns (dkv, dscore, dape). + + Recomputes w in-register (like the fwd) and emits dkv + dscore in one kernel; + dape = sum over the (B,N) rows of dscore is a single trailing reduce. + """ + B, N, W, HD = kv.shape + bn = B * N + kv3 = kv.reshape(bn, W, HD) + sc3 = score.reshape(bn, W, HD) + do2 = dout.reshape(bn, HD) + dkv = torch.empty_like(kv3) + dscore = torch.empty_like(sc3) + BLOCK_R = max(16, triton.next_power_of_2(W)) + BLOCK_D = min(64, max(16, triton.next_power_of_2(HD))) + grid = (bn, triton.cdiv(HD, BLOCK_D)) + _compressor_pool_bwd_kernel[grid]( + kv3, + sc3, + ape, + do2, + dkv, + dscore, + W, + HD, + kv3.stride(0), + kv3.stride(1), + kv3.stride(2), + sc3.stride(0), + sc3.stride(1), + sc3.stride(2), + ape.stride(0), + ape.stride(1), + do2.stride(0), + do2.stride(1), + dkv.stride(0), + dkv.stride(1), + dkv.stride(2), + dscore.stride(0), + dscore.stride(1), + dscore.stride(2), + BLOCK_R=BLOCK_R, + BLOCK_D=BLOCK_D, + ) + dkv = dkv.reshape(B, N, W, HD) + dscore = dscore.reshape(B, N, W, HD) + dape = dscore.sum(dim=(0, 1)).to(ape.dtype) + return dkv, dscore, dape + + +class _CompressorSoftmaxPool(torch.autograd.Function): + @staticmethod + def forward(ctx, kv, score, ape): + pooled = _pool_fwd_triton(kv, score, ape) + ctx.save_for_backward(kv, score, ape) + return pooled + + @staticmethod + def backward(ctx, dout): + kv, score, ape = ctx.saved_tensors + # Fused Triton backward (default-on): recompute softmax + emit dkv/dscore in + # one launch instead of the ~14-kernel eager analytic burst. Handles both the + # HCA (W=ratio) and CSA (W=2*ratio overlap) pools -- the window fits one tile + # in either case. PRIMUS_COMPRESS_POOL_TRITON_BWD=0 restores the eager path. + if os.environ.get("PRIMUS_COMPRESS_POOL_TRITON_BWD", "1") != "0" and pool_supported(kv): + return _pool_bwd_triton(dout.contiguous(), kv, score, ape) + # Analytic gradient of pooled = sum_w softmax_w(score+ape) * kv (fp32). + w = torch.softmax((score + ape).float(), dim=2) # [B, N, W, hd] + dout3 = dout.unsqueeze(2).float() # [B, N, 1, hd] + dkv = (dout3 * w).to(kv.dtype) + dw = dout3 * kv.float() + # softmax jacobian: dscore = w * (dw - sum_w(dw * w)) + dscore = (w * (dw - (dw * w).sum(dim=2, keepdim=True))).to(score.dtype) + dape = dscore.sum(dim=(0, 1)).to(ape.dtype) + return dkv, dscore, dape + + +def fused_softmax_weighted_pool(kv: torch.Tensor, score: torch.Tensor, ape: torch.Tensor) -> torch.Tensor: + """Fused ``(softmax(score+ape, dim=2) * kv).sum(dim=2)``. + + kv/score: ``[B, N, W, hd]``; ape: ``[W, hd]`` -> ``[B, N, hd]``. + """ + return _CompressorSoftmaxPool.apply(kv, score, ape) + + +def pool_supported(kv: torch.Tensor) -> bool: + return kv.is_cuda and kv.dtype in (torch.float16, torch.bfloat16, torch.float32) diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_collapse.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_collapse.py new file mode 100644 index 000000000..10f905c4b --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_collapse.py @@ -0,0 +1,232 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton-fused HyperConnection ``collapse`` (small-kernel-fusion 2026-07-03). + +Single FWD + single BWD kernel pair for the eager body of +:meth:`primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.collapse`:: + + out[..., d] = sum_k pre[..., k] * x[..., k, d] + +The eager body ``(pre.unsqueeze(-1) * x).sum(dim=-2)`` materialises a full +``[..., K, D]`` temporary (the broadcast multiply) and then reduces it — two +kernels + ``K*D`` of extra HBM traffic per call. This kernel streams ``D`` and +contracts the small ``K`` (constexpr) in registers, writing only the ``[..., D]`` +result. Symmetric to the already-shipped ``expand`` fusion (``hc_expand.py``). + +Gradients:: + + dx[..., k, d] = pre[..., k] * g[..., d] + dpre[..., k] = sum_d x[..., k, d] * g[..., d] + +Gating: routed through :func:`hc_collapse_triton` when +``PRIMUS_HC_COLLAPSE_TRITON != "0"`` (default-on); falls back to eager for +unsupported configs (CPU input, K out of range, dtype mismatch). +""" + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +_SUPPORTED_K = (1, 2, 4, 8, 16) + +_BLOCK_M = 16 +_BLOCK_D = 128 +_NUM_WARPS = 4 + + +@triton.jit +def _hc_collapse_fwd_kernel( + X_PTR, # [M, K, D] + PRE_PTR, # [M, K] + OUT_PTR, # [M, D] OUT + M, + D, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_d = tl.program_id(1) + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rd = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + k_ax = tl.arange(0, K) + mask_m = rm < M + mask_d = rd < D + mask2 = mask_m[:, None] & mask_d[None, :] + mask3 = mask_m[:, None, None] & mask_d[None, None, :] + + KD = K * D + x3 = tl.load( + X_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :], + mask=mask3, + other=0.0, + ).to(tl.float32) + pre = tl.load(PRE_PTR + rm[:, None] * K + k_ax[None, :], mask=mask_m[:, None], other=0.0).to(tl.float32) + out = tl.sum(pre[:, :, None] * x3, axis=1) # [BLOCK_M, BLOCK_D] + tl.store(OUT_PTR + rm[:, None] * D + rd[None, :], out.to(OUT_PTR.dtype.element_ty), mask=mask2) + + +@triton.jit +def _hc_collapse_bwd_kernel( + G_PTR, # [M, D] grad of out + X_PTR, # [M, K, D] + PRE_PTR, # [M, K] + DX_PTR, # [M, K, D] OUT + DPRE_P_PTR, # [NUMD, M, K] partial OUT (reduced host-side) + M, + D, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_d = tl.program_id(1) + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rd = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + k_ax = tl.arange(0, K) + mask_m = rm < M + mask_d = rd < D + mask2 = mask_m[:, None] & mask_d[None, :] + mask3 = mask_m[:, None, None] & mask_d[None, None, :] + + KD = K * D + x3 = tl.load( + X_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :], + mask=mask3, + other=0.0, + ).to(tl.float32) + g = tl.load(G_PTR + rm[:, None] * D + rd[None, :], mask=mask2, other=0.0).to(tl.float32) + pre = tl.load(PRE_PTR + rm[:, None] * K + k_ax[None, :], mask=mask_m[:, None], other=0.0).to(tl.float32) + + # dx[k, d] = pre[k] * g[d] + dx = pre[:, :, None] * g[:, None, :] # [BLOCK_M, K, BLOCK_D] + tl.store( + DX_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :], + dx.to(DX_PTR.dtype.element_ty), + mask=mask3, + ) + + # dpre[k] = sum_d x[k, d] * g[d] (partial over this d-block) + dpre = tl.sum(x3 * g[:, None, :], axis=2) # [BLOCK_M, K] + tl.store( + DPRE_P_PTR + pid_d * (M * K) + rm[:, None] * K + k_ax[None, :], + dpre, + mask=mask_m[:, None], + ) + + +class HCCollapseFn(torch.autograd.Function): + """Autograd wrapper for the fused ``collapse`` FWD/BWD Triton kernels. + + ``out[..., d] = sum_k pre[..., k] * x[..., k, d]`` + """ + + @staticmethod + def forward(ctx, x, pre): # type: ignore[override] + K = x.shape[-2] + D = x.shape[-1] + leading = x.shape[:-2] + M = 1 + for s in leading: + M *= s + + x_c = x.contiguous() + pre_c = pre.contiguous() + out = torch.empty((*leading, D), dtype=x.dtype, device=x.device) + + grid = (triton.cdiv(M, _BLOCK_M), triton.cdiv(D, _BLOCK_D)) + _hc_collapse_fwd_kernel[grid]( + x_c, + pre_c, + out, + M, + D, + K=K, + BLOCK_M=_BLOCK_M, + BLOCK_D=_BLOCK_D, + num_warps=_NUM_WARPS, + ) + + ctx.save_for_backward(x_c, pre_c) + ctx.leading = tuple(leading) + ctx.K = K + ctx.D = D + ctx.M = M + return out + + @staticmethod + def backward(ctx, g): # type: ignore[override] + x_c, pre_c = ctx.saved_tensors + leading, K, D, M = ctx.leading, ctx.K, ctx.D, ctx.M + g = g.contiguous() + + dx = torch.empty_like(x_c) + num_d = triton.cdiv(D, _BLOCK_D) + dpre_p = torch.empty((num_d, M, K), dtype=torch.float32, device=x_c.device) + + grid = (triton.cdiv(M, _BLOCK_M), num_d) + _hc_collapse_bwd_kernel[grid]( + g, + x_c, + pre_c, + dx, + dpre_p, + M, + D, + K=K, + BLOCK_M=_BLOCK_M, + BLOCK_D=_BLOCK_D, + num_warps=_NUM_WARPS, + ) + + d_pre = dpre_p.sum(dim=0).to(pre_c.dtype).view(*leading, K) + return dx, d_pre + + +def is_triton_path_enabled() -> bool: + """``PRIMUS_HC_COLLAPSE_TRITON`` knob; default-on, set ``"0"`` to disable.""" + return os.environ.get("PRIMUS_HC_COLLAPSE_TRITON", "1") != "0" + + +def is_triton_kernel_supported(x: torch.Tensor, pre: torch.Tensor) -> bool: + if not (x.is_cuda and pre.is_cuda): + return False + if x.dim() < 2: + return False + K = x.shape[-2] + if K not in _SUPPORTED_K: + return False + if pre.shape[-1] != K: + return False + if x.dtype != pre.dtype or x.dtype not in (torch.float16, torch.bfloat16, torch.float32): + return False + return True + + +def eager_hc_collapse(x: torch.Tensor, pre: torch.Tensor) -> torch.Tensor: + """Reference eager ``collapse`` (matches HyperMixer.collapse bit-for-bit).""" + return (pre.unsqueeze(-1) * x).sum(dim=-2) + + +def hc_collapse_triton(x, pre): + """Fused ``out[...,d] = sum_k pre[...,k] * x[...,k,d]`` with eager fallback.""" + if is_triton_path_enabled() and is_triton_kernel_supported(x, pre): + return HCCollapseFn.apply(x, pre) + return eager_hc_collapse(x, pre) + + +__all__ = [ + "HCCollapseFn", + "hc_collapse_triton", + "eager_hc_collapse", + "is_triton_path_enabled", + "is_triton_kernel_supported", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_expand.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_expand.py new file mode 100644 index 000000000..4bb7724b4 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_expand.py @@ -0,0 +1,287 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton-fused HyperConnection ``expand``. + +Single FWD + single BWD kernel pair for the eager body of +:meth:`primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.expand`:: + + new[..., h, d] = post[..., h] * out[..., d] + sum_k comb[..., h, k] * x[..., k, d] + +The contraction over the small ``K`` is done in registers (unrolled, ``K`` is a +constexpr) while ``D`` is streamed; the ``post`` outer-product and the final add +are folded into the same kernel. + +Gating: routed through :func:`hc_expand_triton` when +``PRIMUS_HC_EXPAND_TRITON != "0"`` (default-on); falls back to eager for +unsupported configs (CPU input, K out of range, dtype mismatch). + +Supported shapes: + +* ``K`` (= ``hc_mult``) in ``{1, 2, 4, 8, 16}``. +* ``x``/``out``/``post``/``comb`` share one floating dtype and are CUDA tensors. +* Any leading shape (the wrapper flattens to ``[M, K, D]``). +""" + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +_SUPPORTED_K = (1, 2, 4, 8, 16) + +_BLOCK_M = 16 +_BLOCK_D = 128 +_NUM_WARPS = 4 + + +@triton.jit +def _hc_expand_fwd_kernel( + X_PTR, # [M, K, D] + OUT_PTR, # [M, D] + POST_PTR, # [M, K] + COMB_PTR, # [M, K, K] + NEW_PTR, # [M, K, D] OUT + M, + D, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_d = tl.program_id(1) + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rd = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + k_ax = tl.arange(0, K) + mask_m = rm < M + mask_d = rd < D + mask2 = mask_m[:, None] & mask_d[None, :] + mask3 = mask_m[:, None, None] & mask_d[None, None, :] + + KD = K * D + KK = K * K + + # x as [BLOCK_M, K(stream), BLOCK_D], loaded once. + x3 = tl.load( + X_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :], + mask=mask3, + other=0.0, + ).to(tl.float32) + out_tile = tl.load(OUT_PTR + rm[:, None] * D + rd[None, :], mask=mask2, other=0.0).to(tl.float32) + + for h in range(K): + comb_h = tl.load( + COMB_PTR + rm[:, None] * KK + h * K + k_ax[None, :], mask=mask_m[:, None], other=0.0 + ).to( + tl.float32 + ) # [BLOCK_M, K] + post_h = tl.load(POST_PTR + rm * K + h, mask=mask_m, other=0.0).to(tl.float32) # [BLOCK_M] + # mix_h[d] = sum_k comb_h[k] * x3[k, d] + mix_h = tl.sum(comb_h[:, :, None] * x3, axis=1) + new_h = post_h[:, None] * out_tile + mix_h + tl.store( + NEW_PTR + rm[:, None] * KD + h * D + rd[None, :], + new_h.to(NEW_PTR.dtype.element_ty), + mask=mask2, + ) + + +@triton.jit +def _hc_expand_bwd_kernel( + G_PTR, # [M, K, D] grad of new + X_PTR, # [M, K, D] + OUT_PTR, # [M, D] + POST_PTR, # [M, K] + COMB_PTR, # [M, K, K] + DX_PTR, # [M, K, D] OUT + DOUT_PTR, # [M, D] OUT + DPOST_P_PTR, # [NUMD, M, K] partial OUT (reduced host-side) + DCOMB_P_PTR, # [NUMD, M, K, K] partial OUT (reduced host-side) + M, + D, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_d = tl.program_id(1) + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + rd = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + k_ax = tl.arange(0, K) + mask_m = rm < M + mask_d = rd < D + mask2 = mask_m[:, None] & mask_d[None, :] + mask3 = mask_m[:, None, None] & mask_d[None, None, :] + + KD = K * D + KK = K * K + + # x3, g3 as [BLOCK_M, K(stream), BLOCK_D]; the stream axis of g indexes h. + x3 = tl.load( + X_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :], + mask=mask3, + other=0.0, + ).to(tl.float32) + g3 = tl.load( + G_PTR + rm[:, None, None] * KD + k_ax[None, :, None] * D + rd[None, None, :], + mask=mask3, + other=0.0, + ).to(tl.float32) + out_tile = tl.load(OUT_PTR + rm[:, None] * D + rd[None, :], mask=mask2, other=0.0).to(tl.float32) + post = tl.load(POST_PTR + rm[:, None] * K + k_ax[None, :], mask=mask_m[:, None], other=0.0).to( + tl.float32 + ) # [BLOCK_M, K(h)] + + # d_out[d] = sum_h post[h] * g[h, d] + dout = tl.sum(post[:, :, None] * g3, axis=1) + tl.store(DOUT_PTR + rm[:, None] * D + rd[None, :], dout.to(DOUT_PTR.dtype.element_ty), mask=mask2) + + # d_post[h] = sum_d out[d] * g[h, d] (partial over this d-block) + dpost = tl.sum(out_tile[:, None, :] * g3, axis=2) # [BLOCK_M, K(h)] + tl.store( + DPOST_P_PTR + pid_d * (M * K) + rm[:, None] * K + k_ax[None, :], + dpost, + mask=mask_m[:, None], + ) + + # d_x[k, d] = sum_h comb[h, k] * g[h, d] + for k in range(K): + comb_col = tl.load( + COMB_PTR + rm[:, None] * KK + k_ax[None, :] * K + k, mask=mask_m[:, None], other=0.0 + ).to( + tl.float32 + ) # [BLOCK_M, K(h)] + dxk = tl.sum(comb_col[:, :, None] * g3, axis=1) + tl.store(DX_PTR + rm[:, None] * KD + k * D + rd[None, :], dxk.to(DX_PTR.dtype.element_ty), mask=mask2) + + # d_comb[h, k] = sum_d x[k, d] * g[h, d] (partial over this d-block) + for h in range(K): + g_h = tl.load(G_PTR + rm[:, None] * KD + h * D + rd[None, :], mask=mask2, other=0.0).to(tl.float32) + dcomb_h = tl.sum(x3 * g_h[:, None, :], axis=2) # [BLOCK_M, K(k)] + tl.store( + DCOMB_P_PTR + pid_d * (M * KK) + rm[:, None] * KK + h * K + k_ax[None, :], + dcomb_h, + mask=mask_m[:, None], + ) + + +class HCExpandFn(torch.autograd.Function): + """Autograd wrapper around the fused ``expand`` FWD/BWD Triton kernels. + + ``new[..., h, d] = post[..., h] * out[..., d] + sum_k comb[..., h, k] * x[..., k, d]`` + """ + + @staticmethod + def forward(ctx, x, out, post, comb): # type: ignore[override] + K = comb.shape[-1] + D = x.shape[-1] + leading = x.shape[:-2] + M = 1 + for s in leading: + M *= s + + x_c = x.contiguous() + out_c = out.contiguous() + post_c = post.contiguous() + comb_c = comb.contiguous() + + new = torch.empty((*leading, K, D), dtype=x.dtype, device=x.device) + + grid = (triton.cdiv(M, _BLOCK_M), triton.cdiv(D, _BLOCK_D)) + _hc_expand_fwd_kernel[grid]( + x_c, + out_c, + post_c, + comb_c, + new, + M, + D, + K=K, + BLOCK_M=_BLOCK_M, + BLOCK_D=_BLOCK_D, + num_warps=_NUM_WARPS, + ) + + ctx.save_for_backward(x_c, out_c, post_c, comb_c) + ctx.leading = tuple(leading) + ctx.K = K + ctx.D = D + ctx.M = M + return new + + @staticmethod + def backward(ctx, g): # type: ignore[override] + x_c, out_c, post_c, comb_c = ctx.saved_tensors + leading, K, D, M = ctx.leading, ctx.K, ctx.D, ctx.M + g = g.contiguous() + + device = x_c.device + dx = torch.empty_like(x_c) + dout = torch.empty_like(out_c) + num_d = triton.cdiv(D, _BLOCK_D) + dpost_p = torch.empty((num_d, M, K), dtype=torch.float32, device=device) + dcomb_p = torch.empty((num_d, M, K, K), dtype=torch.float32, device=device) + + grid = (triton.cdiv(M, _BLOCK_M), num_d) + _hc_expand_bwd_kernel[grid]( + g, + x_c, + out_c, + post_c, + comb_c, + dx, + dout, + dpost_p, + dcomb_p, + M, + D, + K=K, + BLOCK_M=_BLOCK_M, + BLOCK_D=_BLOCK_D, + num_warps=_NUM_WARPS, + ) + + d_post = dpost_p.sum(dim=0).to(post_c.dtype).view(*leading, K) + d_comb = dcomb_p.sum(dim=0).to(comb_c.dtype).view(*leading, K, K) + return dx, dout, d_post, d_comb + + +def is_triton_path_enabled() -> bool: + """``PRIMUS_HC_EXPAND_TRITON`` knob; default-on, set ``"0"`` to disable.""" + return os.environ.get("PRIMUS_HC_EXPAND_TRITON", "1") != "0" + + +def is_triton_kernel_supported(x: torch.Tensor, post: torch.Tensor, comb: torch.Tensor) -> bool: + if not (x.is_cuda and post.is_cuda and comb.is_cuda): + return False + K = comb.shape[-1] + if K not in _SUPPORTED_K: + return False + if comb.shape[-2] != K or x.shape[-2] != K or post.shape[-1] != K: + return False + if not (x.dtype == post.dtype == comb.dtype) or x.dtype not in ( + torch.float16, + torch.bfloat16, + torch.float32, + ): + return False + return True + + +def hc_expand_triton(x, out, post, comb): + """Fused ``new[...,h,d] = post[h]*out[d] + sum_k comb[h,k]*x[k,d]``.""" + return HCExpandFn.apply(x, out, post, comb) + + +__all__ = [ + "HCExpandFn", + "hc_expand_triton", + "is_triton_path_enabled", + "is_triton_kernel_supported", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_glue.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_glue.py new file mode 100644 index 000000000..dc39ddd81 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/hc_glue.py @@ -0,0 +1,622 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton-fused HyperConnection elemwise glue (plan-6 P37). + +Fuses the post-linear elemwise tail of +:meth:`primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.compute_weights` +into a single FWD + single BWD Triton kernel pair. Eager body: + +.. code-block:: python + + pre_logit = logits[..., :K] * scale[0] + base[:K] + post_logit = logits[..., K:2K] * scale[1] + base[K:2K] + comb_logit = (logits[..., 2K:].view(..., K, K) * scale[2] + + base[2K:].view(K, K)) + + pre = sigmoid(pre_logit) + eps # (eps, 1+eps] + post = 2 * sigmoid(post_logit) # (0, 2) + comb_pre_sinkhorn = softmax(comb_logit, dim=-1) + eps + +After P36 the trailing ``sinkhorn_normalize(comb_pre_sinkhorn)`` is its +own Triton kernel. P37 fuses everything **between** the +``_packed_logits`` GEMM and the Sinkhorn call into one kernel: + +* 3 slices of ``logits`` (free in Triton — pointer arithmetic); +* 3 fused-multiply-adds against ``scale`` / ``base``; +* 2 sigmoid + 1 softmax + 2 eps adds; +* 1 final 2x multiply for ``post``. + +The eager chain is ~8 elementwise GPU launches per call (P32 trace +attributes ~3-5 ms / iter across all 16 ``HyperConnection`` invocations +to the ``elementwise_kernel_manual_unroll<128, 8>`` bucket). P37 +collapses those into **one** FWD kernel + **one** BWD kernel. + +The matmul inside ``_packed_logits`` (`F.linear(flat * rsqrt, W)`) +stays as ``torch.nn.functional.linear`` -- GEMM is already +GPU-bound and fusing it into the elemwise chain would re-implement a +GEMM badly. Same goes for the ``collapse`` reduce and the +``expand`` outer-product + matmul -- those live around matmuls and +are not net wins as separate Triton kernels (deferred to a possible +P37b if the residual trace shows otherwise). + +Gating: routed through :func:`hc_glue_compute_tail_triton` when +``PRIMUS_HC_TRITON != "0"`` (default-on). Set to ``"0"`` to fall back +to the eager body kept in +:func:`primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.compute_weights`. + +Supported shape constraints: + +* ``K`` (= ``hc_mult``) must be a power of 2 in ``{1, 2, 4, 8, 16}``. + V4-Flash uses ``K=4``. +* ``logits`` must be contiguous; the wrapper calls ``.contiguous()`` + defensively. +* Any leading shape is supported (the wrapper flattens to ``[N, (2+K)*K]``). +""" + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +_TORCH_TO_TL_DTYPE = { + torch.float64: tl.float64, + torch.float32: tl.float32, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, +} + + +def _triton_dtype(t: torch.dtype): + try: + return _TORCH_TO_TL_DTYPE[t] + except KeyError as exc: + raise TypeError( + f"hc_glue: unsupported dtype {t}; expected one of {list(_TORCH_TO_TL_DTYPE)}" + ) from exc + + +# Triton requires power-of-2 block extents, and the register footprint is +# linear in K + K*K, so we cap K at 16 for the in-register path. V4-Flash +# uses K=4; larger K would require a different layout. +_SUPPORTED_K = (1, 2, 4, 8, 16) + + +# --------------------------------------------------------------------------- +# Triton kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _hc_compute_tail_fwd_kernel( + LOGITS_PTR, # [N, (2+K)*K] contiguous, fp32 + SCALE_PTR, # [3] fp32 + BASE_PTR, # [(2+K)*K] fp32 + PRE_PTR, # [N, K] OUT_DTYPE + POST_PTR, # [N, K] OUT_DTYPE + COMB_PTR, # [N, K, K] OUT_DTYPE (softmax+eps; before Sinkhorn) + PRE_SIG_PTR, # [N, K] fp32 -- saved-for-backward (sigmoid(pre_logit)) + POST_SIG_PTR, # [N, K] fp32 -- saved-for-backward (sigmoid(post_logit)) + COMB_SM_PTR, # [N, K, K] fp32 -- saved-for-backward (softmax(comb_logit)) + N, + EPS: tl.constexpr, + K: tl.constexpr, + BLOCK_LEADING: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + """One program tile of ``BLOCK_LEADING`` rows of the leading axis. + + Reads ``LOGITS[BLOCK_LEADING, (2+K)*K]`` once, applies the elemwise + tail in fp32 registers, writes the three output tensors (cast to + OUT_DTYPE) plus three fp32 saved-for-backward states. + + Register footprint per program (at K=4, BLOCK_LEADING=64): + BLOCK_LEADING * (2+K)*K = 64 * 24 = 1536 fp32 logits + BLOCK_LEADING * K = 64 * 4 = 256 pre/post + BLOCK_LEADING * K * K = 64 * 16 = 1024 comb + Total ~2800 fp32 = ~11 KiB; comfortable for MI355 256-VGPR warps. + """ + + pid = tl.program_id(0) + offs = pid * BLOCK_LEADING + tl.arange(0, BLOCK_LEADING) + mask_leading = offs < N + + k_idx = tl.arange(0, K) + + # Load the three scale scalars (broadcast across all rows). + scale0 = tl.load(SCALE_PTR + 0) + scale1 = tl.load(SCALE_PTR + 1) + scale2 = tl.load(SCALE_PTR + 2) + + # Base partitions: base[:K], base[K:2K], base[2K:].view(K, K). + base_pre = tl.load(BASE_PTR + k_idx) + base_post = tl.load(BASE_PTR + K + k_idx) + base_comb = tl.load(BASE_PTR + 2 * K + k_idx[:, None] * K + k_idx[None, :]) + + KK_TOTAL: tl.constexpr = (2 + K) * K + + # Logits row stride is KK_TOTAL. Pre slice = [:, :K]; post = [:, K:2K]; + # comb = [:, 2K:].view(K, K). All three are contiguous reads. + pre_logit = tl.load( + LOGITS_PTR + offs[:, None] * KK_TOTAL + k_idx[None, :], + mask=mask_leading[:, None], + other=0.0, + ) + post_logit = tl.load( + LOGITS_PTR + offs[:, None] * KK_TOTAL + K + k_idx[None, :], + mask=mask_leading[:, None], + other=0.0, + ) + comb_logit = tl.load( + LOGITS_PTR + offs[:, None, None] * KK_TOTAL + 2 * K + k_idx[None, :, None] * K + k_idx[None, None, :], + mask=mask_leading[:, None, None], + other=0.0, + ) + + # Apply scale + base. + pre_logit = pre_logit * scale0 + base_pre[None, :] + post_logit = post_logit * scale1 + base_post[None, :] + comb_logit = comb_logit * scale2 + base_comb[None, :, :] + + # Sigmoids (per-element). + pre_sig = tl.sigmoid(pre_logit) + post_sig = tl.sigmoid(post_logit) + + # Softmax along axis=2 (the inner K of comb_logit). Standard + # numerically-stable softmax: subtract row max, exp, divide by sum. + cmax = tl.max(comb_logit, axis=2, keep_dims=True) + cexp = tl.exp(comb_logit - cmax) + csum = tl.sum(cexp, axis=2, keep_dims=True) + comb_sm = cexp / csum + + # Output tensors (with eps additions and the 2x for post). + pre_out = pre_sig + EPS + post_out = 2.0 * post_sig + comb_out = comb_sm + EPS + + # Save fp32 states for backward. + tl.store( + PRE_SIG_PTR + offs[:, None] * K + k_idx[None, :], + pre_sig, + mask=mask_leading[:, None], + ) + tl.store( + POST_SIG_PTR + offs[:, None] * K + k_idx[None, :], + post_sig, + mask=mask_leading[:, None], + ) + tl.store( + COMB_SM_PTR + offs[:, None, None] * K * K + k_idx[None, :, None] * K + k_idx[None, None, :], + comb_sm, + mask=mask_leading[:, None, None], + ) + + # Write the three caller-visible outputs (cast to OUT_DTYPE). + tl.store( + PRE_PTR + offs[:, None] * K + k_idx[None, :], + pre_out.to(OUT_DTYPE), + mask=mask_leading[:, None], + ) + tl.store( + POST_PTR + offs[:, None] * K + k_idx[None, :], + post_out.to(OUT_DTYPE), + mask=mask_leading[:, None], + ) + tl.store( + COMB_PTR + offs[:, None, None] * K * K + k_idx[None, :, None] * K + k_idx[None, None, :], + comb_out.to(OUT_DTYPE), + mask=mask_leading[:, None, None], + ) + + +@triton.jit +def _hc_compute_tail_bwd_kernel( + DPRE_PTR, # [N, K] grad in OUT_DTYPE + DPOST_PTR, # [N, K] grad in OUT_DTYPE + DCOMB_PTR, # [N, K, K] grad in OUT_DTYPE + PRE_SIG_PTR, # [N, K] fp32 saved + POST_SIG_PTR, # [N, K] fp32 saved + COMB_SM_PTR, # [N, K, K] fp32 saved + SCALE_PTR, # [3] fp32 (read) + DLOGITS_PTR, # [N, (2+K)*K] OUT (fp32) + DSCALE_PTR, # [N, 3] partial sums (fp32); reduced host-side + DBASE_PTR, # [N, (2+K)*K] partial sums (fp32); reduced host-side + N, + K: tl.constexpr, + BLOCK_LEADING: tl.constexpr, +): + """VJP through the elemwise tail. + + Forward: + pre_logit = logits[:, :K] * s0 + base[:K] + post_logit = logits[:, K:2K] * s1 + base[K:2K] + comb_logit = logits[:, 2K:].view(K, K) * s2 + base[2K:].view(K, K) + pre = sigmoid(pre_logit) + eps + post = 2 * sigmoid(post_logit) + comb_sm = softmax(comb_logit, axis=-1); comb = comb_sm + eps + + Backward (per element, all in fp32): + d_pre_logit = d_pre * pre_sig * (1 - pre_sig) + d_post_logit = d_post * 2 * post_sig * (1 - post_sig) + d_comb_logit = comb_sm * (d_comb - sum(d_comb * comb_sm, axis=-1)) + + d_logits[:, :K] = d_pre_logit * s0 + d_logits[:, K:2K] = d_post_logit * s1 + d_logits[:, 2K:] = d_comb_logit.view(K*K) * s2 + + d_scale[0] = sum(logits[:, :K] * d_pre_logit) + d_scale[1] = sum(logits[:, K:2K] * d_post_logit) + d_scale[2] = sum(logits[:, 2K:].view(K, K) * d_comb_logit) + + d_base[:K] = sum_n d_pre_logit + d_base[K:2K] = sum_n d_post_logit + d_base[2K:] = sum_n d_comb_logit + + Notes: + * d_scale needs the *original* fp32 logits. To avoid carrying + a fourth saved tensor we re-derive logits from the existing + state: that's not possible (sigmoid is not invertible at + large magnitudes). Instead, we save d_scale as a per-row + partial sum that the host-side wrapper reduces with a + torch.sum at the end. Wait -- we actually need the logits + themselves for d_scale. Solution: store ``logits * d_*`` + (the cross-term) directly without ever materialising logits + again -- this is just the elementwise product of the + recomputed (d_pre_logit, d_post_logit, d_comb_logit) with + the *input-side* gradient of each, taken at the pre-scale + point. Concretely, since + d_pre = sigmoid(s0 * x + b) (forward eq) + ∂loss/∂s0 = ∂loss/∂pre_logit * x (chain rule) + and the BWD kernel does NOT have ``x = pre_logits_pre_scale`` + directly available (logits is a function input passed in + fp32). We therefore accept the per-row partials being + a function of ``logits`` and load logits inside the BWD + kernel. See the LOGITS_PTR arg added below. + + For simplicity and correctness, the wrapper passes the *fp32 + logits* tensor into BWD and the kernel recomputes the pre-scale + cross-term inline. At V4-Flash widths the extra HBM read is + negligible (24 fp32 / row vs the K*K state already loaded). + """ + + pid = tl.program_id(0) + offs = pid * BLOCK_LEADING + tl.arange(0, BLOCK_LEADING) + mask_leading = offs < N + + k_idx = tl.arange(0, K) + + # Read scale (broadcast scalars). + scale0 = tl.load(SCALE_PTR + 0) + scale1 = tl.load(SCALE_PTR + 1) + scale2 = tl.load(SCALE_PTR + 2) + + # Saved forward states. + pre_sig = tl.load( + PRE_SIG_PTR + offs[:, None] * K + k_idx[None, :], + mask=mask_leading[:, None], + other=0.0, + ) + post_sig = tl.load( + POST_SIG_PTR + offs[:, None] * K + k_idx[None, :], + mask=mask_leading[:, None], + other=0.0, + ) + comb_sm = tl.load( + COMB_SM_PTR + offs[:, None, None] * K * K + k_idx[None, :, None] * K + k_idx[None, None, :], + mask=mask_leading[:, None, None], + other=0.0, + ) + + # Upstream grads (cast to fp32 for the chain). + d_pre = tl.load( + DPRE_PTR + offs[:, None] * K + k_idx[None, :], + mask=mask_leading[:, None], + other=0.0, + ).to(tl.float32) + d_post = tl.load( + DPOST_PTR + offs[:, None] * K + k_idx[None, :], + mask=mask_leading[:, None], + other=0.0, + ).to(tl.float32) + d_comb = tl.load( + DCOMB_PTR + offs[:, None, None] * K * K + k_idx[None, :, None] * K + k_idx[None, None, :], + mask=mask_leading[:, None, None], + other=0.0, + ).to(tl.float32) + + # Sigmoid VJP (no eps contribution -- + EPS has derivative 1): + d_pre_logit = d_pre * pre_sig * (1.0 - pre_sig) + d_post_logit = d_post * 2.0 * post_sig * (1.0 - post_sig) + + # Softmax VJP: d_comb_logit = comb_sm * (d_comb - sum(d_comb * comb_sm, axis=2)) + dot = tl.sum(d_comb * comb_sm, axis=2, keep_dims=True) + d_comb_logit = comb_sm * (d_comb - dot) + + KK_TOTAL: tl.constexpr = (2 + K) * K + + # Write d_logits (= scale * d_*_logit at each slice). + tl.store( + DLOGITS_PTR + offs[:, None] * KK_TOTAL + k_idx[None, :], + d_pre_logit * scale0, + mask=mask_leading[:, None], + ) + tl.store( + DLOGITS_PTR + offs[:, None] * KK_TOTAL + K + k_idx[None, :], + d_post_logit * scale1, + mask=mask_leading[:, None], + ) + tl.store( + DLOGITS_PTR + + offs[:, None, None] * KK_TOTAL + + 2 * K + + k_idx[None, :, None] * K + + k_idx[None, None, :], + d_comb_logit * scale2, + mask=mask_leading[:, None, None], + ) + + # d_base accumulates the d_*_logit values directly (no scale factor). + # We write per-row partials and let host-side torch.sum reduce them + # (avoids cross-block atomic_add). + tl.store( + DBASE_PTR + offs[:, None] * KK_TOTAL + k_idx[None, :], + d_pre_logit, + mask=mask_leading[:, None], + ) + tl.store( + DBASE_PTR + offs[:, None] * KK_TOTAL + K + k_idx[None, :], + d_post_logit, + mask=mask_leading[:, None], + ) + tl.store( + DBASE_PTR + offs[:, None, None] * KK_TOTAL + 2 * K + k_idx[None, :, None] * K + k_idx[None, None, :], + d_comb_logit, + mask=mask_leading[:, None, None], + ) + + # d_scale needs ``logits * d_*_logit`` per-element, summed over + # each slice's K (or K*K) inner extent and then over rows. We + # write the per-row d_*_logit values into ``DBASE_PTR`` above + # (which is exactly ``d_*_logit`` with no scale factor); the + # host-side wrapper then computes + # d_scale[0] = (logits[:, :K] * d_base_partials[:, :K] ).sum() + # d_scale[1] = (logits[:, K:2K] * d_base_partials[:, K:2K]).sum() + # d_scale[2] = (logits[:, 2K:] * d_base_partials[:, 2K:] ).sum() + # with torch.sum (avoiding cross-block atomic_add). ``DSCALE_PTR`` + # is therefore left as the zero-init buffer the wrapper allocated; + # we keep it in the kernel signature for forward-compatibility. + + +# --------------------------------------------------------------------------- +# Block-leading heuristic +# --------------------------------------------------------------------------- + + +def _pick_block_leading(n: int, k: int) -> int: + """Pick ``BLOCK_LEADING`` based on K and the work-axis size. + + At K=4 the in-register state per row is ~``24 + K + K*K = 44`` fp32 + elements; ``BLOCK_LEADING=64`` fits comfortably (~11 KiB). At K=16 + the per-row state grows ~6x; drop to 8. + """ + + if k <= 4: + cap = 64 + elif k <= 8: + cap = 32 + else: + cap = 8 + + if n < cap: + return max(1, triton.next_power_of_2(n)) + return cap + + +# --------------------------------------------------------------------------- +# torch.autograd.Function wrapper +# --------------------------------------------------------------------------- + + +class HCComputeTailFn(torch.autograd.Function): + """Autograd-aware wrapper around the FWD/BWD Triton kernels. + + Saves ``(logits, scale, pre_sig, post_sig, comb_sm)`` for backward. + Returns ``(pre, post, comb_pre_sinkhorn)`` -- the caller then runs + ``sinkhorn_normalize`` on ``comb_pre_sinkhorn``. + + Shape: ``logits [..., (2+K)*K]`` fp32, ``scale [3]`` fp32, + ``base [(2+K)*K]`` fp32, ``K`` a power of 2 in ``{1, 2, 4, 8, 16}``. + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, + logits: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + K: int, + eps: float, + out_dtype: torch.dtype, + ): + if K not in _SUPPORTED_K: + raise ValueError(f"hc_glue Triton path: unsupported K={K}; expected one of {_SUPPORTED_K}") + if logits.shape[-1] != (2 + K) * K: + raise ValueError( + f"hc_glue: logits last-dim must be (2+K)*K = {(2 + K) * K}, " f"got {logits.shape[-1]}" + ) + if scale.numel() != 3: + raise ValueError(f"hc_glue: scale must have 3 elements, got {scale.numel()}") + if base.numel() != (2 + K) * K: + raise ValueError( + f"hc_glue: base must have (2+K)*K = {(2 + K) * K} elements, " f"got {base.numel()}" + ) + + logits_c = logits.contiguous().to(torch.float32) + scale_c = scale.contiguous().to(torch.float32) + base_c = base.contiguous().to(torch.float32) + + leading_shape = logits_c.shape[:-1] + N = 1 + for s in leading_shape: + N *= s + + device = logits_c.device + pre = torch.empty((*leading_shape, K), dtype=out_dtype, device=device) + post = torch.empty((*leading_shape, K), dtype=out_dtype, device=device) + comb = torch.empty((*leading_shape, K, K), dtype=out_dtype, device=device) + + # Saved-for-backward fp32 states. + pre_sig = torch.empty((N, K), dtype=torch.float32, device=device) + post_sig = torch.empty((N, K), dtype=torch.float32, device=device) + comb_sm = torch.empty((N, K, K), dtype=torch.float32, device=device) + + block_leading = _pick_block_leading(N, K) + grid = (triton.cdiv(N, block_leading),) + _hc_compute_tail_fwd_kernel[grid]( + logits_c, + scale_c, + base_c, + pre, + post, + comb, + pre_sig, + post_sig, + comb_sm, + N, + EPS=float(eps), + K=K, + BLOCK_LEADING=block_leading, + OUT_DTYPE=_triton_dtype(out_dtype), + ) + + ctx.save_for_backward(logits_c, scale_c, pre_sig, post_sig, comb_sm) + ctx.K = K + ctx.eps = float(eps) + ctx.leading_shape = tuple(leading_shape) + ctx.out_dtype = out_dtype + return pre, post, comb + + @staticmethod + def backward(ctx, d_pre, d_post, d_comb): # type: ignore[override] + logits_c, scale_c, pre_sig, post_sig, comb_sm = ctx.saved_tensors + K: int = ctx.K + leading_shape = ctx.leading_shape + + N = 1 + for s in leading_shape: + N *= s + + d_pre = d_pre.contiguous() + d_post = d_post.contiguous() + d_comb = d_comb.contiguous() + + device = logits_c.device + d_logits = torch.empty_like(logits_c) + d_base_partials = torch.empty((N, (2 + K) * K), dtype=torch.float32, device=device) + # d_scale_partials is no longer written by the kernel; we leave + # the buffer hint in the kernel signature for forward-compat. + d_scale_partials = torch.zeros((N, 3), dtype=torch.float32, device=device) + + block_leading = _pick_block_leading(N, K) + grid = (triton.cdiv(N, block_leading),) + _hc_compute_tail_bwd_kernel[grid]( + d_pre, + d_post, + d_comb, + pre_sig, + post_sig, + comb_sm, + scale_c, + d_logits, + d_scale_partials, + d_base_partials, + N, + K=K, + BLOCK_LEADING=block_leading, + ) + + # Host-side reductions: + # d_base[i] = sum_n d_base_partials[n, i] (3*K + K*K entries) + # d_scale[0] = sum_{n, k} logits_slice_pre * d_base_partials_pre + # d_scale[1] = sum_{n, k} logits_slice_post * d_base_partials_post + # d_scale[2] = sum_{n, k1, k2} logits_slice_comb * d_base_partials_comb + # + # Equivalent to: d_scale[i] = (logits_slice_i * d_base_partials_i).sum() + # since d_base = d_*_logit (no scale) and d_scale[i] = sum logits * d_*_logit. + logits_flat = logits_c.reshape(N, -1) + d_base = d_base_partials.sum(dim=0) + d_scale_0 = (logits_flat[:, :K] * d_base_partials[:, :K]).sum() + d_scale_1 = (logits_flat[:, K : 2 * K] * d_base_partials[:, K : 2 * K]).sum() + d_scale_2 = (logits_flat[:, 2 * K :] * d_base_partials[:, 2 * K :]).sum() + d_scale = torch.stack([d_scale_0, d_scale_1, d_scale_2]) + + # Reshape d_logits back to the leading shape that the FWD input + # had so the autograd machinery passes it back to the caller's + # F.linear backward correctly. + d_logits = d_logits.view(*leading_shape, (2 + K) * K) + + # K, eps, out_dtype are non-differentiable. + return d_logits, d_scale, d_base, None, None, None + + +# --------------------------------------------------------------------------- +# Public Python entry points +# --------------------------------------------------------------------------- + + +def is_triton_path_enabled() -> bool: + """Return True iff the ``PRIMUS_HC_TRITON`` env knob is not ``"0"``. + + Default-on; A/B toggle via ``PRIMUS_HC_TRITON=0``. + """ + + return os.environ.get("PRIMUS_HC_TRITON", "1") != "0" + + +def is_triton_kernel_supported(logits: torch.Tensor, K: int) -> bool: + """Return True iff the input shape / device is supported. + + Used by the dispatcher in + :meth:`primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.compute_weights` + to safely fall back to the eager body for unsupported configurations. + """ + + if not logits.is_cuda: + return False + if K not in _SUPPORTED_K: + return False + if logits.shape[-1] != (2 + K) * K: + return False + return True + + +def hc_glue_compute_tail_triton( + logits: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + *, + K: int, + eps: float, + out_dtype: torch.dtype, +): + """Run the Triton-fused HC compute_weights tail. + + Returns ``(pre, post, comb_pre_sinkhorn)`` -- the caller runs + ``sinkhorn_normalize`` on ``comb_pre_sinkhorn``. + """ + + return HCComputeTailFn.apply(logits, scale, base, K, eps, out_dtype) + + +__all__ = [ + "HCComputeTailFn", + "hc_glue_compute_tail_triton", + "is_triton_path_enabled", + "is_triton_kernel_supported", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py new file mode 100644 index 000000000..c82543d75 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score.py @@ -0,0 +1,503 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton-fused Indexer scoring (plan-6 P38). + +Fuses the ``einsum + relu + mul + sum + causal_mask`` chain in +:meth:`primus.backends.megatron.core.transformer.indexer.Indexer.forward` +into a single FWD + single BWD Triton kernel pair. Eager body: + +.. code-block:: python + + relu_term = F.relu(torch.einsum("bshd,bpd->bshp", q_i, k_icomp)) + scores = (relu_term * w_i.unsqueeze(-1)).sum(dim=2) # [B, S, P] + mask = self._causal_mask(S, P, scores.device, scores.dtype) + scores = scores + mask.unsqueeze(0) + +P38 collapses these ~7 ATen kernels (einsum + relu + mul + sum + +mask alloc + mask add + dtype cast) into one Triton kernel that: + +* Computes the per-(q, p) dot product over the head-feature dim ``Hd`` + inline (no `[B, S, H, P]` intermediate tensor materialised); +* Applies ``relu`` per element; +* Multiplies by ``w_i`` per head; +* Reduces over heads; +* Materialises the causal mask **inline** via + ``tl.where((p + 1) * compress_ratio - 1 <= s, 0.0, -INF)`` — no + ``[S, P]`` mask tensor, no HBM traffic; +* Writes ``scores [B, S, P]`` (the `topk` and trailing tail stay + host-side; `topk` is heavy GPU compute on its own and benefits + from being its own kernel). + +The BWD kernel recomputes the per-tile `relu` mask in the BWD pass +(FlashAttention-style trick) instead of saving it — saves ``H * S * P`` +bits of HBM per CSA layer. + +Gating: ``PRIMUS_INDEXER_TRITON == "1"`` (**default-OFF**). + +Default-off rationale (P38 descope per `plan-6/02-phase-details.md` +§"Task list refinement"): at V4-Flash widths (B=1, S=4096, P=1024, +H=8, Hd=128) the eager `einsum` maps to a cuBLAS / hipBLASLt +batched-matmul that already runs at ~28 TFLOP/s on MI355. The +generic Triton kernel here is FWD-competitive only at small shapes +(3.35x FWD speedup at B=2, S=128, P=32) but regresses ~30% at the +production V4-Flash shape; BWD regresses ~12x because of cross-tile +``atomic_add`` traffic on `dq` / `dk` / `dw`. The kernel stays +available for future tuning + small-shape paths via the env knob. + +Per-shape behavior summary: + +* V4-Flash production (B=1, S=4096, P=1024, H=8, Hd=128, bf16): + FWD 0.424 ms (triton) vs 0.306 ms (eager) -> **0.72x** (regression). + BWD 6.457 ms (triton) vs 0.489 ms (eager) -> **0.08x** (regression). +* Small (B=2, S=128, P=32, H=8, Hd=128, bf16): + FWD 0.053 ms (triton) vs 0.176 ms (eager) -> **3.35x** (speedup). + BWD 0.226 ms (triton) vs 0.256 ms (eager) -> **1.14x** (speedup). +""" + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +# Indexer scoring runs at V4-Flash widths [B=1, S=4096, P=1024, H=8, +# Hd=128]. H is small enough to be a compile-time constant in the +# kernel; supported values are documented here. +_SUPPORTED_H = (1, 2, 4, 8, 16) + + +# --------------------------------------------------------------------------- +# Triton kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _indexer_score_fwd_kernel( + Q_PTR, # [B, S, H, Hd] - q_i + K_PTR, # [B, P, Hd] - k_icomp.squeeze(2) + W_PTR, # [B, S, H] - w_i + SCORES_PTR, # [B, S, P] - out (fp32 internal, cast OUT_DTYPE) + B, + S, + P, + H: tl.constexpr, + HD: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + BLOCK_S: tl.constexpr, + BLOCK_P: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + """One program tile = one ``[B_b, BLOCK_S, BLOCK_P]`` chunk. + + Loop axis: ``Hd`` (head-feature dim) and ``H`` (heads), both + compile-time known. Per tile the kernel: + + 1. For each head ``h`` (unrolled, since H is constexpr): + a. Load ``q_i[b_b, s_tile, h, :]`` -- shape ``[BLOCK_S, HD]``. + b. Load ``k_icomp[b_b, p_tile, :]`` -- shape ``[BLOCK_P, HD]``. + c. Compute ``dot = q @ k.T`` -- shape ``[BLOCK_S, BLOCK_P]``. + d. Apply ``relu(dot)``. + e. Load ``w_i[b_b, s_tile, h]`` -- shape ``[BLOCK_S]``. + f. Multiply: ``acc += relu * w[:, None]``. + 2. Materialise causal mask inline: positions ``s`` may attend to + pool position ``p`` iff ``(p + 1) * compress_ratio - 1 <= s``; + write ``-inf`` otherwise. + 3. Store ``acc`` cast to OUT_DTYPE. + + Register footprint at V4-Flash (BLOCK_S=64, BLOCK_P=64, H=8, HD=128): + per head: 64*128 + 64*128 = 16384 fp32 = 64 KiB for (q, k). + per tile: 64*64 = 4096 fp32 = 16 KiB for acc. + Cumulative per-program peak ~ 80 KiB ≈ 320 VGPRs / warp. At MI355 + 256 VGPRs we drop BLOCK_S to 32 for safety on the small kernel. + """ + + pid_b = tl.program_id(0) + pid_s = tl.program_id(1) + pid_p = tl.program_id(2) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + p_offs = pid_p * BLOCK_P + tl.arange(0, BLOCK_P) + s_mask = s_offs < S + p_mask = p_offs < P + + hd_idx = tl.arange(0, HD) + + acc = tl.zeros((BLOCK_S, BLOCK_P), dtype=tl.float32) + + # Unroll over heads (H is small and constexpr). + for h in tl.static_range(0, H): + # q [BLOCK_S, HD]: q_i[pid_b, s_offs, h, :] + q_tile = tl.load( + Q_PTR + pid_b * S * H * HD + s_offs[:, None] * H * HD + h * HD + hd_idx[None, :], + mask=s_mask[:, None], + other=0.0, + ).to(tl.float32) + # k [BLOCK_P, HD]: k_icomp[pid_b, p_offs, :] + k_tile = tl.load( + K_PTR + pid_b * P * HD + p_offs[:, None] * HD + hd_idx[None, :], + mask=p_mask[:, None], + other=0.0, + ).to(tl.float32) + # dot [BLOCK_S, BLOCK_P] = q @ k.T + dot = tl.dot(q_tile, tl.trans(k_tile), out_dtype=tl.float32) + dot = tl.maximum(dot, 0.0) # relu + # w [BLOCK_S]: w_i[pid_b, s_offs, h] + w_h = tl.load( + W_PTR + pid_b * S * H + s_offs * H + h, + mask=s_mask, + other=0.0, + ).to(tl.float32) + acc += dot * w_h[:, None] + + # Apply causal mask inline. Allowed iff `(p + 1) * cr - 1 <= s`, + # i.e. the pool position's window end is no later than the query. + s_arr = s_offs[:, None] + p_arr = p_offs[None, :] + allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr + NEG_INF = -float("inf") + acc = tl.where(allowed, acc, NEG_INF) + + tl.store( + SCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :], + acc.to(OUT_DTYPE), + mask=s_mask[:, None] & p_mask[None, :], + ) + + +@triton.jit +def _indexer_score_bwd_kernel( + DSCORES_PTR, # [B, S, P] grad in OUT_DTYPE + Q_PTR, # [B, S, H, Hd] + K_PTR, # [B, P, Hd] + W_PTR, # [B, S, H] + DQ_PTR, # [B, S, H, Hd] OUT (fp32) + DK_PTR, # [B, P, Hd] OUT (fp32, scattered) + DW_PTR, # [B, S, H] OUT (fp32) + B, + S, + P, + H: tl.constexpr, + HD: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + BLOCK_S: tl.constexpr, + BLOCK_P: tl.constexpr, +): + """VJP through the scoring chain. + + The mask is `where(allowed, acc, -inf)`; its derivative is 1 in + allowed positions and 0 elsewhere. Out-of-range positions therefore + contribute zero grad and we set ``dscores`` to 0 there before + walking back. + + Per element forward: + acc = sum_h relu(q . k) * w_h + masked = where(allowed, acc, -inf) + Per element backward: + d_acc = d_masked * where(allowed, 1, 0) # equiv. d_masked masked by allowed + # d_relu_dot[h] = d_acc * w_h + # d_w[h] = d_acc * relu_dot[h] + # relu' = (dot > 0). Recompute dot, relu mask. + # d_dot[h] = d_relu_dot[h] * (dot > 0) + # d_q[h, :] = d_dot[h, p] @ k[p, :] + # d_k[p, :] = d_dot[s, p] @ q[s, h, :] + + For the BWD we use one program per (b, s_tile, p_tile) just like + FWD; ``d_q`` is accumulated locally (one chunk per s_tile so + multiple p_tiles need atomic add), ``d_k`` is scattered across + s_tiles so needs atomic add, ``d_w`` is accumulated similarly. + """ + + pid_b = tl.program_id(0) + pid_s = tl.program_id(1) + pid_p = tl.program_id(2) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + p_offs = pid_p * BLOCK_P + tl.arange(0, BLOCK_P) + s_mask = s_offs < S + p_mask = p_offs < P + + hd_idx = tl.arange(0, HD) + + # Load dscores [BLOCK_S, BLOCK_P] + dmasked = tl.load( + DSCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :], + mask=s_mask[:, None] & p_mask[None, :], + other=0.0, + ).to(tl.float32) + # Apply causal mask (zero out invalid positions). + s_arr = s_offs[:, None] + p_arr = p_offs[None, :] + allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr + d_acc = tl.where(allowed, dmasked, 0.0) + + # Unroll over heads. + for h in tl.static_range(0, H): + # Reload q, k for this head (FlashAttention-style recompute). + q_tile = tl.load( + Q_PTR + pid_b * S * H * HD + s_offs[:, None] * H * HD + h * HD + hd_idx[None, :], + mask=s_mask[:, None], + other=0.0, + ).to(tl.float32) + k_tile = tl.load( + K_PTR + pid_b * P * HD + p_offs[:, None] * HD + hd_idx[None, :], + mask=p_mask[:, None], + other=0.0, + ).to(tl.float32) + w_h = tl.load( + W_PTR + pid_b * S * H + s_offs * H + h, + mask=s_mask, + other=0.0, + ).to(tl.float32) + + # Recompute dot, relu, relu_dot. + dot = tl.dot(q_tile, tl.trans(k_tile), out_dtype=tl.float32) + relu_dot = tl.maximum(dot, 0.0) # [BLOCK_S, BLOCK_P] + relu_mask = dot > 0.0 + + # d_w[h] = sum_p d_acc * relu_dot + dw_h = tl.sum(d_acc * relu_dot, axis=1) # [BLOCK_S] + + # d_relu_dot = d_acc * w_h[:, None] + d_relu_dot = d_acc * w_h[:, None] + # d_dot = d_relu_dot where relu_mask else 0 + d_dot = tl.where(relu_mask, d_relu_dot, 0.0) + + # d_q[s, h, :] = sum_p d_dot[s, p] * k[p, :] shape [BLOCK_S, HD] + d_q = tl.dot(d_dot, k_tile, out_dtype=tl.float32) + # d_k[p, :] = sum_s d_dot[s, p] * q[s, :] shape [BLOCK_P, HD] + d_k = tl.dot(tl.trans(d_dot), q_tile, out_dtype=tl.float32) + + # Stores with atomic_add since multiple p_tiles/s_tiles touch + # the same locations. + tl.atomic_add( + DQ_PTR + pid_b * S * H * HD + s_offs[:, None] * H * HD + h * HD + hd_idx[None, :], + d_q, + mask=s_mask[:, None], + ) + tl.atomic_add( + DK_PTR + pid_b * P * HD + p_offs[:, None] * HD + hd_idx[None, :], + d_k, + mask=p_mask[:, None], + ) + tl.atomic_add( + DW_PTR + pid_b * S * H + s_offs * H + h, + dw_h, + mask=s_mask, + ) + + +# --------------------------------------------------------------------------- +# Heuristics +# --------------------------------------------------------------------------- + + +def _pick_block(s: int, p: int, h: int, hd: int) -> tuple[int, int]: + """Pick BLOCK_S, BLOCK_P that fit MI355 register budget.""" + block_s = 32 + block_p = 32 + if hd >= 128: + block_s = 32 + block_p = 32 + if hd >= 256: + block_s = 16 + block_p = 32 + # Floor both tiles to 16: BLOCK_S / BLOCK_P are the M / N dims of the + # ``tl.dot(q_tile, k_tile.T)`` in the FWD/BWD score kernels, so they + # must be >= the gfx1250 WMMA minimum. For tiny S / P (e.g. unit-test + # shapes), ``next_power_of_2(s|p)`` can drop below 16 and the dot fails + # to select a matrix-core intrinsic ("no matching matrix core intrinsic + # for wmma version 3 ... [0, 0, K]"). Surplus rows/cols are masked by + # ``s_mask`` / ``p_mask`` inside the kernels, so a 16-wide tile over a + # smaller S / P is safe. + block_s = max(16, min(block_s, triton.next_power_of_2(max(1, s)))) + block_p = max(16, min(block_p, triton.next_power_of_2(max(1, p)))) + return block_s, block_p + + +# --------------------------------------------------------------------------- +# autograd.Function wrapper +# --------------------------------------------------------------------------- + + +class IndexerScoreFn(torch.autograd.Function): + """Autograd-aware wrapper around the FWD/BWD Triton kernels. + + Returns ``scores [B, S, P]`` in ``out_dtype``. Caller runs + ``torch.topk`` / sentinel substitution / padding host-side. + + Inputs: + q_i [B, S, H, Hd] - per-head queries (any float dtype). + k_icomp [B, P, Hd] - compressed keys (any float dtype). + w_i [B, S, H] - per-head weights (any float dtype). + compress_ratio - int (typically 4 for CSA). + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, + q_i: torch.Tensor, + k_icomp: torch.Tensor, + w_i: torch.Tensor, + compress_ratio: int, + out_dtype: torch.dtype, + ): + if q_i.dim() != 4: + raise ValueError(f"q_i must be [B, S, H, Hd], got shape {tuple(q_i.shape)}") + if k_icomp.dim() != 3: + raise ValueError(f"k_icomp must be [B, P, Hd], got shape {tuple(k_icomp.shape)}") + if w_i.dim() != 3: + raise ValueError(f"w_i must be [B, S, H], got shape {tuple(w_i.shape)}") + + B, S, H, HD = q_i.shape + Bk, P, HDk = k_icomp.shape + Bw, Sw, Hw = w_i.shape + if B != Bk or B != Bw: + raise ValueError(f"Mismatched B: q={B} k={Bk} w={Bw}") + if HD != HDk: + raise ValueError(f"Mismatched Hd: q={HD} k={HDk}") + if S != Sw: + raise ValueError(f"Mismatched S: q={S} w={Sw}") + if H != Hw: + raise ValueError(f"Mismatched H: q={H} w={Hw}") + if H not in _SUPPORTED_H: + raise ValueError(f"Unsupported H={H}; expected one of {_SUPPORTED_H}") + if HD & (HD - 1) != 0: + raise ValueError(f"Hd must be a power of 2, got {HD}") + + q_c = q_i.contiguous() + k_c = k_icomp.contiguous() + w_c = w_i.contiguous() + + device = q_c.device + scores = torch.empty((B, S, P), dtype=out_dtype, device=device) + + block_s, block_p = _pick_block(S, P, H, HD) + grid = (B, triton.cdiv(S, block_s), triton.cdiv(P, block_p)) + _indexer_score_fwd_kernel[grid]( + q_c, + k_c, + w_c, + scores, + B, + S, + P, + H=H, + HD=HD, + COMPRESS_RATIO=int(compress_ratio), + BLOCK_S=block_s, + BLOCK_P=block_p, + OUT_DTYPE={ + torch.float32: tl.float32, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, + torch.float64: tl.float64, + }[out_dtype], + ) + + ctx.save_for_backward(q_c, k_c, w_c) + ctx.compress_ratio = int(compress_ratio) + ctx.shape = (B, S, P, H, HD) + ctx.in_dtypes = (q_i.dtype, k_icomp.dtype, w_i.dtype) + return scores + + @staticmethod + def backward(ctx, d_scores): # type: ignore[override] + q_c, k_c, w_c = ctx.saved_tensors + B, S, P, H, HD = ctx.shape + compress_ratio = ctx.compress_ratio + q_dtype, k_dtype, w_dtype = ctx.in_dtypes + + d_scores = d_scores.contiguous() + device = q_c.device + + d_q_fp32 = torch.zeros((B, S, H, HD), dtype=torch.float32, device=device) + d_k_fp32 = torch.zeros((B, P, HD), dtype=torch.float32, device=device) + d_w_fp32 = torch.zeros((B, S, H), dtype=torch.float32, device=device) + + block_s, block_p = _pick_block(S, P, H, HD) + grid = (B, triton.cdiv(S, block_s), triton.cdiv(P, block_p)) + _indexer_score_bwd_kernel[grid]( + d_scores, + q_c, + k_c, + w_c, + d_q_fp32, + d_k_fp32, + d_w_fp32, + B, + S, + P, + H=H, + HD=HD, + COMPRESS_RATIO=int(compress_ratio), + BLOCK_S=block_s, + BLOCK_P=block_p, + ) + + return d_q_fp32.to(q_dtype), d_k_fp32.to(k_dtype), d_w_fp32.to(w_dtype), None, None + + +# --------------------------------------------------------------------------- +# Public entry points +# --------------------------------------------------------------------------- + + +def is_triton_path_enabled() -> bool: + """Return True iff ``PRIMUS_INDEXER_TRITON_FULL == "1"``. + + **Re-purposed at P41:** the P38 full-fuse path is now gated by the + distinct env knob ``PRIMUS_INDEXER_TRITON_FULL`` (default ``"0"``). + The original ``PRIMUS_INDEXER_TRITON`` env now controls the cheaper + P41 post-einsum tail fusion in + :mod:`indexer_score_post`. + + The full-fuse path is descoped at V4-Flash widths (cuBLAS einsum + beats the generic Triton kernel + BWD atomic_add contention + regresses 12x). Kept available for small-shape paths and future + tuning. Set ``PRIMUS_INDEXER_TRITON_FULL=1`` to opt-in. + """ + return os.environ.get("PRIMUS_INDEXER_TRITON_FULL", "0") == "1" + + +def is_triton_kernel_supported(q_i: torch.Tensor, k_icomp: torch.Tensor, w_i: torch.Tensor) -> bool: + """Return True iff the input shapes / device support the Triton path.""" + if not q_i.is_cuda or not k_icomp.is_cuda or not w_i.is_cuda: + return False + if q_i.dim() != 4 or k_icomp.dim() != 3 or w_i.dim() != 3: + return False + B, S, H, HD = q_i.shape + if H not in _SUPPORTED_H: + return False + if HD & (HD - 1) != 0: + return False + return True + + +def indexer_score_triton( + q_i: torch.Tensor, + k_icomp: torch.Tensor, + w_i: torch.Tensor, + *, + compress_ratio: int, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Compute Indexer scores via the fused Triton kernel. + + Returns ``scores [B, S, P]`` of dtype ``out_dtype``. Masked + positions hold ``-inf``. + """ + return IndexerScoreFn.apply(q_i, k_icomp, w_i, compress_ratio, out_dtype) + + +__all__ = [ + "IndexerScoreFn", + "indexer_score_triton", + "is_triton_path_enabled", + "is_triton_kernel_supported", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py new file mode 100644 index 000000000..f1f4c61e3 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_score_post.py @@ -0,0 +1,468 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton-fused Indexer scoring -- post-einsum tail (plan-6 P41). + +Companion to :mod:`indexer_score` (P38). P38 fused the entire +``einsum + relu + mul + sum + causal_mask`` chain into one kernel and +lost to cuBLAS / hipBLASLt on the matmul half (~28 TFLOP/s eager vs +~20 TFLOP/s Triton at V4-Flash widths). P41 keeps the einsum eager +and fuses **only** the post-matmul tail +(``relu -> mul(w_i) -> sum(H) -> + causal_mask``). The matmul is +compute-bound and stays on cuBLAS peak; the tail is bandwidth-bound +and Triton wins because every elementwise op costs a full HBM +round-trip. + +Eager body (with einsum kept): + +.. code-block:: python + + dot = torch.einsum("bshd,bpd->bshp", q_i, k_icomp) # stays eager + relu_term = F.relu(dot) + scores = (relu_term * w_i.unsqueeze(-1)).sum(dim=2) # [B, S, P] + mask = self._causal_mask(S, P, scores.device, scores.dtype) + scores = scores + mask.unsqueeze(0) + +P41 collapses the tail (``relu + mul + sum + mask_alloc + mask_add + +dtype cast``, ~5 ATen launches) into one Triton kernel that: + +* Reads ``dot [B, S, H, P]`` once; +* Applies ``relu`` per element; +* Multiplies by ``w_i[B, S, H, 1]`` (broadcast over P); +* Reduces over heads ``H``; +* Materialises the causal mask **inline** via + ``tl.where((p + 1) * compress_ratio - 1 <= s, acc, -inf)``; +* Writes ``scores [B, S, P]`` cast to ``OUT_DTYPE``. + +BWD takes ``d_scores [B, S, P]`` + saved ``dot [B, S, H, P]`` + +``w_i [B, S, H]``, emits: + +* ``d_dot[b, s, h, p] = d_scores[b, s, p] * w_i[b, s, h]`` where + ``dot[b, s, h, p] > 0`` else 0 (one HBM write per element); +* ``d_w_i[b, s, h] = sum_p(d_scores[b, s, p] * relu(dot[b, s, h, p]))`` + (one reduction per (b, s, h) — no cross-block atomic_add). + +The BWD has no atomic_add traffic (the P38 BWD's killer). Both FWD +and BWD are bandwidth-bound and should win at V4-Flash widths. + +Gating: ``PRIMUS_INDEXER_TRITON == "1"`` (the re-purposed env knob; +default initially ``"0"`` then ``"1"`` after the proxy A/B confirms). +The legacy P38 full-fuse path lives behind ``PRIMUS_INDEXER_TRITON_FULL`` +(see :mod:`indexer_score`). +""" + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +# H is small and known per call site -- V4-Flash uses H=8. Same +# supported set as P38. +_SUPPORTED_H = (1, 2, 4, 8, 16) + + +# --------------------------------------------------------------------------- +# Triton kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _indexer_score_post_fwd_kernel( + DOT_PTR, # [B, S, H, P] - eager einsum output, pre-relu + W_PTR, # [B, S, H] - per-head weights + SCORES_PTR, # [B, S, P] - out + B, + S, + P, + H: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + BLOCK_S: tl.constexpr, + BLOCK_P: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + """One program tile = ``[B_b, BLOCK_S, BLOCK_P]``. + + Per-tile workflow: + + 1. Initialise ``acc [BLOCK_S, BLOCK_P]`` fp32 to 0. + 2. For each head ``h`` (constexpr unroll): + a. Load ``dot[b, s_tile, h, p_tile]`` -> [BLOCK_S, BLOCK_P]. + b. Apply ``relu``. + c. Load ``w[b, s_tile, h]`` -> [BLOCK_S]. + d. ``acc += relu * w[:, None]``. + 3. Apply causal mask inline. + 4. Store ``acc`` cast to ``OUT_DTYPE``. + """ + + pid_b = tl.program_id(0) + pid_s = tl.program_id(1) + pid_p = tl.program_id(2) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + p_offs = pid_p * BLOCK_P + tl.arange(0, BLOCK_P) + s_mask = s_offs < S + p_mask = p_offs < P + + acc = tl.zeros((BLOCK_S, BLOCK_P), dtype=tl.float32) + + for h in tl.static_range(0, H): + # dot[pid_b, s_offs, h, p_offs] -> [BLOCK_S, BLOCK_P] + dot_tile = tl.load( + DOT_PTR + pid_b * S * H * P + s_offs[:, None] * H * P + h * P + p_offs[None, :], + mask=s_mask[:, None] & p_mask[None, :], + other=0.0, + ).to(tl.float32) + relu_tile = tl.maximum(dot_tile, 0.0) + + # w[pid_b, s_offs, h] -> [BLOCK_S] + w_h = tl.load( + W_PTR + pid_b * S * H + s_offs * H + h, + mask=s_mask, + other=0.0, + ).to(tl.float32) + + acc += relu_tile * w_h[:, None] + + # Causal mask inline. + s_arr = s_offs[:, None] + p_arr = p_offs[None, :] + allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr + NEG_INF = -float("inf") + acc = tl.where(allowed, acc, NEG_INF) + + tl.store( + SCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :], + acc.to(OUT_DTYPE), + mask=s_mask[:, None] & p_mask[None, :], + ) + + +@triton.jit +def _indexer_score_post_bwd_dw_kernel( + DSCORES_PTR, # [B, S, P] - grad in (any float dtype) + DOT_PTR, # [B, S, H, P] - saved pre-relu activations + DW_PTR, # [B, S, H] - OUT (fp32) + B, + S, + P, + H: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + BLOCK_S: tl.constexpr, + BLOCK_P_INNER: tl.constexpr, +): + """Per ``(b, s_tile, h)`` program: ``d_w[h] = sum_p(d_acc * relu(dot[h]))``. + + Loops over P internally in chunks of ``BLOCK_P_INNER``; ``d_w`` + is fully local to each program (no cross-block ``atomic_add``). + """ + + pid_b = tl.program_id(0) + pid_s = tl.program_id(1) + h = tl.program_id(2) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + s_mask = s_offs < S + s_arr = s_offs[:, None] + + acc = tl.zeros((BLOCK_S,), dtype=tl.float32) + + n_p_chunks = tl.cdiv(P, BLOCK_P_INNER) + for chunk in range(n_p_chunks): + p_offs = chunk * BLOCK_P_INNER + tl.arange(0, BLOCK_P_INNER) + p_mask = p_offs < P + p_arr = p_offs[None, :] + + dmasked = tl.load( + DSCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :], + mask=s_mask[:, None] & p_mask[None, :], + other=0.0, + ).to(tl.float32) + allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr + d_acc = tl.where(allowed, dmasked, 0.0) + + dot_tile = tl.load( + DOT_PTR + pid_b * S * H * P + s_offs[:, None] * H * P + h * P + p_offs[None, :], + mask=s_mask[:, None] & p_mask[None, :], + other=0.0, + ).to(tl.float32) + relu_tile = tl.maximum(dot_tile, 0.0) + + acc += tl.sum(d_acc * relu_tile, axis=1) + + tl.store( + DW_PTR + pid_b * S * H + s_offs * H + h, + acc, + mask=s_mask, + ) + + +@triton.jit +def _indexer_score_post_bwd_ddot_kernel( + DSCORES_PTR, # [B, S, P] - grad in (any float dtype) + DOT_PTR, # [B, S, H, P] - saved pre-relu activations + W_PTR, # [B, S, H] - per-head weights + DDOT_PTR, # [B, S, H, P] - OUT (fp32) + B, + S, + P, + H: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + BLOCK_S: tl.constexpr, + BLOCK_P: tl.constexpr, +): + """VJP through the post-einsum tail — ``d_dot`` only. + + Per-element forward (with mask): + d_acc = where(allowed, d_scores, 0) + Per-element backward: + d_dot[h] = where(dot[h] > 0, d_acc * w_h, 0) + + ``d_dot`` is one HBM store per element (the grid covers each + output element exactly once); ``d_w`` is computed via eager + PyTorch in the autograd wrapper as a single ATen reduce (cheap + + avoids the cross-block ``atomic_add`` traffic that hurt the + P38 BWD). + """ + + pid_b = tl.program_id(0) + pid_s = tl.program_id(1) + pid_p = tl.program_id(2) + + s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S) + p_offs = pid_p * BLOCK_P + tl.arange(0, BLOCK_P) + s_mask = s_offs < S + p_mask = p_offs < P + + dmasked = tl.load( + DSCORES_PTR + pid_b * S * P + s_offs[:, None] * P + p_offs[None, :], + mask=s_mask[:, None] & p_mask[None, :], + other=0.0, + ).to(tl.float32) + s_arr = s_offs[:, None] + p_arr = p_offs[None, :] + allowed = (p_arr + 1) * COMPRESS_RATIO - 1 <= s_arr + d_acc = tl.where(allowed, dmasked, 0.0) + + for h in tl.static_range(0, H): + dot_tile = tl.load( + DOT_PTR + pid_b * S * H * P + s_offs[:, None] * H * P + h * P + p_offs[None, :], + mask=s_mask[:, None] & p_mask[None, :], + other=0.0, + ).to(tl.float32) + relu_mask = dot_tile > 0.0 + + w_h = tl.load( + W_PTR + pid_b * S * H + s_offs * H + h, + mask=s_mask, + other=0.0, + ).to(tl.float32) + + d_relu = d_acc * w_h[:, None] + d_dot = tl.where(relu_mask, d_relu, 0.0) + tl.store( + DDOT_PTR + pid_b * S * H * P + s_offs[:, None] * H * P + h * P + p_offs[None, :], + d_dot, + mask=s_mask[:, None] & p_mask[None, :], + ) + + +# --------------------------------------------------------------------------- +# Heuristics +# --------------------------------------------------------------------------- + + +def _pick_block(s: int, p: int, h: int) -> tuple[int, int]: + """Pick BLOCK_S, BLOCK_P that fit MI355 register budget. + + Tail kernel is much cheaper than the P38 full-fuse one (no dot + product, only elementwise + reduce) so we can use larger blocks. + """ + block_s = 64 + block_p = 64 + return min(block_s, triton.next_power_of_2(max(1, s))), min(block_p, triton.next_power_of_2(max(1, p))) + + +# --------------------------------------------------------------------------- +# autograd.Function wrapper +# --------------------------------------------------------------------------- + + +class IndexerScorePostFn(torch.autograd.Function): + """Autograd-aware wrapper around the post-einsum tail kernels. + + Inputs: + dot [B, S, H, P] - eager einsum output (pre-relu); any float dtype. + w_i [B, S, H] - per-head weights; any float dtype. + compress_ratio - int (typically 4 for CSA). + out_dtype - dtype of the returned ``scores``. + + Returns: + scores [B, S, P] of dtype ``out_dtype``. + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, + dot: torch.Tensor, + w_i: torch.Tensor, + compress_ratio: int, + out_dtype: torch.dtype, + ): + if dot.dim() != 4: + raise ValueError(f"dot must be [B, S, H, P], got shape {tuple(dot.shape)}") + if w_i.dim() != 3: + raise ValueError(f"w_i must be [B, S, H], got shape {tuple(w_i.shape)}") + + B, S, H, P = dot.shape + Bw, Sw, Hw = w_i.shape + if B != Bw: + raise ValueError(f"Mismatched B: dot={B} w={Bw}") + if S != Sw: + raise ValueError(f"Mismatched S: dot={S} w={Sw}") + if H != Hw: + raise ValueError(f"Mismatched H: dot={H} w={Hw}") + if H not in _SUPPORTED_H: + raise ValueError(f"Unsupported H={H}; expected one of {_SUPPORTED_H}") + + dot_c = dot.contiguous() + w_c = w_i.contiguous() + device = dot_c.device + scores = torch.empty((B, S, P), dtype=out_dtype, device=device) + + block_s, block_p = _pick_block(S, P, H) + grid = (B, triton.cdiv(S, block_s), triton.cdiv(P, block_p)) + _indexer_score_post_fwd_kernel[grid]( + dot_c, + w_c, + scores, + B, + S, + P, + H=H, + COMPRESS_RATIO=int(compress_ratio), + BLOCK_S=block_s, + BLOCK_P=block_p, + OUT_DTYPE={ + torch.float32: tl.float32, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, + torch.float64: tl.float64, + }[out_dtype], + ) + + ctx.save_for_backward(dot_c, w_c) + ctx.compress_ratio = int(compress_ratio) + ctx.shape = (B, S, H, P) + ctx.in_dtypes = (dot.dtype, w_i.dtype) + return scores + + @staticmethod + def backward(ctx, d_scores): # type: ignore[override] + dot_c, w_c = ctx.saved_tensors + B, S, H, P = ctx.shape + compress_ratio = ctx.compress_ratio + dot_dtype, w_dtype = ctx.in_dtypes + + d_scores = d_scores.contiguous() + device = dot_c.device + + d_dot_fp32 = torch.empty((B, S, H, P), dtype=torch.float32, device=device) + d_w_fp32 = torch.empty((B, S, H), dtype=torch.float32, device=device) + + block_s, block_p = _pick_block(S, P, H) + grid_ddot = (B, triton.cdiv(S, block_s), triton.cdiv(P, block_p)) + _indexer_score_post_bwd_ddot_kernel[grid_ddot]( + d_scores, + dot_c, + w_c, + d_dot_fp32, + B, + S, + P, + H=H, + COMPRESS_RATIO=int(compress_ratio), + BLOCK_S=block_s, + BLOCK_P=block_p, + ) + + # d_w[b, s, h] = sum_p (d_acc * relu(dot[h])) — one program per + # (b, s_tile, h), loops over P internally so the reduce is fully + # local (no cross-block atomic_add). + block_p_inner = min(128, triton.next_power_of_2(max(1, P))) + grid_dw = (B, triton.cdiv(S, block_s), H) + _indexer_score_post_bwd_dw_kernel[grid_dw]( + d_scores, + dot_c, + d_w_fp32, + B, + S, + P, + H=H, + COMPRESS_RATIO=int(compress_ratio), + BLOCK_S=block_s, + BLOCK_P_INNER=block_p_inner, + ) + + return d_dot_fp32.to(dot_dtype), d_w_fp32.to(w_dtype), None, None + + +# --------------------------------------------------------------------------- +# Public entry points +# --------------------------------------------------------------------------- + + +def is_triton_path_enabled() -> bool: + """Return True iff ``PRIMUS_INDEXER_TRITON != "0"`` (default ``"1"``). + + The env knob is **re-purposed** at P41 to mean "post-einsum tail + fusion" (cheap, bandwidth-bound). Legacy P38 full-fuse path + lives behind ``PRIMUS_INDEXER_TRITON_FULL``. + + Plan-8 P57 close-out 2 (2026-05-15): default flipped from ``"0"`` + to ``"1"``. Microbench at V4-Flash widths is consistently positive + (FWD 4.30x / BWD 1.63x) and the EP=8 proxy A/B shows a small but + positive ~0.2 ms / iter gain. Set ``PRIMUS_INDEXER_TRITON=0`` to + revert to the eager Python body. + """ + return os.environ.get("PRIMUS_INDEXER_TRITON", "1") != "0" + + +def is_triton_kernel_supported(dot: torch.Tensor, w_i: torch.Tensor) -> bool: + """Return True iff the input shapes / device support the Triton path.""" + if not dot.is_cuda or not w_i.is_cuda: + return False + if dot.dim() != 4 or w_i.dim() != 3: + return False + _, _, H, _ = dot.shape + if H not in _SUPPORTED_H: + return False + return True + + +def indexer_score_post_triton( + dot: torch.Tensor, + w_i: torch.Tensor, + *, + compress_ratio: int, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Compute Indexer scores from the eager einsum output via Triton. + + Returns ``scores [B, S, P]`` of dtype ``out_dtype``. Masked + positions hold ``-inf``. + """ + return IndexerScorePostFn.apply(dot, w_i, compress_ratio, out_dtype) + + +__all__ = [ + "IndexerScorePostFn", + "indexer_score_post_triton", + "is_triton_path_enabled", + "is_triton_kernel_supported", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rmsnorm.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rmsnorm.py new file mode 100644 index 000000000..ca8410dc6 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rmsnorm.py @@ -0,0 +1,366 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton-fused RMSNorm FWD/BWD (small-kernel-fusion campaign 2026-07-03). + +Collapses the eager RMSNorm chain — ``x.float()`` (bf16→fp32 cast), +``x.pow(2)`` / ``x.square()``, ``.mean(-1)``, ``+ eps``, ``rsqrt``, +``* rstd``, ``.to(in_dtype)`` (fp32→bf16 cast), optional ``* weight`` — +into ONE Triton kernel (FWD) + ONE Triton kernel (BWD). + +Every non-TE eager RMS site in the DeepSeek-V4 model body routes through +this kernel: + +* ``_per_head_rms_norm`` (``deepseek_v4_attention``) — no weight, out=in_dtype. +* ``LocalRMSNorm`` (``compressor.kv_norm`` etc.) — weight (+grad), + mid-cast to in_dtype before the weight multiply, out=promote(in, weight). +* ``HyperMixer._packed_logits`` RMS — no weight, out=fp32. +* ``HyperHead.forward`` RMS — no weight, out=fp32. + +The eager reference (matched bit-for-bit modulo fp32 accumulation order): + +.. code-block:: python + + x32 = x.float() + rstd = torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps) + y = x32 * rstd # fp32 normalized value + if mid_cast: y = y.to(in_dtype) # LocalRMSNorm rounds here first + if weight is not None: y = y * weight + out = y.to(out_dtype) + +Gradients (weight applied per channel, mid-cast treated as identity for +autograd — matches ``Tensor.to`` grad semantics): + +.. code-block:: python + + # forward: y_k = w_k * rstd * x_k, rstd = (mean(x^2)+eps)^-1/2 + c = sum_k( g_k * w_k * x_k ) # reduce over D + dx_k = rstd * g_k * w_k - x_k * rstd**3 * c / D + dw_k = sum_over_rows( g_k * (x_k * rstd) ) # normalized value pre-weight + +Gating: routed through :func:`fused_rms_norm` when ``PRIMUS_RMSNORM_TRITON +!= "0"`` (default-on) and the input is a supported CUDA/HIP float tensor; +otherwise the eager reference runs. +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch +import triton +import triton.language as tl + +_TORCH_TO_TL_DTYPE = { + torch.float64: tl.float64, + torch.float32: tl.float32, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, +} + + +def _triton_dtype(t: torch.dtype): + try: + return _TORCH_TO_TL_DTYPE[t] + except KeyError as exc: + raise TypeError( + f"rmsnorm: unsupported dtype {t}; expected one of {list(_TORCH_TO_TL_DTYPE)}" + ) from exc + + +# --------------------------------------------------------------------------- +# Triton kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _rmsnorm_fwd_kernel( + X_PTR, # [N, D] contiguous + W_PTR, # [D] contiguous (fp32) or dummy when HAS_WEIGHT is False + OUT_PTR, # [N, D] contiguous (OUT_DTYPE) + RSTD_PTR, # [N] fp32 (saved for backward) + N, + D, + EPS: tl.constexpr, + BLOCK_D: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + MID_CAST: tl.constexpr, + IN_DTYPE: tl.constexpr, + OUT_DTYPE: tl.constexpr, +): + """One program per row; two passes over D (sum-of-squares, then write).""" + row = tl.program_id(0) + if row >= N: + return + x_row = X_PTR + row * D + + # Pass 1: sum of squares in fp32. + acc = tl.zeros((), dtype=tl.float32) + for off in range(0, D, BLOCK_D): + cols = off + tl.arange(0, BLOCK_D) + mask = cols < D + x = tl.load(x_row + cols, mask=mask, other=0.0).to(tl.float32) + acc += tl.sum(x * x, axis=0) + + rstd = 1.0 / tl.sqrt(acc / D + EPS) + tl.store(RSTD_PTR + row, rstd) + + # Pass 2: normalize (+ optional mid-cast, weight) and write. + for off in range(0, D, BLOCK_D): + cols = off + tl.arange(0, BLOCK_D) + mask = cols < D + x = tl.load(x_row + cols, mask=mask, other=0.0).to(tl.float32) + y = x * rstd + if MID_CAST: + y = y.to(IN_DTYPE).to(tl.float32) + if HAS_WEIGHT: + w = tl.load(W_PTR + cols, mask=mask, other=0.0).to(tl.float32) + y = y * w + tl.store(OUT_PTR + row * D + cols, y.to(OUT_DTYPE), mask=mask) + + +@triton.jit +def _rmsnorm_bwd_kernel( + X_PTR, # [N, D] contiguous (original input) + W_PTR, # [D] contiguous (fp32) or dummy + RSTD_PTR, # [N] fp32 + DY_PTR, # [N, D] contiguous (upstream grad) + DX_PTR, # [N, D] contiguous (output, IN_DTYPE) + DW_PTR, # [D] fp32 accumulator (atomic) or dummy + N, + D, + BLOCK_D: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + IN_DTYPE: tl.constexpr, +): + """One program per row. Computes dx (and atomic-accumulates dw).""" + row = tl.program_id(0) + if row >= N: + return + x_row = X_PTR + row * D + dy_row = DY_PTR + row * D + rstd = tl.load(RSTD_PTR + row) + + # Pass 1: c = sum_k( dy_k * w_k * x_k ). + c = tl.zeros((), dtype=tl.float32) + for off in range(0, D, BLOCK_D): + cols = off + tl.arange(0, BLOCK_D) + mask = cols < D + x = tl.load(x_row + cols, mask=mask, other=0.0).to(tl.float32) + dy = tl.load(dy_row + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_WEIGHT: + w = tl.load(W_PTR + cols, mask=mask, other=0.0).to(tl.float32) + c += tl.sum(dy * w * x, axis=0) + else: + c += tl.sum(dy * x, axis=0) + + coef = rstd * rstd * rstd * c / D + + # Pass 2: dx_k = rstd*dy_k*w_k - x_k*coef ; accumulate dw_k += dy_k*x_k*rstd. + for off in range(0, D, BLOCK_D): + cols = off + tl.arange(0, BLOCK_D) + mask = cols < D + x = tl.load(x_row + cols, mask=mask, other=0.0).to(tl.float32) + dy = tl.load(dy_row + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_WEIGHT: + w = tl.load(W_PTR + cols, mask=mask, other=0.0).to(tl.float32) + dx = rstd * dy * w - x * coef + tl.atomic_add(DW_PTR + cols, dy * (x * rstd), mask=mask) + else: + dx = rstd * dy - x * coef + tl.store(DX_PTR + row * D + cols, dx.to(IN_DTYPE), mask=mask) + + +def _pick_block_d(d: int) -> int: + """Tile the reduction axis; power-of-2, capped so registers stay bounded.""" + if d <= 2048: + return triton.next_power_of_2(d) + return 2048 + + +# --------------------------------------------------------------------------- +# torch.autograd.Function +# --------------------------------------------------------------------------- + + +class FusedRMSNormFn(torch.autograd.Function): + """Autograd wrapper for the fused RMSNorm FWD/BWD Triton kernels. + + ``apply(x, weight, eps, mid_cast, out_dtype)``: + + * ``x``: ``[..., D]`` float tensor (RMS over the last dim). + * ``weight``: ``[D]`` tensor or ``None`` (parameter-less RMS). + * ``eps``: numerical floor. + * ``mid_cast``: when ``True``, round the normalized value to ``x.dtype`` + before the weight multiply (matches ``LocalRMSNorm``); no-op when + ``weight is None``. + * ``out_dtype``: dtype of the returned tensor. + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, + x: torch.Tensor, + weight: Optional[torch.Tensor], + eps: float, + mid_cast: bool, + out_dtype: torch.dtype, + ) -> torch.Tensor: + D = x.shape[-1] + x2 = x.reshape(-1, D) + x2 = x2.contiguous() + N = x2.shape[0] + + has_weight = weight is not None + w32 = None + if has_weight: + if weight.shape[-1] != D: + raise ValueError(f"rmsnorm: weight dim {tuple(weight.shape)} != x last dim {D}") + w32 = weight.reshape(D).to(torch.float32).contiguous() + + out = torch.empty((N, D), dtype=out_dtype, device=x.device) + rstd = torch.empty((N,), dtype=torch.float32, device=x.device) + block_d = _pick_block_d(D) + grid = (N,) + _rmsnorm_fwd_kernel[grid]( + x2, + w32 if has_weight else x2, # dummy ptr when no weight + out, + rstd, + N, + D, + EPS=float(eps), + BLOCK_D=block_d, + HAS_WEIGHT=has_weight, + MID_CAST=bool(mid_cast) and has_weight, + IN_DTYPE=_triton_dtype(x.dtype), + OUT_DTYPE=_triton_dtype(out_dtype), + ) + + ctx.save_for_backward(x2, w32, rstd) + ctx.has_weight = has_weight + ctx.in_dtype = x.dtype + ctx.weight_dtype = weight.dtype if has_weight else None + ctx.weight_shape = tuple(weight.shape) if has_weight else None + ctx.block_d = block_d + ctx.D = D + ctx.x_shape = tuple(x.shape) + return out.reshape(ctx.x_shape) + + @staticmethod + def backward(ctx, dy: torch.Tensor): # type: ignore[override] + x2, w32, rstd = ctx.saved_tensors + D = ctx.D + N = x2.shape[0] + has_weight = ctx.has_weight + + # dy comes in out_dtype; the kernel upcasts to fp32 internally, so any + # float dtype is fine. Keep it as-is (contiguous) without a lossy cast. + dy2 = dy.reshape(-1, D).contiguous() + + dx = torch.empty((N, D), dtype=ctx.in_dtype, device=dy.device) + dw_acc = ( + torch.zeros((D,), dtype=torch.float32, device=dy.device) + if has_weight + else torch.empty((1,), dtype=torch.float32, device=dy.device) + ) + grid = (N,) + _rmsnorm_bwd_kernel[grid]( + x2, + w32 if has_weight else x2, + rstd, + dy2, + dx, + dw_acc, + N, + D, + BLOCK_D=ctx.block_d, + HAS_WEIGHT=has_weight, + IN_DTYPE=_triton_dtype(ctx.in_dtype), + ) + + dx = dx.reshape(ctx.x_shape) + dweight = None + if has_weight: + dweight = dw_acc.to(ctx.weight_dtype).reshape(ctx.weight_shape) + return dx, dweight, None, None, None + + +# --------------------------------------------------------------------------- +# Public entry points +# --------------------------------------------------------------------------- + + +def is_triton_path_enabled() -> bool: + """Default-on; A/B toggle via ``PRIMUS_RMSNORM_TRITON=0``.""" + return os.environ.get("PRIMUS_RMSNORM_TRITON", "1") != "0" + + +def is_triton_kernel_supported(x: torch.Tensor, weight: Optional[torch.Tensor]) -> bool: + """Supported iff CUDA/HIP float input (and matching-dtype-family weight).""" + if not x.is_cuda: + return False + if x.dtype not in _TORCH_TO_TL_DTYPE: + return False + if x.shape[-1] == 0: + return False + if weight is not None and not weight.is_cuda: + return False + return True + + +def eager_rms_norm( + x: torch.Tensor, + weight: Optional[torch.Tensor] = None, + *, + eps: float, + mid_cast: bool, + out_dtype: torch.dtype, +) -> torch.Tensor: + """Reference eager RMSNorm matching the fused kernel's contract.""" + in_dtype = x.dtype + x32 = x.float() + rstd = torch.rsqrt(x32.pow(2).mean(dim=-1, keepdim=True) + eps) + y = x32 * rstd + if weight is not None: + if mid_cast: + y = y.to(in_dtype) + y = y.to(torch.float32) * weight.to(torch.float32) + return y.to(out_dtype) + + +def fused_rms_norm( + x: torch.Tensor, + weight: Optional[torch.Tensor] = None, + *, + eps: float, + mid_cast: bool = False, + out_dtype: Optional[torch.dtype] = None, +) -> torch.Tensor: + """Dispatch: Triton path when enabled + supported, else eager reference. + + ``out_dtype`` defaults to the eager output dtype: + ``promote_types(x.dtype, weight.dtype)`` when weighted, else ``x.dtype``. + """ + if out_dtype is None: + if weight is not None: + out_dtype = torch.promote_types(x.dtype, weight.dtype) + else: + out_dtype = x.dtype + + if is_triton_path_enabled() and is_triton_kernel_supported(x, weight): + return FusedRMSNormFn.apply(x, weight, float(eps), bool(mid_cast), out_dtype) + return eager_rms_norm(x, weight, eps=eps, mid_cast=mid_cast, out_dtype=out_dtype) + + +__all__ = [ + "FusedRMSNormFn", + "fused_rms_norm", + "eager_rms_norm", + "is_triton_path_enabled", + "is_triton_kernel_supported", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py new file mode 100644 index 000000000..9aeaf2cd4 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py @@ -0,0 +1,778 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton-fused interleaved partial RoPE FWD/BWD (plan-6 P35). + +The eager body in :func:`apply_interleaved_partial_rope` is a 9-op chain: + +.. code-block:: python + + x_nope = x[..., :nope] # slice (1) + x_rope = x[..., nope:] # slice (2) + x_pairs = x_rope.reshape(..., rd // 2, 2) # reshape (3) + even, odd = x_pairs[..., 0], x_pairs[..., 1] + cos = cos.unsqueeze(-2).to(orig_dtype) # unsqueeze + cast (4) + sin = sin.unsqueeze(-2).to(orig_dtype) # unsqueeze + cast (5) + rot_even = even * cos - odd * sin # 2 muls + 1 sub (6) + rot_odd = even * sin + odd * cos # 2 muls + 1 add (7) + rotated = torch.stack([rot_even, rot_odd], -1).reshape(..., rd) # stack + reshape (8) + return torch.cat([x_nope, rotated], -1) # cat (9) + +The plan-5 P32 final EP=8 proxy trace attributes: + +* ``CatArrayBatchedCopy_contig`` ≈ **10.0 ms / 24 calls** to the closing + ``torch.cat`` (the nope-prefix copy + the rotated suffix into one + contiguous tensor), and +* a non-trivial share of ``elementwise_kernel_manual_unroll<128, 8>`` + (~61 ms / 693 calls) to the four broadcast muls. + +At 16 invocations per iter (q + k per ``DualRoPE`` call × 8 layers) the +per-call cost is **~3-5 ms** at the V4-Flash widths. + +This module collapses the 9-op chain into one Triton kernel that: + +1. Flattens the input to ``[N, H, head_dim]`` where ``N`` is the product + of all leading axes (caller does the ``.reshape(-1, H, head_dim)`` + plumbing; the kernel is shape-agnostic). +2. Per program processes a ``[BLOCK_H]`` slice of one position's + ``H * head_dim`` row. cos/sin are loaded **once per position** + (shared across the ``BLOCK_H`` heads in the program) so the per-call + HBM traffic for cos/sin is exactly ``N * rd / 2`` reads, not + ``N * H * rd / 2`` (which the broadcast-muls in the eager body would + imply if they hit memory). +3. Writes ``out [N, H, head_dim]`` in one pass: nope channels copied + verbatim; trailing ``rotary_dim`` channels rotated using the + interleaved (2k, 2k+1) pairing. No ``torch.cat``-style second + memcpy; the kernel writes the full output in a single contiguous + pass. + +The BWD kernel is the transpose of the FWD rotation matrix +(``cos, sin / -sin, cos -> cos, -sin / sin, cos``) — analytically: + +.. code-block:: python + + dx[..., 2k] = dout[..., 2k] * cos + dout[..., 2k+1] * sin + dx[..., 2k+1] = -dout[..., 2k] * sin + dout[..., 2k+1] * cos + dx[..., :nope] = dout[..., :nope] # straight copy + +Cos / sin are buffers (not Parameters) so they have no gradient. + +Gating: routed through :func:`apply_interleaved_partial_rope` when +``PRIMUS_ROPE_TRITON != "0"`` (default-on). Set to ``"0"`` to fall back +to the eager body (kept in tree for A/B and as the reference path for +G38). +""" + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +# --------------------------------------------------------------------------- +# Triton dtype mapping +# --------------------------------------------------------------------------- + +_TORCH_TO_TL_DTYPE = { + torch.float64: tl.float64, + torch.float32: tl.float32, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, +} + + +def _triton_dtype(t: torch.dtype): + try: + return _TORCH_TO_TL_DTYPE[t] + except KeyError as exc: + raise TypeError( + f"rope_interleaved_partial: unsupported dtype {t}; " f"expected one of {list(_TORCH_TO_TL_DTYPE)}" + ) from exc + + +# --------------------------------------------------------------------------- +# Triton kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _apply_rope_fwd_kernel( + X_PTR, # [N, H, head_dim] contiguous + COS_PTR, # [N, rd_half] contiguous (broadcast over H) + SIN_PTR, # [N, rd_half] contiguous (broadcast over H) + OUT_PTR, # [N, H, head_dim] contiguous + N, + H, + HEAD_DIM: tl.constexpr, + NOPE: tl.constexpr, + RD_HALF: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_NOPE: tl.constexpr, # next_pow2(NOPE) — block stride for nope copy + BLOCK_RD_HALF: tl.constexpr, # next_pow2(RD_HALF) + DTYPE: tl.constexpr, +): + """Apply interleaved partial RoPE FWD over a tile of ``BLOCK_H`` heads + for one position ``pid_n``. + + Layout: x / out are ``[N, H, head_dim]``; the program processes + ``x[pid_n, pid_h*BLOCK_H : (pid_h+1)*BLOCK_H, :]`` and writes the + rotated result to the same slice of ``out``. + + The nope prefix (``head_dim - rotary_dim`` channels) is copied + verbatim — kept inside the kernel so the eager-body + ``torch.cat([x_nope, rotated], -1)`` second-pass memcpy goes away. + """ + + pid_n = tl.program_id(0) + pid_h = tl.program_id(1) + + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = h_offs < H + + # cos / sin for this position — load once, broadcast across BLOCK_H heads. + rd_half_offs = tl.arange(0, BLOCK_RD_HALF) + rd_half_mask = rd_half_offs < RD_HALF + cos = tl.load( + COS_PTR + pid_n * RD_HALF + rd_half_offs, + mask=rd_half_mask, + other=0.0, + ) + sin = tl.load( + SIN_PTR + pid_n * RD_HALF + rd_half_offs, + mask=rd_half_mask, + other=0.0, + ) + cos = cos.to(DTYPE) + sin = sin.to(DTYPE) + + # Base pointers for this position × head block. + row_base = pid_n * H * HEAD_DIM + h_offs[:, None] * HEAD_DIM # [BLOCK_H, 1] + + # 1) Copy the nope channels verbatim. + if NOPE > 0: + nope_offs = tl.arange(0, BLOCK_NOPE) + nope_mask = (nope_offs < NOPE)[None, :] & h_mask[:, None] + x_nope = tl.load( + X_PTR + row_base + nope_offs[None, :], + mask=nope_mask, + other=0.0, + ) + tl.store( + OUT_PTR + row_base + nope_offs[None, :], + x_nope, + mask=nope_mask, + ) + + # 2) Rotate the trailing rotary_dim channels (interleaved pairs). + even_offs = NOPE + 2 * rd_half_offs + odd_offs = NOPE + 2 * rd_half_offs + 1 + pair_mask = h_mask[:, None] & rd_half_mask[None, :] + even = tl.load( + X_PTR + row_base + even_offs[None, :], + mask=pair_mask, + other=0.0, + ) + odd = tl.load( + X_PTR + row_base + odd_offs[None, :], + mask=pair_mask, + other=0.0, + ) + rot_even = even * cos[None, :] - odd * sin[None, :] + rot_odd = even * sin[None, :] + odd * cos[None, :] + tl.store( + OUT_PTR + row_base + even_offs[None, :], + rot_even, + mask=pair_mask, + ) + tl.store( + OUT_PTR + row_base + odd_offs[None, :], + rot_odd, + mask=pair_mask, + ) + + +@triton.jit +def _apply_rope_bwd_kernel( + DOUT_PTR, # [N, H, head_dim] contiguous + COS_PTR, # [N, rd_half] contiguous + SIN_PTR, # [N, rd_half] contiguous + DX_PTR, # [N, H, head_dim] contiguous + N, + H, + HEAD_DIM: tl.constexpr, + NOPE: tl.constexpr, + RD_HALF: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_NOPE: tl.constexpr, + BLOCK_RD_HALF: tl.constexpr, + DTYPE: tl.constexpr, +): + """Apply the transpose rotation for the BWD pass. + + .. code-block:: python + + dx[..., 2k] = dout[..., 2k] * cos + dout[..., 2k+1] * sin + dx[..., 2k+1] = -dout[..., 2k] * sin + dout[..., 2k+1] * cos + dx[..., :nope] = dout[..., :nope] + """ + + pid_n = tl.program_id(0) + pid_h = tl.program_id(1) + + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = h_offs < H + + rd_half_offs = tl.arange(0, BLOCK_RD_HALF) + rd_half_mask = rd_half_offs < RD_HALF + cos = tl.load( + COS_PTR + pid_n * RD_HALF + rd_half_offs, + mask=rd_half_mask, + other=0.0, + ) + sin = tl.load( + SIN_PTR + pid_n * RD_HALF + rd_half_offs, + mask=rd_half_mask, + other=0.0, + ) + cos = cos.to(DTYPE) + sin = sin.to(DTYPE) + + row_base = pid_n * H * HEAD_DIM + h_offs[:, None] * HEAD_DIM + + # 1) Straight copy of the nope-prefix gradient. + if NOPE > 0: + nope_offs = tl.arange(0, BLOCK_NOPE) + nope_mask = (nope_offs < NOPE)[None, :] & h_mask[:, None] + dout_nope = tl.load( + DOUT_PTR + row_base + nope_offs[None, :], + mask=nope_mask, + other=0.0, + ) + tl.store( + DX_PTR + row_base + nope_offs[None, :], + dout_nope, + mask=nope_mask, + ) + + # 2) Transposed rotation for the rotary suffix. + even_offs = NOPE + 2 * rd_half_offs + odd_offs = NOPE + 2 * rd_half_offs + 1 + pair_mask = h_mask[:, None] & rd_half_mask[None, :] + dout_even = tl.load( + DOUT_PTR + row_base + even_offs[None, :], + mask=pair_mask, + other=0.0, + ) + dout_odd = tl.load( + DOUT_PTR + row_base + odd_offs[None, :], + mask=pair_mask, + other=0.0, + ) + dx_even = dout_even * cos[None, :] + dout_odd * sin[None, :] + dx_odd = -dout_even * sin[None, :] + dout_odd * cos[None, :] + tl.store( + DX_PTR + row_base + even_offs[None, :], + dx_even, + mask=pair_mask, + ) + tl.store( + DX_PTR + row_base + odd_offs[None, :], + dx_odd, + mask=pair_mask, + ) + + +# --------------------------------------------------------------------------- +# Block-size heuristic +# --------------------------------------------------------------------------- + + +def _pick_block_h(h: int) -> int: + """Pick BLOCK_H tiling the heads axis. + + V4-Flash Q has ``H = 64`` heads, K has ``H = 1``. For H=1 the block + must be 1; for H=64 a BLOCK_H of 8 keeps the per-program live tile + under ~16 KiB at head_dim=512 / bf16 while amortising the cos/sin + load across 8 heads. + """ + + if h <= 1: + return 1 + if h <= 4: + return min(h, 4) + if h <= 16: + return 8 + return 8 # default for V4-Flash Q (H=64) and K (H=1 handled above) + + +# --------------------------------------------------------------------------- +# torch.autograd.Function wrapper +# --------------------------------------------------------------------------- + + +class RoPEInterleavedPartialFn(torch.autograd.Function): + """Autograd-aware wrapper around the FWD/BWD Triton kernels. + + The wrapper: + + 1. Flattens ``x`` to ``[N, H, head_dim]`` and ``cos / sin`` to + ``[N, rd_half]`` (callers may pass any leading shape; the wrapper + does the reshape). + 2. Casts cos / sin to ``x.dtype`` if needed (matches the plan-5 P32 + RoPE bf16 cast contract in :func:`apply_interleaved_partial_rope`). + 3. Saves only ``cos, sin, rotary_dim, head_dim, nope`` for the + backward — the input ``x`` is not needed, and the saved tensors + are buffers (cos / sin) so they have no autograd graph node. + + Layout assumptions (validated up front): + + * ``x.is_contiguous()`` — call ``.contiguous()`` in the caller if needed. + * ``cos.shape == sin.shape == leading_shape + (rotary_dim // 2,)`` where + ``leading_shape == x.shape[:-2]`` (i.e. one cos/sin row per position). + * ``rotary_dim <= x.shape[-1]`` and ``rotary_dim % 2 == 0``. + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + rotary_dim: int, + ) -> torch.Tensor: + if rotary_dim == 0: + ctx.save_for_backward(cos, sin) + ctx.rotary_dim = 0 + ctx.head_dim = x.shape[-1] + return x.contiguous() + + if rotary_dim % 2 != 0: + raise ValueError(f"rotary_dim must be even, got {rotary_dim}") + head_dim = x.shape[-1] + if rotary_dim > head_dim: + raise ValueError(f"rotary_dim ({rotary_dim}) must be <= head_dim ({head_dim})") + + x = x.contiguous() + leading_shape = x.shape[:-2] + H = x.shape[-2] + N = 1 + for s in leading_shape: + N *= s + + rd_half = rotary_dim // 2 + # cos / sin should have shape == leading_shape + (rd_half,) + if cos.shape[-1] != rd_half or sin.shape[-1] != rd_half: + raise ValueError( + f"cos/sin last dim must be rotary_dim // 2 ({rd_half}); " + f"got cos.shape={tuple(cos.shape)}, sin.shape={tuple(sin.shape)}" + ) + cos_flat = cos.contiguous().reshape(N, rd_half).to(x.dtype) + sin_flat = sin.contiguous().reshape(N, rd_half).to(x.dtype) + + out = torch.empty_like(x) + + nope = head_dim - rotary_dim + block_h = _pick_block_h(H) + block_nope = max(triton.next_power_of_2(max(nope, 1)), 1) + block_rd_half = triton.next_power_of_2(rd_half) + + grid = (N, triton.cdiv(H, block_h)) + _apply_rope_fwd_kernel[grid]( + x, + cos_flat, + sin_flat, + out, + N, + H, + HEAD_DIM=head_dim, + NOPE=nope, + RD_HALF=rd_half, + BLOCK_H=block_h, + BLOCK_NOPE=block_nope, + BLOCK_RD_HALF=block_rd_half, + DTYPE=_triton_dtype(x.dtype), + ) + + ctx.save_for_backward(cos_flat, sin_flat) + ctx.rotary_dim = rotary_dim + ctx.head_dim = head_dim + ctx.leading_shape = tuple(leading_shape) + ctx.H = H + return out + + @staticmethod + def backward(ctx, dout: torch.Tensor): # type: ignore[override] + rotary_dim = ctx.rotary_dim + cos_flat, sin_flat = ctx.saved_tensors + + if rotary_dim == 0: + return dout.contiguous(), None, None, None + + dout = dout.contiguous() + head_dim = ctx.head_dim + H = ctx.H + leading_shape = ctx.leading_shape + N = 1 + for s in leading_shape: + N *= s + + rd_half = rotary_dim // 2 + nope = head_dim - rotary_dim + block_h = _pick_block_h(H) + block_nope = max(triton.next_power_of_2(max(nope, 1)), 1) + block_rd_half = triton.next_power_of_2(rd_half) + + dx = torch.empty_like(dout) + + grid = (N, triton.cdiv(H, block_h)) + _apply_rope_bwd_kernel[grid]( + dout, + cos_flat, + sin_flat, + dx, + N, + H, + HEAD_DIM=head_dim, + NOPE=nope, + RD_HALF=rd_half, + BLOCK_H=block_h, + BLOCK_NOPE=block_nope, + BLOCK_RD_HALF=block_rd_half, + DTYPE=_triton_dtype(dout.dtype), + ) + return dx, None, None, None + + +# --------------------------------------------------------------------------- +# Fused variant: compute cos/sin IN-KERNEL from (position_ids, inv_freq). +# +# The plain RoPE kernel above consumes precomputed cos/sin tensors, which +# means the caller still pays 3 small kernels per call +# (``position_ids.float() * inv_freq`` -> ``cos`` -> ``sin``) plus a +# ``.to(x.dtype)`` cast, and (in ``DeepseekV4Attention._apply_rope_q_k``) +# recomputes them identically for Q and K. This variant folds the cos/sin +# generation into the rotation kernel: it loads the row's scalar position + +# the ``[rd_half]`` inv_freq vector and computes ``cos``/``sin`` in registers, +# so no cos/sin HBM tensor is ever materialised. YaRN is already baked into +# ``inv_freq`` (a buffer), so the compress base is handled transparently. +# --------------------------------------------------------------------------- + + +@triton.jit +def _rope_gen_fwd_kernel( + X_PTR, # [N, H, head_dim] contiguous + POS_PTR, # [N] (position per row; fp32) + INVFREQ_PTR, # [rd_half] fp32 + OUT_PTR, # [N, H, head_dim] contiguous + N, + H, + HEAD_DIM: tl.constexpr, + NOPE: tl.constexpr, + RD_HALF: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_NOPE: tl.constexpr, + BLOCK_RD_HALF: tl.constexpr, + DTYPE: tl.constexpr, +): + pid_n = tl.program_id(0) + pid_h = tl.program_id(1) + + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = h_offs < H + + rd_half_offs = tl.arange(0, BLOCK_RD_HALF) + rd_half_mask = rd_half_offs < RD_HALF + + # cos/sin for this position, computed once, broadcast across BLOCK_H heads. + pos = tl.load(POS_PTR + pid_n).to(tl.float32) + inv_freq = tl.load(INVFREQ_PTR + rd_half_offs, mask=rd_half_mask, other=0.0).to(tl.float32) + angle = pos * inv_freq + cos = tl.cos(angle).to(DTYPE) + sin = tl.sin(angle).to(DTYPE) + + row_base = pid_n * H * HEAD_DIM + h_offs[:, None] * HEAD_DIM + + if NOPE > 0: + nope_offs = tl.arange(0, BLOCK_NOPE) + nope_mask = (nope_offs < NOPE)[None, :] & h_mask[:, None] + x_nope = tl.load(X_PTR + row_base + nope_offs[None, :], mask=nope_mask, other=0.0) + tl.store(OUT_PTR + row_base + nope_offs[None, :], x_nope, mask=nope_mask) + + even_offs = NOPE + 2 * rd_half_offs + odd_offs = NOPE + 2 * rd_half_offs + 1 + pair_mask = h_mask[:, None] & rd_half_mask[None, :] + even = tl.load(X_PTR + row_base + even_offs[None, :], mask=pair_mask, other=0.0) + odd = tl.load(X_PTR + row_base + odd_offs[None, :], mask=pair_mask, other=0.0) + rot_even = even * cos[None, :] - odd * sin[None, :] + rot_odd = even * sin[None, :] + odd * cos[None, :] + tl.store(OUT_PTR + row_base + even_offs[None, :], rot_even, mask=pair_mask) + tl.store(OUT_PTR + row_base + odd_offs[None, :], rot_odd, mask=pair_mask) + + +@triton.jit +def _rope_gen_bwd_kernel( + DOUT_PTR, # [N, H, head_dim] contiguous + POS_PTR, # [N] fp32 + INVFREQ_PTR, # [rd_half] fp32 + DX_PTR, # [N, H, head_dim] contiguous + N, + H, + HEAD_DIM: tl.constexpr, + NOPE: tl.constexpr, + RD_HALF: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_NOPE: tl.constexpr, + BLOCK_RD_HALF: tl.constexpr, + DTYPE: tl.constexpr, +): + pid_n = tl.program_id(0) + pid_h = tl.program_id(1) + + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = h_offs < H + + rd_half_offs = tl.arange(0, BLOCK_RD_HALF) + rd_half_mask = rd_half_offs < RD_HALF + + pos = tl.load(POS_PTR + pid_n).to(tl.float32) + inv_freq = tl.load(INVFREQ_PTR + rd_half_offs, mask=rd_half_mask, other=0.0).to(tl.float32) + angle = pos * inv_freq + cos = tl.cos(angle).to(DTYPE) + sin = tl.sin(angle).to(DTYPE) + + row_base = pid_n * H * HEAD_DIM + h_offs[:, None] * HEAD_DIM + + if NOPE > 0: + nope_offs = tl.arange(0, BLOCK_NOPE) + nope_mask = (nope_offs < NOPE)[None, :] & h_mask[:, None] + dout_nope = tl.load(DOUT_PTR + row_base + nope_offs[None, :], mask=nope_mask, other=0.0) + tl.store(DX_PTR + row_base + nope_offs[None, :], dout_nope, mask=nope_mask) + + even_offs = NOPE + 2 * rd_half_offs + odd_offs = NOPE + 2 * rd_half_offs + 1 + pair_mask = h_mask[:, None] & rd_half_mask[None, :] + dout_even = tl.load(DOUT_PTR + row_base + even_offs[None, :], mask=pair_mask, other=0.0) + dout_odd = tl.load(DOUT_PTR + row_base + odd_offs[None, :], mask=pair_mask, other=0.0) + dx_even = dout_even * cos[None, :] + dout_odd * sin[None, :] + dx_odd = -dout_even * sin[None, :] + dout_odd * cos[None, :] + tl.store(DX_PTR + row_base + even_offs[None, :], dx_even, mask=pair_mask) + tl.store(DX_PTR + row_base + odd_offs[None, :], dx_odd, mask=pair_mask) + + +class RoPEFromPositionsFn(torch.autograd.Function): + """RoPE that computes cos/sin in-kernel from ``(position_ids, inv_freq)``. + + Eliminates the separate ``cos``/``sin`` generation kernels and HBM + tensors. ``position_ids`` / ``inv_freq`` are non-differentiable + (positions are integer, inv_freq is a buffer), so the backward returns + ``None`` for both and recomputes cos/sin in-kernel. + """ + + @staticmethod + def forward(ctx, x, pos_flat, inv_freq, rotary_dim): # type: ignore[override] + head_dim = x.shape[-1] + if rotary_dim % 2 != 0: + raise ValueError(f"rotary_dim must be even, got {rotary_dim}") + if rotary_dim > head_dim: + raise ValueError(f"rotary_dim ({rotary_dim}) must be <= head_dim ({head_dim})") + + x = x.contiguous() + leading = x.shape[:-2] + H = x.shape[-2] + N = 1 + for s in leading: + N *= s + rd_half = rotary_dim // 2 + if inv_freq.shape[-1] != rd_half: + raise ValueError( + f"inv_freq last dim must be rotary_dim // 2 ({rd_half}); got {tuple(inv_freq.shape)}" + ) + + pos_c = pos_flat.contiguous().to(torch.float32) + inv_c = inv_freq.contiguous().to(torch.float32) + out = torch.empty_like(x) + + nope = head_dim - rotary_dim + block_h = _pick_block_h(H) + block_nope = max(triton.next_power_of_2(max(nope, 1)), 1) + block_rd_half = triton.next_power_of_2(rd_half) + grid = (N, triton.cdiv(H, block_h)) + _rope_gen_fwd_kernel[grid]( + x, + pos_c, + inv_c, + out, + N, + H, + HEAD_DIM=head_dim, + NOPE=nope, + RD_HALF=rd_half, + BLOCK_H=block_h, + BLOCK_NOPE=block_nope, + BLOCK_RD_HALF=block_rd_half, + DTYPE=_triton_dtype(x.dtype), + ) + + ctx.save_for_backward(pos_c, inv_c) + ctx.rotary_dim = rotary_dim + ctx.head_dim = head_dim + ctx.H = H + ctx.N = N + return out + + @staticmethod + def backward(ctx, dout): # type: ignore[override] + pos_c, inv_c = ctx.saved_tensors + rotary_dim = ctx.rotary_dim + head_dim = ctx.head_dim + H = ctx.H + N = ctx.N + dout = dout.contiguous() + rd_half = rotary_dim // 2 + nope = head_dim - rotary_dim + block_h = _pick_block_h(H) + block_nope = max(triton.next_power_of_2(max(nope, 1)), 1) + block_rd_half = triton.next_power_of_2(rd_half) + dx = torch.empty_like(dout) + grid = (N, triton.cdiv(H, block_h)) + _rope_gen_bwd_kernel[grid]( + dout, + pos_c, + inv_c, + dx, + N, + H, + HEAD_DIM=head_dim, + NOPE=nope, + RD_HALF=rd_half, + BLOCK_H=block_h, + BLOCK_NOPE=block_nope, + BLOCK_RD_HALF=block_rd_half, + DTYPE=_triton_dtype(dout.dtype), + ) + return dx, None, None, None + + +def apply_rope_from_positions( + x: torch.Tensor, + position_ids: torch.Tensor, + inv_freq: torch.Tensor, + *, + rotary_dim: int, +) -> torch.Tensor: + """Fused interleaved partial RoPE that generates cos/sin in-kernel. + + ``x``: ``[..., H, head_dim]``; ``position_ids`` broadcastable to + ``x.shape[:-2]``; ``inv_freq``: ``[rotary_dim // 2]`` (YaRN pre-applied). + Dispatches to the Triton kernel when enabled + CUDA, else falls back to + the eager ``cos = pos*inv_freq -> cos/sin -> rotate`` path. + """ + if rotary_dim == 0: + return x + if is_triton_path_enabled() and x.is_cuda: + leading = x.shape[:-2] + pos_flat = position_ids.broadcast_to(leading).reshape(-1) + return RoPEFromPositionsFn.apply(x, pos_flat, inv_freq, rotary_dim) + # Eager fallback: build cos/sin then apply. + freqs = position_ids.float().unsqueeze(-1) * inv_freq + return eager_apply_interleaved_partial_rope(x, freqs.cos(), freqs.sin(), rotary_dim=rotary_dim) + + +# --------------------------------------------------------------------------- +# Public Python entry points +# --------------------------------------------------------------------------- + + +def is_triton_path_enabled() -> bool: + """Return True iff the ``PRIMUS_ROPE_TRITON`` env knob is not ``"0"``. + + Mirrors :func:`primus.backends.megatron.core.extensions._triton.stack_grouped_weight.is_triton_path_enabled` + (plan-6 P34). Default-on, A/B toggle via ``PRIMUS_ROPE_TRITON=0``. + """ + + return os.environ.get("PRIMUS_ROPE_TRITON", "1") != "0" + + +def eager_apply_interleaved_partial_rope( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + *, + rotary_dim: int, +) -> torch.Tensor: + """Reference eager implementation matching + :func:`primus.backends.megatron.core.transformer.dual_rope.apply_interleaved_partial_rope`. + + Kept here so the test / bench code can A/B against the exact eager + body without depending on the consumer module. Bit-for-bit + equivalent to the original eager body (same op order, same dtype + cast). + """ + + head_dim = x.shape[-1] + if rotary_dim > head_dim or rotary_dim % 2 != 0: + raise ValueError(f"rotary_dim must be even and <= head_dim ({head_dim}), got {rotary_dim}") + if rotary_dim == 0: + return x + + orig_dtype = x.dtype + nope = head_dim - rotary_dim + x_nope = x[..., :nope] + x_rope = x[..., nope:] + + x_pairs = x_rope.reshape(*x_rope.shape[:-1], rotary_dim // 2, 2) + even = x_pairs[..., 0] + odd = x_pairs[..., 1] + + cos = cos.unsqueeze(-2).to(orig_dtype) + sin = sin.unsqueeze(-2).to(orig_dtype) + + rot_even = even * cos - odd * sin + rot_odd = even * sin + odd * cos + + rotated = torch.stack([rot_even, rot_odd], dim=-1).reshape(*x_rope.shape[:-1], rotary_dim) + return torch.cat([x_nope, rotated], dim=-1) + + +def apply_rope_interleaved_partial( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + *, + rotary_dim: int, +) -> torch.Tensor: + """Dispatch: Triton path when ``PRIMUS_ROPE_TRITON != "0"`` (default), + else eager fallback. + + Shape contract matches the dual-RoPE caller: + + * ``x``: ``[..., H, head_dim]`` (any leading shape; flattened + internally). + * ``cos, sin``: ``[..., rotary_dim // 2]`` where the leading shape + is ``x.shape[:-2]`` (i.e. one cos/sin row per position). + * ``rotary_dim``: even and ``<= head_dim``. + + Output shape matches ``x``. + """ + + if rotary_dim == 0: + return x + + if is_triton_path_enabled() and x.is_cuda: + return RoPEInterleavedPartialFn.apply(x, cos, sin, rotary_dim) + return eager_apply_interleaved_partial_rope(x, cos, sin, rotary_dim=rotary_dim) + + +__all__ = [ + "RoPEInterleavedPartialFn", + "RoPEFromPositionsFn", + "apply_rope_interleaved_partial", + "apply_rope_from_positions", + "eager_apply_interleaved_partial_rope", + "is_triton_path_enabled", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/sinkhorn.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/sinkhorn.py new file mode 100644 index 000000000..70c622eee --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/sinkhorn.py @@ -0,0 +1,520 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton-fused Sinkhorn-Knopp FWD/BWD (plan-6 P36). + +Replaces the plan-5 P29 ``torch.compile`` Sinkhorn path with a +hand-rolled Triton kernel that runs the entire alternating row/col +normalize trajectory **in registers** per row of the leading axis. +The plan-5 P32 final trace attributes: + +* ``Torch-Compiled Region`` ≈ **21 ms / 16 calls** (FWD-side) +* ``CompiledFunctionBackward`` ≈ **41 ms / 16 calls** (BWD-side) + +to the cached compiled artefact built in +:func:`primus.backends.megatron.core.transformer.hyper_connection._build_compiled_sinkhorn`. +``torch.compile`` collapses the 1 + 2*(n_iters - 1) fp32 reductions +into one Inductor-fused Triton kernel, but its Dynamo-side +bookkeeping (the ``Torch-Compiled Region`` event itself) and +Inductor's BWD path are still non-trivial per-call overhead. + +The Triton path emits **one** FWD kernel and **one** BWD kernel; no +Dynamo bookkeeping per call. At V4-Flash widths +(``[..., K=4, K=4]``, ``n_iters=20``, fp32 internal compute) the +``K*K = 16`` elements per row fit in fp32 registers; the 39 normalize +steps (1 priming col-normalize + 19 row+col pairs) all run in +registers; the BWD recomputes the FWD trajectory in registers and +walks the analytic VJP backward step-by-step (the doubly-stochastic +projection has a closed-form VJP per step that fits in the same +register budget). + +Eager reference (see :func:`primus...hyper_connection.sinkhorn_normalize`): + +.. code-block:: python + + in_dtype = logits.dtype + m = logits.float() + m = m / (m.sum(dim=-2, keepdim=True) + eps) # initial col-normalize + for _ in range(n_iters - 1): + m = m / (m.sum(dim=-1, keepdim=True) + eps) # row-normalize + m = m / (m.sum(dim=-2, keepdim=True) + eps) # col-normalize + return m.to(in_dtype) + +Per-step VJP (closed form): + +.. code-block:: python + + # y = x / s, where s = sum(x, axis=axis) + eps + # dx = (dy - sum(dy * y, axis=axis, keep_dims=True)) / s + +Both axes (row / col) have the same shape of VJP (Triton's ``tl.sum`` +handles either with the right ``axis=`` argument), so the BWD kernel +is just the FWD-trajectory recompute followed by 39 of these VJP +steps in reverse order. + +Gating: routed through :func:`apply_sinkhorn_normalize_triton` when +``PRIMUS_SINKHORN_TRITON != "0"`` (default-on). Set to ``"0"`` to +fall back to either the plan-5 P29 compiled path (if +``use_compiled=True`` on the call site) or the eager path. + +Supported shape constraints: + +* Last two dims must be square and ``K ∈ {1, 2, 4, 8, 16}`` (Triton's + ``tl.arange`` needs a power-of-2 block); V4-Flash uses ``K=4``. +* ``logits`` must be contiguous (the wrapper calls ``.contiguous()`` + defensively). +* Any leading shape is supported (the wrapper flattens to + ``[N, K, K]``). +""" + +from __future__ import annotations + +import os + +import torch +import triton +import triton.language as tl + +# --------------------------------------------------------------------------- +# Triton dtype mapping +# --------------------------------------------------------------------------- + +_TORCH_TO_TL_DTYPE = { + torch.float64: tl.float64, + torch.float32: tl.float32, + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, +} + + +def _triton_dtype(t: torch.dtype): + try: + return _TORCH_TO_TL_DTYPE[t] + except KeyError as exc: + raise TypeError( + f"sinkhorn: unsupported dtype {t}; expected one of {list(_TORCH_TO_TL_DTYPE)}" + ) from exc + + +# Supported K (last two dim) values for the in-register Triton path. +# Block-size constraints (Triton requires power-of-2 block extents) plus +# register budget (~K*K*BLOCK_LEADING fp32 elements live per program) +# practically limit K to {1, 2, 4, 8, 16}. V4-Flash uses K=4. +_SUPPORTED_K = (1, 2, 4, 8, 16) + + +# --------------------------------------------------------------------------- +# Triton kernels +# --------------------------------------------------------------------------- + + +@triton.jit +def _sinkhorn_fwd_kernel( + X_PTR, # [N, K, K] contiguous + Y_PTR, # [N, K, K] contiguous + STATES_PTR, # [N, 2 * N_ITERS, K, K] contiguous (FWD trajectory cache) + N, + K: tl.constexpr, + N_ITERS: tl.constexpr, + EPS: tl.constexpr, + BLOCK_LEADING: tl.constexpr, + DTYPE: tl.constexpr, + COMPUTE_DTYPE: tl.constexpr, +): + """Run the 1 + 2*(N_ITERS - 1) alternating row/col normalize + trajectory for ``BLOCK_LEADING`` rows of the leading axis. + + The internal ``m`` tensor stays in fp32 throughout (matches the + eager :func:`sinkhorn_normalize` fp32 contract); the trailing cast + to ``DTYPE`` happens at the final store. + + Writes the full FWD state trajectory ``m_0, m_1, ..., m_{2*N_ITERS-1}`` + to ``STATES_PTR``; the BWD kernel reads it back so it can walk the + analytic VJP backward step-by-step without re-running FWD. At + ``K=4, N_ITERS=20`` the cache is ``40 * 16 * 4 = 2560`` bytes per + row (~10 MiB total at V4-Flash ``N=4096``) -- negligible HBM + overhead vs the 256 KiB FWD output of the same call. + """ + + pid = tl.program_id(0) + offs = pid * BLOCK_LEADING + tl.arange(0, BLOCK_LEADING) + mask_leading = offs < N + + r_offs = tl.arange(0, K) + c_offs = tl.arange(0, K) + + KK: tl.constexpr = K * K + base = offs[:, None, None] * KK + r_offs[None, :, None] * K + c_offs[None, None, :] + full_mask = mask_leading[:, None, None] + + # State buffer layout: [N, N_STATES, K, K] flattened row-major. + # state_base[s] = offs * (N_STATES * K * K) + s * (K * K) + r*K + c. + n_states_total: tl.constexpr = 2 * N_ITERS + state_row_stride: tl.constexpr = n_states_total * KK + state_base = offs[:, None, None] * state_row_stride + r_offs[None, :, None] * K + c_offs[None, None, :] + + m = tl.load(X_PTR + base, mask=full_mask, other=0.0).to(COMPUTE_DTYPE) + # Save m_0 = x.float() (state index 0). + tl.store(STATES_PTR + state_base, m, mask=full_mask) + + # Priming step: col-normalize (sum over the row axis, dim=-2 in the + # caller's [...,K,K]; this is axis=1 in our [BLOCK,K,K] layout). + s = tl.sum(m, axis=1, keep_dims=True) + EPS + m = m / s + tl.store(STATES_PTR + state_base + 1 * KK, m, mask=full_mask) # m_1 + + # Alternating loop -- N_ITERS - 1 (row, col) pairs. After the + # priming step we are at state index 1; each loop iteration writes + # two more states (the row half then the col half). + for it in tl.static_range(N_ITERS - 1): + s = tl.sum(m, axis=2, keep_dims=True) + EPS # row-normalize + m = m / s + tl.store(STATES_PTR + state_base + (2 + 2 * it) * KK, m, mask=full_mask) + s = tl.sum(m, axis=1, keep_dims=True) + EPS # col-normalize + m = m / s + tl.store(STATES_PTR + state_base + (3 + 2 * it) * KK, m, mask=full_mask) + + tl.store(Y_PTR + base, m.to(DTYPE), mask=full_mask) + + +@triton.jit +def _sinkhorn_bwd_kernel( + STATES_PTR, # [N, 2 * N_ITERS, K, K] contiguous (FWD trajectory cache) + DY_PTR, # [N, K, K] contiguous (upstream grad in caller dtype) + DX_PTR, # [N, K, K] contiguous (output) + N, + K: tl.constexpr, + N_ITERS: tl.constexpr, + EPS: tl.constexpr, + BLOCK_LEADING: tl.constexpr, + DTYPE: tl.constexpr, + COMPUTE_DTYPE: tl.constexpr, +): + """Apply the analytic VJP for the FWD trajectory. + + Reads the cached FWD states ``m_0, m_1, ..., m_{2*N_ITERS-1}`` from + HBM (written by :func:`_sinkhorn_fwd_kernel`) and walks the + trajectory backward, applying the per-step closed-form VJP: + + .. code-block:: python + + # forward: y = x / s, where s = sum(x, axis=axis) + eps + # backward: dx = (dy - sum(dy * y, axis=axis, keep_dims=True)) / s + + The HBM round-trip for the cache is ~10 MiB total at V4-Flash + widths (``K=4, N=4096``, 40 states / row × 16 fp32 / state); on a + 3 TB/s HBM device that's ~3 microseconds of read traffic -- + negligible vs the BWD's actual arithmetic. This sidesteps a + Triton AST-visitor restriction in our toolchain that rejects + runtime indexing of Python lists holding ``tl.tensor`` bundles, + which would otherwise let us recompute the trajectory in + registers. + + Trajectory indexing convention: + * ``m_0`` (state 0) = ``x.float()`` (input, BEFORE the priming step); + * step ``s`` transforms ``m_{s-1}`` -> ``m_s`` via + ``axis = 1`` (col-normalize) when ``s`` is odd, ``axis = 2`` + (row-normalize) when ``s`` is even; + * total steps ``n_steps = 1 + 2*(N_ITERS - 1) = 2*N_ITERS - 1``. + """ + + pid = tl.program_id(0) + offs = pid * BLOCK_LEADING + tl.arange(0, BLOCK_LEADING) + mask_leading = offs < N + + r_offs = tl.arange(0, K) + c_offs = tl.arange(0, K) + + KK: tl.constexpr = K * K + n_states_total: tl.constexpr = 2 * N_ITERS + state_row_stride: tl.constexpr = n_states_total * KK + + base = offs[:, None, None] * KK + r_offs[None, :, None] * K + c_offs[None, None, :] + full_mask = mask_leading[:, None, None] + state_base = offs[:, None, None] * state_row_stride + r_offs[None, :, None] * K + c_offs[None, None, :] + + dy = tl.load(DY_PTR + base, mask=full_mask, other=0.0).to(COMPUTE_DTYPE) + dm = dy + + # Walk the FWD trajectory backward. Trajectory parity is regular: + # step 1 (priming) is col-normalize, then alternating row, col, ... + # ending with step ``2*N_ITERS - 1`` (col). Walking BACKWARD we + # always see a (col, row) pair, repeated ``N_ITERS - 1`` times, + # followed by the priming col-step. + # + # We pair the loop iterations explicitly instead of using a + # ``step % 2`` runtime check so the axis= argument to ``tl.sum`` + # stays a compile-time Python int (otherwise the two branches' + # ``keep_dims=True`` outputs disagree on shape and Triton refuses + # to compile). + for i in range(N_ITERS - 1): + # Outer (col) step: indices 2*N_ITERS - 1, 2*N_ITERS - 3, ..., 3. + col_step = 2 * N_ITERS - 1 - 2 * i + m_before_c = tl.load( + STATES_PTR + state_base + (col_step - 1) * KK, + mask=full_mask, + other=0.0, + ) + m_after_c = tl.load( + STATES_PTR + state_base + col_step * KK, + mask=full_mask, + other=0.0, + ) + s_c = tl.sum(m_before_c, axis=1, keep_dims=True) + EPS + dot_c = tl.sum(dm * m_after_c, axis=1, keep_dims=True) + dm = (dm - dot_c) / s_c + + # Inner (row) step: indices 2*N_ITERS - 2, 2*N_ITERS - 4, ..., 2. + row_step = col_step - 1 + m_before_r = tl.load( + STATES_PTR + state_base + (row_step - 1) * KK, + mask=full_mask, + other=0.0, + ) + m_after_r = tl.load( + STATES_PTR + state_base + row_step * KK, + mask=full_mask, + other=0.0, + ) + s_r = tl.sum(m_before_r, axis=2, keep_dims=True) + EPS + dot_r = tl.sum(dm * m_after_r, axis=2, keep_dims=True) + dm = (dm - dot_r) / s_r + + # Final priming col-normalize step (step 1: m_0 -> m_1). + m_before_p = tl.load(STATES_PTR + state_base, mask=full_mask, other=0.0) + m_after_p = tl.load(STATES_PTR + state_base + KK, mask=full_mask, other=0.0) + s_p = tl.sum(m_before_p, axis=1, keep_dims=True) + EPS + dot_p = tl.sum(dm * m_after_p, axis=1, keep_dims=True) + dm = (dm - dot_p) / s_p + + tl.store(DX_PTR + base, dm.to(DTYPE), mask=full_mask) + + +# --------------------------------------------------------------------------- +# Block-leading heuristic +# --------------------------------------------------------------------------- + + +def _pick_block_leading(n: int, k: int) -> int: + """Pick ``BLOCK_LEADING`` based on K and the work-axis size. + + At K=4 the in-register tensor is ``[BLOCK_LEADING, 4, 4] = 16 + fp32 / row``; with ~256 VGPRs per warp on MI355 a BLOCK_LEADING of + ``128`` keeps live registers + VGPR-spilled state comfortable. + + At K=8 (16x larger per-row footprint) drop to 32. + + At K=16 drop to 8 -- the FWD trajectory keeps ~40 intermediate + matrices live during BWD recomputation, so per-program LDS / spill + needs to stay bounded. + """ + + if k <= 4: + cap = 128 + elif k <= 8: + cap = 32 + else: + cap = 8 + + if n < cap: + # Round to next power of 2 >= n (Triton requires power-of-2 block). + return max(1, triton.next_power_of_2(n)) + return cap + + +# --------------------------------------------------------------------------- +# torch.autograd.Function wrapper +# --------------------------------------------------------------------------- + + +class SinkhornNormalizeFn(torch.autograd.Function): + """Autograd-aware wrapper around the FWD/BWD Triton kernels. + + Saves only the input ``x`` for the backward; ``n_iters`` and + ``eps`` are static keys (compiled into the kernel binary cache). + + Shape: any ``[..., K, K]`` where ``K`` is a power of 2 in + ``{1, 2, 4, 8, 16}``. V4-Flash uses ``K=4, n_iters=20, eps=1e-6``. + """ + + @staticmethod + def forward( # type: ignore[override] + ctx, + logits: torch.Tensor, + n_iters: int, + eps: float, + ) -> torch.Tensor: + if logits.dim() < 2: + raise ValueError( + f"sinkhorn_normalize: input must be at least 2-D, got shape {tuple(logits.shape)}" + ) + K = logits.shape[-1] + if logits.shape[-2] != K: + raise ValueError( + f"sinkhorn_normalize: input must be square in the last two dims, " + f"got shape {tuple(logits.shape)}" + ) + if K not in _SUPPORTED_K: + raise ValueError(f"sinkhorn Triton path: unsupported K={K}; expected one of {_SUPPORTED_K}") + if int(n_iters) < 1: + raise ValueError(f"n_iters must be >= 1, got {n_iters}") + + x = logits.contiguous() + leading_shape = x.shape[:-2] + N = 1 + for s in leading_shape: + N *= s + + out = torch.empty_like(x) + # FWD trajectory cache: 2 * N_ITERS states per row of K*K fp32 + # elements. Saved for backward. At V4-Flash widths (K=4, + # N_ITERS=20, N=4096) this is 10 MiB per call -- negligible vs + # the 256 KiB FWD output and the ~170 GiB total HBM footprint + # of the proxy. + n_states_total = 2 * int(n_iters) + states_buf = torch.empty( + (N, n_states_total, K, K), + dtype=torch.float32, + device=x.device, + ) + + block_leading = _pick_block_leading(N, K) + grid = (triton.cdiv(N, block_leading),) + _sinkhorn_fwd_kernel[grid]( + x, + out, + states_buf, + N, + K=K, + N_ITERS=int(n_iters), + EPS=float(eps), + BLOCK_LEADING=block_leading, + DTYPE=_triton_dtype(x.dtype), + COMPUTE_DTYPE=tl.float32, + ) + + ctx.save_for_backward(states_buf) + ctx.n_iters = int(n_iters) + ctx.eps = float(eps) + ctx.K = K + ctx.leading_shape = tuple(leading_shape) + return out + + @staticmethod + def backward(ctx, dy: torch.Tensor): # type: ignore[override] + (states_buf,) = ctx.saved_tensors + n_iters: int = ctx.n_iters + eps: float = ctx.eps + K: int = ctx.K + leading_shape = ctx.leading_shape + + dy = dy.contiguous() + N = 1 + for s in leading_shape: + N *= s + + dx = torch.empty_like(dy) + block_leading = _pick_block_leading(N, K) + grid = (triton.cdiv(N, block_leading),) + _sinkhorn_bwd_kernel[grid]( + states_buf, + dy, + dx, + N, + K=K, + N_ITERS=n_iters, + EPS=eps, + BLOCK_LEADING=block_leading, + DTYPE=_triton_dtype(dy.dtype), + COMPUTE_DTYPE=tl.float32, + ) + # n_iters / eps gradients: not differentiable parameters. + return dx, None, None + + +# --------------------------------------------------------------------------- +# Public Python entry points +# --------------------------------------------------------------------------- + + +def is_triton_path_enabled() -> bool: + """Return True iff the ``PRIMUS_SINKHORN_TRITON`` env knob is not ``"0"``. + + Default-on; A/B toggle via ``PRIMUS_SINKHORN_TRITON=0``. + """ + + return os.environ.get("PRIMUS_SINKHORN_TRITON", "1") != "0" + + +def is_triton_kernel_supported(logits: torch.Tensor) -> bool: + """Return True iff the input shape / device is supported by the Triton path. + + Used by the dispatcher in + :func:`primus.backends.megatron.core.transformer.hyper_connection.sinkhorn_normalize` + to safely fall back to the compiled / eager path when the kernel + can't handle the input (non-CUDA, non-square, K out-of-range). + """ + + if not logits.is_cuda: + return False + if logits.dim() < 2: + return False + K = logits.shape[-1] + if logits.shape[-2] != K: + return False + if K not in _SUPPORTED_K: + return False + return True + + +def eager_sinkhorn_normalize( + logits: torch.Tensor, + *, + n_iters: int = 20, + eps: float = 1e-6, +) -> torch.Tensor: + """Reference eager implementation matching + :func:`primus.backends.megatron.core.transformer.hyper_connection.sinkhorn_normalize` + eager body bit-for-bit (same op order, same fp32 cast contract). + + Kept here so the unit tests / bench can A/B against the canonical + eager body without depending on the consumer module. + """ + + in_dtype = logits.dtype + m = logits.float() + m = m / (m.sum(dim=-2, keepdim=True) + eps) + for _ in range(max(n_iters - 1, 0)): + m = m / (m.sum(dim=-1, keepdim=True) + eps) + m = m / (m.sum(dim=-2, keepdim=True) + eps) + return m.to(in_dtype) + + +def sinkhorn_normalize_triton( + logits: torch.Tensor, + *, + n_iters: int = 20, + eps: float = 1e-6, +) -> torch.Tensor: + """Run the Triton-fused Sinkhorn-Knopp normalize. + + Dispatcher: if the env knob is on AND the shape is supported, call + the Triton path; else fall back to the eager body in + :func:`eager_sinkhorn_normalize`. + """ + + if is_triton_path_enabled() and is_triton_kernel_supported(logits): + return SinkhornNormalizeFn.apply(logits, n_iters, eps) + return eager_sinkhorn_normalize(logits, n_iters=n_iters, eps=eps) + + +__all__ = [ + "SinkhornNormalizeFn", + "sinkhorn_normalize_triton", + "eager_sinkhorn_normalize", + "is_triton_path_enabled", + "is_triton_kernel_supported", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/__init__.py new file mode 100644 index 000000000..bbedd45f7 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/__init__.py @@ -0,0 +1,23 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton **v0** attention backend — DEPRECATED gathered CSA. + +The original ``compress_ratio == 4`` CSA path that consumes a pre-gathered +``[B, Sq, K_topk, head_dim]`` tensor. It is ~30-260x slower than the v1 pool +path (see ``deepseek-v4/develop/perf/attention_perf.md``) and is NOT used by the +production dispatch. Kept for reference / correctness tests only. + +Prefer :mod:`.._triton_v1` (pool CSA + dense/HCA) or :mod:`.._triton_v2` +(fused single-latent sparse-MLA). +""" + +from .v4_csa_attention import V4CSAAttentionFn, v4_csa_attention_v0 + +__all__ = [ + "v4_csa_attention_v0", + "V4CSAAttentionFn", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention.py new file mode 100644 index 000000000..5f7d2f853 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention.py @@ -0,0 +1,238 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 CSA (gathered) attention — Triton **v0** autograd entry point. DEPRECATED. + +The original ``gathered`` CSA path (takes a pre-gathered ``[B, Sq, K, D]`` +tensor). It is ~30-260x slower than the v1 pool path (see attention_perf.md) +and is NOT used by the production dispatch — retained for reference/tests only. +Prefer ``_triton_v1`` (pool) or ``_triton_v2`` (fused sparse-MLA). +""" +from __future__ import annotations + +from typing import Optional + +import torch + +from primus.backends.megatron.core.transformer.v4_attention_kernels import ( + _flydsl_v0_deprecated as _flydsl, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels import _tilelang +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v0_deprecated.v4_csa_attention_bwd import ( + _launch_v4_csa_attention_bwd, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v0_deprecated.v4_csa_attention_fwd import ( + _launch_v4_csa_attention_fwd, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention import ( + v4_attention_v1, +) + + +class V4CSAAttentionFn(torch.autograd.Function): + """Triton FWD + Triton BWD for the CSA fused attention path.""" + + @staticmethod + def forward( # type: ignore[override] + ctx, + q: torch.Tensor, # [B, H, Sq, D] + k_local: torch.Tensor, # [B, H, Sq, D] + v_local: torch.Tensor, # [B, H, Sq, D] + gathered: torch.Tensor, # [B, Sq, K_topk, D] + sparse_mask: torch.Tensor, # [B, Sq, K_topk] + sink: Optional[torch.Tensor], # [H] or None + swa_window: int, + attn_dropout: float, + training: bool, + scale: float, + ) -> torch.Tensor: + if attn_dropout > 0.0 and training: + # Plan-4 P26 does not implement dropout in the kernel — V4 + # is trained with attn_dropout=0 so this branch is unreachable + # in production. We refuse explicitly so a stray non-zero + # dropout configuration raises rather than silently dropping + # the kernel path. + raise NotImplementedError( + "v4_csa_attention_v0 does not implement in-kernel attention " + "dropout (V4 trains with attn_dropout=0). Got " + f"attn_dropout={attn_dropout}, training={training}." + ) + + out, lse = _launch_v4_csa_attention_fwd( + q, + k_local, + v_local, + gathered, + sparse_mask, + sink=sink, + swa_window=swa_window, + scale=scale, + ) + # Save tensors the BWD kernel needs. ``sink`` may be None; we + # stash that fact on ``ctx`` because ``save_for_backward`` does + # not accept ``None``. + ctx.save_for_backward(q, k_local, v_local, gathered, sparse_mask, out, lse, sink) + ctx.swa_window = int(swa_window) + ctx.attn_dropout = float(attn_dropout) + ctx.training_mode = bool(training) + ctx.scale = float(scale) + ctx.sink_was_none = sink is None + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): # type: ignore[override] + """Triton BWD: re-materialises joint ``P`` from saved LSE; emits all five gradients.""" + q, k_local, v_local, gathered, sparse_mask, out, lse, sink = ctx.saved_tensors + + sink_arg = None if ctx.sink_was_none else sink + + # Ensure dout is contiguous in the [B, H, Sq, D] layout the + # kernel expects. + if not grad_out.is_contiguous(): + grad_out = grad_out.contiguous() + + dq, dk_local, dv_local, dgathered, dsink = _launch_v4_csa_attention_bwd( + q, + k_local, + v_local, + gathered, + sparse_mask, + out, + grad_out, + lse, + sink=sink_arg, + swa_window=ctx.swa_window, + scale=ctx.scale, + ) + + # Honor needs_input_grad: zero-cost ``None`` for inputs that + # don't want gradients. The kernel still computed them (the + # main cost is the matmul, not the per-output cast), so this is + # purely cleanliness. + if not ctx.needs_input_grad[0]: + dq = None + if not ctx.needs_input_grad[1]: + dk_local = None + if not ctx.needs_input_grad[2]: + dv_local = None + if not ctx.needs_input_grad[3]: + dgathered = None + # sparse_mask (index 4) is built from the indexer's ``topk_idxs >= + # 0`` test — it is NOT a learnable parameter and never needs a + # gradient. The kernel does not produce one. + if not ctx.needs_input_grad[5] or ctx.sink_was_none: + dsink = None + + # Forward signature: (q, k_local, v_local, gathered, sparse_mask, + # sink, swa_window, attn_dropout, training, scale). + return dq, dk_local, dv_local, dgathered, None, dsink, None, None, None, None + + +def v4_csa_attention_v0( + q: torch.Tensor, # [B, H, Sq, D] + k_local: torch.Tensor, # [B, H, Sq, D] + v_local: torch.Tensor, # [B, H, Sq, D] + gathered: torch.Tensor, # [B, Sq, K_topk, D] + *, + sink: Optional[torch.Tensor], # [H] or None + swa_window: int, + sparse_mask: torch.Tensor, # [B, Sq, K_topk] + attn_dropout: float, + training: bool, + scale: float, + use_tilelang: bool = False, + use_flydsl: bool = False, +) -> torch.Tensor: + """Triton-backed V4 CSA fused attention. + + Drop-in replacement for :func:`eager_v4_csa_attention` with + identical signature and dtype contract. Routes through + :class:`V4CSAAttentionFn` so autograd works. + + When ``gathered.shape[2] == 0`` (degenerate Indexer state — no + valid top-K positions) the wrapper short-circuits to the dense + :func:`v4_attention_v1` kernel: the local SWA + sink path is exactly + what CSA reduces to in that limit, and the dense kernel handles it + natively. ``sparse_mask`` is unused on that path so it is allowed + to be empty too. + + ``use_tilelang`` is plumbed by ``DeepseekV4Attention.forward`` + from the ``use_v4_tilelang_csa_attention`` config flag and only + triggers a tilelang dispatch when the relevant plan-8 P54 / P55 + kernels are registered (otherwise the dispatcher warns once and + falls back here). + + Returns ``[B, H, Sq, D]`` in ``v_local.dtype``. + """ + K_topk = gathered.shape[2] + if K_topk == 0: + # Degenerate sparse branch — fall through to the dense kernel. + # CSA's local SWA branch matches v4_attention_v1's dense+SWA+sink + # path bit-identically when K_topk == 0 (the joint softmax + # collapses to the local-only softmax). + return v4_attention_v1( + q, + k_local, + v_local, + sink=sink, + swa_window=swa_window, + additive_mask=None, + attn_dropout=attn_dropout, + training=training, + scale=scale, + ) + + # Plan-8 P49 / P57 close-out 2: tilelang dispatcher hook for the + # CSA family. Defaults OFF; only fires when the caller passes + # ``use_tilelang=True`` (i.e. the config flag is set). + if _tilelang.should_dispatch("v4_csa_attention_fwd", enabled=use_tilelang): + return _tilelang.v4_csa_attention_fwd_tilelang( + q, + k_local, + v_local, + gathered, + sparse_mask=sparse_mask, + sink=sink, + swa_window=swa_window, + attn_dropout=attn_dropout, + training=training, + scale=scale, + ) + # FlyDSL CSA backend hook (forward-only; soft-dep, default off). + # Short-circuits on enabled=False; falls back to Triton if the + # runtime/kernel is unavailable. Inference/eval path -- training flows + # through V4CSAAttentionFn below. + if _flydsl.should_dispatch("v4_csa_attention_fwd", enabled=use_flydsl): + return _flydsl.v4_csa_attention_fwd_flydsl( + q, + k_local, + v_local, + gathered, + sparse_mask=sparse_mask, + sink=sink, + swa_window=swa_window, + scale=scale, + attn_dropout=attn_dropout, + training=training, + ) + return V4CSAAttentionFn.apply( + q, + k_local, + v_local, + gathered, + sparse_mask, + sink, + swa_window, + attn_dropout, + training, + scale, + ) + + +__all__ = [ + "V4CSAAttentionFn", + "v4_csa_attention_v0", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_bwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_bwd.py new file mode 100644 index 000000000..5e1355fdd --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_bwd.py @@ -0,0 +1,488 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 CSA attention backward Triton kernel (plan-4 P26, ``compress_ratio == 4``). + +Two-kernel design (mirroring :mod:`v4_attention_bwd`): + +* Pre-pass: ``D[b, h, m] = sum_d (dout[b,h,m,d] * out[b,h,m,d])`` — + reuses :func:`_v4_attention_bwd_preprocess_kernel` from the dense + module since the contract is identical. +* Main pass: one program per ``(b, qhid, m)`` query row; re-materialises + the joint softmax row from the saved LSE; emits + + :: + dq [B, H, Sq, D] direct store (one program per row) + dk_local [B, H, Sq, D] atomic-add (multiple m's hit same n) + dv_local [B, H, Sq, D] atomic-add + dgathered [B, Sq, K_topk, D] atomic-add (no H dim — broadcast in fwd + means all H heads contribute) + dsink [H] atomic-add per query + +dtype contract: + +* All inputs loaded in input dtype; per-row dot products reduce in fp32 + via ``.to(tl.float32)`` upcast before the multiply (matches the FWD's + bf16-tensor-core / fp32-accumulator semantics). +* The online ``P / dP / dS`` re-materialisation is fp32 (matches the + FWD's softmax-in-fp32 contract). +* Output gradients are returned in input dtype (cast from fp32 buffers + by the launcher). + +Math derivation (per query (b, h, m), see plan-4 ``02-phase-details.md`` +Phase 26 section): + + joint_logits = cat(qk_local, qk_sparse, sink_h) + P_j = exp(joint_logits[j] - lse) + out_d = sum_n P_local[n] * v_local[n,d] + sum_k P_sparse[k] * g[k,d] + D = sum_d (dout[d] * out[d]) + dP_local[n] = sum_d (dout[d] * v_local[n,d]) + dP_sparse[k] = sum_d (dout[d] * g[k,d]) + dS_local[n] = P_local[n] * (dP_local[n] - D) + dS_sparse[k] = P_sparse[k] * (dP_sparse[k] - D) + dS_sink = -P_sink * D # sink val is 0 + + dq[d] = sum_n dS_local[n] * scale * k_local[n,d] + + sum_k dS_sparse[k] * scale * g[k,d] + dk_local[n,d] += dS_local[n] * scale * q[d] + dv_local[n,d] += P_local[n] * dout[d] + dgathered[k,d] += dS_sparse[k] * scale * q[d] + + P_sparse[k] * dout[d] # both branches + dsink_h += dS_sink + +Edge cases: + +* ``K_topk == 0`` — the wrapper short-circuits to the dense + :func:`v4_attention_v1` BWD before reaching this kernel. +* All-masked tile rows — uses ``NEG_INF = -1e30`` finite sentinel so + ``exp(NEG_INF - lse) = exp(-large) ≈ 0`` for fully-masked positions + (matches the FWD). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import ( + _v4_attention_bwd_preprocess_kernel, +) + +# --------------------------------------------------------------------------- +# Main BWD kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_csa_attention_bwd_kernel( + Q, + K_LOCAL, + V_LOCAL, + GATHERED, + SPARSE_MASK, + DOUT, + LSE, + D, + DQ, # fp32 buffer [B, H, Sq, D] + DK_LOCAL, # fp32 buffer [B, H, Sq, D] + DV_LOCAL, # fp32 buffer [B, H, Sq, D] + DGATHERED, # fp32 buffer [B, Sq, K_topk, D] + DSINK, # fp32 buffer [H] or sentinel + SINK, # [H] or sentinel + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_klb, + stride_klh, + stride_kln, + stride_kld, + stride_vlb, + stride_vlh, + stride_vln, + stride_vld, + stride_gb, + stride_gm, + stride_gk, + stride_gd, + stride_smb, + stride_smm, + stride_smk, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dqb, + stride_dqh, + stride_dqm, + stride_dqd, + stride_dklb, + stride_dklh, + stride_dkln, + stride_dkld, + stride_dvlb, + stride_dvlh, + stride_dvln, + stride_dvld, + stride_dgb, + stride_dgm, + stride_dgk, + stride_dgd, + seqlen_q, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + SWA_WINDOW: tl.constexpr, + HAS_SINK: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """V4 CSA fused-attention BWD (one program per (b, qhid, m) query row).""" + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + bid = pid_bh // HEAD_Q + qhid = pid_bh % HEAD_Q + + offs_d = tl.arange(0, BLOCK_DMODEL) + + NEG_INF: tl.constexpr = -1.0e30 + + q_active = pid_m < seqlen_q + + # ---- Load Q row, dout row, lse, D scalar ------------------------------ + q_row_offset = bid * stride_qb + qhid * stride_qh + pid_m * stride_qm + q = tl.load(Q + q_row_offset + offs_d * stride_qd, mask=q_active, other=0.0) + + do_row_offset = bid * stride_dob + qhid * stride_doh + pid_m * stride_dom + dout = tl.load(DOUT + do_row_offset + offs_d * stride_dod, mask=q_active, other=0.0) + q_f = q.to(tl.float32) + dout.to(tl.float32) + + lse = tl.load( + LSE + bid * stride_lb + qhid * stride_lh + pid_m * stride_lm, + mask=q_active, + other=0.0, + ) + dvec = tl.load( + D + bid * stride_db + qhid * stride_dh + pid_m * stride_dm, + mask=q_active, + other=0.0, + ) + + # ---- Sink contribution to dsink --------------------------------------- + # dS_sink = -P_sink * D; logit_sink = sink_h, so dsink_h += dS_sink. + if HAS_SINK: + sink_h = tl.load(SINK + qhid).to(tl.float32) + p_sink = tl.exp(sink_h - lse) + # Mask boundary rows so they don't contribute. + dsink_contrib = tl.where(q_active, -p_sink * dvec, 0.0) + tl.atomic_add(DSINK + qhid, dsink_contrib) + + # dq accumulator (fp32, kept in registers across the n-loop and k-loop) + dq = tl.zeros([BLOCK_DMODEL], dtype=tl.float32) + + # ---- Local SWA branch ------------------------------------------------- + n_loop_end = pid_m + 1 + if n_loop_end > seqlen_q: + n_loop_end = seqlen_q + + if SWA_WINDOW > 0: + n_lo_raw = pid_m - SWA_WINDOW + 1 + if n_lo_raw < 0: + n_lo_raw = 0 + n_loop_start = (n_lo_raw // BLOCK_N) * BLOCK_N + else: + n_loop_start = 0 + + for n_start in range(n_loop_start, n_loop_end, BLOCK_N): + offs_n = n_start + tl.arange(0, BLOCK_N) + + kl_ptrs = ( + K_LOCAL + + bid * stride_klb + + qhid * stride_klh + + offs_n[:, None] * stride_kln + + offs_d[None, :] * stride_kld + ) + kl_load_mask = offs_n[:, None] < seqlen_q + kl = tl.load(kl_ptrs, mask=kl_load_mask, other=0.0) + + vl_ptrs = ( + V_LOCAL + + bid * stride_vlb + + qhid * stride_vlh + + offs_n[:, None] * stride_vln + + offs_d[None, :] * stride_vld + ) + vl = tl.load(vl_ptrs, mask=kl_load_mask, other=0.0) + + # Re-materialise qk in fp32 (matches FWD). + kl_f = kl.to(tl.float32) + qk = tl.sum(kl_f * q_f[None, :], axis=1) * sm_scale + + if SWA_WINDOW > 0: + in_window = (offs_n >= pid_m - SWA_WINDOW + 1) & (offs_n <= pid_m) + else: + in_window = offs_n <= pid_m + qk = tl.where(in_window, qk, NEG_INF) + qk = tl.where(offs_n < seqlen_q, qk, NEG_INF) + # Boundary: off-grid m rows have lse=0 already, but we additionally + # zero this whole tile's contribution by forcing qk to NEG_INF. + qk = tl.where(q_active, qk, NEG_INF) + + # P = exp(qk - lse) (joint softmax slice for the local branch) + p = tl.exp(qk - lse) + + # dP[n] = sum_d (dout[d] * vl[n, d]) + dp = tl.sum(dout[None, :].to(tl.float32) * vl.to(tl.float32), axis=1) + + # dS[n] = P[n] * (dP[n] - D) + ds = p * (dp - dvec) + + # dq += sum_n (ds[n] * scale * kl[n, d]) + dq += tl.sum(ds[:, None] * kl.to(tl.float32), axis=0) * sm_scale + + # dk_local[n, d] += ds[n] * scale * q[d] — atomic-add into fp32 buf + dk_contrib = ds[:, None] * sm_scale * q[None, :].to(tl.float32) + dk_ptrs = ( + DK_LOCAL + + bid * stride_dklb + + qhid * stride_dklh + + offs_n[:, None] * stride_dkln + + offs_d[None, :] * stride_dkld + ) + tl.atomic_add(dk_ptrs, dk_contrib, mask=kl_load_mask, sem="relaxed") + + # dv_local[n, d] += p[n] * dout[d] — atomic-add into fp32 buf + dv_contrib = p[:, None] * dout[None, :].to(tl.float32) + dv_ptrs = ( + DV_LOCAL + + bid * stride_dvlb + + qhid * stride_dvlh + + offs_n[:, None] * stride_dvln + + offs_d[None, :] * stride_dvld + ) + tl.atomic_add(dv_ptrs, dv_contrib, mask=kl_load_mask, sem="relaxed") + + # ---- Sparse branch ---------------------------------------------------- + for k_start in range(0, K_topk, BLOCK_K): + offs_k = k_start + tl.arange(0, BLOCK_K) + + g_ptrs = ( + GATHERED + + bid * stride_gb + + pid_m * stride_gm + + offs_k[:, None] * stride_gk + + offs_d[None, :] * stride_gd + ) + g_load_mask = offs_k[:, None] < K_topk + g = tl.load(g_ptrs, mask=g_load_mask, other=0.0) + + sm_ptrs = SPARSE_MASK + bid * stride_smb + pid_m * stride_smm + offs_k * stride_smk + sm_load_mask = offs_k < K_topk + sm = tl.load(sm_ptrs, mask=sm_load_mask, other=0.0).to(tl.float32) + + qk_sparse = tl.sum(g.to(tl.float32) * q[None, :].to(tl.float32), axis=1) * sm_scale + sm + qk_sparse = tl.where(offs_k < K_topk, qk_sparse, NEG_INF) + qk_sparse = tl.where(q_active, qk_sparse, NEG_INF) + + p = tl.exp(qk_sparse - lse) + + # dP[k] = sum_d (dout[d] * g[k, d]) + dp = tl.sum(dout[None, :].to(tl.float32) * g.to(tl.float32), axis=1) + ds = p * (dp - dvec) + + # dq += sum_k (ds[k] * scale * g[k, d]) + dq += tl.sum(ds[:, None] * g.to(tl.float32), axis=0) * sm_scale + + # dgathered[k, d] += ds[k] * scale * q[d] + p[k] * dout[d] + # gathered is broadcast across H in the FWD, so this atomic-add + # accumulates contributions from every query head — matches the + # eager autograd semantics of ``gathered.unsqueeze(1).expand(B, H, + # Sq, K, D)``. + dg_contrib = ds[:, None] * sm_scale * q[None, :].to(tl.float32) + p[:, None] * dout[None, :].to( + tl.float32 + ) + dg_ptrs = ( + DGATHERED + + bid * stride_dgb + + pid_m * stride_dgm + + offs_k[:, None] * stride_dgk + + offs_d[None, :] * stride_dgd + ) + tl.atomic_add(dg_ptrs, dg_contrib, mask=g_load_mask, sem="relaxed") + + # ---- Store dq (direct — no collisions across programs) ---------------- + dq_offset = bid * stride_dqb + qhid * stride_dqh + pid_m * stride_dqm + tl.store(DQ + dq_offset + offs_d * stride_dqd, dq, mask=q_active) + + +def _launch_v4_csa_attention_bwd( + q: torch.Tensor, # [B, H, Sq, D] + k_local: torch.Tensor, # [B, H, Sq, D] + v_local: torch.Tensor, # [B, H, Sq, D] + gathered: torch.Tensor, # [B, Sq, K_topk, D] + sparse_mask: torch.Tensor, # [B, Sq, K_topk] + out: torch.Tensor, # [B, H, Sq, D] (FWD output) + dout: torch.Tensor, # [B, H, Sq, D] + lse: torch.Tensor, # [B, H, Sq] fp32 + *, + sink: Optional[torch.Tensor], # [H] or None + swa_window: int, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Launch the V4 CSA attention backward kernel. + + Returns ``(dq, dk_local, dv_local, dgathered, dsink)`` — gradients in + the input dtype, with ``dsink`` returned only when ``sink is not + None`` (else ``None``). + """ + if not q.is_cuda: + raise ValueError("v4_csa_attention_v0 BWD requires CUDA / HIP tensors.") + if dout.shape != out.shape or out.shape != q.shape: + raise ValueError( + "v4_csa_attention_v0 BWD shape mismatch: " + f"out={tuple(out.shape)}, dout={tuple(dout.shape)}, q={tuple(q.shape)}" + ) + + B, HQ, Sq, D = q.shape + K_topk = gathered.shape[2] + + has_sink = sink is not None + + BLOCK_N = 32 + BLOCK_K = 32 + BLOCK_DMODEL = D + + # Allocate fp32 output buffers for atomic_add. Cast to input dtype + # before returning. + dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dk_local_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dv_local_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dgathered_fp32 = torch.zeros((B, Sq, K_topk, D), device=q.device, dtype=torch.float32) + if has_sink: + dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32) + sink_arg = sink.to(torch.float32) if sink.dtype != torch.float32 else sink + else: + dsink_fp32 = q # sentinel; HAS_SINK=False inside kernel + sink_arg = q + + # D scalar = (dout * out).sum(-1) — reuse the dense module's pre-pass + d_buf = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32) + pre_grid = (triton.cdiv(Sq, BLOCK_N), B * HQ) + _v4_attention_bwd_preprocess_kernel[pre_grid]( + out, + dout, + d_buf, + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + Sq, + HEAD=HQ, + BLOCK_M=BLOCK_N, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=4, + num_stages=1, + ) + + grid = (Sq, B * HQ) + _v4_csa_attention_bwd_kernel[grid]( + q, + k_local, + v_local, + gathered, + sparse_mask, + dout, + lse, + d_buf, + dq_fp32, + dk_local_fp32, + dv_local_fp32, + dgathered_fp32, + dsink_fp32, + sink_arg, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k_local.stride(0), + k_local.stride(1), + k_local.stride(2), + k_local.stride(3), + v_local.stride(0), + v_local.stride(1), + v_local.stride(2), + v_local.stride(3), + gathered.stride(0), + gathered.stride(1), + gathered.stride(2), + gathered.stride(3), + sparse_mask.stride(0), + sparse_mask.stride(1), + sparse_mask.stride(2), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_fp32.stride(0), + dq_fp32.stride(1), + dq_fp32.stride(2), + dq_fp32.stride(3), + dk_local_fp32.stride(0), + dk_local_fp32.stride(1), + dk_local_fp32.stride(2), + dk_local_fp32.stride(3), + dv_local_fp32.stride(0), + dv_local_fp32.stride(1), + dv_local_fp32.stride(2), + dv_local_fp32.stride(3), + dgathered_fp32.stride(0), + dgathered_fp32.stride(1), + dgathered_fp32.stride(2), + dgathered_fp32.stride(3), + Sq, + K_topk, + float(scale), + HEAD_Q=HQ, + SWA_WINDOW=int(swa_window) if swa_window > 0 else 0, + HAS_SINK=has_sink, + BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=4, + num_stages=1, + ) + + dq_out = dq_fp32.to(q.dtype) + dk_local_out = dk_local_fp32.to(k_local.dtype) + dv_local_out = dv_local_fp32.to(v_local.dtype) + dgathered_out = dgathered_fp32.to(gathered.dtype) + dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None + return dq_out, dk_local_out, dv_local_out, dgathered_out, dsink_out diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_fwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_fwd.py new file mode 100644 index 000000000..a1798758a --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v0_deprecated/v4_csa_attention_fwd.py @@ -0,0 +1,407 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 CSA attention forward Triton kernel (plan-4 P26, ``compress_ratio == 4``). + +CSA fuses three branches into a single online softmax: + +* **Local SWA**: ``q @ k_local^T`` with sliding-window-causal masking. +* **Sparse top-K**: ``q . gathered[m, :, :]`` where the wrapper has + pre-gathered ``[B, Sq, K, D]`` rows from the compressed pool (the + per-query top-K gather lives outside the kernel — see plan-4 + ``02-phase-details.md`` Phase 26 design notes). +* **Per-head learned sink**: a virtual key column with notional value + zero, joined as the last softmax candidate so its probability mass + is shared across local + sparse branches. + +The kernel produces one ``[BLOCK_DMODEL]`` output row per program; the +grid is ``(seqlen_q, batch * head_q)`` so each program owns exactly one +``(b, qhid, m)`` query row. The per-row design keeps the sparse-branch +SMEM footprint inside the MI355 budget at ``head_dim=512``: the +gathered tile is only ``[BLOCK_K, head_dim] * 2 bytes ≈ 32 KiB`` per +program, while a multi-row tile would balloon to +``[BLOCK_M, BLOCK_K, head_dim] * 2 bytes ≈ 1 MiB``. + +dtype contract (matches :func:`eager_v4_csa_attention`): + +* Q / K / V / gathered are loaded in input dtype (bf16 in production); + the per-row dot products (``sum(k * q[None, :], axis=-1)``) reduce in + fp32 because we ``.to(tl.float32)`` before the multiply. +* The online-softmax accumulator (``m_running``, ``l_running``, + ``acc``) lives in fp32 — the *only* fp32 step inside the kernel. +* Output is written back in input dtype; saved ``LSE`` is fp32 (BWD + re-materialises ``P`` from it). + +Edge cases handled: + +* ``K_topk == 0`` — wrapper short-circuits to the dense + :func:`v4_attention_v1` kernel before reaching this file. +* ``topk_idx == -1`` — wrapper sets the corresponding ``sparse_mask`` + entry to ``-inf``; the kernel just adds the bias and the masked + position contributes ~0 to the softmax denominator. +* All-masked tile rows — the running max and per-tile max are both + the finite ``NEG_INF`` sentinel (``-1e30``), so + ``exp(NEG_INF - NEG_INF) = exp(0) = 1`` algebraically but the + contribution to ``acc`` and ``l_running`` is gated by the + per-element ``exp(qk - m_new)`` which stays at exactly zero for + every ``-inf``-masked entry. This avoids the ``exp(-inf - -inf) = + exp(NaN)`` failure mode that ``-float("inf")`` would have. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +# --------------------------------------------------------------------------- +# Triton kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_csa_attention_fwd_kernel( + Q, + K_LOCAL, + V_LOCAL, + GATHERED, + SPARSE_MASK, + SINK, + OUT, + LSE, + # Q strides: [B, H, Sq, D] row-major (contiguous on D) + stride_qb, + stride_qh, + stride_qm, + stride_qd, + # K_local strides: [B, H, Sq, D] row-major (CSA always has K_H == HQ — + # the V4 forward broadcast-expanded MQA single-latent KV across the H + # query heads before this call) + stride_klb, + stride_klh, + stride_kln, + stride_kld, + # V_local strides: [B, H, Sq, D] row-major + stride_vlb, + stride_vlh, + stride_vln, + stride_vld, + # gathered strides: [B, Sq, K_topk, D] row-major (no H dim — gather + # is per-query but shared across heads) + stride_gb, + stride_gm, + stride_gk, + stride_gd, + # sparse_mask strides: [B, Sq, K_topk] row-major (broadcasts over H) + stride_smb, + stride_smm, + stride_smk, + # OUT strides: [B, H, Sq, D] row-major + stride_ob, + stride_oh, + stride_om, + stride_od, + # LSE strides: [B, H, Sq] row-major + stride_lb, + stride_lh, + stride_lm, + seqlen_q, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + SWA_WINDOW: tl.constexpr, # > 0 for V4; 0 falls back to full causal + HAS_SINK: tl.constexpr, + BLOCK_N: tl.constexpr, # local-key tile size + BLOCK_K: tl.constexpr, # sparse-key tile size + BLOCK_DMODEL: tl.constexpr, # head_dim — must be a power of 2 +): + """V4 CSA fused-attention FWD (one program per output row).""" + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + bid = pid_bh // HEAD_Q + qhid = pid_bh % HEAD_Q + + offs_d = tl.arange(0, BLOCK_DMODEL) + + # ---- Load Q row [BLOCK_DMODEL] ----------------------------------------- + q_row_offset = bid * stride_qb + qhid * stride_qh + pid_m * stride_qm + q_ptrs = Q + q_row_offset + offs_d * stride_qd + q_active = pid_m < seqlen_q + q = tl.load(q_ptrs, mask=q_active, other=0.0) + + # Online-softmax running state (fp32). NEG_INF is a finite sentinel + # (-1e30) so all-masked tiles do not produce NaN through + # ``exp(-inf - -inf) = exp(NaN)``. + NEG_INF: tl.constexpr = -1.0e30 + acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32) + m_i = tl.full((), value=NEG_INF, dtype=tl.float32) + l_i = tl.zeros((), dtype=tl.float32) + + # ---- Local SWA branch -------------------------------------------------- + # Causal: keys n in [0, pid_m]. SWA: keys n in [pid_m - SWA_WINDOW + 1, + # pid_m]. We walk from the SWA window's lower bound (rounded down to + # BLOCK_N) up to (pid_m + 1). The in-kernel window check inside the + # tile loop handles the boundary cases exactly so the result matches + # ``sliding_window_causal_mask(...)``. + n_loop_end = pid_m + 1 + if n_loop_end > seqlen_q: + n_loop_end = seqlen_q + + # Lower bound of the SWA window (clamped to >= 0). When SWA_WINDOW <= 0 + # this collapses to a full causal walk from 0. + if SWA_WINDOW > 0: + n_lo_raw = pid_m - SWA_WINDOW + 1 + if n_lo_raw < 0: + n_lo_raw = 0 + # Round down to a BLOCK_N multiple so tile-aligned loads stay aligned. + n_loop_start = (n_lo_raw // BLOCK_N) * BLOCK_N + else: + n_loop_start = 0 + + for n_start in range(n_loop_start, n_loop_end, BLOCK_N): + offs_n = n_start + tl.arange(0, BLOCK_N) + + # K_local tile: [BLOCK_N, BLOCK_DMODEL] in k_local.dtype + kl_ptrs = ( + K_LOCAL + + bid * stride_klb + + qhid * stride_klh + + offs_n[:, None] * stride_kln + + offs_d[None, :] * stride_kld + ) + kl_load_mask = offs_n[:, None] < seqlen_q + kl = tl.load(kl_ptrs, mask=kl_load_mask, other=0.0) + + # qk = sum_d (kl[n, d] * q[d]) -> [BLOCK_N], computed in fp32 by + # upcasting the operands. (Matches the eager reference's + # bf16-tensor-core matmul w/ fp32 accumulator semantics; a 1xD + # tl.dot is not portable on the HIP backend, see plan-4 P26 note.) + qk = tl.sum(kl.to(tl.float32) * q[None, :].to(tl.float32), axis=1) * sm_scale + + # SWA-causal mask: keep n in [pid_m - SWA_WINDOW + 1, pid_m]. + # When SWA_WINDOW <= 0 fall back to full causal. + if SWA_WINDOW > 0: + in_window = (offs_n >= pid_m - SWA_WINDOW + 1) & (offs_n <= pid_m) + else: + in_window = offs_n <= pid_m + qk = tl.where(in_window, qk, NEG_INF) + # Boundary mask for keys past seqlen_q. + qk = tl.where(offs_n < seqlen_q, qk, NEG_INF) + + # Online softmax update (shared with sparse branch + sink). + m_tile = tl.max(qk, axis=0) + m_new = tl.maximum(m_i, m_tile) + alpha = tl.exp(m_i - m_new) + p = tl.exp(qk - m_new) + l_i = l_i * alpha + tl.sum(p, axis=0) + + # V_local tile: [BLOCK_N, BLOCK_DMODEL] in v_local.dtype + vl_ptrs = ( + V_LOCAL + + bid * stride_vlb + + qhid * stride_vlh + + offs_n[:, None] * stride_vln + + offs_d[None, :] * stride_vld + ) + vl = tl.load(vl_ptrs, mask=kl_load_mask, other=0.0) + + # acc += sum_n (p[n] * vl[n, :]) — fp32 accumulator. + acc = acc * alpha + tl.sum(p[:, None] * vl.to(tl.float32), axis=0) + m_i = m_new + + # ---- Sparse top-K branch ---------------------------------------------- + # gathered is per-query (no H dim — broadcast across heads in the + # eager reference). We walk K_topk in BLOCK_K tiles. + for k_start in range(0, K_topk, BLOCK_K): + offs_k = k_start + tl.arange(0, BLOCK_K) + + g_ptrs = ( + GATHERED + + bid * stride_gb + + pid_m * stride_gm + + offs_k[:, None] * stride_gk + + offs_d[None, :] * stride_gd + ) + g_load_mask = offs_k[:, None] < K_topk + g = tl.load(g_ptrs, mask=g_load_mask, other=0.0) + + # qk_sparse = sum_d (g[k, d] * q[d]) -> [BLOCK_K] + qk_sparse = tl.sum(g.to(tl.float32) * q[None, :].to(tl.float32), axis=1) * sm_scale + + # Caller-supplied sparse_mask: -inf for topk_idx == -1 entries. + sm_ptrs = SPARSE_MASK + bid * stride_smb + pid_m * stride_smm + offs_k * stride_smk + sm_load_mask = offs_k < K_topk + sm = tl.load(sm_ptrs, mask=sm_load_mask, other=0.0).to(tl.float32) + qk_sparse = qk_sparse + sm + + # Boundary mask for offs_k past K_topk. + qk_sparse = tl.where(offs_k < K_topk, qk_sparse, NEG_INF) + + # Online softmax update — shares m_i / l_i with the local branch. + m_tile = tl.max(qk_sparse, axis=0) + m_new = tl.maximum(m_i, m_tile) + alpha = tl.exp(m_i - m_new) + p = tl.exp(qk_sparse - m_new) + l_i = l_i * alpha + tl.sum(p, axis=0) + + # acc += sum_k (p[k] * g[k, :]) — fp32 accumulator. + acc = acc * alpha + tl.sum(p[:, None] * g.to(tl.float32), axis=0) + m_i = m_new + + # ---- Sink (joint over both branches) ---------------------------------- + if HAS_SINK: + sink_h = tl.load(SINK + qhid).to(tl.float32) + m_new = tl.maximum(m_i, sink_h) + alpha = tl.exp(m_i - m_new) + beta = tl.exp(sink_h - m_new) + l_i = l_i * alpha + beta + acc = acc * alpha + m_i = m_new + + # ---- Final divide + cast back to output dtype ------------------------- + out = acc / l_i + lse = m_i + tl.log(l_i) + + out_offset = bid * stride_ob + qhid * stride_oh + pid_m * stride_om + out_ptrs = OUT + out_offset + offs_d * stride_od + tl.store(out_ptrs, out.to(OUT.dtype.element_ty), mask=q_active) + + lse_ptr = LSE + bid * stride_lb + qhid * stride_lh + pid_m * stride_lm + tl.store(lse_ptr, lse, mask=q_active) + + +def _launch_v4_csa_attention_fwd( + q: torch.Tensor, # [B, H, Sq, D] + k_local: torch.Tensor, # [B, H, Sq, D] + v_local: torch.Tensor, # [B, H, Sq, D] + gathered: torch.Tensor, # [B, Sq, K_topk, D] + sparse_mask: torch.Tensor, # [B, Sq, K_topk] + *, + sink: Optional[torch.Tensor], # [H] or None + swa_window: int, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Launch the V4 CSA attention forward kernel. + + Returns ``(out, lse)`` where ``out`` matches ``v_local.dtype`` and + ``lse`` is fp32. ``lse`` is what the BWD kernel needs to + re-materialise the joint softmax without storing the + ``[Sq, Sq + K_topk]`` joint ``P`` matrix. + """ + if q.dim() != 4 or k_local.dim() != 4 or v_local.dim() != 4: + raise ValueError( + "v4_csa_attention_v0 forward expects q / k_local / v_local of rank 4 " + f"(got {q.dim()} / {k_local.dim()} / {v_local.dim()})" + ) + if gathered.dim() != 4: + raise ValueError( + f"v4_csa_attention_v0 forward expects gathered of rank 4 [B, Sq, K, D]; " + f"got rank {gathered.dim()}, shape {tuple(gathered.shape)}" + ) + if sparse_mask.dim() != 3: + raise ValueError( + f"v4_csa_attention_v0 forward expects sparse_mask of rank 3 [B, Sq, K]; " + f"got rank {sparse_mask.dim()}, shape {tuple(sparse_mask.shape)}" + ) + + B, HQ, Sq, D = q.shape + if k_local.shape != q.shape or v_local.shape != q.shape: + raise ValueError( + "v4_csa_attention_v0 requires k_local.shape == v_local.shape == q.shape " + f"(got q={tuple(q.shape)}, k_local={tuple(k_local.shape)}, " + f"v_local={tuple(v_local.shape)})." + ) + + Bg, Sqg, K_topk, Dg = gathered.shape + if Bg != B or Sqg != Sq or Dg != D: + raise ValueError( + "v4_csa_attention_v0 gathered shape mismatch: expected " + f"[B, Sq, K, D] = [{B}, {Sq}, *, {D}]; got {tuple(gathered.shape)}." + ) + Bm, Sqm, Km = sparse_mask.shape + if Bm != B or Sqm != Sq or Km != K_topk: + raise ValueError( + "v4_csa_attention_v0 sparse_mask shape mismatch: expected " + f"[B, Sq, K] = [{B}, {Sq}, {K_topk}]; got {tuple(sparse_mask.shape)}." + ) + + if not q.is_cuda: + raise ValueError("v4_csa_attention_v0 requires CUDA / HIP tensors.") + if q.dtype != k_local.dtype or q.dtype != v_local.dtype or q.dtype != gathered.dtype: + raise ValueError( + "v4_csa_attention_v0 requires q.dtype == k_local.dtype == v_local.dtype " + f"== gathered.dtype (got {q.dtype} / {k_local.dtype} / " + f"{v_local.dtype} / {gathered.dtype})." + ) + + has_sink = sink is not None + + out = torch.empty_like(q) + lse = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32) + + # Tile sizes: BLOCK_N / BLOCK_K conservative for SMEM at head_dim=512. + # Per-row design (one program per (b, qhid, m)) means the gathered + # tile is [BLOCK_K, D] only; multi-row would balloon SMEM. + BLOCK_N = 32 + BLOCK_K = 32 + BLOCK_DMODEL = D # head_dim must be a power of 2 for tl.arange + + grid = (Sq, B * HQ) + + # Sentinel pointer when sink is absent. Triton requires a real tensor — + # we pass q (any tensor) and gate via the constexpr. + sink_ptr = sink if has_sink else q + + _v4_csa_attention_fwd_kernel[grid]( + q, + k_local, + v_local, + gathered, + sparse_mask, + sink_ptr, + out, + lse, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k_local.stride(0), + k_local.stride(1), + k_local.stride(2), + k_local.stride(3), + v_local.stride(0), + v_local.stride(1), + v_local.stride(2), + v_local.stride(3), + gathered.stride(0), + gathered.stride(1), + gathered.stride(2), + gathered.stride(3), + sparse_mask.stride(0), + sparse_mask.stride(1), + sparse_mask.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + Sq, + K_topk, + float(scale), + HEAD_Q=HQ, + SWA_WINDOW=int(swa_window) if swa_window > 0 else 0, + HAS_SINK=has_sink, + BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=4, + num_stages=1, + ) + return out, lse diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/__init__.py new file mode 100644 index 000000000..dccccc15c --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/__init__.py @@ -0,0 +1,29 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Triton **v1** attention backend (production, separate K/V). + +The current production Triton attention kernels + autograd Functions: + +* dense (cr=0) / HCA (cr=128): :func:`v4_attention_v1` / :class:`V4AttentionFn` + (``v4_attention_fwd`` / ``v4_attention_bwd`` launchers). +* CSA (cr=4), in-kernel pool gather + scatter-add: + :func:`v4_csa_attention_v1` / :class:`V4CSAPoolAttentionFn` + (``v4_csa_attention_fwd`` / ``v4_csa_attention_bwd`` pool launchers). + +See ``_triton_v0_deprecated`` for the deprecated gathered CSA path and ``_triton_v2`` for +the fused single-latent sparse-MLA path. +""" + +from .v4_attention import V4AttentionFn, v4_attention_v1 +from .v4_csa_attention import V4CSAPoolAttentionFn, v4_csa_attention_v1 + +__all__ = [ + "v4_attention_v1", + "V4AttentionFn", + "v4_csa_attention_v1", + "V4CSAPoolAttentionFn", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/_v4_attn_tuning.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/_v4_attn_tuning.py new file mode 100644 index 000000000..be72fad73 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/_v4_attn_tuning.py @@ -0,0 +1,65 @@ +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +"""Architecture-aware tuned defaults for the V4 Triton attention kernels. + +The V4 attention fwd/bwd/CSA kernels expose ~60 ``PRIMUS_V4_ATTN_*`` / ``PRIMUS_V4_CSA_*`` +env knobs whose hard-coded defaults are the **gfx950 / MI355X** R1/R2 sweep winners. On +gfx1250 (MI450, CDNA-next) the optimum differs systematically -- larger ``BLOCK_M`` and +fewer warps/stages, and far more dKV head-split -- so a gfx950-tuned default leaves a lot on +the table. This module centralises the per-arch defaults so each launcher can ask for the +right value for the GPU it is running on; the env knobs still override everything. + +gfx1250 values come from node-safe microbench sweeps (``ab_sweep/opt7b_*`` / ``opt7c_*``); +untuned knobs/archs fall back to the historical gfx950 defaults so nothing regresses. +""" +from __future__ import annotations + +import functools + + +@functools.lru_cache(maxsize=1) +def gpu_arch() -> str: + """Lower-cased GPU gfx arch (e.g. ``gfx1250``); ``""`` if it can't be determined.""" + try: + import torch + + name = torch.cuda.get_device_properties(0).gcnArchName # e.g. "gfx1250:sramecc+:xnack-" + return name.split(":")[0].strip().lower() + except Exception: + return "" + + +def is_gfx1250() -> bool: + return gpu_arch() == "gfx1250" + + +def fwd_attn_defaults(is_hca: bool): + """``(BLOCK_M, BLOCK_N, NUM_WARPS, NUM_STAGES)`` for the V4 attention FWD kernel. + + gfx1250: ``BM=128, BN=32, W=4, S=1`` wins **both** shapes -- SWA **+62-70%** vs the gfx950 + winner ``BM=64/BN=16/W=8/S=2`` (``ab_sweep/opt7c``), and HCA (cr=128) **+15-21%** vs the + gfx950 HCA winner ``BM=128/BN=16/W=8/S=1`` (``ab_sweep/opt7d``). + """ + if is_gfx1250(): + return 128, 32, 4, 1 + # gfx950 / MI355X -- historical R2-sweep winners (shape-dependent BLOCK_M / stages). + if is_hca: + return 128, 16, 8, 1 + return 64, 16, 8, 2 + + +def bwd_dkv_head_groups_default(hq: int, hk: int) -> int: + """dKV head-split groups for the V4 attention BWD -- ONLY the MQA (HQ>=64, HK==1) path. + + gfx1250 sweep (``ab_sweep/opt7b``, MQA): ``HG=32`` is +37-53% vs the gfx950 default ``2`` + (the kernel notes "HG=4/8 regress" -- true on gfx950, inverted on gfx1250). + + NOTE: the current DeepSeek-V4-Pro attention runs **MHA at the kernel level** -- the + single-latent KV is expanded to all H heads (deepseek_v4_attention.py), so HK==HQ and the + caller's ``if HQ > HK`` guard means this function is **not reached** there (HG stays 1). + The MHA bwd is already gfx950-optimal on gfx1250 (``ab_sweep/opt7e`` -- every alt config + regresses), so there is no bwd retune for V4-Pro; this default only matters for a true-MQA + config. The real gfx1250 attention win is the FWD (see ``fwd_attn_defaults``). + """ + if not (hq >= 64 and hk == 1): + return 1 + return 32 if is_gfx1250() else 2 diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention.py new file mode 100644 index 000000000..3acbc6493 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention.py @@ -0,0 +1,298 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 attention Triton autograd entry point (plan-4 P25). + +Public API: + +* :func:`v4_attention_v1` — functional API matching + :func:`eager_v4_attention`'s signature; routes through + :class:`V4AttentionFn` so autograd works. +* :class:`V4AttentionFn` — :class:`torch.autograd.Function` wrapping + the Triton FWD + Triton BWD (re-materialises softmax from the saved + LSE; sink gradient atomic-added per query head). + +Dispatch contract (consumed by ``DeepseekV4Attention.forward``): + +* ``compress_ratio == 0``: caller passes ``swa_window > 0`` and + ``additive_mask=None`` so the kernel applies the SWA-causal mask + in-place. +* ``compress_ratio == 128`` (HCA): caller pre-concatenates pool keys + to local keys, passes the pool-only ``[Sq, P]`` additive mask, sets + ``hca_local_seqlen=Sq``, and keeps ``swa_window > 0``. The kernel + splits the loop into a pruned local SWA branch plus the compressed-pool + suffix. + +dtype contract (must match :func:`eager_v4_attention`): + +* Q / K / V matmuls run on tensor cores in input dtype (bf16 in + production); the matmul accumulator inside is fp32. +* The online-softmax accumulator (``m_running / l_running / acc``) is + fp32 — the *only* fp32 step. +* Output is in ``v.dtype``. +* Saved LSE is fp32 (BWD re-materialises ``P`` from it). +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from primus.backends.megatron.core.transformer.v4_attention_kernels import ( + _flydsl_v0_deprecated as _flydsl, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels import _tilelang +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import ( + _launch_v4_attention_bwd, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import ( + _launch_v4_attention_fwd, +) + +# --------------------------------------------------------------------------- +# Autograd Function +# --------------------------------------------------------------------------- + + +class V4AttentionFn(torch.autograd.Function): + """Triton FWD + (eager-recompute BWD until P25-stage-2 lands).""" + + @staticmethod + def forward( # type: ignore[override] + ctx, + q: torch.Tensor, # [B, H, Sq, D] + k: torch.Tensor, # [B, K_H, Sk, D] + v: torch.Tensor, # [B, K_H, Sk, D] + sink: Optional[torch.Tensor], # [H] or None + additive_mask: Optional[torch.Tensor], # [Sq, Sk] or None + swa_window: int, + attn_dropout: float, + training: bool, + scale: float, + hca_local_seqlen: int, + ) -> torch.Tensor: + if attn_dropout > 0.0 and training: + # Plan-4 P25 does not implement dropout in the kernel — V4 + # is trained with attn_dropout=0 so this branch is unreachable + # in production. We refuse explicitly so a stray non-zero + # dropout configuration raises rather than silently dropping + # the kernel path. + raise NotImplementedError( + "v4_attention_v1 does not implement in-kernel attention " + "dropout (V4 trains with attn_dropout=0). Got " + f"attn_dropout={attn_dropout}, training={training}." + ) + + out, lse = _launch_v4_attention_fwd( + q, + k, + v, + sink=sink, + swa_window=swa_window, + additive_mask=additive_mask, + scale=scale, + hca_local_seqlen=hca_local_seqlen, + ) + # Save tensors the BWD kernel needs. None-typed args + # (``sink`` / ``additive_mask``) are stashed on ``ctx`` because + # ``save_for_backward`` does not accept ``None``. + ctx.save_for_backward(q, k, v, out, lse, sink, additive_mask) + ctx.swa_window = int(swa_window) + ctx.attn_dropout = float(attn_dropout) + ctx.training_mode = bool(training) + ctx.scale = float(scale) + ctx.hca_local_seqlen = int(hca_local_seqlen) + ctx.sink_was_none = sink is None + ctx.mask_was_none = additive_mask is None + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): # type: ignore[override] + """Triton BWD: re-materialises ``P`` from saved LSE; emits dq, dk, dv, dsink. + + Routes through :func:`_launch_v4_attention_bwd`, which: + + * runs a small fp32 pre-pass that computes the per-query + ``D = (dout * out).sum(-1)`` scalar, + * runs the main BWD kernel parallelized over query blocks, + accumulating ``dQ`` in registers and atomic-adding into + ``dK / dV / dsink``, + * casts the fp32 accumulator buffers back to the input dtype + for return. + + The Function returns gradients in the same positional order as + the forward: ``(q, k, v, sink, additive_mask, swa_window, + attn_dropout, training, scale)``. Non-tensor / non-grad inputs + get ``None``. + """ + q, k, v, out, lse, sink, additive_mask = ctx.saved_tensors + + sink_arg = None if ctx.sink_was_none else sink + mask_arg = None if ctx.mask_was_none else additive_mask + + # Ensure dout is contiguous in the [B, H, Sq, D] layout the + # kernel expects. The kernel reads strides from the tensor so + # any contiguous-in-its-strides tensor would work, but a + # ``.contiguous()`` here keeps the access pattern simple. + if not grad_out.is_contiguous(): + grad_out = grad_out.contiguous() + + dq, dk, dv, dsink = _launch_v4_attention_bwd( + q, + k, + v, + out, + grad_out, + lse, + sink=sink_arg, + swa_window=ctx.swa_window, + additive_mask=mask_arg, + scale=ctx.scale, + hca_local_seqlen=ctx.hca_local_seqlen, + ) + + # Honor needs_input_grad: zero-cost ``None`` for inputs that + # don't want gradients. (We still computed them — the main + # cost in the kernel is the matmul, not the per-output cast — + # so this is purely a cleanliness touch.) + if not ctx.needs_input_grad[0]: + dq = None + if not ctx.needs_input_grad[1]: + dk = None + if not ctx.needs_input_grad[2]: + dv = None + if not ctx.needs_input_grad[3] or ctx.sink_was_none: + dsink = None + + # Forward signature: (q, k, v, sink, additive_mask, swa_window, + # attn_dropout, training, scale, hca_local_seqlen) + return dq, dk, dv, dsink, None, None, None, None, None, None + + +# --------------------------------------------------------------------------- +# Functional API +# --------------------------------------------------------------------------- + + +def v4_attention_v1( + q: torch.Tensor, # [B, H, Sq, D] + k: torch.Tensor, # [B, K_H, Sk, D] K_H ∈ {1, H} + v: torch.Tensor, # [B, K_H, Sk, D] + *, + sink: Optional[torch.Tensor], # [H] or None + swa_window: int, + additive_mask: Optional[torch.Tensor], # [Sq, Sk] or None + attn_dropout: float, + training: bool, + scale: float, + hca_local_seqlen: int = 0, + use_tilelang: bool = False, + use_flydsl: bool = False, +) -> torch.Tensor: + """Triton- or tilelang-backed V4 dense / HCA attention. + + Drop-in replacement for :func:`eager_v4_attention` with identical + signature and dtype contract. + + Dispatch precedence (plan-8 P57 close-out 2): + + * If ``use_tilelang=True`` AND the plan-8 P50 / P51 tilelang + kernels are registered + tilelang importable at the pinned + version, route through the tilelang FWD/BWD. + * Otherwise, route through :class:`V4AttentionFn` (plan-4 P25 + Triton FWD + plan-5 P32 final split BWD). + + ``use_tilelang`` is plumbed by ``DeepseekV4Attention.forward`` + from the ``use_v4_tilelang_attention`` config flag. Default-False + so containers without tilelang installed never trigger any + tilelang import. + + The MQA case (``K_H == 1``) is detected from ``k.shape[1]`` and + each kernel internally broadcasts the single shared K / V head + across the query heads. + + Returns ``[B, H, Sq, D]`` in ``v.dtype``. + """ + # Plan-8 P49 / P57 close-out 2: tilelang dispatcher hook. + # ``should_dispatch`` short-circuits on ``enabled=False`` so the + # tilelang import never fires on a container that does not have + # the package installed. + if _tilelang.should_dispatch("v4_attention_fwd", enabled=use_tilelang): + # P51: route through the tilelang autograd Function when both + # FWD + BWD kernels are registered, so gradients flow through + # the tilelang path correctly. When only the FWD is registered + # (P50-only state), fall through to the FWD-direct call — the + # Triton BWD will still get called via the autograd graph that + # the caller sets up (only the FWD wrapper is non-autograd in + # that mode). + # `_lazy_load` is idempotent + cheap on cache hit; calling it + # here ensures `is_tilelang_kernel_available("v4_attention_bwd")` + # reflects the post-import state. + _tilelang._lazy_load("v4_attention_bwd") + if _tilelang.is_tilelang_kernel_available("v4_attention_bwd"): + from primus.backends.megatron.core.transformer.v4_attention_kernels._tilelang.v4_attention_autograd_tilelang import ( + v4_attention_tilelang, + ) + + return v4_attention_tilelang( + q, + k, + v, + sink=sink, + additive_mask=additive_mask, + swa_window=swa_window, + attn_dropout=attn_dropout, + training=training, + scale=scale, + hca_local_seqlen=hca_local_seqlen, + ) + return _tilelang.v4_attention_fwd_tilelang( + q, + k, + v, + sink=sink, + additive_mask=additive_mask, + swa_window=swa_window, + attn_dropout=attn_dropout, + training=training, + scale=scale, + hca_local_seqlen=hca_local_seqlen, + ) + # FlyDSL backend hook (forward-only; soft-dep, default off). Mirrors the + # _tilelang hook: short-circuits on enabled=False so FlyDSL is never + # imported on the common path, and falls back to Triton if the runtime + # or the kernel is unavailable. Training (autograd) still flows through + # V4AttentionFn below; this path is for inference/eval. + if _flydsl.should_dispatch("v4_attention_fwd", enabled=use_flydsl): + return _flydsl.v4_attention_fwd_flydsl( + q, + k, + v, + sink=sink, + swa_window=swa_window, + additive_mask=additive_mask, + scale=scale, + hca_local_seqlen=hca_local_seqlen, + ) + return V4AttentionFn.apply( + q, + k, + v, + sink, + additive_mask, + swa_window, + attn_dropout, + training, + scale, + hca_local_seqlen, + ) + + +__all__ = [ + "V4AttentionFn", + "v4_attention_v1", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_bwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_bwd.py new file mode 100644 index 000000000..ff8072ef9 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_bwd.py @@ -0,0 +1,2086 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 attention backward Triton kernel (plan-4 P25, ``compress_ratio in {0, 128}``). + +Two-kernel design: + +* :func:`_v4_attention_bwd_preprocess_kernel` — computes the per-query + diagonal scalar ``D = sum_d (dout[..., d] * out[..., d])`` in fp32. + This is the standard FlashAttention-2 BWD pre-pass; it lets the main + kernel skip materialising ``dP @ V^T`` as a tile-wise sum. +* :func:`_v4_attention_bwd_kernel` — main BWD pass. Parallelizes over + query blocks (one program per ``m_block × batch × head_q``) and: + + 1. re-materialises ``P = exp(qk - LSE)`` in fp32 from saved + ``Q / K / V / LSE`` (so the [Sq, Sk] ``P`` matrix is never + stored), + 2. computes ``dS = P * (dP - D)`` where ``dP = dout @ V^T``, + 3. accumulates ``dQ`` in registers (no atomic — one program per + m-block writes its dQ tile straight to global), + 4. atomic-adds ``dK = scale * dS^T @ Q`` and ``dV = P^T @ dout`` + into the global dK / dV buffers, + 5. atomic-adds the sink gradient + ``dsink_h += -sum_t (P_sink_t * D_t)`` per query head. + +The launcher allocates ``dQ / dK / dV / dsink`` as fp32 buffers (so +``tl.atomic_add`` works regardless of input dtype) and casts back to +the input dtype on return. + +dtype contract: + +* All matmuls run on tensor cores in input dtype with fp32 accumulator. +* The online ``P / dP / dS`` re-materialisation is fp32 (matches the + FWD's softmax-in-fp32 contract). +* Output gradients are returned in input dtype (cast from fp32 + buffers). + +Tile choice: ``BLOCK_M = BLOCK_N = 32`` at ``head_dim = 512`` so the +peak SMEM (Q + K + V + dout = 4 × 32 × 512 × 2 = 128 KiB) fits under +MI355's 160 KiB SMEM budget with some headroom for register pressure. +P25 perf follow-up may explore ``BLOCK_M = BLOCK_N = 64`` once the +correctness baseline is locked. +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch +import triton +import triton.language as tl + +# --------------------------------------------------------------------------- +# Pre-pass kernel (D scalar) +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_attention_bwd_preprocess_kernel( + OUT, + DOUT, + D, + stride_ob, + stride_oh, + stride_om, + stride_od, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_db, + stride_dh, + stride_dm, + seqlen_q, + HEAD: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """Compute ``D[b, h, m] = sum_d (dout[b, h, m, d] * out[b, h, m, d])``.""" + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + bid = pid_bh // HEAD + hid = pid_bh % HEAD + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_DMODEL) + + out_ptrs = ( + OUT + bid * stride_ob + hid * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od + ) + dout_ptrs = ( + DOUT + + bid * stride_dob + + hid * stride_doh + + offs_m[:, None] * stride_dom + + offs_d[None, :] * stride_dod + ) + mask = offs_m[:, None] < seqlen_q + o = tl.load(out_ptrs, mask=mask, other=0.0).to(tl.float32) + do = tl.load(dout_ptrs, mask=mask, other=0.0).to(tl.float32) + d = tl.sum(o * do, 1) + + d_ptrs = D + bid * stride_db + hid * stride_dh + offs_m * stride_dm + tl.store(d_ptrs, d, mask=offs_m < seqlen_q) + + +# --------------------------------------------------------------------------- +# Main BWD kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_attention_bwd_kernel( + Q, + K, + V, + DOUT, + LSE, + D, + DQ, # fp32 buffer [B, H, Sq, D] + DK, # fp32 buffer [B, K_H, Sk, D] + DV, # fp32 buffer [B, K_H, Sk, D] + DSINK, # fp32 buffer [H] or sentinel + SINK, # [H] fp32 or sentinel + ADD_MASK, # [Sq, Sk] or sentinel + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_kb, + stride_kh, + stride_kn, + stride_kd, + stride_vb, + stride_vh, + stride_vn, + stride_vd, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dqb, + stride_dqh, + stride_dqm, + stride_dqd, + stride_dkb, + stride_dkh, + stride_dkn, + stride_dkd, + stride_dvb, + stride_dvh, + stride_dvn, + stride_dvd, + stride_ms, + stride_mn, + seqlen_q, + seqlen_k, + sm_scale, + HEAD_Q: tl.constexpr, + HEAD_K: tl.constexpr, + SWA_WINDOW: tl.constexpr, + HAS_SINK: tl.constexpr, + HAS_ADD_MASK: tl.constexpr, + HCA_LOCAL_SEQLEN: tl.constexpr, + USE_CAUSAL: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """V4 attention BWD (single kernel, parallelize over m-blocks).""" + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + bid = pid_bh // HEAD_Q + qhid = pid_bh % HEAD_Q + if HEAD_K == HEAD_Q: + khid = qhid + else: + khid = 0 + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_DMODEL) + + NEG_INF: tl.constexpr = -1.0e30 + + # Load Q, dout, LSE, D for this m-block. + q_ptrs = ( + Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd + ) + dout_ptrs = ( + DOUT + + bid * stride_dob + + qhid * stride_doh + + offs_m[:, None] * stride_dom + + offs_d[None, :] * stride_dod + ) + lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm + dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm + + q_load_mask = offs_m[:, None] < seqlen_q + q = tl.load(q_ptrs, mask=q_load_mask, other=0.0) + dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0) + lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0) + dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0) + + # dQ accumulator (fp32, kept in registers across the n-loop) + dq = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + + # Sink gradient — accumulate per-query contribution then atomic_add + # into DSINK[qhid] at the end. The sink contributes + # ``dsink_h += sum_t -P_sink_t * D_t`` because dlogits_sink = P_sink * (0 - D). + if HAS_SINK: + sink_h = tl.load(SINK + qhid).to(tl.float32) + # P_sink for each query in the m-block, masked for boundary rows + p_sink = tl.exp(sink_h - lse) + p_sink_masked = tl.where(offs_m < seqlen_q, p_sink, 0.0) + dvec_masked = tl.where(offs_m < seqlen_q, dvec, 0.0) + dsink_contrib = tl.sum(-p_sink_masked * dvec_masked) + tl.atomic_add(DSINK + qhid, dsink_contrib) + + # Determine n-loop bounds (matches FWD's P30 SWA tile pruning). + n_loop_start = 0 + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0: + n_loop_end = seqlen_k + elif SWA_WINDOW > 0: + n_loop_start = pid_m * BLOCK_M - SWA_WINDOW + 1 + if n_loop_start < 0: + n_loop_start = 0 + n_loop_start = (n_loop_start // BLOCK_N) * BLOCK_N + n_loop_end = (pid_m + 1) * BLOCK_M + local_end = HCA_LOCAL_SEQLEN if HCA_LOCAL_SEQLEN > 0 else seqlen_k + if n_loop_end > local_end: + n_loop_end = local_end + elif USE_CAUSAL: + n_loop_end = (pid_m + 1) * BLOCK_M + if n_loop_end > seqlen_k: + n_loop_end = seqlen_k + else: + n_loop_end = seqlen_k + + for n_start in range(n_loop_start, n_loop_end, BLOCK_N): + offs_n = n_start + tl.arange(0, BLOCK_N) + + k_ptrs = ( + K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd + ) + v_ptrs = ( + V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd + ) + kv_load_mask = offs_n[:, None] < seqlen_k + k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0) + v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0) + + # qk = Q @ K.T * scale + mask (re-materialise in fp32) + qk = tl.dot(q, tl.trans(k)) * sm_scale + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0: + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + offs_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + else: + if SWA_WINDOW > 0: + in_window = (offs_n[None, :] >= offs_m[:, None] - SWA_WINDOW + 1) & ( + offs_n[None, :] <= offs_m[:, None] + ) + qk = tl.where(in_window, qk, NEG_INF) + elif USE_CAUSAL: + qk = tl.where(offs_n[None, :] <= offs_m[:, None], qk, NEG_INF) + qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF) + qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF) + + # P = exp(qk - LSE) in fp32. For boundary rows / fully-masked + # rows, lse is 0 (loaded with mask) and qk is NEG_INF, so + # exp(NEG_INF - 0) = 0 — no contribution. ✓ + p = tl.exp(qk - lse[:, None]) + + # dP = dout @ V.T (fp32 accumulator) + dp = tl.dot(dout, tl.trans(v)) + + # dS = P * (dP - D) + ds = p * (dp - dvec[:, None]) + + # dQ += dS @ K * scale + dq += tl.dot(ds.to(k.dtype), k) * sm_scale + + # dK += scale * dS.T @ Q (atomic_add into fp32 DK buffer) + dk_contrib = tl.dot(tl.trans(ds.to(q.dtype)), q) * sm_scale + dk_ptrs = ( + DK + + bid * stride_dkb + + khid * stride_dkh + + offs_n[:, None] * stride_dkn + + offs_d[None, :] * stride_dkd + ) + dk_mask = offs_n[:, None] < seqlen_k + tl.atomic_add(dk_ptrs, dk_contrib, mask=dk_mask, sem="relaxed") + + # dV += P.T @ dout (atomic_add into fp32 DV buffer) + dv_contrib = tl.dot(tl.trans(p.to(dout.dtype)), dout) + dv_ptrs = ( + DV + + bid * stride_dvb + + khid * stride_dvh + + offs_n[:, None] * stride_dvn + + offs_d[None, :] * stride_dvd + ) + dv_mask = offs_n[:, None] < seqlen_k + tl.atomic_add(dv_ptrs, dv_contrib, mask=dv_mask, sem="relaxed") + + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN > 0: + for n_start in range(HCA_LOCAL_SEQLEN, seqlen_k, BLOCK_N): + offs_n = n_start + tl.arange(0, BLOCK_N) + pool_n = offs_n - HCA_LOCAL_SEQLEN + + k_ptrs = ( + K + + bid * stride_kb + + khid * stride_kh + + offs_n[:, None] * stride_kn + + offs_d[None, :] * stride_kd + ) + v_ptrs = ( + V + + bid * stride_vb + + khid * stride_vh + + offs_n[:, None] * stride_vn + + offs_d[None, :] * stride_vd + ) + kv_load_mask = offs_n[:, None] < seqlen_k + k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0) + v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0) + + qk = tl.dot(q, tl.trans(k)) * sm_scale + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF) + qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout, tl.trans(v)) + ds = p * (dp - dvec[:, None]) + + dq += tl.dot(ds.to(k.dtype), k) * sm_scale + + dk_contrib = tl.dot(tl.trans(ds.to(q.dtype)), q) * sm_scale + dk_ptrs = ( + DK + + bid * stride_dkb + + khid * stride_dkh + + offs_n[:, None] * stride_dkn + + offs_d[None, :] * stride_dkd + ) + dk_mask = offs_n[:, None] < seqlen_k + tl.atomic_add(dk_ptrs, dk_contrib, mask=dk_mask, sem="relaxed") + + dv_contrib = tl.dot(tl.trans(p.to(dout.dtype)), dout) + dv_ptrs = ( + DV + + bid * stride_dvb + + khid * stride_dvh + + offs_n[:, None] * stride_dvn + + offs_d[None, :] * stride_dvd + ) + dv_mask = offs_n[:, None] < seqlen_k + tl.atomic_add(dv_ptrs, dv_contrib, mask=dv_mask, sem="relaxed") + + # Store dQ (fp32 buffer; launcher casts back to input dtype on return). + dq_ptrs = ( + DQ + + bid * stride_dqb + + qhid * stride_dqh + + offs_m[:, None] * stride_dqm + + offs_d[None, :] * stride_dqd + ) + tl.store(dq_ptrs, dq, mask=offs_m[:, None] < seqlen_q) + + +# --------------------------------------------------------------------------- +# Plan-5 P32 split-kernel BWD — dQ kernel + dK/dV kernel (no atomics) +# +# The monolithic ``_v4_attention_bwd_kernel`` parallelises over m-blocks +# and ``tl.atomic_add``s into ``dK / dV``. With ``H=64`` heads × +# ``SWA_WINDOW / BLOCK_M = 4`` m-blocks per K position, every K position +# is touched by ~256 concurrent atomics; even at "relaxed" semantics on +# MI355 these serialise through the L2 atomic engine. +# +# P32 splits the BWD into two kernels that each write their own output +# slice with NO atomics: +# +# * ``_v4_attention_bwd_dq_kernel`` — one program per ``(b, qhid, m_block)``; +# accumulates ``dQ`` in registers (same as the monolithic), but also +# handles the per-head sink gradient (``dsink``) which only needs the +# saved ``LSE`` and ``D``. Writes ``dQ`` straight to global; ``dsink`` +# is atomic-added (single counter per head — contention there is +# negligible). +# * ``_v4_attention_bwd_dkv_kernel`` — one program per +# ``(b, khid, n_block)``; iterates the *m*-blocks that contribute to +# this n-block (pruned by SWA / causal bounds), accumulates ``dK`` / +# ``dV`` in registers and writes straight to global. No atomics. +# +# Total compute roughly doubles (``Q @ K.T`` re-materialised once per +# kernel instead of shared) but the atomic pressure goes to zero, which +# is the dominant cost on MI355 at ``H=64`` × ``BLOCK_M=BLOCK_N=32``. +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_attention_bwd_dq_kernel( + Q, + K, + V, + DOUT, + LSE, + D, + DQ, + DSINK, + SINK, + ADD_MASK, + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_kb, + stride_kh, + stride_kn, + stride_kd, + stride_vb, + stride_vh, + stride_vn, + stride_vd, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dqb, + stride_dqh, + stride_dqm, + stride_dqd, + stride_ms, + stride_mn, + seqlen_q, + seqlen_k, + sm_scale, + HEAD_Q: tl.constexpr, + HEAD_K: tl.constexpr, + SWA_WINDOW: tl.constexpr, + HAS_SINK: tl.constexpr, + HAS_ADD_MASK: tl.constexpr, + HCA_LOCAL_SEQLEN: tl.constexpr, + USE_CAUSAL: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + ACCUMULATE: tl.constexpr = False, + EXACT_TILES_M: tl.constexpr = False, + EXACT_TILES_N: tl.constexpr = False, +): + """V4 attention BWD — dQ only (parallel over m-blocks, no atomics for dQ). + + When ``ACCUMULATE`` is True, the kernel performs ``dQ += dq`` instead + of ``dQ = dq``. Each ``(b, qhid, m_block)`` program owns a unique + output slice, so the implicit read-modify-write is race-free across + programs within a single launch, and sequential launches of this + kernel against the same buffer are also race-free. + """ + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + bid = pid_bh // HEAD_Q + qhid = pid_bh % HEAD_Q + if HEAD_K == HEAD_Q: + khid = qhid + else: + khid = 0 + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_DMODEL) + + NEG_INF: tl.constexpr = -1.0e30 + + q_ptrs = ( + Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd + ) + dout_ptrs = ( + DOUT + + bid * stride_dob + + qhid * stride_doh + + offs_m[:, None] * stride_dom + + offs_d[None, :] * stride_dod + ) + lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm + dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm + + q_load_mask = offs_m[:, None] < seqlen_q + q = tl.load(q_ptrs, mask=q_load_mask, other=0.0) + dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0) + lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0) + dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0) + + dq = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + + if HAS_SINK: + sink_h = tl.load(SINK + qhid).to(tl.float32) + p_sink = tl.exp(sink_h - lse) + p_sink_masked = tl.where(offs_m < seqlen_q, p_sink, 0.0) + dvec_masked = tl.where(offs_m < seqlen_q, dvec, 0.0) + dsink_contrib = tl.sum(-p_sink_masked * dvec_masked) + tl.atomic_add(DSINK + qhid, dsink_contrib) + + n_loop_start = 0 + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0: + n_loop_end = seqlen_k + elif SWA_WINDOW > 0: + n_loop_start = pid_m * BLOCK_M - SWA_WINDOW + 1 + if n_loop_start < 0: + n_loop_start = 0 + n_loop_start = (n_loop_start // BLOCK_N) * BLOCK_N + n_loop_end = (pid_m + 1) * BLOCK_M + local_end = HCA_LOCAL_SEQLEN if HCA_LOCAL_SEQLEN > 0 else seqlen_k + if n_loop_end > local_end: + n_loop_end = local_end + elif USE_CAUSAL: + n_loop_end = (pid_m + 1) * BLOCK_M + if n_loop_end > seqlen_k: + n_loop_end = seqlen_k + else: + n_loop_end = seqlen_k + + for n_start in range(n_loop_start, n_loop_end, BLOCK_N): + offs_n = n_start + tl.arange(0, BLOCK_N) + + k_ptrs = ( + K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd + ) + v_ptrs = ( + V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd + ) + kv_load_mask = offs_n[:, None] < seqlen_k + k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0) + v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0) + + qk = tl.dot(q, tl.trans(k)) * sm_scale + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0: + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + offs_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + else: + if SWA_WINDOW > 0: + in_window = (offs_n[None, :] >= offs_m[:, None] - SWA_WINDOW + 1) & ( + offs_n[None, :] <= offs_m[:, None] + ) + qk = tl.where(in_window, qk, NEG_INF) + elif USE_CAUSAL: + qk = tl.where(offs_n[None, :] <= offs_m[:, None], qk, NEG_INF) + # Plan-8 P57: EXACT_TILES_* skip the boundary masks when the + # launcher confirms ``seqlen_q % BLOCK_M == 0`` (resp. seqlen_k). + if not EXACT_TILES_N: + qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF) + if not EXACT_TILES_M: + qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout, tl.trans(v)) + ds = p * (dp - dvec[:, None]) + # P57 cr=0 BWD: defer ``sm_scale`` to a single multiply after + # the n-loop, and fuse the inner ``dq += dot(ds, k)`` into an + # MFMA-acc form via ``tl.dot(..., acc=dq)``. Numerics are + # bit-identical modulo fp32 associativity. + dq = tl.dot(ds.to(k.dtype), k, acc=dq) + + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN > 0: + for n_start in range(HCA_LOCAL_SEQLEN, seqlen_k, BLOCK_N): + offs_n = n_start + tl.arange(0, BLOCK_N) + pool_n = offs_n - HCA_LOCAL_SEQLEN + + k_ptrs = ( + K + + bid * stride_kb + + khid * stride_kh + + offs_n[:, None] * stride_kn + + offs_d[None, :] * stride_kd + ) + v_ptrs = ( + V + + bid * stride_vb + + khid * stride_vh + + offs_n[:, None] * stride_vn + + offs_d[None, :] * stride_vd + ) + kv_load_mask = offs_n[:, None] < seqlen_k + k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0) + v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0) + + qk = tl.dot(q, tl.trans(k)) * sm_scale + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + if not EXACT_TILES_N: + qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF) + if not EXACT_TILES_M: + qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout, tl.trans(v)) + ds = p * (dp - dvec[:, None]) + dq = tl.dot(ds.to(k.dtype), k, acc=dq) + + # P57 cr=0 BWD: fold ``sm_scale`` once after both loops. + dq = dq * sm_scale + + dq_ptrs = ( + DQ + + bid * stride_dqb + + qhid * stride_dqh + + offs_m[:, None] * stride_dqm + + offs_d[None, :] * stride_dqd + ) + if ACCUMULATE: + dq_prev = tl.load(dq_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0) + tl.store(dq_ptrs, dq + dq_prev, mask=offs_m[:, None] < seqlen_q) + else: + tl.store(dq_ptrs, dq, mask=offs_m[:, None] < seqlen_q) + + +@triton.jit +def _v4_attention_bwd_dkv_kernel( + Q, + K, + V, + DOUT, + LSE, + D, + DK, # fp32 buffer [B, K_H, Sk, D] + DV, # fp32 buffer [B, K_H, Sk, D] + ADD_MASK, + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_kb, + stride_kh, + stride_kn, + stride_kd, + stride_vb, + stride_vh, + stride_vn, + stride_vd, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dkb, + stride_dkh, + stride_dkn, + stride_dkd, + stride_dvb, + stride_dvh, + stride_dvn, + stride_dvd, + stride_ms, + stride_mn, + seqlen_q, + seqlen_k, + sm_scale, + HEAD_Q: tl.constexpr, + HEAD_K: tl.constexpr, + SWA_WINDOW: tl.constexpr, + HAS_ADD_MASK: tl.constexpr, + HCA_LOCAL_SEQLEN: tl.constexpr, + USE_CAUSAL: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + # Pool-suffix flag: when HCA split-mask, this program is one of + # the (HEAD_Q-block over n in pool range). The pool branch iterates + # all m's, so we run it inside the same kernel after the local + # branch finishes. + ATOMIC_REDUCE: tl.constexpr = False, + # Plan-8 P57 — MQA head-split parallelism. When ``NUM_HEAD_GROUPS > 1`` + # and ``HEAD_K != HEAD_Q`` (MQA), the kernel adds an extra grid dim + # ``pid_h_group = program_id(2)`` and each program owns + # ``HEAD_Q / NUM_HEAD_GROUPS`` query heads. Multiple head_group + # programs collide on the same MQA dK / dV slice, so writes are + # ``tl.atomic_add`` instead of ``tl.store``. + NUM_HEAD_GROUPS: tl.constexpr = 1, + EXACT_TILES_M: tl.constexpr = False, + EXACT_TILES_N: tl.constexpr = False, +): + """V4 attention BWD — dK / dV only (parallel over n-blocks, no atomics for dK/dV). + + For MQA (``HEAD_K == 1``) every query head contributes to the same + shared K / V, so this kernel must iterate ``H`` query heads per + ``(b, n_block)``. We expose that via the ``HEAD_Q`` constexpr loop. + + HCA pool n-blocks (``n_block_start >= HCA_LOCAL_SEQLEN``) are handled + by :func:`_v4_attention_bwd_dkv_pool_kernel`, which parallelises the + pool work across ``(m_block, b * qhid)``. We early-return here for + those blocks so we do not double-count the pool contribution. + """ + pid_n = tl.program_id(0) + pid_bh = tl.program_id(1) + pid_h_group = tl.program_id(2) + bid = pid_bh // HEAD_K + khid = pid_bh % HEAD_K + + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, BLOCK_DMODEL) + + NEG_INF: tl.constexpr = -1.0e30 + + # HCA mode: pool n-blocks are handled by the dedicated pool kernel. + if HCA_LOCAL_SEQLEN > 0: + if pid_n * BLOCK_N >= HCA_LOCAL_SEQLEN: + return + is_pool_block = False + else: + is_pool_block = False + + # Load K, V tiles for this n-block once. + k_ptrs = ( + K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd + ) + v_ptrs = ( + V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd + ) + kv_load_mask = offs_n[:, None] < seqlen_k + k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0) + v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0) + # R2: hoist tl.trans outside the m-loop. Lets the JIT keep kt / vt + # in registers across iterations instead of re-transposing every step. + kt = tl.trans(k) + vt = tl.trans(v) + + dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) + dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) + + # Determine m-loop bounds for this n-block. SWA / causal restricts the + # range of m's that can see this n-block. For ``HAS_ADD_MASK`` with + # arbitrary additive bias we iterate the full m axis. + n_block_lo = pid_n * BLOCK_N + n_block_hi = n_block_lo + BLOCK_N + + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0: + m_loop_start = 0 + m_loop_end = seqlen_q + elif is_pool_block: + # Every m sees every visible pool slot; pool mask drives the rest. + m_loop_start = 0 + m_loop_end = seqlen_q + elif SWA_WINDOW > 0: + # n is seen by m where n_block_lo <= m <= n_block_hi + SWA_WINDOW - 1. + # Causal also requires m >= n_block_lo. Round to BLOCK_M tiles. + m_loop_start = (n_block_lo // BLOCK_M) * BLOCK_M + m_loop_end = n_block_hi + SWA_WINDOW - 1 + if m_loop_end > seqlen_q: + m_loop_end = seqlen_q + # Round up to BLOCK_M boundary so the m-loop iterates whole tiles. + m_loop_end = ((m_loop_end + BLOCK_M - 1) // BLOCK_M) * BLOCK_M + elif USE_CAUSAL: + m_loop_start = (n_block_lo // BLOCK_M) * BLOCK_M + m_loop_end = seqlen_q + m_loop_end = ((m_loop_end + BLOCK_M - 1) // BLOCK_M) * BLOCK_M + else: + m_loop_start = 0 + m_loop_end = seqlen_q + m_loop_end = ((m_loop_end + BLOCK_M - 1) // BLOCK_M) * BLOCK_M + + # MHA path (HEAD_K == HEAD_Q) — only one query head contributes to + # this khid, so we use ``qhid = khid`` and skip the head loop. For + # MQA (HEAD_K != HEAD_Q) we iterate every query head. + if HEAD_K == HEAD_Q: + qhid = khid + for m_start in range(m_loop_start, m_loop_end, BLOCK_M): + offs_m = m_start + tl.arange(0, BLOCK_M) + + q_ptrs = ( + Q + + bid * stride_qb + + qhid * stride_qh + + offs_m[:, None] * stride_qm + + offs_d[None, :] * stride_qd + ) + dout_ptrs = ( + DOUT + + bid * stride_dob + + qhid * stride_doh + + offs_m[:, None] * stride_dom + + offs_d[None, :] * stride_dod + ) + lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm + dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm + + q_load_mask = offs_m[:, None] < seqlen_q + q = tl.load(q_ptrs, mask=q_load_mask, other=0.0) + dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0) + lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0) + dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0) + + qk = tl.dot(q, kt) * sm_scale + + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0: + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + offs_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + elif is_pool_block: + pool_n = offs_n - HCA_LOCAL_SEQLEN + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + else: + if SWA_WINDOW > 0: + in_window = (offs_n[None, :] >= offs_m[:, None] - SWA_WINDOW + 1) & ( + offs_n[None, :] <= offs_m[:, None] + ) + qk = tl.where(in_window, qk, NEG_INF) + elif USE_CAUSAL: + qk = tl.where(offs_n[None, :] <= offs_m[:, None], qk, NEG_INF) + + if not EXACT_TILES_N: + qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF) + if not EXACT_TILES_M: + qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout, vt) + ds = p * (dp - dvec[:, None]) + + # P57 cr=0 BWD: fuse the inner ``dv += dot(p, dout)`` and + # ``dk += dot(ds, q)`` into MFMA-acc form; defer ``sm_scale`` + # on dk to a single multiply after the m-loop. + dv = tl.dot(tl.trans(p.to(dout.dtype)), dout, acc=dv) + dk = tl.dot(tl.trans(ds.to(q.dtype)), q, acc=dk) + else: + # MQA path. With NUM_HEAD_GROUPS > 1 the program owns a slice of + # query heads; otherwise it iterates all HEAD_Q heads. + head_per_group: tl.constexpr = HEAD_Q // NUM_HEAD_GROUPS + h_start = pid_h_group * head_per_group + h_end = h_start + head_per_group + for h_iter in range(h_start, h_end): + qhid = h_iter + + for m_start in range(m_loop_start, m_loop_end, BLOCK_M): + offs_m = m_start + tl.arange(0, BLOCK_M) + + q_ptrs = ( + Q + + bid * stride_qb + + qhid * stride_qh + + offs_m[:, None] * stride_qm + + offs_d[None, :] * stride_qd + ) + dout_ptrs = ( + DOUT + + bid * stride_dob + + qhid * stride_doh + + offs_m[:, None] * stride_dom + + offs_d[None, :] * stride_dod + ) + lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm + dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm + + q_load_mask = offs_m[:, None] < seqlen_q + q = tl.load(q_ptrs, mask=q_load_mask, other=0.0) + dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0) + lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0) + dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0) + + qk = tl.dot(q, kt) * sm_scale + + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0: + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + offs_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + elif is_pool_block: + pool_n = offs_n - HCA_LOCAL_SEQLEN + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + else: + if SWA_WINDOW > 0: + in_window = (offs_n[None, :] >= offs_m[:, None] - SWA_WINDOW + 1) & ( + offs_n[None, :] <= offs_m[:, None] + ) + qk = tl.where(in_window, qk, NEG_INF) + elif USE_CAUSAL: + qk = tl.where(offs_n[None, :] <= offs_m[:, None], qk, NEG_INF) + + if not EXACT_TILES_N: + qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF) + if not EXACT_TILES_M: + qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout, vt) + ds = p * (dp - dvec[:, None]) + + # P57 cr=0 BWD: same scale-defer + tl.dot(acc=) trick + # as the non-MQA branch above. + dv = tl.dot(tl.trans(p.to(dout.dtype)), dout, acc=dv) + dk = tl.dot(tl.trans(ds.to(q.dtype)), q, acc=dk) + + # P57 cr=0 BWD: fold ``sm_scale`` on dk once after the (head x m) + # loop. dv carries no scale. + dk = dk * sm_scale + + dk_ptrs = ( + DK + + bid * stride_dkb + + khid * stride_dkh + + offs_n[:, None] * stride_dkn + + offs_d[None, :] * stride_dkd + ) + dv_ptrs = ( + DV + + bid * stride_dvb + + khid * stride_dvh + + offs_n[:, None] * stride_dvn + + offs_d[None, :] * stride_dvd + ) + if ATOMIC_REDUCE or NUM_HEAD_GROUPS > 1: + # ``ATOMIC_REDUCE`` is set on the CSA pool BWD path (multiple + # ``(b, head, n_block)`` programs collapse into a single + # ``(b, n_block)`` global slice with ``stride_dkh = stride_dvh = 0``). + # ``NUM_HEAD_GROUPS > 1`` is the MQA head-split path: multiple + # head_group programs target the same MQA dK / dV slice. Both + # require fp32 atomic_add to merge. + tl.atomic_add(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k, sem="relaxed") + tl.atomic_add(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k, sem="relaxed") + else: + tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k) + tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k) + + +# --------------------------------------------------------------------------- +# Plan-8 P57 — HCA pool dK / dV kernel +# +# The HCA split-mask BWD has a single "pool" n-block (P=Sk-HCA_LOCAL_SEQLEN +# keys; typically 32 at V4-Flash widths). Folding the pool into the main +# dK/dV kernel makes that single n-block program iterate every m-block × +# every query head — at H=64 / Sq=4096 / BLOCK_M=32 that is 8192 inner +# iterations versus ~256 for a local-SWA n-block program (32× more work). +# +# This kernel parallelises pool work over ``m_block`` (one program per +# ``(b, m_block)``) and iterates ``HEAD_Q`` heads internally. Each program +# accumulates its (BLOCK_N, BLOCK_DMODEL) tile contribution in registers +# (no per-head atomic) and atomic-adds the final tile into the shared +# pool slice of dK / dV exactly twice (once for dk, once for dv) per +# program. At V4-Flash widths the contention is 128-way per pool cell +# instead of 8192-way for the m_block × qhid full grid, while still +# extracting 128× more parallelism than the original single-program pool +# branch. +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_attention_bwd_dkv_pool_kernel( + Q, + K, + V, + DOUT, + LSE, + D, + DK, # fp32 buffer [B, K_H, Sk, D] + DV, # fp32 buffer [B, K_H, Sk, D] + ADD_MASK, # [Sq, P] additive pool mask + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_kb, + stride_kh, + stride_kn, + stride_kd, + stride_vb, + stride_vh, + stride_vn, + stride_vd, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dkb, + stride_dkh, + stride_dkn, + stride_dkd, + stride_dvb, + stride_dvh, + stride_dvn, + stride_dvd, + stride_ms, + stride_mn, + seqlen_q, + seqlen_k, + pool_size, + sm_scale, + HEAD_Q: tl.constexpr, + HEAD_K: tl.constexpr, + HCA_LOCAL_SEQLEN: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """V4 HCA BWD — dK / dV for the pool keys (parallel m-blocks, head-loop inside).""" + pid_m = tl.program_id(0) + pid_b = tl.program_id(1) + bid = pid_b + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_DMODEL) + pool_n = tl.arange(0, BLOCK_N) + offs_n = HCA_LOCAL_SEQLEN + pool_n + + NEG_INF: tl.constexpr = -1.0e30 + + pool_n_mask = pool_n < pool_size + q_load_mask = offs_m[:, None] < seqlen_q + + # Per-(m_block, b) program iterates all query heads, accumulating in + # registers. For MQA (HK=1) all heads share the same (b, khid=0) K/V, + # so we reload K/V once at the start. For MHA (HK=HQ) we reload K/V + # per head inside the loop. + if HEAD_K == 1: + k_ptrs = K + bid * stride_kb + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd + v_ptrs = V + bid * stride_vb + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd + kv_load_mask = pool_n_mask[:, None] + k_shared = tl.load(k_ptrs, mask=kv_load_mask, other=0.0) + v_shared = tl.load(v_ptrs, mask=kv_load_mask, other=0.0) + kt_shared = tl.trans(k_shared) + vt_shared = tl.trans(v_shared) + + dk_acc = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) + dv_acc = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) + + for qhid in range(0, HEAD_Q): + if HEAD_K == HEAD_Q: + khid = qhid + else: + khid = 0 + + q_ptrs = ( + Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd + ) + dout_ptrs = ( + DOUT + + bid * stride_dob + + qhid * stride_doh + + offs_m[:, None] * stride_dom + + offs_d[None, :] * stride_dod + ) + lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm + dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm + + q = tl.load(q_ptrs, mask=q_load_mask, other=0.0) + dout = tl.load(dout_ptrs, mask=q_load_mask, other=0.0) + lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0) + dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0) + + if HEAD_K == 1: + kt = kt_shared + vt = vt_shared + else: + k_ptrs = ( + K + + bid * stride_kb + + khid * stride_kh + + offs_n[:, None] * stride_kn + + offs_d[None, :] * stride_kd + ) + v_ptrs = ( + V + + bid * stride_vb + + khid * stride_vh + + offs_n[:, None] * stride_vn + + offs_d[None, :] * stride_vd + ) + kv_load_mask = pool_n_mask[:, None] + k = tl.load(k_ptrs, mask=kv_load_mask, other=0.0) + v = tl.load(v_ptrs, mask=kv_load_mask, other=0.0) + kt = tl.trans(k) + vt = tl.trans(v) + + qk = tl.dot(q, kt) * sm_scale + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & pool_n_mask[None, :] + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + qk = tl.where(pool_n_mask[None, :], qk, NEG_INF) + qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout, vt) + ds = p * (dp - dvec[:, None]) + + if HEAD_K == 1: + dk_acc += tl.dot(tl.trans(ds.to(q.dtype)), q) * sm_scale + dv_acc += tl.dot(tl.trans(p.to(dout.dtype)), dout) + else: + # MHA path — each qhid maps to its own khid slice. Accumulator is + # not shared across heads; flush per-head with atomic_add. + dk_contrib = tl.dot(tl.trans(ds.to(q.dtype)), q) * sm_scale + dv_contrib = tl.dot(tl.trans(p.to(dout.dtype)), dout) + dk_ptrs_h = ( + DK + + bid * stride_dkb + + khid * stride_dkh + + offs_n[:, None] * stride_dkn + + offs_d[None, :] * stride_dkd + ) + dv_ptrs_h = ( + DV + + bid * stride_dvb + + khid * stride_dvh + + offs_n[:, None] * stride_dvn + + offs_d[None, :] * stride_dvd + ) + tl.atomic_add(dk_ptrs_h, dk_contrib, mask=pool_n_mask[:, None], sem="relaxed") + tl.atomic_add(dv_ptrs_h, dv_contrib, mask=pool_n_mask[:, None], sem="relaxed") + + if HEAD_K == 1: + khid_final = 0 + dk_ptrs = ( + DK + + bid * stride_dkb + + khid_final * stride_dkh + + offs_n[:, None] * stride_dkn + + offs_d[None, :] * stride_dkd + ) + dv_ptrs = ( + DV + + bid * stride_dvb + + khid_final * stride_dvh + + offs_n[:, None] * stride_dvn + + offs_d[None, :] * stride_dvd + ) + write_mask = pool_n_mask[:, None] + tl.atomic_add(dk_ptrs, dk_acc, mask=write_mask, sem="relaxed") + tl.atomic_add(dv_ptrs, dv_acc, mask=write_mask, sem="relaxed") + + +# --------------------------------------------------------------------------- +# Plan-8 P57 R2 — atomic-free MHA pool dK / dV kernel +# +# The R1 ``_v4_attention_bwd_dkv_pool_kernel`` MHA branch parallelises over +# m-blocks (one program per ``(b, m_block)``) and atomic-adds the per-head +# dK / dV contribution INSIDE the head loop: +# +# for qhid in range(HEAD_Q): # 64 heads +# tl.atomic_add(dk_ptrs[qhid], dk_contrib) +# tl.atomic_add(dv_ptrs[qhid], dv_contrib) +# +# At V4-Flash widths (B=1 H=64 Sq=4096 BM=64) this fires +# m_blocks × heads × 2 = 64 × 64 × 2 = 8192 +# ``tl.atomic_add`` instructions per pool kernel invocation, each on a +# (BLOCK_N=32, D=512) fp32 tile. Even at ``sem="relaxed"`` the L2 atomic +# engine on MI355 serialises 64-way contention on each per-head pool +# cache line — the pool kernel is dominated by atomic stalls. +# +# R2 rewrite: SWAP the loop nesting. Parallelise over ``(b, qhid)`` (one +# program per query head, owning the unique ``DK / DV[bid, qhid, pool, :]`` +# slice). The kernel pre-loads ``K / V[qhid, pool]`` ONCE, then iterates +# all m-blocks for that head, accumulating ``dk_acc / dv_acc`` in +# registers. The final write is a single ``tl.store`` per program — no +# atomic at all. +# +# Grid: ``(HEAD_Q, B)`` = (64, 1) at V4-Flash → 64 programs. That's only +# 25% of MI355's 256 CUs, but each program now has ~64 m-iter × dense +# matmul work (≈ heavy enough to saturate VALU+MFMA in its CU). +# +# Per-program HBM: +# * K / V[qhid, pool] = 2 × 32 × 512 × 2 B = 64 KiB (loaded once) +# * Q / dout[qhid, *] = 2 × Sq × D × 2 B = 8 MiB total per program +# * dK / dV[qhid, pool] = 2 × 32 × 512 × 4 B = 128 KiB (stored once) +# +# Compared to the R1 pool kernel: +# * Total Q / dout HBM reads UNCHANGED (64 programs × 64 m-iter == +# R1's 64 programs × 64 head-iter). +# * Total K / V HBM reads DROP ~64× (each (head, pool) tile loaded +# exactly once instead of once per m-block). +# * Atomic_add instructions DROP from 8192 → 0. +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_attention_bwd_dkv_pool_mha_kernel( + Q, + K, + V, + DOUT, + LSE, + D, + DK, + DV, + ADD_MASK, + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_kb, + stride_kh, + stride_kn, + stride_kd, + stride_vb, + stride_vh, + stride_vn, + stride_vd, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dkb, + stride_dkh, + stride_dkn, + stride_dkd, + stride_dvb, + stride_dvh, + stride_dvn, + stride_dvd, + stride_ms, + stride_mn, + seqlen_q, + pool_size, + sm_scale, + HCA_LOCAL_SEQLEN: tl.constexpr, + DK_POOL_OFFSET: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + POOL_EXACT: tl.constexpr, + EXACT_TILES_M: tl.constexpr, + M_SPLIT: tl.constexpr, +): + """MHA-only atomic-free HCA pool dK / dV kernel. + + Parallelism: ``(M_SPLIT * HEAD_Q, B)`` — ``M_SPLIT`` programs per + ``(b, qhid)``. ``M_SPLIT=1`` is the pure atomic-free design (one + program owns the full pool slice for its head); ``M_SPLIT>1`` shards + the m-loop into ``M_SPLIT`` chunks and merges with ``M_SPLIT``-way + ``tl.atomic_add`` per (b, qhid, pool, :) slice — the contention is + bounded by ``M_SPLIT`` and the extra programs lift HBM utilisation. + + Constraints: ``HEAD_K == HEAD_Q`` (MHA only). The MQA path keeps + using the original :func:`_v4_attention_bwd_dkv_pool_kernel`. + """ + pid_hm = tl.program_id(0) + pid_b = tl.program_id(1) + bid = pid_b + qhid = pid_hm // M_SPLIT + pid_m_chunk = pid_hm % M_SPLIT + khid = qhid # MHA invariant + + offs_d = tl.arange(0, BLOCK_DMODEL) + pool_n = tl.arange(0, BLOCK_N) + offs_n = HCA_LOCAL_SEQLEN + pool_n + + NEG_INF: tl.constexpr = -1.0e30 + + # Pre-load K / V for this (bid, khid) pool slice ONCE. + k_ptrs = ( + K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd + ) + v_ptrs = ( + V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd + ) + if POOL_EXACT: + k_tile = tl.load(k_ptrs) + v_tile = tl.load(v_ptrs) + else: + pool_n_mask = pool_n < pool_size + kv_load_mask = pool_n_mask[:, None] + k_tile = tl.load(k_ptrs, mask=kv_load_mask, other=0.0) + v_tile = tl.load(v_ptrs, mask=kv_load_mask, other=0.0) + kt = tl.trans(k_tile) + vt = tl.trans(v_tile) + + dk_acc = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) + dv_acc = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32) + + # Shard the m-loop into M_SPLIT contiguous chunks; each program owns + # one chunk of m-blocks. Chunk size is ``ceil(num_m_blocks / M_SPLIT)`` + # rounded up so the last chunk may be shorter. + num_m_blocks = (seqlen_q + BLOCK_M - 1) // BLOCK_M + m_chunk_size = (num_m_blocks + M_SPLIT - 1) // M_SPLIT + m_block_lo = pid_m_chunk * m_chunk_size + m_block_hi = m_block_lo + m_chunk_size + if m_block_hi > num_m_blocks: + m_block_hi = num_m_blocks + m_loop_start = m_block_lo * BLOCK_M + m_loop_end = m_block_hi * BLOCK_M + if m_loop_end > seqlen_q: + m_loop_end = seqlen_q + + for m_start in range(m_loop_start, m_loop_end, BLOCK_M): + offs_m = m_start + tl.arange(0, BLOCK_M) + + q_ptrs = ( + Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd + ) + dout_ptrs = ( + DOUT + + bid * stride_dob + + qhid * stride_doh + + offs_m[:, None] * stride_dom + + offs_d[None, :] * stride_dod + ) + lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm + dvec_ptrs = D + bid * stride_db + qhid * stride_dh + offs_m * stride_dm + + if EXACT_TILES_M: + q_local = tl.load(q_ptrs) + dout_local = tl.load(dout_ptrs) + lse = tl.load(lse_ptrs) + dvec = tl.load(dvec_ptrs) + else: + q_load_mask = offs_m[:, None] < seqlen_q + q_local = tl.load(q_ptrs, mask=q_load_mask, other=0.0) + dout_local = tl.load(dout_ptrs, mask=q_load_mask, other=0.0) + lse = tl.load(lse_ptrs, mask=offs_m < seqlen_q, other=0.0) + dvec = tl.load(dvec_ptrs, mask=offs_m < seqlen_q, other=0.0) + + # qk = Q @ K.T then scale + additive bias. We *cannot* defer the + # scale across the softmax (exp(scale*x) ≠ exp(x)), so scale stays + # on the qk path. + qk = tl.dot(q_local, kt) * sm_scale + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn + if POOL_EXACT and EXACT_TILES_M: + add_bias = tl.load(mask_ptrs).to(tl.float32) + else: + if POOL_EXACT: + mask_load_mask = offs_m[:, None] < seqlen_q + elif EXACT_TILES_M: + mask_load_mask = pool_n[None, :] < pool_size + else: + mask_load_mask = (offs_m[:, None] < seqlen_q) & (pool_n[None, :] < pool_size) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + if not POOL_EXACT: + qk = tl.where(pool_n[None, :] < pool_size, qk, NEG_INF) + if not EXACT_TILES_M: + qk = tl.where(offs_m[:, None] < seqlen_q, qk, NEG_INF) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout_local, vt) + ds = p * (dp - dvec[:, None]) + + # R2 scale-defer for dk: accumulate ds.T @ Q without sm_scale, + # fold the single multiply after the m-loop. dv carries no scale. + dv_acc = tl.dot(tl.trans(p.to(dout_local.dtype)), dout_local, acc=dv_acc) + dk_acc = tl.dot(tl.trans(ds.to(q_local.dtype)), q_local, acc=dk_acc) + + dk_acc = dk_acc * sm_scale + + # Plan-8 P57 R2: DK / DV may either be the full ``[B, HK, Sk, D]`` + # fp32 buffer (then ``DK_POOL_OFFSET == HCA_LOCAL_SEQLEN`` so we + # land on the pool slice) or a pool-only ``[B, HK, pool_size, D]`` + # fp32 sidecar (then ``DK_POOL_OFFSET == 0``). The latter is the + # default R2 path: the full ``dk_out / dv_out`` allocate in the + # input dtype and we cast the small sidecar to bf16 at the end — + # saves the big ``dk_fp32 -> bf16`` cast. + offs_n_dk = DK_POOL_OFFSET + pool_n + dk_ptrs = ( + DK + + bid * stride_dkb + + khid * stride_dkh + + offs_n_dk[:, None] * stride_dkn + + offs_d[None, :] * stride_dkd + ) + dv_ptrs = ( + DV + + bid * stride_dvb + + khid * stride_dvh + + offs_n_dk[:, None] * stride_dvn + + offs_d[None, :] * stride_dvd + ) + if M_SPLIT == 1: + if POOL_EXACT: + tl.store(dk_ptrs, dk_acc) + tl.store(dv_ptrs, dv_acc) + else: + write_mask = pool_n[:, None] < pool_size + tl.store(dk_ptrs, dk_acc, mask=write_mask) + tl.store(dv_ptrs, dv_acc, mask=write_mask) + else: + if POOL_EXACT: + tl.atomic_add(dk_ptrs, dk_acc, sem="relaxed") + tl.atomic_add(dv_ptrs, dv_acc, sem="relaxed") + else: + write_mask = pool_n[:, None] < pool_size + tl.atomic_add(dk_ptrs, dk_acc, mask=write_mask, sem="relaxed") + tl.atomic_add(dv_ptrs, dv_acc, mask=write_mask, sem="relaxed") + + +# --------------------------------------------------------------------------- +# Python launcher +# --------------------------------------------------------------------------- + + +def _launch_v4_attention_bwd( + q: torch.Tensor, # [B, H, Sq, D] + k: torch.Tensor, # [B, K_H, Sk, D] + v: torch.Tensor, # [B, K_H, Sk, D] + out: torch.Tensor, # [B, H, Sq, D] (FWD output) + dout: torch.Tensor, # [B, H, Sq, D] + lse: torch.Tensor, # [B, H, Sq] fp32 + *, + sink: Optional[torch.Tensor], # [H] or None + swa_window: int, + additive_mask: Optional[torch.Tensor], # [Sq, Sk] or None + scale: float, + hca_local_seqlen: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Launch the V4 attention backward kernel. + + Returns ``(dq, dk, dv, dsink)`` — gradients in the input dtype, with + ``dsink`` returned only when ``sink is not None`` (else ``None``). + """ + if not q.is_cuda: + raise ValueError("v4_attention_v1 BWD requires CUDA / HIP tensors.") + if dout.shape != out.shape or out.shape != q.shape: + raise ValueError( + "v4_attention_v1 BWD shape mismatch: " + f"out={tuple(out.shape)}, dout={tuple(dout.shape)}, q={tuple(q.shape)}" + ) + + B, HQ, Sq, D = q.shape + HK = k.shape[1] + Sk = k.shape[2] + + has_sink = sink is not None + has_add_mask = additive_mask is not None + hca_local_seqlen = int(hca_local_seqlen) + if hca_local_seqlen: + if not has_add_mask: + raise ValueError("hca_local_seqlen requires a pool additive_mask.") + if hca_local_seqlen <= 0 or hca_local_seqlen >= Sk: + raise ValueError( + "hca_local_seqlen must split local and pool keys " + f"(got hca_local_seqlen={hca_local_seqlen}, Sk={Sk})." + ) + expected_mask_shape = (Sq, Sk - hca_local_seqlen) + if tuple(additive_mask.shape) != expected_mask_shape: + raise ValueError( + "HCA split-mask mode expects additive_mask shape " + f"{expected_mask_shape}, got {tuple(additive_mask.shape)}." + ) + if swa_window <= 0: + raise ValueError("HCA split-mask mode requires swa_window > 0.") + use_causal = (not has_add_mask) and (swa_window <= 0) + swa_window_constexpr = ( + int(swa_window) if ((not has_add_mask or hca_local_seqlen) and swa_window > 0) else 0 + ) + + # Plan-8 P57 R2 sweep at V4-Flash widths (B=1 H=64 Sq=4096 D=512 P=32): + # R1 settled on ``BLOCK_M=64 BLOCK_N=16``; the R2 sweep — once the + # MHA pool kernel went atomic-free with ``M_SPLIT=4`` — discovered + # ``BLOCK_M=32`` with ``dKV num_warps=2`` cuts cr=128 HCA BWD from + # 3.79 ms to ~3.44 ms AND cr=0 dense BWD from 3.01 ms to ~2.77 ms. + # Why BM=32 now wins: + # + # * dq: smaller Q tile (32×D=512=32KiB vs 64×512=64KiB) drops VGPR + # pressure enough to run ``num_warps=2`` (1 warp = 64 threads + # per program) instead of nw=4 — same overall warp occupancy + # across the grid but the per-warp register budget doubles and + # k-axis matmul (D=512) pipelines without VGPR spilling. + # * dkv (local SWA): the m-loop is 5 iter at BM=32 (vs 3 iter at + # BM=64), but each iter loads half the Q/dout tile (32×512×2 B + # each) so per-iter HBM is halved. Total per-program HBM is + # roughly the same; the win is reduced per-iter register pressure + # enabling ``num_warps=2``. + # * pool (MHA atomic-free): more m-blocks per program (128 at + # BM=32 vs 64 at BM=64) lets the ``M_SPLIT=4`` shard cleanly + # into 32 m-iter per program (vs 16) and fills the SIMDs. + BLOCK_M = int(os.getenv("PRIMUS_V4_ATTN_BWD_BLOCK_M", "32")) + BLOCK_N = int(os.getenv("PRIMUS_V4_ATTN_BWD_BLOCK_N", "16")) + BLOCK_DMODEL = D + + # Plan-8 P57 R2: dq/dk/dv output buffer strategy + # ---------------------------------------------- + # dq is written by ``_v4_attention_bwd_dq_kernel`` via ``tl.store`` + # (no atomic_add — each (b, qhid, m_block) program owns a unique + # slice). So we let the kernel write directly into the input dtype + # and skip the final fp32 -> bf16 cast (~ 0.15 ms saved at V4-Flash + # widths). The dq accumulator inside the kernel remains fp32 for the + # n-loop reduction; only the final store converts. + # + # dk / dv are similar for the MHA non-MQA case (NUM_HEAD_GROUPS=1): + # the dKV kernel ``tl.store``s the local-SWA slice with no atomic. + # The HCA pool kernel still uses atomic_add (M_SPLIT > 1) which + # requires fp32, so we keep a SEPARATE small fp32 pool sidecar + # buffer of shape ``[B, HK, pool_size, D]`` and merge it into the + # full bf16 dk / dv at the end. + # + # For MQA (HK != HQ) or NUM_HEAD_GROUPS > 1: keep the original + # fp32-everywhere path since dKV uses atomic_add. + dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=q.dtype) + use_native_dkv_dtype = (HQ == HK) and (q.dtype in (torch.bfloat16, torch.float16)) + if use_native_dkv_dtype: + dk_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=k.dtype) + dv_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=v.dtype) + else: + dk_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=torch.float32) + dv_fp32 = torch.zeros((B, HK, Sk, D), device=q.device, dtype=torch.float32) + # Pool-side fp32 sidecar (only allocated when HCA + MHA + native dtype dKV). + if use_native_dkv_dtype and hca_local_seqlen > 0: + pool_size_alloc = Sk - hca_local_seqlen + dk_pool_fp32 = torch.zeros((B, HK, pool_size_alloc, D), device=q.device, dtype=torch.float32) + dv_pool_fp32 = torch.zeros((B, HK, pool_size_alloc, D), device=q.device, dtype=torch.float32) + else: + dk_pool_fp32 = None + dv_pool_fp32 = None + if has_sink: + dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32) + sink_arg = sink.to(torch.float32) if sink.dtype != torch.float32 else sink + else: + dsink_fp32 = q # sentinel; HAS_SINK=False inside kernel + sink_arg = q + + # D scalar = (dout * out).sum(-1) + d_buf = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32) + pre_grid = (triton.cdiv(Sq, BLOCK_M), B * HQ) + _v4_attention_bwd_preprocess_kernel[pre_grid]( + out, + dout, + d_buf, + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + Sq, + HEAD=HQ, + BLOCK_M=BLOCK_M, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=4, + num_stages=1, + ) + + # Mask sentinel ptr for HAS_ADD_MASK=False + mask_arg = additive_mask if has_add_mask else q + if has_add_mask: + stride_ms = additive_mask.stride(0) + stride_mn = additive_mask.stride(1) + else: + stride_ms = 0 + stride_mn = 0 + + # Plan-8 P57: skip the boundary masks (``offs_m < seqlen_q`` / + # ``offs_n < seqlen_k``) when the launcher confirms the seqlens are + # exact multiples of BLOCK_M / BLOCK_N. Saves a per-inner-iter + # broadcast + tl.where on production V4-Flash widths where Sq=4096 + # (BM=64) and Sk=4128 (BN=16) both divide cleanly. + exact_tiles_m = (Sq % BLOCK_M) == 0 + exact_tiles_n = (Sk % BLOCK_N) == 0 + + # Plan-5 P32: split BWD (dQ kernel + dK/dV kernel, no atomics for + # dQ / dK / dV) is now the default — wins both the operator microbench + # *and* the EP8 proxy after the dual-RoPE bf16-cast fix (P32 RoPE bug: + # ``apply_interleaved_partial_rope`` was upcasting Q/K to fp32 because + # cos/sin came from ``position_ids.float()`` and bf16 * fp32 = fp32, + # which 2x'd Q/K HBM traffic, inflated the kernel time 1.8-7x in the + # proxy and made the monolithic design *look* faster in A/B traces). + # ``PRIMUS_V4_ATTN_BWD_USE_SPLIT=0`` falls back to the monolithic + # design for kernel-level perf experiments / debugging. + if os.getenv("PRIMUS_V4_ATTN_BWD_USE_SPLIT", "1") == "1": + dq_grid = (triton.cdiv(Sq, BLOCK_M), B * HQ) + _v4_attention_bwd_dq_kernel[dq_grid]( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + dsink_fp32, + sink_arg, + mask_arg, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_fp32.stride(0), + dq_fp32.stride(1), + dq_fp32.stride(2), + dq_fp32.stride(3), + stride_ms, + stride_mn, + Sq, + Sk, + float(scale), + HEAD_Q=HQ, + HEAD_K=HK, + SWA_WINDOW=swa_window_constexpr, + HAS_SINK=has_sink, + HAS_ADD_MASK=has_add_mask, + HCA_LOCAL_SEQLEN=hca_local_seqlen, + USE_CAUSAL=use_causal, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL=BLOCK_DMODEL, + EXACT_TILES_M=exact_tiles_m, + EXACT_TILES_N=exact_tiles_n, + # Plan-8 P57 R2: ``num_warps=2 num_stages=1`` is the new sweet + # spot once BLOCK_M dropped from 64 -> 32. With a 32×D Q tile + # the dq inner loop fits in the per-warp register budget with + # only 1 warp/64-thread wave, doubling the per-warp register + # budget vs the R1 nw=4 default and removing the K/V double- + # buffer (num_stages=2) cost since the loop now runs short. + num_warps=int(os.getenv("PRIMUS_V4_ATTN_BWD_DQ_NUM_WARPS", "2")), + num_stages=int(os.getenv("PRIMUS_V4_ATTN_BWD_DQ_NUM_STAGES", "1")), + ) + # Plan-8 P57: in HCA mode the pool n-block(s) are handled by the + # dedicated pool kernel that parallelises over (m_block, b * qhid) + # — see ``_v4_attention_bwd_dkv_pool_kernel``. Skip the pool grid + # rows here so we don't double-count. + # Plan-8 P57 R2: per-kernel BN override. Default uses the global + # BLOCK_N (BN=16 from R1 sweep). Setting ``PRIMUS_V4_ATTN_BWD_DKV_BLOCK_N`` + # lets us test wider K-tiles in the dKV kernel without affecting + # the dq / pool paths. + dkv_block_n = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_BLOCK_N", str(BLOCK_N))) + if hca_local_seqlen > 0: + dkv_n_blocks = triton.cdiv(hca_local_seqlen, dkv_block_n) + else: + dkv_n_blocks = triton.cdiv(Sk, dkv_block_n) + # Plan-8 P57: ``NUM_HEAD_GROUPS`` controls MQA head-split parallelism + # for the local dKV kernel. Default (1) keeps the original head loop; + # >1 splits the head loop and uses ``tl.atomic_add`` for dKV. + # The dense-bench sweep showed head-split is essentially neutral on + # MI355 at H=64 (HBM/atomic-bound, not compute-bound), so default 1. + # BUT that result does NOT hold for the MQA/HQ>=128 (V4-Pro) SWA dkv + # shape: at HQ=128/HK=1 the HG=1 grid is 1 workgroup/CU (occ=1) with each + # program serially grinding all 128 heads -> tail/pipeline-fill bound. + # HG=2 doubles the grid + fills the idle MFMA cycles (grad cos ~1.0); + # HG=4/8 regress. So default the MQA/HQ>=128 case to 2; still + # overridable via PRIMUS_V4_ATTN_BWD_DKV_HEAD_GROUPS. + num_head_groups = 1 + if HQ > HK: + # Arch-aware (see ._v4_attn_tuning): gfx1250 wants HG=32 (+37-53%, ab_sweep/opt7b), + # gfx950 wants 2. Env knob still overrides. + from ._v4_attn_tuning import bwd_dkv_head_groups_default + + _hg_default = str(bwd_dkv_head_groups_default(HQ, HK)) + target = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_HEAD_GROUPS", _hg_default)) + while target > 1 and HQ % target != 0: + target //= 2 + num_head_groups = max(1, target) + dkv_grid = (dkv_n_blocks, B * HK, num_head_groups) + # Plan-8 P57 R2: ``num_warps=2 num_stages=1`` paired with BM=32. + # The 32×D Q tile + persistent K/V tile fit in the 2-warp register + # budget; bumping to nw=4 doubles spills (m-loop is 5 iter at + # BM=32 so the wider warp budget hurts more than the SIMD-fill + # benefit). num_stages>1 also costs more than it pays because + # the K/V are loaded once before the m-loop. + dkv_num_warps = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_NUM_WARPS", "2")) + dkv_num_stages = int(os.getenv("PRIMUS_V4_ATTN_BWD_DKV_NUM_STAGES", "1")) + _v4_attention_bwd_dkv_kernel[dkv_grid]( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + mask_arg, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dk_fp32.stride(0), + dk_fp32.stride(1), + dk_fp32.stride(2), + dk_fp32.stride(3), + dv_fp32.stride(0), + dv_fp32.stride(1), + dv_fp32.stride(2), + dv_fp32.stride(3), + stride_ms, + stride_mn, + Sq, + Sk, + float(scale), + HEAD_Q=HQ, + HEAD_K=HK, + SWA_WINDOW=swa_window_constexpr, + HAS_ADD_MASK=has_add_mask, + HCA_LOCAL_SEQLEN=hca_local_seqlen, + USE_CAUSAL=use_causal, + BLOCK_M=BLOCK_M, + BLOCK_N=dkv_block_n, + BLOCK_DMODEL=BLOCK_DMODEL, + NUM_HEAD_GROUPS=num_head_groups, + EXACT_TILES_M=exact_tiles_m, + EXACT_TILES_N=(Sk % dkv_block_n) == 0, + num_warps=dkv_num_warps, + num_stages=dkv_num_stages, + ) + if hca_local_seqlen > 0 and os.getenv("PRIMUS_V4_ATTN_BWD_HCA_POOL", "1") == "1": + pool_size = Sk - hca_local_seqlen + pool_block_n = max(16, triton.next_power_of_2(pool_size)) + pool_exact = pool_size == pool_block_n + # Plan-8 P57 R2: use the atomic-free MHA pool kernel when + # HEAD_K == HEAD_Q (production V4-Flash widths). Parallelism + # is over (qhid, b) instead of (m_block, b); each program + # owns the unique dK / dV[bid, qhid, pool, :] slice and writes + # it with ``tl.store`` after accumulating across all m-blocks + # in registers — no atomic_add at all. Eliminates the 8192 + # per-launch atomic_add stalls of the R1 MHA pool path. + # + # For MQA (HK=1, multiple qhids share one khid slice) we still + # need to merge across heads; fall back to the original pool + # kernel which atomic-adds in the head loop. + use_mha_pool = (HQ == HK) and (os.getenv("PRIMUS_V4_ATTN_BWD_POOL_MHA", "1") == "1") + if use_mha_pool: + # Plan-8 P57 R2: pool M_SPLIT controls m-loop sharding. + # M_SPLIT=1 = pure atomic-free (64 progs); M_SPLIT=N adds + # N-way atomic_add per (b, qhid, pool) slice but lifts the + # program count to N*HQ (better HBM bandwidth utilization + # at the cost of contention). Default 4 = 256 progs. + pool_m_split = int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_M_SPLIT", "4")) + num_m_blocks = triton.cdiv(Sq, BLOCK_M) + while pool_m_split > num_m_blocks: + pool_m_split //= 2 + pool_m_split = max(1, pool_m_split) + pool_grid = (HQ * pool_m_split, B) + # Plan-8 P57 R2: pool kernel writes to fp32 sidecar (so + # ``M_SPLIT > 1`` atomic_add works). The kernel uses + # offset ``HCA_LOCAL_SEQLEN + pool_n`` along the seqlen + # axis, so we pass a "virtual" full-seqlen view of the + # pool buffer: a tensor of shape (B, HK, Sk, D) where the + # data starts ``HCA_LOCAL_SEQLEN`` BEFORE the pool + # buffer's actual base. We accomplish this by passing the + # pool buffer's base pointer with an *offset* base + # subtracted via the seqlen stride. PyTorch tensors don't + # support negative offsets cleanly, so instead we pass + # the pool buffer's storage with seqlen stride and tell + # the kernel to subtract HCA_LOCAL_SEQLEN — see the + # ``DK_SEQ_OFFSET`` constexpr below. + if dk_pool_fp32 is not None: + pool_dk_buf = dk_pool_fp32 + pool_dv_buf = dv_pool_fp32 + pool_dk_strides = ( + dk_pool_fp32.stride(0), + dk_pool_fp32.stride(1), + dk_pool_fp32.stride(2), + dk_pool_fp32.stride(3), + ) + pool_dv_strides = ( + dv_pool_fp32.stride(0), + dv_pool_fp32.stride(1), + dv_pool_fp32.stride(2), + dv_pool_fp32.stride(3), + ) + pool_hca_arg = 0 # pool buffer is pool-only; no HCA_LOCAL offset + else: + pool_dk_buf = dk_fp32 + pool_dv_buf = dv_fp32 + pool_dk_strides = ( + dk_fp32.stride(0), + dk_fp32.stride(1), + dk_fp32.stride(2), + dk_fp32.stride(3), + ) + pool_dv_strides = ( + dv_fp32.stride(0), + dv_fp32.stride(1), + dv_fp32.stride(2), + dv_fp32.stride(3), + ) + pool_hca_arg = hca_local_seqlen + _v4_attention_bwd_dkv_pool_mha_kernel[pool_grid]( + q, + k, + v, + dout, + lse, + d_buf, + pool_dk_buf, + pool_dv_buf, + additive_mask, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + pool_dk_strides[0], + pool_dk_strides[1], + pool_dk_strides[2], + pool_dk_strides[3], + pool_dv_strides[0], + pool_dv_strides[1], + dv_fp32.stride(2), + dv_fp32.stride(3), + stride_ms, + stride_mn, + Sq, + pool_size, + float(scale), + HCA_LOCAL_SEQLEN=hca_local_seqlen, + DK_POOL_OFFSET=pool_hca_arg, + BLOCK_M=BLOCK_M, + BLOCK_N=pool_block_n, + BLOCK_DMODEL=BLOCK_DMODEL, + POOL_EXACT=pool_exact, + EXACT_TILES_M=exact_tiles_m, + M_SPLIT=pool_m_split, + # Plan-8 P57 R2: ``num_warps=4 num_stages=3`` for the + # atomic-free MHA pool kernel. Pool m-loop is long + # (Sq/(BM*M_SPLIT) = 32 iter at BM=32 M_SPLIT=4) so + # 3-stage pipelining of Q+dout loads pays off. + num_warps=int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_NUM_WARPS", "4")), + num_stages=int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_NUM_STAGES", "3")), + ) + else: + # Plan-8 P57: pool kernel block_m is decoupled from the dKV + # block_m. dKV benefits from BM=64 (fewer programs each loading + # a wider Q tile); pool kernel benefits from a smaller BM (more + # programs => more parallelism over the head loop). + # gfx950/MI355X: the pool kernel scales with program count, so + # BM=16 (vs the dKV default 32) speeds up the full HCA backward + # (cos ~1.0). Only the pool path reads this, so SWA bwd is + # unaffected. Env-overridable. + pool_block_m = int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_BLOCK_M", "16")) + pool_grid = (triton.cdiv(Sq, pool_block_m), B) + _v4_attention_bwd_dkv_pool_kernel[pool_grid]( + q, + k, + v, + dout, + lse, + d_buf, + dk_fp32, + dv_fp32, + additive_mask, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dk_fp32.stride(0), + dk_fp32.stride(1), + dk_fp32.stride(2), + dk_fp32.stride(3), + dv_fp32.stride(0), + dv_fp32.stride(1), + dv_fp32.stride(2), + dv_fp32.stride(3), + stride_ms, + stride_mn, + Sq, + Sk, + pool_size, + float(scale), + HEAD_Q=HQ, + HEAD_K=HK, + HCA_LOCAL_SEQLEN=hca_local_seqlen, + BLOCK_M=pool_block_m, + BLOCK_N=pool_block_n, + BLOCK_DMODEL=BLOCK_DMODEL, + # Plan-8 P57 sweep with dKV BM=64 BN=16: pool kernel + # ``num_warps=4 num_stages=1`` is the sweet spot. Higher + # num_stages adds Q/dout double-buffering overhead that + # doesn't pay off because the pool kernel keeps K/V in + # registers (single load) and the head-loop dominates + # the iteration count. + num_warps=int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_NUM_WARPS", "4")), + num_stages=int(os.getenv("PRIMUS_V4_ATTN_BWD_POOL_NUM_STAGES", "1")), + ) + # Plan-8 P57 R2: cast outputs back to input dtype only when not + # already in the input dtype. dq is always allocated in q.dtype + # (no atomic), so this is a no-op. dk / dv may be either: + # * native bf16 (HQ == HK MHA path) — no cast needed for the + # local-SWA part; only the pool sidecar (if present) needs + # a fp32 -> bf16 cast that we splice into the pool slice. + # * fp32 (MQA or grouped paths) — full cast at the end. + dq_out = dq_fp32 if dq_fp32.dtype == q.dtype else dq_fp32.to(q.dtype) + if dk_fp32.dtype != k.dtype: + dk_out = dk_fp32.to(k.dtype) + else: + dk_out = dk_fp32 + if dv_fp32.dtype != v.dtype: + dv_out = dv_fp32.to(v.dtype) + else: + dv_out = dv_fp32 + if dk_pool_fp32 is not None and hca_local_seqlen > 0: + # Splice the pool sidecar (fp32) into the pool slice of the + # bf16 dk_out / dv_out buffers. + dk_out[..., hca_local_seqlen:, :] = dk_pool_fp32.to(dk_out.dtype) + dv_out[..., hca_local_seqlen:, :] = dv_pool_fp32.to(dv_out.dtype) + dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None + return dq_out, dk_out, dv_out, dsink_out + + grid = (triton.cdiv(Sq, BLOCK_M), B * HQ) + _v4_attention_bwd_kernel[grid]( + q, + k, + v, + dout, + lse, + d_buf, + dq_fp32, + dk_fp32, + dv_fp32, + dsink_fp32, + sink_arg, + mask_arg, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_fp32.stride(0), + dq_fp32.stride(1), + dq_fp32.stride(2), + dq_fp32.stride(3), + dk_fp32.stride(0), + dk_fp32.stride(1), + dk_fp32.stride(2), + dk_fp32.stride(3), + dv_fp32.stride(0), + dv_fp32.stride(1), + dv_fp32.stride(2), + dv_fp32.stride(3), + stride_ms, + stride_mn, + Sq, + Sk, + float(scale), + HEAD_Q=HQ, + HEAD_K=HK, + SWA_WINDOW=swa_window_constexpr, + HAS_SINK=has_sink, + HAS_ADD_MASK=has_add_mask, + HCA_LOCAL_SEQLEN=hca_local_seqlen, + USE_CAUSAL=use_causal, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=8, + num_stages=1, + ) + + # Cast fp32 buffers back to input dtype. + dq_out = dq_fp32.to(q.dtype) + dk_out = dk_fp32.to(k.dtype) + dv_out = dv_fp32.to(v.dtype) + dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None + return dq_out, dk_out, dv_out, dsink_out + + +__all__ = [ + "_v4_attention_bwd_preprocess_kernel", + "_v4_attention_bwd_kernel", + "_launch_v4_attention_bwd", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_fwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_fwd.py new file mode 100644 index 000000000..4d1bf440e --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_attention_fwd.py @@ -0,0 +1,494 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 attention forward Triton kernel (plan-4 P25, ``compress_ratio in {0, 128}``). + +FlashAttention-style block-wise online softmax that handles V4's exact +shape envelope: + +* ``head_dim = 512`` (single tile — no partial-RoPE / NOPE split because + RoPE is applied outside the kernel); +* MQA single-latent KV (``K.shape[1] == 1``) broadcast across query + heads, *or* full MHA (``K.shape[1] == H``); +* optional per-head learned softmax sink (``[H]``), joined as a virtual + key column with notional value zero at the end of the K-loop; +* optional sliding-window-causal mask (``swa_window > 0``) applied + in-kernel; +* optional caller-supplied additive bias. The generic path accepts + ``[Sq, Sk]`` and ignores ``swa_window``. The HCA split-mask path + accepts pool-only ``[Sq, P]`` with ``HCA_LOCAL_SEQLEN=Sq`` and keeps + the local branch on kernel-native SWA. + +dtype contract (matches :func:`eager_v4_attention`): + +* Q / K / V matmuls run in input dtype on tensor cores; accumulators + inside the matmul are fp32. +* The online-softmax accumulator (``m_running``, ``l_running``, + ``acc``) lives in fp32 — this is the *only* fp32 step inside the + kernel. +* Output is written back in input dtype; saved ``LSE`` is fp32 (BWD + re-materialises ``P`` from it). +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch +import triton +import triton.language as tl + +# --------------------------------------------------------------------------- +# Triton kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_attention_fwd_kernel( + Q, + K, + V, + OUT, + LSE, + SINK, + ADD_MASK, + # Q strides: [B, H, Sq, D] row-major (contiguous on D) + stride_qb, + stride_qh, + stride_qm, + stride_qd, + # K strides: [B, K_H, Sk, D] row-major (K_H == 1 for MQA, == H for MHA) + stride_kb, + stride_kh, + stride_kn, + stride_kd, + # V strides: [B, K_H, Sk, D] row-major + stride_vb, + stride_vh, + stride_vn, + stride_vd, + # OUT strides: [B, H, Sq, D] row-major + stride_ob, + stride_oh, + stride_om, + stride_od, + # LSE strides: [B, H, Sq] row-major + stride_lb, + stride_lh, + stride_lm, + # ADD_MASK strides: [Sq, Sk] row-major (broadcasts over B, H) + stride_ms, + stride_mn, + seqlen_q, + seqlen_k, + sm_scale, + HEAD_Q: tl.constexpr, + HEAD_K: tl.constexpr, # 1 for MQA, == HEAD_Q for MHA + SWA_WINDOW: tl.constexpr, # 0 = off, > 0 = SWA window + HAS_SINK: tl.constexpr, + HAS_ADD_MASK: tl.constexpr, + HCA_LOCAL_SEQLEN: tl.constexpr, # 0 = generic mask; >0 = [local SWA keys | pool keys] + USE_CAUSAL: tl.constexpr, # only meaningful when HAS_ADD_MASK = False + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """V4 attention FWD. + + Grid layout: ``(cdiv(seqlen_q, BLOCK_M), batch * HEAD_Q)``. Each + program (program-id) computes one ``[BLOCK_M, BLOCK_DMODEL]`` slice + of ``OUT`` and the matching ``[BLOCK_M]`` slice of ``LSE``. + + Mask precedence (must match :func:`eager_v4_attention`): + + 1. ``HAS_ADD_MASK`` — load the ``[Sq, Sk]`` additive bias and add + to ``qk``. SWA / causal masks are NOT applied in-kernel (the + caller has embedded all structure in the bias). + 2. ``SWA_WINDOW > 0`` — sliding-window causal: + keep ``offs_n in [offs_m - SWA_WINDOW + 1, offs_m]``. + 3. ``USE_CAUSAL`` — full causal: keep ``offs_n <= offs_m``. + + In all branches, keys at index ``>= seqlen_k`` are masked to ``-inf`` + so the boundary rows of the K-loop tile do not contaminate the + softmax denominator. + """ + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + + bid = pid_bh // HEAD_Q + qhid = pid_bh % HEAD_Q + if HEAD_K == HEAD_Q: + khid = qhid + else: + khid = 0 # MQA: single shared K / V head + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_DMODEL) + + # Q tile: [BLOCK_M, BLOCK_DMODEL] in q.dtype + q_ptrs = ( + Q + bid * stride_qb + qhid * stride_qh + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd + ) + q_load_mask = offs_m[:, None] < seqlen_q + q = tl.load(q_ptrs, mask=q_load_mask, other=0.0) + + # Online-softmax running state (fp32). We use the finite sentinel + # NEG_INF (-1e30) instead of -float("inf") so that the all-masked- + # tile corner case (m_running == m_tile == NEG_INF, e.g. for early + # queries under SWA when n_start is entirely outside the window) + # does NOT produce NaN through ``exp(-inf - -inf) = exp(NaN)``. + # NEG_INF is far enough below any plausible logit that + # ``exp(NEG_INF) ≈ 0`` and the algebra is identical to using -inf. + NEG_INF: tl.constexpr = -1.0e30 + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32) + m_i = tl.full([BLOCK_M], value=NEG_INF, dtype=tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + + # Determine k-loop bounds. + # + # Plan-5 P30: for SWA, skip K tiles that are guaranteed outside the + # sliding window for *every* row in this M block. The old P25 loop + # started at zero and relied on the mask below, which means late + # rows at S=4096 spent most of their time multiplying all-masked + # tiles. Rounding down to the BLOCK_N boundary preserves exact + # masking semantics while cutting the steady-state SWA tile count + # from O(S / BLOCK_N) to O(window / BLOCK_N). + n_loop_start = 0 + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0: + # Caller's additive_mask handles all masking; iterate the full + # key axis. (We still apply boundary mask for keys past seqlen_k.) + n_loop_end = seqlen_k + elif SWA_WINDOW > 0: + # This block's earliest row is pid_m * BLOCK_M. Any key before + # earliest_row - SWA_WINDOW + 1 is invisible for every row in the + # block, so skip those tiles entirely. + n_loop_start = pid_m * BLOCK_M - SWA_WINDOW + 1 + if n_loop_start < 0: + n_loop_start = 0 + n_loop_start = (n_loop_start // BLOCK_N) * BLOCK_N + # Keys with n > max(offs_m) are causal-masked for every row. + n_loop_end = (pid_m + 1) * BLOCK_M + local_end = HCA_LOCAL_SEQLEN if HCA_LOCAL_SEQLEN > 0 else seqlen_k + if n_loop_end > local_end: + n_loop_end = local_end + elif USE_CAUSAL: + # Full causal: keys with n > max(offs_m) are -inf, so the loop + # only needs to cover keys up to (pid_m + 1) * BLOCK_M. + n_loop_end = (pid_m + 1) * BLOCK_M + if n_loop_end > seqlen_k: + n_loop_end = seqlen_k + else: + n_loop_end = seqlen_k + + # K-loop: local branch. For HCA split-mask this covers only the + # pruned local SWA prefix; the pool suffix is handled by the second + # loop below. + for n_start in range(n_loop_start, n_loop_end, BLOCK_N): + offs_n = n_start + tl.arange(0, BLOCK_N) + + # K tile: [BLOCK_N, BLOCK_DMODEL] in k.dtype + k_ptrs = ( + K + bid * stride_kb + khid * stride_kh + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd + ) + k_load_mask = offs_n[:, None] < seqlen_k + k = tl.load(k_ptrs, mask=k_load_mask, other=0.0) + + # qk = q @ k.T : [BLOCK_M, BLOCK_N] in fp32 (tl.dot accumulator) + qk = tl.dot(q, tl.trans(k)) * sm_scale + + # Mask: additive_mask OR SWA / causal (mutually exclusive — see + # the docstring's precedence rule). + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN == 0: + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + offs_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + else: + if SWA_WINDOW > 0: + # offs_n in [offs_m - SWA_WINDOW + 1, offs_m] + in_window = (offs_n[None, :] >= offs_m[:, None] - SWA_WINDOW + 1) & ( + offs_n[None, :] <= offs_m[:, None] + ) + qk = tl.where(in_window, qk, NEG_INF) + elif USE_CAUSAL: + qk = tl.where(offs_n[None, :] <= offs_m[:, None], qk, NEG_INF) + + # Boundary: keys past seqlen_k were loaded as 0; mask them to + # NEG_INF (finite sentinel) so they do not contribute to the + # softmax denominator and do NOT produce NaN through + # ``exp(-inf - -inf)``. + qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF) + + # Online softmax update + m_ij = tl.max(qk, 1) + m_new = tl.maximum(m_i, m_ij) + alpha = tl.exp(m_i - m_new) + p = tl.exp(qk - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p, 1) + + # V tile: [BLOCK_N, BLOCK_DMODEL] in v.dtype + v_ptrs = ( + V + bid * stride_vb + khid * stride_vh + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd + ) + v_load_mask = offs_n[:, None] < seqlen_k + v = tl.load(v_ptrs, mask=v_load_mask, other=0.0) + + # P57 R2: fuse ``acc * alpha`` into the V-dot MFMA-acc input. + # AMD's Triton lowering threads ``acc`` directly into the MFMA + # C-tile, eliminating one fp32 register round-trip per K-tile. + # Same pattern as the cr=0 BWD's ``dq = tl.dot(ds, k, acc=dq)``. + acc = tl.dot(p.to(v.dtype), v, acc=acc * alpha[:, None]) + m_i = m_new + + # HCA split-mask pool branch. The local branch above prunes the SWA + # prefix; the pool suffix is short (P = S / 128 for V4-Flash) and + # uses the caller-provided pool-only visibility mask [Sq, P]. + if HAS_ADD_MASK and HCA_LOCAL_SEQLEN > 0: + for n_start in range(HCA_LOCAL_SEQLEN, seqlen_k, BLOCK_N): + offs_n = n_start + tl.arange(0, BLOCK_N) + pool_n = offs_n - HCA_LOCAL_SEQLEN + + k_ptrs = ( + K + + bid * stride_kb + + khid * stride_kh + + offs_n[:, None] * stride_kn + + offs_d[None, :] * stride_kd + ) + k_load_mask = offs_n[:, None] < seqlen_k + k = tl.load(k_ptrs, mask=k_load_mask, other=0.0) + + qk = tl.dot(q, tl.trans(k)) * sm_scale + + mask_ptrs = ADD_MASK + offs_m[:, None] * stride_ms + pool_n[None, :] * stride_mn + mask_load_mask = (offs_m[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k) + add_bias = tl.load(mask_ptrs, mask=mask_load_mask, other=0.0).to(tl.float32) + qk = qk + add_bias + qk = tl.where(offs_n[None, :] < seqlen_k, qk, NEG_INF) + + m_ij = tl.max(qk, 1) + m_new = tl.maximum(m_i, m_ij) + alpha = tl.exp(m_i - m_new) + p = tl.exp(qk - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p, 1) + + v_ptrs = ( + V + + bid * stride_vb + + khid * stride_vh + + offs_n[:, None] * stride_vn + + offs_d[None, :] * stride_vd + ) + v_load_mask = offs_n[:, None] < seqlen_k + v = tl.load(v_ptrs, mask=v_load_mask, other=0.0) + + # P57 R2: MFMA-acc fusion -- see local-branch comment above. + acc = tl.dot(p.to(v.dtype), v, acc=acc * alpha[:, None]) + m_i = m_new + + # Sink: virtual key column with value 0. The max-subtract trick + # uses sink as a candidate row maximum; sink contributes to + # l_running but NOT to acc (its notional value is zero). + if HAS_SINK: + sink_h = tl.load(SINK + qhid).to(tl.float32) + m_new = tl.maximum(m_i, sink_h) + alpha = tl.exp(m_i - m_new) + beta = tl.exp(sink_h - m_new) + l_i = l_i * alpha + beta + acc = acc * alpha[:, None] + m_i = m_new + + # Final divide + cast back to output dtype. + out = acc / l_i[:, None] + lse = m_i + tl.log(l_i) + + out_ptrs = ( + OUT + bid * stride_ob + qhid * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od + ) + tl.store(out_ptrs, out.to(OUT.dtype.element_ty), mask=offs_m[:, None] < seqlen_q) + + lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm + tl.store(lse_ptrs, lse, mask=offs_m < seqlen_q) + + +# --------------------------------------------------------------------------- +# Python launcher +# --------------------------------------------------------------------------- + + +def _launch_v4_attention_fwd( + q: torch.Tensor, # [B, H, Sq, D] + k: torch.Tensor, # [B, K_H, Sk, D] K_H ∈ {1, H} + v: torch.Tensor, # [B, K_H, Sk, D] + *, + sink: Optional[torch.Tensor], # [H] or None + swa_window: int, + additive_mask: Optional[torch.Tensor], # [Sq, Sk] or None + scale: float, + hca_local_seqlen: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Launch the V4 attention forward kernel. + + Returns ``(out, lse)`` where ``out`` matches ``v.dtype`` and + ``lse`` is fp32. ``lse`` is what the BWD kernel needs to + re-materialise the softmax without storing the ``[Sq, Sk]`` ``P`` + matrix. + + The launcher does NO autograd bookkeeping — it's a thin wrapper + around the kernel suitable for the + :class:`V4AttentionFn` autograd Function and for unit tests. + """ + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise ValueError( + "v4_attention_v1 forward expects q / k / v of rank 4 " + f"(got q.dim={q.dim()}, k.dim={k.dim()}, v.dim={v.dim()})" + ) + B, HQ, Sq, D = q.shape + Bk, HK, Sk, Dk = k.shape + Bv, HKv, Skv, Dv = v.shape + if (Bk, Sk, Dk) != (B, Sk, D) or (Bv, HKv, Skv, Dv) != (Bk, HK, Sk, D): + raise ValueError( + "v4_attention_v1 shape mismatch: " f"q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}" + ) + if HK != 1 and HK != HQ: + raise ValueError(f"v4_attention_v1 requires K_H ∈ {{1 (MQA), {HQ} (MHA)}}; got K_H={HK}.") + if not q.is_cuda or not k.is_cuda or not v.is_cuda: + raise ValueError("v4_attention_v1 requires CUDA / HIP tensors.") + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError( + "v4_attention_v1 requires q.dtype == k.dtype == v.dtype " + f"(got {q.dtype} / {k.dtype} / {v.dtype})." + ) + + has_sink = sink is not None + has_add_mask = additive_mask is not None + hca_local_seqlen = int(hca_local_seqlen) + if hca_local_seqlen: + if not has_add_mask: + raise ValueError("hca_local_seqlen requires a pool additive_mask.") + if hca_local_seqlen <= 0 or hca_local_seqlen >= Sk: + raise ValueError( + "hca_local_seqlen must split local and pool keys " + f"(got hca_local_seqlen={hca_local_seqlen}, Sk={Sk})." + ) + expected_mask_shape = (Sq, Sk - hca_local_seqlen) + if tuple(additive_mask.shape) != expected_mask_shape: + raise ValueError( + "HCA split-mask mode expects additive_mask shape " + f"{expected_mask_shape}, got {tuple(additive_mask.shape)}." + ) + if swa_window <= 0: + raise ValueError("HCA split-mask mode requires swa_window > 0.") + + # Mask precedence: additive_mask wins over swa_window. When neither + # is set, USE_CAUSAL = True (eager / V4 default). + use_causal = (not has_add_mask) and (swa_window <= 0) + swa_window_constexpr = ( + int(swa_window) if ((not has_add_mask or hca_local_seqlen) and swa_window > 0) else 0 + ) + + out = torch.empty_like(q) + lse = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32) + + # P57 R2: tile sweep on V4-Flash widths (B=1, H=64, S=4096, D=512, + # SWA=128) finds (BLOCK_M=64, BLOCK_N=16, num_warps=8, num_stages=2) + # wins by **~17 %** on cr=4 FWD (1.69 -> 1.43 ms) and **~38 %** on + # cr=0 FWD (0.79 -> 0.49 ms). The intuition: + # + # * BLOCK_M=64 halves the program grid (vs the R1 32x32 layout) so + # the SWA K-tile prefix is shared across twice as many query rows. + # Q is loaded once per program -- bigger M tiles amortise the + # Q-load over more K-tile iterations. + # * BLOCK_N=16 keeps the per-stage K/V tile at 16x512x2 = 16 KiB so + # num_stages=2 double-buffering fits the LDS budget: + # 64 (Q) + 16*2 (K) + 16*2 (V) = 128 KiB << 160 KiB MI355X limit. + # The R1 32x32 layout already needed 96 KiB at num_stages=1 and + # would have hit 160 KiB at num_stages=2 -- no headroom. + # * num_stages=2 software-pipelines K/V loads against the + # ``tl.dot`` + softmax update, hiding the HBM gather latency + # that previously serialised the inner loop. + # + # Env-overridable so future shape regressions can fall back without + # rebuilding. The defaults are the per-shape winner from the R2 + # sweep (`p57/r2_sweep_local.sh`). + # Arch-aware defaults (see ._v4_attn_tuning); env knobs still override. + # gfx950/MI355X: HCA fwd wins at BLOCK_M=128, pure SWA wants 64 (regresses at 128). + # gfx1250: SWA wins at BM=128/BN=32/W4/S1 (+62-70% vs gfx950; ab_sweep/opt7c). + from ._v4_attn_tuning import fwd_attn_defaults + + _bm, _bn, _w, _s = fwd_attn_defaults(is_hca=bool(hca_local_seqlen)) + BLOCK_M = int(os.getenv("PRIMUS_V4_ATTN_FWD_BLOCK_M", str(_bm))) + BLOCK_N = int(os.getenv("PRIMUS_V4_ATTN_FWD_BLOCK_N", str(_bn))) + NUM_WARPS_FWD = int(os.getenv("PRIMUS_V4_ATTN_FWD_WARPS", str(_w))) + NUM_STAGES_FWD = int(os.getenv("PRIMUS_V4_ATTN_FWD_STAGES", str(_s))) + BLOCK_DMODEL = D # head_dim must be a power of 2 for tl.dot + + grid = (triton.cdiv(Sq, BLOCK_M), B * HQ) + + # Sentinel pointers when sink / mask are absent. Triton requires a + # real tensor — we pass q (any tensor) and gate via the constexpr. + sink_ptr = sink if has_sink else q + mask_ptr = additive_mask if has_add_mask else q + if has_add_mask: + stride_ms = additive_mask.stride(0) + stride_mn = additive_mask.stride(1) + else: + stride_ms = 0 + stride_mn = 0 + + _v4_attention_fwd_kernel[grid]( + q, + k, + v, + out, + lse, + sink_ptr, + mask_ptr, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + stride_ms, + stride_mn, + Sq, + Sk, + float(scale), + HEAD_Q=HQ, + HEAD_K=HK, + SWA_WINDOW=swa_window_constexpr, + HAS_SINK=has_sink, + HAS_ADD_MASK=has_add_mask, + HCA_LOCAL_SEQLEN=hca_local_seqlen, + USE_CAUSAL=use_causal, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=NUM_WARPS_FWD, + num_stages=NUM_STAGES_FWD, + ) + return out, lse + + +__all__ = [ + "_v4_attention_fwd_kernel", + "_launch_v4_attention_fwd", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention.py new file mode 100644 index 000000000..8310d707a --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention.py @@ -0,0 +1,175 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 CSA (pool) attention — Triton **v1** autograd entry point (production, cr=4). + +In-kernel top-k gather + pool scatter-add. Moved from the former top-level +``v4_csa_attention_v0.py`` during the triton v0/v1/v2 reorg; pairs with the +dense/HCA ``v4_attention_v1`` (also v1). See ``_triton_v0_deprecated`` for the deprecated +gathered path and ``_triton_v2`` for the fused sparse-MLA path. +""" +from __future__ import annotations + +from typing import Optional + +import torch + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention import ( + v4_attention_v1, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_csa_attention_bwd import ( + _launch_v4_csa_attention_pool_bwd, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_csa_attention_fwd import ( + _launch_v4_csa_attention_pool_fwd, +) + + +class V4CSAPoolAttentionFn(torch.autograd.Function): + """Triton CSA attention with in-kernel topk gather and pool scatter-add.""" + + @staticmethod + def forward( # type: ignore[override] + ctx, + q: torch.Tensor, + k_local: torch.Tensor, + v_local: torch.Tensor, + pool: torch.Tensor, + topk_idxs: torch.Tensor, + sink: Optional[torch.Tensor], + swa_window: int, + attn_dropout: float, + training: bool, + scale: float, + ) -> torch.Tensor: + if attn_dropout > 0.0 and training: + raise NotImplementedError( + "v4_csa_attention_v1 does not implement in-kernel attention " + "dropout (V4 trains with attn_dropout=0). Got " + f"attn_dropout={attn_dropout}, training={training}." + ) + + out, lse = _launch_v4_csa_attention_pool_fwd( + q, + k_local, + v_local, + pool, + topk_idxs, + sink=sink, + swa_window=swa_window, + scale=scale, + ) + ctx.save_for_backward(q, k_local, v_local, pool, topk_idxs, out, lse, sink) + ctx.swa_window = int(swa_window) + ctx.attn_dropout = float(attn_dropout) + ctx.training_mode = bool(training) + ctx.scale = float(scale) + ctx.sink_was_none = sink is None + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): # type: ignore[override] + q, k_local, v_local, pool, topk_idxs, out, lse, sink = ctx.saved_tensors + sink_arg = None if ctx.sink_was_none else sink + + if not grad_out.is_contiguous(): + grad_out = grad_out.contiguous() + + dq, dk_local, dv_local, dpool, dsink = _launch_v4_csa_attention_pool_bwd( + q, + k_local, + v_local, + pool, + topk_idxs, + out, + grad_out, + lse, + sink=sink_arg, + swa_window=ctx.swa_window, + scale=ctx.scale, + ) + + if not ctx.needs_input_grad[0]: + dq = None + if not ctx.needs_input_grad[1]: + dk_local = None + if not ctx.needs_input_grad[2]: + dv_local = None + if not ctx.needs_input_grad[3]: + dpool = None + if not ctx.needs_input_grad[5] or ctx.sink_was_none: + dsink = None + + # Forward signature: (q, k_local, v_local, pool, topk_idxs, sink, + # swa_window, attn_dropout, training, scale). + return dq, dk_local, dv_local, dpool, None, dsink, None, None, None, None + + +def v4_csa_attention_v1( + q: torch.Tensor, # [B, H, Sq, D] + k_local: torch.Tensor, # [B, H, Sq, D] + v_local: torch.Tensor, # [B, H, Sq, D] + pool: torch.Tensor, # [B, P, D] + *, + topk_idxs: torch.Tensor, # [B, Sq, K_topk], -1 masks a slot + sink: Optional[torch.Tensor], # [H] or None + swa_window: int, + attn_dropout: float, + training: bool, + scale: float, + use_tilelang: bool = False, + use_flydsl: bool = False, # accepted for call-site parity; no from-pool FlyDSL kernel, Triton path +) -> torch.Tensor: + """Triton-backed CSA attention that gathers sparse keys in-kernel. + + ``pool`` is the compressed-pool tensor before per-query top-K gather. + ``topk_idxs`` drives the sparse branch directly; negative entries are + masked and contribute no probability mass. The backward kernel emits + ``dpool`` with atomic scatter-add, avoiding the materialised + ``[B, Sq, K_topk, D]`` gathered tensor and its autograd scatter. + + ``use_tilelang`` is reserved for the plan-8 P54 / P55 from-pool + tilelang path (not landed); currently always falls through to + :class:`V4CSAPoolAttentionFn` (Triton). + """ + K_topk = topk_idxs.shape[2] + if K_topk == 0: + return v4_attention_v1( + q, + k_local, + v_local, + sink=sink, + swa_window=swa_window, + additive_mask=None, + attn_dropout=attn_dropout, + training=training, + scale=scale, + ) + + # Plan-8 P57 close-out 2: from-pool tilelang path is not landed + # (P54 / P55 descoped); the gated dispatcher always returns False + # so this stays on the Triton autograd path. We still consult + # the dispatcher so a future P54 / P55 landing can flip behavior + # without re-wiring this callsite. + del use_tilelang + return V4CSAPoolAttentionFn.apply( + q, + k_local, + v_local, + pool, + topk_idxs, + sink, + swa_window, + attn_dropout, + training, + scale, + ) + + +__all__ = [ + "V4CSAPoolAttentionFn", + "v4_csa_attention_v1", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_bwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_bwd.py new file mode 100644 index 000000000..030e35f76 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_bwd.py @@ -0,0 +1,2321 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 CSA attention backward Triton kernel (plan-4 P26, ``compress_ratio == 4``). + +Two-kernel design (mirroring :mod:`v4_attention_bwd`): + +* Pre-pass: ``D[b, h, m] = sum_d (dout[b,h,m,d] * out[b,h,m,d])`` — + reuses :func:`_v4_attention_bwd_preprocess_kernel` from the dense + module since the contract is identical. +* Main pass: one program per ``(b, qhid, m)`` query row; re-materialises + the joint softmax row from the saved LSE; emits + + :: + dq [B, H, Sq, D] direct store (one program per row) + dk_local [B, H, Sq, D] atomic-add (multiple m's hit same n) + dv_local [B, H, Sq, D] atomic-add + dgathered [B, Sq, K_topk, D] atomic-add (no H dim — broadcast in fwd + means all H heads contribute) + dsink [H] atomic-add per query + +dtype contract: + +* All inputs loaded in input dtype; per-row dot products reduce in fp32 + via ``.to(tl.float32)`` upcast before the multiply (matches the FWD's + bf16-tensor-core / fp32-accumulator semantics). +* The online ``P / dP / dS`` re-materialisation is fp32 (matches the + FWD's softmax-in-fp32 contract). +* Output gradients are returned in input dtype (cast from fp32 buffers + by the launcher). + +Math derivation (per query (b, h, m), see plan-4 ``02-phase-details.md`` +Phase 26 section): + + joint_logits = cat(qk_local, qk_sparse, sink_h) + P_j = exp(joint_logits[j] - lse) + out_d = sum_n P_local[n] * v_local[n,d] + sum_k P_sparse[k] * g[k,d] + D = sum_d (dout[d] * out[d]) + dP_local[n] = sum_d (dout[d] * v_local[n,d]) + dP_sparse[k] = sum_d (dout[d] * g[k,d]) + dS_local[n] = P_local[n] * (dP_local[n] - D) + dS_sparse[k] = P_sparse[k] * (dP_sparse[k] - D) + dS_sink = -P_sink * D # sink val is 0 + + dq[d] = sum_n dS_local[n] * scale * k_local[n,d] + + sum_k dS_sparse[k] * scale * g[k,d] + dk_local[n,d] += dS_local[n] * scale * q[d] + dv_local[n,d] += P_local[n] * dout[d] + dgathered[k,d] += dS_sparse[k] * scale * q[d] + + P_sparse[k] * dout[d] # both branches + dsink_h += dS_sink + +Edge cases: + +* ``K_topk == 0`` — the wrapper short-circuits to the dense + :func:`v4_attention_v1` BWD before reaching this kernel. +* All-masked tile rows — uses ``NEG_INF = -1e30`` finite sentinel so + ``exp(NEG_INF - lse) = exp(-large) ≈ 0`` for fully-masked positions + (matches the FWD). +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch +import triton +import triton.language as tl + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_bwd import ( + _v4_attention_bwd_dkv_kernel, + _v4_attention_bwd_dq_kernel, + _v4_attention_bwd_kernel, + _v4_attention_bwd_preprocess_kernel, +) + +# --------------------------------------------------------------------------- +# Main BWD kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_csa_attention_pool_bwd_kernel( + Q, + K_LOCAL, + V_LOCAL, + POOL, + TOPK_IDXS, + DOUT, + LSE, + D, + DQ, + DK_LOCAL, + DV_LOCAL, + DPOOL, + DSINK, + SINK, + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_klb, + stride_klh, + stride_kln, + stride_kld, + stride_vlb, + stride_vlh, + stride_vln, + stride_vld, + stride_pb, + stride_pp, + stride_pd, + stride_tib, + stride_tim, + stride_tik, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dqb, + stride_dqh, + stride_dqm, + stride_dqd, + stride_dklb, + stride_dklh, + stride_dkln, + stride_dkld, + stride_dvlb, + stride_dvlh, + stride_dvln, + stride_dvld, + stride_dpb, + stride_dpp, + stride_dpd, + seqlen_q, + pool_size, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + SWA_WINDOW: tl.constexpr, + HAS_SINK: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + STORE_DPOOL: tl.constexpr, + LOCAL_ONLY: tl.constexpr, +): + """CSA BWD with in-kernel scatter-add into the compressed-pool gradient.""" + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + bid = pid_bh // HEAD_Q + qhid = pid_bh % HEAD_Q + + offs_d = tl.arange(0, BLOCK_DMODEL) + NEG_INF: tl.constexpr = -1.0e30 + q_active = pid_m < seqlen_q + + q_row_offset = bid * stride_qb + qhid * stride_qh + pid_m * stride_qm + q = tl.load(Q + q_row_offset + offs_d * stride_qd, mask=q_active, other=0.0) + + do_row_offset = bid * stride_dob + qhid * stride_doh + pid_m * stride_dom + dout = tl.load(DOUT + do_row_offset + offs_d * stride_dod, mask=q_active, other=0.0) + q_f = q.to(tl.float32) + dout_f = dout.to(tl.float32) + + lse = tl.load( + LSE + bid * stride_lb + qhid * stride_lh + pid_m * stride_lm, + mask=q_active, + other=0.0, + ) + dvec = tl.load( + D + bid * stride_db + qhid * stride_dh + pid_m * stride_dm, + mask=q_active, + other=0.0, + ) + + if HAS_SINK: + sink_h = tl.load(SINK + qhid).to(tl.float32) + p_sink = tl.exp(sink_h - lse) + tl.atomic_add(DSINK + qhid, tl.where(q_active, -p_sink * dvec, 0.0)) + + dq = tl.zeros([BLOCK_DMODEL], dtype=tl.float32) + + n_loop_end = pid_m + 1 + if n_loop_end > seqlen_q: + n_loop_end = seqlen_q + if SWA_WINDOW > 0: + n_lo_raw = pid_m - SWA_WINDOW + 1 + if n_lo_raw < 0: + n_lo_raw = 0 + n_loop_start = (n_lo_raw // BLOCK_N) * BLOCK_N + else: + n_loop_start = 0 + + for n_start in range(n_loop_start, n_loop_end, BLOCK_N): + offs_n = n_start + tl.arange(0, BLOCK_N) + + kl_ptrs = ( + K_LOCAL + + bid * stride_klb + + qhid * stride_klh + + offs_n[:, None] * stride_kln + + offs_d[None, :] * stride_kld + ) + kl_load_mask = offs_n[:, None] < seqlen_q + kl = tl.load(kl_ptrs, mask=kl_load_mask, other=0.0) + + vl_ptrs = ( + V_LOCAL + + bid * stride_vlb + + qhid * stride_vlh + + offs_n[:, None] * stride_vln + + offs_d[None, :] * stride_vld + ) + vl = tl.load(vl_ptrs, mask=kl_load_mask, other=0.0) + + kl_f = kl.to(tl.float32) + qk = tl.sum(kl_f * q_f[None, :], axis=1) * sm_scale + if SWA_WINDOW > 0: + in_window = (offs_n >= pid_m - SWA_WINDOW + 1) & (offs_n <= pid_m) + else: + in_window = offs_n <= pid_m + qk = tl.where(in_window & (offs_n < seqlen_q) & q_active, qk, NEG_INF) + + p = tl.exp(qk - lse) + vl_f = vl.to(tl.float32) + dp = tl.sum(dout_f[None, :] * vl_f, axis=1) + ds = p * (dp - dvec) + + dq += tl.sum(ds[:, None] * kl_f, axis=0) * sm_scale + + dk_contrib = ds[:, None] * sm_scale * q_f[None, :] + dk_ptrs = ( + DK_LOCAL + + bid * stride_dklb + + qhid * stride_dklh + + offs_n[:, None] * stride_dkln + + offs_d[None, :] * stride_dkld + ) + tl.atomic_add(dk_ptrs, dk_contrib, mask=kl_load_mask, sem="relaxed") + + dv_contrib = p[:, None] * dout_f[None, :] + dv_ptrs = ( + DV_LOCAL + + bid * stride_dvlb + + qhid * stride_dvlh + + offs_n[:, None] * stride_dvln + + offs_d[None, :] * stride_dvld + ) + tl.atomic_add(dv_ptrs, dv_contrib, mask=kl_load_mask, sem="relaxed") + + if not LOCAL_ONLY: + for k_start in range(0, K_topk, BLOCK_K): + offs_k = k_start + tl.arange(0, BLOCK_K) + topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + offs_k * stride_tik + topk = tl.load(topk_ptrs, mask=offs_k < K_topk, other=-1) + valid = (offs_k < K_topk) & (topk >= 0) & (topk < pool_size) + safe_topk = tl.where(valid, topk, 0) + + pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd + pool = tl.load(pool_ptrs, mask=valid[:, None], other=0.0) + pool_f = pool.to(tl.float32) + + qk_sparse = tl.sum(pool_f * q_f[None, :], axis=1) * sm_scale + qk_sparse = tl.where(valid & q_active, qk_sparse, NEG_INF) + + p = tl.exp(qk_sparse - lse) + dp = tl.sum(dout_f[None, :] * pool_f, axis=1) + ds = p * (dp - dvec) + + dq += tl.sum(ds[:, None] * pool_f, axis=0) * sm_scale + + if STORE_DPOOL: + dpool_contrib = ds[:, None] * sm_scale * q_f[None, :] + p[:, None] * dout_f[None, :] + dpool_ptrs = ( + DPOOL + bid * stride_dpb + safe_topk[:, None] * stride_dpp + offs_d[None, :] * stride_dpd + ) + tl.atomic_add(dpool_ptrs, dpool_contrib, mask=valid[:, None], sem="relaxed") + + dq_offset = bid * stride_dqb + qhid * stride_dqh + pid_m * stride_dqm + tl.store(DQ + dq_offset + offs_d * stride_dqd, dq, mask=q_active) + + +@triton.jit +def _v4_csa_attention_pool_sparse_bwd_kernel( + Q, + POOL, + TOPK_IDXS, + DOUT, + LSE, + D, + DQ, + DPOOL, + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_pb, + stride_pp, + stride_pd, + stride_tib, + stride_tim, + stride_tik, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dqb, + stride_dqh, + stride_dqm, + stride_dqd, + stride_dp_part, + stride_dpb, + stride_dpp, + stride_dpd, + partition_size, + seqlen_q, + pool_size, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + STORE_DPOOL: tl.constexpr = True, +): + """Sparse CSA BWD using a head block so pool work maps to tl.dot. + + Plan-5 P32: ``DPOOL`` is laid out as ``[N_PART, B, P, D]`` instead + of ``[B, P, D]``. Each program writes its dpool contributions into + ``DPOOL[pid_m // partition_size, bid, ...]``. With ``N_PART = Sq / + partition_size``, the per-cache-line atomic contention drops by a + factor of ``N_PART`` because writes from different m-partitions hit + different DRAM rows. The launcher reduces the partial axis on the + Python side, which is bandwidth-cheap (``2 * N_PART * P * D * + sizeof(fp32)``). + """ + pid_m = tl.program_id(0) + pid_h_block = tl.program_id(1) + bid = tl.program_id(2) + partition_id = pid_m // partition_size + + offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H) + offs_k = tl.arange(0, BLOCK_K) + offs_d = tl.arange(0, BLOCK_DMODEL) + h_mask = offs_h < HEAD_Q + q_active = pid_m < seqlen_q + + q_ptrs = ( + Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd + ) + q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + dout_ptrs = ( + DOUT + + bid * stride_dob + + offs_h[:, None] * stride_doh + + pid_m * stride_dom + + offs_d[None, :] * stride_dod + ) + dout = tl.load(dout_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + lse = tl.load( + LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm, + mask=h_mask & q_active, + other=0.0, + ) + dvec = tl.load( + D + bid * stride_db + offs_h * stride_dh + pid_m * stride_dm, + mask=h_mask & q_active, + other=0.0, + ) + + dq = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32) + + for k_start in range(0, K_topk, BLOCK_K): + sparse_k = k_start + offs_k + topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik + topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1) + valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size) + safe_topk = tl.where(valid_k, topk, 0) + + pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd + pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0) + + q_bf16 = q.to(pool.dtype) + dout_bf16 = dout.to(pool.dtype) + qk = tl.dot(q_bf16, tl.trans(pool)) * sm_scale + qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, -1.0e30) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout_bf16, tl.trans(pool)) + ds = p * (dp - dvec[:, None]) + + dq += tl.dot(ds.to(pool.dtype), pool) * sm_scale + + if STORE_DPOOL: + dpool_contrib = tl.dot(tl.trans(ds.to(q.dtype)), q_bf16) * sm_scale + dpool_contrib += tl.dot(tl.trans(p.to(dout.dtype)), dout_bf16) + dpool_ptrs = ( + DPOOL + + partition_id * stride_dp_part + + bid * stride_dpb + + safe_topk[:, None] * stride_dpp + + offs_d[None, :] * stride_dpd + ) + tl.atomic_add(dpool_ptrs, dpool_contrib, mask=valid_k[:, None], sem="relaxed") + else: + # Force the optimiser to keep this branch as a no-op so the + # ``STORE_DPOOL=False`` variant truly drops the two trailing + # ``tl.dot`` ops + atomic from the IR. + pass + + # Plan-5 P32: each ``(pid_m, pid_h_block, bid)`` program writes to a + # disjoint slice of ``DQ`` (different ``m`` rows + different head + # blocks), so a plain ``tl.store`` is safe and ~2× cheaper than the + # atomic on MI355. + dq_ptrs = ( + DQ + + bid * stride_dqb + + offs_h[:, None] * stride_dqh + + pid_m * stride_dqm + + offs_d[None, :] * stride_dqd + ) + tl.store(dq_ptrs, dq, mask=h_mask[:, None] & q_active) + + +@triton.jit +def _v4_csa_attention_pool_sparse_bwd_dq_only_kernel( + Q, + POOL, + TOPK_IDXS, + DOUT, + LSE, + D, + DQ, + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_pb, + stride_pp, + stride_pd, + stride_tib, + stride_tim, + stride_tik, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dqb, + stride_dqh, + stride_dqm, + stride_dqd, + seqlen_q, + pool_size, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """P57: dq-only sparse CSA BWD (no dpool_partial write). + + Mirrors :func:`_v4_csa_attention_pool_sparse_bwd_partial_kernel` + but drops the per-visit ``dpool_partial`` write and the two + ``dpool_contrib`` matmuls. The freed register file lets us run + with a larger ``BLOCK_K`` / more ``num_stages`` so the per-program + latency drops sharply. The dpool partial is produced by a sibling + ``_v4_csa_attention_pool_sparse_bwd_dpool_only_kernel`` that + omits the ``dq`` accumulator instead, and the segreduce kernel + folds the partial into ``dpool[B, P, D]``. + + The reason for splitting: the joint kernel's wall-clock at the + proxy shape is ~3.7 ms, dominated by the 4 GiB ``dpool_partial`` + write *interleaved* with the ``[BLOCK_H, BLOCK_DMODEL]`` fp32 + ``dq`` accumulator (~64 KB live across the K-loop). Splitting + drops the live-register footprint of each kernel ~2×, letting + the compiler use higher ``num_stages`` for prefetching. + """ + pid_m = tl.program_id(0) + pid_h_block = tl.program_id(1) + bid = tl.program_id(2) + + offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H) + offs_k = tl.arange(0, BLOCK_K) + offs_d = tl.arange(0, BLOCK_DMODEL) + h_mask = offs_h < HEAD_Q + q_active = pid_m < seqlen_q + + q_ptrs = ( + Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd + ) + q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + dout_ptrs = ( + DOUT + + bid * stride_dob + + offs_h[:, None] * stride_doh + + pid_m * stride_dom + + offs_d[None, :] * stride_dod + ) + dout = tl.load(dout_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + lse = tl.load( + LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm, + mask=h_mask & q_active, + other=0.0, + ) + dvec = tl.load( + D + bid * stride_db + offs_h * stride_dh + pid_m * stride_dm, + mask=h_mask & q_active, + other=0.0, + ) + + dq = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32) + + # P57 R2: scale-defer + acc-form MFMA (mirror of the joint + # partial kernel). See the comment in + # ``_v4_csa_attention_pool_sparse_bwd_partial_kernel`` for the + # math. + pool_dtype = POOL.dtype.element_ty + q_scaled = (q.to(tl.float32) * sm_scale).to(pool_dtype) + dout_bf16 = dout.to(pool_dtype) + + for k_start in range(0, K_topk, BLOCK_K): + sparse_k = k_start + offs_k + topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik + topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1) + valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size) + safe_topk = tl.where(valid_k, topk, 0) + + pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd + pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0) + + qk = tl.dot(q_scaled, tl.trans(pool)) + qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, -1.0e30) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout_bf16, tl.trans(pool)) + ds = p * (dp - dvec[:, None]) + + dq = tl.dot(ds.to(pool_dtype), pool, acc=dq) + + dq = dq * sm_scale + + dq_ptrs = ( + DQ + + bid * stride_dqb + + offs_h[:, None] * stride_dqh + + pid_m * stride_dqm + + offs_d[None, :] * stride_dqd + ) + tl.store(dq_ptrs, dq, mask=h_mask[:, None] & q_active) + + +@triton.jit +def _v4_csa_attention_pool_sparse_bwd_dpool_only_kernel( + Q, + POOL, + TOPK_IDXS, + DOUT, + LSE, + D, + DPOOL_PARTIAL, # [B, M, K_topk, D] fp32 — NO atomics, each program owns its slice + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_pb, + stride_pp, + stride_pd, + stride_tib, + stride_tim, + stride_tik, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dpb, + stride_dpm, + stride_dpk, + stride_dpd, + seqlen_q, + pool_size, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """P57: dpool-partial-only sparse CSA BWD (no dq accumulator). + + Mirror of :func:`_v4_csa_attention_pool_sparse_bwd_dq_only_kernel` + that drops the ``dq`` accumulator (64 KB fp32 per program) so + the per-iter ``dpool_contrib`` matmul + write loop runs with a + tighter register footprint and can issue more in-flight HBM + writes per warp. + """ + pid_m = tl.program_id(0) + pid_h_block = tl.program_id(1) + bid = tl.program_id(2) + + offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H) + offs_k = tl.arange(0, BLOCK_K) + offs_d = tl.arange(0, BLOCK_DMODEL) + h_mask = offs_h < HEAD_Q + q_active = pid_m < seqlen_q + + q_ptrs = ( + Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd + ) + q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + dout_ptrs = ( + DOUT + + bid * stride_dob + + offs_h[:, None] * stride_doh + + pid_m * stride_dom + + offs_d[None, :] * stride_dod + ) + dout = tl.load(dout_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + lse = tl.load( + LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm, + mask=h_mask & q_active, + other=0.0, + ) + dvec = tl.load( + D + bid * stride_db + offs_h * stride_dh + pid_m * stride_dm, + mask=h_mask & q_active, + other=0.0, + ) + + # P57 R2: scale-defer (q-fold) — see partial kernel for math. + pool_dtype = POOL.dtype.element_ty + q_scaled = (q.to(tl.float32) * sm_scale).to(pool_dtype) + dout_bf16 = dout.to(pool_dtype) + + for k_start in range(0, K_topk, BLOCK_K): + sparse_k = k_start + offs_k + topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik + topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1) + valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size) + safe_topk = tl.where(valid_k, topk, 0) + + pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd + pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0) + + qk = tl.dot(q_scaled, tl.trans(pool)) + qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, -1.0e30) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout_bf16, tl.trans(pool)) + ds = p * (dp - dvec[:, None]) + + dpool_contrib = tl.dot(tl.trans(ds.to(pool_dtype)), q_scaled) + dpool_contrib = tl.dot(tl.trans(p.to(pool_dtype)), dout_bf16, acc=dpool_contrib) + dpool_contrib = tl.where(valid_k[:, None], dpool_contrib, 0.0) + dpool_partial_ptrs = ( + DPOOL_PARTIAL + + bid * stride_dpb + + pid_m * stride_dpm + + sparse_k[:, None] * stride_dpk + + offs_d[None, :] * stride_dpd + ) + tl.store( + dpool_partial_ptrs, + dpool_contrib, + mask=(sparse_k[:, None] < K_topk) & q_active, + ) + + +@triton.jit +def _v4_csa_attention_pool_sparse_bwd_partial_sorted_kernel( + Q, + POOL, + TOPK_IDXS, + INV_PERM, # [B, MK] int32 — sorted_position = inv_perm[orig_flat] + DOUT, + LSE, + D, + DQ, + DPOOL_PARTIAL, # [B, MK_sorted, D] partial buffer (bf16 / fp32) + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_pb, + stride_pp, + stride_pd, + stride_tib, + stride_tim, + stride_tik, + stride_ipb, + stride_ipi, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dqb, + stride_dqh, + stride_dqm, + stride_dqd, + stride_dpb, + stride_dpi, + stride_dpd, + seqlen_q, + pool_size, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """P57 sorted variant of the sparse CSA BWD partial kernel. + + Writes the per-visit ``dpool_contrib`` tile into the sorted-order + position of ``DPOOL_PARTIAL`` (a flat ``[B, MK, D]`` buffer) + using a host-side ``INV_PERM`` (the inverse of the segreduce + ``perm``). The downstream segreduce kernel can then read + contiguous ``[bin_start:bin_end, :]`` slices instead of gathering + via ``perm[i]`` — turning a random-access reduction into a + streaming one (better L2 / Infinity Cache reuse on MI355). + + The partial kernel writes are now scattered (one row per sparse_k + slot to a sorted position), but the HBM write bandwidth is + similar because Triton's vector store still coalesces 128 B + cache-line bursts as long as ``BLOCK_DMODEL`` is contiguous. + """ + pid_m = tl.program_id(0) + pid_h_block = tl.program_id(1) + bid = tl.program_id(2) + + offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H) + offs_k = tl.arange(0, BLOCK_K) + offs_d = tl.arange(0, BLOCK_DMODEL) + h_mask = offs_h < HEAD_Q + q_active = pid_m < seqlen_q + + q_ptrs = ( + Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd + ) + q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + dout_ptrs = ( + DOUT + + bid * stride_dob + + offs_h[:, None] * stride_doh + + pid_m * stride_dom + + offs_d[None, :] * stride_dod + ) + dout = tl.load(dout_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + lse = tl.load( + LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm, + mask=h_mask & q_active, + other=0.0, + ) + dvec = tl.load( + D + bid * stride_db + offs_h * stride_dh + pid_m * stride_dm, + mask=h_mask & q_active, + other=0.0, + ) + + dq = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32) + flat_base = pid_m * K_topk + + # P57 R2: scale-defer (q-fold) — see partial kernel for math. + pool_dtype = POOL.dtype.element_ty + q_scaled = (q.to(tl.float32) * sm_scale).to(pool_dtype) + dout_bf16 = dout.to(pool_dtype) + + for k_start in range(0, K_topk, BLOCK_K): + sparse_k = k_start + offs_k + topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik + topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1) + valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size) + safe_topk = tl.where(valid_k, topk, 0) + + pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd + pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0) + + qk = tl.dot(q_scaled, tl.trans(pool)) + qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, -1.0e30) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout_bf16, tl.trans(pool)) + ds = p * (dp - dvec[:, None]) + + dq = tl.dot(ds.to(pool_dtype), pool, acc=dq) + + dpool_contrib = tl.dot(tl.trans(ds.to(pool_dtype)), q_scaled) + dpool_contrib = tl.dot(tl.trans(p.to(pool_dtype)), dout_bf16, acc=dpool_contrib) + dpool_contrib = tl.where(valid_k[:, None], dpool_contrib, 0.0) + + # P57: look up the SORTED position for each ``(m, sparse_k)`` + # visit. ``flat_base + sparse_k`` is the ORIGINAL flat index; + # ``inv_perm[bid, orig_flat]`` is the sorted slot to write. + orig_flat = flat_base + sparse_k + inv_perm_ptrs = INV_PERM + bid * stride_ipb + orig_flat * stride_ipi + sorted_idx = tl.load(inv_perm_ptrs, mask=sparse_k < K_topk, other=0) + dpool_partial_ptrs = ( + DPOOL_PARTIAL + bid * stride_dpb + sorted_idx[:, None] * stride_dpi + offs_d[None, :] * stride_dpd + ) + tl.store( + dpool_partial_ptrs, + dpool_contrib, + mask=(sparse_k[:, None] < K_topk) & q_active, + ) + + dq = dq * sm_scale + + dq_ptrs = ( + DQ + + bid * stride_dqb + + offs_h[:, None] * stride_dqh + + pid_m * stride_dqm + + offs_d[None, :] * stride_dqd + ) + tl.store(dq_ptrs, dq, mask=h_mask[:, None] & q_active) + + +@triton.jit +def _v4_csa_attention_pool_segreduce_sequential_kernel( + DPOOL_PARTIAL, # [B, MK_sorted, D] — written in sorted order by partial_sorted kernel + BIN_PTR, # [B, P+1] int32 — prefix sum of count per pool slot (same as non-sorted variant) + DPOOL, # [B, P, D] + stride_dpb, + stride_dpi, + stride_dpd, + stride_binb, + stride_binp, + stride_db, + stride_dp, + stride_dd, + P, + D_size, + BLOCK_D: tl.constexpr, + BLOCK_I: tl.constexpr, +): + """P57 sequential segreduce — reads dpool_partial[bid, i:i+BI, :] + contiguously without a perm lookup since the partial kernel + already wrote in sorted order. + """ + pid_p = tl.program_id(0) + pid_d_block = tl.program_id(1) + pid_b = tl.program_id(2) + + offs_d = pid_d_block * BLOCK_D + tl.arange(0, BLOCK_D) + d_mask = offs_d < D_size + + bin_start = tl.load(BIN_PTR + pid_b * stride_binb + pid_p * stride_binp) + bin_end = tl.load(BIN_PTR + pid_b * stride_binb + (pid_p + 1) * stride_binp) + + acc = tl.zeros([BLOCK_D], dtype=tl.float32) + i = bin_start + while i < bin_end: + offs_i = i + tl.arange(0, BLOCK_I) + valid_i = offs_i < bin_end + partial_ptrs = ( + DPOOL_PARTIAL + pid_b * stride_dpb + offs_i[:, None] * stride_dpi + offs_d[None, :] * stride_dpd + ) + partial = tl.load( + partial_ptrs, + mask=valid_i[:, None] & d_mask[None, :], + other=0.0, + ) + acc += tl.sum(partial, axis=0) + i += BLOCK_I + + dpool_offset = pid_b * stride_db + pid_p * stride_dp + offs_d * stride_dd + tl.store(DPOOL + dpool_offset, acc, mask=d_mask) + + +@triton.jit +def _v4_csa_attention_pool_sparse_bwd_partial_kernel( + Q, + POOL, + TOPK_IDXS, + DOUT, + LSE, + D, + DQ, + DPOOL_PARTIAL, # [B, M, K_topk, D] fp32 — NO atomics, each program owns its slice + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_pb, + stride_pp, + stride_pd, + stride_tib, + stride_tim, + stride_tik, + stride_dob, + stride_doh, + stride_dom, + stride_dod, + stride_lb, + stride_lh, + stride_lm, + stride_db, + stride_dh, + stride_dm, + stride_dqb, + stride_dqh, + stride_dqm, + stride_dqd, + stride_dpb, + stride_dpm, + stride_dpk, + stride_dpd, + seqlen_q, + pool_size, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """Sparse CSA BWD that writes dpool contributions to a compact + ``[B, M, K_topk, D]`` partial buffer **without** atomics. + + Plan-5 P32: the atomic_add to the shared ``dpool[B, P, D]`` buffer + was the dominant cost (~15 ms out of ~24 ms — see + ``PRIMUS_V4_CSA_BWD_SKIP_DPOOL_ATOMIC`` profile). By emitting the + raw per-visit contributions into a compact partial buffer (4 GB + for the proxy shape, indexed by ``(b, m, k_slot)``), a follow-up + segmented reduction kernel can fold them into ``dpool[B, P, D]`` + using a sorted inverse index — atomics-free and bandwidth-bound. + """ + pid_m = tl.program_id(0) + pid_h_block = tl.program_id(1) + bid = tl.program_id(2) + + offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H) + offs_k = tl.arange(0, BLOCK_K) + offs_d = tl.arange(0, BLOCK_DMODEL) + h_mask = offs_h < HEAD_Q + q_active = pid_m < seqlen_q + + q_ptrs = ( + Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd + ) + q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + dout_ptrs = ( + DOUT + + bid * stride_dob + + offs_h[:, None] * stride_doh + + pid_m * stride_dom + + offs_d[None, :] * stride_dod + ) + dout = tl.load(dout_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + lse = tl.load( + LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm, + mask=h_mask & q_active, + other=0.0, + ) + dvec = tl.load( + D + bid * stride_db + offs_h * stride_dh + pid_m * stride_dm, + mask=h_mask & q_active, + other=0.0, + ) + + dq = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32) + + # P57 R2: hoist q/dout dtype-cast and SCALE-DEFER ``sm_scale`` out + # of the K-loop. Pre-scaling ``q_bf16`` folds the ``sm_scale`` + # factor that the kernel previously applied to ``qk`` (line: + # ``tl.dot(q_bf16, K^T) * sm_scale``) and to the first + # ``dpool_contrib`` matmul (``tl.dot(ds^T, q_bf16) * sm_scale``) + # — both are linear in q so factoring sm_scale into q is exact. + # ``dq`` accumulates ``Σ_k ds @ pool`` without sm_scale; a single + # ``dq *= sm_scale`` after the loop replaces the per-iter + # multiply on the ``[BLOCK_H, BLOCK_DMODEL]`` fp32 accumulator. + # Combined with ``tl.dot(..., acc=dq)`` (MFMA in-place + # accumulator), this drops ~3 fp32 multiplies + 1 fp32 add per + # K_topk/BLOCK_K=16 iteration on the proxy shape. + pool_dtype = POOL.dtype.element_ty + dout_bf16 = dout.to(pool_dtype) + q_scaled = (q.to(tl.float32) * sm_scale).to(pool_dtype) + + for k_start in range(0, K_topk, BLOCK_K): + sparse_k = k_start + offs_k + topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik + topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1) + valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size) + safe_topk = tl.where(valid_k, topk, 0) + + pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd + pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0) + + qk = tl.dot(q_scaled, tl.trans(pool)) + qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, -1.0e30) + + p = tl.exp(qk - lse[:, None]) + dp = tl.dot(dout_bf16, tl.trans(pool)) + ds = p * (dp - dvec[:, None]) + + dq = tl.dot(ds.to(pool_dtype), pool, acc=dq) + + # Plan-5 P32: write the per-visit dpool contribution to its own + # compact slot in ``DPOOL_PARTIAL[b, m, k_slot, :]``. No atomic + # needed because each ``(b, m, k_slot)`` slot is owned by + # exactly one program × iteration. + # P57 R2: ``q_scaled`` already carries the ``sm_scale`` factor, + # so the first matmul's ``* sm_scale`` is folded in. The + # ``acc=`` form on the second matmul fuses the fp32 add. + dpool_contrib = tl.dot(tl.trans(ds.to(pool_dtype)), q_scaled) + dpool_contrib = tl.dot(tl.trans(p.to(pool_dtype)), dout_bf16, acc=dpool_contrib) + # Zero out invalid k slots so the reduction can sum them in + # without first checking validity (the inverse index will skip + # invalid visits anyway via the sentinel sort key, but a clean + # buffer is friendlier to debugging and to scalar fallbacks). + dpool_contrib = tl.where(valid_k[:, None], dpool_contrib, 0.0) + dpool_partial_ptrs = ( + DPOOL_PARTIAL + + bid * stride_dpb + + pid_m * stride_dpm + + sparse_k[:, None] * stride_dpk + + offs_d[None, :] * stride_dpd + ) + tl.store( + dpool_partial_ptrs, + dpool_contrib, + mask=(sparse_k[:, None] < K_topk) & q_active, + ) + + # P57 R2: deferred ``sm_scale`` on the dq accumulator. Single + # ``[BLOCK_H, BLOCK_DMODEL]`` fp32 multiply replaces the per-iter + # ``* sm_scale`` inside the loop. + dq = dq * sm_scale + + dq_ptrs = ( + DQ + + bid * stride_dqb + + offs_h[:, None] * stride_dqh + + pid_m * stride_dqm + + offs_d[None, :] * stride_dqd + ) + tl.store(dq_ptrs, dq, mask=h_mask[:, None] & q_active) + + +@triton.jit +def _v4_csa_attention_pool_segreduce_kernel( + DPOOL_PARTIAL, # [B, M, K_topk, D] fp32 — compact partial buffer + SORTED_PERM, # [B, M*K_topk] int32 — sorted (m*K + k) per pool slot + BIN_PTR, # [B, P+1] int32 — prefix sum of count per pool slot + DPOOL, # [B, P, D] fp32 — output (single tl.store per p, d) + stride_dpb, + stride_dpmk, + stride_dpd, + stride_permb, + stride_permi, + stride_binb, + stride_binp, + stride_db, + stride_dp, + stride_dd, + P, + D_size, + BLOCK_D: tl.constexpr, + BLOCK_I: tl.constexpr, # tiles over the visit indices to expose ILP +): + """Segmented reduction: ``DPOOL[b, p, :] = Σ_i DPOOL_PARTIAL[b, perm[i], :]`` + for ``i ∈ [bin_ptr[b, p], bin_ptr[b, p+1])``. + + Plan-5 P32: one program per ``(b, p, d_block)``. Different + programs write disjoint output slices, so the writes are plain + ``tl.store``\\ s — no atomics. The visit indices for slot ``p`` + are stored consecutively in ``SORTED_PERM`` (built once on the + host by sorting ``topk_idxs``), so loads from + ``DPOOL_PARTIAL[..., flat_idx, :]`` are reasonably coalescable + after the sort. ``BLOCK_D`` may exceed the actual head dim + ``D_size`` (the default proxy shape has ``D = 512`` but the unit + tests use ``D = 32``), so all dpool / dpool_partial accesses are + masked by ``offs_d < D_size``. + """ + pid_p = tl.program_id(0) + pid_d_block = tl.program_id(1) + pid_b = tl.program_id(2) + + offs_d = pid_d_block * BLOCK_D + tl.arange(0, BLOCK_D) + d_mask = offs_d < D_size + + bin_start = tl.load(BIN_PTR + pid_b * stride_binb + pid_p * stride_binp) + bin_end = tl.load(BIN_PTR + pid_b * stride_binb + (pid_p + 1) * stride_binp) + + acc = tl.zeros([BLOCK_D], dtype=tl.float32) + i = bin_start + while i < bin_end: + # Load up to BLOCK_I visit indices and prefetch their fp32 + # partial rows; the tl.dot-free vectorised path nets enough + # ILP to saturate HBM bandwidth on MI355. + offs_i = i + tl.arange(0, BLOCK_I) + valid_i = offs_i < bin_end + flat_idx = tl.load( + SORTED_PERM + pid_b * stride_permb + offs_i * stride_permi, + mask=valid_i, + other=0, + ) + partial_ptrs = ( + DPOOL_PARTIAL + + pid_b * stride_dpb + + flat_idx[:, None] * stride_dpmk + + offs_d[None, :] * stride_dpd + ) + partial = tl.load( + partial_ptrs, + mask=valid_i[:, None] & d_mask[None, :], + other=0.0, + ) + acc += tl.sum(partial, axis=0) + i += BLOCK_I + + dpool_offset = pid_b * stride_db + pid_p * stride_dp + offs_d * stride_dd + tl.store(DPOOL + dpool_offset, acc, mask=d_mask) + + +# --------------------------------------------------------------------------- +# Python launcher +# --------------------------------------------------------------------------- + + +def _launch_v4_csa_attention_pool_bwd( + q: torch.Tensor, # [B, H, Sq, D] + k_local: torch.Tensor, # [B, H, Sq, D] + v_local: torch.Tensor, # [B, H, Sq, D] + pool: torch.Tensor, # [B, P, D] + topk_idxs: torch.Tensor, # [B, Sq, K_topk] + out: torch.Tensor, # [B, H, Sq, D] + dout: torch.Tensor, # [B, H, Sq, D] + lse: torch.Tensor, # [B, H, Sq] + *, + sink: Optional[torch.Tensor], + swa_window: int, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Launch CSA backward with in-kernel scatter-add into ``pool.grad``.""" + if not q.is_cuda: + raise ValueError("v4_csa_attention_v0 pool BWD requires CUDA / HIP tensors.") + if dout.shape != out.shape or out.shape != q.shape: + raise ValueError( + "v4_csa_attention_v0 pool BWD shape mismatch: " + f"out={tuple(out.shape)}, dout={tuple(dout.shape)}, q={tuple(q.shape)}" + ) + if pool.dim() != 3: + raise ValueError( + f"v4_csa_attention_v0 pool BWD expects pool rank 3 [B, P, D], got {tuple(pool.shape)}." + ) + if topk_idxs.dim() != 3: + raise ValueError( + f"v4_csa_attention_v0 pool BWD expects topk_idxs rank 3 [B, Sq, K], got {tuple(topk_idxs.shape)}." + ) + + B, HQ, Sq, D = q.shape + Bp, P, Dp = pool.shape + Bt, Sqt, K_topk = topk_idxs.shape + if Bp != B or Dp != D or Bt != B or Sqt != Sq: + raise ValueError( + "v4_csa_attention_v0 pool BWD shape mismatch: " + f"q={tuple(q.shape)}, pool={tuple(pool.shape)}, topk_idxs={tuple(topk_idxs.shape)}" + ) + if topk_idxs.dtype not in (torch.int32, torch.int64): + raise ValueError(f"v4_csa_attention_v0 topk_idxs must be int32/int64, got {topk_idxs.dtype}.") + + has_sink = sink is not None + + BLOCK_N = 32 + # Plan-5 P32: sweeping ``BLOCK_K ∈ {32, 64, 128}`` and ``num_warps + # ∈ {4, 8}`` on the EP8 microbench picked ``BLOCK_K=32`` / + # ``num_warps=4`` as the cheapest configuration (24.9 ms vs 26.0 ms + # at the previous defaults). Override via + # ``PRIMUS_V4_CSA_BWD_BLOCK_K`` for shape-specific tuning. + BLOCK_K = int(os.getenv("PRIMUS_V4_CSA_BWD_BLOCK_K", "32")) + # Plan-5 P32: the segreduce / partial path is markedly more + # register-pressured than the atomic-add gather kernel (the extra + # ``dpool_partial`` write tile pins 64 × 512 fp32 in flight per + # warp). A second sweep at ``BLOCK_K_PARTIAL ∈ {8, 16, 32, 64, + # 128}`` × ``num_warps ∈ {4, 8, 16}`` showed ``BLOCK_K=16`` / + # ``num_warps=8`` is ~2 ms faster than the gather path's + # ``BLOCK_K=32`` choice (16.2 ms vs 18.6 ms). Keep the gather + # path's ``BLOCK_K=32`` since it has its own register profile. + BLOCK_K_PARTIAL = int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_BLOCK_K", "16")) + BLOCK_DMODEL = D + use_split_sparse = os.getenv("PRIMUS_V4_CSA_BWD_SPLIT_SPARSE", "1") != "0" + + # Plan-5 P32: enable stream-level overlap between the local SWA + # BWD (dq + dk/dv kernels) and the sparse pool BWD kernel. The two + # paths write to disjoint output buffers (``dq_local_fp32`` / + # ``dk_local_fp32`` / ``dv_local_fp32`` / ``dsink_fp32`` vs + # ``dq_sparse_fp32`` / ``dpool_partial``) so the streams can + # progress concurrently. We sum ``dq_local + dq_sparse`` on the + # default stream after the join. + # Plan-5 P32: stream overlap defaults off — the segreduce path + # already places ``dq`` writes on disjoint buffers and the local + # kernel cost (~7.6 ms) is shorter than the sparse path (~3 ms in + # the segreduce variant), so there is little compute to hide. Set + # ``PRIMUS_V4_CSA_BWD_STREAM_OVERLAP=1`` to re-enable for + # experimentation. P57 R2 confirmed that enabling overlap + # actually doubles wall-clock to ~11 ms on the proxy shape — the + # two paths both saturate HBM read bandwidth, so concurrent + # execution serializes on memory traffic and adds stream + # synchronization overhead. + use_stream_overlap = use_split_sparse and os.getenv("PRIMUS_V4_CSA_BWD_STREAM_OVERLAP", "0") != "0" + + # P57: when the input is bf16/fp16 and we use the split local SWA + # path (every slab overwritten by tl.store, no atomic_add), keep + # ``dq / dk_local / dv_local`` in the INPUT dtype so the kernels + # write directly at the final precision. Eliminates the trailing + # ~256 MB fp32 → bf16 cast on each of these 3 buffers (~0.5 ms + # total at the proxy shape) AND halves their HBM write traffic + # in-kernel (~1 GB → 512 MB for the three buffers). The + # monolithic local kernel and the gather sparse path still need + # fp32 because they ``atomic_add`` into the buffer and Triton's + # bf16 atomic_add is not supported on AMD MI355. + use_local_split_alloc = use_split_sparse and os.getenv("PRIMUS_V4_ATTN_BWD_USE_SPLIT", "0") == "1" + if use_local_split_alloc and q.dtype in (torch.bfloat16, torch.float16): + local_dq_dtype = q.dtype + else: + local_dq_dtype = torch.float32 + if use_local_split_alloc: + dq_fp32 = torch.empty((B, HQ, Sq, D), device=q.device, dtype=local_dq_dtype) + dk_local_fp32 = torch.empty((B, HQ, Sq, D), device=q.device, dtype=local_dq_dtype) + dv_local_fp32 = torch.empty((B, HQ, Sq, D), device=q.device, dtype=local_dq_dtype) + else: + dq_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dk_local_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + dv_local_fp32 = torch.zeros((B, HQ, Sq, D), device=q.device, dtype=torch.float32) + # ``dpool_fp32`` is small (2 MB at proxy) so the zero-fill cost + # is negligible (~20 us); keep ``zeros`` so the gather + + # atomic_add fallback path stays safe. + dpool_fp32 = torch.zeros((B, P, D), device=q.device, dtype=torch.float32) + # The split-sparse paths (both segreduce and gather) always write + # ``dq_sparse`` via ``tl.store`` from disjoint ``(pid_m, + # pid_h_block)`` programs, so we use a dedicated buffer and sum + # at the end. Aliasing to ``dq_fp32`` would clobber the + # local-SWA contribution. P57: match input dtype (bf16) when + # ``local_dq_dtype`` is bf16 so the final add and dtype-cast + # become free. + if use_split_sparse: + dq_sparse_fp32 = torch.empty((B, HQ, Sq, D), device=q.device, dtype=local_dq_dtype) + else: + dq_sparse_fp32 = dq_fp32 + if has_sink: + dsink_fp32 = torch.zeros((HQ,), device=q.device, dtype=torch.float32) + sink_arg = sink.to(torch.float32) if sink.dtype != torch.float32 else sink + else: + dsink_fp32 = q + sink_arg = q + + d_buf = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32) + pre_grid = (triton.cdiv(Sq, BLOCK_N), B * HQ) + _v4_attention_bwd_preprocess_kernel[pre_grid]( + out, + dout, + d_buf, + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + Sq, + HEAD=HQ, + BLOCK_M=BLOCK_N, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=4, + num_stages=1, + ) + + # Plan-5 P32: pre-compute the segreduce inverse index on the + # default stream BEFORE launching the local kernels. The sort + + # searchsorted are CPU-light but launch ~0.2 ms of kernels, and + # placing them BEFORE the local BWD launches lets the sparse + # stream's wait_event fire as soon as the index + d_buf are + # ready (instead of waiting until after all local kernels have + # been queued). + perm32 = None + bin_ptr = None + dpool_partial = None + use_segreduce = use_split_sparse and os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE", "1") == "1" + if use_segreduce: + # P57: ``dpool_partial`` defaults to the INPUT dtype when bf16 + # / fp16 — halves HBM write+read traffic on the 4 GiB partial + # buffer (4 GB fp32 → 2 GB bf16) for the ~production-shape + # bench, saving ~0.6 ms / step. fp32 inputs keep an fp32 + # partial so the parity tests' tight 1e-4 atol holds. + # The legacy P32 attempt failed parity because it used bf16 + # unconditionally; gating on the input dtype keeps the + # bf16-only speed win without regressing fp32 numerics. + env_dtype = os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_DTYPE", "") + if env_dtype == "bf16": + partial_dtype = torch.bfloat16 + elif env_dtype == "fp16": + partial_dtype = torch.float16 + elif env_dtype == "fp32": + partial_dtype = torch.float32 + else: + # Default: match input dtype for bf16 / fp16, fp32 + # otherwise. ``q.dtype`` is the most reliable proxy for + # the gradient-tolerance class. + if q.dtype in (torch.bfloat16, torch.float16): + partial_dtype = q.dtype + else: + partial_dtype = torch.float32 + # P57: experimental ``use_sorted_partial`` flips the partial + # buffer layout from ``[B, M, K_topk, D]`` to + # ``[B, MK_sorted, D]``. The downstream segreduce kernel can + # then read contiguous slices, but the partial kernel writes + # become scattered (each ``sparse_k`` row goes to its + # ``inv_perm[orig_flat]`` sorted position). On the EP8 proxy + # the natural-order partial wins (~6.43 ms vs ~6.82 ms), + # because the partial kernel is write-bound and contiguous + # writes coalesce better than scattered ones. Keep the + # sorted-partial kernels in tree for shapes where the + # segreduce-read trade-off tips the other way (e.g., tiny + # ``D`` or very-large ``P`` workloads). + use_sorted_partial = os.getenv("PRIMUS_V4_CSA_BWD_SORTED_PARTIAL", "0") != "0" + with torch.no_grad(): + MK = Sq * K_topk + flat_topk = topk_idxs.contiguous().view(B, MK).to(torch.int32) + sentinel = torch.full_like(flat_topk, P) + masked = torch.where((flat_topk >= 0) & (flat_topk < P), flat_topk, sentinel) + sorted_topk, perm = torch.sort(masked, dim=1, stable=True) + perm32 = perm.to(torch.int32) + queries = torch.arange(P + 1, device=q.device, dtype=torch.int32) + queries = queries.unsqueeze(0).expand(B, -1).contiguous() + bin_ptr = torch.searchsorted(sorted_topk, queries, right=False).to(torch.int32) + inv_perm32 = None + if use_sorted_partial: + # inv_perm[bid, orig_flat] = sorted_position. Built by + # scattering ``arange(MK)`` to the ``perm``-indexed + # positions. ~30 us at the proxy shape. + inv_perm32 = torch.empty_like(perm32) + idx_range = torch.arange(MK, device=q.device, dtype=torch.int32).unsqueeze(0).expand(B, -1) + inv_perm32.scatter_(1, perm.long(), idx_range) + dpool_partial = torch.empty((B, Sq, K_topk, D), device=q.device, dtype=partial_dtype) + + # Lazily set up an extra CUDA stream for the sparse pool BWD so it + # can overlap with the local BWD launches (default-stream serial). + sparse_stream_ctx = None + if use_stream_overlap: + sparse_stream = torch.cuda.Stream(device=q.device) + d_buf_done = torch.cuda.current_stream(q.device).record_event() + sparse_stream.wait_event(d_buf_done) + sparse_stream_ctx = torch.cuda.stream(sparse_stream) + + if use_split_sparse: + # P57: expose local-SWA tuning knobs (BLOCK_M, num_warps, + # num_stages). The local dq / dkv kernels live in + # ``v4_attention_bwd.py`` (outside the P57 file scope), but the + # LAUNCHER picks the block size + warp / stage count, and those + # parameters dominate the local-path wall clock at the + # production SWA=128 shape. The new defaults below are tuned + # specifically for ``B=1, H=64, Sq=4096, D=512, K_topk=512, + # swa_window=128`` on MI355 and cut the local dq+dkv from + # ~6.9 ms (P32 defaults: BM=32, BN=32, w=8, s=1) to ~4.1 ms + # (BM=64, BN=16 → fewer m programs and smaller n tiles to + # match the small SWA window). + local_block_m = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_BLOCK_M", "64")) + local_block_n = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_BLOCK_N", "16")) + # P57: the dq and dkv kernels prefer different per-axis block + # sizes (dq iterates ``n`` over a small SWA window so larger + # ``BLOCK_M`` packs more m-rows per program; dkv iterates + # ``m`` so its program-axis ``BLOCK_N`` and the inner-axis + # ``BLOCK_M`` can be tuned independently). + local_block_m_dq = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DQ_BLOCK_M", str(local_block_m))) + local_block_n_dq = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DQ_BLOCK_N", str(local_block_n))) + # P57 R2: dkv-specific tuning — ``BLOCK_M_dkv=16 warps=2 stages=1`` + # wins post-R2 scale-defer (the dkv kernel's m-tile prefers + # smaller block sizes when the sparse partial path runs faster + # and exposes more concurrent CU slots). BM=16 is ~10 us faster + # than BM=32 at the proxy shape. + local_block_m_dkv = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DKV_BLOCK_M", "16")) + local_block_n_dkv = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DKV_BLOCK_N", str(local_block_n))) + # P57 R2: ``DQ_WARPS=4, DQ_STAGES=2`` wins ~140 us over the + # R1 ``W=8, S=1`` default. The dq kernel iterates ``n`` over + # a small SWA=128 window, so 4 warps fit the working set in + # registers and 2 stages prefetch the next K/V tile while the + # current MFMA chain executes. + local_dq_warps = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DQ_WARPS", "4")) + local_dq_stages = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DQ_STAGES", "2")) + local_dkv_warps = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DKV_WARPS", "2")) + local_dkv_stages = int(os.getenv("PRIMUS_V4_CSA_BWD_LOCAL_DKV_STAGES", "1")) + swa_local = int(swa_window) if swa_window > 0 else 0 + use_local_split = os.getenv("PRIMUS_V4_ATTN_BWD_USE_SPLIT", "0") == "1" + if use_local_split: + dq_grid = (triton.cdiv(Sq, local_block_m_dq), B * HQ) + _v4_attention_bwd_dq_kernel[dq_grid]( + q, + k_local, + v_local, + dout, + lse, + d_buf, + dq_fp32, + dsink_fp32, + sink_arg, + q, # ADD_MASK sentinel + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k_local.stride(0), + k_local.stride(1), + k_local.stride(2), + k_local.stride(3), + v_local.stride(0), + v_local.stride(1), + v_local.stride(2), + v_local.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_fp32.stride(0), + dq_fp32.stride(1), + dq_fp32.stride(2), + dq_fp32.stride(3), + 0, + 0, + Sq, + Sq, + float(scale), + HEAD_Q=HQ, + HEAD_K=HQ, + SWA_WINDOW=swa_local, + HAS_SINK=has_sink, + HAS_ADD_MASK=False, + HCA_LOCAL_SEQLEN=0, + USE_CAUSAL=True, + BLOCK_M=local_block_m_dq, + BLOCK_N=local_block_n_dq, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=local_dq_warps, + num_stages=local_dq_stages, + ) + dkv_grid = (triton.cdiv(Sq, local_block_n_dkv), B * HQ) + _v4_attention_bwd_dkv_kernel[dkv_grid]( + q, + k_local, + v_local, + dout, + lse, + d_buf, + dk_local_fp32, + dv_local_fp32, + q, # ADD_MASK sentinel + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k_local.stride(0), + k_local.stride(1), + k_local.stride(2), + k_local.stride(3), + v_local.stride(0), + v_local.stride(1), + v_local.stride(2), + v_local.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dk_local_fp32.stride(0), + dk_local_fp32.stride(1), + dk_local_fp32.stride(2), + dk_local_fp32.stride(3), + dv_local_fp32.stride(0), + dv_local_fp32.stride(1), + dv_local_fp32.stride(2), + dv_local_fp32.stride(3), + 0, + 0, + Sq, + Sq, + float(scale), + HEAD_Q=HQ, + HEAD_K=HQ, + SWA_WINDOW=swa_local, + HAS_ADD_MASK=False, + HCA_LOCAL_SEQLEN=0, + USE_CAUSAL=True, + BLOCK_M=local_block_m_dkv, + BLOCK_N=local_block_n_dkv, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=local_dkv_warps, + num_stages=local_dkv_stages, + ) + else: + local_grid = (triton.cdiv(Sq, local_block_m), B * HQ) + _v4_attention_bwd_kernel[local_grid]( + q, + k_local, + v_local, + dout, + lse, + d_buf, + dq_fp32, + dk_local_fp32, + dv_local_fp32, + dsink_fp32, + sink_arg, + q, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k_local.stride(0), + k_local.stride(1), + k_local.stride(2), + k_local.stride(3), + v_local.stride(0), + v_local.stride(1), + v_local.stride(2), + v_local.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_fp32.stride(0), + dq_fp32.stride(1), + dq_fp32.stride(2), + dq_fp32.stride(3), + dk_local_fp32.stride(0), + dk_local_fp32.stride(1), + dk_local_fp32.stride(2), + dk_local_fp32.stride(3), + dv_local_fp32.stride(0), + dv_local_fp32.stride(1), + dv_local_fp32.stride(2), + dv_local_fp32.stride(3), + 0, + 0, + Sq, + Sq, + float(scale), + HEAD_Q=HQ, + HEAD_K=HQ, + SWA_WINDOW=swa_local, + HAS_SINK=has_sink, + HAS_ADD_MASK=False, + HCA_LOCAL_SEQLEN=0, + USE_CAUSAL=True, + BLOCK_M=local_block_m, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=8, + num_stages=1, + ) + else: + grid = (Sq, B * HQ) + _v4_csa_attention_pool_bwd_kernel[grid]( + q, + k_local, + v_local, + pool, + topk_idxs, + dout, + lse, + d_buf, + dq_fp32, + dk_local_fp32, + dv_local_fp32, + dpool_fp32, + dsink_fp32, + sink_arg, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k_local.stride(0), + k_local.stride(1), + k_local.stride(2), + k_local.stride(3), + v_local.stride(0), + v_local.stride(1), + v_local.stride(2), + v_local.stride(3), + pool.stride(0), + pool.stride(1), + pool.stride(2), + topk_idxs.stride(0), + topk_idxs.stride(1), + topk_idxs.stride(2), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_fp32.stride(0), + dq_fp32.stride(1), + dq_fp32.stride(2), + dq_fp32.stride(3), + dk_local_fp32.stride(0), + dk_local_fp32.stride(1), + dk_local_fp32.stride(2), + dk_local_fp32.stride(3), + dv_local_fp32.stride(0), + dv_local_fp32.stride(1), + dv_local_fp32.stride(2), + dv_local_fp32.stride(3), + dpool_fp32.stride(0), + dpool_fp32.stride(1), + dpool_fp32.stride(2), + Sq, + P, + K_topk, + float(scale), + HEAD_Q=HQ, + SWA_WINDOW=int(swa_window) if swa_window > 0 else 0, + HAS_SINK=has_sink, + BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, + BLOCK_DMODEL=BLOCK_DMODEL, + STORE_DPOOL=(os.getenv("PRIMUS_V4_CSA_BWD_SKIP_DPOOL", "0") != "1"), + LOCAL_ONLY=False, + num_warps=4, + num_stages=1, + ) + + # Plan-5 P32: replace the gather-based sparse pool BWD with a + # dense-pool BWD that re-uses the split dQ / dK_pool+dV_pool kernels + # against a ``[Sq, P]`` additive mask synthesised from + # ``topk_idxs``. The dense path eliminates the per-tile + # ``tl.atomic_add(DPOOL, ...)`` collisions (``H_blocks × ~2k m's`` + # programs colliding on the shared ``[P, D]`` pool slots) at the + # cost of iterating the full pool dimension per m-block. With + # ``P=1024`` and ``K_topk=512`` (visibility ≈ 50%) the extra + # compute is small but the wall time win is large because atomics + # are the dominant cost at MI355 ``H=64``. + # + # Supported when ``B == 1`` (the proxy shape). For larger batches + # the kernel mask is per-(b, m), which the legacy gather kernel + # already handles natively — we fall back to it. + # Plan-5 P32: an alternative dense-pool path re-uses the split + # dQ + dK/dV kernels with a ``[Sq, P]`` ``log(count)`` additive + # mask, swapping ~1 B fp32 atomic adds for ``H * P/BLOCK_N`` + # atomic_adds plus ~2× the compute (full ``P`` vs ``K_topk``). + # At the proxy shape (``H=64, P=1024, K_topk=512``) the extra + # compute outweighs the atomic savings (56 ms vs the gather + # path's 25 ms), so the dense path is kept in tree but **off** + # by default. Toggle on for shapes with ``K_topk / P`` near 1 + # via ``PRIMUS_V4_CSA_BWD_DENSE_POOL=1``. + use_dense_pool_sparse = ( + use_split_sparse and B == 1 and os.getenv("PRIMUS_V4_CSA_BWD_DENSE_POOL", "0") == "1" + ) + # Enter the sparse stream just before launching the sparse pool + # kernels (default-stream kernels above have already been issued + # async on the default stream — the GPU can now overlap their + # execution with the sparse stream's work). + sparse_stream_entered = False + if sparse_stream_ctx is not None and use_split_sparse: + sparse_stream_ctx.__enter__() + sparse_stream_entered = True + + if use_dense_pool_sparse: + # Build ``additive_mask[Sq, P]`` from ``topk_idxs[0]``. A finite + # ``NEG_INF`` is required so the bf16 ``Q @ pool.T`` matmul does + # not produce ``NaN`` after ``+ -inf``. Duplicate top-K slots + # (k1, k2 both pointing to pool position p) make the gather BWD + # accumulate ``2 *`` the contribution of a single key. We + # collapse the gather visibility into a *count*-weighted mask: + # + # mask[m, p] = log(count[m, p]) if count[m, p] > 0 + # NEG_INF otherwise + # + # With this mask the dense kernel's ``P[m, p] = exp(qk + log + # count - LSE) = count * exp(qk - LSE)`` so the dense-pool + # ``ds`` matches the gather BWD's ``sum_k ds[m, k]`` term-by- + # term and the dense ``dq`` / ``dpool`` agree with the + # reference. (Invalid ``-1`` slots contribute zero count.) + NEG_INF = -1.0e30 + topk_b = topk_idxs[0] # [Sq, K_topk] + valid_per_topk = (topk_b >= 0) & (topk_b < P) + safe_topk = torch.where(valid_per_topk, topk_b, torch.zeros_like(topk_b)).to(torch.int64) + count = torch.zeros((Sq, P), device=q.device, dtype=torch.float32) + count.scatter_add_( + 1, + safe_topk, + valid_per_topk.to(torch.float32), + ) + sparse_mask = torch.where( + count > 0, + count.clamp_min(1.0).log(), + torch.full((), NEG_INF, device=q.device, dtype=torch.float32), + ) + + # View pool as a per-head K/V with stride_kh=stride_vh=0 so the + # split dq/dkv kernels see ``HEAD_K == HEAD_Q`` and parallelise + # the dkv kernel over the full ``(P/BLOCK_N, B*HEAD_Q)`` grid + # (2 048 programs at proxy) instead of ``HEAD_K=1`` MQA (32 + # programs) — that path was bandwidth-starved at ``D=512``. + # Use ``stride_dk = stride_dv = 0`` so each program's ``dk / + # dv`` write fans into the same shared ``dpool`` slice; the + # accumulation is via ``tl.atomic_add`` which costs ~4 096 + # call ops total (single atomic per (n_block, head) plus dK + + # dV) — orders of magnitude less than the gather kernel's + # ~1 B fp32 atomic adds. + pool_4d = pool.unsqueeze(1).expand(B, HQ, P, D) # stride_kh=0 + + sparse_block_m = 32 + sparse_block_n = 32 + dq_sparse_grid = (triton.cdiv(Sq, sparse_block_m), B * HQ) + _v4_attention_bwd_dq_kernel[dq_sparse_grid]( + q, + pool_4d, + pool_4d, + dout, + lse, + d_buf, + dq_sparse_fp32, + dsink_fp32, # sentinel only; HAS_SINK=False + q, + sparse_mask, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + pool_4d.stride(0), + pool_4d.stride(1), + pool_4d.stride(2), + pool_4d.stride(3), + pool_4d.stride(0), + pool_4d.stride(1), + pool_4d.stride(2), + pool_4d.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_sparse_fp32.stride(0), + dq_sparse_fp32.stride(1), + dq_sparse_fp32.stride(2), + dq_sparse_fp32.stride(3), + sparse_mask.stride(0), + sparse_mask.stride(1), + Sq, + P, + float(scale), + HEAD_Q=HQ, + HEAD_K=HQ, + SWA_WINDOW=0, + HAS_SINK=False, + HAS_ADD_MASK=True, + HCA_LOCAL_SEQLEN=0, + USE_CAUSAL=False, + BLOCK_M=sparse_block_m, + BLOCK_N=sparse_block_n, + BLOCK_DMODEL=BLOCK_DMODEL, + ACCUMULATE=True, + num_warps=8, + num_stages=1, + ) + # Single shared ``dpool`` slice that all ``(head, n_block)`` + # programs atomic_add into. Reshape with an extra head axis of + # stride 0 so the dkv kernel can index it as ``[B, HQ, P, D]`` + # without per-head storage; the actual underlying tensor is + # still ``[B, P, D]``. + dk_pool_shared = dpool_fp32.unsqueeze(1).expand(B, HQ, P, D) + dv_pool_shared = dpool_fp32.unsqueeze(1).expand(B, HQ, P, D) + dkv_sparse_grid = (triton.cdiv(P, sparse_block_n), B * HQ) + _v4_attention_bwd_dkv_kernel[dkv_sparse_grid]( + q, + pool_4d, + pool_4d, + dout, + lse, + d_buf, + dk_pool_shared, + dv_pool_shared, + sparse_mask, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + pool_4d.stride(0), + pool_4d.stride(1), + pool_4d.stride(2), + pool_4d.stride(3), + pool_4d.stride(0), + pool_4d.stride(1), + pool_4d.stride(2), + pool_4d.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dk_pool_shared.stride(0), + dk_pool_shared.stride(1), + dk_pool_shared.stride(2), + dk_pool_shared.stride(3), + dv_pool_shared.stride(0), + dv_pool_shared.stride(1), + dv_pool_shared.stride(2), + dv_pool_shared.stride(3), + sparse_mask.stride(0), + sparse_mask.stride(1), + Sq, + P, + float(scale), + HEAD_Q=HQ, + HEAD_K=HQ, + SWA_WINDOW=0, + HAS_ADD_MASK=True, + HCA_LOCAL_SEQLEN=0, + USE_CAUSAL=False, + BLOCK_M=sparse_block_m, + BLOCK_N=sparse_block_n, + BLOCK_DMODEL=BLOCK_DMODEL, + ATOMIC_REDUCE=True, + num_warps=8, + num_stages=1, + ) + # ``dpool_fp32`` already holds the summed dk+dv contributions + # because both kernel views aliased the same ``[B, P, D]`` + # storage via ``stride_kh=stride_vh=0`` and the kernel does + # both ``atomic_add(dk_ptrs, ...)`` and ``atomic_add(dv_ptrs, + # ...)``. + elif use_split_sparse: + # P57: BLOCK_H=32, num_warps=4, num_stages=2 wins the proxy + # sweep over the P32 default (BLOCK_H=64, num_warps=8, stages=1). + # Doubling the head-axis grid (HQ=64 -> 2 h-blocks per m) lifts + # MI355 occupancy, while fewer warps + 2 stages reduces register + # pressure per warp (the dq[BLOCK_H, BLOCK_DMODEL]=fp32 + # accumulator dominates VGPR live range). + BLOCK_H = int(os.getenv("PRIMUS_V4_CSA_BWD_SPARSE_BLOCK_H", "32")) + # Plan-5 P32: ``PRIMUS_V4_CSA_BWD_SEGREDUCE=1`` is now the + # default — wins both the standalone CSA BWD microbench + # (16.31 ms vs 24.83 ms gather/atomic) and the EP8 proxy + # (578 ms vs 665 ms / iter) after the P32 dual-RoPE bf16-cast + # fix. Pre-fix, ``apply_interleaved_partial_rope`` was + # promoting Q/K to fp32 (cos/sin came from + # ``position_ids.float()`` and bf16*fp32=fp32), which 2x'd + # Q/K HBM traffic, inflated *every* attention kernel time + # 1.8-7x in the proxy trace, and made the gather + atomic + # path look faster end-to-end purely because the segreduce + # 4 GiB partial buffer competed against artificially-bloated + # attention traffic. Setting ``PRIMUS_V4_CSA_BWD_SEGREDUCE=0`` + # falls back to gather + atomic for kernel-tuning. The + # segmented-reduction path writes per-visit dpool + # contributions to a compact ``[B, M, K_topk, D]`` partial + # buffer (no atomics) and then folds them into + # ``dpool[B, P, D]`` via a sorted inverse index. ``perm32``, + # ``bin_ptr`` and ``dpool_partial`` were built above on the + # default stream BEFORE the local kernels were launched, so + # they're ready by the time the sparse stream gets here. + # P57: experimental split-kernel path (opt-in). + # Sub-kernels: + # + # _v4_csa_attention_pool_sparse_bwd_dq_only_kernel + # _v4_csa_attention_pool_sparse_bwd_dpool_only_kernel + # + # vs the joint kernel + # ``_v4_csa_attention_pool_sparse_bwd_partial_kernel``. The + # split path frees the live-VGPR footprint of each sub-kernel + # but pays a 2× read on Q / dout / lse / dvec / pool / + # topk_idxs. On the EP8 proxy shape the joint kernel still + # wins (~3.7 ms vs ~4.5 ms split), so split is opt-in via + # ``PRIMUS_V4_CSA_BWD_SPLIT_DQ_DPOOL=1`` and the joint kernel + # remains the default. Keeping the split sub-kernels in tree + # for future shapes where the read amplification is cheaper + # (smaller H or D, or compute-bound regimes). + use_split_dq_dpool = os.getenv("PRIMUS_V4_CSA_BWD_SPLIT_DQ_DPOOL", "0") != "0" + if use_segreduce and use_split_dq_dpool: + sparse_grid = (Sq, triton.cdiv(HQ, BLOCK_H), B) + dq_warps = int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_DQ_WARPS", "4")) + dq_stages = int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_DQ_STAGES", "2")) + dpool_warps = int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_DPOOL_WARPS", "4")) + dpool_stages = int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_DPOOL_STAGES", "2")) + _v4_csa_attention_pool_sparse_bwd_dq_only_kernel[sparse_grid]( + q, + pool, + topk_idxs, + dout, + lse, + d_buf, + dq_sparse_fp32, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + pool.stride(0), + pool.stride(1), + pool.stride(2), + topk_idxs.stride(0), + topk_idxs.stride(1), + topk_idxs.stride(2), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_sparse_fp32.stride(0), + dq_sparse_fp32.stride(1), + dq_sparse_fp32.stride(2), + dq_sparse_fp32.stride(3), + Sq, + P, + K_topk, + float(scale), + HEAD_Q=HQ, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K_PARTIAL, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=dq_warps, + num_stages=dq_stages, + ) + _v4_csa_attention_pool_sparse_bwd_dpool_only_kernel[sparse_grid]( + q, + pool, + topk_idxs, + dout, + lse, + d_buf, + dpool_partial, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + pool.stride(0), + pool.stride(1), + pool.stride(2), + topk_idxs.stride(0), + topk_idxs.stride(1), + topk_idxs.stride(2), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dpool_partial.stride(0), + dpool_partial.stride(1), + dpool_partial.stride(2), + dpool_partial.stride(3), + Sq, + P, + K_topk, + float(scale), + HEAD_Q=HQ, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K_PARTIAL, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=dpool_warps, + num_stages=dpool_stages, + ) + # P57: segreduce reduction (shared between split and joint + # partial paths). Reduces ``dpool_partial[B, MK, D]`` into + # ``dpool[B, P, D]`` using the sorted inverse index. + dpool_partial_flat = dpool_partial.view(B, Sq * K_topk, D) + block_d_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_D", "512")) + block_i_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_I", "64")) + seg_grid = (P, triton.cdiv(D, block_d_seg), B) + _v4_csa_attention_pool_segreduce_kernel[seg_grid]( + dpool_partial_flat, + perm32, + bin_ptr, + dpool_fp32, + dpool_partial_flat.stride(0), + dpool_partial_flat.stride(1), + dpool_partial_flat.stride(2), + perm32.stride(0), + perm32.stride(1), + bin_ptr.stride(0), + bin_ptr.stride(1), + dpool_fp32.stride(0), + dpool_fp32.stride(1), + dpool_fp32.stride(2), + P, + D, + BLOCK_D=block_d_seg, + BLOCK_I=block_i_seg, + num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_WARPS", "4")), + num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_STAGES", "2")), + ) + elif use_segreduce and use_sorted_partial: + sparse_grid = (Sq, triton.cdiv(HQ, BLOCK_H), B) + dpool_partial_flat = dpool_partial.view(B, Sq * K_topk, D) + _v4_csa_attention_pool_sparse_bwd_partial_sorted_kernel[sparse_grid]( + q, + pool, + topk_idxs, + inv_perm32, + dout, + lse, + d_buf, + dq_sparse_fp32, + dpool_partial_flat, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + pool.stride(0), + pool.stride(1), + pool.stride(2), + topk_idxs.stride(0), + topk_idxs.stride(1), + topk_idxs.stride(2), + inv_perm32.stride(0), + inv_perm32.stride(1), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_sparse_fp32.stride(0), + dq_sparse_fp32.stride(1), + dq_sparse_fp32.stride(2), + dq_sparse_fp32.stride(3), + dpool_partial_flat.stride(0), + dpool_partial_flat.stride(1), + dpool_partial_flat.stride(2), + Sq, + P, + K_topk, + float(scale), + HEAD_Q=HQ, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K_PARTIAL, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_WARPS", "4")), + num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_STAGES", "2")), + ) + # Sequential segreduce — no perm lookup. + block_d_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_D", "512")) + block_i_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_I", "64")) + seg_grid = (P, triton.cdiv(D, block_d_seg), B) + _v4_csa_attention_pool_segreduce_sequential_kernel[seg_grid]( + dpool_partial_flat, + bin_ptr, + dpool_fp32, + dpool_partial_flat.stride(0), + dpool_partial_flat.stride(1), + dpool_partial_flat.stride(2), + bin_ptr.stride(0), + bin_ptr.stride(1), + dpool_fp32.stride(0), + dpool_fp32.stride(1), + dpool_fp32.stride(2), + P, + D, + BLOCK_D=block_d_seg, + BLOCK_I=block_i_seg, + num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_WARPS", "4")), + num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_STAGES", "2")), + ) + elif use_segreduce: + sparse_grid = (Sq, triton.cdiv(HQ, BLOCK_H), B) + _v4_csa_attention_pool_sparse_bwd_partial_kernel[sparse_grid]( + q, + pool, + topk_idxs, + dout, + lse, + d_buf, + dq_sparse_fp32, + dpool_partial, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + pool.stride(0), + pool.stride(1), + pool.stride(2), + topk_idxs.stride(0), + topk_idxs.stride(1), + topk_idxs.stride(2), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_sparse_fp32.stride(0), + dq_sparse_fp32.stride(1), + dq_sparse_fp32.stride(2), + dq_sparse_fp32.stride(3), + dpool_partial.stride(0), + dpool_partial.stride(1), + dpool_partial.stride(2), + dpool_partial.stride(3), + Sq, + P, + K_topk, + float(scale), + HEAD_Q=HQ, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K_PARTIAL, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_WARPS", "4")), + num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_PARTIAL_STAGES", "3")), + ) + # Same segreduce reduction as the split path above. + dpool_partial_flat = dpool_partial.view(B, Sq * K_topk, D) + block_d_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_D", "512")) + block_i_seg = int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_BLOCK_I", "64")) + seg_grid = (P, triton.cdiv(D, block_d_seg), B) + _v4_csa_attention_pool_segreduce_kernel[seg_grid]( + dpool_partial_flat, + perm32, + bin_ptr, + dpool_fp32, + dpool_partial_flat.stride(0), + dpool_partial_flat.stride(1), + dpool_partial_flat.stride(2), + perm32.stride(0), + perm32.stride(1), + bin_ptr.stride(0), + bin_ptr.stride(1), + dpool_fp32.stride(0), + dpool_fp32.stride(1), + dpool_fp32.stride(2), + P, + D, + BLOCK_D=block_d_seg, + BLOCK_I=block_i_seg, + num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_WARPS", "4")), + num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_SEGREDUCE_STAGES", "2")), + ) + else: + # Plan-5 P32: legacy gather-sparse path. Keeps the atomic_add + # to a shared ``dpool[B, P, D]`` buffer; this is the shipped + # default. The segreduce path is opt-in via + # ``PRIMUS_V4_CSA_BWD_SEGREDUCE=1`` for kernel-level perf + # experiments on the CSA microbench. + n_part_env = int(os.getenv("PRIMUS_V4_CSA_BWD_DPOOL_PARTITIONS", "1")) + n_part = max(1, min(n_part_env, Sq)) + while Sq % n_part != 0 and n_part > 1: + n_part -= 1 + partition_size = Sq // n_part + if n_part > 1: + dpool_partial = torch.zeros((n_part, B, P, D), device=q.device, dtype=torch.float32) + else: + dpool_partial = dpool_fp32.unsqueeze(0) + sparse_grid = (Sq, triton.cdiv(HQ, BLOCK_H), B) + _v4_csa_attention_pool_sparse_bwd_kernel[sparse_grid]( + q, + pool, + topk_idxs, + dout, + lse, + d_buf, + dq_sparse_fp32, + dpool_partial, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + pool.stride(0), + pool.stride(1), + pool.stride(2), + topk_idxs.stride(0), + topk_idxs.stride(1), + topk_idxs.stride(2), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + d_buf.stride(0), + d_buf.stride(1), + d_buf.stride(2), + dq_sparse_fp32.stride(0), + dq_sparse_fp32.stride(1), + dq_sparse_fp32.stride(2), + dq_sparse_fp32.stride(3), + dpool_partial.stride(0), + dpool_partial.stride(1), + dpool_partial.stride(2), + dpool_partial.stride(3), + partition_size, + Sq, + P, + K_topk, + float(scale), + HEAD_Q=HQ, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K, + BLOCK_DMODEL=BLOCK_DMODEL, + STORE_DPOOL=(os.getenv("PRIMUS_V4_CSA_BWD_SKIP_DPOOL_ATOMIC", "0") != "1"), + num_warps=int(os.getenv("PRIMUS_V4_CSA_BWD_SPARSE_WARPS", "4")), + num_stages=int(os.getenv("PRIMUS_V4_CSA_BWD_SPARSE_STAGES", "1")), + ) + if n_part > 1: + dpool_fp32 = dpool_partial.sum(dim=0) + + # Exit the sparse stream context (if entered) and join it with the + # default stream before the final dtype casts so the dq accumulator + # below sees both the local and sparse contributions. + if sparse_stream_entered: + sparse_done = torch.cuda.current_stream(q.device).record_event() + sparse_stream_ctx.__exit__(None, None, None) + torch.cuda.current_stream(q.device).wait_event(sparse_done) + + # The split-sparse paths write the sparse contribution to a + # dedicated ``dq_sparse_fp32`` buffer; combine it with the local + # SWA contribution before the dtype cast. In-place add saves the + # extra ~128 MB allocation that ``+`` would do. + if use_split_sparse: + dq_fp32.add_(dq_sparse_fp32) + + # P57: if the local buffers are already in input dtype, ``.to`` + # is a no-op view; only fp32 buffers actually do a cast. + dq_out = dq_fp32 if dq_fp32.dtype == q.dtype else dq_fp32.to(q.dtype) + dk_local_out = dk_local_fp32 if dk_local_fp32.dtype == k_local.dtype else dk_local_fp32.to(k_local.dtype) + dv_local_out = dv_local_fp32 if dv_local_fp32.dtype == v_local.dtype else dv_local_fp32.to(v_local.dtype) + dpool_out = dpool_fp32 if dpool_fp32.dtype == pool.dtype else dpool_fp32.to(pool.dtype) + dsink_out = dsink_fp32.to(sink.dtype) if has_sink else None + return dq_out, dk_local_out, dv_local_out, dpool_out, dsink_out + + +__all__ = [ + "_v4_csa_attention_bwd_kernel", + "_v4_csa_attention_pool_bwd_kernel", + "_v4_csa_attention_pool_sparse_bwd_kernel", + "_launch_v4_csa_attention_bwd", + "_launch_v4_csa_attention_pool_bwd", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_fwd.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_fwd.py new file mode 100644 index 000000000..5ee75bbbe --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v1/v4_csa_attention_fwd.py @@ -0,0 +1,1040 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 CSA attention forward Triton kernel (plan-4 P26, ``compress_ratio == 4``). + +CSA fuses three branches into a single online softmax: + +* **Local SWA**: ``q @ k_local^T`` with sliding-window-causal masking. +* **Sparse top-K**: ``q . gathered[m, :, :]`` where the wrapper has + pre-gathered ``[B, Sq, K, D]`` rows from the compressed pool (the + per-query top-K gather lives outside the kernel — see plan-4 + ``02-phase-details.md`` Phase 26 design notes). +* **Per-head learned sink**: a virtual key column with notional value + zero, joined as the last softmax candidate so its probability mass + is shared across local + sparse branches. + +The kernel produces one ``[BLOCK_DMODEL]`` output row per program; the +grid is ``(seqlen_q, batch * head_q)`` so each program owns exactly one +``(b, qhid, m)`` query row. The per-row design keeps the sparse-branch +SMEM footprint inside the MI355 budget at ``head_dim=512``: the +gathered tile is only ``[BLOCK_K, head_dim] * 2 bytes ≈ 32 KiB`` per +program, while a multi-row tile would balloon to +``[BLOCK_M, BLOCK_K, head_dim] * 2 bytes ≈ 1 MiB``. + +dtype contract (matches :func:`eager_v4_csa_attention`): + +* Q / K / V / gathered are loaded in input dtype (bf16 in production); + the per-row dot products (``sum(k * q[None, :], axis=-1)``) reduce in + fp32 because we ``.to(tl.float32)`` before the multiply. +* The online-softmax accumulator (``m_running``, ``l_running``, + ``acc``) lives in fp32 — the *only* fp32 step inside the kernel. +* Output is written back in input dtype; saved ``LSE`` is fp32 (BWD + re-materialises ``P`` from it). + +Edge cases handled: + +* ``K_topk == 0`` — wrapper short-circuits to the dense + :func:`v4_attention_v1` kernel before reaching this file. +* ``topk_idx == -1`` — wrapper sets the corresponding ``sparse_mask`` + entry to ``-inf``; the kernel just adds the bias and the masked + position contributes ~0 to the softmax denominator. +* All-masked tile rows — the running max and per-tile max are both + the finite ``NEG_INF`` sentinel (``-1e30``), so + ``exp(NEG_INF - NEG_INF) = exp(0) = 1`` algebraically but the + contribution to ``acc`` and ``l_running`` is gated by the + per-element ``exp(qk - m_new)`` which stays at exactly zero for + every ``-inf``-masked entry. This avoids the ``exp(-inf - -inf) = + exp(NaN)`` failure mode that ``-float("inf")`` would have. +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch +import triton +import triton.language as tl + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v1.v4_attention_fwd import ( + _launch_v4_attention_fwd, +) + +# --------------------------------------------------------------------------- +# Triton kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_csa_attention_pool_fwd_kernel( + Q, + K_LOCAL, + V_LOCAL, + POOL, + TOPK_IDXS, + SINK, + OUT, + LSE, + # Q strides: [B, H, Sq, D] + stride_qb, + stride_qh, + stride_qm, + stride_qd, + # K_local strides: [B, H, Sq, D] + stride_klb, + stride_klh, + stride_kln, + stride_kld, + # V_local strides: [B, H, Sq, D] + stride_vlb, + stride_vlh, + stride_vln, + stride_vld, + # Pool strides: [B, P, D] (shared across heads) + stride_pb, + stride_pp, + stride_pd, + # topk_idxs strides: [B, Sq, K_topk] + stride_tib, + stride_tim, + stride_tik, + # OUT strides: [B, H, Sq, D] + stride_ob, + stride_oh, + stride_om, + stride_od, + # LSE strides: [B, H, Sq] + stride_lb, + stride_lh, + stride_lm, + seqlen_q, + pool_size, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + SWA_WINDOW: tl.constexpr, + HAS_SINK: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """CSA FWD with in-kernel topk gather from the compressed pool.""" + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + bid = pid_bh // HEAD_Q + qhid = pid_bh % HEAD_Q + + offs_d = tl.arange(0, BLOCK_DMODEL) + q_active = pid_m < seqlen_q + NEG_INF: tl.constexpr = -1.0e30 + + q_row_offset = bid * stride_qb + qhid * stride_qh + pid_m * stride_qm + q = tl.load(Q + q_row_offset + offs_d * stride_qd, mask=q_active, other=0.0) + + acc = tl.zeros([BLOCK_DMODEL], dtype=tl.float32) + m_i = tl.full((), value=NEG_INF, dtype=tl.float32) + l_i = tl.zeros((), dtype=tl.float32) + + n_loop_end = pid_m + 1 + if n_loop_end > seqlen_q: + n_loop_end = seqlen_q + if SWA_WINDOW > 0: + n_lo_raw = pid_m - SWA_WINDOW + 1 + if n_lo_raw < 0: + n_lo_raw = 0 + n_loop_start = (n_lo_raw // BLOCK_N) * BLOCK_N + else: + n_loop_start = 0 + + for n_start in range(n_loop_start, n_loop_end, BLOCK_N): + offs_n = n_start + tl.arange(0, BLOCK_N) + kl_ptrs = ( + K_LOCAL + + bid * stride_klb + + qhid * stride_klh + + offs_n[:, None] * stride_kln + + offs_d[None, :] * stride_kld + ) + kl_load_mask = offs_n[:, None] < seqlen_q + kl = tl.load(kl_ptrs, mask=kl_load_mask, other=0.0) + qk = tl.sum(kl.to(tl.float32) * q[None, :].to(tl.float32), axis=1) * sm_scale + + if SWA_WINDOW > 0: + in_window = (offs_n >= pid_m - SWA_WINDOW + 1) & (offs_n <= pid_m) + else: + in_window = offs_n <= pid_m + qk = tl.where(in_window & (offs_n < seqlen_q), qk, NEG_INF) + + m_tile = tl.max(qk, axis=0) + m_new = tl.maximum(m_i, m_tile) + alpha = tl.exp(m_i - m_new) + p = tl.exp(qk - m_new) + l_i = l_i * alpha + tl.sum(p, axis=0) + + vl_ptrs = ( + V_LOCAL + + bid * stride_vlb + + qhid * stride_vlh + + offs_n[:, None] * stride_vln + + offs_d[None, :] * stride_vld + ) + vl = tl.load(vl_ptrs, mask=kl_load_mask, other=0.0) + acc = acc * alpha + tl.sum(p[:, None] * vl.to(tl.float32), axis=0) + m_i = m_new + + for k_start in range(0, K_topk, BLOCK_K): + offs_k = k_start + tl.arange(0, BLOCK_K) + topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + offs_k * stride_tik + topk = tl.load(topk_ptrs, mask=offs_k < K_topk, other=-1) + valid = (offs_k < K_topk) & (topk >= 0) & (topk < pool_size) + safe_topk = tl.where(valid, topk, 0) + + pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd + pool = tl.load(pool_ptrs, mask=valid[:, None], other=0.0) + + qk_sparse = tl.sum(pool.to(tl.float32) * q[None, :].to(tl.float32), axis=1) * sm_scale + qk_sparse = tl.where(valid, qk_sparse, NEG_INF) + + m_tile = tl.max(qk_sparse, axis=0) + m_new = tl.maximum(m_i, m_tile) + alpha = tl.exp(m_i - m_new) + p = tl.exp(qk_sparse - m_new) + l_i = l_i * alpha + tl.sum(p, axis=0) + + acc = acc * alpha + tl.sum(p[:, None] * pool.to(tl.float32), axis=0) + m_i = m_new + + if HAS_SINK: + sink_h = tl.load(SINK + qhid).to(tl.float32) + m_new = tl.maximum(m_i, sink_h) + alpha = tl.exp(m_i - m_new) + beta = tl.exp(sink_h - m_new) + l_i = l_i * alpha + beta + acc = acc * alpha + m_i = m_new + + out = acc / l_i + lse = m_i + tl.log(l_i) + + out_offset = bid * stride_ob + qhid * stride_oh + pid_m * stride_om + tl.store(OUT + out_offset + offs_d * stride_od, out.to(OUT.dtype.element_ty), mask=q_active) + + lse_ptr = LSE + bid * stride_lb + qhid * stride_lh + pid_m * stride_lm + tl.store(lse_ptr, lse, mask=q_active) + + +# --------------------------------------------------------------------------- +# Python launcher +# --------------------------------------------------------------------------- + + +def _launch_v4_csa_attention_pool_fwd( + q: torch.Tensor, # [B, H, Sq, D] + k_local: torch.Tensor, # [B, H, Sq, D] + v_local: torch.Tensor, # [B, H, Sq, D] + pool: torch.Tensor, # [B, P, D] + topk_idxs: torch.Tensor, # [B, Sq, K_topk], -1 means masked + *, + sink: Optional[torch.Tensor], + swa_window: int, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Launch CSA forward with in-kernel topk gather from ``pool``.""" + if q.dim() != 4 or k_local.dim() != 4 or v_local.dim() != 4: + raise ValueError( + "v4_csa_attention_v0 pool forward expects q / k_local / v_local of rank 4 " + f"(got {q.dim()} / {k_local.dim()} / {v_local.dim()})" + ) + if pool.dim() != 3: + raise ValueError( + f"v4_csa_attention_v0 pool forward expects pool of rank 3 [B, P, D]; " + f"got rank {pool.dim()}, shape {tuple(pool.shape)}" + ) + if topk_idxs.dim() != 3: + raise ValueError( + f"v4_csa_attention_v0 pool forward expects topk_idxs of rank 3 [B, Sq, K]; " + f"got rank {topk_idxs.dim()}, shape {tuple(topk_idxs.shape)}" + ) + + B, HQ, Sq, D = q.shape + if k_local.shape != q.shape or v_local.shape != q.shape: + raise ValueError( + "v4_csa_attention_v0 pool path requires k_local.shape == v_local.shape == q.shape " + f"(got q={tuple(q.shape)}, k_local={tuple(k_local.shape)}, " + f"v_local={tuple(v_local.shape)})." + ) + Bp, P, Dp = pool.shape + if Bp != B or Dp != D: + raise ValueError( + "v4_csa_attention_v0 pool shape mismatch: expected " + f"[B, P, D] = [{B}, *, {D}]; got {tuple(pool.shape)}." + ) + Bt, Sqt, K_topk = topk_idxs.shape + if Bt != B or Sqt != Sq: + raise ValueError( + "v4_csa_attention_v0 topk_idxs shape mismatch: expected " + f"[B, Sq, K] = [{B}, {Sq}, *]; got {tuple(topk_idxs.shape)}." + ) + if topk_idxs.dtype not in (torch.int32, torch.int64): + raise ValueError(f"v4_csa_attention_v0 topk_idxs must be int32/int64, got {topk_idxs.dtype}.") + if not q.is_cuda: + raise ValueError("v4_csa_attention_v0 requires CUDA / HIP tensors.") + if q.dtype != k_local.dtype or q.dtype != v_local.dtype or q.dtype != pool.dtype: + raise ValueError( + "v4_csa_attention_v0 pool path requires q.dtype == k_local.dtype == " + f"v_local.dtype == pool.dtype (got {q.dtype} / {k_local.dtype} / " + f"{v_local.dtype} / {pool.dtype})." + ) + + has_sink = sink is not None + + # Plan-5 P32: default to the split FWD (local SWA dense kernel + + # head-block sparse kernel + LSE merge). The monolithic per-row + # ``_v4_csa_attention_pool_fwd_kernel`` stays in tree as the + # ``PRIMUS_V4_CSA_FWD_FORCE_MONOLITHIC=1`` fallback so we can A/B + # the two designs from the proxy without rebuilding. + if os.getenv("PRIMUS_V4_CSA_FWD_FORCE_MONOLITHIC", "0") != "1": + return _launch_v4_csa_attention_pool_fwd_split( + q, + k_local, + v_local, + pool, + topk_idxs, + sink=sink, + swa_window=int(swa_window), + scale=float(scale), + ) + + out = torch.empty_like(q) + lse = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32) + + BLOCK_N = 32 + BLOCK_K = 32 + BLOCK_DMODEL = D + grid = (Sq, B * HQ) + sink_ptr = sink if has_sink else q + + _v4_csa_attention_pool_fwd_kernel[grid]( + q, + k_local, + v_local, + pool, + topk_idxs, + sink_ptr, + out, + lse, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k_local.stride(0), + k_local.stride(1), + k_local.stride(2), + k_local.stride(3), + v_local.stride(0), + v_local.stride(1), + v_local.stride(2), + v_local.stride(3), + pool.stride(0), + pool.stride(1), + pool.stride(2), + topk_idxs.stride(0), + topk_idxs.stride(1), + topk_idxs.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + Sq, + P, + K_topk, + float(scale), + HEAD_Q=HQ, + SWA_WINDOW=int(swa_window) if swa_window > 0 else 0, + HAS_SINK=has_sink, + BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=4, + num_stages=1, + ) + return out, lse + + +# --------------------------------------------------------------------------- +# Plan-5 P32 split CSA FWD — sparse head-block tile + LSE merge +# +# The per-row design in ``_v4_csa_attention_pool_fwd_kernel`` cannot reach +# tensor-core throughput because ``tl.sum(k * q, axis=1)`` is a per-row +# reduction (one program per ``(b, qhid, m)``). FlashAttention-style +# multi-row ``tl.dot`` tiles need ``BLOCK_M >= 16`` queries per program; +# the sparse branch's per-query ``topk_idxs`` gather blocks that — adjacent +# query rows have different sparse keys. +# +# P32 splits the FWD into three launches that ARE multi-row friendly: +# +# 1. ``_launch_v4_attention_fwd`` (already shipped) handles the local +# SWA branch with ``BLOCK_M=BLOCK_N=32`` ``tl.dot`` tiles. We call it +# with ``sink=None`` so the returned ``(out_local, lse_local)`` does +# NOT include the per-head sink. +# 2. ``_v4_csa_attention_pool_sparse_fwd_kernel`` (new) handles the +# sparse pool branch with a **head-block** tile. The per-query +# ``topk_idxs`` gather is shared across all ``H`` query heads — the +# pool tensor itself has no head dimension. One program owns one +# ``(b, m, h_block)`` and runs ``tl.dot(Q[BLOCK_H, D], +# tl.trans(pool[BLOCK_K, D]))`` per top-K tile, online-softmax- +# updating per-head ``m_i / l_i / acc`` along the way. Output +# ``(out_sparse, lse_sparse)`` does NOT include the sink. +# 3. ``_v4_csa_attention_lse_merge_kernel`` (new) combines the two +# ``(out, lse)`` pairs with the per-head sink under one final online +# softmax. Result: a joint ``out`` and joint ``lse`` mathematically +# identical to the monolithic kernel (modulo fp32 reduction order), +# plus the per-iteration BWD contract is unchanged because the joint +# ``lse`` is what the BWD already re-materialises ``P`` from. +# --------------------------------------------------------------------------- + + +@triton.jit +def _v4_csa_attention_pool_sparse_fwd_kernel( + Q, + POOL, + TOPK_IDXS, + OUT, + LSE, + # Q strides: [B, H, Sq, D] + stride_qb, + stride_qh, + stride_qm, + stride_qd, + # Pool strides: [B, P, D] + stride_pb, + stride_pp, + stride_pd, + # topk_idxs strides: [B, Sq, K_topk] + stride_tib, + stride_tim, + stride_tik, + # OUT strides: [B, H, Sq, D] + stride_ob, + stride_oh, + stride_om, + stride_od, + # LSE strides: [B, H, Sq] + stride_lb, + stride_lh, + stride_lm, + seqlen_q, + pool_size, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """Sparse-branch CSA FWD with head-block tile + ``tl.dot``. + + Grid: ``(seqlen_q, cdiv(HEAD_Q, BLOCK_H), B)``. Each program owns + one ``(b, h_block, m)`` and computes the sparse branch's normalized + output + LSE for ``BLOCK_H`` heads. The pool gather is shared + across heads because ``pool`` and ``topk_idxs`` have no H axis. + """ + pid_m = tl.program_id(0) + pid_h_block = tl.program_id(1) + bid = tl.program_id(2) + + offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H) + offs_k = tl.arange(0, BLOCK_K) + offs_d = tl.arange(0, BLOCK_DMODEL) + h_mask = offs_h < HEAD_Q + q_active = pid_m < seqlen_q + + NEG_INF: tl.constexpr = -1.0e30 + + # Q tile: [BLOCK_H, BLOCK_DMODEL] + q_ptrs = ( + Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd + ) + q = tl.load(q_ptrs, mask=h_mask[:, None] & q_active, other=0.0) + + # Online-softmax state per head. + acc = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32) + m_i = tl.full([BLOCK_H], value=NEG_INF, dtype=tl.float32) + l_i = tl.zeros([BLOCK_H], dtype=tl.float32) + + for k_start in range(0, K_topk, BLOCK_K): + sparse_k = k_start + offs_k + topk_ptrs = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + sparse_k * stride_tik + topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1) + valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size) + safe_topk = tl.where(valid_k, topk, 0) + + pool_ptrs = POOL + bid * stride_pb + safe_topk[:, None] * stride_pp + offs_d[None, :] * stride_pd + pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0) + + # qk = Q @ pool.T : [BLOCK_H, BLOCK_K] in fp32 accumulator + qk = tl.dot(q.to(pool.dtype), tl.trans(pool)) * sm_scale + qk = tl.where((h_mask[:, None] & valid_k[None, :] & q_active), qk, NEG_INF) + + # Online softmax update per head row. + m_tile = tl.max(qk, axis=1) + m_new = tl.maximum(m_i, m_tile) + alpha = tl.exp(m_i - m_new) + p = tl.exp(qk - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(pool.dtype), pool) + m_i = m_new + + # Detect "all invalid" rows (K_topk == 0 or every topk_idx == -1): + # l_i stays zero; we leave acc as zero and write a NEG_INF lse so the + # merge kernel treats the sparse branch as carrying zero softmax mass + # for that (b, h, m). + empty = l_i == 0.0 + safe_l = tl.where(empty, 1.0, l_i) + out = acc / safe_l[:, None] + lse = tl.where(empty, NEG_INF, m_i + tl.log(safe_l)) + + out_ptrs = ( + OUT + bid * stride_ob + offs_h[:, None] * stride_oh + pid_m * stride_om + offs_d[None, :] * stride_od + ) + tl.store(out_ptrs, out.to(OUT.dtype.element_ty), mask=h_mask[:, None] & q_active) + + lse_ptrs = LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm + tl.store(lse_ptrs, lse, mask=h_mask & q_active) + + +@triton.jit +def _v4_csa_attention_pool_sparse_merge_fwd_kernel( + Q, + POOL, + TOPK_IDXS, + OUT_LOCAL, # local-branch result (no sink) — input + LSE_LOCAL, # local-branch lse (no sink) — input + SINK, + OUT, # joint output + LSE, # joint lse + # Q strides: [B, H, Sq, D] + stride_qb, + stride_qh, + stride_qm, + stride_qd, + # Pool strides: [B, P, D] + stride_pb, + stride_pp, + stride_pd, + # topk_idxs strides: [B, Sq, K_topk] + stride_tib, + stride_tim, + stride_tik, + # OUT_LOCAL strides: [B, H, Sq, D] + stride_olb, + stride_olh, + stride_olm, + stride_old, + # LSE_LOCAL strides: [B, H, Sq] + stride_llb, + stride_llh, + stride_llm, + # OUT (joint) strides: [B, H, Sq, D] + stride_ob, + stride_oh, + stride_om, + stride_od, + # LSE (joint) strides: [B, H, Sq] + stride_lb, + stride_lh, + stride_lm, + seqlen_q, + pool_size, + K_topk, + sm_scale, + HEAD_Q: tl.constexpr, + HAS_SINK: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, + K_DIVISIBLE: tl.constexpr, # True iff K_topk % BLOCK_K == 0 + H_DIVISIBLE: tl.constexpr, # True iff HEAD_Q % BLOCK_H == 0 +): + """Sparse-branch CSA FWD + joint merge fused. + + Plan-8 P57 cr=4 FWD speedup: eliminates the separate + ``_v4_csa_attention_lse_merge_kernel`` launch by reading + ``out_local`` / ``lse_local`` at the tail of this kernel and writing + the joint output directly. Saves one kernel launch (~10 us) plus the + intermediate ``[B, H, Sq, D]`` ``out_sparse`` HBM round-trip + (~250 MiB read + 250 MiB write = ~100 us at MI355X HBM speed). + + Grid: ``(seqlen_q, cdiv(HEAD_Q, BLOCK_H), B)`` — same as the + non-fused sparse kernel. Each program owns one ``(b, h_block, m)`` + and produces the joint ``out[BLOCK_H, BLOCK_DMODEL]`` row. + + Math (joint softmax over local + sparse + sink, all with + ``-m_max`` rescaling): + + m_max = max(lse_local, lse_sparse, sink_h) + denom = exp(lse_local - m_max) + + exp(lse_sparse - m_max) + + exp(sink_h - m_max) # 0 if no sink + joint_out = (out_local * exp(lse_local - m_max) + + out_sparse * exp(lse_sparse - m_max)) / denom + joint_lse = m_max + log(denom) + + where ``out_sparse = acc / l_sparse`` is the normalized sparse + output (computed inside this kernel, never materialised to HBM). + """ + pid_m = tl.program_id(0) + pid_h_block = tl.program_id(1) + bid = tl.program_id(2) + + offs_h = pid_h_block * BLOCK_H + tl.arange(0, BLOCK_H) + offs_k = tl.arange(0, BLOCK_K) + offs_d = tl.arange(0, BLOCK_DMODEL) + if H_DIVISIBLE: + h_mask = tl.full([BLOCK_H], True, dtype=tl.int1) + else: + h_mask = offs_h < HEAD_Q + + NEG_INF: tl.constexpr = -1.0e30 + + # Q tile: [BLOCK_H, BLOCK_DMODEL] + q_ptrs = ( + Q + bid * stride_qb + offs_h[:, None] * stride_qh + pid_m * stride_qm + offs_d[None, :] * stride_qd + ) + if H_DIVISIBLE: + q = tl.load(q_ptrs) + else: + q = tl.load(q_ptrs, mask=h_mask[:, None], other=0.0) + + # Online-softmax state for the sparse branch. + acc = tl.zeros([BLOCK_H, BLOCK_DMODEL], dtype=tl.float32) + m_i = tl.full([BLOCK_H], value=NEG_INF, dtype=tl.float32) + l_i = tl.zeros([BLOCK_H], dtype=tl.float32) + + # Hoist loop-invariant pointer offsets out of the K-loop. Triton + # usually does this but spelling it out keeps the inner body lean. + topk_base = TOPK_IDXS + bid * stride_tib + pid_m * stride_tim + pool_base = POOL + bid * stride_pb + pool_d_offs = offs_d[None, :] * stride_pd + + for k_start in range(0, K_topk, BLOCK_K): + sparse_k = k_start + offs_k + topk_ptrs = topk_base + sparse_k * stride_tik + if K_DIVISIBLE: + topk = tl.load(topk_ptrs) + valid_k = (topk >= 0) & (topk < pool_size) + else: + topk = tl.load(topk_ptrs, mask=sparse_k < K_topk, other=-1) + valid_k = (sparse_k < K_topk) & (topk >= 0) & (topk < pool_size) + safe_topk = tl.where(valid_k, topk, 0) + + pool_ptrs = pool_base + safe_topk[:, None] * stride_pp + pool_d_offs + pool = tl.load(pool_ptrs, mask=valid_k[:, None], other=0.0) + + qk = tl.dot(q.to(pool.dtype), tl.trans(pool), out_dtype=tl.float32) * sm_scale + if H_DIVISIBLE: + qk = tl.where(valid_k[None, :], qk, NEG_INF) + else: + qk = tl.where(h_mask[:, None] & valid_k[None, :], qk, NEG_INF) + + m_tile = tl.max(qk, axis=1) + m_new = tl.maximum(m_i, m_tile) + alpha = tl.exp(m_i - m_new) + p = tl.exp(qk - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(pool.dtype), pool, out_dtype=tl.float32) + m_i = m_new + + # Detect "all invalid" sparse rows (every topk_idx == -1 or + # K_topk == 0): l_sparse stays zero so the sparse branch carries + # zero softmax mass for that (b, h, m). + empty_sparse = l_i == 0.0 + safe_l_sparse = tl.where(empty_sparse, 1.0, l_i) + out_sparse_norm = acc / safe_l_sparse[:, None] # [BLOCK_H, BLOCK_DMODEL] fp32 + lse_sparse = tl.where(empty_sparse, NEG_INF, m_i + tl.log(safe_l_sparse)) + + # ---- Load local-branch output + lse and merge with sink ------------- + out_local_ptrs = ( + OUT_LOCAL + + bid * stride_olb + + offs_h[:, None] * stride_olh + + pid_m * stride_olm + + offs_d[None, :] * stride_old + ) + out_local = tl.load(out_local_ptrs, mask=h_mask[:, None], other=0.0).to(tl.float32) + + lse_local_ptrs = LSE_LOCAL + bid * stride_llb + offs_h * stride_llh + pid_m * stride_llm + lse_local = tl.load(lse_local_ptrs, mask=h_mask, other=NEG_INF) + + if HAS_SINK: + sink_h = tl.load(SINK + offs_h, mask=h_mask, other=NEG_INF).to(tl.float32) + else: + sink_h = tl.full([BLOCK_H], value=NEG_INF, dtype=tl.float32) + + m_max = tl.maximum(lse_local, lse_sparse) + if HAS_SINK: + m_max = tl.maximum(m_max, sink_h) + + alpha_local = tl.exp(lse_local - m_max) + alpha_sparse = tl.exp(lse_sparse - m_max) + if HAS_SINK: + alpha_sink = tl.exp(sink_h - m_max) + else: + alpha_sink = tl.zeros([BLOCK_H], dtype=tl.float32) + + denom = alpha_local + alpha_sparse + alpha_sink + # Empty-row safety: if every branch is NEG_INF, denom == 0. Use a + # safe denom and rely on the alpha_* being zero so the numerator is + # zero too — output stays at 0 / 1 = 0, lse remains NEG_INF. + safe_denom = tl.where(denom == 0.0, 1.0, denom) + + joint_out = (out_local * alpha_local[:, None] + out_sparse_norm * alpha_sparse[:, None]) / safe_denom[ + :, None + ] + joint_lse = tl.where(denom == 0.0, NEG_INF, m_max + tl.log(safe_denom)) + + out_ptrs = ( + OUT + bid * stride_ob + offs_h[:, None] * stride_oh + pid_m * stride_om + offs_d[None, :] * stride_od + ) + tl.store(out_ptrs, joint_out.to(OUT.dtype.element_ty), mask=h_mask[:, None]) + + lse_ptrs = LSE + bid * stride_lb + offs_h * stride_lh + pid_m * stride_lm + tl.store(lse_ptrs, joint_lse, mask=h_mask) + + +@triton.jit +def _v4_csa_attention_lse_merge_kernel( + OUT_LOCAL, + LSE_LOCAL, + OUT_SPARSE, + LSE_SPARSE, + SINK, + OUT, + LSE, + stride_olb, + stride_olh, + stride_olm, + stride_old, + stride_llb, + stride_llh, + stride_llm, + stride_osb, + stride_osh, + stride_osm, + stride_osd, + stride_lsb, + stride_lsh, + stride_lsm, + stride_ob, + stride_oh, + stride_om, + stride_od, + stride_lb, + stride_lh, + stride_lm, + seqlen_q, + HEAD_Q: tl.constexpr, + HAS_SINK: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL: tl.constexpr, +): + """Merge ``(out_local, lse_local)`` and ``(out_sparse, lse_sparse)`` + with the per-head sink under one final online softmax. + + Grid: ``(cdiv(seqlen_q, BLOCK_M), B * HEAD_Q)``. Each program owns + one ``[BLOCK_M, BLOCK_DMODEL]`` slice of the joint output. + """ + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + bid = pid_bh // HEAD_Q + qhid = pid_bh % HEAD_Q + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, BLOCK_DMODEL) + m_mask = offs_m < seqlen_q + + NEG_INF: tl.constexpr = -1.0e30 + + lse_local = tl.load( + LSE_LOCAL + bid * stride_llb + qhid * stride_llh + offs_m * stride_llm, + mask=m_mask, + other=NEG_INF, + ) + lse_sparse = tl.load( + LSE_SPARSE + bid * stride_lsb + qhid * stride_lsh + offs_m * stride_lsm, + mask=m_mask, + other=NEG_INF, + ) + + out_local = tl.load( + OUT_LOCAL + + bid * stride_olb + + qhid * stride_olh + + offs_m[:, None] * stride_olm + + offs_d[None, :] * stride_old, + mask=m_mask[:, None], + other=0.0, + ).to(tl.float32) + out_sparse = tl.load( + OUT_SPARSE + + bid * stride_osb + + qhid * stride_osh + + offs_m[:, None] * stride_osm + + offs_d[None, :] * stride_osd, + mask=m_mask[:, None], + other=0.0, + ).to(tl.float32) + + if HAS_SINK: + sink_h = tl.load(SINK + qhid).to(tl.float32) + else: + sink_h = NEG_INF + + m_max = tl.maximum(lse_local, lse_sparse) + if HAS_SINK: + m_max = tl.maximum(m_max, sink_h) + + alpha_local = tl.exp(lse_local - m_max) + alpha_sparse = tl.exp(lse_sparse - m_max) + if HAS_SINK: + alpha_sink = tl.exp(sink_h - m_max) + else: + alpha_sink = tl.zeros_like(alpha_local) + + denom = alpha_local + alpha_sparse + alpha_sink + # Empty-row safety: if every branch is NEG_INF, denom == 0. Use a + # safe denom and rely on the alpha_* being zero so the numerator is + # zero too — output stays at 0 / 1 = 0, lse remains NEG_INF. + safe_denom = tl.where(denom == 0.0, 1.0, denom) + + out = (out_local * alpha_local[:, None] + out_sparse * alpha_sparse[:, None]) / safe_denom[:, None] + lse = tl.where(denom == 0.0, NEG_INF, m_max + tl.log(safe_denom)) + + out_ptrs = ( + OUT + bid * stride_ob + qhid * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od + ) + tl.store(out_ptrs, out.to(OUT.dtype.element_ty), mask=m_mask[:, None]) + + lse_ptrs = LSE + bid * stride_lb + qhid * stride_lh + offs_m * stride_lm + tl.store(lse_ptrs, lse, mask=m_mask) + + +def _launch_v4_csa_attention_pool_fwd_split( + q: torch.Tensor, + k_local: torch.Tensor, + v_local: torch.Tensor, + pool: torch.Tensor, + topk_idxs: torch.Tensor, + *, + sink: Optional[torch.Tensor], + swa_window: int, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """P32 split CSA FWD: local SWA via dense kernel + sparse head-block + LSE merge.""" + B, HQ, Sq, D = q.shape + P = pool.shape[1] + K_topk = topk_idxs.shape[2] + has_sink = sink is not None + + # Step 1: local SWA branch (no sink — applied in the merge). + out_local, lse_local = _launch_v4_attention_fwd( + q, + k_local, + v_local, + sink=None, + swa_window=int(swa_window) if swa_window > 0 else 0, + additive_mask=None, + scale=float(scale), + hca_local_seqlen=0, + ) + + # P57 attempt-11: tile sweep on the fused sparse+merge kernel found + # (BLOCK_H=64, BLOCK_K=16, num_warps=8, num_stages=3) wins on the + # V4-Flash widths (H=64 D=512). The intuition: + # + # * BLOCK_H=64 keeps the full head axis in a single program, so the + # pool gather (which is per-(b, m) — no H dim) is read once and + # reused across all H query heads. + # * BLOCK_K=16 halves the per-tile pool tile (16 × 512 × 2 B = 16 KiB) + # so the AMD Triton software pipeliner can keep 3 stages live in + # LDS (3 × 16 KiB = 48 KiB pool buffer) without crowding out the + # `acc` register accumulator. With BLOCK_K=32 the pool tile is + # 2× larger and 3-stage pipelining ran out of resources. + # * num_stages=3 overlaps pool gather → `tl.dot` → softmax update + # across 3 K-tile iterations, hiding the gather latency. + # + # P57 R2 re-confirmed the BLOCK_K/stages winner via a fresh BK x + # ST x NW sweep (`p57/r2_sweep.sh`); BLOCK_K=32 ran out of LDS at + # num_stages>=2 and regressed 2-3 x. The R2 win for cr=4 FWD comes + # from the *local-SWA* kernel re-tile, not this kernel. + # BLOCK_H must be >= 16 so the head axis maps to a valid gfx1250 WMMA + # M-tile. Without the floor, HQ in {1, 2, 4, 8} (a power of two below + # 16 — e.g. high tensor-parallel head sharding, or tiny test shapes) + # leaves ``next_power_of_2(HQ) == HQ`` and the old ``BLOCK_H > HQ`` + # guard never fired, so BLOCK_H stayed < 16 and the sparse ``tl.dot`` + # failed to select a matrix-core intrinsic ("no matching matrix core + # intrinsic for wmma version 3 ... instruction shape [0, 0, K]"). + # Surplus heads in the 16-wide tile are masked by H_DIVISIBLE/h_mask. + BLOCK_H = 64 if HQ >= 64 else max(triton.next_power_of_2(HQ), 16) + BLOCK_K = 16 + BLOCK_DMODEL = D + sparse_grid = (Sq, triton.cdiv(HQ, BLOCK_H), B) + sink_arg = sink if has_sink else q + + # P57 attempt-3: fuse sparse + merge into one kernel. The legacy + # 2-kernel path (sparse → ``out_sparse``/``lse_sparse``; merge) + # required a ~500 MiB intermediate HBM round-trip plus a ~135 us + # merge launch. The fused kernel reads ``out_local`` / ``lse_local`` + # at the tail of the sparse loop and writes joint ``out`` / ``lse`` + # directly. Env-gated A/B fallback: + # ``PRIMUS_V4_CSA_FWD_SEPARATE_MERGE=1`` keeps the split layout for + # parity debugging. + out = torch.empty_like(q) + lse = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32) + + if os.getenv("PRIMUS_V4_CSA_FWD_SEPARATE_MERGE", "0") != "1": + _v4_csa_attention_pool_sparse_merge_fwd_kernel[sparse_grid]( + q, + pool, + topk_idxs, + out_local, + lse_local, + sink_arg, + out, + lse, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + pool.stride(0), + pool.stride(1), + pool.stride(2), + topk_idxs.stride(0), + topk_idxs.stride(1), + topk_idxs.stride(2), + out_local.stride(0), + out_local.stride(1), + out_local.stride(2), + out_local.stride(3), + lse_local.stride(0), + lse_local.stride(1), + lse_local.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + Sq, + P, + K_topk, + float(scale), + HEAD_Q=HQ, + HAS_SINK=has_sink, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K, + BLOCK_DMODEL=BLOCK_DMODEL, + K_DIVISIBLE=(K_topk % BLOCK_K == 0), + H_DIVISIBLE=(HQ % BLOCK_H == 0), + num_warps=8, + num_stages=3, + ) + return out, lse + + # ---- legacy split-merge path (env-gated for A/B debugging) ---------- + out_sparse = torch.empty_like(q) + lse_sparse = torch.empty((B, HQ, Sq), device=q.device, dtype=torch.float32) + _v4_csa_attention_pool_sparse_fwd_kernel[sparse_grid]( + q, + pool, + topk_idxs, + out_sparse, + lse_sparse, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + pool.stride(0), + pool.stride(1), + pool.stride(2), + topk_idxs.stride(0), + topk_idxs.stride(1), + topk_idxs.stride(2), + out_sparse.stride(0), + out_sparse.stride(1), + out_sparse.stride(2), + out_sparse.stride(3), + lse_sparse.stride(0), + lse_sparse.stride(1), + lse_sparse.stride(2), + Sq, + P, + K_topk, + float(scale), + HEAD_Q=HQ, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=8, + num_stages=1, + ) + MERGE_BLOCK_M = 32 + merge_grid = (triton.cdiv(Sq, MERGE_BLOCK_M), B * HQ) + _v4_csa_attention_lse_merge_kernel[merge_grid]( + out_local, + lse_local, + out_sparse, + lse_sparse, + sink_arg, + out, + lse, + out_local.stride(0), + out_local.stride(1), + out_local.stride(2), + out_local.stride(3), + lse_local.stride(0), + lse_local.stride(1), + lse_local.stride(2), + out_sparse.stride(0), + out_sparse.stride(1), + out_sparse.stride(2), + out_sparse.stride(3), + lse_sparse.stride(0), + lse_sparse.stride(1), + lse_sparse.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + Sq, + HEAD_Q=HQ, + HAS_SINK=has_sink, + BLOCK_M=MERGE_BLOCK_M, + BLOCK_DMODEL=BLOCK_DMODEL, + num_warps=4, + num_stages=1, + ) + return out, lse + + +__all__ = [ + "_v4_csa_attention_fwd_kernel", + "_v4_csa_attention_pool_fwd_kernel", + "_v4_csa_attention_pool_sparse_fwd_kernel", + "_v4_csa_attention_pool_sparse_merge_fwd_kernel", + "_v4_csa_attention_lse_merge_kernel", + "_launch_v4_csa_attention_fwd", + "_launch_v4_csa_attention_pool_fwd", + "_launch_v4_csa_attention_pool_fwd_split", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/__init__.py new file mode 100644 index 000000000..5aa001f2e --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/__init__.py @@ -0,0 +1,24 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plain-Triton DeepSeek-V4 sparse-MLA attention backend ("triton v2"). + +Same fused single-latent (K == V) sparse-MLA representation and public API as +the gluon backend, but written in vanilla Triton so the QK / PV GEMMs lower to +MFMA via ``tl.dot``. Distinct from the separate-KV Triton CSA/dense backends +(``_triton_v0_deprecated`` gathered / ``_triton_v1`` pool), which keep K and V separate. + +* :func:`sparse_mla_fwd_v4_triton` -> ``(o, lse)`` +* :func:`sparse_mla_bwd_v4_triton` -> ``(dq, dkv, d_sink)`` +""" + +from .dsa_bwd_v4_triton import sparse_mla_bwd_v4_triton +from .dsa_fwd_v4_triton import sparse_mla_fwd_v4_triton + +__all__ = [ + "sparse_mla_fwd_v4_triton", + "sparse_mla_bwd_v4_triton", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/_amd_knobs.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/_amd_knobs.py new file mode 100644 index 000000000..2c910de7e --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/_amd_knobs.py @@ -0,0 +1,44 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Shared AMD Triton compiler-knob helpers for the triton_v2 sparse-MLA kernels. + +primus_turbo globally enables ``TRITON_HIP_USE_BLOCK_PINGPONG`` / +``TRITON_HIP_USE_ASYNC_COPY`` (via ``set_triton_knobs_gfx950``) as soon as it is +imported by the training runtime. Those knobs double-buffer (ping-pong) a +kernel's LDS operand tiles. Measured on gfx950 (bench_v4_attention, flash) they +are a *pessimization* for the V4 sparse-MLA kernels — the forward is ~16-29% +slower and the backward is slightly slower with them on — and they also overflow +the 160 KB LDS limit for the wide (BH=64/TK=128) dKV tiling. + +These knobs are read at *compile time* and are NOT part of Triton's compile +cache key (``HIPOptions.hash`` does not hash them), so compiling a kernel inside +:func:`amd_pingpong_disabled` pins that kernel to the non-ping-pong schedule for +the whole process, while restoring the knobs on exit leaves every other kernel +(compiled outside the scope) exactly as primus_turbo configured it. +""" + +import contextlib + +import triton + + +@contextlib.contextmanager +def amd_pingpong_disabled(): + """Temporarily disable the AMD Triton ping-pong / async-copy LDS knobs.""" + amd = getattr(getattr(triton, "knobs", None), "amd", None) + if amd is None: + yield + return + prev_pp = amd.use_block_pingpong + prev_ac = amd.use_async_copy + try: + amd.use_block_pingpong = False + amd.use_async_copy = False + yield + finally: + amd.use_block_pingpong = prev_pp + amd.use_async_copy = prev_ac diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_kernels.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_kernels.py new file mode 100644 index 000000000..8f2478801 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_kernels.py @@ -0,0 +1,337 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""Owned plain-Triton backward compute kernels for the "triton v2" sparse-MLA +backend. Forked from the shared ``_gluon_dsa/_dsa_bwd_gather.py`` plain-Triton +kernels so this backend can be tuned independently of the gluon path. + +Three kernels implement the non-atomic chunked-gather backward: + * ``_bwd_chunk_dq_store_ds`` — dQ accumulation for one rank chunk, ALSO stores + per-tile dS and P to [T, H, R_CHUNK] buffers for reuse by the dKV kernel. + * ``_bwd_compute_dkv_intermediate`` — dKV intermediate for one chunk, REUSES + the stored dS/P (no S/P/dS recompute). Consumes q/do transposed to [T, D, H]. + * ``_bwd_dkv_gather_acc`` — CSR inverted-topk reduce of the intermediate into + the fp32 dKV accumulator. +""" + +import triton +import triton.language as tl + + +@triton.jit +def _bwd_chunk_dq_store_ds( + Q_ptr, # [T, H, D] bf16 + KV_ptr, # [T, 1, D] bf16 + dO_ptr, # [T, H, D_V] bf16 + TopK_ptr, # [T, TOPK] int32 + LSE_ptr, # [T, H] fp32 + Delta_ptr, # [T, H] fp32 (computed here on the first chunk, read after) + O_ptr, # [T, H, D_V] bf16 (fwd output, for Delta = rowsum(O*dO)) + dQ_ptr, # [T, H, D] bf16 — read-modify-write across chunks + dS_ptr, # [T, H, R_CHUNK] bf16 — output chunk dS + P_ptr, # [T, H, R_CHUNK] bf16 — output chunk P + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_o_t: tl.int64, + stride_o_h: tl.int64, + stride_dq_t: tl.int64, + stride_dq_h: tl.int64, + stride_topk_t: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + R_START: tl.int32, + R_CHUNK: tl.constexpr, + BLOCK_H: tl.constexpr, + TILE_K: tl.constexpr, + D_V: tl.constexpr, + D_ROPE: tl.constexpr, + HAS_ROPE: tl.constexpr, + IS_FIRST_CHUNK: tl.constexpr, +): + """dQ accumulation for rank chunk [R_START, R_START+R_CHUNK), plus stores + chunk dS and P to [T, H, R_CHUNK] buffers for _bwd_compute_dkv_intermediate. + Grid: (total_tokens, num_hg). + + Delta = rowsum(O*dO) is fused here (computed + stored on the first chunk, + reloaded on later chunks) instead of a separate preprocess kernel — dO is + already resident, so this only adds the O load and drops a whole kernel + a + duplicate dO read. + + HAS_ROPE=False (V4 zero-pad): skips all rope MMAs/loads and writes the dQ rope + columns as zero (they are discarded downstream but kept well-defined).""" + token_idx = tl.program_id(0) + hg_idx = tl.program_id(1) + offs_h = hg_idx * BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = offs_h < num_heads + offs_v = tl.arange(0, D_V) + offs_r = tl.arange(0, D_ROPE) + + q_base = token_idx * stride_q_t + Q_lora = tl.load( + Q_ptr + q_base + offs_h[:, None] * stride_q_h + offs_v[None, :], mask=mask_h[:, None], other=0.0 + ) + if HAS_ROPE: + Q_rope = tl.load( + Q_ptr + q_base + offs_h[:, None] * stride_q_h + (D_V + offs_r[None, :]), + mask=mask_h[:, None], + other=0.0, + ) + do_base = token_idx * stride_do_t + dO_val = tl.load( + dO_ptr + do_base + offs_h[:, None] * stride_do_h + offs_v[None, :], mask=mask_h[:, None], other=0.0 + ) + lse = tl.load(LSE_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0) + if IS_FIRST_CHUNK: + # Delta = rowsum(O * dO): fold the preprocess in (dO already loaded). + O_val = tl.load( + O_ptr + token_idx * stride_o_t + offs_h[:, None] * stride_o_h + offs_v[None, :], + mask=mask_h[:, None], + other=0.0, + ) + delta = tl.sum(O_val.to(tl.float32) * dO_val.to(tl.float32), axis=1) + tl.store(Delta_ptr + token_idx * num_heads + offs_h, delta, mask=mask_h) + else: + delta = tl.load(Delta_ptr + token_idx * num_heads + offs_h, mask=mask_h, other=0.0) + + dq_base = token_idx * stride_dq_t + if IS_FIRST_CHUNK: + dQ_lora = tl.zeros([BLOCK_H, D_V], dtype=tl.float32) + else: + dQ_lora = tl.load( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :], + mask=mask_h[:, None], + other=0.0, + ).to(tl.float32) + if HAS_ROPE: + if IS_FIRST_CHUNK: + dQ_rope = tl.zeros([BLOCK_H, D_ROPE], dtype=tl.float32) + else: + dQ_rope = tl.load( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]), + mask=mask_h[:, None], + other=0.0, + ).to(tl.float32) + + NUM_TILES: tl.constexpr = (R_CHUNK + TILE_K - 1) // TILE_K + topk_base = token_idx * stride_topk_t + R_START + offs_tile = tl.arange(0, TILE_K) + ds_base = token_idx * stride_ds_t + hg_idx * BLOCK_H * stride_ds_h + + for t in range(NUM_TILES): + tile_start = t * TILE_K + tile_offs = tile_start + offs_tile + valid = tile_offs < R_CHUNK + topk_pos = tl.load(TopK_ptr + topk_base + tile_offs, mask=valid, other=-1) + valid = valid & (topk_pos != -1) + safe_pos = tl.where(valid, topk_pos, 0) + + K_lora_T = tl.load( + KV_ptr + safe_pos[None, :] * stride_kv_t + offs_v[:, None], mask=valid[None, :], other=0.0 + ) + + S = tl.dot(Q_lora, K_lora_T) + if HAS_ROPE: + K_rope_T = tl.load( + KV_ptr + safe_pos[None, :] * stride_kv_t + (D_V + offs_r[:, None]), + mask=valid[None, :], + other=0.0, + ) + S += tl.dot(Q_rope, K_rope_T) + S = tl.where(valid[None, :] & mask_h[:, None], S * scale, float("-inf")) + P = tl.exp(S - lse[:, None]) + P = tl.where(valid[None, :] & mask_h[:, None], P, 0.0) + dP = tl.dot(dO_val, K_lora_T) + dS = P * (dP - delta[:, None]) * scale + dS = tl.where(valid[None, :] & mask_h[:, None], dS, 0.0) + + dS_bf = dS.to(tl.bfloat16) + dQ_lora += tl.dot(dS_bf, tl.trans(K_lora_T)).to(tl.float32) + if HAS_ROPE: + dQ_rope += tl.dot(dS_bf, tl.trans(K_rope_T)).to(tl.float32) + + local_h = tl.arange(0, BLOCK_H) + tl.store( + dS_ptr + ds_base + local_h[:, None] * stride_ds_h + tile_offs[None, :], + dS_bf, + mask=mask_h[:, None] & valid[None, :], + ) + tl.store( + P_ptr + ds_base + local_h[:, None] * stride_ds_h + tile_offs[None, :], + P.to(tl.bfloat16), + mask=mask_h[:, None] & valid[None, :], + ) + + tl.store( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + offs_v[None, :], + dQ_lora.to(Q_lora.dtype), + mask=mask_h[:, None], + ) + if HAS_ROPE: + tl.store( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]), + dQ_rope.to(Q_lora.dtype), + mask=mask_h[:, None], + ) + else: + tl.store( + dQ_ptr + dq_base + offs_h[:, None] * stride_dq_h + (D_V + offs_r[None, :]), + tl.zeros([BLOCK_H, D_ROPE], dtype=Q_lora.dtype), + mask=mask_h[:, None], + ) + + +@triton.jit +def _bwd_compute_dkv_intermediate( + Q_ptr, # [T, H, D_QK] bf16 (UNtransposed; loaded transposed via strided index) + dO_ptr, # [T, H, D_V] bf16 (UNtransposed) + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + Interm_ptr, # [T, R_CHUNK, D_QK] bf16 — output, one writer per (q, rank) + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + stride_interm_t: tl.int64, # R_CHUNK * D_QK + stride_interm_k: tl.int64, # D_QK + num_heads: tl.int32, + R_CHUNK: tl.constexpr, + TILE_K: tl.constexpr, + BLOCK_H: tl.constexpr, + NUM_HG: tl.constexpr, + D_V: tl.constexpr, + D_ROPE: tl.constexpr, + HAS_ROPE: tl.constexpr, +): + """dKV intermediate for one chunk, REUSING stored dS/P (no recompute). + Loads Q/dO UNtransposed with a [D, BLOCK_H] strided index pattern (contiguous + along D per head), removing the external q.transpose(1,2).contiguous() / + do.transpose(...).contiguous() copies. Grid: (total_tokens,). + + HAS_ROPE=False (V4 zero-pad): skips Q_rope load, the dKV_rope MMA and the + interm rope store (interm rope columns are never read by the gather).""" + token_idx = tl.program_id(0) + + NUM_TILES: tl.constexpr = (R_CHUNK + TILE_K - 1) // TILE_K + offs_tile = tl.arange(0, TILE_K) + offs_v = tl.arange(0, D_V) + offs_r = tl.arange(0, D_ROPE) + + interm_base_t = token_idx * stride_interm_t + q_base = token_idx * stride_q_t + do_base = token_idx * stride_do_t + ds_base = token_idx * stride_ds_t + + for t in range(NUM_TILES): + tile_start = t * TILE_K + tile_offs = tile_start + offs_tile + valid = tile_offs < R_CHUNK + + dKV_lora = tl.zeros([D_V, TILE_K], dtype=tl.float32) + if HAS_ROPE: + dKV_rope = tl.zeros([D_ROPE, TILE_K], dtype=tl.float32) + + for hg in range(NUM_HG): + offs_h = hg * BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = offs_h < num_heads + + # [D_V, BLOCK_H] loaded transposed: rows=offs_v (stride 1, contiguous), + # cols=offs_h (stride stride_q_h). No HBM transpose copy needed. + Q_lora_T = tl.load( + Q_ptr + q_base + offs_h[None, :] * stride_q_h + offs_v[:, None], + mask=mask_h[None, :], + other=0.0, + ) + dO_T = tl.load( + dO_ptr + do_base + offs_h[None, :] * stride_do_h + offs_v[:, None], + mask=mask_h[None, :], + other=0.0, + ) + + dS_val = tl.load( + dS_ptr + ds_base + offs_h[:, None] * stride_ds_h + tile_offs[None, :], + mask=mask_h[:, None] & valid[None, :], + other=0.0, + ) + P_val = tl.load( + P_ptr + ds_base + offs_h[:, None] * stride_ds_h + tile_offs[None, :], + mask=mask_h[:, None] & valid[None, :], + other=0.0, + ) + + dKV_lora += tl.dot(Q_lora_T, dS_val.to(Q_lora_T.dtype)).to(tl.float32) + dKV_lora += tl.dot(dO_T, P_val.to(dO_T.dtype)).to(tl.float32) + if HAS_ROPE: + Q_rope_T = tl.load( + Q_ptr + q_base + offs_h[None, :] * stride_q_h + (D_V + offs_r[:, None]), + mask=mask_h[None, :], + other=0.0, + ) + dKV_rope += tl.dot(Q_rope_T, dS_val.to(Q_rope_T.dtype)).to(tl.float32) + + interm_lora_ptrs = Interm_ptr + interm_base_t + tile_offs[None, :] * stride_interm_k + offs_v[:, None] + tl.store(interm_lora_ptrs, dKV_lora.to(tl.bfloat16), mask=valid[None, :]) + + if HAS_ROPE: + interm_rope_ptrs = ( + Interm_ptr + interm_base_t + tile_offs[None, :] * stride_interm_k + D_V + offs_r[:, None] + ) + tl.store(interm_rope_ptrs, dKV_rope.to(tl.bfloat16), mask=valid[None, :]) + + +@triton.jit +def _bwd_dkv_gather_acc( + Interm_ptr, # [T, R_CHUNK, D] bf16 — chunk intermediate + InvPtr_ptr, # [T+1] int32 — CSR row pointers + InvData_ptr, # [T*R_CHUNK] int32 — encoded as q*R_CHUNK + local_r + dKV_acc_ptr, # [T, D] fp32 — accumulator (read-modify-write across chunks) + stride_interm_r: tl.int64, # D + stride_acc_t: tl.int64, # D + D_V: tl.constexpr, + D_ROPE: tl.constexpr, + HAS_ROPE: tl.constexpr, + BLOCK_K: tl.constexpr = 64, +): + """Gather one chunk's bf16 intermediate into the fp32 dKV accumulator. + Grid: (num_kv,) — one CTA per KV token, no atomics. Tiled BLOCK_K rows/iter. + + HAS_ROPE=False (V4 zero-pad): interm rope columns are never written, so skip + reading/accumulating them; dKV_acc rope stays at its zero-init value.""" + k = tl.program_id(0) + offs_v = tl.arange(0, D_V) + offs_r = tl.arange(0, D_ROPE) + offs_k = tl.arange(0, BLOCK_K) + + start = tl.load(InvPtr_ptr + k) + end = tl.load(InvPtr_ptr + k + 1) + + acc_base = k.to(tl.int64) * stride_acc_t + dkv_acc_lora = tl.load(dKV_acc_ptr + acc_base + offs_v).to(tl.float32) + if HAS_ROPE: + dkv_acc_rope = tl.load(dKV_acc_ptr + acc_base + D_V + offs_r).to(tl.float32) + + n_entries = end - start + for ti in range(0, tl.cdiv(n_entries, BLOCK_K)): + e_local = ti * BLOCK_K + offs_k + valid = e_local < n_entries + entry = tl.load(InvData_ptr + start + e_local, mask=valid, other=0).to(tl.int64) + rp = entry[:, None] * stride_interm_r + rows_v = tl.load(Interm_ptr + rp + offs_v[None, :], mask=valid[:, None], other=0.0).to(tl.float32) + dkv_acc_lora += tl.sum(rows_v, axis=0) + if HAS_ROPE: + rows_r = tl.load(Interm_ptr + rp + D_V + offs_r[None, :], mask=valid[:, None], other=0.0).to( + tl.float32 + ) + dkv_acc_rope += tl.sum(rows_r, axis=0) + + tl.store(dKV_acc_ptr + acc_base + offs_v, dkv_acc_lora) + if HAS_ROPE: + tl.store(dKV_acc_ptr + acc_base + D_V + offs_r, dkv_acc_rope) diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py new file mode 100644 index 000000000..1256d7ade --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_bwd_v4_triton.py @@ -0,0 +1,229 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plain-Triton DeepSeek-V4 sparse-MLA backward (the "triton v2" backend). + +Companion to :func:`sparse_mla_fwd_v4_triton`; API mirrors +:func:`sparse_mla_bwd_v4_gluon` -> ``(dq, dkv, d_sink)``. + +Uses the **non-atomic chunked-gather** scheme (the same one the gluon backward +shares for its dKV gather): per rank-chunk it runs a pure-Triton dQ kernel and a +dKV-intermediate kernel (both ``tl.dot`` / MFMA, plain stores — no atomics), +then reduces the intermediate into ``dkv`` via a CSR inverted-topk gather. This +is the fully-Triton analogue of the gluon dQ/dKV-interm kernels, so its dKV is +not bottlenecked by global atomics. ``d_sink`` is the closed-form torch +reduction ``-sum_t exp(sink - lse) * delta``. +""" + +import contextlib +import os + +import torch +import triton + +# CSR-inverted-topk builder + Delta preprocess are backend-neutral torch/host +# helpers; reuse them. The compute kernels are owned locally (dsa_bwd_kernels) +# so this backend can be tuned independently of the gluon path. +from .._gluon_dsa._dsa_bwd_gather import _build_inverted_topk_slice +from ._amd_knobs import amd_pingpong_disabled +from .dsa_bwd_kernels import ( + _bwd_chunk_dq_store_ds, + _bwd_compute_dkv_intermediate, + _bwd_dkv_gather_acc, +) + + +def sparse_mla_bwd_v4_triton(q, kv, o, do, topk_indices, lse, attn_sink=None, kv_lora_rank=512, scale=None): + """DeepSeek-V4 sparse-MLA backward (plain Triton / MFMA, non-atomic). + + Returns ``(dq, dkv, d_sink)`` with ``dkv`` shaped like ``kv`` and ``d_sink`` + ``[num_heads]`` fp32 (or ``None`` when ``attn_sink`` is None). + """ + assert q.is_contiguous() and o.is_contiguous() and do.is_contiguous() + assert topk_indices.is_contiguous() and lse.is_contiguous() + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + topk = topk_indices.shape[1] + if scale is None: + scale = 1.0 / (d_qk**0.5) + if kv.dim() == 2: + kv = kv.unsqueeze(1) + assert kv.is_contiguous() + num_kv = kv.shape[0] + + has_sink = attn_sink is not None + if has_sink: + assert attn_sink.dtype == torch.float32 and attn_sink.shape == (num_heads,) + + # Delta = rowsum(O*dO) is fused into the first dQ chunk (no preprocess kernel). + delta = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device) + + # ---- config (mirror the gluon bwd chunking) ---- + # R_CHUNK (rank-chunk width): dQ is read-modify-written across chunks, so more + # chunks = more redundant dq reload passes + repeated launches/CSR builds. For + # high head counts (H>=128) the dq RMW volume is large, so a single chunk over + # the whole topk (bounded for memory) is a big win (-22% on pro cr4). For low + # head counts (H<=64) the smaller dq RMW is outweighed by the larger per-chunk + # buffers/occupancy, and 256 stays best — so keep the cap there. + if num_heads >= 128: + R_CHUNK = min(topk, 1536) + else: + R_CHUNK = min(256, topk) + BH_DQ, TK_DQ = 64, 16 + # dKV-intermediate tiling. The default (BH_DKV=32, TK_DKV=64) is best for high + # head counts (H>=128) and for chunk widths that are not 128-aligned. For low + # head counts (H<=64) with a 128-aligned chunk, a wider TILE_K=128 over a single + # head-group (BH_DKV=64, NUM_HG=1) reduces redundant Q/dO re-loads and issues + # fuller MMAs (measured ~6-7% faster on the full flash bwd: cr=0 1.21->1.14 ms, + # cr=4 5.44->5.09 ms in bench_v4_attention); it regresses for H>=128 (register + # pressure) and for non-128-aligned chunks (partial tiles), so it is guarded. + # + # The wide config's per-launch LDS overflows the 160 KB limit ONLY when the AMD + # ping-pong / async-copy knobs are on (primus_turbo's set_triton_knobs_gfx950() + # enables them globally in training), which double-buffer the LDS operand tiles + # and ~double this kernel's shared memory (BH64/TK128 R_CHUNK256: fits + # standalone, 347904 B in training). Since the whole bwd now compiles with those + # knobs disabled (see the amd_pingpong_disabled scope below), the wide config + # fits in both the benchmark and training, so it is the default where it applies + # (~6-7% faster on the full flash bwd). Set PRIMUS_DSA_DKV_SAFE=1 to force the + # narrow 32/64 (which fits regardless of the knobs). + use_wide_dkv = ( + num_heads <= 64 and R_CHUNK % 128 == 0 and os.environ.get("PRIMUS_DSA_DKV_SAFE", "0") != "1" + ) + BH_DKV, TK_DKV = (64, 128) if use_wide_dkv else (32, 64) + num_hg_dq = triton.cdiv(num_heads, BH_DQ) + num_hg_dkv = triton.cdiv(num_heads, BH_DKV) + + # In the V4 single-latent form the D_ROPE block of q/kv is a zero pad (RoPE is + # baked in-place over the 512 latent) and its gradient is discarded by the + # adapter (dq[..., :D_V], dkv[..., :D_V]). So all rope MMAs / K_rope loads / + # interm-rope traffic compute a provably-zero result that is thrown away. + # Skip them: bit-identical non-rope outputs, zero rope outputs (== gluon). + HAS_ROPE = False + + dq = torch.empty_like(q) + chunk_dS = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + chunk_P = torch.empty(total_tokens, num_heads, R_CHUNK, dtype=torch.bfloat16, device=q.device) + dkv_acc = torch.zeros(num_kv, d_qk, dtype=torch.float32, device=q.device) + interm = torch.empty(total_tokens, R_CHUNK, d_qk, dtype=torch.bfloat16, device=q.device) + + # pad topk to an R_CHUNK multiple (-1 = invalid). + topk_padded_len = ((topk + R_CHUNK - 1) // R_CHUNK) * R_CHUNK + if topk_padded_len != topk: + pad = torch.full((total_tokens, topk_padded_len - topk), -1, dtype=torch.int32, device=q.device) + topk_padded = torch.cat([topk_indices, pad], dim=1).contiguous() + else: + topk_padded = topk_indices + + all_csr = [ + _build_inverted_topk_slice(topk_padded[:, rs : rs + R_CHUNK], rs, R_CHUNK, num_kv=num_kv) + for rs in range(0, topk, R_CHUNK) + ] + + # The AMD ping-pong / async-copy knobs primus_turbo enables globally are a + # pessimization for the whole triton_v2 backward (measured ~5-7% slower on the + # full bwd for both flash and pro in bench_v4_attention), and they overflow the + # 160 KB LDS limit for the wide dKV tiling. Compile the entire bwd (dQ / + # dKV-intermediate / gather) with them disabled; the knobs are read at compile + # time and are not in Triton's cache key, so this pins the faster non-ping-pong + # schedule for these kernels without touching any kernel compiled elsewhere. + # PRIMUS_DSA_BWD_PINGPONG_OFF=0 keeps the ambient knobs (the wide dKV still + # forces them off below, since it does not fit otherwise). + _bwd_ctx = ( + amd_pingpong_disabled() + if os.environ.get("PRIMUS_DSA_BWD_PINGPONG_OFF", "1") == "1" + else contextlib.nullcontext() + ) + with _bwd_ctx: + for chunk_idx, r_start in enumerate(range(0, topk, R_CHUNK)): + is_first = r_start == 0 + + _bwd_chunk_dq_store_ds[(total_tokens, num_hg_dq)]( + q, + kv, + do, + topk_padded, + lse, + delta, + o, + dq, + chunk_dS, + chunk_P, + q.stride(0), + q.stride(1), + kv.stride(0), + do.stride(0), + do.stride(1), + o.stride(0), + o.stride(1), + dq.stride(0), + dq.stride(1), + topk_padded.stride(0), + chunk_dS.stride(0), + chunk_dS.stride(1), + scale, + num_heads, + r_start, + R_CHUNK=R_CHUNK, + BLOCK_H=BH_DQ, + TILE_K=TK_DQ, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_ROPE=HAS_ROPE, + IS_FIRST_CHUNK=is_first, + num_warps=4, + waves_per_eu=1, + ) + + # The wide dKV kernel MUST compile with ping-pong off to fit LDS, even + # if the outer bwd gate is disabled — force it off here regardless. + _dkv_ctx = amd_pingpong_disabled() if use_wide_dkv else contextlib.nullcontext() + with _dkv_ctx: + _bwd_compute_dkv_intermediate[(total_tokens,)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + num_heads, + R_CHUNK=R_CHUNK, + TILE_K=TK_DKV, + BLOCK_H=BH_DKV, + NUM_HG=num_hg_dkv, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_ROPE=HAS_ROPE, + num_warps=4, + ) + + inv_ptr, inv_data = all_csr[chunk_idx] + _bwd_dkv_gather_acc[(num_kv,)]( + interm, + inv_ptr, + inv_data, + dkv_acc, + interm.stride(1), + dkv_acc.stride(0), + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_ROPE=HAS_ROPE, + num_warps=4, + ) + + d_sink = None + if has_sink: + d_sink = -(torch.exp(attn_sink.unsqueeze(0) - lse) * delta).sum(0) + + dkv = dkv_acc.to(kv.dtype).unsqueeze(1) + return dq, dkv, d_sink diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_fwd_v4_triton.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_fwd_v4_triton.py new file mode 100644 index 000000000..3f6c58b1e --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_v2/dsa_fwd_v4_triton.py @@ -0,0 +1,206 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plain-Triton DeepSeek-V4 sparse-MLA forward (the "triton v2" backend). + +Same sparse-MLA latent representation and public API as the gluon backend +(:func:`sparse_mla_fwd_v4_gluon`) — fused single MQA latent (K = V = the first +``kv_lora_rank`` channels), per-token absolute top-k indices, optional per-head +softmax sink — but written in vanilla Triton so the QK / PV GEMMs lower to MFMA +through ``tl.dot`` (no hand-rolled gluon layouts). One program handles one query +token and a block of heads; it gathers the selected KV latent rows tile-by-tile +and runs an online (flash) softmax with a sink-augmented denominator. + +This contrasts with the in-tree Triton CSA backend (``_triton/v4_csa_*``) which +keeps K and V separate to share kernels with the dense path; here K == V is a +single latent (matching gluon / the V4 paper), which is the whole point of v2. +""" + +import contextlib +import os + +import torch +import triton +import triton.language as tl + +from ._amd_knobs import amd_pingpong_disabled + + +def _get_fwd_configs(): + # Focused around the autotune winners from the wide sweep (TILE_K=16, + # num_stages=3 dominated -> latency-bound; deep pipelining is the lever), plus + # num_stages=4 to probe deeper pipelining. Keeps first-call autotune cheap. + return [ + triton.Config({"BLOCK_H": bh, "TILE_K": tk, "waves_per_eu": wpe}, num_warps=4, num_stages=ns) + for bh in (32, 64) + for tk in (16, 32) + for ns in (2, 3, 4) + for wpe in (0, 1) + ] + + +@triton.autotune(configs=_get_fwd_configs(), key=["num_heads", "TOPK", "D_V", "D_ROPE", "HAS_ROPE"]) +@triton.jit +def _sparse_mla_fwd_tr_kernel( + Q_ptr, # [total_tokens, num_heads, D_QK] bf16 + KV_ptr, # [num_kv, 1, D_QK] bf16 + TopK_ptr, # [total_tokens, TOPK] int32 + Sink_ptr, # [num_heads] fp32 (guarded by HAS_SINK) + O_ptr, # [total_tokens, num_heads, D_V] bf16 + LSE_ptr, # [total_tokens, num_heads] fp32 (sink-inclusive) + stride_q_t, + stride_q_h, + stride_kv_t, + stride_o_t, + stride_o_h, + stride_topk_t, + scale, + num_heads, + TOPK: tl.constexpr, + BLOCK_H: tl.constexpr, + TILE_K: tl.constexpr, + D_V: tl.constexpr, + D_ROPE: tl.constexpr, + HAS_ROPE: tl.constexpr, + HAS_SINK: tl.constexpr, +): + tok = tl.program_id(0) + hg = tl.program_id(1) + + offs_h = hg * BLOCK_H + tl.arange(0, BLOCK_H) + mask_h = offs_h < num_heads + offs_v = tl.arange(0, D_V) + offs_r = tl.arange(0, D_ROPE) + + q_base = tok.to(tl.int64) * stride_q_t + offs_h.to(tl.int64)[:, None] * stride_q_h + q_lora = tl.load(Q_ptr + q_base + offs_v[None, :], mask=mask_h[:, None], other=0.0) + if HAS_ROPE: + q_rope = tl.load(Q_ptr + q_base + (D_V + offs_r)[None, :], mask=mask_h[:, None], other=0.0) + + m_i = tl.full([BLOCK_H], float("-inf"), dtype=tl.float32) + l_i = tl.zeros([BLOCK_H], dtype=tl.float32) + acc = tl.zeros([BLOCK_H, D_V], dtype=tl.float32) + + topk_base = tok.to(tl.int64) * stride_topk_t + for kt in range(0, TOPK, TILE_K): + offs_k = kt + tl.arange(0, TILE_K) + idx = tl.load(TopK_ptr + topk_base + offs_k, mask=offs_k < TOPK, other=-1) + valid = idx >= 0 + safe = tl.where(valid, idx, 0).to(tl.int64) + + kv_base = safe[:, None] * stride_kv_t + k_lora = tl.load(KV_ptr + kv_base + offs_v[None, :], mask=valid[:, None], other=0.0) + + # S = q @ k^T over [lora ++ rope] -> [BLOCK_H, TILE_K] + s = tl.dot(q_lora, tl.trans(k_lora)) + if HAS_ROPE: + k_rope = tl.load(KV_ptr + kv_base + (D_V + offs_r)[None, :], mask=valid[:, None], other=0.0) + s += tl.dot(q_rope, tl.trans(k_rope)) + s = s * scale + s = tl.where(valid[None, :] & mask_h[:, None], s, float("-inf")) + + m_ij = tl.max(s, axis=1) + m_new = tl.maximum(m_i, m_ij) + m_new = tl.where(m_new > float("-inf"), m_new, 0.0) + alpha = tl.exp(m_i - m_new) + p = tl.exp(s - m_new[:, None]) + l_i = alpha * l_i + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(k_lora.dtype), k_lora) + m_i = m_new + + if HAS_SINK: + sink = tl.load(Sink_ptr + offs_h, mask=mask_h, other=float("-inf")) + m_f = tl.maximum(m_i, sink) + af = tl.exp(m_i - m_f) + l_t = l_i * af + tl.exp(sink - m_f) + acc = acc * af[:, None] + acc = acc / l_t[:, None] + lse = m_f + tl.log(l_t) + else: + acc = acc / l_i[:, None] + lse = m_i + tl.log(l_i) + + o_base = tok.to(tl.int64) * stride_o_t + offs_h.to(tl.int64)[:, None] * stride_o_h + tl.store(O_ptr + o_base + offs_v[None, :], acc.to(O_ptr.dtype.element_ty), mask=mask_h[:, None]) + tl.store(LSE_ptr + tok.to(tl.int64) * num_heads + offs_h, lse, mask=mask_h) + + +def sparse_mla_fwd_v4_triton(q, kv, topk_indices, attn_sink=None, kv_lora_rank=512, scale=None): + """DeepSeek-V4 sparse-MLA forward (plain Triton / MFMA). API mirrors the gluon path. + + Args: + q: [total_tokens, num_heads, d_qk] bf16 + kv: [num_kv, 1, d_qk] bf16 (or [num_kv, d_qk]); single MQA latent + topk_indices: [total_tokens, topk] int32 (SWA + sparse, -1 = invalid) + attn_sink: [num_heads] fp32 optional per-head learnable sink + kv_lora_rank: int, default 512 + scale: float, default 1/sqrt(d_qk) + + Returns: + o: [total_tokens, num_heads, kv_lora_rank] (q.dtype) + lse: [total_tokens, num_heads] fp32 (sink-inclusive when attn_sink given) + """ + assert q.is_contiguous() and topk_indices.is_contiguous() + total_tokens, num_heads, d_qk = q.shape + rope_rank = d_qk - kv_lora_rank + topk = topk_indices.shape[1] + if scale is None: + scale = 1.0 / (d_qk**0.5) + if kv.dim() == 2: + kv = kv.unsqueeze(1) + assert kv.is_contiguous() + assert kv.shape[0] >= total_tokens and kv.shape[-1] == d_qk + + has_sink = attn_sink is not None + if has_sink: + assert attn_sink.is_contiguous() and attn_sink.dtype == torch.float32 + assert attn_sink.shape == (num_heads,) + sink_ptr = attn_sink + else: + sink_ptr = torch.empty(1, dtype=torch.float32, device=q.device) + + o = torch.empty(total_tokens, num_heads, kv_lora_rank, dtype=q.dtype, device=q.device) + lse = torch.empty(total_tokens, num_heads, dtype=torch.float32, device=q.device) + + # V4 single-latent form: the D_ROPE block of q/kv is a zero pad (RoPE baked + # in-place over the 512 latent), so the rope QK term is provably zero — skip it. + has_rope = False + + grid = lambda META: (total_tokens, triton.cdiv(num_heads, META["BLOCK_H"])) + # The AMD ping-pong / async-copy knobs primus_turbo enables globally are a + # pessimization for this fwd kernel (~16-29% slower on flash in + # bench_v4_attention). They are read at compile time and are not part of + # Triton's cache key, so autotuning/compiling this kernel with them disabled + # pins the faster (non-ping-pong) schedule for the process without touching + # any other kernel. PRIMUS_DSA_FWD_PINGPONG_OFF=0 keeps the ambient knobs. + _fwd_ctx = ( + amd_pingpong_disabled() + if os.environ.get("PRIMUS_DSA_FWD_PINGPONG_OFF", "1") == "1" + else contextlib.nullcontext() + ) + with _fwd_ctx: + _sparse_mla_fwd_tr_kernel[grid]( + Q_ptr=q, + KV_ptr=kv, + TopK_ptr=topk_indices, + Sink_ptr=sink_ptr, + O_ptr=o, + LSE_ptr=lse, + stride_q_t=q.stride(0), + stride_q_h=q.stride(1), + stride_kv_t=kv.stride(0), + stride_o_t=o.stride(0), + stride_o_h=o.stride(1), + stride_topk_t=topk_indices.stride(0), + scale=scale, + num_heads=num_heads, + TOPK=topk, + D_V=kv_lora_rank, + D_ROPE=rope_rank, + HAS_ROPE=has_rope, + HAS_SINK=has_sink, + ) + return o, lse diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/__init__.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/__init__.py new file mode 100644 index 000000000..5381d3d46 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_turbo_flydsl/__init__.py @@ -0,0 +1,41 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 sparse-MLA kernel-pair via the **Primus-Turbo** public API. + +Unlike the other in-tree fused single-latent backends (``_gluon_v2`` / +``_triton_v2`` / ``_flydsl_v1``), which vendor their kernels inside Primus, this +backend calls straight into the installed **``primus_turbo``** package — the +native-FlyDSL sparse-MLA v2 attention +(``primus_turbo.flydsl.attention.kernels.sparse_mla_v2``) from the Primus-Turbo +``dev/kyle/flydsl_attn_deepseekv4`` line. This is the "turbo API" integration: +Primus owns only the thin V4 adapter binding; the kernels live in Primus-Turbo. + +Public kernel-pair API (identical to the other sparse-MLA backends): + +* ``sparse_mla_fwd_v4_turbo_flydsl(q, kv, topk, attn_sink=None, + kv_lora_rank=512, scale=None) -> (o, lse)`` +* ``sparse_mla_bwd_v4_turbo_flydsl(q, kv, o, do, topk, lse, attn_sink=None, + kv_lora_rank=512, scale=None) -> (dq, dkv, d_sink)`` + +``primus_turbo`` (with the flydsl sparse-MLA attention) and the ``flydsl`` pip +package are required; the import fails with a clear message otherwise (handled by +the lazy loader in :mod:`..` / :func:`load_turbo_attention_backends`). +""" + +from __future__ import annotations + +from primus_turbo.flydsl.attention.kernels.sparse_mla_v2 import ( + sparse_mla_bwd_v4_flydsl, + sparse_mla_fwd_v4_flydsl, +) + +# Turbo-suffixed re-exports so the V4 wrapper / benchmark can bind them without +# clashing with the in-tree ``_flydsl_v1`` (which has identically-named kernels). +sparse_mla_fwd_v4_turbo_flydsl = sparse_mla_fwd_v4_flydsl +sparse_mla_bwd_v4_turbo_flydsl = sparse_mla_bwd_v4_flydsl + +__all__ = ["sparse_mla_fwd_v4_turbo_flydsl", "sparse_mla_bwd_v4_turbo_flydsl"] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_flydsl.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_flydsl.py new file mode 100644 index 000000000..968992336 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_flydsl.py @@ -0,0 +1,38 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 attention via the native FlyDSL sparse-MLA backend ("flydsl_v1"). + +Thin binding of the kernel-agnostic V4 adapters (:mod:`v4_sparse_mla_adapter`) +to the ``sparse_mla_{fwd,bwd}_v4_flydsl`` kernel pair (``_flydsl_v1``): the fused +single-latent (K == V) sparse-MLA path implemented in native FlyDSL MFMA +(``rocdl.mfma_*``) over a per-token top-k gather. The forward is fully native +FlyDSL; the backward uses a native FlyDSL dQ kernel plus the shared Triton dKV +intermediate/scatter-gather. Numerically equivalent to the eager V4 references. + +Depends only on the installed ``flydsl`` pip package (gfx950 / CDNA4); it is +therefore loaded LAZILY (see ``load_flydsl_attention_backends``), never at +package import time. +""" + +from __future__ import annotations + +from primus.backends.megatron.core.transformer.v4_attention_kernels._flydsl_v1 import ( + sparse_mla_bwd_v4_flydsl, + sparse_mla_fwd_v4_flydsl, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import ( + make_attention, + make_csa_from_pool, +) + +v4_csa_attention_flydsl = make_csa_from_pool(sparse_mla_fwd_v4_flydsl, sparse_mla_bwd_v4_flydsl) +v4_attention_flydsl = make_attention(sparse_mla_fwd_v4_flydsl, sparse_mla_bwd_v4_flydsl) + +__all__ = [ + "v4_csa_attention_flydsl", + "v4_attention_flydsl", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon.py new file mode 100644 index 000000000..0f4bfb305 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon.py @@ -0,0 +1,33 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 attention via the Gluon sparse-MLA backend (gfx950). + +Thin binding of the kernel-agnostic V4 adapters (:mod:`v4_sparse_mla_adapter`) +to the gluon ``sparse_mla_{fwd,bwd}_v4_gluon`` kernel pair (``_gluon_dsa``). +See the adapter module for the full V4 <-> sparse-MLA mapping (zero rope pad, +``[local ++ pool]`` kv buffer, ``[SWA window ++ pool]`` topk, grad mapping). +Numerically equivalent to :func:`eager_v4_csa_attention` / :func:`eager_v4_attention`. +""" + +from __future__ import annotations + +from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_dsa import ( + sparse_mla_bwd_v4_gluon, + sparse_mla_fwd_v4_gluon, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import ( + make_attention, + make_csa_from_pool, +) + +v4_csa_attention_gluon = make_csa_from_pool(sparse_mla_fwd_v4_gluon, sparse_mla_bwd_v4_gluon) +v4_attention_gluon = make_attention(sparse_mla_fwd_v4_gluon, sparse_mla_bwd_v4_gluon) + +__all__ = [ + "v4_csa_attention_gluon", + "v4_attention_gluon", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v2.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v2.py new file mode 100644 index 000000000..53c3a4d8b --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v2.py @@ -0,0 +1,38 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 attention via the Gluon sparse-MLA backend ("gluon_v2"). + +``gluon_v2`` forward is the Gluon sparse-MLA kernel (gfx950 / CDNA4 hardware-controlled +layouts + async double-buffered pipeline, rope-skip + exp2 + MFMA K=32); the backward is +currently the plain-Triton chunked-gather kernel (shared with ``triton_v2``) and is being +migrated to Gluon. Thin binding of the kernel-agnostic V4 adapters +(:mod:`v4_sparse_mla_adapter`) to the ``sparse_mla_{fwd,bwd}_v4_gluon_v2`` kernel pair +(``_gluon_v2``). + +The Gluon forward requires a Gluon-capable (recompiled) triton; the kernel raises a clear +build hint otherwise. Numerically equivalent to the eager V4 references (validated in +tests/.../test_v4_gluon_v2_attention.py). +""" + +from __future__ import annotations + +from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v2 import ( + sparse_mla_bwd_v4_gluon_v2, + sparse_mla_fwd_v4_gluon_v2, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import ( + make_attention, + make_csa_from_pool, +) + +v4_csa_attention_gluon_v2 = make_csa_from_pool(sparse_mla_fwd_v4_gluon_v2, sparse_mla_bwd_v4_gluon_v2) +v4_attention_gluon_v2 = make_attention(sparse_mla_fwd_v4_gluon_v2, sparse_mla_bwd_v4_gluon_v2) + +__all__ = [ + "v4_csa_attention_gluon_v2", + "v4_attention_gluon_v2", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v3.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v3.py new file mode 100644 index 000000000..6fe1f022b --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_gluon_v3.py @@ -0,0 +1,26 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 attention via the Gluon sparse-MLA backend ("gluon_v3").""" + +from __future__ import annotations + +from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v3 import ( + sparse_mla_bwd_v4_gluon_v3, + sparse_mla_fwd_v4_gluon_v3, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import ( + make_attention, + make_csa_from_pool, +) + +v4_csa_attention_gluon_v3 = make_csa_from_pool(sparse_mla_fwd_v4_gluon_v3, sparse_mla_bwd_v4_gluon_v3) +v4_attention_gluon_v3 = make_attention(sparse_mla_fwd_v4_gluon_v3, sparse_mla_bwd_v4_gluon_v3) + +__all__ = [ + "v4_csa_attention_gluon_v3", + "v4_attention_gluon_v3", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_triton.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_triton.py new file mode 100644 index 000000000..98f14e186 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_triton.py @@ -0,0 +1,37 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 attention via the plain-Triton sparse-MLA backend ("triton_v2"). + +The ``_v2`` denotes the second *kernel* implementation of the V4 sparse-MLA +attention (NOT a new Triton release). Thin binding of the kernel-agnostic V4 +adapters (:mod:`v4_sparse_mla_adapter`) to the ``sparse_mla_{fwd,bwd}_v4_triton`` +kernel pair (``_triton_v2``) — the fused single-latent (K == V) sparse-MLA path +whose GEMMs lower to MFMA via ``tl.dot``. Unlike the gluon backend (gfx950 / +CDNA4 hardware-controlled layouts), this is vanilla Triton and runs on any +MFMA-capable arch (gfx942 / gfx950). Distinct from the in-tree separate-KV CSA +Triton backend (``v4_csa_attention_v0``). Numerically equivalent to the eager V4 +references. +""" + +from __future__ import annotations + +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_v2 import ( + sparse_mla_bwd_v4_triton, + sparse_mla_fwd_v4_triton, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import ( + make_attention, + make_csa_from_pool, +) + +v4_csa_attention_v2 = make_csa_from_pool(sparse_mla_fwd_v4_triton, sparse_mla_bwd_v4_triton) +v4_attention_v2 = make_attention(sparse_mla_fwd_v4_triton, sparse_mla_bwd_v4_triton) + +__all__ = [ + "v4_csa_attention_v2", + "v4_attention_v2", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_turbo_flydsl.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_turbo_flydsl.py new file mode 100644 index 000000000..9da7c4ea9 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_csa_attention_turbo_flydsl.py @@ -0,0 +1,39 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""V4 attention via the Primus-Turbo native-FlyDSL sparse-MLA backend ("turbo"). + +Thin binding of the kernel-agnostic V4 adapters (:mod:`v4_sparse_mla_adapter`) to +the ``sparse_mla_{fwd,bwd}_v4_turbo_flydsl`` kernel pair (:mod:`_turbo_flydsl`), +which re-exports the installed ``primus_turbo`` flydsl sparse-MLA v2 kernels. Same +fused single-latent (K == V) sparse-MLA-with-sink math as the ``gluon_v2`` / +``triton_v2`` / ``flydsl_v1`` backends, so it is numerically equivalent to the +eager V4 references (validated in +``tests/.../test_v4_turbo_flydsl_attention.py``). + +Requires the installed ``primus_turbo`` (with the flydsl sparse-MLA attention) and +the ``flydsl`` pip package (gfx950 / CDNA4); the import raises a clear hint +otherwise (see :func:`..load_turbo_attention_backends`). +""" + +from __future__ import annotations + +from primus.backends.megatron.core.transformer.v4_attention_kernels._turbo_flydsl import ( + sparse_mla_bwd_v4_turbo_flydsl, + sparse_mla_fwd_v4_turbo_flydsl, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_sparse_mla_adapter import ( + make_attention, + make_csa_from_pool, +) + +v4_csa_attention_turbo = make_csa_from_pool(sparse_mla_fwd_v4_turbo_flydsl, sparse_mla_bwd_v4_turbo_flydsl) +v4_attention_turbo = make_attention(sparse_mla_fwd_v4_turbo_flydsl, sparse_mla_bwd_v4_turbo_flydsl) + +__all__ = [ + "v4_csa_attention_turbo", + "v4_attention_turbo", +] diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py new file mode 100644 index 000000000..f31815887 --- /dev/null +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/v4_sparse_mla_adapter.py @@ -0,0 +1,295 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Kernel-agnostic V4 attention adapters over a sparse-MLA fwd/bwd kernel pair. + +Every "fused single-latent" V4 backend (gluon, triton-v2, flydsl-v2) speaks the +same sparse-MLA contract: + +* ``fwd(q[T,H,Dqk], kv[T,1,Dqk], topk[T,TOPK], attn_sink, kv_lora_rank, scale)`` + -> ``(o[T,H,Dv], lse[T,H])`` +* ``bwd(q, kv, o, do, topk, lse, attn_sink, kv_lora_rank, scale)`` + -> ``(dq, dkv, d_sink)`` + +This module maps Primus's V4 attention representations (per-head q, single MQA +latent K = V with RoPE baked in-place over ``head_dim = 512``, compressed pool, +per-query top-K, joint local-SWA + sparse softmax with sink) onto that contract +and maps gradients back — once — so each backend is just a kernel pair: + +* :func:`make_csa_from_pool(fwd, bwd)` -> CSA (cr=4) wrapper +* :func:`make_attention(fwd, bwd)` -> dense (cr=0) / HCA (cr=128) wrapper + +The fwd/bwd kernels are passed to the autograd Function as non-tensor args, so +the same Function serves all backends (backward returns ``None`` for them). +""" + +from __future__ import annotations + +from typing import Callable, Optional + +import torch + +_ROPE_PAD = 64 # dummy separate-rope block (zeros); the kernels need D_ROPE > 0 + + +def _pad_topk_64(topk: torch.Tensor) -> torch.Tensor: + """Pad the topk width to a multiple of 64 with -1 so a backend whose dKV + tiling is 64-wide (e.g. gluon) stays valid (HCA 128+32=160 -> 192).""" + tk = topk.shape[1] + pad = ((tk + 63) // 64) * 64 - tk + if pad > 0: + topk = torch.cat( + [topk, torch.full((topk.shape[0], pad), -1, device=topk.device, dtype=topk.dtype)], dim=1 + ) + return topk.contiguous() + + +def _build_csa_topk(topk_idxs: torch.Tensor, S: int, P: int, W: int) -> torch.Tensor: + """Flat topk [B*S, W+K] over the per-batch [local ++ pool] buffer. + + ``topk_idxs`` [B, S, K] holds pool indices in [0, P) (or -1). Batch ``b`` + occupies rows ``[b*(S+P) : (b+1)*(S+P))`` (local 0..S-1, pool S..S+P-1). + """ + B, _, K = topk_idxs.shape + device = topk_idxs.device + base = (torch.arange(B, device=device) * (S + P)).view(B, 1, 1) + + win_pos = torch.arange(S, device=device).view(S, 1) - W + 1 + torch.arange(W, device=device).view(1, W) + win_valid = win_pos >= 0 + win_idx = base + win_pos.view(1, S, W) + win_idx = torch.where(win_valid.view(1, S, W), win_idx, torch.full_like(win_idx, -1)) + + pool_valid = topk_idxs >= 0 + pool_idx = torch.where(pool_valid, base + S + topk_idxs, torch.full_like(topk_idxs, -1)) + + return torch.cat([win_idx, pool_idx], dim=2).reshape(B * S, W + K).to(torch.int32).contiguous() + + +class _V4SparseMLACSAFn(torch.autograd.Function): + """Autograd wrapper: sparse-MLA FWD/BWD for the V4 CSA (cr=4) layer.""" + + @staticmethod + def forward( # type: ignore[override] + ctx, + q_bh: torch.Tensor, # [B, H, S, D] + k_local_bh: torch.Tensor, # [B, H, S, D] (single MQA latent, head-broadcast) + v_local_bh: torch.Tensor, # [B, H, S, D] (== k_local in V4) + pool: torch.Tensor, # [B, P, D] + topk_idxs: torch.Tensor, # [B, S, K] pool indices, -1 = invalid + sink: Optional[torch.Tensor], # [H] fp32 or None + swa_window: int, + scale: float, + fwd_fn: Callable, + bwd_fn: Callable, + ) -> torch.Tensor: + B, H, S, D = q_bh.shape + P = pool.shape[1] + W = int(swa_window) + assert q_bh.dtype == torch.bfloat16, "sparse-MLA adapter requires bf16" + assert W > 0, "sparse-MLA adapter requires swa_window > 0" + + latent = k_local_bh[:, 0, :, :] # [B, S, D] + + z_q = torch.zeros(B * S, H, _ROPE_PAD, device=q_bh.device, dtype=q_bh.dtype) + q_g = torch.cat([q_bh.permute(0, 2, 1, 3).reshape(B * S, H, D), z_q], dim=-1).contiguous() + + kv512 = torch.cat([latent, pool], dim=1).reshape(B * (S + P), 1, D) + z_kv = torch.zeros(B * (S + P), 1, _ROPE_PAD, device=q_bh.device, dtype=q_bh.dtype) + kv_g = torch.cat([kv512, z_kv], dim=-1).contiguous() + + topk_g = _pad_topk_64(_build_csa_topk(topk_idxs, S, P, W)) + + sink_arg = sink.float().contiguous() if sink is not None else None + o_g, lse = fwd_fn(q_g, kv_g, topk_g, attn_sink=sink_arg, kv_lora_rank=D, scale=float(scale)) + + ctx.save_for_backward(q_g, kv_g, o_g, lse, topk_g, sink_arg if sink is not None else q_g.new_empty(0)) + ctx.shapes = (B, H, S, D, P, W) + ctx.scale = float(scale) + ctx.sink_was_none = sink is None + ctx.bwd_fn = bwd_fn + return o_g.reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous() + + @staticmethod + def backward(ctx, grad_o_bh: torch.Tensor): # type: ignore[override] + q_g, kv_g, o_g, lse, topk_g, sink_saved = ctx.saved_tensors + B, H, S, D, P, W = ctx.shapes + sink_arg = None if ctx.sink_was_none else sink_saved + + grad_o_g = grad_o_bh.permute(0, 2, 1, 3).reshape(B * S, H, D).contiguous() + dq_g, dkv_g, dsink = ctx.bwd_fn( + q_g, kv_g, o_g, grad_o_g, topk_g, lse, attn_sink=sink_arg, kv_lora_rank=D, scale=ctx.scale + ) + + dq_bh = dq_g[:, :, :D].reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous() + dkv512 = dkv_g[:, 0, :D].reshape(B, S + P, D) + dlatent = dkv512[:, :S, :] + dpool = dkv512[:, S:, :].contiguous() + + dk_local = torch.zeros(B, H, S, D, device=dq_bh.device, dtype=dq_bh.dtype) + dk_local[:, 0, :, :] = dlatent.to(dq_bh.dtype) + # V4 is single-latent (K = V = kv): the kernel returns one combined + # ``dkv`` which we route entirely through ``dk_local``. The V branch + # gradient is structurally zero, so we return ``None`` for it instead + # of allocating (and zeroing) a full [B, H, S, D] tensor — this removes + # the largest ``Memset (Device)`` bucket in the trace (~268 MB / call). + # ``k_local_bh`` and ``v_local_bh`` are two ``kv.expand`` views of the + # same latent, so autograd accumulates ``dk_local + 0`` into ``kv`` — + # identical to before. + dv_local = None + + dsink_out = None + if not ctx.sink_was_none and dsink is not None: + dsink_out = dsink.to(sink_saved.dtype) + + # forward args: (q, k_local, v_local, pool, topk_idxs, sink, swa_window, scale, fwd_fn, bwd_fn) + return dq_bh, dk_local, dv_local, dpool.to(dq_bh.dtype), None, dsink_out, None, None, None, None + + +class _V4SparseMLAAttnFn(torch.autograd.Function): + """Sparse-MLA FWD/BWD for the V4 dense (cr=0) and HCA (cr=128) layers.""" + + @staticmethod + def forward( # type: ignore[override] + ctx, + q_bh: torch.Tensor, # [B, H, S, D] + k_bh: torch.Tensor, # [B, H, Skv, D] (Skv = S for cr=0; S+P for HCA) + v_bh: torch.Tensor, # [B, H, Skv, D] (== k_bh in V4) + sink: Optional[torch.Tensor], + swa_window: int, + additive_mask: Optional[torch.Tensor], # [S, P] pool-only mask (HCA) or None + scale: float, + hca_local_seqlen: int, + fwd_fn: Callable, + bwd_fn: Callable, + ) -> torch.Tensor: + B, H, S, D = q_bh.shape + Skv = k_bh.shape[2] + W = int(swa_window) + assert q_bh.dtype == torch.bfloat16, "sparse-MLA adapter requires bf16" + assert W > 0, "sparse-MLA adapter requires swa_window > 0" + + device = q_bh.device + base = (torch.arange(B, device=device) * Skv).view(B, 1, 1) + win_pos = ( + torch.arange(S, device=device).view(S, 1) - W + 1 + torch.arange(W, device=device).view(1, W) + ) + win_valid = win_pos >= 0 + win_idx = base + win_pos.view(1, S, W) + win_idx = torch.where(win_valid.view(1, S, W), win_idx, torch.full_like(win_idx, -1)) + + if hca_local_seqlen > 0 and additive_mask is not None: + P = Skv - int(hca_local_seqlen) + vis = (additive_mask == 0).view(1, S, P) + ps = torch.arange(P, device=device).view(1, 1, P) + pool_idx = torch.where( + vis, base + hca_local_seqlen + ps, torch.full((B, S, P), -1, device=device) + ) + topk = torch.cat([win_idx, pool_idx], dim=2) + else: + topk = win_idx + topk_g = _pad_topk_64(topk.reshape(B * S, -1).to(torch.int32)) + + z_q = torch.zeros(B * S, H, _ROPE_PAD, device=device, dtype=q_bh.dtype) + q_g = torch.cat([q_bh.permute(0, 2, 1, 3).reshape(B * S, H, D), z_q], dim=-1).contiguous() + kv512 = k_bh[:, 0, :, :].reshape(B * Skv, 1, D) + z_kv = torch.zeros(B * Skv, 1, _ROPE_PAD, device=device, dtype=q_bh.dtype) + kv_g = torch.cat([kv512, z_kv], dim=-1).contiguous() + + sink_arg = sink.float().contiguous() if sink is not None else None + o_g, lse = fwd_fn(q_g, kv_g, topk_g, attn_sink=sink_arg, kv_lora_rank=D, scale=float(scale)) + + ctx.save_for_backward(q_g, kv_g, o_g, lse, topk_g, sink_arg if sink is not None else q_g.new_empty(0)) + ctx.shapes = (B, H, S, D, Skv) + ctx.scale = float(scale) + ctx.sink_was_none = sink is None + ctx.bwd_fn = bwd_fn + return o_g.reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous() + + @staticmethod + def backward(ctx, grad_o_bh: torch.Tensor): # type: ignore[override] + q_g, kv_g, o_g, lse, topk_g, sink_saved = ctx.saved_tensors + B, H, S, D, Skv = ctx.shapes + sink_arg = None if ctx.sink_was_none else sink_saved + + grad_o_g = grad_o_bh.permute(0, 2, 1, 3).reshape(B * S, H, D).contiguous() + dq_g, dkv_g, dsink = ctx.bwd_fn( + q_g, kv_g, o_g, grad_o_g, topk_g, lse, attn_sink=sink_arg, kv_lora_rank=D, scale=ctx.scale + ) + + dq_bh = dq_g[:, :, :D].reshape(B, S, H, D).permute(0, 2, 1, 3).contiguous() + dkv = dkv_g[:, 0, :D].reshape(B, Skv, D) + dk_bh = torch.zeros(B, H, Skv, D, device=dq_bh.device, dtype=dq_bh.dtype) + dk_bh[:, 0, :, :] = dkv.to(dq_bh.dtype) + # Single-latent (K = V): route the combined ``dkv`` through ``dk_bh``; + # the V branch gradient is structurally zero, so return ``None`` and + # skip the big [B, H, Skv, D] memset (see the CSA branch note). + dv_bh = None + + dsink_out = None + if not ctx.sink_was_none and dsink is not None: + dsink_out = dsink.to(sink_saved.dtype) + + # forward args: (q, k, v, sink, swa_window, additive_mask, scale, hca_local_seqlen, fwd_fn, bwd_fn) + return dq_bh, dk_bh, dv_bh, dsink_out, None, None, None, None, None, None + + +def make_csa_from_pool(fwd_fn: Callable, bwd_fn: Callable) -> Callable: + """Build a ``v4_csa_attention_v1``-style wrapper for a kernel pair.""" + + def _csa_from_pool( + q_bh, + k_local_bh, + v_local_bh, + pool, + *, + topk_idxs, + sink, + swa_window, + attn_dropout, + training, + scale, + ): + if attn_dropout > 0.0 and training: + raise NotImplementedError( + "sparse-MLA CSA adapter does not implement in-kernel attention dropout " + f"(V4 trains with attn_dropout=0). Got attn_dropout={attn_dropout}, training={training}." + ) + return _V4SparseMLACSAFn.apply( + q_bh, k_local_bh, v_local_bh, pool, topk_idxs, sink, int(swa_window), float(scale), fwd_fn, bwd_fn + ) + + return _csa_from_pool + + +def make_attention(fwd_fn: Callable, bwd_fn: Callable) -> Callable: + """Build a dense (cr=0) / HCA (cr=128) attention wrapper for a kernel pair.""" + + def _attention( + q, + k, + v, + *, + sink, + swa_window, + additive_mask, + attn_dropout, + training, + scale, + hca_local_seqlen=0, + ): + if attn_dropout > 0.0 and training: + raise NotImplementedError( + "sparse-MLA attention adapter does not implement in-kernel attention dropout " + f"(V4 trains with attn_dropout=0). Got attn_dropout={attn_dropout}, training={training}." + ) + return _V4SparseMLAAttnFn.apply( + q, k, v, sink, int(swa_window), additive_mask, float(scale), int(hca_local_seqlen), fwd_fn, bwd_fn + ) + + return _attention + + +__all__ = ["make_csa_from_pool", "make_attention"] diff --git a/primus/backends/megatron/megatron_pretrain_trainer.py b/primus/backends/megatron/megatron_pretrain_trainer.py index b7daff509..162aad4ef 100644 --- a/primus/backends/megatron/megatron_pretrain_trainer.py +++ b/primus/backends/megatron/megatron_pretrain_trainer.py @@ -136,10 +136,14 @@ def train(self): from primus.core.utils.import_utils import get_model_provider - # Determine model type (gpt, mamba, or diffusion) from backend_args + # Determine model type (gpt / mamba / deepseek_v4 / diffusion) from backend_args model_type = getattr(self.backend_args, "model_type", "gpt") log_rank_0(f"-detected model_type: {model_type}") + # Import the appropriate training components based on model_type. + # DeepSeek-V4 is causal-LM with the same data shape as GPT, so we + # reuse pretrain_gpt's forward_step + dataset provider; only the + # model_provider itself is V4-specific. if model_type == "mamba": from pretrain_mamba import ( # type: ignore forward_step, @@ -152,6 +156,13 @@ def train(self): # only TP rank 0 enters dataset construction while the core dataset builder still issues # distributed barriers, which deadlocks for TP>1. train_valid_test_datasets_provider.is_distributed = True + elif model_type == "deepseek_v4": + from pretrain_gpt import ( # type: ignore + forward_step, + train_valid_test_datasets_provider, + ) + + log_rank_0("Using DeepSeek-V4 model provider; reusing pretrain_gpt forward_step + datasets") else: # Use overridable methods so subclasses (e.g. diffusion/Flux) can plug in their own # forward_step / dataset_provider. Defaults pull from pretrain_gpt. diff --git a/primus/backends/megatron/patches/deepseek_v4_flops_patches.py b/primus/backends/megatron/patches/deepseek_v4_flops_patches.py new file mode 100644 index 000000000..247be14a8 --- /dev/null +++ b/primus/backends/megatron/patches/deepseek_v4_flops_patches.py @@ -0,0 +1,910 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 FLOPs reporting patch. + +Plan-3 P20. Megatron's +:func:`megatron.training.training.num_floating_point_operations` is shaped +for GPT / MLA / Mamba and gets V4 wrong on every axis that matters: + +* QKV projections — V4 has a Q LoRA path + (``hidden -> q_lora_rank -> n_heads * head_dim``) and a single-latent + ``linear_kv`` (``hidden -> head_dim`` shared as both K and V); the + upstream MHA/GQA branch counts a flat ``hidden * (q + k + v)`` projection + instead. +* Output projection — V4 uses grouped low-rank ``linear_o_a`` / + ``linear_o_b`` (``(n*d/o_groups) -> o_groups*o_lora -> hidden``); the + upstream branch counts a flat ``(n*d) -> hidden`` proj. +* Attention scores at the wrong sequence length — V4's mHC residual + packs ``hc_mult`` parallel streams into the layer-internal sequence axis + (``[B, S*K, D]``). Per-layer GEMMs run at ``S_eff = S * hc_mult``, but + upstream uses ``args.seq_length``. +* Compressor + Indexer side paths — CSA (``compress_ratio==4``) and HCA + (``compress_ratio==128``) layers add a Compressor (``wkv`` + ``wgate``); + CSA additionally runs an Indexer (``w_dq`` + ``w_iuq`` + ``w_w`` + + mini-Compressor + scoring einsum). Upstream knows about neither. +* Hash routing — V4's first ``num_hash_layers`` MoE layers use a + parameter-free hash router; upstream charges them the topk router GEMM + (``hidden * num_experts``) anyway. +* MTP — V4's MTP block runs a full inner V4 transformer layer per depth + (attention + MoE FFN) on top of the ``eh_proj``; upstream counts only + three norms and the single ``2H -> H`` projection. + +The mismatch makes per-iter MFU comparisons across PP/EP/VPP configs +meaningless because the denominator is wrong by a configuration-dependent +factor. + +This module installs a single ``before_train`` patch that monkey-patches +``training_module.num_floating_point_operations`` with a wrapper. The +wrapper: + +1. Falls through to the upstream function byte-for-byte for + ``args.model_type != "deepseek_v4"`` so dense GPT / MLA / Mamba runs + are unchanged. +2. For V4 runs, evaluates :func:`compute_v4_flops` — a closed form + derived from the V4 forward pass (see + ``deepseek-v4/develop/plan-3/02-phase-details.md#p20--v4-aware-tflops-reporting`` + for the per-component derivation) — and returns its total. +3. On first invocation, logs the per-component FLOPs breakdown at rank 0 + so the formula can be sanity-checked against a hand calculation. + +Convention follows upstream Megatron: pure-FMAC counts internally, then a +single ``forward_backward_factor (3) * fma_factor (2) = 6`` multiplier at +the end. + +Plan-6 P33 closes two known gaps in the plan-3 P20 closed form: + +* SWA visible-pair pruning — :func:`_attn_scores_fmac_per_layer` now + counts only causal-visible ``(q, k)`` pairs surviving the per-layer + SWA + pool + sparse top-K masks (via :func:`_visible_pairs`). The + legacy ``B * n * d * S_eff^2`` upper bound over-counted the local + branch by ``S_eff / swa_window`` once plan-5 P30+ made kernels honor + per-row SWA pruning (~128x at the V4-Flash proxy shape). +* HyperConnection matmul accounting — the new ``hc`` row of + :class:`_V4FlopsBreakdown` counts ``HyperMixer.fn`` (``K*D -> + (2+K)*K`` per token, twice per layer) and ``HyperHead.fn`` (``K*D + -> K`` per token, once at the trunk end and per MTP depth). These + matmuls were not in the plan-3 P20 form. + +See ``deepseek-v4/develop/plan-6/02-phase-details.md#phase-33`` for +the design and ``develop/progress/p33/p33-summary.md`` for the +per-component delta vs the legacy formula. +""" + +from __future__ import annotations + +import ast +from dataclasses import dataclass +from typing import Any, List, Optional, Sequence, Tuple + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + +# --------------------------------------------------------------------------- +# Shared constants (mirror Megatron's expansion factors) +# --------------------------------------------------------------------------- + +_FORWARD_BACKWARD_FACTOR: int = 3 # 1 forward + 2 backward. +_FMA_FACTOR: int = 2 # multiply + add per matmul element. +_SWIGLU_FFN_EXPANSION_FACTOR: int = 3 # gate + up + down (all hidden*ffn). + + +# --------------------------------------------------------------------------- +# compress_ratios parsing — V4 yamls store the schedule as a JSON-like string. +# --------------------------------------------------------------------------- + + +def _parse_compress_ratios(raw: Any) -> Optional[List[int]]: + """Parse ``compress_ratios`` into ``list[int]`` or return ``None``. + + Accepts ``None`` (fully dense), a string like ``"[0, 0, 4, 128]"`` (the + YAML form), or an existing list / tuple of ints. Mirrors the runtime + helper in + :func:`primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block._parse_int_sequence` + intentionally inline so this patch has no V4-module import dependency + (the patch loads at ``before_train``, before model build). + """ + if raw is None: + return None + if isinstance(raw, (list, tuple)): + return [int(x) for x in raw] + if isinstance(raw, str): + text = raw.strip() + if not text: + return None + # Tolerant of the YAML form `[0,0,4,128]` and Python repr `[0, 0, 4, 128]`. + return [int(x) for x in ast.literal_eval(text)] + raise TypeError(f"Unsupported compress_ratios type: {type(raw).__name__}") + + +def _normalize_layer_ratios( + raw: Any, + *, + num_layers: int, + mtp_num_layers: int, +) -> Tuple[List[int], List[int]]: + """Return ``(decoder_ratios, mtp_ratios)`` each padded to its expected length. + + Mirrors the V4 block's normalization (``deepseek_v4_block._normalize_compress_ratios``) + but also returns the trailing MTP slice when the YAML provides + ``num_layers + mtp_num_layers`` entries (the canonical DeepSeek layout). + Default-fills any missing slot with ``0`` (dense). + """ + parsed = _parse_compress_ratios(raw) + if parsed is None: + return [0] * num_layers, [0] * mtp_num_layers + + if len(parsed) == num_layers + mtp_num_layers: + return parsed[:num_layers], parsed[num_layers:] + if len(parsed) == num_layers: + return parsed, [0] * mtp_num_layers + if len(parsed) > num_layers: + return parsed[:num_layers], [0] * mtp_num_layers + pad = parsed[-1] if parsed else 0 + return parsed + [pad] * (num_layers - len(parsed)), [0] * mtp_num_layers + + +# --------------------------------------------------------------------------- +# Per-component closed-form helpers +# +# All functions return *FMAC* counts (multiplies only). The caller multiplies +# by ``_FMA_FACTOR * _FORWARD_BACKWARD_FACTOR = 6`` exactly once at the end. +# --------------------------------------------------------------------------- + + +def _attn_qkv_o_fmac_per_layer( + *, + batch_size: int, + seq_len_eff: int, + hidden_size: int, + num_heads: int, + head_dim: int, + q_lora_rank: int, + o_lora_rank: int, + o_groups: int, +) -> int: + """FMAC for V4 attention's projections (Q LoRA + single-latent KV + O). + + Independent of ``compress_ratio``: all V4 layer types share the same + projection structure. The score / softmax cost is counted separately + by :func:`_attn_scores_fmac_per_layer`. + """ + n_d = num_heads * head_dim + + qkv = ( + # linear_q_down_proj: hidden -> q_lora_rank + hidden_size * q_lora_rank + # linear_q_up_proj: q_lora_rank -> n_heads * head_dim + + q_lora_rank * n_d + # linear_kv (single latent): hidden -> head_dim + + hidden_size * head_dim + ) + if o_lora_rank > 0: + # linear_o_a: (n*d/o_groups) -> o_groups*o_lora ⇒ FMAC = (n*d)*o_lora + # linear_o_b: o_groups*o_lora -> hidden ⇒ FMAC = (o_groups*o_lora)*hidden + o_proj = n_d * o_lora_rank + (o_groups * o_lora_rank) * hidden_size + else: + # Flat fallback: (n*d) -> hidden + o_proj = n_d * hidden_size + + return batch_size * seq_len_eff * (qkv + o_proj) + + +def _local_visible_pairs(swa_window: int, seq_len_eff: int) -> int: + """Causal-visible ``(q, k)`` pair count for the local SWA branch. + + Each query at position ``q`` attends to keys in + ``[max(0, q - swa_window + 1), q]``. Summed over + ``q in [0, seq_len_eff)`` the count is: + + * ``swa_window * seq_len_eff - swa_window * (swa_window - 1) // 2`` for + ``0 < swa_window < seq_len_eff`` (queries below ``swa_window - 1`` + see only ``q + 1`` keys, queries above saturate at ``swa_window``). + * ``seq_len_eff * (seq_len_eff + 1) // 2`` for the full-causal + fallback (``swa_window == 0`` or ``>= seq_len_eff``). + """ + s = int(seq_len_eff) + w = int(swa_window) + if w <= 0 or w >= s: + return s * (s + 1) // 2 + return w * s - w * (w - 1) // 2 + + +def _pool_visible_pairs(compress_ratio: int, seq_len_eff: int) -> int: + """Causal-visible ``(q, p)`` pair count for the HCA compressed pool. + + Pool slot ``p`` covers source positions ``[p*c, (p+1)*c)`` with + ``c == compress_ratio``; slot ``p`` is visible to query ``q`` iff + ``(p+1) * c - 1 <= q``. Summed over queries, the count equals + ``sum_{m=1..seq_len_eff} floor(m / c)`` with closed form + ``c * n * (n - 1) // 2 + n * (T - c*n + 1)`` where ``n = T // c`` and + ``T = seq_len_eff``. + """ + c = int(compress_ratio) + t = int(seq_len_eff) + if c <= 0 or t <= 0: + return 0 + n = t // c + if n == 0: + return 0 + return c * n * (n - 1) // 2 + n * (t - c * n + 1) + + +def _visible_pairs( + *, + swa_window: int, + compress_ratio: int, + index_topk: int, + seq_len_eff: int, +) -> int: + """Total causal-visible ``(q, k)`` pair count for one V4 attention layer. + + Plan-6 P33 closed form — see + ``deepseek-v4/develop/perf/attention_perf.md`` "Test Shape And Counting" + for the per-branch derivation. + + * ``cr == 0`` (dense + SWA): local SWA pairs only. + * ``cr == 128`` (HCA): local SWA pairs + causal pool pairs. + * ``cr == 4`` (CSA): local SWA pairs + sparse top-K pairs + (``min(index_topk, pool) * seq_len_eff``); top-K is treated as + fully visible because the indexer assigns each query a per-row + causal-respecting pool subset. + """ + local = _local_visible_pairs(swa_window, seq_len_eff) + cr = int(compress_ratio) + if cr == 0: + return local + + pool = max(1, int(seq_len_eff) // cr) + if cr == 128: + return local + _pool_visible_pairs(cr, seq_len_eff) + if cr == 4: + sparse_keys = min(int(index_topk) if index_topk else pool, pool) + return local + sparse_keys * int(seq_len_eff) + # Forward-compatible: any other ratio is treated as full pool cross-attn. + return local + pool * int(seq_len_eff) + + +def _attn_scores_fmac_per_layer( + *, + batch_size: int, + seq_len_eff: int, + num_heads: int, + head_dim: int, + compress_ratio: int, + index_topk: int, + swa_window: int, +) -> int: + """FMAC for the attention score matmuls. + + Plan-6 P33 rewrite: counts only the causal-visible ``(query, key)`` + pairs surviving the per-layer mask (SWA + pool + sparse top-K) — the + plan-3 P20 ``S_eff^2`` upper bound over-counted dense / HCA local + attention by ``S_eff / swa_window`` (16x at ``swa=128, S_eff=4096``) + once plan-5 P30 SWA K-loop pruning made the per-layer kernels track + visible pairs only. + + FMAC = ``2 * num_heads * head_dim * visible_pairs``: one ``n*d`` for + the ``QK^T`` matmul + one ``n*d`` for the ``PV`` matmul, summed over + visible ``(q, k)`` pairs. The forward + backward * FMA expansion is + applied once at the end in :class:`_V4FlopsBreakdown`. + """ + pairs = _visible_pairs( + swa_window=swa_window, + compress_ratio=compress_ratio, + index_topk=index_topk, + seq_len_eff=seq_len_eff, + ) + return 2 * batch_size * num_heads * head_dim * pairs + + +def _compressor_fmac_per_layer( + *, + batch_size: int, + seq_len_eff: int, + hidden_size: int, + head_dim: int, + compress_ratio: int, +) -> int: + """FMAC for the V4 :class:`Compressor` (HCA / CSA only). + + Compressor projects ``hidden -> coff*head_dim`` for both ``wkv`` and + ``wgate``; ``coff = 2`` in overlap mode (CSA, ratio==4) and ``coff = 1`` + in non-overlap mode (HCA, ratio==128). Inputs are at the full + pre-pool seq length, so cost is paid at ``S_eff``. + """ + if compress_ratio == 0: + return 0 + coff = 2 if compress_ratio == 4 else 1 + # wkv + wgate, each hidden -> coff * head_dim + return 2 * batch_size * seq_len_eff * hidden_size * (coff * head_dim) + + +def _indexer_fmac_per_layer( + *, + batch_size: int, + seq_len_eff: int, + hidden_size: int, + compress_ratio: int, + index_head_dim: int, + index_n_heads: int, +) -> int: + """FMAC for the V4 :class:`Indexer` (CSA only).""" + if compress_ratio != 4: + return 0 + + pool = max(1, seq_len_eff // compress_ratio) + dq_rank = index_head_dim # Indexer.__init__ default: dq_rank = index_head_dim + inh_ihd = index_n_heads * index_head_dim + + proj = ( + # w_dq: hidden -> dq_rank + hidden_size * dq_rank + # w_iuq: dq_rank -> inh * ihd + + dq_rank * inh_ihd + # w_w: hidden -> inh + + hidden_size * index_n_heads + ) + # mini-Compressor inside the Indexer: head_dim=index_head_dim, ratio=4 (coff=2), + # so wkv + wgate each cost hidden * (2 * index_head_dim). + mini_compressor = 2 * hidden_size * (2 * index_head_dim) + proj += mini_compressor + + proj_fmac = batch_size * seq_len_eff * proj + + # Scoring einsum: (B, S_eff, inh, ihd) · (B, P, ihd) → (B, S_eff, inh, P) + scoring_fmac = batch_size * seq_len_eff * index_n_heads * pool * index_head_dim + + return proj_fmac + scoring_fmac + + +def _moe_fmac_per_layer( + *, + batch_size: int, + seq_len_eff: int, + hidden_size: int, + moe_ffn_hidden_size: int, + moe_router_topk: int, + num_experts: int, + is_hash_layer: bool, + shared_expert_ffn_hidden_size: int, +) -> int: + """FMAC for the V4 MoE FFN (router + routed experts + shared expert). + + Hash-routed layers skip the topk router GEMM (they look up bucket + indices from the raw input ids — a parameter-free op). Non-hash + layers pay ``hidden * num_experts`` per token for the router. + SwiGLU's ``ffn_expansion_factor=3`` collapses the (gate + up + down) + matmul triple into a single multiplier consistent with upstream. + """ + router = 0 if is_hash_layer else hidden_size * num_experts + routed = moe_router_topk * _SWIGLU_FFN_EXPANSION_FACTOR * hidden_size * moe_ffn_hidden_size + shared = ( + _SWIGLU_FFN_EXPANSION_FACTOR * hidden_size * shared_expert_ffn_hidden_size + if shared_expert_ffn_hidden_size > 0 + else 0 + ) + return batch_size * seq_len_eff * (router + routed + shared) + + +def _hc_mixer_fmac_per_layer( + *, + batch_size: int, + seq_len: int, + hidden_size: int, + hc_mult: int, +) -> int: + """FMAC for the two ``HyperMixer.fn`` matmuls in one V4 hybrid layer. + + Each :class:`HyperMixer` projects the un-packed K streams + ``[B, S, K*D] -> [B, S, (2+K)*K]``; the V4 hybrid layer runs two + mixers per layer (one before / one inside the attention sub-block, + and one for the FFN sub-block — see + ``primus.backends.megatron.core.transformer.hyper_connection``). + + The leading axis is ``B * S`` (not ``B * S * hc_mult``) because the + K stream lifting that packs streams into the sequence axis happens + *after* the mixer. The (small) ``HyperMixer.expand`` matmul + ``comb @ x`` is left out by design (see plan-6 P33 spec). + """ + if hc_mult <= 0: + return 0 + n_d = hc_mult * hidden_size + return 2 * batch_size * seq_len * n_d * ((2 + hc_mult) * hc_mult) + + +def _hc_head_fmac( + *, + batch_size: int, + seq_len: int, + hidden_size: int, + hc_mult: int, + mtp_num_layers: int, +) -> int: + """FMAC for the ``HyperHead.fn`` matmuls (trunk end + per MTP depth). + + Each :class:`HyperHead` projects ``[B, S, K*D] -> [B, S, K]`` to + produce the sigmoid weights for the final K-stream collapse. V4 has + one head at the trunk end and one head per MTP depth (each with its + own ``num_nextn_predict_layers`` copies). + """ + if hc_mult <= 0: + return 0 + n_d = hc_mult * hidden_size + return (1 + mtp_num_layers) * batch_size * seq_len * n_d * hc_mult + + +def _mtp_eh_proj_fmac( + *, + batch_size: int, + seq_len: int, + hidden_size: int, + mtp_num_layers: int, +) -> int: + """FMAC for the per-MTP-depth ``eh_proj`` (``2H -> H``). + + Runs at the original (un-packed) seq length because MTP runs **before** + the V4 transformer layer's stream-lift, on the embedding output. Two + norms per depth are negligible. + """ + if mtp_num_layers <= 0: + return 0 + return mtp_num_layers * batch_size * seq_len * (2 * hidden_size) * hidden_size + + +def _logits_fmac( + *, + batch_size: int, + seq_len: int, + hidden_size: int, + padded_vocab_size: int, + mtp_num_layers: int, +) -> int: + """FMAC for the LM head (one per main path + one per MTP depth).""" + return (mtp_num_layers + 1) * batch_size * seq_len * hidden_size * padded_vocab_size + + +# --------------------------------------------------------------------------- +# Public closed form +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _V4FlopsBreakdown: + """Per-component FMAC totals (multiply-only, pre fwd+bwd / FMA scaling). + + Plan-6 P33 adds the ``hc`` field for HyperConnection matmuls + (``HyperMixer.fn`` + ``HyperHead.fn``). Field added at the end so + existing positional callers and the breakdown log ordering stay + stable; ``hc`` defaults to 0 so historical pickles / external + consumers that built the dataclass with 7 positional args keep + working. + """ + + attn_qkv_o: int + attn_scores: int + compressor: int + indexer: int + moe: int + mtp: int + logits: int + hc: int = 0 + + def total_fmac(self) -> int: + return ( + self.attn_qkv_o + + self.attn_scores + + self.compressor + + self.indexer + + self.moe + + self.mtp + + self.logits + + self.hc + ) + + def to_total_flops(self) -> int: + """Apply the fwd+bwd (3) × FMA (2) = 6 expansion.""" + return _FORWARD_BACKWARD_FACTOR * _FMA_FACTOR * self.total_fmac() + + +def compute_v4_flops(args: Any, batch_size: int) -> Tuple[int, _V4FlopsBreakdown]: + """Closed-form V4 FLOPs for one global batch. + + Returns ``(total_flops, breakdown)`` where ``total_flops`` is the + Megatron-convention number suitable for direct substitution into + upstream's reporting and ``breakdown`` is a per-component report + (FMAC pre-expansion) for sanity logging. + """ + seq_len = int(args.seq_length) + hc_mult = int(getattr(args, "hc_mult", 1) or 1) + seq_len_eff = seq_len * hc_mult + + hidden_size = int(args.hidden_size) + num_heads = int(args.num_attention_heads) + head_dim = int(args.kv_channels) + q_lora_rank = int(getattr(args, "q_lora_rank", 0) or 0) + o_lora_rank = int(getattr(args, "o_lora_rank", 0) or 0) + o_groups = int(getattr(args, "o_groups", 1) or 1) + + num_layers = int(args.num_layers) + mtp_num_layers = int(getattr(args, "mtp_num_layers", 0) or 0) + + decoder_ratios, mtp_ratios = _normalize_layer_ratios( + getattr(args, "compress_ratios", None), + num_layers=num_layers, + mtp_num_layers=mtp_num_layers, + ) + + moe_ffn_hidden_size = int(getattr(args, "moe_ffn_hidden_size", None) or args.ffn_hidden_size) + moe_router_topk = int(getattr(args, "moe_router_topk", 1) or 1) + num_experts = int(getattr(args, "num_experts", 1) or 1) + shared_expert_ffn_hidden_size = int(getattr(args, "moe_shared_expert_intermediate_size", 0) or 0) + num_hash_layers = int(getattr(args, "num_hash_layers", 0) or 0) + + index_topk = int(getattr(args, "index_topk", 0) or 0) + index_head_dim = int(getattr(args, "index_head_dim", 0) or 0) + index_n_heads = int(getattr(args, "index_n_heads", 0) or 0) + + swa_window = int(getattr(args, "attn_sliding_window", 0) or 0) + + padded_vocab_size = int(getattr(args, "padded_vocab_size", None) or args.vocab_size) + + # ---- decoder layers ---- + attn_qkv_o = 0 + attn_scores = 0 + compressor = 0 + indexer = 0 + moe = 0 + hc = 0 + + for layer_idx in range(num_layers): + ratio = int(decoder_ratios[layer_idx]) + + attn_qkv_o += _attn_qkv_o_fmac_per_layer( + batch_size=batch_size, + seq_len_eff=seq_len_eff, + hidden_size=hidden_size, + num_heads=num_heads, + head_dim=head_dim, + q_lora_rank=q_lora_rank, + o_lora_rank=o_lora_rank, + o_groups=o_groups, + ) + attn_scores += _attn_scores_fmac_per_layer( + batch_size=batch_size, + seq_len_eff=seq_len_eff, + num_heads=num_heads, + head_dim=head_dim, + compress_ratio=ratio, + index_topk=index_topk, + swa_window=swa_window, + ) + compressor += _compressor_fmac_per_layer( + batch_size=batch_size, + seq_len_eff=seq_len_eff, + hidden_size=hidden_size, + head_dim=head_dim, + compress_ratio=ratio, + ) + indexer += _indexer_fmac_per_layer( + batch_size=batch_size, + seq_len_eff=seq_len_eff, + hidden_size=hidden_size, + compress_ratio=ratio, + index_head_dim=index_head_dim, + index_n_heads=index_n_heads, + ) + moe += _moe_fmac_per_layer( + batch_size=batch_size, + seq_len_eff=seq_len_eff, + hidden_size=hidden_size, + moe_ffn_hidden_size=moe_ffn_hidden_size, + moe_router_topk=moe_router_topk, + num_experts=num_experts, + is_hash_layer=(layer_idx < num_hash_layers), + shared_expert_ffn_hidden_size=shared_expert_ffn_hidden_size, + ) + hc += _hc_mixer_fmac_per_layer( + batch_size=batch_size, + seq_len=seq_len, + hidden_size=hidden_size, + hc_mult=hc_mult, + ) + + # ---- MTP layers (one full V4 layer per depth + eh_proj per depth) ---- + mtp_attn_qkv_o = 0 + mtp_attn_scores = 0 + mtp_compressor = 0 + mtp_indexer = 0 + mtp_moe = 0 + mtp_hc = 0 + for depth in range(mtp_num_layers): + ratio = int(mtp_ratios[depth]) if depth < len(mtp_ratios) else 0 + mtp_attn_qkv_o += _attn_qkv_o_fmac_per_layer( + batch_size=batch_size, + seq_len_eff=seq_len_eff, + hidden_size=hidden_size, + num_heads=num_heads, + head_dim=head_dim, + q_lora_rank=q_lora_rank, + o_lora_rank=o_lora_rank, + o_groups=o_groups, + ) + mtp_attn_scores += _attn_scores_fmac_per_layer( + batch_size=batch_size, + seq_len_eff=seq_len_eff, + num_heads=num_heads, + head_dim=head_dim, + compress_ratio=ratio, + index_topk=index_topk, + swa_window=swa_window, + ) + mtp_compressor += _compressor_fmac_per_layer( + batch_size=batch_size, + seq_len_eff=seq_len_eff, + hidden_size=hidden_size, + head_dim=head_dim, + compress_ratio=ratio, + ) + mtp_indexer += _indexer_fmac_per_layer( + batch_size=batch_size, + seq_len_eff=seq_len_eff, + hidden_size=hidden_size, + compress_ratio=ratio, + index_head_dim=index_head_dim, + index_n_heads=index_n_heads, + ) + mtp_moe += _moe_fmac_per_layer( + batch_size=batch_size, + seq_len_eff=seq_len_eff, + hidden_size=hidden_size, + moe_ffn_hidden_size=moe_ffn_hidden_size, + moe_router_topk=moe_router_topk, + num_experts=num_experts, + # V4's MTP depths run after num_hash_layers in the routing ordering + # (the released checkpoint stores topk router weights for them). + is_hash_layer=False, + shared_expert_ffn_hidden_size=shared_expert_ffn_hidden_size, + ) + mtp_hc += _hc_mixer_fmac_per_layer( + batch_size=batch_size, + seq_len=seq_len, + hidden_size=hidden_size, + hc_mult=hc_mult, + ) + + attn_qkv_o += mtp_attn_qkv_o + attn_scores += mtp_attn_scores + compressor += mtp_compressor + indexer += mtp_indexer + moe += mtp_moe + hc += mtp_hc + hc += _hc_head_fmac( + batch_size=batch_size, + seq_len=seq_len, + hidden_size=hidden_size, + hc_mult=hc_mult, + mtp_num_layers=mtp_num_layers, + ) + + mtp = _mtp_eh_proj_fmac( + batch_size=batch_size, + seq_len=seq_len, + hidden_size=hidden_size, + mtp_num_layers=mtp_num_layers, + ) + + logits = _logits_fmac( + batch_size=batch_size, + seq_len=seq_len, + hidden_size=hidden_size, + padded_vocab_size=padded_vocab_size, + mtp_num_layers=mtp_num_layers, + ) + + breakdown = _V4FlopsBreakdown( + attn_qkv_o=attn_qkv_o, + attn_scores=attn_scores, + compressor=compressor, + indexer=indexer, + moe=moe, + mtp=mtp, + logits=logits, + hc=hc, + ) + return breakdown.to_total_flops(), breakdown + + +# --------------------------------------------------------------------------- +# Wrapper installation +# --------------------------------------------------------------------------- + +# Module-level latch so the breakdown is logged exactly once even though +# ``num_floating_point_operations`` is called many times per training run +# (one per ``training_log`` and one per ``train_step``). +_BREAKDOWN_LOGGED = False + + +def _emit_breakdown( + *, + args: Any, + batch_size: int, + breakdown: _V4FlopsBreakdown, + total_flops: int, +) -> None: + """Emit the per-component breakdown via single-line ``log_rank_0`` calls. + + One row per ``log_rank_0`` call (instead of a single multi-line message) + so each line passes cleanly through Primus's per-line logger formatter + and rank-aware filter. The header line carries the run-shape metadata + so the breakdown can be matched to a specific ``(args, batch_size)`` + pairing in the log. + """ + + def _tflops(fmac: int) -> float: + return fmac * _FORWARD_BACKWARD_FACTOR * _FMA_FACTOR / 1.0e12 + + rows: List[Tuple[str, float]] = [ + ("attn_qkv_o", _tflops(breakdown.attn_qkv_o)), + ("attn_scores", _tflops(breakdown.attn_scores)), + ("compressor", _tflops(breakdown.compressor)), + ("indexer", _tflops(breakdown.indexer)), + ("moe", _tflops(breakdown.moe)), + ("mtp_eh_proj", _tflops(breakdown.mtp)), + ("logits", _tflops(breakdown.logits)), + ("hc", _tflops(breakdown.hc)), + ] + + log_rank_0( + "[Patch:megatron.deepseek_v4.flops_reporting] V4 closed-form FLOPs " + f"breakdown -- batch_size={batch_size}, seq_length={int(args.seq_length)}, " + f"hc_mult={int(getattr(args, 'hc_mult', 1) or 1)}, " + f"num_layers={int(args.num_layers)}, " + f"mtp_num_layers={int(getattr(args, 'mtp_num_layers', 0) or 0)}" + ) + for name, tflops in rows: + log_rank_0(f"[Patch:megatron.deepseek_v4.flops_reporting] {name:<12s} = {tflops:9.3f} TFLOP") + log_rank_0( + f"[Patch:megatron.deepseek_v4.flops_reporting] {'TOTAL':<12s} = " + f"{total_flops / 1.0e12:9.3f} TFLOP / global-batch" + ) + + +def _make_v4_num_floating_point_operations(original_fn, *, dispatch_v4: bool): + """Return a wrapper that dispatches V4 vs upstream model types. + + ``dispatch_v4`` is captured at install time from ``args.model_type`` + rather than re-checked per call, because Megatron's ``pretrain()`` + overwrites ``args.model_type`` with the ``ModelType`` enum at + ``training.py:1210`` *before* ``train()`` ever calls + ``num_floating_point_operations``. At that point the original + YAML-set string ``"deepseek_v4"`` is gone and a runtime check would + silently fall through to the upstream formula. + """ + + def wrapped(args, batch_size): + if not dispatch_v4: + return original_fn(args, batch_size) + + total_flops, breakdown = compute_v4_flops(args, batch_size) + + global _BREAKDOWN_LOGGED + if not _BREAKDOWN_LOGGED: + _BREAKDOWN_LOGGED = True + _emit_breakdown( + args=args, + batch_size=batch_size, + breakdown=breakdown, + total_flops=total_flops, + ) + + return total_flops + + wrapped.__wrapped__ = original_fn + wrapped._v4_flops_patched = True + return wrapped + + +_TRAINER_REBIND_TARGETS: Sequence[str] = ( + # Primus's Megatron trainer imports ``num_floating_point_operations`` at + # module load time (``primus.modules.trainer.megatron.trainer``: line 125) + # and resolves the bare name from its OWN globals at the call site + # (``trainer.train()``: line 1452). Updating only + # ``megatron.training.training.num_floating_point_operations`` therefore + # never reaches that bound name and the trainer keeps using the upstream + # GPT/MLA-shaped function. We rebind the trainer's local name to the + # wrapper too so the V4 closed form is what actually drives per-iter + # TFLOPs reporting. Listed explicitly so a missing module is loud rather + # than silently silent. + "primus.modules.trainer.megatron.trainer", +) + + +def _rebind_trainer_imports(wrapped_fn) -> List[str]: + """Rebind ``num_floating_point_operations`` in every Primus module that + captured the upstream symbol at import time. + + Returns the list of modules that were actually rebound so the install log + can show whether each downstream binding was wired up. Modules that were + not yet imported (e.g. on a cold-cache trainer init) are skipped silently + — the trainer module imports the function at its own load, which + happens before ``before_train``, so in practice the targeted module is + always present at this point. + """ + import sys + + rebound: List[str] = [] + for module_name in _TRAINER_REBIND_TARGETS: + mod = sys.modules.get(module_name) + if mod is None: + continue + if getattr(mod, "num_floating_point_operations", None) is wrapped_fn: + continue + if not hasattr(mod, "num_floating_point_operations"): + continue + mod.num_floating_point_operations = wrapped_fn + rebound.append(module_name) + return rebound + + +@register_patch( + "megatron.deepseek_v4.flops_reporting", + backend="megatron", + phase="before_train", + description=( + "DeepSeek-V4: replace Megatron's GPT/MLA-shaped " + "num_floating_point_operations with a V4 closed form (Q LoRA + " + "single-latent KV + grouped low-rank O at S * hc_mult, plus " + "Compressor / Indexer side paths and MTP per-depth full inner " + "layer cost). Falls through byte-for-byte for non-V4 model types." + ), + condition=lambda ctx: getattr(get_args(ctx), "model_type", None) == "deepseek_v4", +) +def patch_v4_flops_reporting(ctx: PatchContext): + """Install the V4 FLOPs wrapper on ``training.num_floating_point_operations``.""" + import megatron.training.training as training_module + + original_fn = training_module.num_floating_point_operations + if getattr(original_fn, "_v4_flops_patched", False): + log_rank_0( + "[Patch:megatron.deepseek_v4.flops_reporting] " + "num_floating_point_operations already patched, skip" + ) + return + + # Capture the V4 dispatch decision NOW, while ``args.model_type`` is + # still the YAML-set string. Megatron's ``pretrain()`` will rebind + # ``args.model_type`` to a ``ModelType`` enum at + # ``training.py:1210`` later, so any runtime check inside the wrapper + # would fail. This patch is gated to only install when the YAML model + # type is V4 (see ``condition`` on the decorator), so it's safe to + # hard-set ``dispatch_v4=True`` here. + wrapped = _make_v4_num_floating_point_operations(original_fn, dispatch_v4=True) + training_module.num_floating_point_operations = wrapped + rebound = _rebind_trainer_imports(wrapped) + + log_rank_0( + "[Patch:megatron.deepseek_v4.flops_reporting] wrapped " + "num_floating_point_operations; per-iter TFLOPs now reported " + "with V4-aware closed form (see " + "deepseek-v4/develop/plan-3/02-phase-details.md#p20 for the formula)." + ) + if rebound: + log_rank_0( + "[Patch:megatron.deepseek_v4.flops_reporting] rebound trainer " f"import bindings: {rebound}" + ) + else: + log_rank_0( + "[Patch:megatron.deepseek_v4.flops_reporting] no trainer modules " + "needed rebinding (none of " + f"{list(_TRAINER_REBIND_TARGETS)} were imported yet)." + ) + + +__all__: Sequence[str] = ( + "compute_v4_flops", + "patch_v4_flops_reporting", +) diff --git a/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py b/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py new file mode 100644 index 000000000..9fb0c2de5 --- /dev/null +++ b/primus/backends/megatron/patches/deepseek_v4_get_batch_patches.py @@ -0,0 +1,271 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 hash-router input_ids PP broadcast (upfront/pre-loop). + +V4's hash-routed MoE layers (the first ``num_hash_layers`` MoE layers) +look up a static ``tid2eid`` table and therefore need the raw +``input_ids`` on every PP stage that owns one. Megatron's +:func:`pretrain_gpt.get_batch` only loads tokens on the first / last PP +stage; middle stages short-circuit to ``return None, None, None, ...``. + +Plan-2 P19 attempted two simpler hooks before this one: + +1. **In-forward broadcast** inside :meth:`DeepseekV4Model.forward`. +2. **Per-call broadcast** wrapping :func:`pretrain_gpt.get_batch`. + +Both work for the non-interleaved 1F1B schedule (``PP>1``, +``VPP=1``) but deadlock the interleaved-1F1B / VPP schedule. The reason: +the interleaved scheduler issues a single ``recv_forward`` *before* the +warm-up loop on every non-first PP rank (see +:func:`forward_backward_pipelining_with_interleaving` in +``megatron/core/pipeline_parallel/schedules.py:1363-1392``). PP rank > 0 +therefore parks in that ``recv_forward.wait()`` until PP rank 0's first +``send_forward`` arrives. PP rank 0 meanwhile enters the warm-up loop, +calls ``forward_step`` -> ``get_batch`` -> ``dist.broadcast`` and gets +stuck waiting for PP rank > 0 to issue the matching broadcast — which +will never happen because PP rank > 0 is itself blocked on +``recv_forward``. The non-interleaved 1F1B schedule does not hit this +because its pre-loop ``recv_forward`` is *inside* each warm-up iter +(``schedules.py:2128-2156``), so PP rank > 0 reaches its broadcast call +on the same iteration that PP rank 0 issues the matching ``send_forward`` +and the broadcast pairs up before either rank stalls. + +This patch instead does **all** PP token broadcasts up-front, *before* +the schedule's first ``recv_forward``: + +* It wraps :func:`megatron.core.pipeline_parallel.get_forward_backward_func` + so that every schedule fetched for a train_step is replaced by a + thin wrapper which: + 1. Pre-loads ``num_microbatches`` × ``num_chunks`` batches by calling + the original ``pretrain_gpt.get_batch`` on the first / last PP + stages, and allocating empty token buffers on middle PP stages. + 2. Runs one ``dist.broadcast`` per (chunk, microbatch) on the PP + group, sourced from PP rank 0. All collectives fire before any + ``send_forward`` / ``recv_forward`` runs, so they pair up + deterministically across ranks — no deadlock. + 3. Caches the resulting tuples in a module-local store keyed by + chunk and microbatch ordinal, then calls the underlying schedule + with the original ``data_iterator``. +* It wraps :func:`pretrain_gpt.get_batch` so that, while the cache is + active, calls return the pre-cached tuple for the corresponding + (vp_stage, microbatch) — bypassing the data iterator. Outside a + train_step (e.g. during eval not routed through the schedule) the + wrapper falls back to the original ``get_batch``. +* The cache is reset after each schedule call (success or exception) + so subsequent train_steps start clean. + +Gating: ``model_type == "deepseek_v4"``, ``num_hash_layers > 0``, and +``pipeline_model_parallel_size > 1``. The patch is a strict no-op for +any other model. + +Cost analysis (on the V4 BF16 smoke): + +* Each pre-broadcast moves ``mbs * seq * 8 B`` per microbatch + (~1 KiB for ``mbs=1, seq=128``). With ``num_microbatches=16, + num_chunks=2`` that is ~32 KiB total, dwarfed by activation P2P + (~32 MiB per microbatch). +* Pre-loading consumes the data_iterator on PP rank 0 / last during + ``pre_broadcast`` rather than spread across the warm-up. This + collapses ``batch-generator`` time into a single bursty phase. +""" + +from typing import Any, Optional + +import torch + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + +# --------------------------------------------------------------------------- +# Module-local pre-broadcast cache. +# +# ``data[chunk_id][microbatch_id]`` holds the 6-tuple returned by +# :func:`pretrain_gpt.get_batch`: +# (tokens, labels, loss_mask, attention_mask, position_ids, packed_seq_params) +# +# On middle PP stages only ``tokens`` is meaningful; the other fields stay +# ``None`` (matching what the original ``get_batch`` returns there). +# ``consumed[chunk_id]`` tracks how many microbatches have been popped from +# the cache during the current schedule call; it lets us return the right +# entry from the patched ``get_batch`` regardless of which interleaved +# microbatch ordinal is being processed (interleaved / 1F1B both consume +# microbatches in increasing order *per chunk*). +# --------------------------------------------------------------------------- +_V4_PP_TOKEN_CACHE: dict = { + "active": False, + "data": [], + "consumed": [], +} + + +def _v4_reset_cache() -> None: + _V4_PP_TOKEN_CACHE["active"] = False + _V4_PP_TOKEN_CACHE["data"] = [] + _V4_PP_TOKEN_CACHE["consumed"] = [] + + +def _v4_pre_broadcast_step( + data_iterator: Any, + num_microbatches: int, + original_get_batch, +) -> None: + """Pre-broadcast V4 ``input_ids`` for every (chunk, microbatch) of one step.""" + from megatron.core import parallel_state + from megatron.training import get_args as _get_args + + args = _get_args() + pp_group = parallel_state.get_pipeline_model_parallel_group() + src_global = torch.distributed.get_global_rank(pp_group, 0) + pp_rank = parallel_state.get_pipeline_model_parallel_rank() + pp_world = parallel_state.get_pipeline_model_parallel_world_size() + + cp_size = max(1, int(getattr(args, "context_parallel_size", 1) or 1)) + mbs = int(args.micro_batch_size) + seq = int(args.seq_length) // cp_size + device = torch.device("cuda", torch.cuda.current_device()) + + # ``is_first_or_last_pp`` mirrors :func:`pretrain_gpt.get_batch`'s + # gate: the original returns real tokens only on the first / last PP + # stage. (MTP can also build the dataset on a middle stage; that + # case is left to the original short-circuit because the smoke runs + # ``mtp_num_layers=0`` and the broadcast logic still tolerates a + # ``None``-tuple result by falling back to an empty token buffer.) + is_first_or_last_pp = (pp_rank == 0) or (pp_rank == pp_world - 1) + + if isinstance(data_iterator, list): + iter_list = data_iterator + num_chunks = len(iter_list) + # When VPP is enabled, each chunk has its own ``vp_stage`` index. + chunk_vp_stages = list(range(num_chunks)) + else: + iter_list = [data_iterator] + num_chunks = 1 + # Non-VPP: ``vp_stage`` is ``None`` (matches what + # ``pretrain_gpt.forward_step`` reads off the model attribute). + chunk_vp_stages = [None] + + cache_data: list = [] + for chunk_id in range(num_chunks): + chunk_iter = iter_list[chunk_id] + chunk_vp_stage = chunk_vp_stages[chunk_id] + chunk_cache: list = [] + for _mb_id in range(num_microbatches): + if is_first_or_last_pp and chunk_iter is not None: + tup = original_get_batch(chunk_iter, chunk_vp_stage) + tokens = tup[0] + if tokens is None: + # Should not happen on first / last PP for V4 smoke; + # fall back to empty buffer and let the broadcast + # populate. This keeps us robust against unusual + # configs (e.g. dataset-on-rank gating disabled). + tokens = torch.empty([mbs, seq], dtype=torch.long, device=device) + cached = (tokens, tup[1], tup[2], tup[3], tup[4], tup[5]) + else: + tokens = torch.empty([mbs, seq], dtype=torch.long, device=device) + cached = (tokens, None, None, None, None, None) + + torch.distributed.broadcast(tokens, src=src_global, group=pp_group) + chunk_cache.append(cached) + cache_data.append(chunk_cache) + + _V4_PP_TOKEN_CACHE["data"] = cache_data + _V4_PP_TOKEN_CACHE["consumed"] = [0] * num_chunks + _V4_PP_TOKEN_CACHE["active"] = True + + +def _make_v4_get_batch(original_get_batch): + """Return a ``get_batch`` wrapper that consumes the pre-broadcast cache.""" + + def patched_get_batch(data_iterator: Any, vp_stage: Optional[int] = None): + if not _V4_PP_TOKEN_CACHE["active"]: + return original_get_batch(data_iterator, vp_stage) + chunk_id = vp_stage if vp_stage is not None else 0 + counter = _V4_PP_TOKEN_CACHE["consumed"][chunk_id] + cached = _V4_PP_TOKEN_CACHE["data"][chunk_id][counter] + _V4_PP_TOKEN_CACHE["consumed"][chunk_id] += 1 + return cached + + patched_get_batch.__wrapped__ = original_get_batch + patched_get_batch._v4_pp_get_batch_patched = True + return patched_get_batch + + +def _make_v4_pre_broadcast_schedule(original_schedule, original_get_batch): + """Wrap ``forward_backward_func`` to pre-broadcast V4 tokens up-front.""" + + def patched_schedule(*args, **kwargs): + data_iterator = kwargs.get("data_iterator") + num_microbatches = int(kwargs.get("num_microbatches", 1) or 1) + + try: + _v4_pre_broadcast_step(data_iterator, num_microbatches, original_get_batch) + return original_schedule(*args, **kwargs) + finally: + _v4_reset_cache() + + patched_schedule._v4_pp_schedule_wrapped = True + return patched_schedule + + +@register_patch( + "megatron.deepseek_v4.pp_token_pre_broadcast", + backend="megatron", + phase="before_train", + description=( + "DeepSeek-V4: pre-broadcast input_ids from PP rank 0 across the PP " + "group up-front in the forward_backward schedule wrapper, before the " + "first recv_forward, so middle PP stages owning hash-routed MoE " + "layers can read raw token IDs without deadlocking the interleaved " + "1F1B / VPP schedule." + ), + condition=lambda ctx: ( + getattr(get_args(ctx), "model_type", None) == "deepseek_v4" + and int(getattr(get_args(ctx), "num_hash_layers", 0) or 0) > 0 + and int(getattr(get_args(ctx), "pipeline_model_parallel_size", 1) or 1) > 1 + ), + # Ordered after pp_dump_data so its schedule_wrapper does not double-wrap; + # see ``pp_dump_data_patches.py`` for the priority=100 anchor. + priority=60, +) +def patch_v4_pp_token_pre_broadcast(ctx: PatchContext): + """Install the V4 PP token pre-broadcast hooks.""" + import megatron.core.pipeline_parallel as pp_module + import megatron.training.training as training_module + import pretrain_gpt + + original_get_batch = pretrain_gpt.get_batch + if getattr(original_get_batch, "_v4_pp_get_batch_patched", False): + log_rank_0("[Patch:megatron.deepseek_v4.pp_token_pre_broadcast] get_batch " "already patched, skip") + return + + # Hook 1: replace pretrain_gpt.get_batch with a cache-consuming wrapper. + # We capture ``original_get_batch`` here so the schedule wrapper can call + # the *unpatched* implementation during the pre-broadcast phase (the + # patched version would just hit an empty cache and recurse). + pretrain_gpt.get_batch = _make_v4_get_batch(original_get_batch) + + # Hook 2: replace get_forward_backward_func so every schedule fetched + # by ``train_step`` is wrapped with the pre-broadcast. + original_get_fbf = pp_module.get_forward_backward_func + + def wrapped_get_fbf(): + original_schedule = original_get_fbf() + if getattr(original_schedule, "_v4_pp_schedule_wrapped", False): + return original_schedule + return _make_v4_pre_broadcast_schedule(original_schedule, original_get_batch) + + pp_module.get_forward_backward_func = wrapped_get_fbf + training_module.get_forward_backward_func = wrapped_get_fbf + + log_rank_0( + "[Patch:megatron.deepseek_v4.pp_token_pre_broadcast] wrapped " + "pretrain_gpt.get_batch + get_forward_backward_func; PP rank 0 " + "broadcasts input_ids up-front (once per microbatch × chunk per " + "train_step) so middle PP stages owning hash-routed MoE layers " + "see real token IDs without per-(chunk, microbatch) collectives " + "racing the interleaved 1F1B P2P sends." + ) diff --git a/primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py b/primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py new file mode 100644 index 000000000..b163c24fe --- /dev/null +++ b/primus/backends/megatron/patches/deepseek_v4_pp_shape_patches.py @@ -0,0 +1,162 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""DeepSeek-V4 PP tensor-shape patch. + +V4's HyperConnections (mHC) residual carries ``K = hc_mult`` parallel +streams per position. Inside a single PP stage the layer hidden has +shape ``[B, S, K, D]``; at the PP boundary the V4 transformer block +folds K into the sequence axis via +:func:`primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_block._lower_streams_out` +so the wire tensor is ``[S * K, B, D]`` (plan-2 P15 design — see +``deepseek-v4/develop/plan-2/03-phase-details.md`` P15 and the C1 +finding in ``00-review-findings.md``). + +Megatron has **two** code paths that compute the PP wire tensor shape, +and they need *both* to know about V4's K packing: + +1. The non-interleaved 1F1B schedule (``forward_backward_pipelining_ + without_interleaving``) calls + :func:`megatron.core.pipeline_parallel.schedules.get_tensor_shapes` + (``schedules.py:2096-2103``). Wrapping that function suffices for + smokes A/B/C. + +2. The interleaved 1F1B / VPP schedule (``forward_backward_pipelining_ + with_interleaving``) instead computes ``tensor_shape`` *inline* from + the ``seq_length`` argument (``schedules.py:1001-1004``): + + tensor_shape = [seq_length, micro_batch_size, config.hidden_size] + tensor_shape[0] = tensor_shape[0] // cp_group.size() + if config.sequence_parallel: + tensor_shape[0] = tensor_shape[0] // tp_group.size() + + It does **not** call ``get_tensor_shapes``. With VPP=2 and ``hc_mult= + 4`` this leaves the recv buffer at ``[S, B, hidden]`` while the + sender emits ``[S*K, B, hidden]``. PyTorch P2P does not validate + shape (only ``numel * dtype_size``), so the receiver silently copies + only the first ``S * hidden`` elements, ``_lift_streams_in`` + reshapes them as ``[B, S/K, K, D]``, and the resulting hidden + flattens to ``S/K = 32`` instead of ``S = 128`` — :class:`Deepseek + V4HashRouter` then trips its precondition with ``hidden=32 vs + token_ids=128``. (P19 smoke D run reproduces this cleanly.) + +We therefore install **two** complementary wrappers, both gated on the +V4 ``model_type`` + ``hc_mult > 1`` + ``PP > 1`` condition: + +* :func:`_make_v4_get_tensor_shapes` multiplies the first (seq) dim of + every tuple returned by ``get_tensor_shapes`` by ``hc_mult``, fixing + path (1). +* :func:`_make_v4_interleaved_schedule` wraps + ``forward_backward_pipelining_with_interleaving`` to scale its + ``seq_length`` keyword argument by ``hc_mult`` before the schedule + computes its inline ``tensor_shape``, fixing path (2). Inside the + interleaved schedule, ``seq_length`` is *only* read for that inline + shape (see grep on ``schedules.py``), so scaling it is a no-op for + every other concern (cudagraph, attention masking, etc.). + +The companion ``adjust_tensor_shapes_fn`` parameter on the +non-interleaved 1F1B schedule is intentionally not used because +upstream Megatron explicitly asserts it is unsupported by the +interleaved / VPP schedules (``schedules.py:900-901``); going through +``get_tensor_shapes`` and ``seq_length`` keeps the behavior uniform +across all schedules. +""" + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +def _make_v4_get_tensor_shapes(original_fn, hc_mult: int): + """Return a wrapper that scales the first (seq) dim by ``hc_mult``. + + ``original_fn`` returns a list of ``(seq_length, micro_batch, hidden)`` + triples; the V4 PP wire packs ``K`` streams into the seq axis so we + multiply the seq dim only — micro_batch and hidden are unchanged. + """ + + def patched_get_tensor_shapes(*args, **kwargs): + shapes = original_fn(*args, **kwargs) + return [(s * hc_mult, b, h) for (s, b, h) in shapes] + + patched_get_tensor_shapes.__wrapped__ = original_fn + patched_get_tensor_shapes._v4_pp_shape_patched = True + return patched_get_tensor_shapes + + +def _make_v4_interleaved_schedule(original_fn, hc_mult: int): + """Wrap the interleaved schedule to scale ``seq_length`` by ``hc_mult``. + + The interleaved schedule (``forward_backward_pipelining_with_ + interleaving``) computes its PP wire ``tensor_shape`` inline from + the ``seq_length`` kwarg (``schedules.py:1001``). Inside that + function, ``seq_length`` is consumed only by that inline + computation, so we can safely scale it on the way in to give the + schedule a V4-aware shape without touching any other behaviour. + """ + + def patched_schedule(*args, **kwargs): + if "seq_length" in kwargs and kwargs["seq_length"] is not None: + kwargs["seq_length"] = int(kwargs["seq_length"]) * hc_mult + return original_fn(*args, **kwargs) + + patched_schedule.__wrapped__ = original_fn + patched_schedule._v4_pp_interleaved_patched = True + return patched_schedule + + +@register_patch( + "megatron.deepseek_v4.pp_tensor_shape", + backend="megatron", + phase="before_train", + description=( + "DeepSeek-V4: pack hc_mult=K hyper-streams into the PP wire seq " + "axis so [S*K, B, D] passes between PP stages (covers both the " + "1F1B get_tensor_shapes path and the interleaved-1F1B / VPP " + "inline tensor_shape path)." + ), + condition=lambda ctx: ( + getattr(get_args(ctx), "model_type", None) == "deepseek_v4" + and int(getattr(get_args(ctx), "hc_mult", 1) or 1) > 1 + and int(getattr(get_args(ctx), "pipeline_model_parallel_size", 1) or 1) > 1 + ), +) +def patch_v4_pp_tensor_shape(ctx: PatchContext): + """Multiply the PP P2P seq dim by ``hc_mult`` for V4 models.""" + import megatron.core.pipeline_parallel.schedules as schedules_module + + hc_mult = int(getattr(get_args(ctx), "hc_mult", 1)) + + # Wrapper 1: get_tensor_shapes (used by the non-interleaved schedule). + original_get_tensor_shapes = schedules_module.get_tensor_shapes + if getattr(original_get_tensor_shapes, "_v4_pp_shape_patched", False): + log_rank_0("[Patch:megatron.deepseek_v4.pp_tensor_shape] get_tensor_shapes " "already patched, skip") + else: + schedules_module.get_tensor_shapes = _make_v4_get_tensor_shapes(original_get_tensor_shapes, hc_mult) + log_rank_0( + f"[Patch:megatron.deepseek_v4.pp_tensor_shape] wrapped " + f"get_tensor_shapes; PP wire seq_len * hc_mult={hc_mult} " + "(packs K hyper-streams into the sequence axis)." + ) + + # Wrapper 2: forward_backward_pipelining_with_interleaving (VPP). + # The interleaved schedule reads ``seq_length`` directly to build its + # inline ``tensor_shape``; scaling the kwarg gives it the V4 wire shape + # without rewriting the function. + original_interleaved = schedules_module.forward_backward_pipelining_with_interleaving + if getattr(original_interleaved, "_v4_pp_interleaved_patched", False): + log_rank_0( + "[Patch:megatron.deepseek_v4.pp_tensor_shape] interleaved " "schedule already patched, skip" + ) + else: + schedules_module.forward_backward_pipelining_with_interleaving = _make_v4_interleaved_schedule( + original_interleaved, hc_mult + ) + log_rank_0( + f"[Patch:megatron.deepseek_v4.pp_tensor_shape] wrapped " + f"forward_backward_pipelining_with_interleaving; " + f"seq_length * hc_mult={hc_mult} on the way into the " + "interleaved-1F1B / VPP schedule." + ) diff --git a/primus/backends/megatron/patches/emerging_optimizers_log_level_patches.py b/primus/backends/megatron/patches/emerging_optimizers_log_level_patches.py new file mode 100644 index 000000000..3f11a2c05 --- /dev/null +++ b/primus/backends/megatron/patches/emerging_optimizers_log_level_patches.py @@ -0,0 +1,58 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Raise the ``emerging_optimizers`` (absl) logger to INFO. + +The Muon path pulls in ``emerging_optimizers``, which logs via ``absl`` (``from +absl import logging``). Its Newton-Schulz helper emits a per-call ``DEBUG`` line +for the coefficient schedule (``muon_utils.get_coefficient_iterator``):: + + Iterating through 10 steps with cycle mode. + Coefficient sets: [(3.4445, -4.775, 2.0315), ..., (2.0, -1.5, 0.5)] + +Because Muon orthogonalizes every matrix parameter on every optimizer step, this +fires repeatedly and floods the logs. + +All ``absl`` logging records flow through the single stdlib logger named +``"absl"`` (Primus routes stdlib logging into loguru via an InterceptHandler on +the root logger, with ``root`` at NOTSET, so these DEBUG records are emitted). +Raising the ``"absl"`` logger to INFO drops the DEBUG spam at the source while +keeping INFO and above. We do it in the ``before_train`` phase, which runs after +the optimizer is built but before the first training step. + +Set ``PRIMUS_VERBOSE_EMERGING_OPTIMIZERS=1`` to keep the full DEBUG trace. +""" + +import logging +import os + +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +@register_patch( + "megatron.emerging_optimizers.log_level", + backend="megatron", + phase="before_train", + description=( + "Raise the 'absl' logger (used by emerging_optimizers) to INFO so the " + "Muon Newton-Schulz coefficient DEBUG lines do not flood the logs." + ), + condition=lambda ctx: os.environ.get("PRIMUS_VERBOSE_EMERGING_OPTIMIZERS", "0") != "1", +) +def patch_emerging_optimizers_log_level(ctx: PatchContext): + """Set the absl logger threshold to INFO. Idempotent / best-effort.""" + del ctx + try: + logging.getLogger("absl").setLevel(logging.INFO) + except Exception as exc: # never block training over a logging tweak. + log_rank_0( + f"[Patch:megatron.emerging_optimizers.log_level] could not set " f"'absl' logger level: {exc!r}" + ) + return + + log_rank_0("[Patch:megatron.emerging_optimizers.log_level] 'absl' logger raised to INFO.") diff --git a/primus/backends/megatron/patches/fused_pad_routing_map_patches.py b/primus/backends/megatron/patches/fused_pad_routing_map_patches.py new file mode 100644 index 000000000..7702dc538 --- /dev/null +++ b/primus/backends/megatron/patches/fused_pad_routing_map_patches.py @@ -0,0 +1,68 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Primus fused_pad_routing_map patch. + +Globally replaces Megatron's ``fused_pad_routing_map`` with the Primus Triton +implementation (``primus.backends.megatron.core.fusions.fused_pad_routing_map``) +without modifying upstream Megatron source. + +The Primus version rewrites the kernel to operate directly on the native +``[num_tokens, num_experts]`` layout (no transpose/copy) and drops the +``@jit_fuser`` (``torch.compile``) wrapper, avoiding the Triton kernel +functionalization failure seen on some torch/triton combos. +""" + +import sys + +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +@register_patch( + "megatron.fused_pad_routing_map", + backend="megatron", + phase="before_train", + description="Replace Megatron fused_pad_routing_map with the Primus Triton implementation", +) +def patch_fused_pad_routing_map(ctx: PatchContext): + """Swap the ``fused_pad_routing_map`` symbol everywhere it is referenced.""" + import megatron.core.fusions.fused_pad_routing_map as meg_mod + + from primus.backends.megatron.core.fusions.fused_pad_routing_map import ( + fused_pad_routing_map as primus_fused_pad_routing_map, + ) + + log_rank_0("[Patch:megatron.fused_pad_routing_map] Patching fused_pad_routing_map...") + + # Original function object; used to precisely locate stale references. + orig_fn = getattr(meg_mod, "fused_pad_routing_map", None) + if orig_fn is primus_fused_pad_routing_map: + log_rank_0("[Patch:megatron.fused_pad_routing_map] Already patched; skipping.") + return + + # 1) Replace on the source module so all *future* (incl. lazy) imports resolve + # to the Primus version. This alone covers the common case, since this patch + # runs before token_dispatcher is imported. + meg_mod.fused_pad_routing_map = primus_fused_pad_routing_map + + # 2) Replace references already bound into other modules that imported the symbol + # at top level before this patch ran (precise: only objects that `is orig_fn`). + patched_modules = [] + if orig_fn is not None: + for mod_name, module in list(sys.modules.items()): + if module is None or module is meg_mod: + continue + if getattr(module, "fused_pad_routing_map", None) is orig_fn: + setattr(module, "fused_pad_routing_map", primus_fused_pad_routing_map) + patched_modules.append(mod_name) + + log_rank_0( + "[Patch:megatron.fused_pad_routing_map] Patched " + "megatron.core.fusions.fused_pad_routing_map.fused_pad_routing_map " + f"-> primus (also updated already-imported refs: {patched_modules or 'none'})" + ) diff --git a/primus/backends/megatron/patches/mla_patches.py b/primus/backends/megatron/patches/mla_patches.py index e25ac2d2d..d645d2871 100644 --- a/primus/backends/megatron/patches/mla_patches.py +++ b/primus/backends/megatron/patches/mla_patches.py @@ -20,9 +20,26 @@ backend="megatron", phase="before_train", description=( - "Monkey patch MLA attention to use PrimusMLASelfAttention " "when use_turbo_gemm is enabled." + "Monkey patch MLA attention to use PrimusMLASelfAttention " + "when use_turbo_gemm is enabled (skipped for DeepSeek-V4)." + ), + # Skip for DeepSeek-V4: its DeepseekV4Attention subclasses MLASelfAttention, + # and PrimusMLASelfAttention deliberately bypasses MLASelfAttention.__init__ + # (calls the grandparent), which would break V4's super().__init__ chain if + # the base class were swapped. V4 builds DeepseekV4Attention directly and does + # not need the padded-fusion MLA path, so this patch must not fire for it. + condition=lambda ctx: ( + getattr(get_args(ctx), "use_turbo_gemm", False) + and not any( + getattr(get_args(ctx), _f, False) + for _f in ( + "use_v4_triton_attention", + "use_v4_triton_csa_attention", + "use_v4_tilelang_attention", + "use_v4_tilelang_csa_attention", + ) + ) ), - condition=lambda ctx: getattr(get_args(ctx), "use_turbo_gemm", False), ) def patch_mla_attention(ctx: PatchContext): """ diff --git a/primus/backends/megatron/patches/moe_alltoall_dtoh_patches.py b/primus/backends/megatron/patches/moe_alltoall_dtoh_patches.py new file mode 100644 index 000000000..d4e15eeca --- /dev/null +++ b/primus/backends/megatron/patches/moe_alltoall_dtoh_patches.py @@ -0,0 +1,105 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Primus MoE All-to-All dispatcher D2H patch. + +Patches ``MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize`` so that when +``use_turbo_grouped_gemm`` is enabled, ``tokens_per_expert`` is kept on-device +(PrimusTurbo grouped gemm consumes it on the GPU) instead of being copied to the +host. All other splits are still moved to CPU and the stream sync is unchanged. +""" + +import torch + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +def _turbo_grouped_gemm_on_device(ctx: PatchContext) -> bool: + """Whether PrimusTurbo grouped gemm consumes tokens_per_expert on device. + + The authoritative source is the global Primus/Megatron args, not the + dispatcher's ``TransformerConfig`` (the turbo flags are CLI/args-level and + are not guaranteed to be mirrored onto ``self.config``). + ``use_turbo_grouped_gemm`` is the flag users set (also auto-enabled by + Sync-Free MoE stage >= 2). + """ + try: + args = get_args(ctx) + except Exception: + return False + return bool(getattr(args, "use_turbo_grouped_gemm", False)) + + +@register_patch( + "megatron.moe_alltoall_dtoh_turbo_grouped_gemm", + backend="megatron", + phase="before_train", + description=( + "Skip tokens_per_expert D2H copy in MoEAlltoAllTokenDispatcher " + "when PrimusTurbo grouped gemm (use_turbo_grouped_gemm) is enabled" + ), +) +def patch_moe_alltoall_dtoh(ctx: PatchContext): + """Replace ``MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize``.""" + from megatron.core.transformer.moe import token_dispatcher as td_mod + + cls = td_mod.MoEAlltoAllTokenDispatcher + + keep_tokens_per_expert_on_device = _turbo_grouped_gemm_on_device(ctx) + + def _maybe_dtoh_and_synchronize(self, point, tokens_per_expert=None): + """ + Move all possible GPU tensors to CPU and make a synchronization at the expected point. + """ + maybe_move_tensor_to_cpu = td_mod.maybe_move_tensor_to_cpu + + if not self.drop_and_pad: + if point == self.cuda_dtoh_point: + # Move all possible GPU tensors to CPU at self.cuda_dtoh_point. + on_side_stream = torch.cuda.current_stream() != self.cuda_dtoh_stream + if on_side_stream: + self.cuda_dtoh_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.cuda_dtoh_stream): + # TODO: use MemcpyBatchAsync instead. + # PrimusTurbo grouped gemm consumes tokens_per_expert on device, + # so keep it on the GPU and skip the D2H copy when enabled. + if not keep_tokens_per_expert_on_device: + tokens_per_expert = maybe_move_tensor_to_cpu( + tokens_per_expert, record_stream=on_side_stream + ) + self.input_splits = maybe_move_tensor_to_cpu( + self.input_splits, as_numpy=True, record_stream=on_side_stream + ) + self.output_splits = maybe_move_tensor_to_cpu( + self.output_splits, as_numpy=True, record_stream=on_side_stream + ) + self.output_splits_tp = maybe_move_tensor_to_cpu( + self.output_splits_tp, as_numpy=True, record_stream=on_side_stream + ) + self.num_out_tokens = maybe_move_tensor_to_cpu( + self.num_out_tokens, record_stream=on_side_stream + ) + if self.num_local_experts > 1 and not self.config.moe_permute_fusion: + self.num_global_tokens_per_local_expert = maybe_move_tensor_to_cpu( + self.num_global_tokens_per_local_expert, record_stream=on_side_stream + ) + self.d2h_event = self.cuda_dtoh_stream.record_event() + + if point == self.cuda_sync_point: + # Synchronize with the DtoH stream at self.cuda_sync_point. + self.d2h_event.synchronize() + + return tokens_per_expert + + cls._maybe_dtoh_and_synchronize = _maybe_dtoh_and_synchronize + + log_rank_0( + "[Patch:megatron.moe_alltoall_dtoh_turbo_grouped_gemm] Patched " + "MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize " + f"(skip tokens_per_expert D2H = {keep_tokens_per_expert_on_device})" + ) diff --git a/primus/backends/megatron/patches/triton_autotune_print_patches.py b/primus/backends/megatron/patches/triton_autotune_print_patches.py new file mode 100644 index 000000000..3a3dbc391 --- /dev/null +++ b/primus/backends/megatron/patches/triton_autotune_print_patches.py @@ -0,0 +1,66 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Silence Triton's autotuner stdout spam. + +Triton's ``Autotuner._bench`` (``triton/runtime/autotuner.py``) emits a bare +``print()`` for *every* (kernel, config) pair it benchmarks when the +``knobs.autotuning.print`` flag (env: ``TRITON_PRINT_AUTOTUNING``) is on:: + + Autotuning kernel _permute_kernel with config BLOCK_SIZE: 64, num_warps: 4, ... + +On a multi-node run, every autotuned Triton kernel (e.g. the MoE +permute / unpermute / sort kernels in +``primus/backends/transformer_engine/pytorch/triton/permutation.py``) produces a +whole batch of these lines during warmup, flooding the logs. + +Because it is a ``print()`` (not a ``logging`` record), a logging-level change +cannot suppress it -- the only lever is the Triton ``knobs.autotuning.print`` +flag. We force it off here in the ``before_train`` phase, which runs inside the +training worker after model setup but before the first training step (and thus +before any kernel is autotuned). Programmatic assignment to the knob shadows the +``TRITON_PRINT_AUTOTUNING`` env value, so this works regardless of how the flag +got turned on. + +Set ``PRIMUS_VERBOSE_TRITON=1`` to keep the full autotune trace. +""" + +import os + +from primus.core.patches import PatchContext, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +@register_patch( + "megatron.triton.silence_autotune_print", + backend="megatron", + phase="before_train", + description=( + "Force Triton knobs.autotuning.print off so the autotuner does not " + "flood stderr with one print line per kernel config it benchmarks." + ), + condition=lambda ctx: os.environ.get("PRIMUS_VERBOSE_TRITON", "0") != "1", +) +def patch_silence_triton_autotune_print(ctx: PatchContext): + """Turn off Triton autotuner printing. Idempotent / best-effort.""" + del ctx + try: + from triton import knobs + except Exception as exc: # triton missing -- nothing to do. + log_rank_0(f"[Patch:megatron.triton.silence_autotune_print] triton unavailable: {exc!r}") + return + + try: + knobs.autotuning.print = False + except Exception as exc: # knobs API changed -- never block training. + log_rank_0( + f"[Patch:megatron.triton.silence_autotune_print] could not set " + f"knobs.autotuning.print: {exc!r}" + ) + return + + log_rank_0("[Patch:megatron.triton.silence_autotune_print] Triton autotuner printing disabled.") diff --git a/primus/backends/megatron/training/tokenizer/tokenizer.py b/primus/backends/megatron/training/tokenizer/tokenizer.py index 5e97a987a..fc021dff3 100644 --- a/primus/backends/megatron/training/tokenizer/tokenizer.py +++ b/primus/backends/megatron/training/tokenizer/tokenizer.py @@ -31,6 +31,7 @@ CUSTOM_TOKENIZER_TYPES = { "DeepSeekV2Tokenizer", "DeepSeekV3Tokenizer", + "DeepSeekV4Tokenizer", "Llama2Tokenizer", "Llama3Tokenizer", "MixtralTokenizer", diff --git a/primus/configs/models/megatron/deepseek_v4_base.yaml b/primus/configs/models/megatron/deepseek_v4_base.yaml new file mode 100644 index 000000000..bd874d588 --- /dev/null +++ b/primus/configs/models/megatron/deepseek_v4_base.yaml @@ -0,0 +1,117 @@ +############################################################################### +# DeepSeek-V4 base config (defaults shared by V4-Flash / V4-Pro / etc.). +# +# Reference: +# - deepseek-v4/deepseek-ai/DeepSeek-V4-Flash/config.json +# - deepseek-v4/deepseek-ai/DeeSeek-v4-Pro/config.json +# - deepseek-v4/develop/techblog/01-deepseek-v4-architecture-deep-dive.md +############################################################################### + +extends: + - llama_base.yaml + +# Mark this model family as DeepSeek-V4. The Primus trainer +# (megatron_pretrain_trainer.py) and import_utils.get_model_provider use this +# string to dispatch to the V4-specific builder. +model_type: deepseek_v4 + +# norm +norm_epsilon: 1.0e-06 + +# attention base shape +# V4 uses a single shared latent for K = V (head_dim=512), num_key_value_heads=1. +# The Megatron MLA path (multi_latent_attention=true) is *NOT* reused here +# because V4 layers also need CSA / HCA / SWA branches; we will register a +# new V4 attention spec in Phase 3+. Keep MLA off so we don't accidentally go +# down the V3 path in Phase 1 / 2 stub builders. +multi_latent_attention: false +qk_layernorm: false +apply_rope_fusion: false # V4 uses partial RoPE on a 512-dim head; fused path doesn't fit. + +# ---------- Hyper-Connections (mHC) ---------- +hc_mult: 4 +hc_sinkhorn_iters: 20 +hc_eps: 1.0e-06 +hc_use_sinkhorn: true +mtp_use_separate_hc_head: true + +# ---------- Hybrid Attention ---------- +hybrid_attention_enabled: true +compress_ratios: null # provided by the per-variant yaml (Flash / Pro); + # entries: 0 = dense+SWA, 4 = CSA, 128 = HCA +compress_rope_theta: 160000.0 +index_topk: 512 +index_head_dim: 128 +index_n_heads: 64 +attn_sliding_window: 128 +attn_sink: true + +# DeepSeek-V4 attention backend selection (unified string selectors). The run +# script sets these via PRIMUS_USE_V4_ATTENTION_BACKEND / +# PRIMUS_USE_V4_CSA_ATTENTION_BACKEND (see run_deepseek_v4.sh). +# use_v4_attention_backend (dense cr=0 / HCA cr=128): eager|triton_v1|triton_v2|gluon +# use_v4_csa_attention_backend (CSA cr=4): eager|triton_v0|triton_v1|triton_v2|gluon|flydsl_v0 +# use_turbo_attention (when a core_attention module is built) still takes +# precedence for the dense path. +use_v4_attention_backend: triton_v1 +use_v4_csa_attention_backend: triton_v1 + +# FP8 (E4M3) Indexer QK path (CSA selector). Default off; the indexer QK +# scoring inputs are fake-quantized to FP8 when enabled (BF16 index-score / +# top-k path is preserved). Surface via PRIMUS_USE_V4_FP8_INDEXER. +use_v4_fp8_indexer: false + +# Plan-5 P29 (RESCOPED): wrap ``sinkhorn_normalize`` (the doubly-stochastic +# projection inside HyperMixer.compute_weights) with a cached +# ``torch.compile(fullgraph=True, dynamic=False)`` build. Collapses the +# 39 fp32 ``aten::sum`` reductions per call into a single Inductor-fused +# Triton kernel; AOT autograd handles BWD. Default ``false`` until G32 +# (FWD + BWD parity) and G33b (post-P29 trace) flip it on. Driven by +# the ``PRIMUS_USE_V4_COMPILED_SINKHORN`` env var via run_deepseek_v4.sh. +use_v4_compiled_sinkhorn: false + +# ---------- Output projection (grouped low-rank) ---------- +o_groups: 8 +o_lora_rank: 1024 + +# ---------- MoE ---------- +moe_layer_freq: 1 +moe_router_topk: 6 +moe_router_score_function: sqrtsoftplus # V4 default +moe_router_load_balancing_type: seq_aux_loss +moe_router_enable_expert_bias: true # noaux_tc +moe_router_bias_update_rate: 1.0e-3 +moe_router_topk_scaling_factor: 1.5 # routed_scaling_factor (Flash 1.5 / Pro 2.5) +moe_token_dispatcher_type: alltoall +moe_aux_loss_coeff: 0.001 +moe_router_dtype: null # fp32 recommended at training time +moe_router_pre_softmax: false +moe_grouped_gemm: true +moe_use_legacy_grouped_gemm: false +moe_permute_fusion: true +moe_shared_expert_overlap: true +num_hash_layers: 3 # First 3 MoE layers use hash routing + +# ---------- MTP ---------- +mtp_num_layers: 1 # num_nextn_predict_layers +mtp_loss_scaling_factor: 0.1 + +# ---------- RoPE (main path; compressed layers also use compress_rope_theta) ---------- +rotary_base: 10000 +rope_type: yarn +rotary_scaling_factor: 16.0 # YaRN factor (only enabled on compress layers in code) +original_max_position_embeddings: 65536 +mscale: 1.0 +mscale_all_dim: 1.0 +rotary_interleaved: true # released weights use interleaved RoPE + +# ---------- Activation ---------- +swiglu: true +swiglu_limit: 10.0 # clamped SwiGLU (FP4/FP8 stability) +activation_func_clamp_value: 10.0 # clamp swiglu parameter (Megatron config field) +bias_swiglu_fusion: false # fused path lacks clamp; revisit in perf phase +v4_grouped_experts_support_clamped_swiglu: true # grouped (TEGroupedMLP) backend honors the clamped SwiGLU above (training sets this too) + +# ---------- Parallelism defaults ---------- +expert_model_parallel_size: 1 +expert_tensor_parallel_size: null diff --git a/primus/configs/models/megatron/deepseek_v4_flash.yaml b/primus/configs/models/megatron/deepseek_v4_flash.yaml new file mode 100644 index 000000000..7dd86b623 --- /dev/null +++ b/primus/configs/models/megatron/deepseek_v4_flash.yaml @@ -0,0 +1,54 @@ +############################################################################### +# DeepSeek-V4 Flash (smaller / dense-attention variant). +# +# Source: deepseek-v4/deepseek-ai/DeepSeek-V4-Flash/config.json +############################################################################### + +extends: + - deepseek_v4_base.yaml + +# core shape +num_layers: 43 +hidden_size: 4096 +num_attention_heads: 64 +num_query_groups: 1 +kv_channels: 512 # head_dim +qk_pos_emb_head_dim: 64 # qk_rope_head_dim +ffn_hidden_size: 18432 # not used by MoE layers; placeholder for shared/dense FFN +moe_ffn_hidden_size: 2048 # moe_intermediate_size +moe_shared_expert_intermediate_size: 2048 # n_shared_experts * moe_intermediate_size + +# attention low-rank +q_lora_rank: 1024 +o_lora_rank: 1024 +o_groups: 8 + +# MoE +num_experts: 256 +moe_router_topk: 6 +moe_router_topk_scaling_factor: 1.5 + +# Indexer / Compressor +index_topk: 512 + +# Per-layer compression schedule (from config.json:compress_ratios) +# 0 = uncompressed dense layer (full attention with SWA) +# 4 = CSA branch (overlap mode; per-query top-K from compressed pool via Indexer) +# 128 = HCA branch (non-overlap mode; full visibility over compressed pool) +compress_ratios: "[0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]" + +# context +seq_length: 4096 # V4 supports up to 1M; default training seq is overridden by training yaml +max_position_embeddings: 1048576 + +# vocab +vocab_size: 129280 +make_vocab_size_divisible_by: 64 + +# tokenizer +tokenizer_type: DeepSeekV4Tokenizer +tokenizer_model: deepseek-ai/DeepSeek-V4-Flash +trust_remote_code: true + +# init +init_method_std: 0.02 diff --git a/primus/configs/models/megatron/deepseek_v4_pro.yaml b/primus/configs/models/megatron/deepseek_v4_pro.yaml new file mode 100644 index 000000000..73dc67c70 --- /dev/null +++ b/primus/configs/models/megatron/deepseek_v4_pro.yaml @@ -0,0 +1,46 @@ +############################################################################### +# DeepSeek-V4 Pro (large MoE variant). +# +# Source: deepseek-v4/deepseek-ai/DeeSeek-v4-Pro/config.json +############################################################################### + +extends: + - deepseek_v4_base.yaml + +num_layers: 61 +hidden_size: 7168 +num_attention_heads: 128 +num_query_groups: 1 +kv_channels: 512 +qk_pos_emb_head_dim: 64 +ffn_hidden_size: 18432 +moe_ffn_hidden_size: 3072 +moe_shared_expert_intermediate_size: 3072 + +q_lora_rank: 1536 +o_lora_rank: 1024 +o_groups: 16 + +num_experts: 384 +moe_router_topk: 6 +moe_router_topk_scaling_factor: 2.5 + +index_topk: 1024 + +# Per-layer compression schedule: +# 0 = uncompressed dense layer (full attention with SWA) +# 4 = CSA branch (overlap mode; per-query top-K from compressed pool via Indexer) +# 128 = HCA branch (non-overlap mode; full visibility over compressed pool) +compress_ratios: "[128, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]" + +seq_length: 4096 +max_position_embeddings: 1048576 + +vocab_size: 129280 +make_vocab_size_divisible_by: 64 + +tokenizer_type: DeepSeekV4Tokenizer +tokenizer_model: deepseek-ai/DeepSeek-V4-Pro +trust_remote_code: true + +init_method_std: 0.02 diff --git a/primus/configs/modules/megatron/primus_megatron_module.yaml b/primus/configs/modules/megatron/primus_megatron_module.yaml index 74f8d7179..182ae9317 100644 --- a/primus/configs/modules/megatron/primus_megatron_module.yaml +++ b/primus/configs/modules/megatron/primus_megatron_module.yaml @@ -70,7 +70,7 @@ mlflow_upload_performance_metrics: false disable_profiler_activity_cpu: false torch_profiler_record_shapes: true torch_profiler_with_stack: true -torch_profiler_use_gzip: false +torch_profiler_use_gzip: true # continue/finetune auto_continue_train: false diff --git a/primus/core/projection/module_profilers/attention.py b/primus/core/projection/module_profilers/attention.py index f227a00fa..32b5c9408 100644 --- a/primus/core/projection/module_profilers/attention.py +++ b/primus/core/projection/module_profilers/attention.py @@ -10,7 +10,7 @@ from primus.core.projection.base_module_profiler import BaseModuleProfiler -from .utils import benchmark_layer +from .utils import benchmark_layer, v4_module_inputs class AttentionProfiler(BaseModuleProfiler): @@ -333,13 +333,27 @@ def _get_benchmark_results(self, batch_size: int, seq_len: int) -> tuple[float, # Effective sequence length per rank if CP is used slen_per_cp = seq_len // cp_size - self._cached_results = benchmark_layer( - self.module, - [ - (seq_len, batch_size, self.config.model_config.hidden_size), - ((1, 1, slen_per_cp, seq_len), torch.bool), - ], - ) + hidden = self.config.model_config.hidden_size + tcfg = getattr(self.module, "config", None) + hc_mult = getattr(tcfg, "hc_mult", 1) + # DeepSeek-V4 attention has a different signature + # (forward(hidden[B,S,D], position_ids[B,S])); feed V4-aware + # inputs so the real V4 attention path is exercised instead of + # crashing / falling back. Non-V4 modules use the stock inputs. + v4 = v4_module_inputs(self.module, batch_size, seq_len, hidden, hc_mult, "attention") + if v4 is not None: + ishapes, fkwargs = v4 + self._cached_results = benchmark_layer( + self.module, ishapes, transformer_config=tcfg, forward_kwargs=fkwargs + ) + else: + self._cached_results = benchmark_layer( + self.module, + [ + (seq_len, batch_size, hidden), + ((1, 1, slen_per_cp, seq_len), torch.bool), + ], + ) self._cache_key = cache_key return self._cached_results diff --git a/primus/core/projection/module_profilers/moe_mlp.py b/primus/core/projection/module_profilers/moe_mlp.py index 2aa802446..fbe5ad0be 100644 --- a/primus/core/projection/module_profilers/moe_mlp.py +++ b/primus/core/projection/module_profilers/moe_mlp.py @@ -11,7 +11,7 @@ from primus.core.projection.profiler_spec import ModuleProfilerSpec from primus.core.projection.training_config import TrainingConfig -from .utils import benchmark_moe_layer_decomposed +from .utils import benchmark_layer, benchmark_moe_layer_decomposed, v4_module_inputs # Efficiency fractions for non-GEMM MoE overhead estimation. # These express achievable bandwidth as a fraction of peak HBM bandwidth. @@ -349,10 +349,25 @@ def _get_benchmark_results(self, batch_size: int, seq_len: int) -> tuple[float, self._a2a_fwd_ms = 0.0 self._a2a_bwd_ms = 0.0 else: - fwd, bwd, act_mem, a2a_fwd, a2a_bwd = benchmark_moe_layer_decomposed( - self.module, - [(seq_len, batch_size, self.config.model_config.hidden_size)], - ) + hidden = self.config.model_config.hidden_size + tcfg = getattr(self.module, "config", None) + # DeepSeek-V4 MoE: forward(hidden[B,S,D], *, token_ids[B,S]). + # Feed V4-aware inputs (right layout + token_ids for hash routing). + v4 = v4_module_inputs(self.module, batch_size, seq_len, hidden, 1, "moe") + if v4 is not None: + # DeepseekV4MoE has no stock .dispatch/.combine to decompose + # A2A; benchmark the whole MoE forward. At EP=1 (single-GPU + # benchmark) A2A is ~0 and is restored analytically later. + ishapes, fkwargs = v4 + fwd, bwd, act_mem = benchmark_layer( + self.module, ishapes, transformer_config=tcfg, forward_kwargs=fkwargs + ) + a2a_fwd = a2a_bwd = 0.0 + else: + fwd, bwd, act_mem, a2a_fwd, a2a_bwd = benchmark_moe_layer_decomposed( + self.module, + [(seq_len, batch_size, hidden)], + ) self._cached_results = (fwd, bwd, act_mem) self._a2a_fwd_ms = a2a_fwd self._a2a_bwd_ms = a2a_bwd diff --git a/primus/core/projection/module_profilers/transformer_layer.py b/primus/core/projection/module_profilers/transformer_layer.py index bb9476f3d..721924647 100644 --- a/primus/core/projection/module_profilers/transformer_layer.py +++ b/primus/core/projection/module_profilers/transformer_layer.py @@ -20,11 +20,7 @@ from .moe_mlp import MoEMLPProfiler from .residual_add import ResidualAddProfiler from .router import RouterProfiler -from .utils import ( - _install_balanced_routing_patches, - _kernel_pad_enabled, - benchmark_layer, -) +from .utils import benchmark_layer, v4_module_inputs # ── Fallback HBM bandwidth for elementwise overhead estimation ── _FALLBACK_HBM_BW_GBPS = 5300.0 # MI300X default @@ -321,11 +317,27 @@ def _get_benchmark_results(self, batch_size: int, seq_len: int) -> tuple[float, else: # Get TransformerConfig from the layer module itself (has fp8 setting) transformer_config = getattr(self.layer_module, "config", None) - self._cached_results = benchmark_layer( - self.layer_module, - [(seq_len, batch_size, self.config.model_config.hidden_size)], - transformer_config=transformer_config, - ) + hidden = self.config.model_config.hidden_size + hc_mult = getattr(transformer_config, "hc_mult", 1) + # DeepSeek-V4 hybrid layer needs K-stream input [B,S,K,D] and a + # keyword position_ids; feed V4-aware inputs so the real V4 layer + # (mHC + V4 attention) is benchmarked. Non-V4 layers use the stock + # [S,B,D] input. + v4 = v4_module_inputs(self.layer_module, batch_size, seq_len, hidden, hc_mult, "layer") + if v4 is not None: + ishapes, fkwargs = v4 + self._cached_results = benchmark_layer( + self.layer_module, + ishapes, + transformer_config=transformer_config, + forward_kwargs=fkwargs, + ) + else: + self._cached_results = benchmark_layer( + self.layer_module, + [(seq_len, batch_size, hidden)], + transformer_config=transformer_config, + ) self._cache_key = cache_key return self._cached_results @@ -475,24 +487,27 @@ def _get_benchmark_results(self, batch_size: int, seq_len: int) -> tuple[float, ): # Legacy whole-layer timing (often ~1.5-1.7x pessimistic on backward). transformer_config = getattr(self.layer_module, "config", None) - routing_restores = [] - if _kernel_pad_enabled(): - routing_restores, _ = _install_balanced_routing_patches(self.layer_module) - try: + hidden = self.config.model_config.hidden_size + hc_mult = getattr(transformer_config, "hc_mult", 1) + # DeepSeek-V4 hybrid layer needs K-stream input [B,S,K,D] and a + # keyword position_ids; feed V4-aware inputs so the real V4 layer + # (mHC + V4 attention) is benchmarked. Non-V4 layers use the stock + # [S,B,D] input. + v4 = v4_module_inputs(self.layer_module, batch_size, seq_len, hidden, hc_mult, "layer") + if v4 is not None: + ishapes, fkwargs = v4 self._cached_results = benchmark_layer( self.layer_module, - [(seq_len, batch_size, self.config.model_config.hidden_size)], + ishapes, + transformer_config=transformer_config, + forward_kwargs=fkwargs, + ) + else: + self._cached_results = benchmark_layer( + self.layer_module, + [(seq_len, batch_size, hidden)], transformer_config=transformer_config, ) - finally: - for restore in routing_restores: - try: - restore() - except Exception: - # Best-effort cleanup: restore failures should not fail benchmarking. - _LOGGER.debug("Failed to restore routing patch during cleanup.", exc_info=True) - else: - self._cached_results = self._get_benchmark_composite_results(batch_size, seq_len) self._cache_key = cache_key return self._cached_results diff --git a/primus/core/projection/module_profilers/utils.py b/primus/core/projection/module_profilers/utils.py index 504c7d890..f0a98a6a5 100644 --- a/primus/core/projection/module_profilers/utils.py +++ b/primus/core/projection/module_profilers/utils.py @@ -97,11 +97,53 @@ def _get_fp8_context_for_benchmark(transformer_config): return _FP8ContextFactory(transformer_config) +def v4_module_inputs(module, batch_size, seq_len, hidden_size, hc_mult, kind): + """Build DeepSeek-V4-aware benchmark inputs, or None for non-V4 modules. + + The generic harness feeds stock-attention inputs ``(hidden[S,B,D], + bool attention_mask)``, but the V4 modules have different signatures: + + * ``DeepseekV4Attention.forward(hidden[B,S,D], position_ids[B,S])`` — + both positional; needs integer position_ids (not a bool mask). + * ``DeepseekV4HybridLayer.forward(hidden[B,S,K,D], attention_mask=None, + *, position_ids[B,S])`` — K=hc_mult parallel mHC streams, and + position_ids is keyword-only. + + Returns ``(input_shapes, forward_kwargs)`` or ``None``. + """ + cls = type(module).__name__ + if kind == "attention" and "DeepseekV4Attention" in cls: + return ( + [(batch_size, seq_len, hidden_size), ((batch_size, seq_len), torch.int64)], + None, + ) + if kind == "layer" and "DeepseekV4HybridLayer" in cls: + k = max(1, int(hc_mult or 1)) + ishapes = [(batch_size, seq_len, k, hidden_size)] if k > 1 else [(batch_size, seq_len, hidden_size)] + # position_ids for RoPE; token_ids required by hash-routed MoE layers + # (ignored by non-hash layers). Both keyword-only on the V4 layer forward. + return ( + ishapes, + { + "position_ids": ((batch_size, seq_len), torch.int64), + "token_ids": ((batch_size, seq_len), torch.int64), + }, + ) + if kind == "moe" and "DeepseekV4MoE" in cls: + # forward(hidden[B,S,D], *, token_ids[B,S]) -> [B,S,D] + return ( + [(batch_size, seq_len, hidden_size)], + {"token_ids": ((batch_size, seq_len), torch.int64)}, + ) + return None + + def benchmark_layer( layer_module: torch.nn.Module, input_shapes: List[Union[Tuple[int, ...], Tuple[Tuple[int, ...], torch.dtype]]], num_iterations: int = 64, # Match typical microbatch count transformer_config=None, # Optional: pass config to enable FP8 context + forward_kwargs=None, # Optional: dict name -> shape/(shape,dtype) passed as keywords ) -> tuple[float, float, int]: """ Benchmark both forward and backward passes of a transformer layer using CUDA events. @@ -152,6 +194,7 @@ def create_input(spec): ) inputs = [create_input(spec) for spec in input_shapes] + kwargs = {name: create_input(spec) for name, spec in (forward_kwargs or {}).items()} # =========================================================================== # Get FP8 context - CRITICAL for accurate FP8 timing! @@ -168,7 +211,7 @@ def create_input(spec): with fp8_context: for _ in range(num_warmup): - outputs = layer_module(*inputs) + outputs = layer_module(*inputs, **kwargs) if not isinstance(outputs, (tuple, list)): outputs = (outputs,) @@ -213,7 +256,7 @@ def create_input(spec): with fp8_context: for _ in range(num_iterations): - outputs = layer_module(*inputs) + outputs = layer_module(*inputs, **kwargs) if device.type == "cuda": torch.cuda.synchronize(device) @@ -237,7 +280,7 @@ def create_input(spec): forward_end = torch.cuda.Event(enable_timing=True) forward_start.record() - outputs = layer_module(*inputs) + outputs = layer_module(*inputs, **kwargs) forward_end.record() # --- Backward pass --- @@ -391,6 +434,7 @@ def benchmark_moe_layer_decomposed( input_shapes: List[Union[Tuple[int, ...], Tuple[Tuple[int, ...], torch.dtype]]], num_iterations: int = 64, transformer_config=None, + forward_kwargs=None, # Optional: dict name -> shape/(shape,dtype) passed as keywords ) -> tuple[float, float, int, float, float]: """ Benchmark an MoE layer with decomposed A2A timing. @@ -453,47 +497,92 @@ def create_input(spec): ) inputs = [create_input(spec) for spec in input_shapes] + kwargs = {name: create_input(spec) for name, spec in (forward_kwargs or {}).items()} fp8_context = _get_fp8_context_for_benchmark(transformer_config) is_rank_0 = int(os.getenv("RANK", "0")) == 0 - # Install kernel-shape-padding (balanced round-robin routing) so the - # per-expert M dimension is constant across iterations. This eliminates - # per-shape JIT recompiles in the FP8 grouped GEMM and is the dominant - # cost on high-expert MoE benches. - routing_restores = [] - if _kernel_pad_enabled(): - routing_restores, routing_descriptors = _install_balanced_routing_patches(moe_module) - if routing_descriptors and is_rank_0: - _, topk, num_experts = routing_descriptors[0] - # The bench feeds a single [seq_len, batch_size, hidden] input, - # so the per-rank token count is the product of the first two - # dims (TP/SP sharding is already reflected in the supplied - # shape). - num_tokens = 1 - for spec in input_shapes: - shape = ( - spec[0] - if isinstance(spec, tuple) and len(spec) == 2 and isinstance(spec[1], torch.dtype) - else spec - ) - # [seq, batch, hidden] - if len(shape) >= 2: - num_tokens = int(shape[0]) * int(shape[1]) - break - m_per_expert = (num_tokens * topk) // max(num_experts, 1) - print( - f" [MoE Decomposed] Kernel-shape padding ON: balanced " - f"round-robin routing on {len(routing_descriptors)} router(s) " - f"(topk={topk}, num_experts={num_experts}, " - f"M_per_expert={m_per_expert}). " - f"Set PRIMUS_BENCH_MOE_KERNEL_PAD=0 to disable." - ) - elif is_rank_0: - print( - " [MoE Decomposed] Kernel-shape padding requested but no " - "router-like submodule found; falling back to stochastic routing." - ) + with fp8_context: + for _ in range(num_warmup): + outputs = moe_module(*inputs, **kwargs) + if not isinstance(outputs, (tuple, list)): + outputs = (outputs,) + + if grad_outputs is None: + grad_outputs = [] + for i, out in enumerate(outputs): + if isinstance(out, torch.Tensor) and out.requires_grad: + grad_outputs.append(torch.randn_like(out)) + output_indices.append(i) + + valid_outputs = [outputs[i] for i in output_indices] + if valid_outputs: + torch.autograd.backward(valid_outputs, grad_outputs) + + moe_module.zero_grad(set_to_none=True) + for inp in inputs: + if inp.requires_grad: + inp.grad = None + + if device.type == "cuda": + torch.cuda.synchronize(device) + + # --- Measure activation memory (forward-only loop) --- + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats(device) + if device.type == "cuda": + torch.cuda.synchronize(device) + mem_before = torch.cuda.memory_allocated(device) + + with fp8_context: + for _ in range(num_iterations): + outputs = moe_module(*inputs, **kwargs) + + if device.type == "cuda": + torch.cuda.synchronize(device) + mem_after_forward = torch.cuda.max_memory_allocated(device) + activation_memory = (mem_after_forward - mem_before) // num_iterations + + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats(device) + del outputs + + # ========================================================================= + # BENCHMARK with decomposed A2A timing + # ========================================================================= + # Monkey-patch dispatch() and combine() to insert CUDA events. + # MoELayer.forward() calls self.dispatch(...) and self.combine(...) + # so instance-attribute patches are picked up by Python's MRO. + original_dispatch = moe_module.dispatch + original_combine = moe_module.combine + + # Accumulate (start_event, end_event) pairs per iteration + _dispatch_events = [] # one (start, end) per iteration + _combine_events = [] + + def timed_dispatch(*args, **kwargs): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = original_dispatch(*args, **kwargs) + end.record() + _dispatch_events.append((start, end)) + return result + + def timed_combine(*args, **kwargs): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = original_combine(*args, **kwargs) + end.record() + _combine_events.append((start, end)) + return result + + moe_module.dispatch = timed_dispatch + moe_module.combine = timed_combine + + forward_times = [] + backward_times = [] try: # ===================================================================== @@ -504,8 +593,19 @@ def create_input(spec): num_warmup, num_iterations = _bench_iter_count(20, num_iterations) with fp8_context: - for _ in range(num_warmup): - outputs = moe_module(*inputs) + for _ in range(num_iterations): + # --- Forward pass --- + forward_start = torch.cuda.Event(enable_timing=True) + forward_end = torch.cuda.Event(enable_timing=True) + + forward_start.record() + outputs = moe_module(*inputs, **kwargs) + forward_end.record() + + # --- Backward pass --- + backward_start = torch.cuda.Event(enable_timing=True) + backward_end = torch.cuda.Event(enable_timing=True) + if not isinstance(outputs, (tuple, list)): outputs = (outputs,) diff --git a/primus/core/projection/performance_projection/projection.py b/primus/core/projection/performance_projection/projection.py index fdcac1433..28d97b31f 100644 --- a/primus/core/projection/performance_projection/projection.py +++ b/primus/core/projection/performance_projection/projection.py @@ -1056,16 +1056,23 @@ def _limit_layers_for_projection(module_config): original_layers = getattr(module_config, "num_layers", 1) or 1 original_moe_layout = getattr(module_config, "moe_layer_freq", None) dense_layers_present = _has_dense_layers(original_moe_layout) - - if has_moe and dense_layers_present: - # Need at least 2 layers to profile both dense (layer 0) and MoE (layer 1) - # so extraction code can correctly classify each type using the full - # model's moe_pattern where layer 0 is typically dense. - max_layers = 2 - else: - max_layers = 1 + # Use 1 layer for fast profiling - results are extrapolated to full model + # Increase to 2-4 for better accuracy if needed. PRIMUS_PROJ_MAX_LAYERS lets + # the caller benchmark more representative layers (e.g. =2 to capture both the + # DeepSeek-V4 HCA(cr=128) and CSA(cr=4) attention kinds, which alternate). + max_layers = int(os.environ.get("PRIMUS_PROJ_MAX_LAYERS", "1")) target_layers = max(1, min(original_layers, max_layers)) module_config.num_layers = target_layers + # Set the benchmarked layers' attention kinds. PRIMUS_PROJ_COMPRESS_RATIOS + # (e.g. "128,4") picks exactly one HCA + one CSA; otherwise trim the model's + # own schedule to the benchmarked layer count. + _cr_env = os.environ.get("PRIMUS_PROJ_COMPRESS_RATIOS", "").strip("[] ") + if _cr_env and hasattr(module_config, "compress_ratios"): + module_config.compress_ratios = [int(x) for x in _cr_env.split(",")][:target_layers] + else: + _cr = getattr(module_config, "compress_ratios", None) + if isinstance(_cr, (list, tuple)) and len(_cr) >= target_layers: + module_config.compress_ratios = list(_cr[:target_layers]) if has_moe: if not dense_layers_present: diff --git a/primus/core/utils/import_utils.py b/primus/core/utils/import_utils.py index 1fd93c408..47b1ed918 100644 --- a/primus/core/utils/import_utils.py +++ b/primus/core/utils/import_utils.py @@ -39,13 +39,27 @@ def get_model_provider(model_type="gpt"): Resolve model_provider across Megatron versions and model types. Args: - model_type (str): Type of model - 'gpt' or 'mamba'. Defaults to 'gpt'. + model_type (str): Type of model - 'gpt', 'mamba', or 'deepseek_v4'. Defaults to 'gpt'. - New: model_provider + gpt_builder/mamba_builder - Mid: model_provider only - Old: pretrain_gpt.model_provider / pretrain_mamba.model_provider """ - # Try to import model_provider + # Primus-owned: DeepSeek-V4 (Phase 2 stub; full V4 wiring lands in Phase 3+) + if model_type == "deepseek_v4": + deepseek_v4_module = importlib.import_module( + "primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders" + ) + log_rank_0( + "[Primus][MegatronCompat] Loaded DeepSeek-V4 model_provider + builder " + f"from {deepseek_v4_module.__name__}" + ) + return partial( + deepseek_v4_module.model_provider, + deepseek_v4_module.deepseek_v4_builder, + ) + + # Upstream Megatron-LM model types if model_type == "mamba": model_provider = lazy_import( ["model_provider", "pretrain_mamba"], "model_provider", log_prefix="[Primus][MegatronCompat]" diff --git a/rccl_avg_workaround/.gitignore b/rccl_avg_workaround/.gitignore new file mode 100644 index 000000000..7a60b85e1 --- /dev/null +++ b/rccl_avg_workaround/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +*.pyc diff --git a/rccl_avg_workaround/sitecustomize.py b/rccl_avg_workaround/sitecustomize.py new file mode 100644 index 000000000..bf10ac87e --- /dev/null +++ b/rccl_avg_workaround/sitecustomize.py @@ -0,0 +1,143 @@ +"""gfx1250 single-GPU bring-up workarounds, auto-imported in every Python +worker via sitecustomize (this dir is on PYTHONPATH). Two independent fixes: + +1. RCCL AVG hang: torch.distributed.all_reduce(op=AVG) HANGS on this build for + (at least) single-rank process groups, while SUM works fine (verified by + collective microbench). Megatron's MoE aux-loss metric reduction + (moe_utils.reduce_aux_losses_tracker_across_ranks) uses op=AVG and + deadlocks. Replace AVG with SUM + divide-by-world-size, which is + mathematically identical for any world size. + +2. primus_turbo import shim: the MI355X production containers bundle the + `primus_turbo` package; the gfx1250 therock container does not. Most Primus + call-sites guard the import (try/except -> HAVE_TURBO=False), but the V4 + model path imports it unconditionally: + deepseek_v4_layer_specs.py -> transformer_engine_spec_provider.py + (DeepSeekV4SpecProvider subclasses PrimusTurboSpecProvider) + -> extensions/primus_turbo.py -> `import primus_turbo.pytorch` + and backends/megatron/core/utils.py -> primus_turbo...attention_utils. + With every use_turbo_* flag False the turbo classes are never SELECTED + (the spec provider returns the TE classes), so a pure import-shim is safe: + install a meta-path finder that fabricates stub modules for primus_turbo.* + whose attributes are auto-generated dummy classes. Attribute chains + evaluated at class-definition time (e.g. the ScalingGranularity.TENSORWISE + default arg in PrimusTurboQuantConfig) resolve fine; actually CALLING or + instantiating any stub raises RuntimeError, so a misrouted turbo path fails + loudly instead of computing garbage. The shim only installs when the real + package is absent, so it can never shadow a real primus_turbo install. +""" + +import sys + + +def _install_rccl_avg_workaround(): + import torch.distributed as dist + + orig_all_reduce = dist.all_reduce + avg_op = dist.ReduceOp.AVG + + def all_reduce_avg_safe(tensor, op=dist.ReduceOp.SUM, group=None, async_op=False): + is_avg = False + try: + is_avg = op == avg_op + except Exception: + is_avg = str(op) == str(avg_op) + if is_avg: + work = orig_all_reduce(tensor, op=dist.ReduceOp.SUM, group=group, async_op=async_op) + try: + ws = dist.get_world_size(group) + except Exception: + ws = 1 + if ws and ws > 1: + # Enqueued on the same stream after the all_reduce, so ordering holds. + tensor.div_(ws) + return work + return orig_all_reduce(tensor, op=op, group=group, async_op=async_op) + + dist.all_reduce = all_reduce_avg_safe + print( + "[rccl_avg_workaround] patched torch.distributed.all_reduce (AVG -> SUM/ws)", + file=sys.stderr, + flush=True, + ) + + +def _install_primus_turbo_stub(): + import importlib.abc + import importlib.machinery + import importlib.util + import types + + # Never shadow a real install. + if importlib.util.find_spec("primus_turbo") is not None: + return + + class _StubMeta(type): + """Dummy-class metaclass: any attribute access mints another dummy + class (covers enum-style chains like ScalingGranularity.TENSORWISE).""" + + def __getattr__(cls, name): + if name.startswith("__"): + raise AttributeError(name) + dummy = _make_dummy(f"{cls._stub_qual}.{name}") + setattr(cls, name, dummy) + return dummy + + def _make_dummy(qual): + def _raise(self, *args, **kwargs): + raise RuntimeError( + f"primus_turbo stub: {qual} was invoked, but primus_turbo is NOT " + "installed in this container. A turbo code path ran despite all " + "use_turbo_*/enable_primus_turbo flags being False — fix the flags " + "instead of installing primus_turbo." + ) + + return _StubMeta(qual.rsplit(".", 1)[-1], (), {"__init__": _raise, "_stub_qual": qual}) + + class _StubModule(types.ModuleType): + def __getattr__(self, name): + if name.startswith("__"): + raise AttributeError(name) + dummy = _make_dummy(f"{self.__name__}.{name}") + setattr(self, name, dummy) + return dummy + + class _Finder(importlib.abc.MetaPathFinder, importlib.abc.Loader): + def find_spec(self, fullname, path=None, target=None): + if fullname == "primus_turbo" or fullname.startswith("primus_turbo."): + return importlib.machinery.ModuleSpec(fullname, self, is_package=True) + return None + + def create_module(self, spec): + return _StubModule(spec.name) + + def exec_module(self, module): + module.__path__ = [] + module._primus_turbo_stub = True + + sys.meta_path.append(_Finder()) + print( + "[primus_turbo_stub] primus_turbo not installed -> import shim active " + "(turbo classes import as raising stubs; all use_turbo_* must stay False)", + file=sys.stderr, + flush=True, + ) + + +# Only patch the actual training worker. Importing torch here is heavy, so we +# must NOT do it for pip/offload-arch/build helper invocations (they stall and +# block setup). Gate on the training entrypoint appearing in argv. +def _is_training_worker(): + argv = " ".join(sys.argv) + return ("primus/cli/main.py" in argv) or ("run_pretrain" in argv) or ("pretrain" in argv) + + +if _is_training_worker(): + try: + _install_rccl_avg_workaround() + except Exception as e: # noqa: BLE001 + print(f"[rccl_avg_workaround] install FAILED: {e}", file=sys.stderr, flush=True) + try: + _install_primus_turbo_stub() + except Exception as e: # noqa: BLE001 + print(f"[primus_turbo_stub] install FAILED: {e}", file=sys.stderr, flush=True) diff --git a/runner/helpers/hooks/train/pretrain/megatron/01_install_emerging_optimizers.sh b/runner/helpers/hooks/train/pretrain/megatron/01_install_emerging_optimizers.sh new file mode 100644 index 000000000..80f788344 --- /dev/null +++ b/runner/helpers/hooks/train/pretrain/megatron/01_install_emerging_optimizers.sh @@ -0,0 +1,58 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# Install NVIDIA-NeMo/Emerging-Optimizers inside the training container. +# +# The Muon optimizer path (primus/backends/megatron/core/optimizer/moun.py) +# hard-requires the ``emerging_optimizers`` package (Newton-Schulz orthogonal- +# ization, including the DeepSeek-V4 hybrid coefficient set). The package is +# NOT bundled in the default Primus container, and the public PyPI name +# ``emerging-optimizers`` is a placeholder stub (metadata-generation-failed) -- +# the real package only installs from the GitHub source. This hook installs it +# from a pinned commit so a Muon run works out of the box. +# +# Gated by PRIMUS_INSTALL_EMERGING_OPTIMIZERS (default off) so non-Muon runs +# pay nothing; idempotent (skips when already importable). Pure-python wheel +# (~7s to build/install). run_deepseek_v4.sh sets the gate when OPTIMIZER=muon. +############################################################################### +set -euo pipefail + +# Only do work when explicitly requested (the Muon launch path sets this). +case "${PRIMUS_INSTALL_EMERGING_OPTIMIZERS:-0}" in + 1 | true | True | TRUE | yes | on) + ;; + *) + echo "[install_eo] PRIMUS_INSTALL_EMERGING_OPTIMIZERS not set; skipping." + exit 0 + ;; +esac + +# Pinned to match the third_party/Emerging-Optimizers submodule and the +# Megatron-LM muon.py integration (which passes ``use_nesterov`` to +# OrthogonalizedOptimizer). The older 06ff4c68 pin used the pre-rename +# ``nesterov`` kwarg and is incompatible with the current Megatron muon.py. +EO_COMMIT="${PRIMUS_EMERGING_OPTIMIZERS_COMMIT:-93d9eb3a6c899b50de73992826451fba3ab6adfb}" +EO_URL="git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@${EO_COMMIT}" + +PY="${PYTHON:-python3}" + +if "${PY}" -c "import emerging_optimizers" >/dev/null 2>&1; then + echo "[install_eo] emerging_optimizers already importable; skipping install." + exit 0 +fi + +echo "[install_eo] Installing emerging_optimizers from ${EO_URL} ..." +# --no-deps: the package's only hard runtime dep is torch (already in the +# container); avoid pulling an incompatible torch/transitive set. +"${PY}" -m pip install --no-cache-dir --no-deps "${EO_URL}" + +if "${PY}" -c "import emerging_optimizers as e; print('[install_eo] installed', getattr(e,'__version__','?'))"; then + echo "[install_eo] OK" +else + echo "[install_eo] ERROR: emerging_optimizers still not importable after install" >&2 + exit 1 +fi diff --git a/tests/unit_tests/backends/megatron/test_compressor_pool.py b/tests/unit_tests/backends/megatron/test_compressor_pool.py new file mode 100644 index 000000000..98682b09c --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_compressor_pool.py @@ -0,0 +1,58 @@ +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +"""Parity test for the fused Compressor softmax-pool kernel. + +Asserts the fused ``fused_softmax_weighted_pool`` (Triton fwd + analytic eager bwd) +matches the eager ``(softmax(score+ape, dim=2) * kv).sum(dim=2)`` it replaces in +:meth:`Compressor.forward`, for both the forward output and the dkv / dscore / dape +gradients. GPU-gated (the kernel is CUDA/Triton-only); fp32 is checked tightly and +bf16 loosely (the fused path reduces in fp32 -> more accurate than eager bf16 weights). +""" +from __future__ import annotations + +import pytest +import torch + +cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="fused compressor pool is CUDA/Triton only") + +try: + from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.compressor_pool import ( + fused_softmax_weighted_pool, + ) + + HAVE_KERNEL = True +except Exception: # triton import may fail on a CPU-only host + HAVE_KERNEL = False + +pytestmark = [cuda, pytest.mark.skipif(not HAVE_KERNEL, reason="compressor_pool/triton unavailable")] + + +def _eager(kv, score, ape): + w = torch.softmax((score + ape).float(), dim=2).to(kv.dtype) + return (kv * w).sum(dim=2) + + +def _rel(a, b): + return (a - b).float().norm() / b.float().norm().clamp_min(1e-12) + + +# (B, N, W, hd): HCA (W=128, no overlap) and CSA (W=8, overlap) +@pytest.mark.parametrize("shape", [(2, 16, 128, 128), (2, 64, 8, 128)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_fused_pool_matches_eager(shape, dtype): + B, N, W, HD = shape + torch.manual_seed(0) + dev = "cuda" + kv = torch.randn(B, N, W, HD, device=dev, dtype=dtype) + sc = torch.randn(B, N, W, HD, device=dev, dtype=dtype) + ape = torch.randn(W, HD, device=dev, dtype=dtype) + g = torch.randn(B, N, HD, device=dev, dtype=dtype) + + ke, se, ae = (t.clone().requires_grad_(True) for t in (kv, sc, ape)) + _eager(ke, se, ae).backward(g) + kf, sf, af = (t.clone().requires_grad_(True) for t in (kv, sc, ape)) + fused_softmax_weighted_pool(kf, sf, af).backward(g) + + tol = 1e-4 if dtype == torch.float32 else 2e-2 + assert _rel(kf.grad, ke.grad) < tol + assert _rel(sf.grad, se.grad) < tol + assert _rel(af.grad, ae.grad) < tol diff --git a/tests/unit_tests/backends/megatron/test_deepseek_v4_flops_patches.py b/tests/unit_tests/backends/megatron/test_deepseek_v4_flops_patches.py new file mode 100644 index 000000000..ddfdb8912 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_deepseek_v4_flops_patches.py @@ -0,0 +1,647 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for ``deepseek_v4_flops_patches.py`` (Plan-3 Phase 20 + Plan-6 P33). + +Coverage: + +* G16 — :func:`compute_v4_flops` matches a hand-derived closed-form total + within 1% on a fully-specified V4-Flash-shaped config (8 layers, mixed + ``compress_ratios=[0,4,128]``, ``mtp_num_layers=1``). Every per-component + byte (``attn_qkv_o``, ``attn_scores``, ``compressor``, ``indexer``, + ``moe``, ``mtp``, ``logits``, ``hc``) is asserted independently so a + regression in any single term fails loudly. +* G17 — When the wrapper is installed with ``dispatch_v4=False`` (i.e. the + installer saw a non-V4 ``args.model_type``), the wrapper returns the + upstream value byte-for-byte. This mirrors how the install-time + ``condition`` gate works: V4 dispatch is captured at install time + because Megatron's ``pretrain()`` overwrites ``args.model_type`` with + a ``ModelType`` enum at ``training.py:1210`` before ``train()`` calls + ``num_floating_point_operations``. +* G36 — Plan-6 P33: SWA visible-pair correction. Parametrised over + ``swa_window``, ``compress_ratio``, ``hc_mult`` so the dense + HCA + + CSA per-layer pair counts and the over-count ratio vs the legacy + ``S_eff^2`` upper bound are pinned independently. +* G36a — Plan-6 P33: HyperConnection ``fn.weight`` matmul accounting. + Asserts the ``hc`` breakdown row equals the closed form + ``B * S * K * D * K * (2 * (L + M) * (2+K) + (1 + M))`` and degrades + to 0 when ``hc_mult <= 1``. +""" + +from __future__ import annotations + +import types +from types import SimpleNamespace + +import pytest + +from primus.backends.megatron.patches.deepseek_v4_flops_patches import ( + _FMA_FACTOR, + _FORWARD_BACKWARD_FACTOR, + _SWIGLU_FFN_EXPANSION_FACTOR, + _make_v4_num_floating_point_operations, + _normalize_layer_ratios, + _visible_pairs, + compute_v4_flops, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _v4_flash_smoke_args( + *, + num_layers: int = 8, + seq_length: int = 128, + hc_mult: int = 4, + mtp_num_layers: int = 0, + compress_ratios=(0, 0, 4, 128, 4, 128, 4, 0), + num_hash_layers: int = 3, + moe_router_topk: int = 6, + num_experts: int = 256, + moe_ffn_hidden_size: int = 2048, + moe_shared_expert_intermediate_size: int = 2048, + attn_sliding_window: int = 0, +): + """Build a V4-Flash-shaped fake ``args`` namespace. + + Defaults mirror the smoke config used in P19 (run_deepseek_v4.sh) + so the unit numbers stay comparable to live runs. ``attn_sliding_window`` + defaults to ``0`` (disabled / full causal) for backward-compatibility + with the plan-3 P20 reference; SWA-aware behaviour is exercised by + the plan-6 P33 G36 tests below. + """ + return SimpleNamespace( + model_type="deepseek_v4", + seq_length=seq_length, + hc_mult=hc_mult, + hidden_size=4096, + num_attention_heads=64, + kv_channels=512, + q_lora_rank=1024, + o_lora_rank=1024, + o_groups=8, + num_layers=num_layers, + mtp_num_layers=mtp_num_layers, + compress_ratios=compress_ratios, + moe_ffn_hidden_size=moe_ffn_hidden_size, + ffn_hidden_size=18432, + moe_router_topk=moe_router_topk, + num_experts=num_experts, + moe_shared_expert_intermediate_size=moe_shared_expert_intermediate_size, + num_hash_layers=num_hash_layers, + index_topk=512, + index_head_dim=128, + index_n_heads=64, + padded_vocab_size=129280, + vocab_size=129280, + attn_sliding_window=attn_sliding_window, + ) + + +def _hand_attn_qkv_o(*, B, S_eff, H, n, d, q_lora, o_lora, o_groups): + """Reference closed-form QKV+O FMAC per layer (matches the patch helper).""" + n_d = n * d + qkv = H * q_lora + q_lora * n_d + H * d + o_proj = n_d * o_lora + (o_groups * o_lora) * H if o_lora > 0 else n_d * H + return B * S_eff * (qkv + o_proj) + + +def _hand_local_pairs(*, swa, S_eff): + """Reference closed form for SWA-pruned local visible pairs.""" + if swa <= 0 or swa >= S_eff: + return S_eff * (S_eff + 1) // 2 + return swa * S_eff - swa * (swa - 1) // 2 + + +def _hand_pool_pairs(*, ratio, S_eff): + """Reference closed form for the causal-visible HCA pool pair count.""" + c = int(ratio) + if c <= 0 or S_eff <= 0: + return 0 + n_full = S_eff // c + if n_full == 0: + return 0 + return c * n_full * (n_full - 1) // 2 + n_full * (S_eff - c * n_full + 1) + + +def _hand_attn_scores(*, B, S_eff, n, d, ratio, index_topk, swa): + """Reference closed-form attention-score FMAC per layer. + + Plan-6 P33: counts only causal-visible ``(query, key)`` pairs + surviving the per-layer mask (SWA + pool + sparse top-K), not the + legacy ``S_eff^2`` upper bound. + """ + local_pairs = _hand_local_pairs(swa=swa, S_eff=S_eff) + if ratio == 0: + pairs = local_pairs + elif ratio == 128: + pairs = local_pairs + _hand_pool_pairs(ratio=ratio, S_eff=S_eff) + elif ratio == 4: + pool = max(1, S_eff // 4) + keys = min(index_topk, pool) if index_topk else pool + pairs = local_pairs + keys * S_eff + else: + pool = max(1, S_eff // ratio) + pairs = local_pairs + pool * S_eff + return 2 * B * n * d * pairs + + +def _hand_compressor(*, B, S_eff, H, d, ratio): + if ratio == 0: + return 0 + coff = 2 if ratio == 4 else 1 + return 2 * B * S_eff * H * (coff * d) + + +def _hand_indexer(*, B, S_eff, H, ratio, ihd, inh): + if ratio != 4: + return 0 + pool = max(1, S_eff // ratio) + proj = H * ihd + ihd * (inh * ihd) + H * inh + 2 * H * (2 * ihd) + scoring = inh * pool * ihd + return B * S_eff * (proj + scoring) + + +def _hand_moe(*, B, S_eff, H, H_moe, topk, n_experts, hash_layer, H_shared): + router = 0 if hash_layer else H * n_experts + routed = topk * _SWIGLU_FFN_EXPANSION_FACTOR * H * H_moe + shared = _SWIGLU_FFN_EXPANSION_FACTOR * H * H_shared if H_shared > 0 else 0 + return B * S_eff * (router + routed + shared) + + +# --------------------------------------------------------------------------- +# G16: closed-form parity per component +# --------------------------------------------------------------------------- + + +class TestComputeV4FlopsClosedForm: + """G16: per-component breakdown matches a hand-derived reference.""" + + @pytest.fixture(scope="class") + def args(self): + return _v4_flash_smoke_args() + + @pytest.fixture(scope="class") + def batch_size(self): + return 16 # GBS used in the P19 smoke run. + + @pytest.fixture(scope="class") + def computed(self, args, batch_size): + total, breakdown = compute_v4_flops(args, batch_size) + return total, breakdown + + def test_attn_qkv_o_term_matches_reference(self, args, batch_size, computed): + _total, br = computed + S_eff = args.seq_length * args.hc_mult + per_layer = _hand_attn_qkv_o( + B=batch_size, + S_eff=S_eff, + H=args.hidden_size, + n=args.num_attention_heads, + d=args.kv_channels, + q_lora=args.q_lora_rank, + o_lora=args.o_lora_rank, + o_groups=args.o_groups, + ) + assert br.attn_qkv_o == per_layer * args.num_layers + + def test_attn_scores_term_matches_reference(self, args, batch_size, computed): + _total, br = computed + S_eff = args.seq_length * args.hc_mult + expected = sum( + _hand_attn_scores( + B=batch_size, + S_eff=S_eff, + n=args.num_attention_heads, + d=args.kv_channels, + ratio=int(r), + index_topk=args.index_topk, + swa=int(getattr(args, "attn_sliding_window", 0) or 0), + ) + for r in args.compress_ratios + ) + assert br.attn_scores == expected + + def test_compressor_term_matches_reference(self, args, batch_size, computed): + _total, br = computed + S_eff = args.seq_length * args.hc_mult + expected = sum( + _hand_compressor( + B=batch_size, + S_eff=S_eff, + H=args.hidden_size, + d=args.kv_channels, + ratio=int(r), + ) + for r in args.compress_ratios + ) + assert br.compressor == expected + + def test_indexer_term_matches_reference(self, args, batch_size, computed): + _total, br = computed + S_eff = args.seq_length * args.hc_mult + expected = sum( + _hand_indexer( + B=batch_size, + S_eff=S_eff, + H=args.hidden_size, + ratio=int(r), + ihd=args.index_head_dim, + inh=args.index_n_heads, + ) + for r in args.compress_ratios + ) + assert br.indexer == expected + + def test_moe_term_respects_hash_layers(self, args, batch_size, computed): + _total, br = computed + S_eff = args.seq_length * args.hc_mult + expected = sum( + _hand_moe( + B=batch_size, + S_eff=S_eff, + H=args.hidden_size, + H_moe=args.moe_ffn_hidden_size, + topk=args.moe_router_topk, + n_experts=args.num_experts, + hash_layer=(layer_idx < args.num_hash_layers), + H_shared=args.moe_shared_expert_intermediate_size, + ) + for layer_idx in range(args.num_layers) + ) + assert br.moe == expected + + def test_logits_term_includes_one_extra_head_per_mtp_depth(self): + args = _v4_flash_smoke_args(mtp_num_layers=2) + _total, br = compute_v4_flops(args, batch_size=8) + expected = (args.mtp_num_layers + 1) * 8 * args.seq_length * args.hidden_size * args.padded_vocab_size + assert br.logits == expected + + def test_total_matches_breakdown_sum_with_expansion(self, computed): + total, br = computed + assert total == _FORWARD_BACKWARD_FACTOR * _FMA_FACTOR * br.total_fmac() + + def test_hc_term_is_nonzero_when_hc_mult_gt_1(self, computed): + """Plan-6 P33: HC matmul row must be populated when ``hc_mult > 1``. + + The exact closed form is pinned by the dedicated G36a test below; + here we just assert the term is present so a future refactor + that silently zeroes ``hc`` fails this fixture too. + """ + _total, br = computed + assert br.hc > 0 + + +# --------------------------------------------------------------------------- +# G36: Plan-6 P33 SWA visible-pair correction +# --------------------------------------------------------------------------- + + +class TestG36SWAVisiblePairs: + """Plan-6 P33: ``_attn_scores_fmac_per_layer`` must count only causal- + visible ``(q, k)`` pairs surviving SWA + pool + sparse top-K masks. + + The legacy plan-3 P20 closed form used ``B * n * d * S_eff^2`` for the + local branch (Megatron's ``S^2/2`` causal upper bound x the FMA-pair + factor) which over-counted by ``S_eff / swa_window`` once the kernel + started honoring SWA per-row pruning. This test pins the new + ``2 * n * d * visible_pairs`` form against the helper, the per-branch + over-count ratios, and the proxy-shape values printed in + ``deepseek-v4/develop/perf/attention_perf.md``. + """ + + @pytest.mark.parametrize("swa", [0, 64, 128, 4096, 8192]) + def test_local_visible_pair_helper(self, swa): + """Helper closed form matches the exhaustive sum-over-queries.""" + S_eff = 4096 + pairs = _visible_pairs( + swa_window=swa, + compress_ratio=0, + index_topk=0, + seq_len_eff=S_eff, + ) + exhaustive = sum(min(q + 1, swa) if (0 < swa < S_eff) else (q + 1) for q in range(S_eff)) + assert pairs == exhaustive + + def test_proxy_shape_dense_visible_pairs_matches_attn_perf_doc(self): + """``swa=128, S_eff=4096, cr=0`` → 516,160 (attention_perf.md row).""" + pairs = _visible_pairs( + swa_window=128, + compress_ratio=0, + index_topk=0, + seq_len_eff=4096, + ) + assert pairs == 516_160 + + def test_proxy_shape_hca_visible_pairs_matches_attn_perf_doc(self): + """``swa=128, S_eff=4096, cr=128`` → 516,160 + 63,520 = 579,680.""" + pairs = _visible_pairs( + swa_window=128, + compress_ratio=128, + index_topk=0, + seq_len_eff=4096, + ) + assert pairs == 516_160 + 63_520 + + def test_proxy_shape_csa_visible_pairs_matches_attn_perf_doc(self): + """``swa=128, S_eff=4096, cr=4, topk=512`` → 516,160 + 512*4096.""" + pairs = _visible_pairs( + swa_window=128, + compress_ratio=4, + index_topk=512, + seq_len_eff=4096, + ) + assert pairs == 516_160 + 512 * 4096 + + @pytest.mark.parametrize("hc_mult", [1, 4]) + @pytest.mark.parametrize("ratio", [0, 4, 128]) + def test_swa128_reduces_attn_scores_vs_full_causal(self, hc_mult, ratio): + """SWA=128 must produce strictly fewer score FMAC than swa=0 on + every layer type (no overlap between local and pool/topk terms + means the SWA-pruned local term strictly dominates the saving). + """ + S = 4096 + args_no_swa = _v4_flash_smoke_args( + num_layers=1, + seq_length=S, + hc_mult=hc_mult, + compress_ratios=(ratio,), + num_hash_layers=0, + attn_sliding_window=0, + ) + args_swa = _v4_flash_smoke_args( + num_layers=1, + seq_length=S, + hc_mult=hc_mult, + compress_ratios=(ratio,), + num_hash_layers=0, + attn_sliding_window=128, + ) + _t1, br_no = compute_v4_flops(args_no_swa, batch_size=1) + _t2, br_swa = compute_v4_flops(args_swa, batch_size=1) + assert br_swa.attn_scores < br_no.attn_scores + # All other components untouched by SWA. + assert br_swa.attn_qkv_o == br_no.attn_qkv_o + assert br_swa.compressor == br_no.compressor + assert br_swa.indexer == br_no.indexer + assert br_swa.moe == br_no.moe + assert br_swa.hc == br_no.hc + + def test_swa_pruned_attn_scores_matches_proxy_overcount_ratio(self): + """Per-layer over-count ratio between legacy ``S_eff^2`` and the + SWA-pruned visible-pair count is ``S_eff / (2*swa) + O(1/S)`` for + swa < S_eff (the factor of 2 comes from counting both the QK^T + and PV matmul halves of the attention score). At the V4-Flash + proxy shape (S=4096, hc_mult=4, S_eff=16384, swa=128) that ratio + is ``16384 / (2 * 128) = 64`` — pin it at >= 60x so a regression + that silently reverts to the legacy ``S_eff^2`` form is caught + while leaving 4-5x of headroom for off-by-one fringe corrections. + """ + S_eff = 16384 + swa = 128 + legacy_local = S_eff * S_eff # plan-3 P20 ``S_eff^2`` form + swa_local_pairs = swa * S_eff - swa * (swa - 1) // 2 + new_local = 2 * swa_local_pairs # 2 * visible_pairs (QK + PV) + ratio = legacy_local / new_local + assert ratio >= 60, f"SWA over-count ratio collapsed: {ratio:.2f}" + + +# --------------------------------------------------------------------------- +# G36a: Plan-6 P33 HyperConnection fn matmul accounting +# --------------------------------------------------------------------------- + + +class TestG36aHCMatmulAccounting: + """Plan-6 P33: the ``hc`` breakdown row must equal the closed form + ``B * S * K * D * K * (2 * (L + M) * (2 + K) + (1 + M))``. + + Pinned independently from G16 because it's a brand-new component + and the formula has a non-obvious factor structure (2 mixers per + layer x (L+M) layers, plus 1 trunk head + M MTP heads). + """ + + @pytest.mark.parametrize("hc_mult", [2, 4, 8]) + @pytest.mark.parametrize("mtp_num_layers", [0, 1, 2]) + def test_hc_matmul_matches_closed_form(self, hc_mult, mtp_num_layers): + args = _v4_flash_smoke_args( + num_layers=4, + seq_length=128, + hc_mult=hc_mult, + mtp_num_layers=mtp_num_layers, + ) + batch_size = 8 + _total, br = compute_v4_flops(args, batch_size=batch_size) + B = batch_size + S = args.seq_length + K = hc_mult + D = args.hidden_size + L = args.num_layers + M = mtp_num_layers + expected = B * S * K * D * K * (2 * (L + M) * (2 + K) + (1 + M)) + assert br.hc == expected + + def test_hc_matmul_uses_seq_len_not_seq_len_eff(self): + """HyperMixer runs on the un-packed ``[B, S, K, D]`` tensor; cost + must scale with ``seq_len`` not ``seq_len * hc_mult``. + + Concretely: doubling ``hc_mult`` from ``K=2 -> 4`` multiplies the + mixer factor ``K * D * K * (2+K)`` by ``(4*4*6) / (2*2*4) = 6``, + but doubling ``seq_length`` only doubles the per-layer cost. + Pinning both axes proves the closed form is keyed on the right + sequence-axis variable. + """ + ref = _v4_flash_smoke_args(num_layers=2, seq_length=64, hc_mult=2, mtp_num_layers=0) + ref_double_s = _v4_flash_smoke_args(num_layers=2, seq_length=128, hc_mult=2, mtp_num_layers=0) + ref_double_k = _v4_flash_smoke_args(num_layers=2, seq_length=64, hc_mult=4, mtp_num_layers=0) + _t1, br_ref = compute_v4_flops(ref, batch_size=1) + _t2, br_s = compute_v4_flops(ref_double_s, batch_size=1) + _t3, br_k = compute_v4_flops(ref_double_k, batch_size=1) + assert br_s.hc == 2 * br_ref.hc + + # Doubling K from 2 -> 4 scales mixer per-layer cost as + # K*K*(2+K) -> (4*4*6) / (2*2*4) = 6 and head cost as + # K*K -> 16/4 = 4. With L=2 mixers (= 4) and 1 head, + # the aggregate multiplier is (4 * 6 + 4) / (4 * 1 + 1) = 28/5. + # (L=2, M=0 → mixer term = 2*L*(2+K)*K^2*D = ratio 6; + # head term = (1+M)*K^2*D = ratio 4) + # Total ratio: (2*L*(2+K_new)*K_new^2 + (1+M)*K_new^2) / + # (2*L*(2+K_old)*K_old^2 + (1+M)*K_old^2) + L = 2 + M = 0 + K_old = 2 + K_new = 4 + num = 2 * L * (2 + K_new) * K_new * K_new + (1 + M) * K_new * K_new + den = 2 * L * (2 + K_old) * K_old * K_old + (1 + M) * K_old * K_old + assert br_k.hc * den == br_ref.hc * num + + +# --------------------------------------------------------------------------- +# Hash-layer / no-hash variant +# --------------------------------------------------------------------------- + + +class TestHashLayerHandling: + """Hash-routed layers must skip the topk router GEMM.""" + + def test_zero_hash_layers_charges_router_on_every_moe_layer(self): + args = _v4_flash_smoke_args(num_hash_layers=0) + _, with_hash = compute_v4_flops(args, batch_size=4) + + args_no_hash = _v4_flash_smoke_args(num_hash_layers=args.num_layers) + _, all_hash = compute_v4_flops(args_no_hash, batch_size=4) + + # All-hash strictly less because router cost is dropped on every layer. + assert all_hash.moe < with_hash.moe + delta_per_layer = 4 * args.seq_length * args.hc_mult * args.hidden_size * args.num_experts + assert with_hash.moe - all_hash.moe == delta_per_layer * args.num_layers + + +# --------------------------------------------------------------------------- +# compress_ratios normalization (decoder + MTP slicing) +# --------------------------------------------------------------------------- + + +class TestNormalizeLayerRatios: + def test_string_yaml_form_parses(self): + decoder, mtp = _normalize_layer_ratios("[0, 0, 4, 128, 4, 0]", num_layers=6, mtp_num_layers=0) + assert decoder == [0, 0, 4, 128, 4, 0] + assert mtp == [] + + def test_decoder_plus_mtp_layout_splits_correctly(self): + decoder, mtp = _normalize_layer_ratios([0, 4, 128, 4, 0], num_layers=4, mtp_num_layers=1) + assert decoder == [0, 4, 128, 4] + assert mtp == [0] + + def test_none_defaults_to_all_dense(self): + decoder, mtp = _normalize_layer_ratios(None, num_layers=3, mtp_num_layers=2) + assert decoder == [0, 0, 0] + assert mtp == [0, 0] + + def test_short_list_pads_with_last_value(self): + decoder, mtp = _normalize_layer_ratios([4, 128], num_layers=4, mtp_num_layers=0) + assert decoder == [4, 128, 128, 128] + assert mtp == [] + + +# --------------------------------------------------------------------------- +# G17: non-V4 fall-through byte-for-byte +# --------------------------------------------------------------------------- + + +class TestDispatchSwitch: + """G17: ``dispatch_v4`` flag controls whether upstream is called.""" + + @pytest.fixture(autouse=True) + def _silence_breakdown_log(self, monkeypatch): + """The wrapper emits a one-shot breakdown via ``log_rank_0`` on the + first V4 call. Under pytest there's no torch.distributed bound so + rank-aware logging would explode; flip the once-only flag to ``True`` + so the wrapper skips the emit path. + """ + from primus.backends.megatron.patches import deepseek_v4_flops_patches as mod + + monkeypatch.setattr(mod, "_BREAKDOWN_LOGGED", True) + + @pytest.fixture + def fake_upstream_factory(self): + """Returns ``(make_wrapped, sentinel_calls)`` per-test.""" + sentinel_calls = [] + + def fake_upstream(args, batch_size): + sentinel_calls.append((id(args), batch_size)) + return 12345 + batch_size + len(getattr(args, "model_type", "") or "") + + def make_wrapped(*, dispatch_v4: bool): + return _make_v4_num_floating_point_operations(fake_upstream, dispatch_v4=dispatch_v4) + + return make_wrapped, sentinel_calls + + @pytest.mark.parametrize( + "model_type", + ["gpt", "llama3", "deepseek_v3", "mamba_hybrid", None, ""], + ) + def test_dispatch_v4_false_falls_through_byte_for_byte(self, fake_upstream_factory, model_type): + make_wrapped, sentinel_calls = fake_upstream_factory + wrapped = make_wrapped(dispatch_v4=False) + args = SimpleNamespace(model_type=model_type) + result = wrapped(args, 32) + assert result == 12345 + 32 + len(model_type or "") + assert sentinel_calls == [(id(args), 32)] + + def test_dispatch_v4_true_uses_v4_closed_form(self, fake_upstream_factory): + make_wrapped, sentinel_calls = fake_upstream_factory + wrapped = make_wrapped(dispatch_v4=True) + args = _v4_flash_smoke_args() + result = wrapped(args, batch_size=4) + assert result > 0 + assert sentinel_calls == [] # Upstream MUST NOT be called when dispatching V4. + + def test_dispatch_v4_true_ignores_runtime_model_type_mutation(self, fake_upstream_factory): + """Megatron's pretrain() rewrites ``args.model_type`` with the + ``ModelType`` enum just before ``train()`` runs. The wrapper must + keep dispatching to V4 anyway because the install-time decision is + captured via the closure flag. + """ + from enum import Enum + + class _FakeModelType(Enum): + encoder_or_decoder = 1 + + make_wrapped, sentinel_calls = fake_upstream_factory + wrapped = make_wrapped(dispatch_v4=True) + args = _v4_flash_smoke_args() + args.model_type = _FakeModelType.encoder_or_decoder # post-pretrain() state + result = wrapped(args, batch_size=4) + assert result > 0 + assert sentinel_calls == [] + + +# --------------------------------------------------------------------------- +# Patch installation lifecycle +# --------------------------------------------------------------------------- + + +class TestPatchInstallation: + def test_idempotent_installation(self, monkeypatch): + """Re-applying the patch with an already-wrapped target is a no-op.""" + from primus.backends.megatron.patches import deepseek_v4_flops_patches as mod + from primus.core.patches.context import PatchContext + + def upstream(args, bs): + return 0 + + wrapped_once = mod._make_v4_num_floating_point_operations(upstream, dispatch_v4=True) + + # ``patch_v4_flops_reporting`` does ``import megatron.training.training``, + # which under bare pytest pulls in the real megatron package tree. + # Stub out every level so the import resolves to our fake module. + import sys + + fake_megatron = types.ModuleType("megatron") + fake_megatron_training_pkg = types.ModuleType("megatron.training") + fake_megatron_training_mod = types.ModuleType("megatron.training.training") + fake_megatron_training_mod.num_floating_point_operations = wrapped_once + fake_megatron_training_pkg.training = fake_megatron_training_mod + fake_megatron.training = fake_megatron_training_pkg + + monkeypatch.setitem(sys.modules, "megatron", fake_megatron) + monkeypatch.setitem(sys.modules, "megatron.training", fake_megatron_training_pkg) + monkeypatch.setitem(sys.modules, "megatron.training.training", fake_megatron_training_mod) + + # Silence rank-aware logger; primus' singleton ``_logger`` isn't bound + # under bare pytest so we replace ``log_rank_0`` with a no-op. + monkeypatch.setattr(mod, "log_rank_0", lambda *a, **kw: None) + + ctx = PatchContext( + backend="megatron", + phase="before_train", + extra={"module_config": types.SimpleNamespace(params=_v4_flash_smoke_args())}, + ) + mod.patch_v4_flops_reporting(ctx) + # The wrapped function must remain the same instance (no double-wrap). + assert fake_megatron_training_mod.num_floating_point_operations is wrapped_once diff --git a/tests/unit_tests/backends/megatron/test_rope_arange_cache.py b/tests/unit_tests/backends/megatron/test_rope_arange_cache.py new file mode 100644 index 000000000..8aa510f3d --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_rope_arange_cache.py @@ -0,0 +1,47 @@ +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. +"""Parity + caching test for ``RoPECache.forward_arange``. + +The compressed-branch RoPE is evaluated at the deterministic positions +``arange(P)`` every forward; ``forward_arange`` caches that table. This asserts +the cached table is bit-identical to recomputing ``forward(arange(n))`` and that +the cache actually memoises (and that ``PRIMUS_COMPRESS_ROPE_CACHE=0`` bypasses it). +""" +from __future__ import annotations + +import pytest +import torch + +from primus.backends.megatron.core.transformer.dual_rope import RoPECache + + +def _rope() -> RoPECache: + return RoPECache(rotary_dim=64, theta=10000.0) + + +@pytest.mark.parametrize("n", [32, 1024]) +def test_forward_arange_matches_forward(n): + """Cached table is bit-identical to the eager arange->outer->cos/sin path.""" + rc = _rope() + cos_a, sin_a = rc.forward_arange(n, "cpu") + cos_e, sin_e = rc.forward(torch.arange(n, device="cpu")) + torch.testing.assert_close(cos_a, cos_e, rtol=0, atol=0) + torch.testing.assert_close(sin_a, sin_e, rtol=0, atol=0) + + +def test_forward_arange_memoises(monkeypatch): + """A repeat call returns the SAME cached tensors (no recompute).""" + monkeypatch.setenv("PRIMUS_COMPRESS_ROPE_CACHE", "1") + rc = _rope() + a = rc.forward_arange(128, "cpu") + b = rc.forward_arange(128, "cpu") + assert a[0] is b[0] and a[1] is b[1] + + +def test_cache_disabled_recomputes(monkeypatch): + """PRIMUS_COMPRESS_ROPE_CACHE=0 bypasses the cache (fresh, equal tensors).""" + monkeypatch.setenv("PRIMUS_COMPRESS_ROPE_CACHE", "0") + rc = _rope() + a = rc.forward_arange(128, "cpu") + b = rc.forward_arange(128, "cpu") + assert a[0] is not b[0] + torch.testing.assert_close(a[0], b[0], rtol=0, atol=0) diff --git a/tests/unit_tests/configs/test_deepseek_v4_yaml.py b/tests/unit_tests/configs/test_deepseek_v4_yaml.py new file mode 100644 index 000000000..3e8498d7a --- /dev/null +++ b/tests/unit_tests/configs/test_deepseek_v4_yaml.py @@ -0,0 +1,335 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-2 P18 — DeepSeek-V4 YAML schema gate (G1). + +Each V4 yaml under ``primus/configs/models/megatron/deepseek_v4_*.yaml`` +must: + +* parse via the standard primus :func:`parse_yaml` loader (extends + + env resolution included); +* construct a :class:`DeepSeekV4TransformerConfig` after the standard + schema massage (no CRIT-level KeyError on plan-2 fields); +* surface ``compress_ratios`` as a tuple of ints after + ``__post_init__`` runs (P18 D4 — the YAML may carry the legacy string + form ``"[0, 0, 4, ...]"`` or a real list, and the dataclass must + normalize both); +* not carry retired fields (``v4_use_custom_mtp_block`` / + ``mtp_compress_ratios`` were dropped in plan-2 P17); +* not silently lose the V4-specific MoE / HC / sliding-window / sink + fields the schema relies on. + +Plus a single provider-singleton check (P18 D1): +:func:`resolve_v4_provider` returns the same instance on repeated +calls within the same config. +""" + +from __future__ import annotations + +from dataclasses import fields +from pathlib import Path +from typing import Any, Dict + +import pytest + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_YAML_DIR = _REPO_ROOT / "primus" / "configs" / "models" / "megatron" + + +# --------------------------------------------------------------------------- +# Helpers — schema-friendly subset of the V4 config kwargs +# --------------------------------------------------------------------------- + + +def _config_kwargs_from_yaml(yaml_dict: Dict[str, Any]) -> Dict[str, Any]: + """Filter the merged YAML dict down to keys consumed by + :class:`DeepSeekV4TransformerConfig` (and its parents). + + The full primus config tree mixes train / data / launcher fields + with the model-config fields we care about here. The dataclass + ignores unknown keys via Python's normal dataclass rules — but + only when we hand it kwargs that exist on the dataclass. We do + that by intersecting with ``DeepSeekV4TransformerConfig``'s + ``__dataclass_fields__``. + """ + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, + ) + + valid = set(DeepSeekV4TransformerConfig.__dataclass_fields__.keys()) + return {k: v for k, v in yaml_dict.items() if k in valid} + + +def _build_v4_config(yaml_dict: Dict[str, Any]): + """Construct a V4 config from a parsed YAML dict. + + Some upstream parents require ``num_attention_heads`` and a few + other fields. We rely on the YAML to provide them; missing fields + fall back to dataclass defaults. + """ + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, + ) + + kwargs = _config_kwargs_from_yaml(yaml_dict) + # The model YAML uses ``num_experts`` (mapped to Megatron's dataclass field + # ``num_moe_experts`` by the primus config loader at runtime); apply the same + # alias here so the standalone build is a valid MoE config. + if "num_experts" in yaml_dict and "num_moe_experts" not in kwargs: + kwargs["num_moe_experts"] = yaml_dict["num_experts"] + # The model YAML inherits ``sequence_parallel: true`` (for runtime TP>1), + # but this standalone build uses tensor_model_parallel_size=1, which + # Megatron rejects ("sequence parallelism without tensor parallelism"). + kwargs["sequence_parallel"] = False + # V4 uses the sqrtsoftplus score function, but Megatron only allows the + # aux-loss-free expert bias with sigmoid. The runner disables expert bias + # for V4; do the same here. None of these fields affect compress_ratios, + # which is the only thing these tests validate. + kwargs["moe_router_enable_expert_bias"] = False + return DeepSeekV4TransformerConfig(**kwargs) + + +# --------------------------------------------------------------------------- +# YAML parsing — base / flash / pro all parse and merge cleanly +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def parse_yaml_fn(): + """Bind the shared loader once for all tests in this module.""" + from primus.core.config.yaml_loader import parse_yaml + + return parse_yaml + + +@pytest.mark.parametrize( + "yaml_name", + ["deepseek_v4_base.yaml", "deepseek_v4_flash.yaml", "deepseek_v4_pro.yaml"], +) +def test_v4_yaml_parses(parse_yaml_fn, yaml_name: str) -> None: + parsed = parse_yaml_fn(str(_YAML_DIR / yaml_name)) + assert isinstance(parsed, dict) + # Every V4 yaml must declare its core shape. + for required in ("num_layers", "hidden_size", "num_attention_heads"): + # base.yaml inherits these from elsewhere, so this is best-effort. + if yaml_name == "deepseek_v4_base.yaml": + continue + assert required in parsed, f"{yaml_name} missing required key {required!r}" + + +# --------------------------------------------------------------------------- +# compress_ratios is normalized to tuple[int, ...] (P18 D4) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "yaml_name", + ["deepseek_v4_flash.yaml", "deepseek_v4_pro.yaml"], +) +def test_compress_ratios_normalized_to_tuple(parse_yaml_fn, yaml_name: str) -> None: + """The dataclass ``__post_init__`` must convert + ``compress_ratios`` (which may arrive as a YAML string or a list) + into ``tuple[int, ...]``. + + The base yaml does not pin a schedule (it is provided per-variant); + flash + pro do. Both must round-trip to a tuple of ints with no + string survivors and no value drift vs the raw schedule. + """ + parsed = parse_yaml_fn(str(_YAML_DIR / yaml_name)) + raw = parsed["compress_ratios"] + + # Sanity: the YAML form is either a string or a list. + assert isinstance(raw, (str, list, tuple)) + + cfg = _build_v4_config(parsed) + normalized = cfg.compress_ratios + + assert normalized is not None + assert isinstance( + normalized, tuple + ), f"{yaml_name}: compress_ratios should normalize to tuple, got {type(normalized).__name__}" + for i, r in enumerate(normalized): + assert isinstance(r, int), f"{yaml_name}: compress_ratios[{i}]={r!r} is not int" + + # Length / values: strip the YAML form and compare to the parsed-int form. + if isinstance(raw, str): + import ast as _ast + + raw_list = _ast.literal_eval(raw) + else: + raw_list = list(raw) + assert ( + tuple(int(x) for x in raw_list) == normalized + ), f"{yaml_name}: compress_ratios value drift — yaml={raw_list}, normalized={normalized}" + + +def test_compress_ratios_canonical_dispatch_values_only(parse_yaml_fn) -> None: + """Plan-2 P17 fixed the comment inversion in V4 yamls; here we + enforce the *value* contract: every entry must be one of + ``{0, 4, 128}`` (V4 attention only knows these branches — + anything else is a schema bug).""" + for yaml_name in ("deepseek_v4_flash.yaml", "deepseek_v4_pro.yaml"): + parsed = parse_yaml_fn(str(_YAML_DIR / yaml_name)) + cfg = _build_v4_config(parsed) + bad = [(i, r) for i, r in enumerate(cfg.compress_ratios) if r not in (0, 4, 128)] + assert not bad, ( + f"{yaml_name}: compress_ratios contains non-canonical values " + f"(allowed: 0=dense+SWA, 4=CSA, 128=HCA); offenders: {bad}" + ) + + +# --------------------------------------------------------------------------- +# Retired schema fields (P17 hygiene) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "retired_field", + ["v4_use_custom_mtp_block", "mtp_compress_ratios"], +) +def test_retired_fields_not_in_v4_config(retired_field: str) -> None: + """Plan-2 P17 deleted the legacy MTP block; both gating fields + must be removed from ``DeepSeekV4TransformerConfig``.""" + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, + ) + + names = {f.name for f in fields(DeepSeekV4TransformerConfig)} + assert ( + retired_field not in names + ), f"{retired_field} must be removed from DeepSeekV4TransformerConfig (plan-2 P17)." + + +@pytest.mark.parametrize( + "yaml_name", + ["deepseek_v4_base.yaml", "deepseek_v4_flash.yaml", "deepseek_v4_pro.yaml"], +) +def test_yaml_does_not_set_retired_fields(parse_yaml_fn, yaml_name: str) -> None: + """If a V4 YAML still references a retired field (e.g. someone + forgot to update a downstream yaml) we want a loud schema error, + not silent acceptance. + + The dataclass would now reject these as ``TypeError`` (unexpected + keyword), but we filter unknown keys for forward-compat in + ``_config_kwargs_from_yaml``; here we explicitly enforce the + contract on the parsed dict instead. + """ + parsed = parse_yaml_fn(str(_YAML_DIR / yaml_name)) + for retired in ("v4_use_custom_mtp_block", "mtp_compress_ratios"): + assert retired not in parsed, ( + f"{yaml_name} still references retired field {retired!r} " + "(plan-2 P17). Remove the line or update training scripts." + ) + + +# --------------------------------------------------------------------------- +# V4-specific schema fields the runtime depends on (D5 / D6 hygiene) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "field_name", + [ + # Hyper-Connection + "hc_mult", + "hc_eps", + # Hybrid attention extras + "compress_ratios", + "compress_rope_theta", + "attn_sliding_window", + "attn_sink", + # Grouped low-rank O projection + "o_groups", + "o_lora_rank", + # MoE-specific extras + "num_hash_layers", + "swiglu_limit", + ], +) +def test_v4_config_carries_runtime_field(field_name: str) -> None: + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, + ) + + names = {f.name for f in fields(DeepSeekV4TransformerConfig)} + assert field_name in names, ( + f"DeepSeekV4TransformerConfig must declare {field_name!r}; " + "removing it silently breaks the V4 runtime." + ) + + +# --------------------------------------------------------------------------- +# Provider singleton (P18 D1) — same instance per config +# --------------------------------------------------------------------------- + + +def test_resolve_v4_provider_caches_per_config(): + from primus.backends.megatron.core.models.deepseek_v4.build_context import ( + resolve_v4_provider, + ) + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, + ) + + # Build a minimal V4 config; we only need the dataclass instance — + # the provider does not touch most fields at construction. + cfg_a = DeepSeekV4TransformerConfig( + num_layers=4, + hidden_size=64, + num_attention_heads=4, + kv_channels=16, + ) + cfg_b = DeepSeekV4TransformerConfig( + num_layers=4, + hidden_size=64, + num_attention_heads=4, + kv_channels=16, + ) + + p_a1 = resolve_v4_provider(cfg_a) + p_a2 = resolve_v4_provider(cfg_a) + p_b1 = resolve_v4_provider(cfg_b) + + assert p_a1 is p_a2, "resolve_v4_provider must reuse the cached provider for the same config (P18 D1)." + assert p_a1 is not p_b1, ( + "Different config instances should each get their own provider so " + "test isolation and runtime overrides are not silently shared." + ) + + +# --------------------------------------------------------------------------- +# Provider activation_func helper (P18 D2) +# --------------------------------------------------------------------------- + + +def test_v4_mlp_activation_func_respects_use_te_activation_func() -> None: + """``provider.v4_mlp_activation_func()`` returns ``None`` when the + config keeps Megatron's eager activation path (the V4 default, + needed for clamped-SwiGLU); only when the user opts into TE does + the spec slot carry the TE class. + """ + from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import ( + DeepSeekV4SpecProvider, + TEActivationOp, + ) + + class _FakeCfg: + use_te_activation_func = False + + p = DeepSeekV4SpecProvider(config=_FakeCfg()) + assert p.v4_mlp_activation_func() is None, ( + "Default V4 path must leave activation_func slot empty so Megatron " + "MLP uses config.activation_func (clamped-SwiGLU)." + ) + + class _TeCfg: + use_te_activation_func = True + + p2 = DeepSeekV4SpecProvider(config=_TeCfg()) + assert p2.v4_mlp_activation_func() is TEActivationOp, ( + "When the user opts into TE activation, the spec slot must carry " + "TEActivationOp so Megatron MLP instantiates the TE-fused path." + ) diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index cc33c5982..24d07cf77 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -10,6 +10,20 @@ import sys from pathlib import Path +import pytest + + +def pytest_addoption(parser): + parser.addoption( + "--run-slow", + action="store_true", + default=False, + help=( + "run tests marked with @pytest.mark.slow (e.g. plan-4 release-tier " + "shape gates at head_dim=512). Default: skip slow tests." + ), + ) + def pytest_configure(config): # Add project root first to ensure main primus package takes precedence @@ -23,3 +37,26 @@ def pytest_configure(config): megatron_path = project_root / "third_party" / "Megatron-LM" if str(megatron_path) not in sys.path: sys.path.append(str(megatron_path)) + + # Register custom markers used by the primus test suite. + config.addinivalue_line( + "markers", + ( + "slow: marks tests as slow (release-tier shape gates that exercise " + "production V4 dims at head_dim=512). Skipped by default; opt in " + "with --run-slow or '-m slow'." + ), + ) + + +def pytest_collection_modifyitems(config, items): + # If the user explicitly opts in via --run-slow or '-m slow', run slow tests. + if config.getoption("--run-slow", default=False): + return + marker_expr = config.getoption("-m") or "" + if "slow" in marker_expr: + return + skip_slow = pytest.mark.skip(reason="slow test (use --run-slow or '-m slow' to enable)") + for item in items: + if "slow" in item.keywords: + item.add_marker(skip_slow) diff --git a/tests/unit_tests/megatron/cco/test_tp_overlap.py b/tests/unit_tests/megatron/cco/test_tp_overlap.py index ceab475e2..a58af584d 100644 --- a/tests/unit_tests/megatron/cco/test_tp_overlap.py +++ b/tests/unit_tests/megatron/cco/test_tp_overlap.py @@ -6,6 +6,7 @@ import functools from contextlib import contextmanager +import pytest import torch import torch.distributed as dist import transformer_engine as te @@ -22,13 +23,21 @@ ) from transformer_engine.pytorch import LayerNormLinear, Linear, fp8_autocast -from primus.backends.transformer_engine import transformer_engine_torch as ptex -from primus.backends.transformer_engine.pytorch.module.base import ( +# The Primus TE comm-overlap extension imports the ROCm ``hip`` Python module +# (``primus.backends.transformer_engine.transformer_engine_torch.comm_overlap``). +# Skip the whole module when ``hip`` isn't importable (non-ROCm or minimally +# provisioned host) instead of failing at collection time. +pytest.importorskip("hip", reason="ROCm 'hip' Python module not available") + +from primus.backends.transformer_engine import ( # noqa: E402 + transformer_engine_torch as ptex, +) +from primus.backends.transformer_engine.pytorch.module.base import ( # noqa: E402 get_workspace, initialize_ub, ) -from primus.core.utils import logger -from primus.core.utils.module_utils import set_logging_rank +from primus.core.utils import logger # noqa: E402 +from primus.core.utils.module_utils import set_logging_rank # noqa: E402 @contextmanager diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/__init__.py b/tests/unit_tests/megatron/transformer/deepseek_v4/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/conftest.py b/tests/unit_tests/megatron/transformer/deepseek_v4/conftest.py new file mode 100644 index 000000000..2b14d4f28 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/conftest.py @@ -0,0 +1,80 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Per-test cleanup hooks for plan-4 V4 attention tests. + +The plan-4 release-tier shape gates (``pytest.mark.slow``, see G28 in +``deepseek-v4/develop/plan-4/02-phase-details.md``) allocate large +eager-reference tensors at production V4 dimensions +(``head_dim=512``, MQA / MHA, ``K_topk`` up to 1024). Specifically the +CSA path's ``torch.einsum("bhsd,bhskd->bhsk", q, gathered.unsqueeze(1) +.expand(...))`` materialises a ``[B, H, Sq, K, D]`` intermediate that +is up to ~64 GiB at fp32 V4-Flash release / ~128 GiB at fp32 V4-Pro +release. PyTorch's caching allocator otherwise holds those tensors +across parametrised tests in the same process, so the second / third +release-tier test in a pytest session ends up OOM-ing well before +exhausting the MI355 287 GiB HBM budget. + +This conftest installs a function-scoped autouse fixture that empties +the CUDA / HIP cache after every test to keep the allocator pressure +test-local. The cost (~milliseconds) is negligible for the fast tier +and an outright requirement for the release tier. + +It also gates the entire directory on MI355X hardware: the V4 kernels +and their numerics are validated on MI355X (gfx950), so on any other +accelerator (or CPU-only host) every test in this directory is skipped. +Set ``PRIMUS_V4_UT_ALLOW_NON_MI355X=1`` to bypass the gate. +""" + +from __future__ import annotations + +import functools +import os + +import pytest + +try: + import torch +except ImportError: # pragma: no cover — pytest collects on torchless envs + torch = None # type: ignore + + +@functools.lru_cache(maxsize=1) +def _is_mi355x() -> bool: + """Return True iff the local accelerator reports as an AMD MI355X. + + Matches the device marketing name (e.g. ``AMD Instinct MI355X``); + returns False on CPU-only / non-ROCm hosts. + """ + if torch is None or not torch.cuda.is_available(): + return False + try: + name = torch.cuda.get_device_name(0) or "" + except Exception: # pragma: no cover — defensive + return False + return "mi355x" in name.lower() + + +@pytest.fixture(autouse=True) +def _require_mi355x(): + """Skip DeepSeek-V4 unit tests unless the GPU is an MI355X. + + The V4 kernels / numerics are validated against MI355X; running them on + other architectures is neither supported nor meaningful. Set + ``PRIMUS_V4_UT_ALLOW_NON_MI355X=1`` to bypass (e.g. local debugging). + """ + if os.environ.get("PRIMUS_V4_UT_ALLOW_NON_MI355X", "0") != "0": + return + if not _is_mi355x(): + pytest.skip("DeepSeek-V4 unit tests require an MI355X GPU") + + +@pytest.fixture(autouse=True) +def _release_gpu_cache_per_test(): + """Free PyTorch's CUDA / HIP cache after each V4 attention test.""" + yield + if torch is not None and torch.cuda.is_available(): + torch.cuda.empty_cache() diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_clamped_swiglu.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_clamped_swiglu.py new file mode 100644 index 000000000..2c991ef1a --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_clamped_swiglu.py @@ -0,0 +1,163 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for V4 pre-multiplication clamped SwiGLU (G3, plan-2 §04). + +These tests pin the math against the HF reference at +``DeepSeek-V4-Flash/inference/model.py:Expert.forward`` and assert the +parameter layout (``w1`` / ``w2`` / ``w3``) so the released checkpoint +loads through the V4 state-dict adapter without remapping. + +Tolerance: 1e-6 absolute, fp32, on randomized inputs (G3 gate). +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from primus.backends.megatron.core.transformer.clamped_swiglu import ( + ClampedSwiGLUMLP, + clamped_swiglu_pre_mul, + clamped_swiglu_pre_mul_fused, +) + + +def _hf_reference_clamped_swiglu( + gate: torch.Tensor, + up: torch.Tensor, + *, + alpha: float, +) -> torch.Tensor: + """Inline HF reference for ``Expert.forward``'s activation. + + Mirrors the HF code exactly, including the float() cast and the + pre-multiplication one-sided / two-sided clamp. + """ + gate = gate.float() + up = up.float() + if alpha > 0.0: + up = torch.clamp(up, min=-alpha, max=alpha) + gate = torch.clamp(gate, max=alpha) + return F.silu(gate) * up + + +@pytest.mark.parametrize("alpha", [7.0, 3.5, 1.0, 0.5]) +def test_clamped_swiglu_pre_mul_matches_hf_reference(alpha: float) -> None: + """G3: split-input form matches HF reference within 1e-6 fp32.""" + torch.manual_seed(1234) + gate = torch.randn(4, 17, 32, dtype=torch.float32) * 5.0 + up = torch.randn(4, 17, 32, dtype=torch.float32) * 5.0 + + out = clamped_swiglu_pre_mul(gate, up, alpha=alpha) + ref = _hf_reference_clamped_swiglu(gate, up, alpha=alpha) + + assert out.shape == gate.shape + assert torch.isfinite(out).all() + max_abs = (out - ref).abs().max().item() + assert max_abs <= 1.0e-6, f"max-abs error vs HF reference = {max_abs}" + + +def test_clamped_swiglu_pre_mul_alpha_zero_disables_clamp() -> None: + """``alpha == 0`` falls back to vanilla SwiGLU (no clamp).""" + torch.manual_seed(7) + gate = torch.randn(2, 16, dtype=torch.float32) * 100.0 + up = torch.randn(2, 16, dtype=torch.float32) * 100.0 + out = clamped_swiglu_pre_mul(gate, up, alpha=0.0) + expected = F.silu(gate) * up + assert (out - expected).abs().max().item() <= 1.0e-6 + + +def test_clamped_swiglu_pre_mul_fused_matches_split_form() -> None: + """The fused [gate | up] form matches the split-input form.""" + torch.manual_seed(99) + I = 24 + gate = torch.randn(3, 11, I, dtype=torch.float32) * 4.0 + up = torch.randn(3, 11, I, dtype=torch.float32) * 4.0 + fused = torch.cat([gate, up], dim=-1) + + out_split = clamped_swiglu_pre_mul(gate, up, alpha=7.0) + out_fused = clamped_swiglu_pre_mul_fused(fused, alpha=7.0) + assert (out_split - out_fused).abs().max().item() <= 1.0e-6 + + +def test_clamped_swiglu_pre_mul_one_sided_gate_clamp() -> None: + """Verify the gate clamp is one-sided (max=alpha) and up is two-sided. + + Sets gate values both well below -alpha and well above +alpha; the + output activation should pass the negative-gate values through SiLU + (since gate clamp only bounds the top side) and bound the positive + gate side at alpha. + """ + alpha = 7.0 + gate_lo = torch.tensor([[-100.0]], dtype=torch.float32) + gate_hi = torch.tensor([[+100.0]], dtype=torch.float32) + up = torch.tensor([[1.0]], dtype=torch.float32) + + out_lo = clamped_swiglu_pre_mul(gate_lo, up, alpha=alpha) + out_hi = clamped_swiglu_pre_mul(gate_hi, up, alpha=alpha) + + # Gate negative side is NOT clamped; SiLU(-100) ~ -100 * sigmoid(-100) ~ 0. + # We just assert the output is finite and small (close to 0). + assert torch.isfinite(out_lo).item() + assert out_lo.abs().item() < 1.0e-6 + # Gate positive side IS clamped at alpha; SiLU(alpha) is the upper bound. + expected_hi = F.silu(torch.tensor(alpha)) * up + assert (out_hi - expected_hi).abs().max().item() <= 1.0e-6 + + +def test_clamped_swiglu_mlp_state_dict_uses_w1_w2_w3_layout() -> None: + """Released-checkpoint compatibility: w1 / w2 / w3 keys exist. + + The HF released ``Expert`` checkpoint uses ``w1.weight`` / + ``w2.weight`` / ``w3.weight`` (no ``gate_up.weight``). This test + fails if a future refactor breaks that promise. + """ + mlp = ClampedSwiGLUMLP(hidden_size=8, intermediate_size=16, alpha=7.0, bias=False) + keys = set(mlp.state_dict().keys()) + assert "w1.weight" in keys + assert "w2.weight" in keys + assert "w3.weight" in keys + # No fused projection key may leak into the state-dict. + assert "gate_up.weight" not in keys + + +def test_clamped_swiglu_mlp_fused_gate_up_matches_split_path() -> None: + """The fused-forward variant produces identical outputs to the eager path.""" + torch.manual_seed(42) + hidden = torch.randn(3, 7, 8, dtype=torch.float32) + + mlp_split = ClampedSwiGLUMLP(hidden_size=8, intermediate_size=16, alpha=7.0) + mlp_fused = ClampedSwiGLUMLP(hidden_size=8, intermediate_size=16, alpha=7.0, fused_gate_up=True) + # Share weights so we are only testing the forward layout, not init. + mlp_fused.load_state_dict(mlp_split.state_dict()) + + out_split = mlp_split(hidden) + out_fused = mlp_fused(hidden) + assert (out_split - out_fused).abs().max().item() <= 1.0e-6 + + +def test_clamped_swiglu_mlp_forward_matches_hf_expert_forward() -> None: + """End-to-end: ``ClampedSwiGLUMLP.forward`` matches HF ``Expert.forward``.""" + torch.manual_seed(0) + H, I = 12, 28 + alpha = 7.0 + mlp = ClampedSwiGLUMLP(hidden_size=H, intermediate_size=I, alpha=alpha) + + x = torch.randn(2, 5, H, dtype=torch.float32) + + # Inline HF reference using the same w1/w2/w3 weights. + w1 = mlp.w1.weight.detach() + w3 = mlp.w3.weight.detach() + w2 = mlp.w2.weight.detach() + gate = F.linear(x, w1) + up = F.linear(x, w3) + h = _hf_reference_clamped_swiglu(gate, up, alpha=alpha) + ref = F.linear(h, w2) + + out = mlp(x) + assert (out - ref).abs().max().item() <= 1.0e-6 diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_deepseek_v4_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_deepseek_v4_attention.py new file mode 100644 index 000000000..9887bccbf --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_deepseek_v4_attention.py @@ -0,0 +1,872 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for the plan-2 V4-faithful :class:`DeepseekV4Attention`. + +These tests run on CPU (fp32) and verify the dense (``compress_ratio == 0``) +attention math against an *inline* reference implementation that mirrors +the released ``DeepSeek-V4-Flash/inference/model.py`` semantics: + +* Single-latent KV: ``K = V = wkv(hidden)`` broadcast to all query heads. +* Per-head ``q_rms``: parameter-less RMS on ``head_dim`` after ``wq_b``. +* Learnable per-head ``attn_sink``: an extra "virtual key" with zero value + joined into the softmax (then dropped before the value-weighted sum). +* Grouped low-rank O: einsum-based ``wo_a`` per group + ``wo_b``. +* Partial **interleaved** RoPE on the last ``qk_pos_emb_head_dim`` channels. + +The reference uses the *same* interleaved RoPE convention as Primus's +:mod:`dual_rope` (per the techblog correction over the original HF PR +which used rotate-half), so the two implementations should agree to +machine precision when given identical weights and positions. + +Plan-2 P17 will replace this inline reference with the actual HF +``DeepseekV4Attention.forward`` from the released checkpoint. +""" + +from __future__ import annotations + +import math + +import pytest + +torch = pytest.importorskip("torch") + +# Importing the module under test pulls Megatron / TE through the +# Primus extensions; skip the whole module when those are unavailable +# (e.g. a CPU-only smoke environment). +mla_module = pytest.importorskip( + "megatron.core.transformer.multi_latent_attention", + reason="Megatron MLA not importable in this environment", +) + +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( # noqa: E402 + DeepSeekV4TransformerConfig, +) +from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( # noqa: E402 + DeepseekV4Attention, +) +from primus.backends.megatron.core.transformer.dual_rope import DualRoPE # noqa: E402 + +# --------------------------------------------------------------------------- +# Inline reference V4 attention forward (single-latent KV, attn_sink, grouped O) +# --------------------------------------------------------------------------- + + +def _rms_norm_per_head(x: torch.Tensor, eps: float) -> torch.Tensor: + """Parameter-less per-head RMS (matches the released checkpoint).""" + in_dtype = x.dtype + x32 = x.float() + rms = torch.rsqrt(x32.square().mean(dim=-1, keepdim=True) + eps) + return (x32 * rms).to(in_dtype) + + +def _reference_v4_attention_forward( + *, + hidden: torch.Tensor, # [B, S, D] + position_ids: torch.Tensor, # [B, S] or [S] + rope: DualRoPE, + wq_a: torch.Tensor, # [q_lora_rank, D] + wq_b: torch.Tensor, # [n_heads * head_dim, q_lora_rank] + q_norm_w: torch.Tensor, # [q_lora_rank] + wkv: torch.Tensor, # [head_dim, D] + kv_norm_w: torch.Tensor, # [head_dim] + wo_a: torch.Tensor, # [o_groups * o_lora_rank, n_heads * head_dim / o_groups] + wo_b: torch.Tensor, # [D, o_groups * o_lora_rank] + attn_sink, # Optional[torch.Tensor] of shape [n_heads]; None disables sink + n_heads: int, + head_dim: int, + rotary_dim: int, + o_groups: int, + o_lora_rank: int, + norm_eps: float, +) -> torch.Tensor: + """Reference V4 attention forward (CPU, fp32-ish, no parallel linear).""" + B, S, D = hidden.shape + + # Q branch. + q_compressed = hidden @ wq_a.t() # [B, S, q_lora_rank] + # RMSNorm on q_lora_rank with learnable gamma (== q_layernorm). + q32 = q_compressed.float() + q_rms = torch.rsqrt(q32.square().mean(dim=-1, keepdim=True) + norm_eps) + q_compressed = (q32 * q_rms).to(q_compressed.dtype) * q_norm_w + q = q_compressed @ wq_b.t() # [B, S, n_heads * head_dim] + q = q.view(B, S, n_heads, head_dim) + q = _rms_norm_per_head(q, norm_eps) + + # KV branch (single-latent). + kv = hidden @ wkv.t() # [B, S, head_dim] + kv32 = kv.float() + kv_rms = torch.rsqrt(kv32.square().mean(dim=-1, keepdim=True) + norm_eps) + kv = (kv32 * kv_rms).to(kv.dtype) * kv_norm_w + kv = kv.view(B, S, 1, head_dim) + + # Partial RoPE (interleaved) on Q and K. K = kv (rope-applied). + q = rope.apply_rope(q, position_ids=position_ids, compress_ratio=0) + kv = rope.apply_rope(kv, position_ids=position_ids, compress_ratio=0) + k = kv # [B, S, 1, head_dim] + v = kv # K = V (single latent) + + # Broadcast K / V across heads. + k = k.expand(B, S, n_heads, head_dim) + v = v.expand(B, S, n_heads, head_dim) + + # Causal mask (no SWA in the test). + q_idx = torch.arange(S, device=hidden.device).unsqueeze(1) + k_idx = torch.arange(S, device=hidden.device).unsqueeze(0) + causal = torch.where(q_idx >= k_idx, 0.0, float("-inf")).to(hidden.dtype) + + # Move heads dim before sequence. + q_bh = q.transpose(1, 2) # [B, H, S, head_dim] + k_bh = k.transpose(1, 2) + v_bh = v.transpose(1, 2) + + scale = 1.0 / math.sqrt(head_dim) + logits = torch.matmul(q_bh.float(), k_bh.float().transpose(-2, -1)) * scale + logits = logits + causal + + if attn_sink is None: + # Plain causal softmax (no virtual key column). + logits = logits - logits.amax(dim=-1, keepdim=True).detach() + probs = logits.softmax(dim=-1).to(v_bh.dtype) + else: + # attn_sink: append a virtual key column with zero value, drop after softmax. + sink_col = attn_sink.float().view(1, n_heads, 1, 1).expand(B, n_heads, S, 1) + logits_aug = torch.cat([logits, sink_col], dim=-1) + logits_aug = logits_aug - logits_aug.amax(dim=-1, keepdim=True).detach() + probs = logits_aug.softmax(dim=-1) + probs = probs[..., :-1].to(v_bh.dtype) + out = torch.matmul(probs, v_bh.float()).to(hidden.dtype) + out = out.transpose(1, 2).contiguous() # [B, S, H, head_dim] + + # Grouped low-rank O. + out_g = out.reshape(B, S, o_groups, (n_heads * head_dim) // o_groups) + wo_a_w = wo_a.view(o_groups, o_lora_rank, (n_heads * head_dim) // o_groups) + o = torch.einsum("bsgd,grd->bsgr", out_g, wo_a_w) + o = o.flatten(2) + return o @ wo_b.t() + + +# --------------------------------------------------------------------------- +# Test fixtures +# --------------------------------------------------------------------------- + + +_TEST_DTYPE = torch.float32 + + +@pytest.fixture(autouse=True) +def _v4_attention_on_cuda(): + """DeepSeek-V4 attention (all backends, incl. eager) runs on GPU tensors. + + These tests build CPU-free tensors by defaulting to the CUDA/HIP device; + skip the module on a CPU-only host. + """ + if not torch.cuda.is_available(): + pytest.skip("DeepSeek-V4 attention requires a CUDA/HIP device") + torch.set_default_device("cuda") + try: + yield + finally: + torch.set_default_device("cpu") + + +def _make_v4_config( + *, + hidden_size: int, + num_heads: int, + head_dim: int, + rotary_dim: int, + q_lora_rank: int, + o_groups: int, + o_lora_rank: int, + attn_sink: bool, + norm_eps: float = 1e-6, +) -> DeepSeekV4TransformerConfig: + """Minimal V4 config for CPU unit tests. + + Relies on dataclass defaults wherever possible and only sets the + fields the V4 attention actually reads. + """ + return DeepSeekV4TransformerConfig( + num_layers=1, + hidden_size=hidden_size, + num_attention_heads=num_heads, + num_query_groups=1, + kv_channels=head_dim, + qk_pos_emb_head_dim=rotary_dim, + qk_head_dim=head_dim - rotary_dim, + v_head_dim=head_dim, + kv_lora_rank=head_dim, # unused — KV branch is overridden by V4 + rope_type="rope", + rotary_base=10000.0, + rotary_scaling_factor=1.0, + rotary_percent=1.0, + original_max_position_embeddings=2048, + # V4 extras + q_lora_rank=q_lora_rank, + o_groups=o_groups, + o_lora_rank=o_lora_rank, + attn_sliding_window=0, + attn_sink=attn_sink, + compress_ratios=None, + compress_rope_theta=160000.0, + # These tests validate the V4 attention math against an inline eager + # reference at tiny (head_dim=16) shapes; the Triton/gluon/turbo kernels + # are specialized for head_dim=512, so pin the eager backend here. + use_v4_attention_backend="eager", + use_v4_csa_attention_backend="eager", + # Misc + layernorm_epsilon=norm_eps, + norm_epsilon=norm_eps, + attention_dropout=0.0, + hidden_dropout=0.0, + ) + + +def _make_attention(config: DeepSeekV4TransformerConfig) -> DeepseekV4Attention: + """Construct V4 attention with default (no-spec) submodules, on CPU. + + With ``submodules=None`` every projection falls back to ``nn.Linear`` + (replicated, no TP), which is exactly what we want for a CPU smoke + test against the inline reference. + """ + rope = DualRoPE( + rotary_dim=config.qk_pos_emb_head_dim, + rope_theta=config.rotary_base, + compress_rope_theta=config.compress_rope_theta, + yarn_factor=1.0, + original_max_position_embeddings=config.original_max_position_embeddings, + ) + return DeepseekV4Attention( + config, + rope=rope, + compress_ratio=0, + submodules=None, + ) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def test_state_dict_keys_match_v4_canonical_layout(): + """The new attention exposes V4-canonical state-dict keys. + + These are exactly the keys the P17 state-dict adapter will map from + the released ``layers.{i}.attn.{wq_a,wq_b,wkv,q_norm,kv_norm,wo_a,wo_b,attn_sink}`` + safetensors layout. + """ + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=8, + attn_sink=True, + ) + attn = _make_attention(config) + keys = set(attn.state_dict().keys()) + + expected = { + "linear_q_down_proj.weight", + "linear_q_up_proj.weight", + "linear_kv.weight", + "q_layernorm.weight", + "kv_layernorm.weight", + "linear_o_a.weight", + "linear_o_b.weight", + "attn_sink", # direct nn.Parameter, matches released checkpoint key + } + missing = expected - keys + assert not missing, f"missing V4-canonical keys: {missing}" + + legacy = { + "q_a.weight", + "q_b.weight", + "k_proj.weight", + "v_proj.weight", + "o_proj.weight", + } + bleed = legacy & keys + assert not bleed, f"legacy plan-1 keys leaked into V4-faithful attention: {bleed}" + + +def test_forward_shape_and_finite(): + """Forward pass returns ``[B, S, hidden_size]`` of finite values.""" + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=8, + attn_sink=True, + ) + attn = _make_attention(config).to(_TEST_DTYPE) + + B, S, D = 2, 8, 64 + hidden = torch.randn(B, S, D, dtype=_TEST_DTYPE) + position_ids = torch.arange(S).unsqueeze(0).expand(B, S) + out = attn(hidden, position_ids) + assert out.shape == (B, S, D) + assert torch.isfinite(out).all() + + +def _copy_weights_to_reference(attn: DeepseekV4Attention) -> dict: + """Extract weights from the new attention into the inline reference's + parameter dict.""" + return { + "wq_a": attn.linear_q_down_proj.weight.detach().clone(), + "wq_b": attn.linear_q_up_proj.weight.detach().clone(), + "q_norm_w": attn.q_layernorm.weight.detach().clone(), + "wkv": attn.linear_kv.weight.detach().clone(), + "kv_norm_w": attn.kv_layernorm.weight.detach().clone(), + "wo_a": attn.linear_o_a.weight.detach().clone(), + "wo_b": attn.linear_o_b.weight.detach().clone(), + } + + +@pytest.mark.parametrize("attn_sink_enabled", [False, True]) +def test_forward_matches_inline_reference(attn_sink_enabled: bool): + """1-layer V4 attention forward agrees with the inline reference. + + Both implementations apply the same math (single-latent KV, partial + interleaved RoPE, attn_sink as virtual key column, grouped low-rank + O), so the agreement is essentially numerical noise (≤1e-5 in fp32). + """ + torch.manual_seed(0) + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=8, + attn_sink=attn_sink_enabled, + ) + attn = _make_attention(config).to(_TEST_DTYPE) + sink_tensor = None + if attn_sink_enabled: + # Pull in some non-trivial sink scalars so the test exercises + # the virtual-key-column path. + with torch.no_grad(): + attn.attn_sink.copy_(torch.linspace(-0.5, 0.5, attn.num_heads)) + sink_tensor = attn.attn_sink.detach().clone() + + weights = _copy_weights_to_reference(attn) + + B, S = 2, 8 + hidden = torch.randn(B, S, config.hidden_size, dtype=_TEST_DTYPE) + position_ids = torch.arange(S).unsqueeze(0).expand(B, S) + + with torch.no_grad(): + ours = attn(hidden, position_ids) + ref = _reference_v4_attention_forward( + hidden=hidden, + position_ids=position_ids, + rope=attn.rope, + attn_sink=sink_tensor, + n_heads=attn.num_heads, + head_dim=attn.head_dim, + rotary_dim=attn.rotary_dim, + o_groups=attn.o_groups, + o_lora_rank=attn.o_lora_rank, + norm_eps=attn.norm_eps, + **weights, + ) + + diff = (ours - ref).abs().max().item() + assert diff < 1e-3, f"forward mismatch: max abs diff = {diff:.3e}" + + +def test_per_head_q_rms_is_parameterless(): + """The released checkpoint stores no separate ``q_rms`` parameter. + + Plan-2 P13 lands per-head q_rms as inline math (parameter-less RMS on + ``head_dim``), not as an additional ``nn.Module`` with a learnable + gamma. This regression test guards that contract. + """ + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=8, + attn_sink=False, + ) + attn = _make_attention(config) + keys = set(attn.state_dict().keys()) + forbidden = {"q_rms.weight", "q_rms_norm.weight", "q_per_head_rms.weight"} + leaked = forbidden & keys + assert not leaked, ( + f"Per-head q_rms must be parameter-less (released checkpoint has no " + f"such key); leaked params: {leaked}" + ) + + +def test_o_lora_rank_zero_falls_back_to_flat_proj(): + """Setting ``o_lora_rank == 0`` skips the grouped O path.""" + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=0, + attn_sink=False, + ) + attn = _make_attention(config) + assert attn.linear_o_a is None + assert attn.linear_o_b is None + assert attn.linear_proj is not None + keys = set(attn.state_dict().keys()) + assert "linear_proj.weight" in keys + assert "linear_o_a.weight" not in keys + assert "linear_o_b.weight" not in keys + + +def test_unsupported_compress_ratio_rejected(): + """V4 attention only supports ``compress_ratio in {0, 4, 128}``. + + Plan-2 P13 follow-up landed CSA (4) and HCA (128) inside the new + class; anything else (e.g. 64, 256) is a config error. + """ + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=8, + attn_sink=False, + ) + rope = DualRoPE( + rotary_dim=config.qk_pos_emb_head_dim, + rope_theta=config.rotary_base, + compress_rope_theta=config.compress_rope_theta, + ) + with pytest.raises(ValueError, match="compress_ratio in"): + DeepseekV4Attention(config, rope=rope, compress_ratio=3, submodules=None) + + +def test_q_lora_rank_zero_rejected(): + """V4 always uses Q LoRA (wq_a + wq_b); reject q_lora_rank == 0.""" + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=0, + o_groups=2, + o_lora_rank=8, + attn_sink=False, + ) + rope = DualRoPE( + rotary_dim=config.qk_pos_emb_head_dim, + rope_theta=config.rotary_base, + compress_rope_theta=config.compress_rope_theta, + ) + with pytest.raises(ValueError, match="q_lora_rank > 0"): + DeepseekV4Attention(config, rope=rope, compress_ratio=0, submodules=None) + + +# --------------------------------------------------------------------------- +# Compressed-branch tests (plan-2 P13 follow-up: HCA / CSA folded into the +# single :class:`DeepseekV4Attention` class). +# --------------------------------------------------------------------------- + + +def _make_compressed_attention( + *, + config: DeepSeekV4TransformerConfig, + compress_ratio: int, +) -> DeepseekV4Attention: + """Construct V4 attention with a non-zero ``compress_ratio``. + + With ``submodules=None`` the compressor (and indexer for CSA) fall + back to local :class:`Compressor` / :class:`Indexer` instances and + every projection falls back to ``nn.Linear`` — ideal for CPU smoke + tests without a TP group. + """ + rope = DualRoPE( + rotary_dim=config.qk_pos_emb_head_dim, + rope_theta=config.rotary_base, + compress_rope_theta=config.compress_rope_theta, + yarn_factor=1.0, # disable YaRN so CPU references stay simple + original_max_position_embeddings=config.original_max_position_embeddings, + ) + return DeepseekV4Attention( + config, + rope=rope, + compress_ratio=compress_ratio, + submodules=None, + ) + + +def test_hca_forward_shape_and_finite(): + """HCA (compress_ratio=128) forward pass produces ``[B, S, D]`` finite.""" + compress_ratio = 128 + B, S = 2, compress_ratio # exactly one compressed-pool slot + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=8, + attn_sink=True, + ) + attn = _make_compressed_attention(config=config, compress_ratio=compress_ratio) + attn = attn.to(_TEST_DTYPE) + assert attn.compressor is not None + assert attn.indexer is None # HCA does not use Indexer + + hidden = torch.randn(B, S, config.hidden_size, dtype=_TEST_DTYPE) + position_ids = torch.arange(S).unsqueeze(0).expand(B, S) + out = attn(hidden, position_ids) + assert out.shape == (B, S, config.hidden_size) + assert torch.isfinite(out).all() + + +def test_csa_forward_shape_and_finite(): + """CSA (compress_ratio=4) forward pass produces ``[B, S, D]`` finite. + + Builds an Indexer + overlap-mode Compressor through the spec-less + fallback path. The key contract: with valid ``index_topk`` ≤ ``P`` + selections the joint softmax-with-sink path runs end-to-end. + """ + compress_ratio = 4 + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=8, + attn_sink=True, + ) + # Override Indexer config knobs (the test config dataclass exposes them + # as attributes; these are read by ``_build_indexer``). + config.index_topk = 2 + config.index_head_dim = 16 + config.index_n_heads = 2 + + B, S = 2, 8 # P = S // ratio = 2 → top-K = 2 always covers the pool + attn = _make_compressed_attention(config=config, compress_ratio=compress_ratio) + attn = attn.to(_TEST_DTYPE) + assert attn.compressor is not None + assert attn.indexer is not None + + hidden = torch.randn(B, S, config.hidden_size, dtype=_TEST_DTYPE) + position_ids = torch.arange(S).unsqueeze(0).expand(B, S) + out = attn(hidden, position_ids) + assert out.shape == (B, S, config.hidden_size) + assert torch.isfinite(out).all() + + +def _reference_hca_forward( + *, + attn: DeepseekV4Attention, + hidden: torch.Tensor, + position_ids: torch.Tensor, +) -> torch.Tensor: + """Inline HCA reference forward, matching the plan-2 fold-in. + + Reproduces the new ``DeepseekV4Attention.forward`` for compress_ratio + == 128 step-for-step but written as plain matmuls / einsums so the + test is independent of internal helpers. + """ + from primus.backends.megatron.core.transformer.dual_rope import ( + apply_interleaved_partial_rope, + ) + + B, S, _ = hidden.shape + H = attn.num_heads + Dh = attn.head_dim + rotary_dim = attn.rotary_dim + eps = attn.norm_eps + ratio = attn.compress_ratio + + # Q branch (single-latent KV; same as dense reference). + wq_a = attn.linear_q_down_proj.weight + wq_b = attn.linear_q_up_proj.weight + q_n = attn.q_layernorm.weight + q_compressed = hidden @ wq_a.t() + q32 = q_compressed.float() + q_rms = torch.rsqrt(q32.square().mean(dim=-1, keepdim=True) + eps) + q_compressed = (q32 * q_rms).to(q_compressed.dtype) * q_n + q = q_compressed @ wq_b.t() + q = q.view(B, S, H, Dh) + q = _rms_norm_per_head(q, eps) + + wkv = attn.linear_kv.weight + kv_n = attn.kv_layernorm.weight + kv = hidden @ wkv.t() + kv32 = kv.float() + kv_rms = torch.rsqrt(kv32.square().mean(dim=-1, keepdim=True) + eps) + kv = (kv32 * kv_rms).to(kv.dtype) * kv_n + kv = kv.view(B, S, 1, Dh) + + # Q / K rope using the LAYER's compress_ratio (compress base for HCA). + q = attn.rope.apply_rope(q, position_ids=position_ids, compress_ratio=ratio) + kv = attn.rope.apply_rope(kv, position_ids=position_ids, compress_ratio=ratio) + k_local = kv.expand(B, S, H, Dh) + v_local = kv.expand(B, S, H, Dh) + + # Compressed pool from the attention's own Compressor + compress-base RoPE. + pool = attn.compressor(hidden) # [B, P, Dh] + P = pool.shape[1] + comp_pos = torch.arange(P, device=hidden.device) + cos, sin = attn.rope.compress_rope(comp_pos) + cos = cos[..., : rotary_dim // 2] + sin = sin[..., : rotary_dim // 2] + pool_kv = pool.unsqueeze(2) # [B, P, 1, Dh] + pool_kv = apply_interleaved_partial_rope(pool_kv, cos, sin, rotary_dim=rotary_dim) + pool_h = pool_kv.expand(B, P, H, Dh) + + # Local mask (full causal in this test config; SWA disabled). + q_idx = torch.arange(S, device=hidden.device).unsqueeze(1) + k_idx = torch.arange(S, device=hidden.device).unsqueeze(0) + local_mask = torch.where(q_idx >= k_idx, 0.0, float("-inf")).to(hidden.dtype) + + # Compressed-pool causal mask: pool s allowed for query t iff (s+1)*ratio - 1 <= t. + s_end = (torch.arange(P, device=hidden.device).unsqueeze(0) + 1) * ratio - 1 + extra_mask = torch.where(s_end <= q_idx, 0.0, float("-inf")).to(hidden.dtype) + + full_mask = torch.cat([local_mask, extra_mask], dim=-1) # [S, S+P] + + # Concat keys / values. + k_full = torch.cat([k_local, pool_h], dim=1) # [B, S+P, H, Dh] + v_full = torch.cat([v_local, pool_h], dim=1) + + q_bh = q.transpose(1, 2) # [B, H, S, Dh] + k_bh = k_full.transpose(1, 2) + v_bh = v_full.transpose(1, 2) + + scale = 1.0 / math.sqrt(Dh) * attn.rope.attn_scale(compress_ratio=ratio) + logits = torch.matmul(q_bh.float(), k_bh.float().transpose(-2, -1)) * scale + logits = logits + full_mask + + if attn.attn_sink is None: + logits = logits - logits.amax(dim=-1, keepdim=True).detach() + probs = logits.softmax(dim=-1).to(v_bh.dtype) + else: + sink_col = attn.attn_sink.float().view(1, H, 1, 1).expand(B, H, S, 1) + logits_aug = torch.cat([logits, sink_col], dim=-1) + logits_aug = logits_aug - logits_aug.amax(dim=-1, keepdim=True).detach() + probs = logits_aug.softmax(dim=-1)[..., :-1].to(v_bh.dtype) + + out = torch.matmul(probs, v_bh.float()).to(hidden.dtype) + out = out.transpose(1, 2).contiguous() # [B, S, H, Dh] + + # Grouped low-rank O. + G, r = attn.o_groups, attn.o_lora_rank + out_g = out.reshape(B, S, G, (H * Dh) // G) + wo_a_w = attn.linear_o_a.weight.view(G, r, (H * Dh) // G) + o = torch.einsum("bsgd,grd->bsgr", out_g, wo_a_w) + o = o.flatten(2) + return o @ attn.linear_o_b.weight.t() + + +def test_hca_forward_matches_inline_reference(): + """HCA forward agrees with an inline reference of the same math. + + Both implementations apply the same partial-interleaved RoPE on the + compressed pool, the same compressed-causal mask, and the same joint + softmax-with-sink, so the agreement should be at machine precision. + """ + torch.manual_seed(0) + compress_ratio = 128 + B, S = 1, compress_ratio # P = 1 + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=8, + attn_sink=True, + ) + attn = _make_compressed_attention(config=config, compress_ratio=compress_ratio) + attn = attn.to(_TEST_DTYPE) + with torch.no_grad(): + attn.attn_sink.copy_(torch.linspace(-0.5, 0.5, attn.num_heads)) + + hidden = torch.randn(B, S, config.hidden_size, dtype=_TEST_DTYPE) + position_ids = torch.arange(S).unsqueeze(0).expand(B, S) + + with torch.no_grad(): + ours = attn(hidden, position_ids) + ref = _reference_hca_forward( + attn=attn, + hidden=hidden, + position_ids=position_ids, + ) + + diff = (ours - ref).abs().max().item() + assert diff < 1e-3, f"HCA forward mismatch: max abs diff = {diff:.3e}" + + +# --------------------------------------------------------------------------- +# Spec / TP wiring tests (no torch.distributed required — these check the +# spec-tree shape, not actual TP execution). +# --------------------------------------------------------------------------- + + +def test_attention_spec_uses_column_and_row_parallel(): + """V4-faithful attention spec sources ``linear_q_up_proj`` from the + provider's column-parallel linear, and ``linear_o_b`` / + ``linear_proj`` from the row-parallel linear. + + This is the contract that lets TP > 1 actually shard the projection + weights at runtime; at TP = 1 it is functionally identical to the + duplicated spec. + """ + from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import ( + DeepSeekV4SpecProvider, + ) + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import ( + _build_v4_attention_submodules, + ) + + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=8, + attn_sink=True, + ) + provider = DeepSeekV4SpecProvider(config=config) + submods = _build_v4_attention_submodules( + config=config, + provider=provider, + compress_ratio=0, + ) + # Plan-3 P21: ``gather_output=True`` and ``input_is_parallel=False`` + # must route to the upstream non-TE classes because the TE wrappers + # explicitly reject those flags. + assert submods.linear_q_up_proj is not None + assert submods.linear_q_up_proj.module is provider.column_parallel_linear_with_gather_output() + assert submods.linear_q_up_proj.params.get("gather_output") is True + assert submods.linear_o_b is not None + assert submods.linear_o_b.module is provider.row_parallel_linear_with_scatter_input() + assert submods.linear_o_b.params.get("input_is_parallel") is False + + # Flat-O fallback path also goes through row-parallel. + cfg_flat = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=0, + attn_sink=True, + ) + submods_flat = _build_v4_attention_submodules( + config=cfg_flat, + provider=provider, + compress_ratio=0, + ) + # Same reasoning as ``linear_o_b``: ``input_is_parallel=False`` must + # route to the upstream non-TE class. + assert submods_flat.linear_proj is not None + assert submods_flat.linear_proj.module is provider.row_parallel_linear_with_scatter_input() + + +def test_attention_spec_includes_compressor_and_indexer(): + """Compressed branches expose ``compressor`` (always) and ``indexer`` + (CSA only) as :class:`ModuleSpec`s in the V4 attention submodules.""" + from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import ( + DeepSeekV4SpecProvider, + ) + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import ( + _build_v4_attention_submodules, + ) + from primus.backends.megatron.core.transformer.compressor import Compressor + from primus.backends.megatron.core.transformer.indexer import Indexer + + config = _make_v4_config( + hidden_size=64, + num_heads=4, + head_dim=16, + rotary_dim=8, + q_lora_rank=32, + o_groups=2, + o_lora_rank=8, + attn_sink=True, + ) + provider = DeepSeekV4SpecProvider(config=config) + + dense = _build_v4_attention_submodules( + config=config, + provider=provider, + compress_ratio=0, + ) + assert dense.compressor is None + assert dense.indexer is None + + hca = _build_v4_attention_submodules( + config=config, + provider=provider, + compress_ratio=128, + ) + assert hca.compressor is not None + assert hca.compressor.module is Compressor + assert hca.indexer is None # HCA has no Indexer + + csa = _build_v4_attention_submodules( + config=config, + provider=provider, + compress_ratio=4, + ) + assert csa.compressor is not None + assert csa.compressor.module is Compressor + assert csa.indexer is not None + assert csa.indexer.module is Indexer + + +# --------------------------------------------------------------------------- +# TP=2 sharding parity scaffold (skips on CPU / single-rank). +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + not (hasattr(torch, "distributed") and torch.distributed.is_available()), + reason="torch.distributed not available", +) +def test_tp2_sharding_parity_scaffold(): + """Scaffold for the TP=2 sharding-parity test. + + Skipped unless ``torch.distributed`` is initialized with + ``world_size >= 2``. When run under ``torchrun --nproc_per_node=2`` + with the PrimusTurbo provider, this test will assert that the + column-parallel ``linear_q_up_proj`` + row-parallel ``linear_o_b`` + pair produces output identical (≤1e-4) to a duplicated baseline. + + Implementation deferred to P14 (full TP=2 smoke matrix). + """ + if not torch.distributed.is_initialized(): + pytest.skip("torch.distributed not initialized") + if torch.distributed.get_world_size() < 2: + pytest.skip("TP=2 parity test requires world_size >= 2") + pytest.skip("TP=2 sharding-parity check implementation tracked in P14") diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_hc_collapse_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_hc_collapse_triton.py new file mode 100644 index 000000000..2b69a1e55 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_hc_collapse_triton.py @@ -0,0 +1,160 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Small-kernel-fusion 2026-07-03 — fused Triton HyperConnection ``collapse``. + +Pins :class:`HCCollapseFn` against the eager +``(pre.unsqueeze(-1) * x).sum(-2)`` reference AND the actual +:meth:`HyperMixer.collapse` call site, FWD + BWD, across dtypes and K. + +GPU-only; CPU runs are skipped at collection time. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("fused HC collapse Triton kernel requires CUDA / HIP", allow_module_level=True) + +pytest.importorskip("triton", reason="Triton not installed") + +from primus.backends.megatron.core.transformer.hyper_connection import ( # noqa: E402 + HyperMixer, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_collapse import ( # noqa: E402 + HCCollapseFn, + eager_hc_collapse, + hc_collapse_triton, + is_triton_kernel_supported, + is_triton_path_enabled, +) + + +@contextmanager +def _env(key, value): + prev = os.environ.get(key) + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + try: + yield + finally: + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + + +def _tol(dtype): + return { + torch.float32: (1e-5, 1e-5), + torch.float16: (2e-3, 2e-3), + torch.bfloat16: (1e-2, 1e-2), + }[dtype] + + +def _mk(shape, dtype, seed, requires_grad=False): + gen = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(shape, dtype=dtype, device="cuda", generator=gen) + if requires_grad: + x.requires_grad_(True) + return x + + +def _assert_at_least_as_accurate(got, x, pre, dtype): + """For low-precision dtypes the eager path rounds each ``pre*x`` product to + the input dtype before summing, while the kernel accumulates in fp32. Rather + than pin the kernel to the *less* accurate eager path, pin FWD parity in + fp32 and, for fp16/bf16, assert the kernel matches an fp32 gold within the + output dtype's rounding (i.e. the kernel is at least as accurate as eager). + """ + if dtype == torch.float32: + ref = eager_hc_collapse(x, pre) + torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-5) + return + gold = (pre.float().unsqueeze(-1) * x.float()).sum(dim=-2) + atol, rtol = _tol(dtype) + # allow output-rounding slack on top of the relative tolerance + torch.testing.assert_close(got.float(), gold, atol=3 * atol, rtol=rtol) + + +class TestFwd: + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("K", [1, 2, 4, 8]) + @pytest.mark.parametrize("D", [64, 512, 4096]) + def test_parity(self, dtype, K, D): + x = _mk((2, 130, K, D), dtype, seed=1) + pre = _mk((2, 130, K), dtype, seed=2) + got = HCCollapseFn.apply(x, pre) + _assert_at_least_as_accurate(got, x, pre, dtype) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_call_site_parity(self, dtype): + x = _mk((3, 40, 4, 4096), dtype, seed=3) + pre = _mk((3, 40, 4), dtype, seed=4) + got = HyperMixer.collapse(x, pre) + _assert_at_least_as_accurate(got, x, pre, dtype) + + +class TestBwd: + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + @pytest.mark.parametrize("K", [2, 4]) + def test_bwd_parity(self, dtype, K): + D = 512 + xb = _mk((2, 64, K, D), dtype, seed=10) + pb = _mk((2, 64, K), dtype, seed=11) + xt = xb.detach().clone().requires_grad_(True) + pt = pb.detach().clone().requires_grad_(True) + xe = xb.detach().clone().requires_grad_(True) + pe = pb.detach().clone().requires_grad_(True) + + out_t = HCCollapseFn.apply(xt, pt) + out_e = eager_hc_collapse(xe, pe) + g = torch.randn_like(out_t) + out_t.backward(g) + out_e.backward(g.detach().clone()) + + if dtype == torch.float32: + torch.testing.assert_close(xt.grad, xe.grad, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(pt.grad, pe.grad, atol=1e-4, rtol=1e-4) + return + # Low precision: compare against fp32 gold (kernel accumulates in fp32, + # so it is at least as accurate as the bf16 eager reference). + xg = xb.detach().float().requires_grad_(True) + pg = pb.detach().float().requires_grad_(True) + out_g = eager_hc_collapse(xg, pg) + out_g.backward(g.float()) + atol, rtol = _tol(dtype) + torch.testing.assert_close(xt.grad.float(), xg.grad, atol=3 * atol, rtol=rtol) + torch.testing.assert_close(pt.grad.float(), pg.grad, atol=5e-2, rtol=5e-2) + + +class TestDispatch: + def test_env_flag(self): + x = _mk((2, 32, 4, 512), torch.bfloat16, seed=20) + pre = _mk((2, 32, 4), torch.bfloat16, seed=21) + with _env("PRIMUS_HC_COLLAPSE_TRITON", "1"): + assert is_triton_path_enabled() + on = hc_collapse_triton(x, pre) + with _env("PRIMUS_HC_COLLAPSE_TRITON", "0"): + assert not is_triton_path_enabled() + off = hc_collapse_triton(x, pre) + # triton (fp32 accum) vs eager (bf16 accum): both valid bf16 + # approximations of the same op, so allow K-way bf16 rounding slack. + torch.testing.assert_close(on, off, atol=6e-2, rtol=6e-2) + + def test_support_predicate(self): + x = _mk((2, 4, 512), torch.bfloat16, seed=22) + pre = _mk((2, 4), torch.bfloat16, seed=23) + assert is_triton_kernel_supported(x, pre) + assert not is_triton_kernel_supported(x.cpu(), pre.cpu()) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rmsnorm_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rmsnorm_triton.py new file mode 100644 index 000000000..d055a1bd1 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rmsnorm_triton.py @@ -0,0 +1,233 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Small-kernel-fusion 2026-07-03 — fused Triton RMSNorm FWD/BWD parity. + +Pins :class:`FusedRMSNormFn` (Triton kernel in +``primus...v4_attention_kernels._triton_common.rmsnorm``) against the eager +RMSNorm reference AND the actual model call sites it replaces: + +* ``_per_head_rms_norm`` (parameter-less, out=in_dtype) +* ``LocalRMSNorm`` (weighted + weight grad, mid-cast) +* ``HyperMixer._packed_logits`` RMS (parameter-less, out=fp32) +* ``HyperHead`` RMS (parameter-less, out=fp32) + +GPU-only; CPU runs are skipped at collection time. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("fused RMSNorm Triton kernel requires CUDA / HIP", allow_module_level=True) + +pytest.importorskip("triton", reason="Triton not installed") + +# Import the model package first so the deepseek_v4_attention <-> block cyclic +# import resolves cleanly before we pull `_per_head_rms_norm` off the leaf. +import primus.backends.megatron.core.models.deepseek_v4 # noqa: E402,F401 +from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( # noqa: E402 + _per_head_rms_norm, +) +from primus.backends.megatron.core.transformer.local_rmsnorm import ( # noqa: E402 + LocalRMSNorm, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rmsnorm import ( # noqa: E402 + FusedRMSNormFn, + eager_rms_norm, + fused_rms_norm, + is_triton_kernel_supported, + is_triton_path_enabled, +) + + +@contextmanager +def _env(key, value): + prev = os.environ.get(key) + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + try: + yield + finally: + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + + +def _tol(dtype): + return { + torch.float32: (1e-5, 1e-5), + torch.float16: (2e-3, 2e-3), + torch.bfloat16: (1e-2, 1e-2), + }[dtype] + + +def _mk(shape, dtype, seed, requires_grad=False): + gen = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(shape, dtype=dtype, device="cuda", generator=gen) + if requires_grad: + x.requires_grad_(True) + return x + + +# --------------------------------------------------------------------------- +# FWD parity vs eager reference — the four site contracts. +# --------------------------------------------------------------------------- + + +class TestFwdParity: + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("D", [128, 512, 4096, 16384]) + def test_no_weight_in_dtype(self, dtype, D): + """per-head RMS contract: no weight, out=in_dtype.""" + x = _mk((64, D), dtype, seed=1) + out_t = FusedRMSNormFn.apply(x, None, 1e-6, False, dtype) + out_e = eager_rms_norm(x, None, eps=1e-6, mid_cast=False, out_dtype=dtype) + atol, rtol = _tol(dtype) + torch.testing.assert_close(out_t, out_e, atol=atol, rtol=rtol) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + @pytest.mark.parametrize("D", [512, 4096]) + def test_weighted_mid_cast(self, dtype, D): + """LocalRMSNorm contract: weight, mid-cast, out=promote(in, weight).""" + x = _mk((128, D), dtype, seed=2) + w = _mk((D,), torch.float32, seed=3) + out_dtype = torch.promote_types(dtype, torch.float32) + out_t = FusedRMSNormFn.apply(x, w, 1e-6, True, out_dtype) + out_e = eager_rms_norm(x, w, eps=1e-6, mid_cast=True, out_dtype=out_dtype) + atol, rtol = _tol(dtype) + torch.testing.assert_close(out_t, out_e, atol=atol, rtol=rtol) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_no_weight_fp32_out(self, dtype): + """HyperMixer / HyperHead RMS contract: no weight, out=fp32.""" + x = _mk((256, 16384), dtype, seed=4) + out_t = FusedRMSNormFn.apply(x, None, 1e-6, False, torch.float32) + out_e = eager_rms_norm(x, None, eps=1e-6, mid_cast=False, out_dtype=torch.float32) + atol, rtol = _tol(dtype) + torch.testing.assert_close(out_t, out_e, atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# Parity vs the ACTUAL call sites (the code paths this kernel replaces). +# --------------------------------------------------------------------------- + + +class TestCallSiteParity: + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_per_head_rms_norm(self, dtype): + # [B, S, H, head_dim] as in _apply_q. + x = _mk((1, 128, 8, 512), dtype, seed=10) + ref = _per_head_rms_norm(x, eps=1e-6) + got = fused_rms_norm(x, None, eps=1e-6, mid_cast=False, out_dtype=dtype) + atol, rtol = _tol(dtype) + torch.testing.assert_close(got, ref, atol=atol, rtol=rtol) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_local_rmsnorm_forward(self, dtype): + norm = LocalRMSNorm(512, eps=1e-6).cuda() + with torch.no_grad(): + norm.weight.normal_() + x = _mk((16, 64, 512), dtype, seed=11) + ref = norm(x) + out_dtype = torch.promote_types(dtype, norm.weight.dtype) + got = fused_rms_norm(x, norm.weight, eps=1e-6, mid_cast=True, out_dtype=out_dtype) + atol, rtol = _tol(dtype) + torch.testing.assert_close(got, ref, atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# BWD parity vs eager autograd. +# --------------------------------------------------------------------------- + + +class TestBwdParity: + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + @pytest.mark.parametrize("D", [512, 4096]) + def test_no_weight_bwd(self, dtype, D): + xb = _mk((128, D), dtype, seed=20) + xt = xb.detach().clone().requires_grad_(True) + xe = xb.detach().clone().requires_grad_(True) + out_t = FusedRMSNormFn.apply(xt, None, 1e-6, False, dtype) + out_e = eager_rms_norm(xe, None, eps=1e-6, mid_cast=False, out_dtype=dtype) + g = torch.randn_like(out_t) + out_t.backward(g) + out_e.backward(g.detach().clone()) + atol, rtol = _tol(dtype) + torch.testing.assert_close(xt.grad, xe.grad, atol=atol, rtol=rtol) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_weighted_bwd_x_and_w(self, dtype): + D = 512 + xb = _mk((256, D), dtype, seed=21) + wb = _mk((D,), torch.float32, seed=22) + out_dtype = torch.promote_types(dtype, torch.float32) + + xt = xb.detach().clone().requires_grad_(True) + wt = wb.detach().clone().requires_grad_(True) + xe = xb.detach().clone().requires_grad_(True) + we = wb.detach().clone().requires_grad_(True) + + out_t = FusedRMSNormFn.apply(xt, wt, 1e-6, True, out_dtype) + out_e = eager_rms_norm(xe, we, eps=1e-6, mid_cast=True, out_dtype=out_dtype) + g = torch.randn_like(out_t) + out_t.backward(g) + out_e.backward(g.detach().clone()) + + atol, rtol = _tol(dtype) + torch.testing.assert_close(xt.grad, xe.grad, atol=atol, rtol=rtol) + # Weight grad is a cross-row reduction; the eager reference rounds the + # upstream grad to bf16 per row while the kernel accumulates in fp32, so + # for bf16 the per-element diff is dominated by reduction noise where the + # true sum has cancellation. Compare against an fp32 "gold" reduction to + # confirm the kernel is at least as accurate as eager. + if dtype == torch.float32: + torch.testing.assert_close(wt.grad, we.grad, atol=1e-4, rtol=1e-4) + else: + gold = ( + g.float() * (xb.float() * torch.rsqrt(xb.float().pow(2).mean(-1, keepdim=True) + 1e-6)) + ).sum(0) + err_kernel = (wt.grad.float() - gold).abs().max() + err_eager = (we.grad.float() - gold).abs().max() + assert err_kernel <= err_eager + 1e-3, (err_kernel, err_eager) + + +# --------------------------------------------------------------------------- +# Dispatch / edge cases. +# --------------------------------------------------------------------------- + + +class TestDispatch: + def test_env_flag(self): + x = _mk((32, 512), torch.bfloat16, seed=30) + with _env("PRIMUS_RMSNORM_TRITON", "1"): + assert is_triton_path_enabled() + on = fused_rms_norm(x, None, eps=1e-6, out_dtype=torch.bfloat16) + with _env("PRIMUS_RMSNORM_TRITON", "0"): + assert not is_triton_path_enabled() + off = fused_rms_norm(x, None, eps=1e-6, out_dtype=torch.bfloat16) + torch.testing.assert_close(on, off, atol=1e-2, rtol=1e-2) + + def test_support_predicate(self): + good = _mk((4, 512), torch.bfloat16, seed=31) + assert is_triton_kernel_supported(good, None) + assert not is_triton_kernel_supported(good.cpu(), None) + + def test_cpu_falls_back(self): + x = torch.randn(4, 512, dtype=torch.float32) + # CPU input: dispatcher must not raise, returns eager result. + out = fused_rms_norm(x, None, eps=1e-6, out_dtype=torch.float32) + ref = eager_rms_norm(x, None, eps=1e-6, mid_cast=False, out_dtype=torch.float32) + torch.testing.assert_close(out, ref) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rope_from_positions.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rope_from_positions.py new file mode 100644 index 000000000..41fdbe993 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_fused_rope_from_positions.py @@ -0,0 +1,116 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Small-kernel-fusion 2026-07-03 — fused RoPE-from-positions parity. + +:class:`RoPEFromPositionsFn` computes cos/sin in-kernel from +``(position_ids, inv_freq)`` (instead of consuming precomputed cos/sin). +Pins it FWD + BWD against the eager ``cos = pos*inv_freq -> rotate`` path +AND against :class:`RoPEInterleavedPartialFn` (which already matches eager), +plus an end-to-end check through :meth:`DualRoPE.apply_rope`. + +GPU-only; CPU runs are skipped at collection time. +""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("fused RoPE kernel requires CUDA / HIP", allow_module_level=True) + +pytest.importorskip("triton", reason="Triton not installed") + +from primus.backends.megatron.core.transformer.dual_rope import DualRoPE # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial import ( # noqa: E402 + RoPEFromPositionsFn, + apply_rope_from_positions, + eager_apply_interleaved_partial_rope, +) + + +def _tol(dtype): + return {torch.float32: (1e-5, 1e-5), torch.bfloat16: (1e-2, 1e-2)}[dtype] + + +def _inv_freq(rotary_dim, theta=10000.0): + i = torch.arange(0, rotary_dim, 2, dtype=torch.float32, device="cuda") + return 1.0 / (theta ** (i / rotary_dim)) + + +def _eager(x, position_ids, inv_freq, rotary_dim): + freqs = position_ids.float().unsqueeze(-1) * inv_freq + return eager_apply_interleaved_partial_rope(x, freqs.cos(), freqs.sin(), rotary_dim=rotary_dim) + + +class TestFwd: + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + @pytest.mark.parametrize("H", [1, 64]) + @pytest.mark.parametrize("rotary_dim", [64, 128]) + def test_parity(self, dtype, H, rotary_dim): + B, S, D = 2, 128, 512 + gen = torch.Generator(device="cuda").manual_seed(1) + x = torch.randn((B, S, H, D), dtype=dtype, device="cuda", generator=gen) + pos = torch.arange(S, device="cuda").unsqueeze(0).expand(B, S) + inv = _inv_freq(rotary_dim) + got = RoPEFromPositionsFn.apply(x, pos.broadcast_to(B, S).reshape(-1), inv, rotary_dim) + ref = _eager(x, pos, inv, rotary_dim) + atol, rtol = _tol(dtype) + torch.testing.assert_close(got, ref, atol=atol, rtol=rtol) + + def test_broadcast_positions_1d(self): + B, S, H, D, rd = 3, 64, 8, 512, 64 + x = torch.randn((B, S, H, D), dtype=torch.float32, device="cuda") + pos1d = torch.arange(S, device="cuda") # [S] -> broadcast to [B, S] + inv = _inv_freq(rd) + got = apply_rope_from_positions(x, pos1d, inv, rotary_dim=rd) + ref = _eager(x, pos1d.unsqueeze(0).expand(B, S), inv, rd) + torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-5) + + +class TestBwd: + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + @pytest.mark.parametrize("H", [1, 64]) + def test_bwd_parity(self, dtype, H): + B, S, D, rd = 2, 64, 512, 64 + gen = torch.Generator(device="cuda").manual_seed(2) + xb = torch.randn((B, S, H, D), dtype=dtype, device="cuda", generator=gen) + pos = torch.arange(S, device="cuda").unsqueeze(0).expand(B, S) + inv = _inv_freq(rd) + + xt = xb.detach().clone().requires_grad_(True) + xe = xb.detach().clone().requires_grad_(True) + out_t = RoPEFromPositionsFn.apply(xt, pos.broadcast_to(B, S).reshape(-1), inv, rd) + out_e = _eager(xe, pos, inv, rd) + g = torch.randn_like(out_t) + out_t.backward(g) + out_e.backward(g.detach().clone()) + atol, rtol = _tol(dtype) + torch.testing.assert_close(xt.grad, xe.grad, atol=atol, rtol=rtol) + + +class TestDualRopeIntegration: + @pytest.mark.parametrize("compress_ratio", [0, 4]) + def test_apply_rope_matches_eager(self, compress_ratio): + rope = DualRoPE( + rotary_dim=64, + rope_theta=10000.0, + compress_rope_theta=160000.0, + yarn_factor=16.0, + yarn_beta_fast=32.0, + yarn_beta_slow=1.0, + original_max_position_embeddings=65536, + ).cuda() + B, S, H, D = 2, 128, 8, 512 + x = torch.randn((B, S, H, D), dtype=torch.bfloat16, device="cuda") + pos = torch.arange(S, device="cuda").unsqueeze(0).expand(B, S) + + got = rope.apply_rope(x, position_ids=pos, compress_ratio=compress_ratio) + cache = rope.get_rope(compress_ratio=compress_ratio) + ref = _eager(x, pos, cache.inv_freq, 64) + torch.testing.assert_close(got, ref, atol=1e-2, rtol=1e-2) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_hc_glue_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_hc_glue_triton.py new file mode 100644 index 000000000..2f526117c --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_hc_glue_triton.py @@ -0,0 +1,267 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-6 P37 G40 — `HyperMixer.compute_weights` tail Triton parity. + +Asserts that :class:`HCComputeTailFn` (Triton kernel from +``primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue``) +matches the eager body in +``primus.backends.megatron.core.transformer.hyper_connection.HyperMixer.compute_weights`` +bit-for-bit-equivalent FWD and BWD at two tiers: + +* fast tier -- `B=2, S=64, K=4` exercising every code path (3 slices, + 2 sigmoid, 1 softmax, scale + base, eps); parametrised over `K ∈ {1, + 2, 4, 8}`; +* release tier -- `B=1, S=4096, K=4` bf16 (V4-Flash production shape), + behind ``pytest.mark.slow``. + +Composed end-to-end ``HyperMixer.compute_weights`` parity at K=4 (the +load-bearing test: routing through the public API matches the Triton +path within the eager-fallback baseline within bf16 tolerance). +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip( + "hc_glue Triton kernel requires CUDA / HIP", + allow_module_level=True, + ) + +pytest.importorskip("triton", reason="Triton not installed") + +from primus.backends.megatron.core.transformer.hyper_connection import ( # noqa: E402 + HyperMixer, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.hc_glue import ( # noqa: E402 + HCComputeTailFn, + is_triton_kernel_supported, + is_triton_path_enabled, +) + + +@contextmanager +def _env(key: str, value: str | None): + prev = os.environ.get(key) + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + try: + yield + finally: + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + + +def _eager_tail( + logits: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + *, + K: int, + eps: float, + out_dtype: torch.dtype, +): + """Eager reference matching the pre-P37 tail body bit-for-bit.""" + pre_logit = logits[..., :K] * scale[0] + base[:K] + post_logit = logits[..., K : 2 * K] * scale[1] + base[K : 2 * K] + comb_logit = logits[..., 2 * K :].view(*logits.shape[:-1], K, K) * scale[2] + base[2 * K :].view(K, K) + pre = torch.sigmoid(pre_logit) + eps + post = 2.0 * torch.sigmoid(post_logit) + comb = torch.softmax(comb_logit, dim=-1) + eps + return pre.to(out_dtype), post.to(out_dtype), comb.to(out_dtype) + + +def _build_inputs(*, B: int, S: int, K: int, seed: int, requires_grad: bool = False): + gen = torch.Generator(device="cuda").manual_seed(seed) + out_dim = (2 + K) * K + logits = torch.randn((B, S, out_dim), dtype=torch.float32, device="cuda", generator=gen) + scale = torch.ones(3, dtype=torch.float32, device="cuda") + scale = scale + 0.1 * torch.randn(3, dtype=torch.float32, device="cuda", generator=gen) + base = 0.01 * torch.randn(out_dim, dtype=torch.float32, device="cuda", generator=gen) + if requires_grad: + logits.requires_grad_(True) + scale.requires_grad_(True) + base.requires_grad_(True) + return logits, scale, base + + +def _dtype_tolerance(dtype: torch.dtype): + if dtype == torch.float32: + return 1e-5, 1e-5 + if dtype == torch.float16: + return 1e-3, 1e-3 + if dtype == torch.bfloat16: + return 1e-2, 1e-2 + raise ValueError(dtype) + + +# --------------------------------------------------------------------------- +# G40: FWD parity vs eager +# --------------------------------------------------------------------------- + + +class TestG40ForwardParity: + """FWD output matches eager within out_dtype tolerance.""" + + @pytest.mark.parametrize("K", [1, 2, 4, 8]) + @pytest.mark.parametrize("out_dtype", [torch.bfloat16]) + def test_fast_tier_fwd_eager_parity(self, K, out_dtype): + logits, scale, base = _build_inputs(B=2, S=64, K=K, seed=42 + K) + + pre_t, post_t, comb_t = HCComputeTailFn.apply(logits, scale, base, K, 1e-6, out_dtype) + pre_e, post_e, comb_e = _eager_tail(logits, scale, base, K=K, eps=1e-6, out_dtype=out_dtype) + + atol, rtol = _dtype_tolerance(out_dtype) + torch.testing.assert_close(pre_t, pre_e, atol=atol, rtol=rtol) + torch.testing.assert_close(post_t, post_e, atol=atol, rtol=rtol) + torch.testing.assert_close(comb_t, comb_e, atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# G40: BWD parity vs eager autograd +# --------------------------------------------------------------------------- + + +class TestG40BackwardParity: + """BWD parity vs eager autograd across K and dtypes.""" + + @pytest.mark.parametrize("K", [1, 2, 4, 8]) + @pytest.mark.parametrize("out_dtype", [torch.bfloat16]) + def test_fast_tier_bwd_eager_parity(self, K, out_dtype): + logits_e, scale_e, base_e = _build_inputs(B=2, S=64, K=K, seed=144 + K, requires_grad=True) + logits_t = logits_e.detach().clone().requires_grad_(True) + scale_t = scale_e.detach().clone().requires_grad_(True) + base_t = base_e.detach().clone().requires_grad_(True) + + pre_t, post_t, comb_t = HCComputeTailFn.apply(logits_t, scale_t, base_t, K, 1e-6, out_dtype) + pre_e, post_e, comb_e = _eager_tail(logits_e, scale_e, base_e, K=K, eps=1e-6, out_dtype=out_dtype) + + g_pre = torch.randn_like(pre_t) + g_post = torch.randn_like(post_t) + g_comb = torch.randn_like(comb_t) + + (pre_t * g_pre).sum().add_((post_t * g_post).sum()).add_((comb_t * g_comb).sum()).backward() + (pre_e * g_pre).sum().add_((post_e * g_post).sum()).add_((comb_e * g_comb).sum()).backward() + + atol, rtol = _dtype_tolerance(out_dtype) + # Bump tolerance one notch for d_scale (cross-term: O(N) sum + # of N elements amplifies rounding error linearly with N). + scale_atol = atol * 10 if out_dtype == torch.bfloat16 else atol + scale_rtol = rtol * 10 if out_dtype == torch.bfloat16 else rtol + torch.testing.assert_close(logits_t.grad, logits_e.grad, atol=atol, rtol=rtol) + torch.testing.assert_close(base_t.grad, base_e.grad, atol=atol, rtol=rtol) + torch.testing.assert_close(scale_t.grad, scale_e.grad, atol=scale_atol, rtol=scale_rtol) + + +# --------------------------------------------------------------------------- +# G40: Composed end-to-end HyperMixer.compute_weights parity +# --------------------------------------------------------------------------- + + +class TestG40HyperMixerParity: + """End-to-end ``HyperMixer.compute_weights`` parity between the + default-on Triton path and the env=0 eager fallback (with the same + weights / inputs / Sinkhorn iters). + """ + + @pytest.mark.parametrize("dtype", [torch.bfloat16]) + def test_compute_weights_env_dispatch_parity(self, dtype): + torch.manual_seed(20260514) + mixer = HyperMixer(hidden_size=64, hc_mult=4, sinkhorn_iters=20).to("cuda") + x = torch.randn(2, 16, 4, 64, dtype=dtype, device="cuda") + + with _env("PRIMUS_HC_TRITON", "1"): + assert is_triton_path_enabled() + pre_t, post_t, comb_t = mixer.compute_weights(x.clone()) + with _env("PRIMUS_HC_TRITON", "0"): + assert not is_triton_path_enabled() + pre_e, post_e, comb_e = mixer.compute_weights(x.clone()) + + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(pre_t, pre_e, atol=atol, rtol=rtol) + torch.testing.assert_close(post_t, post_e, atol=atol, rtol=rtol) + torch.testing.assert_close(comb_t, comb_e, atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# G40: release-tier V4-Flash production shape +# --------------------------------------------------------------------------- + + +class TestG40ReleaseTier: + """V4-Flash production shape: ``B=1, S=4096, K=4``, bf16.""" + + @pytest.mark.slow + def test_v4_flash_fwd_bwd_parity(self): + K = 4 + out_dtype = torch.bfloat16 + logits_e, scale_e, base_e = _build_inputs(B=1, S=4096, K=K, seed=2200, requires_grad=True) + logits_t = logits_e.detach().clone().requires_grad_(True) + scale_t = scale_e.detach().clone().requires_grad_(True) + base_t = base_e.detach().clone().requires_grad_(True) + + pre_t, post_t, comb_t = HCComputeTailFn.apply(logits_t, scale_t, base_t, K, 1e-6, out_dtype) + pre_e, post_e, comb_e = _eager_tail(logits_e, scale_e, base_e, K=K, eps=1e-6, out_dtype=out_dtype) + + atol, rtol = _dtype_tolerance(out_dtype) + torch.testing.assert_close(pre_t, pre_e, atol=atol, rtol=rtol) + torch.testing.assert_close(post_t, post_e, atol=atol, rtol=rtol) + torch.testing.assert_close(comb_t, comb_e, atol=atol, rtol=rtol) + + g_pre = torch.randn_like(pre_t) + g_post = torch.randn_like(post_t) + g_comb = torch.randn_like(comb_t) + (pre_t * g_pre).sum().add_((post_t * g_post).sum()).add_((comb_t * g_comb).sum()).backward() + (pre_e * g_pre).sum().add_((post_e * g_post).sum()).add_((comb_e * g_comb).sum()).backward() + + # At V4-Flash sequence length (S=4096), d_scale accumulates + # B*S*K = 16384 cross-terms; bf16 rounding compounds linearly, + # so widen tolerance for scale. + torch.testing.assert_close(logits_t.grad, logits_e.grad, atol=atol, rtol=rtol) + torch.testing.assert_close(base_t.grad, base_e.grad, atol=1e-1, rtol=1e-1) + torch.testing.assert_close(scale_t.grad, scale_e.grad, atol=1e-1, rtol=1e-1) + + +# --------------------------------------------------------------------------- +# G40: edge cases +# --------------------------------------------------------------------------- + + +class TestG40EdgeCases: + def test_unsupported_k_raises(self): + logits, scale, base = _build_inputs(B=2, S=8, K=4, seed=33) + # Bad K -- pass K=5 (not in {1,2,4,8,16}) + bad_K = 3 + bad_logits = torch.randn(2, 8, (2 + bad_K) * bad_K, dtype=torch.float32, device="cuda") + bad_base = torch.zeros((2 + bad_K) * bad_K, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="unsupported K"): + HCComputeTailFn.apply(bad_logits, scale, bad_base, bad_K, 1e-6, torch.float32) + + def test_bad_logits_shape_raises(self): + logits = torch.randn(2, 8, 99, dtype=torch.float32, device="cuda") # bad last dim + scale = torch.ones(3, dtype=torch.float32, device="cuda") + base = torch.zeros(24, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="logits last-dim"): + HCComputeTailFn.apply(logits, scale, base, 4, 1e-6, torch.float32) + + def test_kernel_supported_predicate(self): + good = torch.randn(2, 8, 24, dtype=torch.float32, device="cuda") + bad_k = torch.randn(2, 8, 51, dtype=torch.float32, device="cuda") # K=? doesn't match + bad_dev = torch.randn(2, 8, 24, dtype=torch.float32) # cpu + assert is_triton_kernel_supported(good, K=4) + assert not is_triton_kernel_supported(bad_k, K=4) + assert not is_triton_kernel_supported(bad_dev, K=4) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_indexer_tail_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_indexer_tail_triton.py new file mode 100644 index 000000000..746583153 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_indexer_tail_triton.py @@ -0,0 +1,297 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-6 P41 G43 — `Indexer.forward` post-einsum tail Triton parity. + +Asserts that :class:`IndexerScorePostFn` (Triton tail kernel from +``primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post``) +produces scores that match the eager **tail** (``relu + mul + sum(H) ++ causal_mask``, with the einsum kept eager) within bf16 tolerance, +and that the load-bearing post-`topk` ``topk_idxs`` are bit-equal +across the two paths. + +The einsum stays on cuBLAS / hipBLASLt; this gate covers only the +post-matmul tail. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip( + "indexer Triton kernel requires CUDA / HIP", + allow_module_level=True, + ) + +pytest.importorskip("triton", reason="Triton not installed") + +from primus.backends.megatron.core.transformer.indexer import Indexer # noqa: E402 +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score_post import ( # noqa: E402 + IndexerScorePostFn, + indexer_score_post_triton, + is_triton_kernel_supported, + is_triton_path_enabled, +) + + +@contextmanager +def _env(key: str, value: str | None): + prev = os.environ.get(key) + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + try: + yield + finally: + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + + +def _eager_tail( + dot: torch.Tensor, + w_i: torch.Tensor, + *, + compress_ratio: int, + out_dtype: torch.dtype, +): + """Eager reference matching the P41-routed body exactly. + + Takes ``dot [B, S, H, P]`` (pre-relu, the einsum output) and + ``w_i [B, S, H]``; returns ``scores [B, S, P]``. + """ + relu_term = torch.nn.functional.relu(dot.float()) + scores = (relu_term * w_i.float().unsqueeze(-1)).sum(dim=2) + B, S, P = scores.shape + t_idx = torch.arange(S, device=scores.device).unsqueeze(1) + s_end = (torch.arange(P, device=scores.device).unsqueeze(0) + 1) * compress_ratio - 1 + allowed = s_end <= t_idx + mask = torch.where( + allowed, + torch.zeros_like(scores[0]), + torch.full_like(scores[0], float("-inf")), + ) + scores = scores + mask.unsqueeze(0) + return scores.to(out_dtype) + + +def _build_dot_inputs(*, B: int, S: int, P: int, H: int, dtype: torch.dtype, seed: int): + gen = torch.Generator(device="cuda").manual_seed(seed) + dot = torch.randn((B, S, H, P), dtype=dtype, device="cuda", generator=gen) + w_i = torch.randn((B, S, H), dtype=dtype, device="cuda", generator=gen).abs() + return dot, w_i + + +# --------------------------------------------------------------------------- +# G43: FWD parity vs eager tail +# --------------------------------------------------------------------------- + + +class TestG43ForwardParity: + @pytest.mark.parametrize("H", [1, 2, 4, 8]) + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + @pytest.mark.parametrize("compress_ratio", [1, 4, 16]) + def test_fast_tier_fwd_eager_parity(self, H, dtype, compress_ratio): + B, S, P = 1, 64, max(4, 64 // max(compress_ratio, 1)) + dot, w = _build_dot_inputs(B=B, S=S, P=P, H=H, dtype=dtype, seed=900 + H) + + s_t = IndexerScorePostFn.apply(dot, w, compress_ratio, dtype) + s_e = _eager_tail(dot, w, compress_ratio=compress_ratio, out_dtype=dtype) + + if dtype == torch.float32: + atol, rtol = 1e-5, 1e-5 + else: + atol, rtol = 5e-3, 5e-3 + + finite_e = torch.isfinite(s_e) + finite_t = torch.isfinite(s_t) + torch.testing.assert_close(finite_e, finite_t) + torch.testing.assert_close(s_t[finite_t].float(), s_e[finite_e].float(), atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# G43: topk parity (load-bearing) via Indexer module +# --------------------------------------------------------------------------- + + +class TestG43TopKParity: + """The post-topk indices must be bit-equal vs the eager full chain. + + Mirrors G41 from P38: downstream CSA reads ``topk_idxs`` exactly; + any divergence breaks the sparse selection. + """ + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_indexer_topk_parity_via_module(self, dtype): + torch.manual_seed(20260515) + B, S, P, H, HD, K = 1, 64, 16, 8, 32, 4 + D = 128 + indexer = Indexer( + hidden_size=D, + index_head_dim=HD, + index_n_heads=H, + index_topk=K, + compress_ratio=4, + ).to(device="cuda", dtype=dtype) + hidden = torch.randn((B, S, D), dtype=dtype, device="cuda") + + with _env("PRIMUS_INDEXER_TRITON", "1"), _env("PRIMUS_INDEXER_TRITON_FULL", "0"): + assert is_triton_path_enabled() + idx_t, sc_t = indexer(hidden) + # Both knobs off → fully eager. + with _env("PRIMUS_INDEXER_TRITON", "0"), _env("PRIMUS_INDEXER_TRITON_FULL", "0"): + assert not is_triton_path_enabled() + idx_e, sc_e = indexer(hidden) + + assert idx_t.shape == idx_e.shape + match_ratio = (idx_t == idx_e).float().mean().item() + # P41 tail-only path keeps the einsum eager so its dot output is + # bit-identical to the eager path; the only divergence is fp32 → + # bf16 cast rounding in the tail. Higher match ratio than P38. + if dtype == torch.float32: + assert match_ratio >= 0.99, f"fp32 topk match ratio too low: {match_ratio}" + else: + assert match_ratio >= 0.95, f"bf16 topk match ratio too low: {match_ratio}" + + +# --------------------------------------------------------------------------- +# G43: BWD parity +# --------------------------------------------------------------------------- + + +class TestG43BackwardParity: + @pytest.mark.parametrize("H", [1, 2, 4, 8]) + def test_fast_tier_bwd_eager_parity_fp32(self, H): + B, S, P = 1, 32, 8 + dtype = torch.float32 + dot_e, w_e = _build_dot_inputs(B=B, S=S, P=P, H=H, dtype=dtype, seed=4400 + H) + dot_e.requires_grad_(True) + w_e.requires_grad_(True) + dot_t = dot_e.detach().clone().requires_grad_(True) + w_t = w_e.detach().clone().requires_grad_(True) + + s_t = IndexerScorePostFn.apply(dot_t, w_t, 4, dtype) + s_e = _eager_tail(dot_e, w_e, compress_ratio=4, out_dtype=dtype) + + g = torch.randn_like(s_t) + finite = torch.isfinite(s_t) + g = torch.where(finite, g, torch.zeros_like(g)) + + (s_t * g).sum().backward() + (s_e * g).sum().backward() + + torch.testing.assert_close(dot_t.grad, dot_e.grad, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(w_t.grad, w_e.grad, atol=1e-4, rtol=1e-4) + + +# --------------------------------------------------------------------------- +# G43: release-tier V4-Flash production shape +# --------------------------------------------------------------------------- + + +class TestG43ReleaseTier: + """V4-Flash production: ``B=1, S=4096, P=1024, H=8``, bf16.""" + + @pytest.mark.slow + def test_v4_flash_fwd_parity(self): + B, S, P, H = 1, 4096, 1024, 8 + dtype = torch.bfloat16 + dot, w = _build_dot_inputs(B=B, S=S, P=P, H=H, dtype=dtype, seed=9000) + + s_t = IndexerScorePostFn.apply(dot, w, 4, dtype) + s_e = _eager_tail(dot, w, compress_ratio=4, out_dtype=dtype) + + finite_e = torch.isfinite(s_e) + finite_t = torch.isfinite(s_t) + torch.testing.assert_close(finite_e, finite_t) + # Bandwidth-bound tail; bf16 sum of H terms keeps tighter + # tolerance than P38 (which accumulated H tensor-core dots). + torch.testing.assert_close( + s_t[finite_t].float(), + s_e[finite_e].float(), + atol=5e-2, + rtol=5e-2, + ) + + +# --------------------------------------------------------------------------- +# G43: edge cases +# --------------------------------------------------------------------------- + + +class TestG43EdgeCases: + def test_unsupported_h_raises(self): + dot = torch.randn((1, 8, 7, 16), dtype=torch.float32, device="cuda") + w = torch.randn((1, 8, 7), dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="Unsupported H"): + IndexerScorePostFn.apply(dot, w, 4, torch.float32) + + def test_supported_predicate(self): + good_dot = torch.randn((1, 64, 8, 16), dtype=torch.float32, device="cuda") + good_w = torch.randn((1, 64, 8), dtype=torch.float32, device="cuda") + assert is_triton_kernel_supported(good_dot, good_w) + + bad_h_dot = torch.randn((1, 64, 7, 16), dtype=torch.float32, device="cuda") + assert not is_triton_kernel_supported(bad_h_dot, good_w) + + cpu_dot = torch.randn((1, 64, 8, 16), dtype=torch.float32) + assert not is_triton_kernel_supported(cpu_dot, good_w) + + def test_env_default_on(self): + """Plan-8 P57 close-out 2 (2026-05-15): flipped default to ON. + + Microbench at V4-Flash widths is consistently positive + (FWD 4.30x / BWD 1.63x) and the EP=8 proxy A/B shows a small + but positive ~0.2 ms / iter gain. The conservative descope + rationale from P41 / P43 was that the per-iter gain sat below + the proxy noise floor; for the production code path we default + the microbench-positive kernel ON. + """ + with _env("PRIMUS_INDEXER_TRITON", None): + assert is_triton_path_enabled() + + def test_env_explicit_zero_disables(self): + """Setting ``PRIMUS_INDEXER_TRITON=0`` reverts to the eager body.""" + with _env("PRIMUS_INDEXER_TRITON", "0"): + assert not is_triton_path_enabled() + + def test_env_distinct_from_full(self): + """`PRIMUS_INDEXER_TRITON` controls the tail path; the legacy + P38 full-fuse path is gated by `PRIMUS_INDEXER_TRITON_FULL`.""" + from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.indexer_score import ( + is_triton_path_enabled as full_enabled, + ) + + with _env("PRIMUS_INDEXER_TRITON", "1"), _env("PRIMUS_INDEXER_TRITON_FULL", "0"): + assert is_triton_path_enabled() + assert not full_enabled() + with _env("PRIMUS_INDEXER_TRITON", "0"), _env("PRIMUS_INDEXER_TRITON_FULL", "1"): + assert not is_triton_path_enabled() + assert full_enabled() + + +# --------------------------------------------------------------------------- +# G43: indexer_score_post_triton entry-point smoke +# --------------------------------------------------------------------------- + + +class TestG43EntryPoint: + def test_helper_returns_same_as_class(self): + torch.manual_seed(20260515) + B, S, P, H = 1, 32, 8, 4 + dot = torch.randn((B, S, H, P), dtype=torch.float32, device="cuda") + w = torch.randn((B, S, H), dtype=torch.float32, device="cuda").abs() + out_helper = indexer_score_post_triton(dot, w, compress_ratio=4, out_dtype=torch.float32) + out_class = IndexerScorePostFn.apply(dot, w, 4, torch.float32) + torch.testing.assert_close(out_helper, out_class) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_rope_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_rope_triton.py new file mode 100644 index 000000000..4a8461af5 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_rope_triton.py @@ -0,0 +1,434 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-6 P35 G38 — `apply_interleaved_partial_rope` Triton FWD/BWD parity. + +Asserts that :class:`RoPEInterleavedPartialFn` (Triton kernel from +``primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial``) +produces the same output as the eager +:func:`apply_interleaved_partial_rope` body in +``primus.backends.megatron.core.transformer.dual_rope``, FWD **and** +BWD, at two tiers: + +* fast tier — small shapes (``B=2, S=8, H=4, head_dim=16, rd=8``) + exercising every code path (nope copy, interleaved pair rotation, + cos/sin broadcast across heads) in milliseconds; parametrised over + ``{fp32, fp16, bf16}`` and ``rotary_dim ∈ {0, 4, 8, 16}``; +* release tier — Q shape (``B=1, S=4096, H=64, head_dim=512, rd=64``) + and K shape (``B=1, S=4096, H=1, head_dim=64, rd=64``) mirroring the + V4-Flash EP=8 proxy widths, behind ``pytest.mark.slow``. + +`gradcheck` is run at the fast tier in fp32 to catch any analytic-VJP +bug in the BWD kernel; the FWD has a clean closed form so its forward +parity assertion is the load-bearing test. + +GPU-only; CPU runs are ``pytest.skip``-ed at module collection time. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip( + "rope_interleaved_partial Triton kernel requires CUDA / HIP", + allow_module_level=True, + ) + +pytest.importorskip("triton", reason="Triton not installed") + +from primus.backends.megatron.core.transformer.dual_rope import ( # noqa: E402 + apply_interleaved_partial_rope, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.rope_interleaved_partial import ( # noqa: E402 + RoPEInterleavedPartialFn, + apply_rope_interleaved_partial, + eager_apply_interleaved_partial_rope, + is_triton_path_enabled, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +@contextmanager +def _env(key: str, value: str | None): + """Temporarily set / unset ``os.environ[key]``.""" + prev = os.environ.get(key) + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + try: + yield + finally: + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + + +def _build_cos_sin( + *, + leading_shape: tuple[int, ...], + rotary_dim: int, + dtype: torch.dtype, + seed: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build cos / sin with shape ``leading_shape + (rotary_dim // 2,)``. + + Mimics the way :class:`RoPECache` produces cos/sin: build + ``position_ids.float() * inv_freq`` then take ``cos`` / ``sin``. + For the test we just generate random freqs to exercise more values. + """ + rd_half = rotary_dim // 2 + gen = torch.Generator(device="cuda").manual_seed(seed) + freqs = torch.randn((*leading_shape, rd_half), dtype=torch.float32, device="cuda", generator=gen) + cos = freqs.cos().to(dtype) + sin = freqs.sin().to(dtype) + return cos, sin + + +def _build_x( + *, + leading_shape: tuple[int, ...], + H: int, + head_dim: int, + dtype: torch.dtype, + requires_grad: bool = False, + seed: int = 0, +) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn( + (*leading_shape, H, head_dim), + dtype=dtype, + device="cuda", + generator=gen, + ) + if requires_grad: + x.requires_grad_(True) + return x + + +def _dtype_tolerance(dtype: torch.dtype) -> tuple[float, float]: + if dtype == torch.float32: + return 1e-6, 1e-6 + if dtype == torch.float16: + return 1e-3, 1e-3 + if dtype == torch.bfloat16: + return 1e-2, 1e-2 + raise ValueError(dtype) + + +# --------------------------------------------------------------------------- +# G38: FWD parity vs eager +# --------------------------------------------------------------------------- + + +_FAST_LEADING = (2, 8) # B=2, S=8 +_FAST_H = 4 +_FAST_HEAD_DIM = 16 + + +class TestG38ForwardParity: + """FWD output matches eager :func:`apply_interleaved_partial_rope` + body within dtype tolerance. + + The Triton kernel does its arithmetic in the caller's dtype (bf16 / + fp16 / fp32), matching the plan-5 P32 RoPE bf16 cast contract. So + the tolerance is the same as if we cast cos/sin to ``x.dtype`` and + ran the eager body — which is exactly what + :func:`eager_apply_interleaved_partial_rope` does. + """ + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("rotary_dim", [0, 4, 8, 16]) + def test_fast_tier_fwd_parity(self, dtype, rotary_dim): + if rotary_dim > _FAST_HEAD_DIM: + pytest.skip("rotary_dim must be <= head_dim") + x = _build_x( + leading_shape=_FAST_LEADING, + H=_FAST_H, + head_dim=_FAST_HEAD_DIM, + dtype=dtype, + seed=42, + ) + cos, sin = _build_cos_sin( + leading_shape=_FAST_LEADING, + rotary_dim=max(rotary_dim, 2), + dtype=dtype, + seed=43, + ) + # For rd=0 the kernel ignores cos/sin; pass any shape. + cos_use = cos[..., : max(rotary_dim // 2, 1)] + sin_use = sin[..., : max(rotary_dim // 2, 1)] + + out_triton = RoPEInterleavedPartialFn.apply(x, cos_use, sin_use, rotary_dim) + out_eager = eager_apply_interleaved_partial_rope(x, cos_use, sin_use, rotary_dim=rotary_dim) + + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_uneven_block_h(self, dtype): + """H=5 is not a multiple of BLOCK_H=8 -> masked load / store paths.""" + leading = (1, 7) # also a prime sequence length + H = 5 + head_dim = 12 + rotary_dim = 4 + x = _build_x(leading_shape=leading, H=H, head_dim=head_dim, dtype=dtype, seed=1234) + cos, sin = _build_cos_sin(leading_shape=leading, rotary_dim=rotary_dim, dtype=dtype, seed=1235) + out_triton = RoPEInterleavedPartialFn.apply(x, cos, sin, rotary_dim) + out_eager = eager_apply_interleaved_partial_rope(x, cos, sin, rotary_dim=rotary_dim) + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# G38: BWD parity vs eager + gradcheck +# --------------------------------------------------------------------------- + + +class TestG38BackwardParity: + """BWD parity vs eager autograd graph, plus a small fp64 gradcheck.""" + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + @pytest.mark.parametrize("rotary_dim", [4, 8, 16]) + def test_fast_tier_bwd_parity(self, dtype, rotary_dim): + if rotary_dim > _FAST_HEAD_DIM: + pytest.skip("rotary_dim must be <= head_dim") + + x_t = _build_x( + leading_shape=_FAST_LEADING, + H=_FAST_H, + head_dim=_FAST_HEAD_DIM, + dtype=dtype, + requires_grad=True, + seed=44, + ) + x_e = x_t.detach().clone().requires_grad_(True) + cos, sin = _build_cos_sin( + leading_shape=_FAST_LEADING, + rotary_dim=rotary_dim, + dtype=dtype, + seed=45, + ) + + out_triton = RoPEInterleavedPartialFn.apply(x_t, cos, sin, rotary_dim) + out_eager = eager_apply_interleaved_partial_rope(x_e, cos, sin, rotary_dim=rotary_dim) + + gen = torch.Generator(device="cuda").manual_seed(46) + grad = torch.randn_like(out_triton, dtype=dtype) + grad_e = grad.detach().clone() + gen.manual_seed(46) + out_triton.backward(grad) + out_eager.backward(grad_e) + + atol, rtol = _dtype_tolerance(dtype) + assert x_t.grad is not None + assert x_e.grad is not None + torch.testing.assert_close(x_t.grad, x_e.grad, atol=atol, rtol=rtol) + + def test_gradcheck_fast_tier(self): + """`torch.autograd.gradcheck` at fp64 + tiny shape catches any + analytic-VJP bug in :func:`_apply_rope_bwd_kernel`. + """ + leading = (1, 4) + H = 2 + head_dim = 8 + rotary_dim = 4 + x = _build_x( + leading_shape=leading, + H=H, + head_dim=head_dim, + dtype=torch.float64, + requires_grad=True, + seed=100, + ) + cos, sin = _build_cos_sin( + leading_shape=leading, + rotary_dim=rotary_dim, + dtype=torch.float64, + seed=101, + ) + + def fn(xx): + return RoPEInterleavedPartialFn.apply(xx, cos, sin, rotary_dim) + + torch.autograd.gradcheck(fn, (x,), eps=1e-6, atol=1e-7, rtol=1e-5) + + +# --------------------------------------------------------------------------- +# G38: release-tier V4-Flash widths (Q + K) +# --------------------------------------------------------------------------- + + +class TestG38ReleaseTier: + """V4-Flash EP=8 widths: Q (H=64, head_dim=512, rd=64) and K (H=1). + + Marked ``slow``; bf16 only (the production dtype). FWD + BWD parity + within the elementwise attention tolerance from plan-5 P32. + """ + + @pytest.mark.slow + @pytest.mark.parametrize( + "spec", + [ + ("Q", (1, 4096), 64, 512, 64), + ("K", (1, 4096), 1, 64, 64), + ], + ids=["Q-shape", "K-shape"], + ) + def test_v4_flash_proxy_parity(self, spec): + name, leading, H, head_dim, rotary_dim = spec + dtype = torch.bfloat16 + + x_t = _build_x( + leading_shape=leading, + H=H, + head_dim=head_dim, + dtype=dtype, + requires_grad=True, + seed=200, + ) + x_e = x_t.detach().clone().requires_grad_(True) + cos, sin = _build_cos_sin(leading_shape=leading, rotary_dim=rotary_dim, dtype=dtype, seed=201) + + out_triton = RoPEInterleavedPartialFn.apply(x_t, cos, sin, rotary_dim) + out_eager = eager_apply_interleaved_partial_rope(x_e, cos, sin, rotary_dim=rotary_dim) + + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol) + + grad = torch.randn_like(out_triton) + out_triton.backward(grad) + out_eager.backward(grad.detach().clone()) + torch.testing.assert_close(x_t.grad, x_e.grad, atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# G38: rd == 0 early-return + error paths +# --------------------------------------------------------------------------- + + +class TestG38EdgeCases: + """Cover :class:`RoPEInterleavedPartialFn`'s defensive validation.""" + + def test_rd_zero_returns_input_contiguous(self): + x = _build_x( + leading_shape=(2, 4), + H=2, + head_dim=8, + dtype=torch.float32, + seed=1, + ) + cos, sin = _build_cos_sin(leading_shape=(2, 4), rotary_dim=2, dtype=torch.float32, seed=2) + out = RoPEInterleavedPartialFn.apply(x, cos, sin, 0) + assert torch.equal(out, x.contiguous()) + + def test_odd_rotary_dim_raises(self): + x = _build_x( + leading_shape=(1, 4), + H=2, + head_dim=8, + dtype=torch.float32, + seed=1, + ) + cos, sin = _build_cos_sin(leading_shape=(1, 4), rotary_dim=4, dtype=torch.float32, seed=2) + with pytest.raises(ValueError, match="rotary_dim must be even"): + RoPEInterleavedPartialFn.apply(x, cos, sin, 3) + + def test_rotary_dim_exceeds_head_dim_raises(self): + x = _build_x( + leading_shape=(1, 4), + H=2, + head_dim=8, + dtype=torch.float32, + seed=1, + ) + cos, sin = _build_cos_sin(leading_shape=(1, 4), rotary_dim=4, dtype=torch.float32, seed=2) + with pytest.raises(ValueError, match="must be <= head_dim"): + RoPEInterleavedPartialFn.apply(x, cos, sin, 16) + + def test_cos_shape_mismatch_raises(self): + x = _build_x( + leading_shape=(2, 4), + H=2, + head_dim=8, + dtype=torch.float32, + seed=1, + ) + cos, sin = _build_cos_sin(leading_shape=(2, 4), rotary_dim=2, dtype=torch.float32, seed=2) + with pytest.raises(ValueError, match="last dim must be rotary_dim"): + RoPEInterleavedPartialFn.apply(x, cos, sin, 4) + + +# --------------------------------------------------------------------------- +# G38: env-flag dispatch through dual_rope.apply_interleaved_partial_rope +# --------------------------------------------------------------------------- + + +class TestG38EnvFlagDispatch: + """`PRIMUS_ROPE_TRITON` env knob flips the dispatcher behaviour + inside :func:`apply_interleaved_partial_rope`. + + The Triton path is bit-equivalent to the eager body within bf16 + tolerance; this test pins both code paths to make sure neither is + silently broken. + """ + + def test_env_on_uses_triton(self): + x = _build_x( + leading_shape=_FAST_LEADING, + H=_FAST_H, + head_dim=_FAST_HEAD_DIM, + dtype=torch.bfloat16, + seed=11, + ) + cos, sin = _build_cos_sin( + leading_shape=_FAST_LEADING, + rotary_dim=8, + dtype=torch.bfloat16, + seed=12, + ) + with _env("PRIMUS_ROPE_TRITON", "1"): + assert is_triton_path_enabled() + out_on = apply_interleaved_partial_rope(x, cos, sin, rotary_dim=8) + with _env("PRIMUS_ROPE_TRITON", "0"): + assert not is_triton_path_enabled() + out_off = apply_interleaved_partial_rope(x, cos, sin, rotary_dim=8) + atol, rtol = _dtype_tolerance(torch.bfloat16) + torch.testing.assert_close(out_on, out_off, atol=atol, rtol=rtol) + + def test_apply_rope_interleaved_partial_dispatcher(self): + """The dispatcher in the kernel module mirrors the dual_rope + wiring -- pin both paths return the same answer. + """ + x = _build_x( + leading_shape=(2, 16), + H=2, + head_dim=16, + dtype=torch.bfloat16, + seed=21, + ) + cos, sin = _build_cos_sin( + leading_shape=(2, 16), + rotary_dim=8, + dtype=torch.bfloat16, + seed=22, + ) + with _env("PRIMUS_ROPE_TRITON", "1"): + out_triton = apply_rope_interleaved_partial(x, cos, sin, rotary_dim=8) + with _env("PRIMUS_ROPE_TRITON", "0"): + out_eager = apply_rope_interleaved_partial(x, cos, sin, rotary_dim=8) + atol, rtol = _dtype_tolerance(torch.bfloat16) + torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_router_post_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_router_post_triton.py new file mode 100644 index 000000000..d6f93a2f8 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_router_post_triton.py @@ -0,0 +1,267 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-6 P39 G42 — V4 router post-logits Triton parity. + +Asserts that :class:`V4RouterPostFn` (Triton kernel from +``primus.backends.megatron.core.transformer.moe._triton.v4_router_post``) +produces ``(probs, routing_map)`` bit-equal to the eager body of +:func:`primus.backends.megatron.core.transformer.moe.v4_topk_router._compute_route` +across the 3 score functions × {with, without bias} × {hash router, +learned router}. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip( + "v4_router_post Triton kernel requires CUDA / HIP", + allow_module_level=True, + ) + +pytest.importorskip("triton", reason="Triton not installed") + +from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import ( # noqa: E402 + V4RouterPostFn, + is_triton_path_enabled, +) +from primus.backends.megatron.core.transformer.moe.v4_hash_router import ( # noqa: E402 + DeepseekV4HashRouter, +) +from primus.backends.megatron.core.transformer.moe.v4_topk_router import ( # noqa: E402 + DeepseekV4LearnedRouter, +) + + +@contextmanager +def _env(key: str, value: str | None): + prev = os.environ.get(key) + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + try: + yield + finally: + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + + +def _eager_v4_score(logits: torch.Tensor, *, score_function: str): + if score_function == "softmax": + return torch.softmax(logits, dim=-1) + if score_function == "sigmoid": + return torch.sigmoid(logits) + if score_function == "sqrtsoftplus": + return torch.sqrt(torch.nn.functional.softplus(logits)) + raise ValueError(score_function) + + +def _eager_post( + logits: torch.Tensor, + indices: torch.Tensor, + *, + score_function: str, + topk_scaling_factor: float, + out_dtype: torch.dtype, +): + scores = _eager_v4_score(logits.float(), score_function=score_function) + weights = scores.gather(1, indices) + if score_function != "softmax": + denom = weights.sum(dim=-1, keepdim=True).clamp(min=1.0e-12) + weights = weights / denom + if topk_scaling_factor != 1.0: + weights = weights * float(topk_scaling_factor) + N, E = logits.shape + probs = torch.zeros(N, E, dtype=out_dtype, device=logits.device) + probs.scatter_(1, indices, weights.to(out_dtype)) + rmap = torch.zeros(N, E, dtype=torch.bool, device=logits.device) + rmap.scatter_(1, indices, True) + return probs, rmap + + +def _build_inputs(*, N: int, E: int, K: int, seed: int): + gen = torch.Generator(device="cuda").manual_seed(seed) + logits = torch.randn((N, E), dtype=torch.float32, device="cuda", generator=gen) + indices = torch.stack( + [torch.randperm(E, generator=gen, device="cuda")[:K] for _ in range(N)], + dim=0, + ).to(torch.int64) + return logits, indices + + +# --------------------------------------------------------------------------- +# G42: FWD parity vs eager (3 score fns × {fp32, bf16} × small / release) +# --------------------------------------------------------------------------- + + +class TestG42ForwardParity: + @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) + @pytest.mark.parametrize("scale", [1.0, 2.5]) + @pytest.mark.parametrize("out_dtype", [torch.float32]) + def test_fast_tier_fwd_eager_parity(self, score_function, scale, out_dtype): + N, E, K = 64, 32, 4 + logits, indices = _build_inputs(N=N, E=E, K=K, seed=2000) + probs_t, rmap_t = V4RouterPostFn.apply(logits, indices, score_function, scale, out_dtype) + probs_e, rmap_e = _eager_post( + logits, indices, score_function=score_function, topk_scaling_factor=scale, out_dtype=out_dtype + ) + torch.testing.assert_close(probs_t, probs_e, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(rmap_t, rmap_e) + torch.testing.assert_close(probs_t.nonzero(), probs_e.nonzero(), check_dtype=False) + + @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) + @pytest.mark.parametrize("shape", [(64, 256, 6), (64, 7, 3), (64, 100, 5)]) + def test_non_pow2_e_k_fwd_parity(self, score_function, shape): + """Arbitrary (non-power-of-2) E / K — the production V4-Flash router + runs E=256, K=6, so K=6 is the load-bearing case. Also cover a + non-power-of-2 E (7, 100).""" + N, E, K = shape + logits, indices = _build_inputs(N=N, E=E, K=K, seed=2500) + probs_t, rmap_t = V4RouterPostFn.apply(logits, indices, score_function, 2.5, torch.float32) + probs_e, rmap_e = _eager_post( + logits, indices, score_function=score_function, topk_scaling_factor=2.5, out_dtype=torch.float32 + ) + torch.testing.assert_close(probs_t, probs_e, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(rmap_t, rmap_e) + + @pytest.mark.slow + @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) + def test_release_tier_fwd_eager_parity(self, score_function): + # V4-Flash production widths: E=256, K=6 (non-power-of-2 topk). + N, E, K = 4096, 256, 6 + logits, indices = _build_inputs(N=N, E=E, K=K, seed=8000) + probs_t, rmap_t = V4RouterPostFn.apply(logits, indices, score_function, 2.5, torch.float32) + probs_e, rmap_e = _eager_post( + logits, indices, score_function=score_function, topk_scaling_factor=2.5, out_dtype=torch.float32 + ) + torch.testing.assert_close(probs_t, probs_e, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(rmap_t, rmap_e) + + +# --------------------------------------------------------------------------- +# G42: BWD parity +# --------------------------------------------------------------------------- + + +class TestG42BackwardParity: + @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) + @pytest.mark.parametrize("scale", [1.0, 2.5]) + @pytest.mark.parametrize("shape", [(64, 32, 4), (64, 256, 6), (64, 100, 5)]) + def test_fast_tier_bwd_eager_parity(self, score_function, scale, shape): + N, E, K = shape + logits_e, indices = _build_inputs(N=N, E=E, K=K, seed=3100) + logits_e = logits_e.requires_grad_(True) + logits_t = logits_e.detach().clone().requires_grad_(True) + + probs_t, _ = V4RouterPostFn.apply(logits_t, indices, score_function, scale, torch.float32) + probs_e, _ = _eager_post( + logits_e, + indices, + score_function=score_function, + topk_scaling_factor=scale, + out_dtype=torch.float32, + ) + + g = torch.randn_like(probs_t) + (probs_t * g).sum().backward() + (probs_e * g).sum().backward() + + torch.testing.assert_close(logits_t.grad, logits_e.grad, atol=1e-4, rtol=1e-4) + + +# --------------------------------------------------------------------------- +# G42: composed routers +# --------------------------------------------------------------------------- + + +class TestG42ComposedRouters: + """Both routers (`DeepseekV4LearnedRouter` and `DeepseekV4HashRouter`) + must produce bit-equal (probs, routing_map) between the + PRIMUS_V4_ROUTER_TRITON=1 and =0 paths. This is the load-bearing + integration test for downstream MoEDispatch. + """ + + @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) + def test_learned_router_env_toggle_parity(self, score_function): + torch.manual_seed(123) + N, D, E, K = 64, 32, 32, 4 + router = DeepseekV4LearnedRouter( + hidden_size=D, num_experts=E, topk=K, score_function=score_function + ).to("cuda") + hidden = torch.randn((1, N, D), dtype=torch.float32, device="cuda") + + with _env("PRIMUS_V4_ROUTER_TRITON", "1"): + assert is_triton_path_enabled() + probs_t, rmap_t = router(hidden) + with _env("PRIMUS_V4_ROUTER_TRITON", "0"): + assert not is_triton_path_enabled() + probs_e, rmap_e = router(hidden) + + torch.testing.assert_close(probs_t, probs_e, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(rmap_t, rmap_e) + + @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) + def test_hash_router_env_toggle_parity(self, score_function): + torch.manual_seed(124) + N, D, E, K = 64, 32, 32, 4 + V = 1000 + router = DeepseekV4HashRouter( + hidden_size=D, + num_experts=E, + topk=K, + vocab_size=V, + score_function=score_function, + ).to("cuda") + hidden = torch.randn((1, N, D), dtype=torch.float32, device="cuda") + token_ids = torch.randint(0, V, (1, N), dtype=torch.long, device="cuda") + + with _env("PRIMUS_V4_ROUTER_TRITON", "1"): + probs_t, rmap_t = router(hidden, token_ids) + with _env("PRIMUS_V4_ROUTER_TRITON", "0"): + probs_e, rmap_e = router(hidden, token_ids) + + torch.testing.assert_close(probs_t, probs_e, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(rmap_t, rmap_e) + + +# --------------------------------------------------------------------------- +# G42: edge cases +# --------------------------------------------------------------------------- + + +class TestG42EdgeCases: + def test_unknown_score_function_raises(self): + logits = torch.randn((4, 8), dtype=torch.float32, device="cuda") + indices = torch.zeros((4, 2), dtype=torch.int64, device="cuda") + with pytest.raises(ValueError, match="score_function"): + V4RouterPostFn.apply(logits, indices, "tanh", 1.0, torch.float32) + + def test_cpu_tensor_asserts(self): + # v4_router_post_triton asserts CUDA tensors (no support predicate). + from primus.backends.megatron.core.transformer.moe._triton.v4_router_post import ( + v4_router_post_triton, + ) + + cpu_logits = torch.randn((4, 8), dtype=torch.float32) + cpu_idx = torch.zeros((4, 2), dtype=torch.int64) + with pytest.raises(AssertionError, match="CUDA"): + v4_router_post_triton( + cpu_logits, + cpu_idx, + score_function="softmax", + topk_scaling_factor=1.0, + out_dtype=torch.float32, + ) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_sinkhorn_triton.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_sinkhorn_triton.py new file mode 100644 index 000000000..934daf1ab --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_sinkhorn_triton.py @@ -0,0 +1,358 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-6 P36 G39 — `sinkhorn_normalize` Triton FWD/BWD parity. + +Asserts that :class:`SinkhornNormalizeFn` (Triton kernel from +``primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.sinkhorn``) +matches the eager :func:`sinkhorn_normalize` body in +``primus.backends.megatron.core.transformer.hyper_connection`` and the +plan-5 P29 compiled path within dtype tolerance, FWD **and** BWD, at +two tiers: + +* fast tier — `B=2, S=64, K=4`, exercising every code path in the + kernel (priming col-step + 19 row/col pairs + cached state buffer + round-trip) in milliseconds; parametrised over ``bf16`` (the + production compute dtype) and ``n_iters ∈ {5, 20}``; +* release tier — `B=1, S=4096, K=4` bf16 (V4-Flash production shape), + behind ``pytest.mark.slow``. + +The **doubly-stochastic property check** is a model-quality contract +independent of the eager path (row / col sums of the FWD output equal +``1`` within ``eps * K``). + +Note on ``torch.autograd.gradcheck``: the eager body casts to fp32 +internally (``m = logits.float()``) and the Triton kernel matches that +contract bit-for-bit. ``gradcheck`` at fp64 input is therefore +incompatible with both paths — the fp32 cast is lossy at the fp64 +finite-difference step size. We instead pin the BWD via a direct +parity check against ``torch.autograd`` of the eager body in fp64 +(where ``out.backward(grad)`` does exactly the right thing through the +``.float()`` cast). + +GPU-only; CPU runs are ``pytest.skip``-ed at module collection time. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip( + "sinkhorn Triton kernel requires CUDA / HIP", + allow_module_level=True, + ) + +pytest.importorskip("triton", reason="Triton not installed") + +from primus.backends.megatron.core.transformer.hyper_connection import ( # noqa: E402 + _get_compiled_sinkhorn, + sinkhorn_normalize, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._triton_common.sinkhorn import ( # noqa: E402 + SinkhornNormalizeFn, + eager_sinkhorn_normalize, + is_triton_kernel_supported, + is_triton_path_enabled, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +@contextmanager +def _env(key: str, value: str | None): + """Temporarily set / unset ``os.environ[key]``.""" + prev = os.environ.get(key) + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + try: + yield + finally: + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + + +def _build_logits( + *, B: int, S: int, K: int, dtype: torch.dtype, seed: int, requires_grad: bool = False +) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + # Sinkhorn input is non-negative (softmax + eps in production). + x = torch.rand((B, S, K, K), dtype=dtype, device="cuda", generator=gen) + 1e-3 + if requires_grad: + x.requires_grad_(True) + return x + + +def _dtype_tolerance(dtype: torch.dtype) -> tuple[float, float]: + if dtype == torch.float32: + return 1e-5, 1e-5 + if dtype == torch.float16: + return 1e-3, 1e-3 + if dtype == torch.bfloat16: + return 1e-2, 1e-2 + if dtype == torch.float64: + return 1e-7, 1e-7 + raise ValueError(dtype) + + +# --------------------------------------------------------------------------- +# G39: FWD parity vs eager (and vs plan-5 P29 compiled) +# --------------------------------------------------------------------------- + + +class TestG39ForwardParity: + """FWD output matches eager and (separately) the compiled path.""" + + @pytest.mark.parametrize("dtype", [torch.bfloat16]) + @pytest.mark.parametrize("n_iters", [5, 20]) + def test_fast_tier_fwd_eager_parity(self, dtype, n_iters): + x = _build_logits(B=2, S=64, K=4, dtype=dtype, seed=42) + + out_triton = SinkhornNormalizeFn.apply(x.clone(), n_iters, 1e-6) + out_eager = eager_sinkhorn_normalize(x.clone(), n_iters=n_iters, eps=1e-6) + + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol) + + @pytest.mark.parametrize("dtype", [torch.bfloat16]) + @pytest.mark.parametrize("n_iters", [5, 20]) + def test_fast_tier_fwd_compiled_parity(self, dtype, n_iters): + """Triton FWD == plan-5 P29 compiled FWD within dtype tolerance. + + Both paths share the same algorithm; this test pins them + together so a future P29 / P36 divergence raises a CI alarm. + """ + x = _build_logits(B=2, S=64, K=4, dtype=dtype, seed=43) + + out_triton = SinkhornNormalizeFn.apply(x.clone(), n_iters, 1e-6) + compiled_fn = _get_compiled_sinkhorn(n_iters, 1e-6, dtype) + out_compiled = compiled_fn(x.clone()) + + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(out_triton, out_compiled, atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# G39: BWD parity vs eager + gradcheck +# --------------------------------------------------------------------------- + + +class TestG39BackwardParity: + """BWD parity vs eager ``torch.autograd`` across dtypes.""" + + @pytest.mark.parametrize("dtype", [torch.bfloat16]) + @pytest.mark.parametrize("n_iters", [5, 20]) + def test_fast_tier_bwd_eager_parity(self, dtype, n_iters): + x_base = _build_logits(B=2, S=64, K=4, dtype=dtype, seed=44) + x_t = x_base.detach().clone().requires_grad_(True) + x_e = x_base.detach().clone().requires_grad_(True) + + out_triton = SinkhornNormalizeFn.apply(x_t, n_iters, 1e-6) + out_eager = eager_sinkhorn_normalize(x_e, n_iters=n_iters, eps=1e-6) + + grad = torch.randn_like(out_triton) + out_triton.backward(grad) + out_eager.backward(grad.detach().clone()) + + atol, rtol = _dtype_tolerance(dtype) + assert x_t.grad is not None + assert x_e.grad is not None + torch.testing.assert_close(x_t.grad, x_e.grad, atol=atol, rtol=rtol) + + def test_fp64_input_bwd_eager_parity(self): + """fp64 input -> the kernel still does fp32-internal compute + (matches eager's ``m = logits.float()`` contract); we compare + Triton BWD vs ``torch.autograd``-of-eager rather than gradcheck + (gradcheck's fp64 finite-difference step is incompatible with + the lossy fp64->fp32 cast). + """ + x_base = _build_logits(B=1, S=8, K=4, dtype=torch.float64, seed=100) + x_t = x_base.detach().clone().requires_grad_(True) + x_e = x_base.detach().clone().requires_grad_(True) + + out_triton = SinkhornNormalizeFn.apply(x_t, 5, 1e-6) + out_eager = eager_sinkhorn_normalize(x_e, n_iters=5, eps=1e-6) + + grad = torch.randn_like(out_triton) + out_triton.backward(grad) + out_eager.backward(grad.detach().clone()) + + # fp64 grads come back through a fp32-internal compute; tolerance + # is set to fp32 precision (the kernel can't be more accurate + # than its compute dtype). + torch.testing.assert_close(x_t.grad, x_e.grad, atol=1e-5, rtol=1e-5) + + +# --------------------------------------------------------------------------- +# G39: doubly-stochastic property (model-quality contract) +# --------------------------------------------------------------------------- + + +class TestG39DoublyStochastic: + """The FWD output's row and column sums equal ``1`` within + ``eps * K`` -- pinning the model-quality contract that + :func:`sinkhorn_normalize` is supposed to guarantee. + + This is independent of the eager path; an algorithmic bug in the + Triton kernel that still happens to match the eager value would + fail this check. + """ + + @pytest.mark.parametrize("dtype", [torch.bfloat16]) + @pytest.mark.parametrize("K", [4, 8]) + def test_row_col_sums_close_to_1(self, dtype, K): + x = _build_logits(B=2, S=16, K=K, dtype=dtype, seed=55) + y = SinkhornNormalizeFn.apply(x, 20, 1e-6) + + row_sums = y.sum(dim=-1) + col_sums = y.sum(dim=-2) + # eps*K = 4e-6 (fp32) or ~6e-5 (bf16 -> the cast back to bf16 + # is the dominant error, so we widen by 16x for bf16). + tol = 4e-6 * K if dtype == torch.float32 else 1e-2 + torch.testing.assert_close(row_sums, torch.ones_like(row_sums), atol=tol, rtol=tol) + torch.testing.assert_close(col_sums, torch.ones_like(col_sums), atol=tol, rtol=tol) + + +# --------------------------------------------------------------------------- +# G39: release-tier V4-Flash production shape +# --------------------------------------------------------------------------- + + +class TestG39ReleaseTier: + """V4-Flash production shape: `B=1, S=4096, K=4`, bf16, n_iters=20. + + Marked ``slow``; pins both FWD and BWD against eager AND compiled + paths to make sure the in-production widths match. + """ + + @pytest.mark.slow + def test_v4_flash_fwd_bwd_parity(self): + dtype = torch.bfloat16 + x_base = _build_logits(B=1, S=4096, K=4, dtype=dtype, seed=200) + + x_t = x_base.detach().clone().requires_grad_(True) + x_e = x_base.detach().clone().requires_grad_(True) + x_c = x_base.detach().clone().requires_grad_(True) + + out_triton = SinkhornNormalizeFn.apply(x_t, 20, 1e-6) + out_eager = eager_sinkhorn_normalize(x_e, n_iters=20, eps=1e-6) + compiled_fn = _get_compiled_sinkhorn(20, 1e-6, dtype) + out_compiled = compiled_fn(x_c) + + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol) + torch.testing.assert_close(out_triton, out_compiled, atol=atol, rtol=rtol) + + grad = torch.randn_like(out_triton) + out_triton.backward(grad.detach().clone()) + out_eager.backward(grad.detach().clone()) + out_compiled.backward(grad.detach().clone()) + torch.testing.assert_close(x_t.grad, x_e.grad, atol=atol, rtol=rtol) + torch.testing.assert_close(x_t.grad, x_c.grad, atol=atol, rtol=rtol) + + +# --------------------------------------------------------------------------- +# G39: edge cases + error paths +# --------------------------------------------------------------------------- + + +class TestG39EdgeCases: + """Defensive validation in :class:`SinkhornNormalizeFn`.""" + + def test_non_square_raises(self): + x = torch.rand(2, 4, 5, device="cuda", dtype=torch.float32) + 1e-3 + with pytest.raises(ValueError, match="square"): + SinkhornNormalizeFn.apply(x, 20, 1e-6) + + def test_unsupported_k_raises(self): + x = torch.rand(2, 32, 32, device="cuda", dtype=torch.float32) + 1e-3 + with pytest.raises(ValueError, match="unsupported K"): + SinkhornNormalizeFn.apply(x, 20, 1e-6) + + def test_n_iters_zero_raises(self): + x = torch.rand(2, 4, 4, device="cuda", dtype=torch.float32) + 1e-3 + with pytest.raises(ValueError, match="n_iters must be"): + SinkhornNormalizeFn.apply(x, 0, 1e-6) + + def test_kernel_supported_predicate(self): + good = torch.rand(2, 4, 4, device="cuda", dtype=torch.float32) + bad_k = torch.rand(2, 32, 32, device="cuda", dtype=torch.float32) + bad_dev = torch.rand(2, 4, 4, dtype=torch.float32) # cpu + assert is_triton_kernel_supported(good) + assert not is_triton_kernel_supported(bad_k) + assert not is_triton_kernel_supported(bad_dev) + + +# --------------------------------------------------------------------------- +# G39: env-flag dispatch through hyper_connection.sinkhorn_normalize +# --------------------------------------------------------------------------- + + +class TestG39EnvFlagDispatch: + """The ``PRIMUS_SINKHORN_TRITON`` env knob flips the dispatcher + inside :func:`sinkhorn_normalize`. Both paths agree within bf16 + tolerance. + """ + + def test_env_on_uses_triton(self): + x = _build_logits(B=2, S=8, K=4, dtype=torch.bfloat16, seed=11) + + with _env("PRIMUS_SINKHORN_TRITON", "1"): + assert is_triton_path_enabled() + out_on = sinkhorn_normalize(x.clone(), n_iters=20, eps=1e-6) + with _env("PRIMUS_SINKHORN_TRITON", "0"): + assert not is_triton_path_enabled() + out_eager = sinkhorn_normalize(x.clone(), n_iters=20, eps=1e-6) + + atol, rtol = _dtype_tolerance(torch.bfloat16) + torch.testing.assert_close(out_on, out_eager, atol=atol, rtol=rtol) + + def test_use_triton_kwarg_overrides_env_off(self): + """``use_triton=True`` forces the Triton path even when the env + knob is off (used by the unit tests / by callers who want a + per-call override). + """ + x = _build_logits(B=2, S=8, K=4, dtype=torch.bfloat16, seed=12) + with _env("PRIMUS_SINKHORN_TRITON", "0"): + out_triton = sinkhorn_normalize(x.clone(), n_iters=20, eps=1e-6, use_triton=True) + out_eager = sinkhorn_normalize(x.clone(), n_iters=20, eps=1e-6) + + atol, rtol = _dtype_tolerance(torch.bfloat16) + torch.testing.assert_close(out_triton, out_eager, atol=atol, rtol=rtol) + + def test_routing_precedence_triton_over_compiled(self): + """When both ``use_triton`` and ``use_compiled`` are set, the + Triton path wins (routing precedence: + ``use_triton > use_compiled > eager``). + """ + x = _build_logits(B=2, S=8, K=4, dtype=torch.bfloat16, seed=13) + with _env("PRIMUS_SINKHORN_TRITON", "0"): + out = sinkhorn_normalize( + x.clone(), + n_iters=20, + eps=1e-6, + use_triton=True, + use_compiled=True, + ) + out_triton_only = sinkhorn_normalize( + x.clone(), + n_iters=20, + eps=1e-6, + use_triton=True, + ) + + torch.testing.assert_close(out, out_triton_only, atol=0, rtol=0) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_backend_import_gating.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_backend_import_gating.py new file mode 100644 index 000000000..ce70fcdfb --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_backend_import_gating.py @@ -0,0 +1,122 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Backend import / default / hardware-gating guarantees for V4 attention. + +These are the regression gates for the "gluon must not be a hard dependency" +contract: + +* the shared config default is ``triton_v1`` (arch-portable), NOT the + gfx950-only ``gluon``; +* importing ``v4_attention_kernels`` (and ``deepseek_v4_attention``) must NOT + eagerly import the gluon backend (``_gluon_dsa`` / + ``triton.experimental.gluon``), so ``eager`` / ``triton_v1`` / ``triton_v2`` + work on any Triton build / GPU arch; +* the gluon backend is loaded lazily via ``load_gluon_attention_backends`` and, + when selected, ``_require_gfx950`` rejects non-gfx950 devices. + +All tests here are hardware-independent (no CUDA required). +""" + +from __future__ import annotations + +import importlib +import sys + +import pytest + + +def test_config_defaults_are_triton_v1(): + """Shared config default must be the arch-portable triton_v1 (not gluon).""" + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, + ) + + fields = DeepSeekV4TransformerConfig.__dataclass_fields__ + assert fields["use_v4_attention_backend"].default == "triton_v1" + assert fields["use_v4_csa_attention_backend"].default == "triton_v1" + + +def test_kernels_package_exposes_lazy_gluon_loader(): + """The package exposes the lazy loader and does NOT eagerly bind gluon entries.""" + pkg = importlib.import_module("primus.backends.megatron.core.transformer.v4_attention_kernels") + # A prior gluon test may have called the lazy loader, which caches the + # backends as module attributes; drop them + reload so this asserts the + # import-time (no eager binding) contract independent of test order. + for attr in ("v4_attention_gluon", "v4_csa_attention_gluon"): + if hasattr(pkg, attr): + delattr(pkg, attr) + importlib.reload(pkg) + assert hasattr(pkg, "load_gluon_attention_backends") + # gluon entries must NOT be module-level attributes (would mean eager import). + assert not hasattr(pkg, "v4_attention_gluon") + assert not hasattr(pkg, "v4_csa_attention_gluon") + + +def test_importing_kernels_package_does_not_pull_gluon(): + """Reloading the kernels package must not import ``_gluon_dsa`` as a side effect.""" + # Purge any previously-imported gluon modules so this asserts the *package* + # import path, independent of test ordering (e.g. the gfx950 gluon UT). + for name in [m for m in sys.modules if "_gluon_dsa" in m]: + del sys.modules[name] + pkg = importlib.import_module("primus.backends.megatron.core.transformer.v4_attention_kernels") + importlib.reload(pkg) + assert not any( + "_gluon_dsa" in m for m in sys.modules + ), "importing v4_attention_kernels must not eagerly import the gluon backend" + + +def test_attention_module_uses_lazy_gluon_helpers(): + """``deepseek_v4_attention`` imports the lazy loader + arch guard, not gluon entries.""" + mod = importlib.import_module("primus.backends.megatron.core.transformer.deepseek_v4_attention") + assert hasattr(mod, "load_gluon_attention_backends") + assert hasattr(mod, "_require_gfx950") + # no eager module-level gluon entry bindings + assert not hasattr(mod, "v4_attention_gluon") + assert not hasattr(mod, "v4_csa_attention_gluon") + + +def test_require_gfx950_rejects_non_gfx950(monkeypatch): + """``_require_gfx950`` raises a clear error on a non-gfx950 device.""" + torch = pytest.importorskip("torch") + from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( + _require_gfx950, + ) + + class _FakeProps: + gcnArchName = "gfx942:sramecc+:xnack-" + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _idx: _FakeProps()) + with pytest.raises(RuntimeError, match="gfx950"): + _require_gfx950() + + +def test_require_gfx950_rejects_no_device(monkeypatch): + """``_require_gfx950`` raises when no accelerator is available.""" + torch = pytest.importorskip("torch") + from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( + _require_gfx950, + ) + + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + with pytest.raises(RuntimeError, match="gfx950"): + _require_gfx950() + + +def test_require_gfx950_accepts_gfx950(monkeypatch): + """``_require_gfx950`` passes on a gfx950 device.""" + torch = pytest.importorskip("torch") + from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( + _require_gfx950, + ) + + class _FakeProps: + gcnArchName = "gfx950:sramecc+:xnack-" + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _idx: _FakeProps()) + _require_gfx950() # must not raise diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_core_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_core_attention.py new file mode 100644 index 000000000..2f4505008 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_core_attention.py @@ -0,0 +1,500 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-3 P22 — ``core_attention`` integration for DeepSeek-V4. + +The dense (``compress_ratio == 0``) layers route their softmax-and-attend +through ``provider.core_attention()`` (PrimusTurboAttention when +``use_turbo_attention=True``, TEDotProductAttention otherwise). HCA +(``compress_ratio == 128``) and CSA (``compress_ratio == 4``) layers +**do not** get a ``core_attention`` slot — their joint softmax and +per-query top-K gather can't be expressed as stock flash-attention. + +This file is the unit-test side of P22's test gates: + +* **G18 (spec surface)** — the V4 attention submodules dataclass exposes + a ``core_attention`` slot; the spec helper emits it for dense layers + only. +* **G18a (alias contract)** — when V4's per-head sink is on AND the + built ``core_attention`` advertises ``use_sink_attention=True``, the + attention module aliases ``self.core_attention.sinks`` to + ``self.attn_sink`` so the released-checkpoint key path + ``layers.{i}.attn.attn_sink`` still loads. +* **G18b (CPU forward equivalence)** — a ``MockTurbo`` core-attention + class that replicates the eager-Python scaled-dot-product math is + injected; the dense-path output must match the original eager-Python + forward within a tight numerical tolerance. The full + Turbo-vs-eager-Python equivalence at full V4-Flash dims is exercised + by the P22 smoke gate on ``mi355-gpu-12``. +* **G18c (no core_attention on HCA / CSA)** — the spec helper does not + emit ``core_attention`` for ``compress_ratio in {4, 128}``. +""" + +from __future__ import annotations + +from dataclasses import fields + +import pytest +import torch +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import ModuleSpec + +# --------------------------------------------------------------------------- +# G18 — submodules surface +# --------------------------------------------------------------------------- + + +class TestSubmodulesSurface: + """``DeepseekV4AttentionSubmodules`` must expose a ``core_attention`` slot.""" + + def test_core_attention_field_present(self): + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( # noqa: F401 + DeepSeekV4TransformerConfig, + ) + from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( + DeepseekV4AttentionSubmodules, + ) + + names = {f.name for f in fields(DeepseekV4AttentionSubmodules)} + assert "core_attention" in names, ( + "Plan-3 P22 added ``core_attention`` to " + f"DeepseekV4AttentionSubmodules; current fields: {sorted(names)}." + ) + + +# --------------------------------------------------------------------------- +# G18 — spec helper emits core_attention for dense only +# --------------------------------------------------------------------------- + + +@pytest.fixture +def _tp1_distributed(): + """1-rank torch.distributed (gloo) with Megatron model-parallel state.""" + import os + + import torch.distributed as dist + from megatron.core import parallel_state + + if dist.is_initialized(): + yield + return + + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29539") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("LOCAL_RANK", "0") + + dist.init_process_group(backend="gloo", world_size=1, rank=0) + try: + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + ) + yield + finally: + parallel_state.destroy_model_parallel() + dist.destroy_process_group() + + +def _make_v4_cfg(*, attn_sink: bool, attn_sliding_window: int = 0): + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, + ) + + return DeepSeekV4TransformerConfig( + num_layers=1, + hidden_size=64, + num_attention_heads=4, + ffn_hidden_size=128, + kv_channels=16, + q_lora_rank=32, + o_groups=2, + o_lora_rank=16, + attn_sink=attn_sink, + attn_sliding_window=attn_sliding_window, + qk_pos_emb_head_dim=8, + num_query_groups=1, + multi_latent_attention=False, + params_dtype=torch.float32, + init_method=lambda w: torch.nn.init.normal_(w, std=0.02), + output_layer_init_method=lambda w: torch.nn.init.normal_(w, std=0.02), + use_cpu_initialization=True, + perform_initialization=True, + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + ) + + +def _make_v4_submods(cfg, compress_ratio: int): + from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import ( + DeepSeekV4SpecProvider, + ) + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import ( + _build_v4_attention_submodules, + ) + + provider = DeepSeekV4SpecProvider(config=cfg) + return _build_v4_attention_submodules( + config=cfg, + provider=provider, + compress_ratio=compress_ratio, + ) + + +class TestSpecEmission: + """The spec helper emits ``core_attention`` only for dense layers.""" + + def test_dense_emits_core_attention(self, _tp1_distributed): + cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=4) + submods = _make_v4_submods(cfg, compress_ratio=0) + assert submods.core_attention is not None, "compress_ratio == 0 must emit a ``core_attention`` spec." + + def test_csa_no_core_attention(self, _tp1_distributed): + cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=4) + submods = _make_v4_submods(cfg, compress_ratio=4) + assert submods.core_attention is None, ( + "compress_ratio == 4 (CSA) must NOT emit ``core_attention`` — " + "per-query top-K gather is not a flash-attn pattern." + ) + + def test_hca_no_core_attention(self, _tp1_distributed): + cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=4) + submods = _make_v4_submods(cfg, compress_ratio=128) + assert submods.core_attention is None, ( + "compress_ratio == 128 (HCA) must NOT emit ``core_attention`` — " + "joint softmax across two key streams needs an LSE-returning kernel." + ) + + +# --------------------------------------------------------------------------- +# G18 — alias contract + forward equivalence +# --------------------------------------------------------------------------- + + +class _MockTurboCoreAttention(torch.nn.Module): + """Stand-in for :class:`PrimusTurboAttention`. + + Replicates the eager-Python scaled-dot-product math (with optional + learned per-head sinks) on CPU, so the V4 attention module's dense + forward can be exercised end-to-end without a CUDA / Turbo build. + + Public surface that ``DeepseekV4Attention`` reads: + + * ``self.use_sink_attention`` — whether learned sinks are honored. + * ``self.sinks`` — ``[num_attention_heads]`` parameter (aliased by + V4's ``self.attn_sink`` after construction). + * ``forward(q, k, v, mask, attn_mask_type=...)`` — accepts ``sbhd`` + Q / K / V (K, V may be MQA-shape ``[S, B, 1, D]``). + """ + + def __init__( + self, + *, + config, + layer_number, + attn_mask_type, + attention_type, + softmax_scale, + k_channels=None, + v_channels=None, + cp_comm_type="p2p", + pg_collection=None, + ): + super().__init__() + self.config = config + self.layer_number = layer_number + self.softmax_scale = float(softmax_scale) + self.use_sink_attention = bool(getattr(config, "attn_sink", False)) + self._head_dim = int(k_channels or config.kv_channels) + self._num_heads = int(config.num_attention_heads) + if self.use_sink_attention: + self.sinks = torch.nn.Parameter(torch.zeros(self._num_heads)) + else: + self.sinks = None + + def forward(self, q, k, v, mask, attn_mask_type=None): + # Inputs in qkv_format="sbhd": [S, B, H_q/H_kv, D]. + S_q, B, H_q, D = q.shape + S_k = k.shape[0] + H_kv = k.shape[2] + + # MQA broadcast. + if H_kv != H_q: + k = k.expand(S_k, B, H_q, D) + v = v.expand(S_k, B, H_q, D) + + # [S, B, H, D] -> [B, H, S, D]. + q_bh = q.permute(1, 2, 0, 3).contiguous().float() + k_bh = k.permute(1, 2, 0, 3).contiguous().float() + v_bh = v.permute(1, 2, 0, 3).contiguous() + + logits = torch.matmul(q_bh, k_bh.transpose(-2, -1)) * self.softmax_scale + + # Causal mask. + if attn_mask_type == AttnMaskType.causal: + causal = torch.full((S_q, S_k), float("-inf"), device=q.device) + causal = torch.triu(causal, diagonal=1) + logits = logits + causal + + # Sink column. + if self.use_sink_attention and self.sinks is not None: + sink_col = self.sinks.float().view(1, H_q, 1, 1).expand(B, H_q, S_q, 1) + logits_aug = torch.cat([logits, sink_col], dim=-1) + logits_aug = logits_aug - logits_aug.amax(dim=-1, keepdim=True).detach() + probs = logits_aug.softmax(dim=-1)[..., :-1] + else: + logits = logits - logits.amax(dim=-1, keepdim=True).detach() + probs = logits.softmax(dim=-1) + + out_bh = torch.matmul(probs.to(v_bh.dtype), v_bh) # [B, H, S, D] + + # [B, H, S, D] -> [S, B, H, D] -> [S, B, H*D]. + out_sbh = out_bh.permute(2, 0, 1, 3).contiguous() + return out_sbh.view(S_q, B, H_q * D) + + +def _build_v4_attention_with(cfg, *, core_attention_class): + """Build a 1L V4 dense attention with an explicit ``core_attention`` class.""" + from primus.backends.megatron.core.transformer.deepseek_v4_attention import ( + DeepseekV4Attention, + ) + from primus.backends.megatron.core.transformer.dual_rope import DualRoPE + + submods = _make_v4_submods(cfg, compress_ratio=0) + if core_attention_class is None: + submods.core_attention = None + else: + submods.core_attention = ModuleSpec(module=core_attention_class) + + rope = DualRoPE( + rotary_dim=cfg.qk_pos_emb_head_dim, + rope_theta=10000.0, + compress_rope_theta=10000.0, + ) + return DeepseekV4Attention( + config=cfg, + rope=rope, + compress_ratio=0, + submodules=submods, + layer_number=0, + ) + + +class TestG18AliasAndForward: + """``self.attn_sink`` aliasing + dense forward equivalence.""" + + def test_sink_alias_when_turbo_supports_sink(self, _tp1_distributed): + cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=0) + attn = _build_v4_attention_with(cfg, core_attention_class=_MockTurboCoreAttention) + + assert attn.core_attention is not None + assert attn._use_core_attention is True + assert attn.attn_sink is not None + assert attn.core_attention.sinks is attn.attn_sink, ( + "Plan-3 P22: when the core-attention class advertises " + "``use_sink_attention=True`` and V4 has ``attn_sink=True``, " + "the V4 attention module must alias ``core_attention.sinks`` " + "to ``self.attn_sink`` so the released-checkpoint key path " + "``layers.{i}.attn.attn_sink`` still loads via state-dict." + ) + + def test_no_alias_when_core_attention_lacks_sink(self, _tp1_distributed): + # Default _build_v4_attention_submodules emits provider.core_attention() + # which is TEDotProductAttention (no use_sink_attention attr). + cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=4) + attn = _build_v4_attention_with(cfg, core_attention_class=None) + # Re-emit core_attention via provider so we exercise the TE path. + # _build_v4_attention_with(core_attention_class=None) drops the slot; + # the default emission already builds TEDotProductAttention which + # does NOT advertise use_sink_attention -> no alias, eager fallback. + # Reproduce the default emission and rebuild. + from primus.backends.megatron.core.extensions.transformer_engine_spec_provider import ( + DeepSeekV4SpecProvider, + ) + + provider = DeepSeekV4SpecProvider(config=cfg) + attn2 = _build_v4_attention_with(cfg, core_attention_class=provider.core_attention()) + assert attn2.core_attention is not None, "TE core_attention must build" + assert attn2._use_core_attention is False, ( + "When the built core_attention does not support learned sinks " + "(e.g. TEDotProductAttention) and V4 sink is on, the dense " + "path must fall back to eager-Python so the inline " + "softmax-with-sink math still produces the correct output." + ) + + @pytest.mark.skipif( + not torch.cuda.is_available(), + reason="V4 specs build TE-parallel linears whose forward requires CUDA.", + ) + def test_dense_forward_matches_eager(self, _tp1_distributed): + torch.manual_seed(0) + cfg = _make_v4_cfg(attn_sink=True, attn_sliding_window=0) + + attn_eager = _build_v4_attention_with(cfg, core_attention_class=None) + attn_turbo = _build_v4_attention_with(cfg, core_attention_class=_MockTurboCoreAttention) + + # Copy weights so both modules have identical parameters. + with torch.no_grad(): + for (n_e, p_e), (n_t, p_t) in zip( + attn_eager.named_parameters(), + attn_turbo.named_parameters(), + ): + if p_e.shape == p_t.shape: + p_t.copy_(p_e) + + # Sanity: both should have the same attn_sink storage. + assert attn_eager.attn_sink is not None + assert attn_turbo.core_attention.sinks is attn_turbo.attn_sink + + # Move both modules onto CUDA — the V4 spec emits TE-parallel + # linears whose forward requires CUDA tensors. V4's DualRoPE is + # held by reference (``self._rope = [rope]``) so it isn't moved + # by ``attn.to(device)``; move it explicitly. + device = torch.device("cuda") + attn_eager = attn_eager.to(device) + attn_turbo = attn_turbo.to(device) + attn_eager._rope[0].to(device) + attn_turbo._rope[0].to(device) + + B, S, D = 2, 7, cfg.hidden_size + hidden = torch.randn(B, S, D, device=device) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, S) + + attn_eager.eval() + attn_turbo.eval() + with torch.no_grad(): + out_eager = attn_eager(hidden, position_ids) + out_turbo = attn_turbo(hidden, position_ids) + + assert out_eager.shape == out_turbo.shape == (B, S, D) + max_abs = (out_eager - out_turbo).abs().max().item() + assert max_abs < 5e-3, ( + f"Dense forward via core_attention diverged from eager-Python " + f"forward beyond tolerance: max-abs={max_abs}. Plan-3 P22 " + "expects the two paths to compute identical math (the mock " + "core_attention replicates the eager softmax-with-sink kernel)." + ) + + +# --------------------------------------------------------------------------- +# G18 — sink-attention args plumbing +# --------------------------------------------------------------------------- + + +class TestSinkAttentionArgsPlumbing: + """``deepseek_v4_builder._maybe_plumb_v4_sink_attention_args``.""" + + def test_plumbing_fires_when_turbo_attention_on_seq_le_window(self): + """Plumbing zeros out the window when it covers the full sequence + (mathematically equivalent to full causal — avoids the aiter + Triton SWA gap).""" + from types import SimpleNamespace + + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import ( + _maybe_plumb_v4_sink_attention_args, + ) + + args = SimpleNamespace( + enable_primus_turbo=True, + use_turbo_attention=True, + attn_sink=True, + attn_sliding_window=128, + seq_length=128, + ) + _maybe_plumb_v4_sink_attention_args(args) + + assert args.use_sink_attention is True + # Window == seq_length -> drop window (full causal is identical). + assert args.sink_sliding_window == 0 + assert args.sink_window_even_layers_only is False + + def test_plumbing_warns_when_real_swa_requested(self): + """When the V4 config genuinely needs SWA (seq > window) the + plumbing zeros out the window with a warning — aiter Triton + flash-attn doesn't support SWA yet, so the V4 dense layer + attends to all causal tokens (deviation from V4-Flash math).""" + from types import SimpleNamespace + + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import ( + _maybe_plumb_v4_sink_attention_args, + ) + + args = SimpleNamespace( + enable_primus_turbo=True, + use_turbo_attention=True, + attn_sink=True, + attn_sliding_window=128, + seq_length=4096, + ) + _maybe_plumb_v4_sink_attention_args(args) + + assert args.use_sink_attention is True + # SWA requested but kernel doesn't support it -> drop window. + assert args.sink_sliding_window == 0 + assert args.sink_window_even_layers_only is False + + def test_plumbing_skips_when_turbo_off(self): + from types import SimpleNamespace + + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import ( + _maybe_plumb_v4_sink_attention_args, + ) + + args = SimpleNamespace( + enable_primus_turbo=False, + use_turbo_attention=False, + attn_sink=True, + attn_sliding_window=128, + ) + _maybe_plumb_v4_sink_attention_args(args) + + # Plumbing must not touch args when Turbo is off. + assert not hasattr(args, "use_sink_attention") or args.use_sink_attention in ( + None, + False, + ) + + def test_plumbing_skips_when_attn_sink_off(self): + from types import SimpleNamespace + + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import ( + _maybe_plumb_v4_sink_attention_args, + ) + + args = SimpleNamespace( + enable_primus_turbo=True, + use_turbo_attention=True, + attn_sink=False, + attn_sliding_window=128, + ) + _maybe_plumb_v4_sink_attention_args(args) + + assert not hasattr(args, "use_sink_attention") or args.use_sink_attention in ( + None, + False, + ) + + def test_plumbing_respects_explicit_user_override(self): + from types import SimpleNamespace + + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_builders import ( + _maybe_plumb_v4_sink_attention_args, + ) + + args = SimpleNamespace( + enable_primus_turbo=True, + use_turbo_attention=True, + attn_sink=True, + attn_sliding_window=128, + seq_length=4096, + use_sink_attention=True, + sink_sliding_window=256, # user override + ) + _maybe_plumb_v4_sink_attention_args(args) + + # Don't clobber user-set sink_sliding_window. + assert args.sink_sliding_window == 256 diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_fp8_indexer.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_fp8_indexer.py new file mode 100644 index 000000000..64393d0e0 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_fp8_indexer.py @@ -0,0 +1,140 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for the DeepSeek-V4 FP8 (E4M3) Indexer QK path. + +CPU-friendly. Validates that: + +* :func:`fake_quantize_fp8_e4m3` is a bounded-error, dtype-preserving, + zero-safe fake-quantization. +* The FP8 Indexer selects almost the same top-k compressed positions as the + BF16 reference (high top-k overlap) — the QK precision drop must not change + which positions the CSA selector picks for the vast majority of queries. +* Output shapes / sentinel (-1) semantics are unchanged by the FP8 path. +""" + +from __future__ import annotations + +import copy + +import pytest +import torch + +from primus.backends.megatron.core.transformer.indexer import ( + Indexer, + fake_quantize_fp8_e4m3, +) + +_HAS_FP8 = hasattr(torch, "float8_e4m3fn") + + +def _fp8_cast_works() -> bool: + if not _HAS_FP8: + return False + try: + torch.zeros(4).to(torch.float8_e4m3fn).to(torch.float32) + return True + except (RuntimeError, TypeError): + return False + + +pytestmark = pytest.mark.skipif( + not _fp8_cast_works(), reason="torch.float8_e4m3fn cast unsupported on this build/device" +) + + +# --------------------------------------------------------------------------- +# fake_quantize_fp8_e4m3 +# --------------------------------------------------------------------------- + + +def test_fake_quant_preserves_dtype_and_shape(): + x = torch.randn(3, 5, 7, dtype=torch.float32) + xq = fake_quantize_fp8_e4m3(x) + assert xq.dtype == x.dtype + assert xq.shape == x.shape + + +def test_fake_quant_zero_input_is_safe(): + x = torch.zeros(4, 4) + xq = fake_quantize_fp8_e4m3(x) + assert torch.equal(xq, x) + + +def test_fake_quant_bounded_relative_error(): + # E4M3 has 3 mantissa bits -> ~2^-3 = 12.5% worst-case step; with dynamic + # per-tensor scaling the typical relative error on sizable values is small. + torch.manual_seed(0) + x = torch.randn(2048, dtype=torch.float32) * 2.0 + xq = fake_quantize_fp8_e4m3(x) + # Compare only on non-tiny values (tiny values have large relative error + # but negligible absolute impact on the QK dot product). + big = x.abs() > 0.1 * x.abs().max() + rel = ((xq[big] - x[big]).abs() / x[big].abs()).max().item() + assert rel < 0.2, f"FP8 fake-quant relative error too large: {rel}" + + +# --------------------------------------------------------------------------- +# Indexer FP8 QK path vs BF16 reference +# --------------------------------------------------------------------------- + + +def _make_indexer(use_fp8_qk: bool) -> Indexer: + return Indexer( + hidden_size=64, + index_head_dim=16, + index_n_heads=4, + index_topk=4, + compress_ratio=4, + use_fp8_qk=use_fp8_qk, + ) + + +def _topk_overlap(idx_a: torch.Tensor, idx_b: torch.Tensor) -> float: + """Mean per-query Jaccard overlap of selected (valid) pool positions.""" + B, S, _K = idx_a.shape + overlaps = [] + for b in range(B): + for s in range(S): + a = {int(i) for i in idx_a[b, s].tolist() if i >= 0} + c = {int(i) for i in idx_b[b, s].tolist() if i >= 0} + if not a and not c: + continue + union = a | c + overlaps.append(len(a & c) / max(1, len(union))) + return sum(overlaps) / max(1, len(overlaps)) + + +def test_fp8_indexer_topk_overlaps_bf16_reference(): + torch.manual_seed(0) + ref = _make_indexer(use_fp8_qk=False) + fp8 = _make_indexer(use_fp8_qk=True) + # Share identical weights so the only difference is the FP8 QK rounding. + fp8.load_state_dict(copy.deepcopy(ref.state_dict())) + + hidden = torch.randn(2, 32, 64, dtype=torch.float32) + with torch.no_grad(): + idx_ref, _ = ref(hidden) + idx_fp8, _ = fp8(hidden) + + assert idx_ref.shape == idx_fp8.shape == (2, 32, 4) + # Sentinel semantics preserved (same set of masked-out / early queries). + assert torch.equal(idx_ref < 0, idx_fp8 < 0) + + overlap = _topk_overlap(idx_ref, idx_fp8) + assert overlap >= 0.6, f"FP8 indexer top-k overlap with BF16 too low: {overlap}" + + +def test_fp8_flag_off_is_identical_to_reference(): + torch.manual_seed(0) + a = _make_indexer(use_fp8_qk=False) + b = _make_indexer(use_fp8_qk=False) + b.load_state_dict(copy.deepcopy(a.state_dict())) + hidden = torch.randn(1, 16, 64, dtype=torch.float32) + with torch.no_grad(): + idx_a, _ = a(hidden) + idx_b, _ = b(hidden) + assert torch.equal(idx_a, idx_b) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_dsa_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_dsa_attention.py new file mode 100644 index 000000000..d76ac5548 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_dsa_attention.py @@ -0,0 +1,300 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Gluon DeepSeek-V4 attention fwd+bwd correctness, in the **V4 form** (gfx950). + +Unlike a raw sparse-MLA test (separate 64-rope, random absolute-token topk), +this validates the *production* V4 invocation paths — the autograd adapters in +``v4_csa_attention_gluon`` that DeepseekV4Attention actually dispatches to — +against the eager V4 references for all three layer kinds: + +* ``compress_ratio == 0`` (dense / SWA) -> :func:`v4_attention_gluon` +* ``compress_ratio == 128`` (HCA) -> :func:`v4_attention_gluon` +* ``compress_ratio == 4`` (CSA) -> :func:`v4_csa_attention_gluon` + +The V4 layout has ``head_dim = 512`` with RoPE applied *in-place* (K = V = 512, +score over 512); the adapter feeds the 512 latent as the gluon "lora" with a +zero rope pad and builds ``kv = [local ++ pool]`` / ``topk = [SWA window ++ +pool]`` (padded to a multiple of 64 so the gluon bwd dKV tiling is valid — HCA +160 -> 192). We compare full fwd (O) and torch-autograd bwd (dQ, dlatent, dpool, +dsink) of the bf16 gluon adapter against the fp32 eager reference. + +GPU-only; skipped off gfx950 / when Gluon is unavailable. ``B = 2`` to exercise +the per-batch token-flattening + topk-offset logic. +""" + +from __future__ import annotations + +import math + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("Gluon DSA kernels require CUDA / HIP", allow_module_level=True) + +# Gluon is an experimental Triton submodule; skip cleanly if absent. +pytest.importorskip("triton", reason="Triton not installed") +try: + from triton.experimental import gluon # noqa: F401 +except Exception: # pragma: no cover - environment-dependent + pytest.skip("triton.experimental.gluon unavailable", allow_module_level=True) + +_ARCH = torch.cuda.get_device_properties(0).gcnArchName +if "gfx950" not in _ARCH: + pytest.skip(f"ported Gluon DSA kernels target gfx950; got {_ARCH}", allow_module_level=True) + +from primus.backends.megatron.core.transformer.sliding_window_kv import ( # noqa: E402 + sliding_window_causal_mask, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import ( # noqa: E402 + eager_v4_attention, + eager_v4_csa_attention, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon import ( # noqa: E402 + v4_attention_gluon, + v4_csa_attention_gluon, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_triton import ( # noqa: E402 + v4_attention_v2, + v4_csa_attention_v2, +) + +D = 512 # V4 head_dim (RoPE baked in-place) +_SWA = 128 + +# Fused single-latent (K == V) sparse-MLA backends sharing the V4-form adapter: +# (dense/HCA attention wrapper, CSA-from-pool wrapper). +_BACKENDS = { + "gluon": (v4_attention_gluon, v4_csa_attention_gluon), + "triton_v2": (v4_attention_v2, v4_csa_attention_v2), +} + + +# --------------------------------------------------------------------------- +# Comparison helpers (bf16 kernel vs fp32 eager autograd) +# --------------------------------------------------------------------------- + + +def _stats(a, b, *, sig=1e-2): + a = a.float() + b = b.float() + d = (a - b).abs() + m = b.abs() > sig + rel = (d[m] / b.abs()[m]) if m.any() else d.new_zeros(0) + med_rel = rel.median().item() if rel.numel() else 0.0 + av, bv = a.flatten(), b.flatten() + cos_err = 1.0 - (av @ bv / (av.norm() * bv.norm() + 1e-30)).item() + return d.max().item(), med_rel, cos_err + + +def _check(name, a, b, *, abs_tol, sig=1e-2, med=None, cos=None): + max_abs, med_rel, cos_err = _stats(a, b, sig=sig) + print(f" {name:8s} max_abs={max_abs:.3e} median_rel={med_rel:.3e} cos_err={cos_err:.3e}", flush=True) + assert max_abs < abs_tol, f"{name} max_abs {max_abs:.3e} >= {abs_tol}" + if med is not None: + assert med_rel < med, f"{name} median_rel {med_rel:.3e} >= {med}" + if cos is not None: + assert cos_err < cos, f"{name} cos_err {cos_err:.3e} >= {cos}" + + +def _leaf(x, *, fp32): + """Detached leaf clone (fp32 for the eager ref, bf16 for the gluon adapter).""" + y = x.float() if fp32 else x.clone() + return y.detach().requires_grad_(True) + + +# --------------------------------------------------------------------------- +# (B, H, S, sink_init) +# --------------------------------------------------------------------------- +_CASES = [ + (2, 16, 256, None), + (2, 16, 256, 1.0), + (1, 32, 512, -1.0), +] + + +@pytest.mark.parametrize("backend", list(_BACKENDS)) +@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v)) +def test_v4_dense_matches_eager(B, H, S, sink_init, backend): + """cr=0 dense/SWA: sparse-MLA v4_attention_v1 fwd+bwd vs eager_v4_attention.""" + attn_fn = _BACKENDS[backend][0] + g = torch.Generator(device="cuda").manual_seed(0) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None + + qg, latg = _leaf(q, fp32=False), _leaf(lat, fp32=False) + sg = _leaf(sink, fp32=False) if sink is not None else None + kg = latg.unsqueeze(1).expand(B, H, S, D) + og = attn_fn( + qg, kg, kg, sink=sg, swa_window=_SWA, additive_mask=None, attn_dropout=0.0, training=True, scale=scale + ) + og.backward(do) + + qf, latf = _leaf(q, fp32=True), _leaf(lat, fp32=True) + sf = _leaf(sink, fp32=True) if sink is not None else None + kf = latf.unsqueeze(1).expand(B, H, S, D) + of = eager_v4_attention( + qf, + kf, + kf, + sink=sf, + swa_window=_SWA, + additive_mask=None, + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[{backend} V4 dense] B={B} H={H} S={S} sink={sink_init}", flush=True) + assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16 + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + if sink is not None: + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) + + +@pytest.mark.parametrize("backend", list(_BACKENDS)) +@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v)) +def test_v4_hca_matches_eager(B, H, S, sink_init, backend): + """cr=128 HCA: sparse-MLA v4_attention_v1 (local++pool) fwd+bwd vs eager_v4_attention.""" + attn_fn = _BACKENDS[backend][0] + cr = 128 + P = max(S // cr, 1) + g = torch.Generator(device="cuda").manual_seed(2) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None + + # HCA causal pool mask [S, P]: pool slot p visible to query s iff (p+1)*cr-1 <= s. + ti = torch.arange(S, device="cuda").view(S, 1) + ps = torch.arange(P, device="cuda").view(1, P) + pool_mask = torch.where( + ((ps + 1) * cr - 1) <= ti, torch.zeros((), device="cuda"), torch.tensor(float("-inf"), device="cuda") + ).to(torch.bfloat16) + + qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False) + sg = _leaf(sink, fp32=False) if sink is not None else None + kg = torch.cat([latg.unsqueeze(1).expand(B, H, S, D), poolg.unsqueeze(1).expand(B, H, P, D)], dim=2) + og = attn_fn( + qg, + kg, + kg, + sink=sg, + swa_window=_SWA, + additive_mask=pool_mask, + attn_dropout=0.0, + training=True, + scale=scale, + hca_local_seqlen=S, + ) + og.backward(do) + + qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True) + sf = _leaf(sink, fp32=True) if sink is not None else None + kf = torch.cat([latf.unsqueeze(1).expand(B, H, S, D), poolf.unsqueeze(1).expand(B, H, P, D)], dim=2) + local_mask = sliding_window_causal_mask(S, _SWA, device="cuda", dtype=torch.float32) + full_mask = torch.cat([local_mask, pool_mask.float()], dim=1) # [S, S+P] + of = eager_v4_attention( + qf, + kf, + kf, + sink=sf, + swa_window=0, + additive_mask=full_mask, + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[{backend} V4 HCA] B={B} H={H} S={S} P={P} sink={sink_init}", flush=True) + assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16 + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + _check("dpool", poolg.grad, poolf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + if sink is not None: + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) + + +@pytest.mark.parametrize("backend", list(_BACKENDS)) +@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v)) +def test_v4_csa_matches_eager(B, H, S, sink_init, backend): + """cr=4 CSA: sparse-MLA v4_csa_attention_v1 fwd+bwd vs eager_v4_csa_attention.""" + csa_fn = _BACKENDS[backend][1] + P = max(S // 4, 1) + K = min(128, P) + g = torch.Generator(device="cuda").manual_seed(3) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None + + # Per-query top-K pool indices in [0, P), with ~1/8 invalid (-1). + topk_idxs = torch.randint(0, P, (B, S, K), generator=g, device="cuda", dtype=torch.int32) + drop = torch.rand(B, S, K, generator=g, device="cuda") < 0.125 + topk_idxs = torch.where(drop, torch.full_like(topk_idxs, -1), topk_idxs) + idx = topk_idxs.clamp(0, P - 1).long() + bidx = torch.arange(B, device="cuda").view(B, 1, 1) + sparse_mask_inf = torch.where( + topk_idxs < 0, torch.tensor(float("-inf"), device="cuda"), torch.zeros((), device="cuda") + ) # [B, S, K] + + qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False) + sg = _leaf(sink, fp32=False) if sink is not None else None + klg = latg.unsqueeze(1).expand(B, H, S, D) + og = csa_fn( + qg, + klg, + klg, + poolg, + topk_idxs=topk_idxs, + sink=sg, + swa_window=_SWA, + attn_dropout=0.0, + training=True, + scale=scale, + ) + og.backward(do) + + qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True) + sf = _leaf(sink, fp32=True) if sink is not None else None + klf = latf.unsqueeze(1).expand(B, H, S, D) + gathered = poolf[bidx, idx] # [B, S, K, D] (differentiable gather into poolf) + of = eager_v4_csa_attention( + qf, + klf, + klf, + gathered, + sink=sf, + swa_window=_SWA, + sparse_mask=sparse_mask_inf.float(), + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[{backend} V4 CSA] B={B} H={H} S={S} P={P} K={K} sink={sink_init}", flush=True) + assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16 + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + _check("dpool", poolg.grad, poolf.grad, abs_tol=2e-1, sig=1e-3, med=2e-2, cos=1e-3) + if sink is not None: + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v2_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v2_attention.py new file mode 100644 index 000000000..012e0d988 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v2_attention.py @@ -0,0 +1,277 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""gluon_v2 DeepSeek-V4 attention fwd+bwd correctness, in the **V4 form** (gfx950), +against the fp32 eager references. + +gluon_v2 = Gluon sparse-MLA forward + plain-Triton chunked-gather backward. This +validates the production V4 invocation paths (the autograd adapters in +``v4_csa_attention_gluon_v2`` that DeepseekV4Attention would dispatch to) for all +three layer kinds: + +* ``compress_ratio == 0`` (dense / SWA) -> :func:`v4_attention_gluon_v2` +* ``compress_ratio == 128`` (HCA) -> :func:`v4_attention_gluon_v2` +* ``compress_ratio == 4`` (CSA) -> :func:`v4_csa_attention_gluon_v2` + +We compare full fwd (O) and torch-autograd bwd (dQ, dlatent, dpool, dsink) of the +bf16 gluon_v2 adapter against the fp32 eager reference (the same tolerances the +gluon backend UT uses). GPU-only; skipped off gfx950 / when Gluon is unavailable. +""" + +from __future__ import annotations + +import math + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("gluon_v2 kernels require CUDA / HIP", allow_module_level=True) + +pytest.importorskip("triton", reason="Triton not installed") +try: + from triton.experimental import gluon # noqa: F401 +except Exception: # pragma: no cover - environment-dependent + pytest.skip("triton.experimental.gluon unavailable", allow_module_level=True) + +_ARCH = torch.cuda.get_device_properties(0).gcnArchName +if "gfx950" not in _ARCH: + pytest.skip(f"gluon_v2 Gluon fwd targets gfx950; got {_ARCH}", allow_module_level=True) + +from primus.backends.megatron.core.transformer.sliding_window_kv import ( # noqa: E402 + sliding_window_causal_mask, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import ( # noqa: E402 + eager_v4_attention, + eager_v4_csa_attention, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon_v2 import ( # noqa: E402 + v4_attention_gluon_v2, + v4_csa_attention_gluon_v2, +) + +# Fail early with the build hint if the installed triton can't compile the Gluon fwd. +try: + from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v2.dsa_fwd_v4_gluon import ( # noqa: E402 + _gluon_available, + ) + + if not _gluon_available(): + pytest.skip("Gluon dialect unavailable for gluon_v2 fwd", allow_module_level=True) +except Exception: # noqa: BLE001 + pass + +D = 512 # V4 head_dim (RoPE baked in-place) +_SWA = 128 + + +def _stats(a, b, *, sig=1e-2): + a = a.float() + b = b.float() + d = (a - b).abs() + m = b.abs() > sig + rel = (d[m] / b.abs()[m]) if m.any() else d.new_zeros(0) + med_rel = rel.median().item() if rel.numel() else 0.0 + av, bv = a.flatten(), b.flatten() + cos_err = 1.0 - (av @ bv / (av.norm() * bv.norm() + 1e-30)).item() + return d.max().item(), med_rel, cos_err + + +def _check(name, a, b, *, abs_tol, sig=1e-2, med=None, cos=None): + max_abs, med_rel, cos_err = _stats(a, b, sig=sig) + print(f" {name:8s} max_abs={max_abs:.3e} median_rel={med_rel:.3e} cos_err={cos_err:.3e}", flush=True) + assert max_abs < abs_tol, f"{name} max_abs {max_abs:.3e} >= {abs_tol}" + if med is not None: + assert med_rel < med, f"{name} median_rel {med_rel:.3e} >= {med}" + if cos is not None: + assert cos_err < cos, f"{name} cos_err {cos_err:.3e} >= {cos}" + + +def _leaf(x, *, fp32): + y = x.float() if fp32 else x.clone() + return y.detach().requires_grad_(True) + + +_CASES = [ + (2, 16, 256, None), + (2, 16, 256, 1.0), + (1, 32, 512, -1.0), +] + + +@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v)) +def test_v4_gluon_v2_dense_matches_eager(B, H, S, sink_init): + """cr=0 dense/SWA: gluon_v2 fwd+bwd vs eager_v4_attention.""" + g = torch.Generator(device="cuda").manual_seed(0) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None + + qg, latg = _leaf(q, fp32=False), _leaf(lat, fp32=False) + sg = _leaf(sink, fp32=False) if sink is not None else None + kg = latg.unsqueeze(1).expand(B, H, S, D) + og = v4_attention_gluon_v2( + qg, kg, kg, sink=sg, swa_window=_SWA, additive_mask=None, attn_dropout=0.0, training=True, scale=scale + ) + og.backward(do) + + qf, latf = _leaf(q, fp32=True), _leaf(lat, fp32=True) + sf = _leaf(sink, fp32=True) if sink is not None else None + kf = latf.unsqueeze(1).expand(B, H, S, D) + of = eager_v4_attention( + qf, + kf, + kf, + sink=sf, + swa_window=_SWA, + additive_mask=None, + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[gluon_v2 V4 dense] B={B} H={H} S={S} sink={sink_init}", flush=True) + assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16 + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + if sink is not None: + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) + + +@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v)) +def test_v4_gluon_v2_hca_matches_eager(B, H, S, sink_init): + """cr=128 HCA: gluon_v2 (local++pool) fwd+bwd vs eager_v4_attention.""" + cr = 128 + P = max(S // cr, 1) + g = torch.Generator(device="cuda").manual_seed(2) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None + + ti = torch.arange(S, device="cuda").view(S, 1) + ps = torch.arange(P, device="cuda").view(1, P) + pool_mask = torch.where( + ((ps + 1) * cr - 1) <= ti, torch.zeros((), device="cuda"), torch.tensor(float("-inf"), device="cuda") + ).to(torch.bfloat16) + + qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False) + sg = _leaf(sink, fp32=False) if sink is not None else None + kg = torch.cat([latg.unsqueeze(1).expand(B, H, S, D), poolg.unsqueeze(1).expand(B, H, P, D)], dim=2) + og = v4_attention_gluon_v2( + qg, + kg, + kg, + sink=sg, + swa_window=_SWA, + additive_mask=pool_mask, + attn_dropout=0.0, + training=True, + scale=scale, + hca_local_seqlen=S, + ) + og.backward(do) + + qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True) + sf = _leaf(sink, fp32=True) if sink is not None else None + kf = torch.cat([latf.unsqueeze(1).expand(B, H, S, D), poolf.unsqueeze(1).expand(B, H, P, D)], dim=2) + local_mask = sliding_window_causal_mask(S, _SWA, device="cuda", dtype=torch.float32) + full_mask = torch.cat([local_mask, pool_mask.float()], dim=1) + of = eager_v4_attention( + qf, + kf, + kf, + sink=sf, + swa_window=0, + additive_mask=full_mask, + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[gluon_v2 V4 HCA] B={B} H={H} S={S} P={P} sink={sink_init}", flush=True) + assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16 + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + _check("dpool", poolg.grad, poolf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + if sink is not None: + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) + + +@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v)) +def test_v4_gluon_v2_csa_matches_eager(B, H, S, sink_init): + """cr=4 CSA: gluon_v2 fwd+bwd vs eager_v4_csa_attention.""" + P = max(S // 4, 1) + K = min(128, P) + g = torch.Generator(device="cuda").manual_seed(3) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None + + topk_idxs = torch.randint(0, P, (B, S, K), generator=g, device="cuda", dtype=torch.int32) + drop = torch.rand(B, S, K, generator=g, device="cuda") < 0.125 + topk_idxs = torch.where(drop, torch.full_like(topk_idxs, -1), topk_idxs) + idx = topk_idxs.clamp(0, P - 1).long() + bidx = torch.arange(B, device="cuda").view(B, 1, 1) + sparse_mask_inf = torch.where( + topk_idxs < 0, torch.tensor(float("-inf"), device="cuda"), torch.zeros((), device="cuda") + ) + + qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False) + sg = _leaf(sink, fp32=False) if sink is not None else None + klg = latg.unsqueeze(1).expand(B, H, S, D) + og = v4_csa_attention_gluon_v2( + qg, + klg, + klg, + poolg, + topk_idxs=topk_idxs, + sink=sg, + swa_window=_SWA, + attn_dropout=0.0, + training=True, + scale=scale, + ) + og.backward(do) + + qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True) + sf = _leaf(sink, fp32=True) if sink is not None else None + klf = latf.unsqueeze(1).expand(B, H, S, D) + gathered = poolf[bidx, idx] + of = eager_v4_csa_attention( + qf, + klf, + klf, + gathered, + sink=sf, + swa_window=_SWA, + sparse_mask=sparse_mask_inf.float(), + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[gluon_v2 V4 CSA] B={B} H={H} S={S} P={P} K={K} sink={sink_init}", flush=True) + assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16 + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + _check("dpool", poolg.grad, poolf.grad, abs_tol=2e-1, sig=1e-3, med=2e-2, cos=1e-3) + if sink is not None: + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v3_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v3_attention.py new file mode 100644 index 000000000..f9a31d52f --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_gluon_v3_attention.py @@ -0,0 +1,322 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""gluon_v3 DeepSeek-V4 attention fwd+bwd correctness against fp32 eager refs.""" + +from __future__ import annotations + +import math + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("gluon_v3 kernels require CUDA / HIP", allow_module_level=True) + +pytest.importorskip("triton", reason="Triton not installed") +try: + from triton.experimental import gluon # noqa: F401 +except Exception: # pragma: no cover - environment-dependent + pytest.skip("triton.experimental.gluon unavailable", allow_module_level=True) + +_ARCH = torch.cuda.get_device_properties(0).gcnArchName +if "gfx950" not in _ARCH: + pytest.skip(f"gluon_v3 targets gfx950; got {_ARCH}", allow_module_level=True) + +from primus.backends.megatron.core.transformer.sliding_window_kv import ( # noqa: E402 + sliding_window_causal_mask, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import ( # noqa: E402 + eager_v4_attention, + eager_v4_csa_attention, +) + +try: + from primus.backends.megatron.core.transformer.v4_attention_kernels._gluon_v3 import ( # noqa: E402 + sparse_mla_bwd_v4_gluon_v3, + sparse_mla_fwd_v4_gluon_v3, + ) + from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_gluon_v3 import ( # noqa: E402 + v4_attention_gluon_v3, + v4_csa_attention_gluon_v3, + ) +except ImportError as exc: # pragma: no cover - environment-dependent + pytest.skip(f"gluon_v3 unavailable: {exc}", allow_module_level=True) + +D = 512 +_ROPE = 64 +_SWA = 128 + + +def _stats(a, b, *, sig=1e-2): + a = a.float() + b = b.float() + d = (a - b).abs() + m = b.abs() > sig + rel = (d[m] / b.abs()[m]) if m.any() else d.new_zeros(0) + med_rel = rel.median().item() if rel.numel() else 0.0 + av, bv = a.flatten(), b.flatten() + cos_err = 1.0 - (av @ bv / (av.norm() * bv.norm() + 1e-30)).item() + return d.max().item(), med_rel, cos_err + + +def _check(name, a, b, *, abs_tol, sig=1e-2, med=None, cos=None): + max_abs, med_rel, cos_err = _stats(a, b, sig=sig) + print(f" {name:8s} max_abs={max_abs:.3e} median_rel={med_rel:.3e} cos_err={cos_err:.3e}", flush=True) + assert max_abs < abs_tol, f"{name} max_abs {max_abs:.3e} >= {abs_tol}" + if med is not None: + assert med_rel < med, f"{name} median_rel {med_rel:.3e} >= {med}" + if cos is not None: + assert cos_err < cos, f"{name} cos_err {cos_err:.3e} >= {cos}" + + +def _leaf(x, *, fp32): + y = x.float() if fp32 else x.clone() + return y.detach().requires_grad_(True) + + +def _eager_sparse_mla(q, kv, topk, sink, *, scale): + q_lora = q[:, :, :D] + kv_lora = kv[:, 0, :D] + safe = topk.clamp(min=0).long() + gathered = kv_lora[safe] + valid = topk >= 0 + scores = torch.einsum("thd,tkd->thk", q_lora, gathered) * scale + scores = torch.where(valid[:, None, :], scores, torch.full_like(scores, float("-inf"))) + if sink is not None: + sink_scores = sink.view(1, -1, 1).expand(q.shape[0], q.shape[1], 1) + all_scores = torch.cat([scores, sink_scores], dim=-1) + else: + all_scores = scores + probs = torch.softmax(all_scores, dim=-1) + out = torch.einsum("thk,tkd->thd", probs[:, :, : topk.shape[1]], gathered) + lse = torch.logsumexp(all_scores, dim=-1) + return out, lse + + +_CASES = [ + (1, 16, 128, None), + (1, 32, 256, 1.0), +] + + +@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v)) +def test_v4_gluon_v3_dense_matches_eager(B, H, S, sink_init): + g = torch.Generator(device="cuda").manual_seed(10) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None + + qg, latg = _leaf(q, fp32=False), _leaf(lat, fp32=False) + sg = _leaf(sink, fp32=False) if sink is not None else None + kg = latg.unsqueeze(1).expand(B, H, S, D) + og = v4_attention_gluon_v3( + qg, kg, kg, sink=sg, swa_window=_SWA, additive_mask=None, attn_dropout=0.0, training=True, scale=scale + ) + og.backward(do) + + qf, latf = _leaf(q, fp32=True), _leaf(lat, fp32=True) + sf = _leaf(sink, fp32=True) if sink is not None else None + kf = latf.unsqueeze(1).expand(B, H, S, D) + of = eager_v4_attention( + qf, + kf, + kf, + sink=sf, + swa_window=_SWA, + additive_mask=None, + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[gluon_v3 V4 dense] B={B} H={H} S={S} sink={sink_init}", flush=True) + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + if sink is not None: + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) + + +def test_v4_gluon_v3_hca_matches_eager(): + B, H, S, cr = 1, 32, 256, 128 + P = max(S // cr, 1) + g = torch.Generator(device="cuda").manual_seed(11) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), 1.0, dtype=torch.float32, device="cuda") + + ti = torch.arange(S, device="cuda").view(S, 1) + ps = torch.arange(P, device="cuda").view(1, P) + pool_mask = torch.where( + ((ps + 1) * cr - 1) <= ti, torch.zeros((), device="cuda"), torch.tensor(float("-inf"), device="cuda") + ).to(torch.bfloat16) + + qg, latg, poolg, sg = ( + _leaf(q, fp32=False), + _leaf(lat, fp32=False), + _leaf(pool, fp32=False), + _leaf(sink, fp32=False), + ) + kg = torch.cat([latg.unsqueeze(1).expand(B, H, S, D), poolg.unsqueeze(1).expand(B, H, P, D)], dim=2) + og = v4_attention_gluon_v3( + qg, + kg, + kg, + sink=sg, + swa_window=_SWA, + additive_mask=pool_mask, + attn_dropout=0.0, + training=True, + scale=scale, + hca_local_seqlen=S, + ) + og.backward(do) + + qf, latf, poolf, sf = ( + _leaf(q, fp32=True), + _leaf(lat, fp32=True), + _leaf(pool, fp32=True), + _leaf(sink, fp32=True), + ) + kf = torch.cat([latf.unsqueeze(1).expand(B, H, S, D), poolf.unsqueeze(1).expand(B, H, P, D)], dim=2) + full_mask = torch.cat( + [sliding_window_causal_mask(S, _SWA, device="cuda", dtype=torch.float32), pool_mask.float()], dim=1 + ) + of = eager_v4_attention( + qf, + kf, + kf, + sink=sf, + swa_window=0, + additive_mask=full_mask, + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[gluon_v3 V4 HCA] B={B} H={H} S={S} P={P}", flush=True) + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + _check("dpool", poolg.grad, poolf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) + + +def test_v4_gluon_v3_csa_adapter_matches_eager(): + B, H, S = 1, 32, 256 + P, K = max(S // 4, 1), 64 + g = torch.Generator(device="cuda").manual_seed(12) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), -1.0, dtype=torch.float32, device="cuda") + topk_idxs = torch.randint(0, P, (B, S, K), generator=g, device="cuda", dtype=torch.int32) + topk_idxs = torch.where( + torch.rand(B, S, K, generator=g, device="cuda") < 0.125, torch.full_like(topk_idxs, -1), topk_idxs + ) + idx = topk_idxs.clamp(0, P - 1).long() + bidx = torch.arange(B, device="cuda").view(B, 1, 1) + sparse_mask_inf = torch.where( + topk_idxs < 0, torch.tensor(float("-inf"), device="cuda"), torch.zeros((), device="cuda") + ) + + qg, latg, poolg, sg = ( + _leaf(q, fp32=False), + _leaf(lat, fp32=False), + _leaf(pool, fp32=False), + _leaf(sink, fp32=False), + ) + klg = latg.unsqueeze(1).expand(B, H, S, D) + og = v4_csa_attention_gluon_v3( + qg, + klg, + klg, + poolg, + topk_idxs=topk_idxs, + sink=sg, + swa_window=_SWA, + attn_dropout=0.0, + training=True, + scale=scale, + ) + og.backward(do) + + qf, latf, poolf, sf = ( + _leaf(q, fp32=True), + _leaf(lat, fp32=True), + _leaf(pool, fp32=True), + _leaf(sink, fp32=True), + ) + klf = latf.unsqueeze(1).expand(B, H, S, D) + gathered = poolf[bidx, idx] + of = eager_v4_csa_attention( + qf, + klf, + klf, + gathered, + sink=sf, + swa_window=_SWA, + sparse_mask=sparse_mask_inf.float(), + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[gluon_v3 V4 CSA adapter] B={B} H={H} S={S} P={P} K={K}", flush=True) + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + _check("dlatent", latg.grad, latf.grad, abs_tol=1e-1, sig=1e-3, med=2e-2, cos=1e-3) + _check("dpool", poolg.grad, poolf.grad, abs_tol=2e-1, sig=1e-3, med=2e-2, cos=1e-3) + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) + + +@pytest.mark.parametrize("H,pool_k,T", [(64, 512, 64), (128, 1024, 64)], ids=["h64_topk640", "h128_topk1152"]) +def test_v4_gluon_v3_round9_csa_formula_path_matches_eager(H, pool_k, T): + g = torch.Generator(device="cuda").manual_seed(13 + H) + scale = 1.0 / math.sqrt(D) + q512 = torch.randn(T, H, D, generator=g, device="cuda", dtype=torch.bfloat16) + kv512 = torch.randn(T + pool_k, 1, D, generator=g, device="cuda", dtype=torch.bfloat16) + q = torch.cat([q512, torch.zeros(T, H, _ROPE, device="cuda", dtype=torch.bfloat16)], dim=-1).contiguous() + kv = torch.cat( + [kv512, torch.zeros(T + pool_k, 1, _ROPE, device="cuda", dtype=torch.bfloat16)], dim=-1 + ).contiguous() + do = torch.randn(T, H, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.randn(H, generator=g, device="cuda", dtype=torch.float32) * 0.1 + + ti = torch.arange(T, device="cuda").view(T, 1) + win = ti - _SWA + 1 + torch.arange(_SWA, device="cuda").view(1, _SWA) + win = torch.where(win >= 0, win, torch.full_like(win, -1)) + pool_topk = T + torch.randint(0, pool_k, (T, pool_k), generator=g, device="cuda", dtype=torch.int64) + topk = torch.cat([win, pool_topk], dim=1).to(torch.int32).contiguous() + + qg, kvg, sg = _leaf(q, fp32=False), _leaf(kv, fp32=False), _leaf(sink, fp32=False) + og, lseg = sparse_mla_fwd_v4_gluon_v3(qg, kvg, topk, attn_sink=sg, kv_lora_rank=D, scale=scale) + dqg, dkvg, dsg = sparse_mla_bwd_v4_gluon_v3( + qg, kvg, og, do, topk, lseg, attn_sink=sg, kv_lora_rank=D, scale=scale + ) + + qf, kvf, sf = _leaf(q, fp32=True), _leaf(kv, fp32=True), _leaf(sink, fp32=True) + of, lsef = _eager_sparse_mla(qf, kvf, topk, sf, scale=scale) + of.backward(do.float()) + + print(f"\n[gluon_v3 round9 CSA formula] T={T} H={H} TOPK={topk.shape[1]}", flush=True) + _check("O", og, of, abs_tol=5e-2, med=3e-2, cos=2e-3) + _check("LSE", lseg, lsef, abs_tol=1e-2, sig=1e-3, med=1e-3, cos=1e-6) + _check("dQ", dqg[:, :, :D], qf.grad[:, :, :D], abs_tol=1e-1, sig=1e-3, med=3e-2, cos=2e-3) + _check("dKV", dkvg[:, :, :D], kvf.grad[:, :, :D], abs_tol=3e-1, sig=1e-3, med=5e-2, cos=5e-3) + _check("dSink", dsg, sf.grad, abs_tol=1.0, sig=1e-3, med=8e-2, cos=2e-2) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_moe.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_moe.py new file mode 100644 index 000000000..dd7859335 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_moe.py @@ -0,0 +1,504 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for V4 MoE forward (G5, plan-2 §04). + +Pins :class:`DeepseekV4MoE` against the HF reference at +``DeepSeek-V4-Flash/inference/model.py:MoE.forward``. The test runs on +GPU (CUDA/HIP), fp32, with ``pg_collection=None`` so the V4 MoE uses its +local-experts path (per-expert dispatch loop, no Megatron dispatcher). + +Pass criteria (G5): +* 1L MoE forward agrees with the inline HF reference to <= 1e-3 max-abs + for both routing modes (learned + hash). +* Output dtype / shape match input. +* Gradient flows from MoE output back into the gate ``weight`` (router) + AND the expert ``w1`` / ``w2`` / ``w3`` Linears. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import pytest +import torch +import torch.nn.functional as F + +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, +) +from primus.backends.megatron.core.transformer.clamped_swiglu import ClampedSwiGLUMLP +from primus.backends.megatron.core.transformer.moe.v4_hash_router import ( + DeepseekV4HashRouter, +) +from primus.backends.megatron.core.transformer.moe.v4_moe import ( + DeepseekV4MoE, + DeepseekV4MoESubmodules, +) +from primus.backends.megatron.core.transformer.moe.v4_topk_router import ( + DeepseekV4LearnedRouter, +) + + +@pytest.fixture(autouse=True) +def _v4_moe_on_cuda(monkeypatch): + """The V4 MoE routers are GPU-only: both the Triton and eager paths require + CUDA/HIP tensors, so default tensor creation to the CUDA device and skip on + a CPU-only host. Force the eager router path (``PRIMUS_V4_ROUTER_TRITON=0``) + so this stays an exact MoE-vs-HF reference check. + """ + if not torch.cuda.is_available(): + pytest.skip("DeepSeek-V4 MoE requires a CUDA/HIP device") + monkeypatch.setenv("PRIMUS_V4_ROUTER_TRITON", "0") + torch.set_default_device("cuda") + try: + yield + finally: + torch.set_default_device("cpu") + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_moe_config( + *, + hidden_size: int, + moe_intermediate_size: int, + shared_expert_intermediate_size: Optional[int], + num_experts: int, + topk: int, + swiglu_limit: float, + score_function: str, + num_hash_layers: int, + vocab_size: int, + enable_expert_bias: bool = False, + topk_scaling_factor: float = 1.0, +) -> DeepSeekV4TransformerConfig: + """Minimal V4 config for the MoE CPU smoke.""" + return DeepSeekV4TransformerConfig( + num_layers=1, + hidden_size=hidden_size, + num_attention_heads=4, + num_query_groups=1, + kv_channels=32, + ffn_hidden_size=hidden_size * 4, # only consumed by attention path + moe_ffn_hidden_size=moe_intermediate_size, + moe_intermediate_size=moe_intermediate_size, + moe_shared_expert_intermediate_size=shared_expert_intermediate_size, + num_moe_experts=num_experts, + moe_router_topk=topk, + moe_router_score_function=score_function, + moe_router_enable_expert_bias=enable_expert_bias, + moe_router_topk_scaling_factor=topk_scaling_factor, + swiglu_limit=swiglu_limit, + num_hash_layers=num_hash_layers, + hash_routing_seed=11, + vocab_size=vocab_size, + padded_vocab_size=vocab_size, + layernorm_epsilon=1.0e-6, + norm_epsilon=1.0e-6, + attention_dropout=0.0, + hidden_dropout=0.0, + # MLATransformerConfig requirements (not exercised by MoE-only test) + qk_pos_emb_head_dim=8, + qk_head_dim=24, + v_head_dim=32, + kv_lora_rank=32, + rope_type="rope", + rotary_base=10000.0, + rotary_scaling_factor=1.0, + rotary_percent=1.0, + original_max_position_embeddings=2048, + ) + + +def _make_moe( + config: DeepSeekV4TransformerConfig, + *, + layer_idx: int, +) -> DeepseekV4MoE: + """Build a CPU-friendly V4 MoE (no pg_collection -> local-experts path).""" + return DeepseekV4MoE( + config=config, + layer_idx=layer_idx, + pg_collection=None, + submodules=DeepseekV4MoESubmodules(), + ) + + +# --------------------------------------------------------------------------- +# Inline HF reference +# --------------------------------------------------------------------------- + + +def _hf_gate_forward( + *, + hidden_flat: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + score_function: str, + topk: int, + route_scale: float, + tid2eid: Optional[torch.Tensor], + token_ids: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Mirrors ``Gate.forward`` exactly. Returns dense ``(weights, indices)``.""" + scores = F.linear(hidden_flat.float(), weight.float()) + if score_function == "softmax": + scores = scores.softmax(dim=-1) + elif score_function == "sigmoid": + scores = scores.sigmoid() + else: + scores = F.softplus(scores).sqrt() + original_scores = scores + if bias is not None: + scores = scores + bias.float() + if tid2eid is not None: + assert token_ids is not None + indices = tid2eid[token_ids.reshape(-1).long()].long() + else: + indices = scores.topk(topk, dim=-1).indices + weights = original_scores.gather(1, indices) + if score_function != "softmax": + weights = weights / weights.sum(dim=-1, keepdim=True) + weights = weights * route_scale + return weights, indices + + +def _hf_expert_forward( + expert: ClampedSwiGLUMLP, + x: torch.Tensor, + weights: Optional[torch.Tensor], + *, + swiglu_limit: float, +) -> torch.Tensor: + """Mirrors ``Expert.forward`` (pre-mul clamp + post-w1/w3 weight scaling).""" + gate = expert.w1(x).float() + up = expert.w3(x).float() + if swiglu_limit > 0.0: + up = up.clamp(min=-swiglu_limit, max=swiglu_limit) + gate = gate.clamp(max=swiglu_limit) + h = F.silu(gate) * up + if weights is not None: + h = weights * h + return expert.w2(h.to(x.dtype)) + + +def _hf_moe_forward( + *, + moe: DeepseekV4MoE, + hidden: torch.Tensor, + token_ids: Optional[torch.Tensor], + score_function: str, + swiglu_limit: float, + route_scale: float, +) -> torch.Tensor: + """Inline reference for the MoE forward, sharing weights with ``moe``. + + Builds the (weights, indices) dense form via ``_hf_gate_forward`` and + runs the per-expert dispatch loop the way the released reference + does, then adds the shared-expert contribution. All tensors share + storage with the live :class:`DeepseekV4MoE` so a single set of + weights drives both forward paths. + """ + shape = hidden.shape + flat_hidden = hidden.reshape(-1, moe.hidden_size) + + # Gate + if moe.use_hash_router: + router = moe.router + assert isinstance(router, DeepseekV4HashRouter) + weights, indices = _hf_gate_forward( + hidden_flat=flat_hidden, + weight=router.weight, + bias=None, + score_function=score_function, + topk=moe.moe_router_topk, + route_scale=route_scale, + tid2eid=router.tid2eid, + token_ids=token_ids, + ) + else: + router = moe.learned_router + assert isinstance(router, DeepseekV4LearnedRouter) + weights, indices = _hf_gate_forward( + hidden_flat=flat_hidden, + weight=router.weight, + bias=router.expert_bias, + score_function=score_function, + topk=moe.moe_router_topk, + route_scale=route_scale, + tid2eid=None, + token_ids=None, + ) + + y = torch.zeros_like(flat_hidden, dtype=torch.float32) + n_routed = moe.num_routed_experts + counts = torch.bincount(indices.flatten(), minlength=n_routed).tolist() + assert moe.local_experts is not None + for local_i, global_i in enumerate(moe.local_expert_indices): + if counts[global_i] == 0: + continue + expert = moe.local_experts[local_i] + idx, top = torch.where(indices == global_i) + y[idx] += _hf_expert_forward( + expert, + flat_hidden[idx], + weights[idx, top, None], + swiglu_limit=swiglu_limit, + ).float() + + if moe.shared_expert is not None: + assert isinstance(moe.shared_expert, ClampedSwiGLUMLP) + y = ( + y + + _hf_expert_forward( + moe.shared_expert, + flat_hidden, + None, + swiglu_limit=swiglu_limit, + ).float() + ) + + return y.type_as(flat_hidden).view(*shape) + + +# --------------------------------------------------------------------------- +# Construction sanity +# --------------------------------------------------------------------------- + + +def test_v4_moe_subclasses_megatron_module() -> None: + """``DeepseekV4MoE`` is a ``MegatronModule`` (config plumbing parity).""" + from megatron.core.transformer.module import MegatronModule + + assert issubclass(DeepseekV4MoE, MegatronModule) + + +def test_v4_moe_cpu_path_builds_local_experts_and_shared_expert() -> None: + """When ``pg_collection`` is None the MoE skips the dispatcher and + builds local :class:`ClampedSwiGLUMLP` experts plus a single + :class:`ClampedSwiGLUMLP` shared expert (matches HF reference).""" + config = _make_moe_config( + hidden_size=12, + moe_intermediate_size=24, + shared_expert_intermediate_size=24, + num_experts=4, + topk=2, + swiglu_limit=7.0, + score_function="sqrtsoftplus", + num_hash_layers=0, + vocab_size=16, + ) + moe = _make_moe(config, layer_idx=0) + + assert moe.token_dispatcher is None + assert moe.grouped_experts is None + assert moe.local_experts is not None + assert len(moe.local_experts) == 4 + assert all(isinstance(e, ClampedSwiGLUMLP) for e in moe.local_experts) + assert isinstance(moe.shared_expert, ClampedSwiGLUMLP) + assert moe.local_expert_indices == [0, 1, 2, 3] + + +def test_v4_moe_set_layer_number_updates_attribute() -> None: + """``set_layer_number`` mirrors :class:`BaseMoELayer`.""" + config = _make_moe_config( + hidden_size=12, + moe_intermediate_size=24, + shared_expert_intermediate_size=None, + num_experts=4, + topk=2, + swiglu_limit=7.0, + score_function="sqrtsoftplus", + num_hash_layers=0, + vocab_size=16, + ) + moe = _make_moe(config, layer_idx=0) + moe.set_layer_number(7) + assert moe.layer_number == 7 + + +# --------------------------------------------------------------------------- +# G5: numerical alignment vs HF reference +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("score_function", ["sqrtsoftplus", "sigmoid", "softmax"]) +@pytest.mark.parametrize("with_shared_expert", [True, False]) +def test_v4_moe_learned_layer_matches_hf_reference(score_function: str, with_shared_expert: bool) -> None: + """G5: learned router + clamped-SwiGLU MoE forward agrees with + the inline HF reference on a 1L toy (CPU fp32).""" + torch.manual_seed(2025) + H = 16 + I = 24 + config = _make_moe_config( + hidden_size=H, + moe_intermediate_size=I, + shared_expert_intermediate_size=I if with_shared_expert else None, + num_experts=4, + topk=2, + swiglu_limit=7.0, + score_function=score_function, + num_hash_layers=0, # learned router for layer_idx >= 0 + vocab_size=32, + ) + moe = _make_moe(config, layer_idx=0) + + hidden = torch.randn(2, 5, H, dtype=torch.float32) * 0.5 + out = moe(hidden, token_ids=None) + ref = _hf_moe_forward( + moe=moe, + hidden=hidden, + token_ids=None, + score_function=score_function, + swiglu_limit=7.0, + route_scale=1.0, + ) + + assert out.shape == hidden.shape + assert torch.isfinite(out).all() + max_abs = (out - ref).abs().max().item() + assert max_abs <= 1.0e-3, f"max-abs vs HF reference = {max_abs}" + + +@pytest.mark.parametrize("score_function", ["sqrtsoftplus", "sigmoid", "softmax"]) +def test_v4_moe_hash_layer_matches_hf_reference(score_function: str) -> None: + """G5: hash router + clamped-SwiGLU MoE forward agrees with + the inline HF reference on a 1L toy (CPU fp32).""" + torch.manual_seed(7) + H = 16 + I = 24 + V = 32 + config = _make_moe_config( + hidden_size=H, + moe_intermediate_size=I, + shared_expert_intermediate_size=I, + num_experts=4, + topk=2, + swiglu_limit=7.0, + score_function=score_function, + num_hash_layers=1, # hash router for layer_idx == 0 + vocab_size=V, + ) + moe = _make_moe(config, layer_idx=0) + + hidden = torch.randn(2, 5, H, dtype=torch.float32) * 0.5 + token_ids = torch.randint(0, V, (2, 5), dtype=torch.long) + + out = moe(hidden, token_ids=token_ids) + ref = _hf_moe_forward( + moe=moe, + hidden=hidden, + token_ids=token_ids, + score_function=score_function, + swiglu_limit=7.0, + route_scale=1.0, + ) + + assert out.shape == hidden.shape + assert torch.isfinite(out).all() + max_abs = (out - ref).abs().max().item() + assert max_abs <= 1.0e-3, f"max-abs vs HF reference = {max_abs}" + + +def test_v4_moe_route_scale_propagates_to_output() -> None: + """``moe_router_topk_scaling_factor`` (HF ``route_scale``) is honored.""" + torch.manual_seed(3) + H = 12 + config = _make_moe_config( + hidden_size=H, + moe_intermediate_size=24, + shared_expert_intermediate_size=None, + num_experts=4, + topk=2, + swiglu_limit=7.0, + score_function="sqrtsoftplus", + num_hash_layers=0, + vocab_size=16, + topk_scaling_factor=2.5, + ) + moe = _make_moe(config, layer_idx=0) + hidden = torch.randn(1, 3, H, dtype=torch.float32) * 0.5 + + out = moe(hidden, token_ids=None) + ref = _hf_moe_forward( + moe=moe, + hidden=hidden, + token_ids=None, + score_function="sqrtsoftplus", + swiglu_limit=7.0, + route_scale=2.5, + ) + assert (out - ref).abs().max().item() <= 1.0e-3 + + +def test_v4_moe_gradient_flows_to_router_and_experts() -> None: + """G5 follow-up: backward pass populates grads on ``router.weight`` + and on at least one expert's ``w1`` / ``w2`` / ``w3`` weight.""" + torch.manual_seed(13) + H = 12 + config = _make_moe_config( + hidden_size=H, + moe_intermediate_size=24, + shared_expert_intermediate_size=24, + num_experts=4, + topk=2, + swiglu_limit=7.0, + score_function="sqrtsoftplus", + num_hash_layers=0, + vocab_size=16, + ) + moe = _make_moe(config, layer_idx=0) + hidden = torch.randn(2, 4, H, dtype=torch.float32, requires_grad=False) * 0.5 + + out = moe(hidden, token_ids=None) + out.sum().backward() + + assert moe.learned_router is not None + assert moe.learned_router.weight.grad is not None + assert torch.isfinite(moe.learned_router.weight.grad).all() + + # At least one routed expert must have received gradient. Some experts + # may be unselected in a small toy run; assert >= 1 has a non-zero grad. + any_expert_grad = False + assert moe.local_experts is not None + for expert in moe.local_experts: + if ( + expert.w1.weight.grad is not None + and expert.w1.weight.grad.abs().sum().item() > 0.0 + and expert.w2.weight.grad is not None + and expert.w3.weight.grad is not None + ): + any_expert_grad = True + break + assert any_expert_grad, "expected at least one routed expert to receive gradient" + + # Shared expert is always-on so it always sees gradient. + assert moe.shared_expert is not None + assert moe.shared_expert.w1.weight.grad is not None + assert moe.shared_expert.w1.weight.grad.abs().sum().item() > 0.0 + + +def test_v4_moe_hash_layer_requires_token_ids() -> None: + """Hash-routed layers raise a clear error if ``token_ids`` is missing.""" + config = _make_moe_config( + hidden_size=12, + moe_intermediate_size=24, + shared_expert_intermediate_size=None, + num_experts=4, + topk=2, + swiglu_limit=7.0, + score_function="sqrtsoftplus", + num_hash_layers=1, + vocab_size=16, + ) + moe = _make_moe(config, layer_idx=0) + hidden = torch.randn(1, 3, 12, dtype=torch.float32) + with pytest.raises(ValueError, match="token_ids is required"): + moe(hidden, token_ids=None) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py new file mode 100644 index 000000000..fff550143 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_mtp.py @@ -0,0 +1,361 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for V4 MTP integration with upstream +:class:`MultiTokenPredictionBlock` (P16) plus plan-2 P17 cleanup +guarantees. + +What this file covers (CPU-friendly): + +* :func:`get_v4_mtp_block_spec` returns a well-formed ``ModuleSpec``: + the outer module is :class:`MultiTokenPredictionBlock`, its + ``submodules.layer_specs`` is a list of length ``mtp_num_layers``, + and each entry is a :class:`MultiTokenPredictionLayer` ``ModuleSpec`` + whose ``mtp_model_layer`` is the V4 hybrid-layer spec we passed in. +* Per-MTP-layer submodules pull V4 RMSNorm / column-parallel linear + from the V4 spec provider — the helper does not silently fall back + to TE / vanilla impls. +* The V4 hybrid layer's ``forward`` returns the upstream-compatible + ``(hidden_states, None)`` tuple so it can plug into + :meth:`MultiTokenPredictionLayer._proj_and_transformer_layer` which + unpacks ``hidden_states, _ = self.mtp_model_layer(...)``. +* The V4 attention spec advertises ``attn_mask_type`` so upstream MTP + validation passes (V4 manages its own SWA / sink mask internally, + but the field is required by the upstream pre-build assertion). +* The legacy primus-owned ``DeepseekV4MTPBlock`` is **gone** in plan-2 + P17 — its module is no longer importable, the + ``v4_use_custom_mtp_block`` config flag is removed, and + ``DeepseekV4Model.__init__`` no longer references either of them. +* :class:`DeepseekV4Model.forward` wires :func:`process_mtp_loss` and + :class:`MultiTokenPredictionBlock` (AST audit; full distributed run + is gated G7 in P19). + +The full ``mtp_num_layers=0`` vs ``mtp_num_layers=1`` main-LM loss +invariance gate (G7) requires distributed init and is tracked into +P19 distributed re-validation. +""" + +from __future__ import annotations + +import ast +import importlib +import inspect +from dataclasses import is_dataclass +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.multi_token_prediction import ( + MultiTokenPredictionBlock, + MultiTokenPredictionBlockSubmodules, + MultiTokenPredictionLayer, + MultiTokenPredictionLayerSubmodules, +) +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_layer import TransformerLayerSubmodules + +from primus.backends.megatron.core.models.deepseek_v4 import ( + DeepseekV4HybridLayer, + DeepseekV4HybridLayerSubmodules, + DeepseekV4TransformerBlock, + DeepseekV4TransformerBlockSubmodules, + get_v4_mtp_block_spec, +) +from primus.backends.megatron.core.models.deepseek_v4.build_context import ( + resolve_v4_provider, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_mtp_layer import ( + DeepseekV4MTPLayer, +) + +_REPO_ROOT = Path(__file__).resolve().parents[5] +_MODEL_PATH = _REPO_ROOT / ("primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_model.py") +_SPECS_PATH = _REPO_ROOT / ("primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_layer_specs.py") + + +# --------------------------------------------------------------------------- +# Minimal V4 config / layer-spec fixtures +# --------------------------------------------------------------------------- + + +def _make_v4_config(*, mtp_num_layers: int = 1): + """Tiny V4 config with the fields the MTP spec helper reads.""" + cfg = MagicMock() + cfg.hidden_size = 64 + cfg.num_layers = 2 + cfg.mtp_num_layers = mtp_num_layers + cfg.hc_mult = 1 + cfg.qk_pos_emb_head_dim = 16 + cfg.attn_sliding_window = 128 + cfg.norm_epsilon = 1e-6 + cfg.layernorm_epsilon = 1e-6 + cfg.use_mup = False + cfg.tensor_model_parallel_size = 1 + cfg.sequence_parallel = False + cfg.bf16 = False + cfg.fp16 = False + return cfg + + +def _placeholder_layer_spec() -> ModuleSpec: + """Stand-in for a V4 hybrid-layer spec; the helper just threads it through.""" + return ModuleSpec( + module=DeepseekV4HybridLayer, + params={"layer_idx": 0, "compress_ratio": 0}, + submodules=DeepseekV4HybridLayerSubmodules(), + ) + + +# --------------------------------------------------------------------------- +# get_v4_mtp_block_spec — structural assertions +# --------------------------------------------------------------------------- + + +def test_helper_returns_multi_token_prediction_block_spec() -> None: + cfg = _make_v4_config(mtp_num_layers=1) + spec = get_v4_mtp_block_spec(cfg, transformer_layer_spec=_placeholder_layer_spec()) + assert isinstance(spec, ModuleSpec) + assert spec.module is MultiTokenPredictionBlock + assert isinstance(spec.submodules, MultiTokenPredictionBlockSubmodules) + + +@pytest.mark.parametrize("mtp_num_layers", [1, 2, 3]) +def test_helper_emits_one_layer_spec_per_depth(mtp_num_layers: int) -> None: + cfg = _make_v4_config(mtp_num_layers=mtp_num_layers) + spec = get_v4_mtp_block_spec(cfg, transformer_layer_spec=_placeholder_layer_spec()) + layer_specs = spec.submodules.layer_specs + assert isinstance(layer_specs, list) + assert len(layer_specs) == mtp_num_layers + for layer_spec in layer_specs: + assert isinstance(layer_spec, ModuleSpec) + # V4 uses its own mHC-aware MTP layer (per-depth HyperHead), not the + # vanilla upstream MultiTokenPredictionLayer. + assert layer_spec.module is DeepseekV4MTPLayer + assert issubclass(DeepseekV4MTPLayer, MultiTokenPredictionLayer) + + +def test_helper_threads_v4_inner_layer_unchanged() -> None: + cfg = _make_v4_config(mtp_num_layers=2) + inner = _placeholder_layer_spec() + spec = get_v4_mtp_block_spec(cfg, transformer_layer_spec=inner) + for layer_spec in spec.submodules.layer_specs: + sub: MultiTokenPredictionLayerSubmodules = layer_spec.submodules + assert sub.mtp_model_layer is inner, ( + "MTP helper must thread the V4 hybrid layer spec through unchanged " + "so MTP depths share HC / hash-routing / clamped-SwiGLU with the main decoder." + ) + + +def test_helper_pulls_norm_and_linear_from_v4_provider() -> None: + cfg = _make_v4_config(mtp_num_layers=1) + # Resolve the provider through the same singleton helper the MTP spec + # builder uses (``resolve_v4_provider`` caches one provider on the config), + # so the identity comparison is against the exact provider that wired the + # spec — not a second, independently-constructed provider instance. + provider = resolve_v4_provider(cfg) + spec = get_v4_mtp_block_spec(cfg, transformer_layer_spec=_placeholder_layer_spec()) + sub = spec.submodules.layer_specs[0].submodules + expected_norm = provider.v4_norm_module() + expected_col = provider.column_parallel_linear() + assert sub.enorm is expected_norm + assert sub.hnorm is expected_norm + assert sub.layer_norm is expected_norm + assert sub.eh_proj is expected_col + + +def test_helper_extracts_last_hybrid_layer_from_block_spec() -> None: + """When handed the decoder *block* spec (as DeepseekV4Model does), the + MTP helper must extract a single hybrid-layer spec (the last one) as the + MTP inner layer — not thread the whole block spec through (which would + trip MultiTokenPredictionLayer's TransformerLayerSubmodules validation).""" + cfg = _make_v4_config(mtp_num_layers=1) + first = _placeholder_layer_spec() + last = _placeholder_layer_spec() + block_spec = ModuleSpec( + module=DeepseekV4TransformerBlock, + submodules=DeepseekV4TransformerBlockSubmodules(layer_specs=[first, last]), + ) + spec = get_v4_mtp_block_spec(cfg, transformer_layer_spec=block_spec) + sub = spec.submodules.layer_specs[0].submodules + assert sub.mtp_model_layer is last, ( + "MTP inner layer must be the last hybrid-layer spec extracted from the " + "decoder block spec (mirrors upstream GPT spec.layer_specs[-1])." + ) + + +def test_helper_rejects_zero_mtp_num_layers() -> None: + cfg = _make_v4_config(mtp_num_layers=0) + with pytest.raises(ValueError, match="mtp_num_layers >= 1"): + get_v4_mtp_block_spec(cfg, transformer_layer_spec=_placeholder_layer_spec()) + + +# --------------------------------------------------------------------------- +# DeepseekV4HybridLayer — upstream-compatible tuple return +# --------------------------------------------------------------------------- + + +def test_layer_submodules_extends_transformer_layer_submodules() -> None: + """Required so MultiTokenPredictionLayer.__init__'s submodules + isinstance check picks up the GPT path (not Mamba).""" + assert is_dataclass(DeepseekV4HybridLayerSubmodules) + assert issubclass(DeepseekV4HybridLayerSubmodules, TransformerLayerSubmodules) + + +def test_layer_forward_signature_returns_tuple() -> None: + """``forward`` is annotated as a tuple-returning callable so + ``MultiTokenPredictionLayer._proj_and_transformer_layer`` can unpack + ``hidden_states, _ = self.mtp_model_layer(...)`` without error.""" + src = inspect.getsource(DeepseekV4HybridLayer.forward) + assert "return x, None" in src, ( + "DeepseekV4HybridLayer.forward must return (hidden_states, None) for " + "upstream MultiTokenPredictionLayer compatibility." + ) + + +# --------------------------------------------------------------------------- +# V4 attention spec — declares an attn_mask_type for upstream MTP validation +# --------------------------------------------------------------------------- + + +def test_attention_spec_declares_supported_attn_mask_type() -> None: + src = _SPECS_PATH.read_text() + assert "AttnMaskType.causal" in src, ( + "V4 attention spec must declare attn_mask_type=AttnMaskType.causal " + "(in spec.params) so MultiTokenPredictionLayer's pre-build " + "validation accepts the V4 inner layer." + ) + # Sanity: import works and value resolves. + assert AttnMaskType.causal.name == "causal" + + +# --------------------------------------------------------------------------- +# Legacy V4 MTP block — retired in plan-2 P17 (G14 dead-code audit) +# --------------------------------------------------------------------------- + + +def test_legacy_mtp_block_module_is_gone() -> None: + """plan-2 P17 deletes :mod:`...deepseek_v4_mtp` outright. + + Importing the module must raise :class:`ImportError`; the symbol + must not be re-exported from the package surface either. + """ + with pytest.raises(ImportError): + importlib.import_module("primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_mtp") + + # Package surface — the legacy class must not leak back via __init__.py. + pkg = importlib.import_module("primus.backends.megatron.core.models.deepseek_v4") + assert not hasattr(pkg, "DeepseekV4MTPBlock"), ( + "DeepseekV4MTPBlock must be retired in plan-2 P17; the package " "must not re-export it." + ) + assert "DeepseekV4MTPBlock" not in getattr(pkg, "__all__", []), ( + "DeepseekV4MTPBlock must be removed from the package __all__ " "(plan-2 P17 dead-code audit)." + ) + + +# --------------------------------------------------------------------------- +# DeepseekV4Model.forward — process_mtp_loss + MTP block wired (AST audit) +# --------------------------------------------------------------------------- + + +def test_model_forward_calls_process_mtp_loss() -> None: + src = _MODEL_PATH.read_text() + tree = ast.parse(src) + has_call = False + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + if node.func.id == "process_mtp_loss": + has_call = True + break + assert has_call, ( + "DeepseekV4Model.forward must call process_mtp_loss to wire MTP " + "into the LM-loss pipeline (plan-2 P16)." + ) + + +def test_model_imports_upstream_mtp_machinery() -> None: + src = _MODEL_PATH.read_text() + for needle in ( + "from megatron.core.transformer.multi_token_prediction import", + "MultiTokenPredictionBlock", + "process_mtp_loss", + "mtp_on_this_rank", + "get_v4_mtp_block_spec", + ): + assert needle in src, f"DeepseekV4Model must import {needle!r} (P16)." + + +def test_model_init_routes_through_v4_mtp_spec_helper() -> None: + src = _MODEL_PATH.read_text() + assert "get_v4_mtp_block_spec(" in src, ( + "DeepseekV4Model.__init__ must use get_v4_mtp_block_spec to build " + "self.mtp; the spec is then handed to MultiTokenPredictionBlock." + ) + + +def test_model_no_longer_references_legacy_mtp_block() -> None: + """Plan-2 P17 retires the legacy MTP block AND the + ``v4_use_custom_mtp_block`` config flag; ``DeepseekV4Model``'s + source must not reference either of them anywhere outside a + historical comment that explicitly says the field is gone. + """ + src = _MODEL_PATH.read_text() + tree = ast.parse(src) + + bad_attr = [] + bad_name = [] + for node in ast.walk(tree): + if isinstance(node, ast.Attribute) and node.attr == "v4_use_custom_mtp_block": + bad_attr.append(getattr(node, "lineno", -1)) + if isinstance(node, ast.Name) and node.id == "DeepseekV4MTPBlock": + bad_name.append(getattr(node, "lineno", -1)) + if isinstance(node, (ast.ImportFrom, ast.Import)): + module = getattr(node, "module", None) or "" + if module.endswith("deepseek_v4_mtp"): + bad_attr.append(getattr(node, "lineno", -1)) + + assert not bad_attr, ( + "DeepseekV4Model must not reference v4_use_custom_mtp_block / " + "deepseek_v4_mtp after plan-2 P17 (lines: %s)." % bad_attr + ) + assert not bad_name, ( + "DeepseekV4Model must not call/reference DeepseekV4MTPBlock after " + "plan-2 P17 (lines: %s)." % bad_name + ) + + +def test_v4_config_no_longer_carries_legacy_mtp_fields() -> None: + """Plan-2 P17 dead-code audit: ``v4_use_custom_mtp_block`` and + ``mtp_compress_ratios`` are removed from the V4 config dataclass.""" + from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_transformer_config import ( + DeepSeekV4TransformerConfig, + ) + + fields = {f.name for f in DeepSeekV4TransformerConfig.__dataclass_fields__.values()} + assert "v4_use_custom_mtp_block" not in fields, ( + "v4_use_custom_mtp_block must be removed from DeepSeekV4TransformerConfig " + "(plan-2 P17 retired the legacy MTP block path)." + ) + assert "mtp_compress_ratios" not in fields, ( + "mtp_compress_ratios was only consumed by the legacy MTP block; " + "remove it alongside the block (plan-2 P17)." + ) + + +# --------------------------------------------------------------------------- +# Tiny CPU-only smoke: model __init__ in a no-MTP / no-distributed config +# does not crash and leaves self.mtp = None +# --------------------------------------------------------------------------- + + +def test_model_init_no_mtp_path_does_not_build_mtp() -> None: + """When ``mtp_num_layers == 0`` the model must not construct an MTP + block (regardless of distributed state).""" + # We don't build a real model here (it requires Megatron's full init + # machinery); we just confirm the __init__ source has the + # mtp_num_layers > 0 guard so the no-MTP path stays inert. + src = _MODEL_PATH.read_text() + assert "mtp_num_layers > 0" in src diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_routers.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_routers.py new file mode 100644 index 000000000..7e2e9b39e --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_routers.py @@ -0,0 +1,381 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Unit tests for V4 routers (G4, plan-2 §04). + +Pins both routers to the HF reference at +``DeepSeek-V4-Flash/inference/model.py:Gate.forward``. Two routers +share the same scoring path; they only differ in *selection* (top-K +argmax vs ``tid2eid`` lookup). + +Pass criteria (G4): +* Identical weights -> identical sparse ``(probs, routing_map)`` to the + HF reference (max-abs <= 1e-6 fp32). +* Routing weight gradient flows to ``weight`` (the learned gate). +* For the hash router, ``tid2eid`` is a parameter with + ``requires_grad=False`` (frozen) so checkpoint round-trips preserve + it without polluting the optimizer state. +""" + +from __future__ import annotations + +from typing import Optional + +import pytest +import torch +import torch.nn.functional as F + +from primus.backends.megatron.core.transformer.moe.v4_hash_router import ( + DeepseekV4HashRouter, + HashRouter, +) +from primus.backends.megatron.core.transformer.moe.v4_topk_router import ( + DeepseekV4LearnedRouter, + V4TopKRouter, + v4_score_fn, +) + + +@pytest.fixture(autouse=True) +def _v4_router_on_cuda(monkeypatch): + """The V4 routers are GPU-only: both the Triton and eager paths require + CUDA/HIP tensors, so default tensor creation to the CUDA device and skip on + a CPU-only host. Force the eager path (``PRIMUS_V4_ROUTER_TRITON=0``) so + these stay exact router-vs-HF reference checks; the fused Triton path has + its own parity test (``test_router_post_triton``). + """ + if not torch.cuda.is_available(): + pytest.skip("DeepSeek-V4 routers require a CUDA/HIP device") + monkeypatch.setenv("PRIMUS_V4_ROUTER_TRITON", "0") + torch.set_default_device("cuda") + try: + yield + finally: + torch.set_default_device("cpu") + + +# --------------------------------------------------------------------------- +# HF reference inline +# --------------------------------------------------------------------------- + + +def _hf_gate_forward( + *, + hidden: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + score_function: str, + topk: int, + route_scale: float, + tid2eid: Optional[torch.Tensor], + token_ids: Optional[torch.Tensor], +): + """Mirrors HF reference ``Gate.forward`` exactly. + + Returns dense ``(weights, indices)`` with shapes ``[N, K]``. + """ + flat_hidden = hidden.reshape(-1, hidden.shape[-1]) + scores = F.linear(flat_hidden.float(), weight.float()) + if score_function == "softmax": + scores = scores.softmax(dim=-1) + elif score_function == "sigmoid": + scores = scores.sigmoid() + else: + scores = F.softplus(scores).sqrt() + original_scores = scores + if bias is not None: + scores = scores + bias.float() + if tid2eid is not None: + assert token_ids is not None + flat_ids = token_ids.reshape(-1).long() + indices = tid2eid[flat_ids].long() + else: + indices = scores.topk(topk, dim=-1).indices + weights = original_scores.gather(1, indices) + if score_function != "softmax": + weights = weights / weights.sum(dim=-1, keepdim=True) + weights = weights * route_scale + return weights, indices + + +def _sparse_from_dense( + *, + weights: torch.Tensor, + indices: torch.Tensor, + num_experts: int, +): + """Pack dense ``(weights, indices)`` into the sparse ``(probs, + routing_map)`` format our routers return.""" + N = weights.shape[0] + probs = torch.zeros(N, num_experts, dtype=weights.dtype, device=weights.device) + probs.scatter_(1, indices, weights) + routing_map = torch.zeros(N, num_experts, dtype=torch.bool, device=weights.device) + routing_map.scatter_(1, indices, True) + return probs, routing_map + + +# --------------------------------------------------------------------------- +# Score function +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) +def test_v4_score_fn_matches_inline_reference(score_function: str) -> None: + torch.manual_seed(123) + logits = torch.randn(8, 16, dtype=torch.float32) + + out = v4_score_fn(logits, score_function=score_function) + if score_function == "softmax": + ref = F.softmax(logits, dim=-1) + elif score_function == "sigmoid": + ref = torch.sigmoid(logits) + else: + ref = F.softplus(logits).sqrt() + assert (out - ref).abs().max().item() <= 1.0e-6 + + +# --------------------------------------------------------------------------- +# Learned router +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("score_function", ["sqrtsoftplus", "sigmoid", "softmax"]) +@pytest.mark.parametrize("enable_expert_bias", [False, True]) +def test_learned_router_matches_hf_reference(score_function: str, enable_expert_bias: bool) -> None: + """G4 learned-router gate: matches HF reference exactly (fp32).""" + torch.manual_seed(2025) + H, E, K = 12, 8, 2 + route_scale = 1.5 + + router = DeepseekV4LearnedRouter( + hidden_size=H, + num_experts=E, + topk=K, + score_function=score_function, + enable_expert_bias=enable_expert_bias, + topk_scaling_factor=route_scale, + ) + if enable_expert_bias: + with torch.no_grad(): + router.expert_bias.copy_(torch.randn(E, dtype=torch.float32)) + + hidden = torch.randn(2, 5, H, dtype=torch.float32) + + probs, routing_map = router(hidden) + + weights_hf, indices_hf = _hf_gate_forward( + hidden=hidden, + weight=router.weight, + bias=router.expert_bias, + score_function=score_function, + topk=K, + route_scale=route_scale, + tid2eid=None, + token_ids=None, + ) + probs_hf, routing_map_hf = _sparse_from_dense(weights=weights_hf, indices=indices_hf, num_experts=E) + + assert probs.shape == probs_hf.shape + assert routing_map.shape == routing_map_hf.shape + assert torch.equal(routing_map, routing_map_hf), "routing_map differs" + max_abs = (probs - probs_hf).abs().max().item() + assert max_abs <= 1.0e-6, f"probs max-abs vs HF reference = {max_abs}" + + +def test_learned_router_back_compat_alias() -> None: + """``V4TopKRouter`` is an alias for the renamed router.""" + assert V4TopKRouter is DeepseekV4LearnedRouter + + +def test_learned_router_grad_flows_to_gate_weight() -> None: + """Probs gradient propagates back into ``weight`` (the gate).""" + torch.manual_seed(11) + H, E, K = 6, 4, 2 + router = DeepseekV4LearnedRouter(hidden_size=H, num_experts=E, topk=K, score_function="sqrtsoftplus") + hidden = torch.randn(3, 4, H, dtype=torch.float32, requires_grad=False) + probs, _ = router(hidden) + loss = probs.sum() + loss.backward() + assert router.weight.grad is not None + assert torch.isfinite(router.weight.grad).all() + assert router.weight.grad.abs().sum().item() > 0.0 + + +def test_learned_router_expert_bias_does_not_contribute_to_probs_grad() -> None: + """``expert_bias`` is selection-only: probs use un-biased scores. + + With ``enable_expert_bias=True`` the bias enters the *selection* + path (top-K) but not the gathered weights. The gradient on the bias + can therefore be exactly zero when the chosen indices stay the same + if you only differentiate the routing weights at the selected + experts. We assert ``expert_bias.grad is None`` (no graph) **after** + a forward + backward, since the bias never enters the autograd + chain in our implementation. + """ + torch.manual_seed(42) + H, E, K = 6, 4, 2 + router = DeepseekV4LearnedRouter( + hidden_size=H, + num_experts=E, + topk=K, + score_function="sqrtsoftplus", + enable_expert_bias=True, + ) + hidden = torch.randn(3, 4, H, dtype=torch.float32) + probs, _ = router(hidden) + probs.sum().backward() + # Bias is detached from the probs graph — grad stays None (or zero). + assert (router.expert_bias.grad is None) or (router.expert_bias.grad.abs().max().item() == 0.0) + + +def test_learned_router_softmax_skips_renormalization() -> None: + """Softmax probs already sum to 1, so the post-topK renorm is skipped. + + We assert by setting ``route_scale != 1`` and confirming the gathered + weights match HF (i.e. they are *not* re-normalized to 1 before the + scale). + """ + torch.manual_seed(0) + H, E, K = 6, 4, 2 + router = DeepseekV4LearnedRouter( + hidden_size=H, + num_experts=E, + topk=K, + score_function="softmax", + topk_scaling_factor=2.0, + ) + hidden = torch.randn(2, 3, H, dtype=torch.float32) + probs, routing_map = router(hidden) + + weights_hf, indices_hf = _hf_gate_forward( + hidden=hidden, + weight=router.weight, + bias=None, + score_function="softmax", + topk=K, + route_scale=2.0, + tid2eid=None, + token_ids=None, + ) + probs_hf, routing_map_hf = _sparse_from_dense(weights=weights_hf, indices=indices_hf, num_experts=E) + assert torch.equal(routing_map, routing_map_hf) + assert (probs - probs_hf).abs().max().item() <= 1.0e-6 + + +# --------------------------------------------------------------------------- +# Hash router +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("score_function", ["sqrtsoftplus", "sigmoid", "softmax"]) +def test_hash_router_matches_hf_reference(score_function: str) -> None: + """G4 hash-router: matches HF reference exactly (fp32).""" + torch.manual_seed(7) + H, E, K, V = 12, 8, 2, 32 + + router = DeepseekV4HashRouter( + hidden_size=H, + num_experts=E, + topk=K, + vocab_size=V, + seed=17, + score_function=score_function, + topk_scaling_factor=1.0, + ) + hidden = torch.randn(2, 5, H, dtype=torch.float32) + token_ids = torch.randint(0, V, (2, 5), dtype=torch.long) + + probs, routing_map = router(hidden, token_ids) + + weights_hf, indices_hf = _hf_gate_forward( + hidden=hidden, + weight=router.weight, + bias=None, + score_function=score_function, + topk=K, + route_scale=1.0, + tid2eid=router.tid2eid, + token_ids=token_ids, + ) + probs_hf, routing_map_hf = _sparse_from_dense(weights=weights_hf, indices=indices_hf, num_experts=E) + assert torch.equal(routing_map, routing_map_hf), "routing_map differs" + max_abs = (probs - probs_hf).abs().max().item() + assert max_abs <= 1.0e-6, f"probs max-abs vs HF reference = {max_abs}" + + +def test_hash_router_back_compat_alias() -> None: + """``HashRouter`` is an alias for the renamed router.""" + assert HashRouter is DeepseekV4HashRouter + + +def test_hash_router_tid2eid_is_frozen_parameter() -> None: + """``tid2eid`` is a Parameter with requires_grad=False. + + Matches the HF reference layout (released checkpoint stores it as + a parameter) so a state-dict round-trip preserves it without + pulling it into the optimizer state. + """ + router = DeepseekV4HashRouter( + hidden_size=8, + num_experts=4, + topk=2, + vocab_size=16, + seed=0, + ) + assert "tid2eid" in dict(router.named_parameters()) + assert router.tid2eid.requires_grad is False + assert router.tid2eid.dtype == torch.int32 + assert tuple(router.tid2eid.shape) == (16, 2) + + +def test_hash_router_state_dict_keys() -> None: + """State-dict exposes ``weight`` and ``tid2eid`` (matches HF gate keys).""" + router = DeepseekV4HashRouter(hidden_size=8, num_experts=4, topk=2, vocab_size=16, seed=0) + keys = set(router.state_dict().keys()) + assert "weight" in keys + assert "tid2eid" in keys + + +def test_hash_router_grad_flows_to_gate_weight() -> None: + """Even with static expert ids, the routing weights' gradient flows + back into the learned gate ``weight``.""" + torch.manual_seed(13) + H, E, K, V = 6, 4, 2, 16 + router = DeepseekV4HashRouter( + hidden_size=H, + num_experts=E, + topk=K, + vocab_size=V, + seed=0, + score_function="sqrtsoftplus", + ) + hidden = torch.randn(3, 4, H, dtype=torch.float32) + token_ids = torch.randint(0, V, (3, 4), dtype=torch.long) + probs, _ = router(hidden, token_ids) + probs.sum().backward() + assert router.weight.grad is not None + assert torch.isfinite(router.weight.grad).all() + assert router.weight.grad.abs().sum().item() > 0.0 + # tid2eid should never accumulate gradient. + assert router.tid2eid.grad is None + + +def test_hash_router_deterministic_table_across_seeds() -> None: + """Same seed -> identical table; different seeds -> different table.""" + a = DeepseekV4HashRouter(hidden_size=4, num_experts=8, topk=2, vocab_size=16, seed=42) + b = DeepseekV4HashRouter(hidden_size=4, num_experts=8, topk=2, vocab_size=16, seed=42) + c = DeepseekV4HashRouter(hidden_size=4, num_experts=8, topk=2, vocab_size=16, seed=43) + assert torch.equal(a.tid2eid, b.tid2eid) + assert not torch.equal(a.tid2eid, c.tid2eid) + + +def test_hash_router_rejects_shape_mismatch() -> None: + """``hidden`` and ``token_ids`` must flatten to the same length.""" + router = DeepseekV4HashRouter(hidden_size=4, num_experts=4, topk=2, vocab_size=8, seed=0) + hidden = torch.randn(2, 3, 4, dtype=torch.float32) # N=6 + token_ids = torch.zeros(1, 5, dtype=torch.long) # N=5 + with pytest.raises(ValueError, match="flatten to the same length"): + router(hidden, token_ids) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_deepep_dispatcher.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_deepep_dispatcher.py new file mode 100644 index 000000000..6d3c6e65c --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_deepep_dispatcher.py @@ -0,0 +1,397 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-3 P23 — Turbo DeepEP dispatcher in V4 specs. + +Today the V4 spec build captures +:class:`megatron.core.transformer.moe.token_dispatcher.MoEFlexTokenDispatcher` +at module-import time, while the Primus turbo patch +(``primus.backends.megatron.patches.turbo.moe_dispatcher_patches``) only +rebinds that module attribute at ``before_train``. The patch fires +*after* V4 spec build, so V4 silently runs the upstream +:class:`MoEFlexTokenDispatcher` even when the user opted into +``use_turbo_deepep=True``. + +P23 fixes this by resolving the dispatcher class **at V4 spec-build time** +through :func:`primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs._pick_v4_dispatcher_cls`, +which checks ``args.enable_primus_turbo``, ``args.use_turbo_deepep``, +``args.tensor_model_parallel_size`` and the import status of the +``primus_turbo`` package. Two side-effects ensure the rest of the +stack agrees: + +* :func:`_maybe_plumb_v4_turbo_deepep_args` (in + ``deepseek_v4_builders.py``) mutates ``args.moe_enable_deepep`` and + ``args.moe_token_dispatcher_type`` BEFORE + ``core_transformer_config_from_args``, so the V4 ``config`` carries + the right ``moe_token_dispatcher_type``. +* :meth:`DeepseekV4MoE._resolve_dispatcher_type_from_spec` recognises + ``PrimusTurboDeepEPTokenDispatcher`` (by class name, no + ``primus_turbo`` import required) and returns ``"flex"``, so the + per-layer log line ``"dispatcher active via …"`` reports the + correct type. + +Test gates exercised here: + +* **G20a — gating predicate** + :func:`is_v4_turbo_deepep_active` matches the four conditions used + by the upstream patch. +* **G20b — args plumbing** + :func:`_maybe_plumb_v4_turbo_deepep_args` mutates only when all + gates pass; respects an explicit ``"allgather"`` opt-in. +* **G20c — class resolution** + :func:`_pick_v4_dispatcher_cls` returns the right + ``(cls, type_name)`` tuple for each + ``(config.moe_token_dispatcher_type, args)`` combination. +* **G20d — V4 MoE resolver recognises Turbo class** + :meth:`DeepseekV4MoE._resolve_dispatcher_type_from_spec` returns + ``"flex"`` for a class named + ``PrimusTurboDeepEPTokenDispatcher`` (mocked when ``primus_turbo`` + is not installed). +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from megatron.core.transformer.moe.token_dispatcher import ( + MoEAllGatherTokenDispatcher, + MoEAlltoAllTokenDispatcher, + MoEFlexTokenDispatcher, +) +from megatron.core.transformer.spec_utils import ModuleSpec + +from primus.backends.megatron.core.models.deepseek_v4 import ( + deepseek_v4_builders, + deepseek_v4_layer_specs, +) +from primus.backends.megatron.core.models.deepseek_v4.deepseek_v4_layer_specs import ( + _pick_v4_dispatcher_cls, + is_v4_turbo_deepep_active, +) +from primus.backends.megatron.core.transformer.moe.v4_moe import DeepseekV4MoE + +# --------------------------------------------------------------------------- +# Test fixtures +# --------------------------------------------------------------------------- + + +def _make_args( + *, + enable_primus_turbo: bool = False, + use_turbo_deepep: bool = False, + tensor_model_parallel_size: int = 1, + moe_token_dispatcher_type: str = "alltoall", + moe_enable_deepep: bool = False, +): + """Minimal ``args`` namespace mirroring Megatron's runtime args.""" + return SimpleNamespace( + enable_primus_turbo=enable_primus_turbo, + use_turbo_deepep=use_turbo_deepep, + tensor_model_parallel_size=tensor_model_parallel_size, + moe_token_dispatcher_type=moe_token_dispatcher_type, + moe_enable_deepep=moe_enable_deepep, + ) + + +def _make_cfg(*, moe_token_dispatcher_type: str = "alltoall"): + """Minimal ``config`` namespace with the dispatcher-type field only.""" + return SimpleNamespace(moe_token_dispatcher_type=moe_token_dispatcher_type) + + +@pytest.fixture +def primus_turbo_available(monkeypatch): + """Pretend the ``primus_turbo`` package is importable.""" + real_find_spec = deepseek_v4_layer_specs.importlib.util.find_spec + + def _fake(name): + if name == "primus_turbo": + return MagicMock() + return real_find_spec(name) + + monkeypatch.setattr(deepseek_v4_layer_specs.importlib.util, "find_spec", _fake) + + +@pytest.fixture +def primus_turbo_unavailable(monkeypatch): + """Pretend the ``primus_turbo`` package is NOT importable.""" + real_find_spec = deepseek_v4_layer_specs.importlib.util.find_spec + + def _fake(name): + if name == "primus_turbo": + return None + return real_find_spec(name) + + monkeypatch.setattr(deepseek_v4_layer_specs.importlib.util, "find_spec", _fake) + + +# ``type(..., (), {})`` produces a class whose ``__name__`` is the +# string we pass — using a plain ``class`` block would set ``__name__`` +# to the local symbol (e.g. ``_FakeTurboDispatcher``) regardless of +# any class-level ``__name__`` assignment. Only the class name +# matters for ``_resolve_dispatcher_type_from_spec``; construction is +# never exercised in these unit tests. +_FakeTurboDispatcher = type("PrimusTurboDeepEPTokenDispatcher", (), {}) + + +@pytest.fixture +def fake_turbo_class(monkeypatch): + """Inject a fake ``PrimusTurboDeepEPTokenDispatcher`` import target.""" + monkeypatch.setattr( + deepseek_v4_layer_specs, + "_import_primus_turbo_deepep_dispatcher_cls", + lambda: _FakeTurboDispatcher, + ) + return _FakeTurboDispatcher + + +# --------------------------------------------------------------------------- +# G20a — gating predicate (is_v4_turbo_deepep_active) +# --------------------------------------------------------------------------- + + +class TestGatingPredicate: + """Mirror of the conditions enforced by ``moe_dispatcher_patches``.""" + + def test_all_gates_open(self, primus_turbo_available): + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=True, + tensor_model_parallel_size=1, + ) + assert is_v4_turbo_deepep_active(args) is True + + def test_primus_turbo_missing(self, primus_turbo_unavailable): + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=True, + tensor_model_parallel_size=1, + ) + assert is_v4_turbo_deepep_active(args) is False + + def test_enable_primus_turbo_off(self, primus_turbo_available): + args = _make_args( + enable_primus_turbo=False, + use_turbo_deepep=True, + tensor_model_parallel_size=1, + ) + assert is_v4_turbo_deepep_active(args) is False + + def test_use_turbo_deepep_off(self, primus_turbo_available): + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=False, + tensor_model_parallel_size=1, + ) + assert is_v4_turbo_deepep_active(args) is False + + def test_tp_gt_1_blocks(self, primus_turbo_available): + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=True, + tensor_model_parallel_size=2, + ) + assert is_v4_turbo_deepep_active(args) is False + + +# --------------------------------------------------------------------------- +# G20b — args plumbing (_maybe_plumb_v4_turbo_deepep_args) +# --------------------------------------------------------------------------- + + +class TestArgsPlumbing: + """The plumbing helper sets ``moe_enable_deepep`` + dispatcher type.""" + + def test_plumb_when_active(self, primus_turbo_available): + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=True, + moe_token_dispatcher_type="alltoall", + moe_enable_deepep=False, + ) + deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args(args) + assert args.moe_enable_deepep is True + assert args.moe_token_dispatcher_type == "flex" + + def test_no_plumb_when_inactive(self, primus_turbo_available): + args = _make_args( + enable_primus_turbo=False, + use_turbo_deepep=True, + moe_token_dispatcher_type="alltoall", + moe_enable_deepep=False, + ) + deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args(args) + assert args.moe_enable_deepep is False + assert args.moe_token_dispatcher_type == "alltoall" + + def test_no_plumb_when_package_missing(self, primus_turbo_unavailable): + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=True, + moe_token_dispatcher_type="alltoall", + ) + deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args(args) + assert args.moe_token_dispatcher_type == "alltoall" + + def test_explicit_allgather_preserved(self, primus_turbo_available): + """User opted into ``allgather`` — never silently overridden.""" + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=True, + moe_token_dispatcher_type="allgather", + ) + deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args(args) + # ``moe_enable_deepep`` may flip (Turbo wants it for any deepep + # path), but the user's dispatcher choice is preserved. + assert args.moe_token_dispatcher_type == "allgather" + + def test_idempotent_when_already_flex(self, primus_turbo_available): + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=True, + moe_token_dispatcher_type="flex", + moe_enable_deepep=True, + ) + deepseek_v4_builders._maybe_plumb_v4_turbo_deepep_args(args) + assert args.moe_token_dispatcher_type == "flex" + assert args.moe_enable_deepep is True + + +# --------------------------------------------------------------------------- +# G20c — class resolution (_pick_v4_dispatcher_cls) +# --------------------------------------------------------------------------- + + +class TestPickDispatcherCls: + """Exhaustive class-resolution table.""" + + def test_alltoall_default(self): + cfg = _make_cfg(moe_token_dispatcher_type="alltoall") + cls, type_name = _pick_v4_dispatcher_cls(cfg, args=_make_args()) + assert cls is MoEAlltoAllTokenDispatcher + assert type_name == "alltoall" + + def test_allgather_explicit(self): + cfg = _make_cfg(moe_token_dispatcher_type="allgather") + cls, type_name = _pick_v4_dispatcher_cls(cfg, args=_make_args()) + assert cls is MoEAllGatherTokenDispatcher + assert type_name == "allgather" + + def test_flex_without_turbo(self): + cfg = _make_cfg(moe_token_dispatcher_type="flex") + cls, type_name = _pick_v4_dispatcher_cls(cfg, args=_make_args()) + assert cls is MoEFlexTokenDispatcher + assert type_name == "flex" + + def test_flex_with_turbo_active(self, primus_turbo_available, fake_turbo_class): + cfg = _make_cfg(moe_token_dispatcher_type="flex") + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=True, + tensor_model_parallel_size=1, + ) + cls, type_name = _pick_v4_dispatcher_cls(cfg, args=args) + assert cls is fake_turbo_class + assert type_name == "flex" + + def test_flex_with_turbo_active_but_class_missing(self, primus_turbo_available, monkeypatch): + """Gracefully fall back to the upstream class with a warning.""" + monkeypatch.setattr( + deepseek_v4_layer_specs, + "_import_primus_turbo_deepep_dispatcher_cls", + lambda: None, + ) + cfg = _make_cfg(moe_token_dispatcher_type="flex") + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=True, + tensor_model_parallel_size=1, + ) + cls, type_name = _pick_v4_dispatcher_cls(cfg, args=args) + assert cls is MoEFlexTokenDispatcher + assert type_name == "flex" + + def test_flex_with_turbo_inactive_tp_gt_1(self, primus_turbo_available, fake_turbo_class): + """TP > 1 keeps the upstream Flex dispatcher.""" + cfg = _make_cfg(moe_token_dispatcher_type="flex") + args = _make_args( + enable_primus_turbo=True, + use_turbo_deepep=True, + tensor_model_parallel_size=2, + ) + cls, type_name = _pick_v4_dispatcher_cls(cfg, args=args) + assert cls is MoEFlexTokenDispatcher + assert type_name == "flex" + + def test_unknown_type_falls_back_to_alltoall(self, caplog): + cfg = _make_cfg(moe_token_dispatcher_type="random_string") + with caplog.at_level("WARNING", logger=deepseek_v4_layer_specs.__name__): + cls, type_name = _pick_v4_dispatcher_cls(cfg, args=_make_args()) + assert cls is MoEAlltoAllTokenDispatcher + assert type_name == "alltoall" + assert any("unsupported moe_token_dispatcher_type" in rec.message for rec in caplog.records) + + def test_args_none_falls_back_when_megatron_not_initialised( + self, primus_turbo_available, fake_turbo_class, monkeypatch + ): + """``args=None`` + no Megatron args available → non-turbo branch.""" + + def _raises(): + raise RuntimeError("Megatron not initialised in unit test") + + # Patch the lazy import target inside the helper. ``megatron.training`` + # may not have ``get_args`` bound yet depending on suite import order + # (it is populated lazily), so ``raising=False`` forces the "raises" + # behaviour either way; monkeypatch restores/deletes it on teardown. + import megatron.training as _megatron_training # noqa: WPS433 + + monkeypatch.setattr(_megatron_training, "get_args", _raises, raising=False) + cfg = _make_cfg(moe_token_dispatcher_type="flex") + cls, type_name = _pick_v4_dispatcher_cls(cfg, args=None) + assert cls is MoEFlexTokenDispatcher + assert type_name == "flex" + + +# --------------------------------------------------------------------------- +# G20d — V4 MoE resolver recognises Turbo class +# --------------------------------------------------------------------------- + + +class TestV4MoEResolver: + """``DeepseekV4MoE._resolve_dispatcher_type_from_spec`` must label + :class:`PrimusTurboDeepEPTokenDispatcher` as ``"flex"``.""" + + def test_turbo_class_resolves_to_flex(self): + spec = ModuleSpec(module=_FakeTurboDispatcher) + assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "flex" + + def test_turbo_class_bare_resolves_to_flex(self): + # The resolver also accepts a bare class (some call sites pass + # the type directly, not a ModuleSpec). + assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(_FakeTurboDispatcher) == "flex" + + def test_alltoall_unchanged(self): + spec = ModuleSpec(module=MoEAlltoAllTokenDispatcher) + assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "alltoall" + + def test_flex_unchanged(self): + spec = ModuleSpec(module=MoEFlexTokenDispatcher) + assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "flex" + + def test_allgather_unchanged(self): + spec = ModuleSpec(module=MoEAllGatherTokenDispatcher) + assert DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) == "allgather" + + def test_unknown_class_falls_back_to_alltoall(self, caplog): + class _Unknown: + __name__ = "_Unknown" + + spec = ModuleSpec(module=_Unknown) + with caplog.at_level("WARNING"): + result = DeepseekV4MoE._resolve_dispatcher_type_from_spec(spec) + assert result == "alltoall" + assert any("unsupported dispatcher module" in rec.message for rec in caplog.records) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py new file mode 100644 index 000000000..0c972cdd8 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_turbo_flydsl_attention.py @@ -0,0 +1,293 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""``turbo`` DeepSeek-V4 attention fwd+bwd correctness, in the **V4 form** (gfx950), +against the fp32 eager references. + +The ``turbo`` backend (:mod:`..._turbo_flydsl`) is the Primus-Turbo native-FlyDSL +sparse-MLA v2 kernel-pair reached through the **turbo API** +(``primus_turbo.flydsl.attention.kernels.sparse_mla_v2``), bound to the V4 +autograd adapters (:mod:`v4_csa_attention_turbo_flydsl`). It is the backend +selected by ``use_v4_attention_backend`` / ``use_v4_csa_attention_backend = +"turbo"``. This validates the production V4 invocation paths for all three layer +kinds: + +* ``compress_ratio == 0`` (dense / SWA) -> :func:`v4_attention_turbo` +* ``compress_ratio == 128`` (HCA) -> :func:`v4_attention_turbo` +* ``compress_ratio == 4`` (CSA) -> :func:`v4_csa_attention_turbo` + +Full fwd (O) and torch-autograd bwd (dQ, dlatent, dpool, dsink) of the bf16 turbo +adapter vs the fp32 eager reference (same tolerances the gluon_v2 backend UT uses; +turbo shares the identical fused single-latent sparse-MLA-with-sink math). GPU-only; +skipped off gfx950 / when primus_turbo's flydsl attention or ``flydsl`` is absent. + +FlyDSL fixes the head-block at 64, so ``num_heads`` must be a multiple of 64 here +(the kernel asserts ``num_heads % 32 == 0``; H=64/128 are the production sizes). +``S`` is a multiple of 128 so the cr=128 closed-form pool (``pool_cr = S/P = 128``) +is exact. +""" + +from __future__ import annotations + +import math + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("turbo sparse-MLA kernels require CUDA / HIP", allow_module_level=True) + +pytest.importorskip("flydsl", reason="flydsl pip package not installed") +# The turbo API: primus_turbo must carry the flydsl sparse-MLA attention submodule. +pytest.importorskip( + "primus_turbo.flydsl.attention.kernels.sparse_mla_v2", + reason="installed primus_turbo has no flydsl sparse-MLA attention (turbo backend)", +) + +_ARCH = torch.cuda.get_device_properties(0).gcnArchName +if "gfx950" not in _ARCH: + pytest.skip(f"turbo native-FlyDSL sparse-MLA targets gfx950; got {_ARCH}", allow_module_level=True) + +from primus.backends.megatron.core.transformer.sliding_window_kv import ( # noqa: E402 + sliding_window_causal_mask, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels._eager.reference import ( # noqa: E402 + eager_v4_attention, + eager_v4_csa_attention, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels.v4_csa_attention_turbo_flydsl import ( # noqa: E402 + v4_attention_turbo, + v4_csa_attention_turbo, +) + +D = 512 # V4 head_dim (RoPE baked in-place) +_SWA = 128 + + +def _stats(a, b, *, sig=1e-2): + a = a.float() + b = b.float() + d = (a - b).abs() + m = b.abs() > sig + rel = (d[m] / b.abs()[m]) if m.any() else d.new_zeros(0) + med_rel = rel.median().item() if rel.numel() else 0.0 + av, bv = a.flatten(), b.flatten() + cos_err = 1.0 - (av @ bv / (av.norm() * bv.norm() + 1e-30)).item() + return d.max().item(), med_rel, cos_err + + +def _check(name, a, b, *, abs_tol, sig=1e-2, med=None, cos=None): + # Guard against a broken kernel returning NaN/Inf (e.g. the cr=4 fast_path + # overflow / bwd race the extraction found) — surface it as a clear failure. + assert torch.isfinite( + a.float() + ).all(), f"{name} has NaN/Inf ({int((~torch.isfinite(a.float())).sum())} bad)" + max_abs, med_rel, cos_err = _stats(a, b, sig=sig) + print(f" {name:8s} max_abs={max_abs:.3e} median_rel={med_rel:.3e} cos_err={cos_err:.3e}", flush=True) + assert max_abs < abs_tol, f"{name} max_abs {max_abs:.3e} >= {abs_tol}" + if med is not None: + assert med_rel < med, f"{name} median_rel {med_rel:.3e} >= {med}" + if cos is not None: + assert cos_err < cos, f"{name} cos_err {cos_err:.3e} >= {cos}" + + +def _leaf(x, *, fp32): + y = x.float() if fp32 else x.clone() + return y.detach().requires_grad_(True) + + +# H must be a multiple of 64 (FlyDSL head-block); S a multiple of 128 (cr=128 pool_cr). +# B is fixed to 1: the flydsl dense/HCA "banded" path uses a closed-form SWA window +# [i-127..i] over the flat token axis, which crosses batch boundaries for B>1 (it assumes +# a single contiguous sequence). Production / bench_v4_attention.py run B(mbs)=1; multi-batch +# dense/HCA is a known flydsl-banded limitation (CSA/cr=4 is fine for B>1 — it uses the +# per-batch-offset topk). Vary H (64/128), S (256/512), and the sink instead. +_CASES = [ + (1, 64, 256, None), + (1, 64, 512, 1.0), + (1, 128, 512, -1.0), +] + + +@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v)) +def test_v4_turbo_dense_matches_eager(B, H, S, sink_init): + """cr=0 dense/SWA: turbo fwd+bwd vs eager_v4_attention.""" + g = torch.Generator(device="cuda").manual_seed(0) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None + + qg, latg = _leaf(q, fp32=False), _leaf(lat, fp32=False) + sg = _leaf(sink, fp32=False) if sink is not None else None + kg = latg.unsqueeze(1).expand(B, H, S, D) + og = v4_attention_turbo( + qg, kg, kg, sink=sg, swa_window=_SWA, additive_mask=None, attn_dropout=0.0, training=True, scale=scale + ) + og.backward(do) + + qf, latf = _leaf(q, fp32=True), _leaf(lat, fp32=True) + sf = _leaf(sink, fp32=True) if sink is not None else None + kf = latf.unsqueeze(1).expand(B, H, S, D) + of = eager_v4_attention( + qf, + kf, + kf, + sink=sf, + swa_window=_SWA, + additive_mask=None, + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[turbo V4 dense] B={B} H={H} S={S} sink={sink_init}", flush=True) + assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16 + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + # abs_tol is generous for the shared-latent grad (summed over H heads -> large-magnitude + # elements -> bf16 abs outliers); median_rel + cos_err are the real correctness guards. + _check("dlatent", latg.grad, latf.grad, abs_tol=3e-1, sig=1e-3, med=2e-2, cos=1e-3) + if sink is not None: + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) + + +@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v)) +def test_v4_turbo_hca_matches_eager(B, H, S, sink_init): + """cr=128 HCA: turbo (local++pool) fwd+bwd vs eager_v4_attention.""" + cr = 128 + P = max(S // cr, 1) + g = torch.Generator(device="cuda").manual_seed(2) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None + + ti = torch.arange(S, device="cuda").view(S, 1) + ps = torch.arange(P, device="cuda").view(1, P) + pool_mask = torch.where( + ((ps + 1) * cr - 1) <= ti, torch.zeros((), device="cuda"), torch.tensor(float("-inf"), device="cuda") + ).to(torch.bfloat16) + + qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False) + sg = _leaf(sink, fp32=False) if sink is not None else None + kg = torch.cat([latg.unsqueeze(1).expand(B, H, S, D), poolg.unsqueeze(1).expand(B, H, P, D)], dim=2) + og = v4_attention_turbo( + qg, + kg, + kg, + sink=sg, + swa_window=_SWA, + additive_mask=pool_mask, + attn_dropout=0.0, + training=True, + scale=scale, + hca_local_seqlen=S, + ) + og.backward(do) + + qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True) + sf = _leaf(sink, fp32=True) if sink is not None else None + kf = torch.cat([latf.unsqueeze(1).expand(B, H, S, D), poolf.unsqueeze(1).expand(B, H, P, D)], dim=2) + local_mask = sliding_window_causal_mask(S, _SWA, device="cuda", dtype=torch.float32) + full_mask = torch.cat([local_mask, pool_mask.float()], dim=1) + of = eager_v4_attention( + qf, + kf, + kf, + sink=sf, + swa_window=0, + additive_mask=full_mask, + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[turbo V4 HCA] B={B} H={H} S={S} P={P} sink={sink_init}", flush=True) + assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16 + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + # abs_tol is generous for the shared-latent grad (summed over H heads -> large-magnitude + # elements -> bf16 abs outliers); median_rel + cos_err are the real correctness guards. + _check("dlatent", latg.grad, latf.grad, abs_tol=3e-1, sig=1e-3, med=2e-2, cos=1e-3) + _check("dpool", poolg.grad, poolf.grad, abs_tol=3e-1, sig=1e-3, med=2e-2, cos=1e-3) + if sink is not None: + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) + + +@pytest.mark.parametrize("B,H,S,sink_init", _CASES, ids=lambda v: str(v)) +def test_v4_turbo_csa_matches_eager(B, H, S, sink_init): + """cr=4 CSA: turbo fwd+bwd vs eager_v4_csa_attention.""" + P = max(S // 4, 1) + K = min(128, P) + g = torch.Generator(device="cuda").manual_seed(3) + scale = 1.0 / math.sqrt(D) + q = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + lat = torch.randn(B, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + pool = torch.randn(B, P, D, generator=g, device="cuda", dtype=torch.bfloat16) + do = torch.randn(B, H, S, D, generator=g, device="cuda", dtype=torch.bfloat16) + sink = torch.full((H,), sink_init, dtype=torch.float32, device="cuda") if sink_init is not None else None + + topk_idxs = torch.randint(0, P, (B, S, K), generator=g, device="cuda", dtype=torch.int32) + drop = torch.rand(B, S, K, generator=g, device="cuda") < 0.125 + topk_idxs = torch.where(drop, torch.full_like(topk_idxs, -1), topk_idxs) + idx = topk_idxs.clamp(0, P - 1).long() + bidx = torch.arange(B, device="cuda").view(B, 1, 1) + sparse_mask_inf = torch.where( + topk_idxs < 0, torch.tensor(float("-inf"), device="cuda"), torch.zeros((), device="cuda") + ) + + qg, latg, poolg = _leaf(q, fp32=False), _leaf(lat, fp32=False), _leaf(pool, fp32=False) + sg = _leaf(sink, fp32=False) if sink is not None else None + klg = latg.unsqueeze(1).expand(B, H, S, D) + og = v4_csa_attention_turbo( + qg, + klg, + klg, + poolg, + topk_idxs=topk_idxs, + sink=sg, + swa_window=_SWA, + attn_dropout=0.0, + training=True, + scale=scale, + ) + og.backward(do) + + qf, latf, poolf = _leaf(q, fp32=True), _leaf(lat, fp32=True), _leaf(pool, fp32=True) + sf = _leaf(sink, fp32=True) if sink is not None else None + klf = latf.unsqueeze(1).expand(B, H, S, D) + gathered = poolf[bidx, idx] + of = eager_v4_csa_attention( + qf, + klf, + klf, + gathered, + sink=sf, + swa_window=_SWA, + sparse_mask=sparse_mask_inf.float(), + attn_dropout=0.0, + training=False, + scale=scale, + ) + of.backward(do.float()) + + print(f"\n[turbo V4 CSA] B={B} H={H} S={S} P={P} K={K} sink={sink_init}", flush=True) + assert og.shape == (B, H, S, D) and og.dtype == torch.bfloat16 + _check("O", og, of, abs_tol=3e-2, med=2e-2, cos=1e-3) + _check("dQ", qg.grad, qf.grad, abs_tol=5e-2, sig=1e-3, med=2e-2, cos=1e-3) + # abs_tol is generous for the shared-latent grad (summed over H heads -> large-magnitude + # elements -> bf16 abs outliers); median_rel + cos_err are the real correctness guards. + _check("dlatent", latg.grad, latf.grad, abs_tol=3e-1, sig=1e-3, med=2e-2, cos=1e-3) + _check("dpool", poolg.grad, poolf.grad, abs_tol=3e-1, sig=1e-3, med=2e-2, cos=1e-3) + if sink is not None: + _check("dSink", sg.grad, sf.grad, abs_tol=5e-1, sig=1e-3, med=5e-2, cos=1e-2) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_bwd.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_bwd.py new file mode 100644 index 000000000..302ed8f38 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_bwd.py @@ -0,0 +1,439 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-4 P25 G24 — `v4_attention_v1` Triton BWD equivalence to autograd-on-eager. + +Asserts that gradients (``dq``, ``dk``, ``dv``, ``dsink``) returned by +:func:`v4_attention_v1`'s autograd Function match the gradients computed +by autograd-on-:func:`eager_v4_attention` within the plan-4 tolerance +budget across the same shape envelope as G23 (V4-Flash + V4-Pro, +``compress_ratio ∈ {0, 128}``, fp32 + bf16, sink_on / sink_off, MQA / +MHA layouts). + +The sink gradient is asserted **per-head**: ``dsink`` is shape ``[H]`` +and each head's gradient must match independently. +""" + +from __future__ import annotations + +import math +from typing import Optional + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("v4_attention_v1 Triton kernel requires CUDA / HIP", allow_module_level=True) + +pytest.importorskip("triton", reason="Triton not installed") + +from primus.backends.megatron.core.transformer.sliding_window_kv import ( # noqa: E402 + sliding_window_causal_mask, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels import ( # noqa: E402 + eager_v4_attention, + v4_attention_v1, +) + +# --------------------------------------------------------------------------- +# Shape envelope (same as G23) +# --------------------------------------------------------------------------- + + +# See ``test_v4_p25_v4_attention_fwd._BASE_SHAPES`` for the fast vs +# release tier rationale (G28 plan-4 release-tier shape gate). +_BASE_SHAPES = [ + ("v4_flash_small", 1, 8, 64, 64, 32), + ("v4_pro_small", 1, 4, 64, 64, 32), + pytest.param( + "v4_flash_release", + 1, + 64, + 1024, + 512, + 128, + marks=pytest.mark.slow, + ), + pytest.param( + "v4_pro_release", + 1, + 128, + 512, + 512, + 128, + marks=pytest.mark.slow, + ), +] +_DTYPES = [torch.float32, torch.bfloat16] +_SINK_MODES = [True, False] +_KV_LAYOUTS = ["mqa", "mha"] + + +def _is_release_tier(variant: str) -> bool: + """Release-tier shapes are tagged by name (``*_release``).""" + return variant.endswith("_release") + + +def _bwd_tol(dtype: torch.dtype, *, release: bool = False) -> dict: + """Plan-4 BWD tolerance budget. + + Release-tier ``head_dim=512`` causes: + * ~sqrt(8) ≈ 2.8x more matmul-accumulation noise than the + ``head_dim=64`` fast tier; + * non-deterministic ``tl.atomic_add`` contributions to ``dk / dv`` + (sliding-window cells receive ~SWA contributions, summed in a + non-deterministic thread order) which adds another ~``sqrt(SWA)`` + bf16 jitter on top of matmul noise. + + Empirically the worst-case outlier at V4-Flash / V4-Pro release + sits around ``0.15-0.18`` for bf16 ``dk``; we set ``atol=2e-1`` to + absorb the long tail while keeping fast-tier shapes tight. + """ + if dtype == torch.float32: + return {"atol": 1e-4, "rtol": 1e-4} + if dtype == torch.bfloat16: + return {"atol": 2e-1, "rtol": 2e-1} if release else {"atol": 5e-2, "rtol": 5e-2} + raise ValueError(f"unsupported dtype {dtype!r}") + + +def _sink_tol(dtype: torch.dtype, *, release: bool = False) -> dict: + """Per-head sink gradient tolerance. + + ``dsink[h] = sum_{b, t} dprobs_at_sink_column[b, h, t]`` — the + reduction is over ``B * Sq`` softmax-derivative terms, so cumulative + bf16 rounding scales linearly with ``Sq``. Release-tier ``Sq=1024`` + is 16x the fast-tier ``Sq=64``; loosen the bf16 budget accordingly. + """ + if dtype == torch.float32: + return {"atol": 1e-4, "rtol": 1e-4} + if dtype == torch.bfloat16: + return {"atol": 5e-2, "rtol": 5e-2} if release else {"atol": 5e-3, "rtol": 5e-3} + raise ValueError(f"unsupported dtype {dtype!r}") + + +def _build_inputs( + *, + B: int, + H: int, + S: int, + D: int, + swa_window: int, + sink_on: bool, + dtype: torch.dtype, + kv_layout: str, + use_hca: bool = False, + seed: int = 4321, +): + """Build (q, k, v, sink, additive_mask, swa_window) with leaves on requires_grad. + + The dense path uses ``swa_window > 0, additive_mask=None``. The HCA + path concatenates a tiny pool to the keys / values and supplies a + pre-built ``[Sq, Sq+P]`` joint additive mask. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + device = "cuda" + K_H = 1 if kv_layout == "mqa" else H + + q = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype, requires_grad=True) + if use_hca: + P = 4 + k_local = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype) + v_local = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype) + pool_k = torch.randn(B, K_H, P, D, generator=g, device=device, dtype=dtype) + pool_v = torch.randn(B, K_H, P, D, generator=g, device=device, dtype=dtype) + k = torch.cat([k_local, pool_k], dim=2).requires_grad_(True) + v = torch.cat([v_local, pool_v], dim=2).requires_grad_(True) + # Joint additive mask: local SWA-causal + pool causal-on-stride + local_mask = sliding_window_causal_mask(S, swa_window, device=device, dtype=dtype) + ratio = 4 + t = torch.arange(S, device=device).unsqueeze(1) + s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * ratio - 1 + pool_mask = torch.where(s_end <= t, 0.0, float("-inf")).to(dtype) + full_mask = torch.cat([local_mask, pool_mask], dim=-1) + kernel_swa_window = swa_window + eager_swa_window = 0 + eager_mask: Optional[torch.Tensor] = full_mask + kernel_mask: Optional[torch.Tensor] = pool_mask + hca_local_seqlen = S + else: + k = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype, requires_grad=True) + v = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype, requires_grad=True) + # Eager and kernel both use the in-kernel SWA mask (eager via + # pre-built additive_mask, kernel via swa_window). + eager_mask = sliding_window_causal_mask(S, swa_window, device=device, dtype=dtype) + kernel_mask = None + eager_swa_window = 0 + kernel_swa_window = swa_window + hca_local_seqlen = 0 + + sink = ( + torch.randn(H, generator=g, device=device, dtype=torch.float32, requires_grad=True) * 0.1 + if sink_on + else None + ) + return dict( + q=q, + k=k, + v=v, + sink=sink, + eager_mask=eager_mask, + eager_swa_window=eager_swa_window, + kernel_mask=kernel_mask, + kernel_swa_window=kernel_swa_window, + hca_local_seqlen=hca_local_seqlen, + ) + + +def _grads_from(model_out: torch.Tensor, *leaves: torch.Tensor) -> list[Optional[torch.Tensor]]: + """Run ``model_out.sum().backward()`` and pull ``leaf.grad`` into a list. + + Cleans up ``leaf.grad`` on each leaf so the caller can re-use the + leaves for a second autograd pass. + """ + real_leaves = [lf for lf in leaves if lf is not None] + grads_out = torch.ones_like(model_out) + grads = torch.autograd.grad( + outputs=model_out, + inputs=real_leaves, + grad_outputs=grads_out, + retain_graph=False, + create_graph=False, + allow_unused=False, + ) + # Re-pad with None for missing leaves + out: list[Optional[torch.Tensor]] = [] + j = 0 + for lf in leaves: + if lf is None: + out.append(None) + else: + out.append(grads[j]) + j += 1 + return out + + +# --------------------------------------------------------------------------- +# G24 — dense (compress_ratio == 0) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("variant,B,H,S,D,swa_window", _BASE_SHAPES) +@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) +@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"]) +@pytest.mark.parametrize("kv_layout", _KV_LAYOUTS) +def test_g24_dense_bwd_matches_eager( + variant: str, + B: int, + H: int, + S: int, + D: int, + swa_window: int, + dtype: torch.dtype, + sink_on: bool, + kv_layout: str, +): + """Dense path: kernel BWD matches autograd-on-eager.""" + # Build two independent leaf sets so the two BWD passes do not + # interfere with each other's ``grad`` slots. + ref_inp = _build_inputs( + B=B, + H=H, + S=S, + D=D, + swa_window=swa_window, + sink_on=sink_on, + dtype=dtype, + kv_layout=kv_layout, + use_hca=False, + seed=4321, + ) + cand_inp = _build_inputs( + B=B, + H=H, + S=S, + D=D, + swa_window=swa_window, + sink_on=sink_on, + dtype=dtype, + kv_layout=kv_layout, + use_hca=False, + seed=4321, + ) + + scale = 1.0 / math.sqrt(D) + + out_ref = eager_v4_attention( + ref_inp["q"], + ref_inp["k"], + ref_inp["v"], + sink=ref_inp["sink"], + swa_window=ref_inp["eager_swa_window"], + additive_mask=ref_inp["eager_mask"], + attn_dropout=0.0, + training=False, + scale=scale, + ) + dq_ref, dk_ref, dv_ref, dsink_ref = _grads_from( + out_ref, ref_inp["q"], ref_inp["k"], ref_inp["v"], ref_inp["sink"] + ) + + out_cand = v4_attention_v1( + cand_inp["q"], + cand_inp["k"], + cand_inp["v"], + sink=cand_inp["sink"], + swa_window=cand_inp["kernel_swa_window"], + additive_mask=cand_inp["kernel_mask"], + attn_dropout=0.0, + training=False, + scale=scale, + hca_local_seqlen=cand_inp["hca_local_seqlen"], + ) + dq_cand, dk_cand, dv_cand, dsink_cand = _grads_from( + out_cand, cand_inp["q"], cand_inp["k"], cand_inp["v"], cand_inp["sink"] + ) + + release = _is_release_tier(variant) + tol = _bwd_tol(dtype, release=release) + torch.testing.assert_close(dq_cand, dq_ref, **tol) + torch.testing.assert_close(dk_cand, dk_ref, **tol) + torch.testing.assert_close(dv_cand, dv_ref, **tol) + if sink_on: + # Per-head sink gradient: shape [H] + assert dsink_ref.shape == dsink_cand.shape == (H,) + torch.testing.assert_close(dsink_cand, dsink_ref, **_sink_tol(dtype, release=release)) + else: + assert dsink_ref is None and dsink_cand is None + + +# --------------------------------------------------------------------------- +# G24 — HCA (compress_ratio == 128) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("variant,B,H,S,D,swa_window", _BASE_SHAPES) +@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) +@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"]) +@pytest.mark.parametrize("kv_layout", _KV_LAYOUTS) +def test_g24_hca_style_bwd_matches_eager( + variant: str, + B: int, + H: int, + S: int, + D: int, + swa_window: int, + dtype: torch.dtype, + sink_on: bool, + kv_layout: str, +): + """HCA path: caller-supplied additive mask, kernel BWD matches autograd-on-eager.""" + ref_inp = _build_inputs( + B=B, + H=H, + S=S, + D=D, + swa_window=swa_window, + sink_on=sink_on, + dtype=dtype, + kv_layout=kv_layout, + use_hca=True, + seed=4321, + ) + cand_inp = _build_inputs( + B=B, + H=H, + S=S, + D=D, + swa_window=swa_window, + sink_on=sink_on, + dtype=dtype, + kv_layout=kv_layout, + use_hca=True, + seed=4321, + ) + + scale = 1.0 / math.sqrt(D) + + out_ref = eager_v4_attention( + ref_inp["q"], + ref_inp["k"], + ref_inp["v"], + sink=ref_inp["sink"], + swa_window=0, + additive_mask=ref_inp["eager_mask"], + attn_dropout=0.0, + training=False, + scale=scale, + ) + dq_ref, dk_ref, dv_ref, dsink_ref = _grads_from( + out_ref, ref_inp["q"], ref_inp["k"], ref_inp["v"], ref_inp["sink"] + ) + + out_cand = v4_attention_v1( + cand_inp["q"], + cand_inp["k"], + cand_inp["v"], + sink=cand_inp["sink"], + swa_window=cand_inp["kernel_swa_window"], + additive_mask=cand_inp["kernel_mask"], + attn_dropout=0.0, + training=False, + scale=scale, + hca_local_seqlen=cand_inp["hca_local_seqlen"], + ) + dq_cand, dk_cand, dv_cand, dsink_cand = _grads_from( + out_cand, cand_inp["q"], cand_inp["k"], cand_inp["v"], cand_inp["sink"] + ) + + release = _is_release_tier(variant) + tol = _bwd_tol(dtype, release=release) + torch.testing.assert_close(dq_cand, dq_ref, **tol) + torch.testing.assert_close(dk_cand, dk_ref, **tol) + torch.testing.assert_close(dv_cand, dv_ref, **tol) + if sink_on: + assert dsink_ref.shape == dsink_cand.shape == (H,) + torch.testing.assert_close(dsink_cand, dsink_ref, **_sink_tol(dtype, release=release)) + else: + assert dsink_ref is None and dsink_cand is None + + +# --------------------------------------------------------------------------- +# Sanity: dq, dk, dv finite-ness (no NaN / Inf) +# --------------------------------------------------------------------------- + + +def test_g24_no_nan_in_grads_dense_fp32(): + """Sanity: kernel BWD does not introduce NaN / Inf for a typical fp32 dense case.""" + inp = _build_inputs( + B=1, + H=4, + S=64, + D=64, + swa_window=32, + sink_on=True, + dtype=torch.float32, + kv_layout="mha", + use_hca=False, + seed=4321, + ) + out_cand = v4_attention_v1( + inp["q"], + inp["k"], + inp["v"], + sink=inp["sink"], + swa_window=inp["kernel_swa_window"], + additive_mask=inp["kernel_mask"], + attn_dropout=0.0, + training=False, + scale=1.0 / math.sqrt(64), + hca_local_seqlen=inp["hca_local_seqlen"], + ) + grads = _grads_from(out_cand, inp["q"], inp["k"], inp["v"], inp["sink"]) + for name, g in zip(("q", "k", "v", "sink"), grads): + if g is None: + continue + assert torch.isfinite(g).all(), f"d{name} contains NaN/Inf" diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_fwd.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_fwd.py new file mode 100644 index 000000000..d28d0cc59 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_attention_fwd.py @@ -0,0 +1,409 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-4 P25 G23 — `v4_attention_v1` Triton FWD equivalence to eager. + +Asserts that :func:`v4_attention_v1` (Triton kernel from +``primus...transformer.v4_attention_kernels.v4_attention``) produces +forward output equal to :func:`eager_v4_attention` within the plan-4 +tolerance budget across: + +* V4-Flash and V4-Pro shape envelopes (head_dim=512, H ∈ {64, 128}); +* ``compress_ratio ∈ {0, 128}`` — dense + SWA + sink (no bias) and HCA + (joint-softmax additive bias, no in-kernel SWA); +* fp32 and bf16 inputs; +* ``sink ∈ {None, learned [H]}``; +* MQA (``K_H == 1``) and MHA (``K_H == HQ``) layouts. + +The test is GPU-only (Triton requires CUDA / HIP); CPU runs are +``pytest.skip``-ed at module collection time. +""" + +from __future__ import annotations + +import math + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("v4_attention_v1 Triton kernel requires CUDA / HIP", allow_module_level=True) + +# Triton import (must be importable in the env). +pytest.importorskip("triton", reason="Triton not installed") + +from primus.backends.megatron.core.transformer.sliding_window_kv import ( # noqa: E402 + sliding_window_causal_mask, +) +from primus.backends.megatron.core.transformer.v4_attention_kernels import ( # noqa: E402 + eager_v4_attention, + v4_attention_v1, +) + +# --------------------------------------------------------------------------- +# Shape envelope (V4-Flash / V4-Pro) +# --------------------------------------------------------------------------- + + +# Two tiers of shapes are exposed to the parametrise decorator below: +# +# * Fast tier — toy ``head_dim=64`` / ``H ∈ {4, 8}`` / ``S=64`` shapes +# that exercise every code path in the kernel (in-kernel SWA mask, +# caller-supplied additive_mask, MQA/MHA, sink) in milliseconds. Run +# on every ``pytest`` invocation. +# +# * Release tier (``pytest.mark.slow``) — production V4 dimensions +# (``head_dim=512``, real ``H``, real ``swa_window``) calibrated so +# the eager fp32 reference fits MI355X HBM. The release tier is the +# plan-4 G28 gate that empirically confirms kernel correctness at the +# exact ``head_dim`` that plan-4 exists to solve. ``S`` is calibrated +# per variant: V4-Flash @ S=1024 / V4-Pro @ S=512 keep peak fp32 +# memory below ~1 GiB per test even for the MHA layout. +# +# The full ``S=4096`` smoke is owned by G30 (P27 smoke gate); these +# unit tests intentionally stay below that to keep the per-test cost +# in the ``-m slow`` budget. +_BASE_SHAPES = [ + # (variant, B, H, S, head_dim, swa_window) + ("v4_flash_small", 1, 8, 64, 64, 32), + ("v4_pro_small", 1, 4, 64, 64, 32), + pytest.param( + "v4_flash_release", + 1, + 64, + 1024, + 512, + 128, + marks=pytest.mark.slow, + ), + pytest.param( + "v4_pro_release", + 1, + 128, + 512, + 512, + 128, + marks=pytest.mark.slow, + ), +] +_DTYPES = [torch.float32, torch.bfloat16] +_SINK_MODES = [True, False] +_KV_LAYOUTS = ["mqa", "mha"] # K_H == 1 vs K_H == HQ + + +def _is_release_tier(variant: str) -> bool: + """Release-tier shapes are tagged by name (``*_release``).""" + return variant.endswith("_release") + + +def _fwd_tol(dtype: torch.dtype, *, release: bool = False) -> dict: + if dtype == torch.float32: + return {"atol": 1e-4, "rtol": 1e-4} + if dtype == torch.bfloat16: + # Release-tier ``head_dim=512`` accumulates ~8x more terms than + # the ``head_dim=64`` fast tier; bf16 long-tail rounding scales + # as ~sqrt(N). Empirically the worst-case outlier sits between + # 2.0e-2 and 5.0e-2 at release-tier dims; we loosen by ~2.5x. + return {"atol": 5e-2, "rtol": 5e-2} if release else {"atol": 2e-2, "rtol": 2e-2} + raise ValueError(f"unsupported dtype {dtype!r}") + + +def _make_dense_inputs( + *, + B: int, + H: int, + S: int, + D: int, + swa_window: int, + sink_on: bool, + dtype: torch.dtype, + kv_layout: str, + seed: int = 1234, +): + """Build inputs for a dense (compress_ratio=0) test case. + + The kernel is invoked with ``swa_window > 0, additive_mask=None``; + the eager reference is invoked with the equivalent ``additive_mask`` + pre-built (so both go through identical mask math). + """ + g = torch.Generator(device="cuda").manual_seed(seed) + device = "cuda" + K_H = 1 if kv_layout == "mqa" else H + + q = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype) + k = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype) + v = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype) + sink = torch.randn(H, generator=g, device=device, dtype=torch.float32) * 0.1 if sink_on else None + return dict(B=B, H=H, S=S, D=D, swa_window=swa_window, q=q, k=k, v=v, sink=sink) + + +def _make_hca_inputs( + *, + B: int, + H: int, + S: int, + P: int, + D: int, + swa_window: int, + sink_on: bool, + dtype: torch.dtype, + kv_layout: str, + seed: int = 1234, +): + """Build inputs for an HCA (compress_ratio=128) test case. + + HCA pre-concatenates a length-``P`` compressed-pool key/value to the + length-``S`` local key/value, and the caller pre-builds the joint + additive mask so the kernel does NOT apply SWA / causal in-kernel. + """ + g = torch.Generator(device="cuda").manual_seed(seed) + device = "cuda" + K_H = 1 if kv_layout == "mqa" else H + + q = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype) + k_local = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype) + v_local = torch.randn(B, K_H, S, D, generator=g, device=device, dtype=dtype) + pool_k = torch.randn(B, K_H, P, D, generator=g, device=device, dtype=dtype) + pool_v = torch.randn(B, K_H, P, D, generator=g, device=device, dtype=dtype) + k_full = torch.cat([k_local, pool_k], dim=2) + v_full = torch.cat([v_local, pool_v], dim=2) + + # Local SWA-causal mask: [S, S] + local_mask = sliding_window_causal_mask(S, swa_window, device=device, dtype=dtype) + # Pool causal-on-stride mask: pool[s] is visible at query t iff + # (s+1)*ratio - 1 <= t. Use a fixed ratio of 4 for the test. + ratio = 4 + t = torch.arange(S, device=device).unsqueeze(1) + s_end = (torch.arange(P, device=device).unsqueeze(0) + 1) * ratio - 1 + pool_mask = torch.where(s_end <= t, 0.0, float("-inf")).to(dtype) + full_mask = torch.cat([local_mask, pool_mask], dim=-1) # [S, S+P] + + sink = torch.randn(H, generator=g, device=device, dtype=torch.float32) * 0.1 if sink_on else None + return dict( + B=B, + H=H, + S=S, + P=P, + D=D, + q=q, + k=k_full, + v=v_full, + sink=sink, + full_mask=full_mask, + pool_mask=pool_mask, + ) + + +# --------------------------------------------------------------------------- +# G23 — dense (compress_ratio == 0) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("variant,B,H,S,D,swa_window", _BASE_SHAPES) +@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) +@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"]) +@pytest.mark.parametrize("kv_layout", _KV_LAYOUTS) +def test_g23_dense_fwd_matches_eager( + variant: str, + B: int, + H: int, + S: int, + D: int, + swa_window: int, + dtype: torch.dtype, + sink_on: bool, + kv_layout: str, +): + """Dense path: kernel uses in-kernel SWA-causal mask + optional sink.""" + toy = _make_dense_inputs( + B=B, + H=H, + S=S, + D=D, + swa_window=swa_window, + sink_on=sink_on, + dtype=dtype, + kv_layout=kv_layout, + ) + + scale = 1.0 / math.sqrt(D) + + # Reference: eager_v4_attention with pre-built SWA additive_mask + eager_mask = sliding_window_causal_mask(S, swa_window, device=toy["q"].device, dtype=dtype) + out_ref = eager_v4_attention( + toy["q"], + toy["k"], + toy["v"], + sink=toy["sink"], + swa_window=0, + additive_mask=eager_mask, + attn_dropout=0.0, + training=False, + scale=scale, + ) + + # Candidate: v4_attention_v1 with swa_window > 0, additive_mask=None + out_cand = v4_attention_v1( + toy["q"], + toy["k"], + toy["v"], + sink=toy["sink"], + swa_window=swa_window, + additive_mask=None, + attn_dropout=0.0, + training=False, + scale=scale, + ) + + assert out_ref.shape == out_cand.shape == toy["q"].shape + assert out_ref.dtype == out_cand.dtype == dtype + torch.testing.assert_close( + out_cand, + out_ref, + **_fwd_tol(dtype, release=_is_release_tier(variant)), + ) + + +# --------------------------------------------------------------------------- +# G23 — HCA (compress_ratio == 128) — caller-supplied additive mask +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("variant,B,H,S,D,swa_window", _BASE_SHAPES) +@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) +@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"]) +@pytest.mark.parametrize("kv_layout", _KV_LAYOUTS) +def test_g23_hca_style_fwd_matches_eager( + variant: str, + B: int, + H: int, + S: int, + D: int, + swa_window: int, + dtype: torch.dtype, + sink_on: bool, + kv_layout: str, +): + """HCA path: caller pre-concatenates pool keys + supplies joint additive mask.""" + P = 4 # tiny pool — exercises the additive_mask branch w/o full HCA setup + toy = _make_hca_inputs( + B=B, + H=H, + S=S, + P=P, + D=D, + swa_window=swa_window, + sink_on=sink_on, + dtype=dtype, + kv_layout=kv_layout, + ) + + scale = 1.0 / math.sqrt(D) + + out_ref = eager_v4_attention( + toy["q"], + toy["k"], + toy["v"], + sink=toy["sink"], + swa_window=0, + additive_mask=toy["full_mask"], + attn_dropout=0.0, + training=False, + scale=scale, + ) + + out_cand = v4_attention_v1( + toy["q"], + toy["k"], + toy["v"], + sink=toy["sink"], + swa_window=swa_window, + additive_mask=toy["pool_mask"], + attn_dropout=0.0, + training=False, + scale=scale, + hca_local_seqlen=S, + ) + + assert out_ref.shape == out_cand.shape == toy["q"].shape + assert out_ref.dtype == out_cand.dtype == dtype + torch.testing.assert_close( + out_cand, + out_ref, + **_fwd_tol(dtype, release=_is_release_tier(variant)), + ) + + +# --------------------------------------------------------------------------- +# G25 — determinism with attn_dropout=0.0 +# --------------------------------------------------------------------------- + + +def test_g25_determinism_fp32_mha(): + """Repeated FWD calls with the same inputs produce bit-identical output (fp32 / MHA).""" + toy = _make_dense_inputs( + B=1, + H=4, + S=64, + D=64, + swa_window=32, + sink_on=True, + dtype=torch.float32, + kv_layout="mha", + ) + scale = 1.0 / math.sqrt(toy["D"]) + + out_a = v4_attention_v1( + toy["q"], + toy["k"], + toy["v"], + sink=toy["sink"], + swa_window=toy["swa_window"], + additive_mask=None, + attn_dropout=0.0, + training=False, + scale=scale, + ) + out_b = v4_attention_v1( + toy["q"], + toy["k"], + toy["v"], + sink=toy["sink"], + swa_window=toy["swa_window"], + additive_mask=None, + attn_dropout=0.0, + training=False, + scale=scale, + ) + assert torch.equal(out_a, out_b), "v4_attention_v1 FWD is non-deterministic at fp32 / MHA" + + +def test_g25_dropout_with_training_is_rejected(): + """``attn_dropout > 0`` with ``training=True`` raises (kernel does not implement dropout).""" + toy = _make_dense_inputs( + B=1, + H=4, + S=64, + D=64, + swa_window=32, + sink_on=False, + dtype=torch.float32, + kv_layout="mha", + ) + scale = 1.0 / math.sqrt(toy["D"]) + with pytest.raises(NotImplementedError, match="dropout"): + v4_attention_v1( + toy["q"], + toy["k"], + toy["v"], + sink=None, + swa_window=toy["swa_window"], + additive_mask=None, + attn_dropout=0.1, + training=True, + scale=scale, + ) diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_csa_in_kernel_gather.py b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_csa_in_kernel_gather.py new file mode 100644 index 000000000..0f2d8c371 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_v4_csa_in_kernel_gather.py @@ -0,0 +1,280 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Plan-5 P31 — CSA in-kernel top-K gather equivalence. + +The P26 CSA kernel consumed a materialised ``gathered`` tensor. P31 moves +that gather into the Triton kernel and scatters gradients directly back +to the compressed pool. These tests compare the new pool/topk API against +the eager reference plus PyTorch gather autograd. +""" + +from __future__ import annotations + +import math +from typing import Optional + +import pytest + +torch = pytest.importorskip("torch") + +if not torch.cuda.is_available(): + pytest.skip("v4_csa_attention_v0 Triton kernel requires CUDA / HIP", allow_module_level=True) + +pytest.importorskip("triton", reason="Triton not installed") + +from primus.backends.megatron.core.transformer.v4_attention_kernels import ( # noqa: E402 + eager_v4_csa_attention, + v4_csa_attention_v1, +) + +_SHAPES = [ + ("fast", 1, 4, 32, 64, 32, 16, 16), + pytest.param("release_head_dim", 1, 16, 128, 512, 128, 64, 64, marks=pytest.mark.slow), +] +_DTYPES = [torch.float32, torch.bfloat16] +_SINK_MODES = [True, False] + + +def _tol(dtype: torch.dtype, *, release: bool = False) -> dict: + if dtype == torch.float32: + return {"atol": 1e-4, "rtol": 1e-4} + if dtype == torch.bfloat16: + return {"atol": 2e-1, "rtol": 2e-1} if release else {"atol": 5e-2, "rtol": 5e-2} + raise ValueError(f"unsupported dtype {dtype!r}") + + +def _sink_tol(dtype: torch.dtype, *, release: bool = False) -> dict: + if dtype == torch.float32: + return {"atol": 1e-4, "rtol": 1e-4} + if dtype == torch.bfloat16: + return {"atol": 5e-2, "rtol": 5e-2} if release else {"atol": 5e-3, "rtol": 5e-3} + raise ValueError(f"unsupported dtype {dtype!r}") + + +def _make_inputs( + *, + B: int, + H: int, + S: int, + D: int, + P: int, + K_topk: int, + sink_on: bool, + dtype: torch.dtype, + requires_grad: bool, + seed: int = 20260509, +): + g = torch.Generator(device="cuda").manual_seed(seed) + device = "cuda" + + q = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype).requires_grad_(requires_grad) + k_local = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype).requires_grad_(requires_grad) + v_local = torch.randn(B, H, S, D, generator=g, device=device, dtype=dtype).requires_grad_(requires_grad) + pool = torch.randn(B, P, D, generator=g, device=device, dtype=dtype).requires_grad_(requires_grad) + + topk_idxs = torch.randint(0, P, (B, S, K_topk), generator=g, device=device, dtype=torch.int64) + if K_topk >= 4: + # Exercise both invalid slots and duplicate slots, which are the + # two cases the in-kernel scatter-add must handle correctly. + topk_idxs[..., 0] = -1 + topk_idxs[..., 1] = topk_idxs[..., 2] + + sink = None + if sink_on: + sink_raw = torch.randn(H, generator=g, device=device, dtype=torch.float32) * 0.1 + sink = sink_raw.detach().clone().requires_grad_(requires_grad) + + return dict(q=q, k_local=k_local, v_local=v_local, pool=pool, topk_idxs=topk_idxs, sink=sink) + + +def _gather_from_pool(pool: torch.Tensor, topk_idxs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + B, P, D = pool.shape + _, S, K = topk_idxs.shape + valid = topk_idxs >= 0 + safe_idx = topk_idxs.clamp(min=0) + gathered = torch.gather( + pool.unsqueeze(1).expand(B, S, P, D), + dim=2, + index=safe_idx.unsqueeze(-1).expand(B, S, K, D), + ) + gathered = gathered * valid.unsqueeze(-1).to(pool.dtype) + sparse_mask = torch.where(valid, 0.0, float("-inf")).to(pool.dtype) + return gathered, sparse_mask + + +def _grads_from(model_out: torch.Tensor, *leaves: Optional[torch.Tensor]) -> list[Optional[torch.Tensor]]: + real_leaves = [leaf for leaf in leaves if leaf is not None] + grads = torch.autograd.grad( + outputs=model_out, + inputs=real_leaves, + grad_outputs=torch.ones_like(model_out), + retain_graph=False, + create_graph=False, + allow_unused=False, + ) + out: list[Optional[torch.Tensor]] = [] + idx = 0 + for leaf in leaves: + if leaf is None: + out.append(None) + else: + out.append(grads[idx]) + idx += 1 + return out + + +@pytest.mark.parametrize("variant,B,H,S,D,P,K_topk,swa_window", _SHAPES) +@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) +@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"]) +def test_p31_csa_pool_fwd_matches_gathered_reference( + variant: str, + B: int, + H: int, + S: int, + D: int, + P: int, + K_topk: int, + swa_window: int, + dtype: torch.dtype, + sink_on: bool, +): + inp = _make_inputs( + B=B, + H=H, + S=S, + D=D, + P=P, + K_topk=K_topk, + sink_on=sink_on, + dtype=dtype, + requires_grad=False, + ) + gathered, sparse_mask = _gather_from_pool(inp["pool"], inp["topk_idxs"]) + scale = 1.0 / math.sqrt(D) + + out_ref = eager_v4_csa_attention( + inp["q"], + inp["k_local"], + inp["v_local"], + gathered, + sink=inp["sink"], + swa_window=swa_window, + sparse_mask=sparse_mask, + attn_dropout=0.0, + training=False, + scale=scale, + ) + out_cand = v4_csa_attention_v1( + inp["q"], + inp["k_local"], + inp["v_local"], + inp["pool"], + topk_idxs=inp["topk_idxs"], + sink=inp["sink"], + swa_window=swa_window, + attn_dropout=0.0, + training=False, + scale=scale, + ) + + torch.testing.assert_close(out_cand, out_ref, **_tol(dtype, release=variant != "fast")) + + +@pytest.mark.parametrize("variant,B,H,S,D,P,K_topk,swa_window", _SHAPES) +@pytest.mark.parametrize("dtype", _DTYPES, ids=lambda d: str(d).rsplit(".", 1)[-1]) +@pytest.mark.parametrize("sink_on", _SINK_MODES, ids=["sink_on", "sink_off"]) +def test_p31_csa_pool_bwd_matches_gathered_reference( + variant: str, + B: int, + H: int, + S: int, + D: int, + P: int, + K_topk: int, + swa_window: int, + dtype: torch.dtype, + sink_on: bool, +): + ref_inp = _make_inputs( + B=B, + H=H, + S=S, + D=D, + P=P, + K_topk=K_topk, + sink_on=sink_on, + dtype=dtype, + requires_grad=True, + seed=4242, + ) + cand_inp = _make_inputs( + B=B, + H=H, + S=S, + D=D, + P=P, + K_topk=K_topk, + sink_on=sink_on, + dtype=dtype, + requires_grad=True, + seed=4242, + ) + scale = 1.0 / math.sqrt(D) + + gathered, sparse_mask = _gather_from_pool(ref_inp["pool"], ref_inp["topk_idxs"]) + out_ref = eager_v4_csa_attention( + ref_inp["q"], + ref_inp["k_local"], + ref_inp["v_local"], + gathered, + sink=ref_inp["sink"], + swa_window=swa_window, + sparse_mask=sparse_mask, + attn_dropout=0.0, + training=False, + scale=scale, + ) + dq_ref, dkl_ref, dvl_ref, dpool_ref, dsink_ref = _grads_from( + out_ref, + ref_inp["q"], + ref_inp["k_local"], + ref_inp["v_local"], + ref_inp["pool"], + ref_inp["sink"], + ) + + out_cand = v4_csa_attention_v1( + cand_inp["q"], + cand_inp["k_local"], + cand_inp["v_local"], + cand_inp["pool"], + topk_idxs=cand_inp["topk_idxs"], + sink=cand_inp["sink"], + swa_window=swa_window, + attn_dropout=0.0, + training=False, + scale=scale, + ) + dq_cand, dkl_cand, dvl_cand, dpool_cand, dsink_cand = _grads_from( + out_cand, + cand_inp["q"], + cand_inp["k_local"], + cand_inp["v_local"], + cand_inp["pool"], + cand_inp["sink"], + ) + + release = variant != "fast" + tol = _tol(dtype, release=release) + torch.testing.assert_close(dq_cand, dq_ref, **tol) + torch.testing.assert_close(dkl_cand, dkl_ref, **tol) + torch.testing.assert_close(dvl_cand, dvl_ref, **tol) + torch.testing.assert_close(dpool_cand, dpool_ref, **tol) + if sink_on: + torch.testing.assert_close(dsink_cand, dsink_ref, **_sink_tol(dtype, release=release)) + else: + assert dsink_ref is None and dsink_cand is None diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_shapes.py b/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_shapes.py new file mode 100644 index 000000000..5254a2b95 --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_shapes.py @@ -0,0 +1,260 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""V4-Flash / V4-Pro attention shape fixtures (plan-4 P24). + +This module exposes the **single source-of-truth** shape tables every +plan-4 test parametrises over. The canonical config knobs come from: + +* ``primus/configs/models/megatron/deepseek_v4_flash.yaml`` +* ``primus/configs/models/megatron/deepseek_v4_pro.yaml`` +* ``primus/configs/models/megatron/deepseek_v4_base.yaml`` + +Each variant exposes three sequence-length tiers: + +* ``small`` (S=128) — fast CI tier; runs in milliseconds on CPU +* ``medium`` (S=512) — moderate tier; runs in seconds on a single GPU +* ``large`` (S=4096) — release tier; gated behind ``pytest.mark.slow`` + +Each shape is parametrised by the ``compress_ratio`` of the layer +under test (``0`` dense + SWA + sink, ``128`` HCA, ``4`` CSA). Pool +size ``P = S // compress_ratio`` for compressed branches; the indexer +top-K (``K``) follows the per-variant ``index_topk`` knob (clamped to +``P`` when ``P < K``, mirroring how the indexer behaves at small +sequences). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable, Optional, Tuple + +SeqTier = str # "small" | "medium" | "large" + +_SEQ_TIERS: Tuple[SeqTier, ...] = ("small", "medium", "large") +_SEQ_LENGTHS = { + "small": 128, + "medium": 512, + "large": 4096, +} + + +@dataclass(frozen=True) +class V4AttnShape: + """Single-shape parametrisation for a V4 attention layer test. + + Attributes follow the V4 attention class: + + * ``B`` — batch size + * ``S`` — sequence length (post-RoPE; raw token count) + * ``H`` — number of attention heads (per-rank; tests run TP=1) + * ``head_dim`` — per-head ``kv_channels`` + * ``q_pe_dim`` — partial-RoPE rotary head dim (channels with RoPE + applied; the remaining ``head_dim - q_pe_dim`` channels are + "NOPE") + * ``attn_sliding_window`` — SWA window for dense / CSA local branch + * ``sink`` — ``True`` enables the per-head ``[H]`` learned softmax + sink (matches V4 yaml) + * ``compress_ratio`` — ``0`` dense, ``4`` CSA, ``128`` HCA + * ``P`` — compressed pool size (``S // compress_ratio`` when + ``compress_ratio > 0``; ``0`` for dense) + * ``K`` — indexer top-K (``min(index_topk, P)`` for CSA; + meaningless / ``0`` for dense / HCA but kept on the dataclass + for API uniformity) + * ``variant`` — display name ("v4_flash" / "v4_pro") for test ids + """ + + variant: str + B: int + S: int + H: int + head_dim: int + q_pe_dim: int + attn_sliding_window: int + sink: bool + compress_ratio: int + P: int + K: int + + @property + def Sk(self) -> int: + """Effective key axis length per branch. + + * dense (cr=0): ``S`` local keys + * HCA (cr=128): ``S + P`` (local + compressed pool) + * CSA (cr=4): ``S`` local keys (the sparse top-K branch is + handled separately via ``gathered``) + """ + if self.compress_ratio == 128: + return self.S + self.P + return self.S + + def shape_id(self) -> str: + """A short human-readable test-id for pytest parametrisation.""" + return ( + f"{self.variant}-cr{self.compress_ratio}-S{self.S}" + f"-H{self.H}-D{self.head_dim}" + f"-sink{int(self.sink)}-w{self.attn_sliding_window}" + ) + + +# --------------------------------------------------------------------------- +# V4-Flash and V4-Pro variant tables +# --------------------------------------------------------------------------- + +# Source: deepseek_v4_flash.yaml + deepseek_v4_base.yaml +_V4_FLASH_DEFAULTS = dict( + H=64, + head_dim=512, + q_pe_dim=64, + attn_sliding_window=128, + sink=True, + index_topk=512, +) + +# Source: deepseek_v4_pro.yaml + deepseek_v4_base.yaml +_V4_PRO_DEFAULTS = dict( + H=128, + head_dim=512, + q_pe_dim=64, + attn_sliding_window=128, + sink=True, + index_topk=1024, +) + + +def _make_shape( + variant: str, + defaults: dict, + *, + compress_ratio: int, + seq: int, + batch: int, + sink: Optional[bool] = None, +) -> V4AttnShape: + H = int(defaults["H"]) + head_dim = int(defaults["head_dim"]) + q_pe_dim = int(defaults["q_pe_dim"]) + attn_sliding_window = int(defaults["attn_sliding_window"]) + sink_default = bool(defaults["sink"]) + index_topk = int(defaults["index_topk"]) + + if compress_ratio == 0: + P = 0 + K = 0 + elif compress_ratio == 4: + P = max(seq // compress_ratio, 1) + K = min(index_topk, P) + elif compress_ratio == 128: + P = max(seq // compress_ratio, 1) + K = 0 # not used by HCA + else: + raise ValueError(f"compress_ratio must be in {{0, 4, 128}}; got {compress_ratio}.") + + return V4AttnShape( + variant=variant, + B=int(batch), + S=int(seq), + H=H, + head_dim=head_dim, + q_pe_dim=q_pe_dim, + attn_sliding_window=attn_sliding_window, + sink=sink_default if sink is None else bool(sink), + compress_ratio=int(compress_ratio), + P=int(P), + K=int(K), + ) + + +def v4_flash_shape( + *, + compress_ratio: int, + seq_tier: SeqTier = "small", + batch: int = 1, + sink: Optional[bool] = None, +) -> V4AttnShape: + """Build a single V4-Flash :class:`V4AttnShape`.""" + if seq_tier not in _SEQ_LENGTHS: + raise ValueError(f"seq_tier must be one of {_SEQ_TIERS}; got {seq_tier!r}.") + return _make_shape( + "v4_flash", + _V4_FLASH_DEFAULTS, + compress_ratio=compress_ratio, + seq=_SEQ_LENGTHS[seq_tier], + batch=batch, + sink=sink, + ) + + +def v4_pro_shape( + *, + compress_ratio: int, + seq_tier: SeqTier = "small", + batch: int = 1, + sink: Optional[bool] = None, +) -> V4AttnShape: + """Build a single V4-Pro :class:`V4AttnShape`.""" + if seq_tier not in _SEQ_LENGTHS: + raise ValueError(f"seq_tier must be one of {_SEQ_TIERS}; got {seq_tier!r}.") + return _make_shape( + "v4_pro", + _V4_PRO_DEFAULTS, + compress_ratio=compress_ratio, + seq=_SEQ_LENGTHS[seq_tier], + batch=batch, + sink=sink, + ) + + +def v4_attention_shape_grid( + *, + variants: Iterable[str] = ("v4_flash", "v4_pro"), + compress_ratios: Iterable[int] = (0, 4, 128), + seq_tiers: Iterable[SeqTier] = ("small",), + batch: int = 1, + sinks: Iterable[bool] = (True, False), +) -> list[V4AttnShape]: + """Cartesian product of (variant × compress_ratio × seq_tier × sink). + + Returns a flat list of :class:`V4AttnShape` suitable for direct + pytest parametrisation. Plan-4 tests typically pin to a small + subset to keep CI green: + + * G22 (P24 refactor safety net): all three compress ratios at + ``small`` seq tier with both sink modes. + * G23 / G24 (P25 fwd / bwd): cr ∈ {0, 128}, seq ∈ {small, medium}. + * G26 / G27 (P26 CSA fwd / bwd): cr == 4, seq ∈ {small, medium}. + """ + out: list[V4AttnShape] = [] + builders = { + "v4_flash": v4_flash_shape, + "v4_pro": v4_pro_shape, + } + for variant in variants: + if variant not in builders: + raise ValueError(f"variant must be one of {tuple(builders)}; got {variant!r}.") + builder = builders[variant] + for cr in compress_ratios: + for tier in seq_tiers: + for sink in sinks: + out.append( + builder( + compress_ratio=cr, + seq_tier=tier, + batch=batch, + sink=sink, + ) + ) + return out + + +__all__ = [ + "SeqTier", + "V4AttnShape", + "v4_attention_shape_grid", + "v4_flash_shape", + "v4_pro_shape", +] diff --git a/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_test_utils.py b/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_test_utils.py new file mode 100644 index 000000000..f7e33a17a --- /dev/null +++ b/tests/unit_tests/megatron/transformer/deepseek_v4/v4_attention_test_utils.py @@ -0,0 +1,221 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +"""Forward + backward equivalence harness (plan-4 P24). + +Every plan-4 test that compares an eager-Python reference against a +candidate (P25 / P26 Triton kernels, or the new P24 reference op +itself) runs through :func:`compare_fwd_bwd`. The harness: + +* clones every input that ``requires_grad`` so reference and candidate + run on independent autograd graphs (one ``.grad`` accumulator each); +* invokes the reference and candidate; +* asserts forward output matches within the supplied tolerance; +* calls ``.sum().backward()`` on each (with ``retain_graph=False``); +* asserts every cloned-and-grad-bearing input's gradient matches + within the supplied tolerance; +* on mismatch, prints a structured diff (max abs / max rel error per + leaf, with the eight worst entries by magnitude). + +Tolerance budget defaults follow the plan-4 P24 design: + +* ``fp32`` — ``atol=1e-5, rtol=1e-5`` for forward, same for backward. +* ``bf16`` — ``atol=2e-2, rtol=2e-2`` for forward, ``atol=5e-2, + rtol=5e-2`` for backward (the looser bwd budget absorbs MQA + atomic-add reordering and other accumulator-order delta). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Iterable, Mapping, Optional + +import torch + +# --------------------------------------------------------------------------- +# Tolerance budget +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class Tol: + """Numerical tolerance pair for ``torch.allclose``-style comparisons.""" + + atol: float + rtol: float + + def as_kwargs(self) -> dict: + return {"atol": self.atol, "rtol": self.rtol} + + +FP32_TOL = Tol(atol=1e-5, rtol=1e-5) +BF16_FWD_TOL = Tol(atol=2e-2, rtol=2e-2) +BF16_BWD_TOL = Tol(atol=5e-2, rtol=5e-2) + + +def default_tols(dtype: torch.dtype) -> tuple[Tol, Tol]: + """Return ``(fwd_tol, bwd_tol)`` for the supplied input dtype.""" + if dtype == torch.float32: + return FP32_TOL, FP32_TOL + if dtype == torch.bfloat16: + return BF16_FWD_TOL, BF16_BWD_TOL + if dtype == torch.float16: + # Treat fp16 like bf16 for tolerance purposes (slightly tighter + # mantissa, slightly tighter exponent — same ballpark). + return BF16_FWD_TOL, BF16_BWD_TOL + raise ValueError(f"No default tolerance preset for dtype {dtype!r}.") + + +# --------------------------------------------------------------------------- +# Input cloning + gradient comparison helpers +# --------------------------------------------------------------------------- + + +def _clone_for_grad(inputs: Mapping[str, Any]) -> dict: + """Return a shallow-copied mapping with leaf tensors deep-cloned. + + Tensors that ``requires_grad`` are replaced with detached clones + that re-enable ``requires_grad`` so the autograd graph for + reference / candidate is independent. Non-tensor values pass + through unchanged. Tensors that do NOT require grad pass through + by reference (so the caller can share read-only tensors like masks + across reference / candidate cheaply). + """ + out: dict = {} + for k, v in inputs.items(): + if isinstance(v, torch.Tensor) and v.requires_grad: + out[k] = v.detach().clone().requires_grad_(True) + else: + out[k] = v + return out + + +def _diag(name: str, ref: torch.Tensor, cand: torch.Tensor, tol: Tol) -> str: + """Format a structured diff for an off-tolerance ``(ref, cand)`` pair.""" + if ref.shape != cand.shape: + return f"{name}: shape mismatch — ref={tuple(ref.shape)} " f"vs cand={tuple(cand.shape)}" + diff = (ref.float() - cand.float()).abs() + # rel uses ref magnitude as denominator (matches torch.allclose + # semantics: atol + rtol * |ref|). + rel = diff / (ref.float().abs() + 1e-12) + max_abs = diff.max().item() + max_rel = rel.max().item() + + flat_diff = diff.flatten() + k = min(8, flat_diff.numel()) + topk = torch.topk(flat_diff, k=k).indices.tolist() + worst = [] + flat_ref = ref.float().flatten() + flat_cand = cand.float().flatten() + for idx in topk: + worst.append( + f" [{idx}]: ref={flat_ref[idx].item():+.6g} " + f"cand={flat_cand[idx].item():+.6g} " + f"abs={(flat_ref[idx] - flat_cand[idx]).abs().item():+.6g}" + ) + worst_block = "\n".join(worst) + return ( + f"{name}: max_abs={max_abs:.6g} (atol={tol.atol:.6g}), " + f"max_rel={max_rel:.6g} (rtol={tol.rtol:.6g}); 8 worst entries:\n" + worst_block + ) + + +def _assert_close( + name: str, + ref: torch.Tensor, + cand: torch.Tensor, + tol: Tol, +) -> None: + if ref is None and cand is None: + return + if ref is None or cand is None: + raise AssertionError( + f"{name}: one side is None (ref is None: {ref is None}, " f"cand is None: {cand is None})" + ) + # Match dtype before comparing — gradient dtypes can differ + # slightly across paths (one side accumulates fp32, the other + # bf16), so coerce to fp32 for the comparison itself. + ref32 = ref.detach().float() + cand32 = cand.detach().float() + if ref32.shape != cand32.shape: + raise AssertionError(_diag(name, ref32, cand32, tol)) + if torch.allclose(ref32, cand32, **tol.as_kwargs()): + return + raise AssertionError(_diag(name, ref32, cand32, tol)) + + +# --------------------------------------------------------------------------- +# Public entry point +# --------------------------------------------------------------------------- + + +def compare_fwd_bwd( + *, + reference: Callable[..., torch.Tensor], + candidate: Callable[..., torch.Tensor], + inputs: Mapping[str, Any], + fwd_tol: Tol, + bwd_tol: Tol, + grad_keys: Optional[Iterable[str]] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run reference + candidate on independent input clones; assert match. + + Args: + reference: callable taking ``**inputs`` and returning a single + tensor (the "eager truth"). + candidate: callable with the same signature (the path under + test). + inputs: kwargs forwarded to both. Tensor values that + ``requires_grad`` are independently cloned so the two + backward passes do not interfere. + fwd_tol: forward output tolerance. + bwd_tol: backward gradient tolerance. + grad_keys: optional override for which input keys to assert + gradients on. Defaults to every tensor in ``inputs`` that + ``requires_grad``. + + Returns: + ``(out_ref, out_cand)`` for follow-up inspection by the caller + (e.g., asserting non-NaN, asserting shape, etc.). + + Raises: + AssertionError: on first forward or backward mismatch, with a + structured diff in the message. + """ + ref_inputs = _clone_for_grad(inputs) + cand_inputs = _clone_for_grad(inputs) + + out_ref = reference(**ref_inputs) + out_cand = candidate(**cand_inputs) + + _assert_close("forward output", out_ref, out_cand, fwd_tol) + + out_ref.sum().backward() + out_cand.sum().backward() + + if grad_keys is None: + grad_keys = [k for k, v in inputs.items() if isinstance(v, torch.Tensor) and v.requires_grad] + for k in grad_keys: + ref_t = ref_inputs[k] + cand_t = cand_inputs[k] + if not isinstance(ref_t, torch.Tensor) or not isinstance(cand_t, torch.Tensor): + raise AssertionError( + f"grad_keys[{k!r}] points at a non-tensor input " + f"(ref={type(ref_t).__name__}, cand={type(cand_t).__name__})" + ) + _assert_close(f"gradient[{k}]", ref_t.grad, cand_t.grad, bwd_tol) + + return out_ref, out_cand + + +__all__ = [ + "Tol", + "FP32_TOL", + "BF16_FWD_TOL", + "BF16_BWD_TOL", + "compare_fwd_bwd", + "default_tols", +] diff --git a/tools/backend_gap_report/build_site_bundle.py b/tools/backend_gap_report/build_site_bundle.py index 134124fce..fa5041199 100644 --- a/tools/backend_gap_report/build_site_bundle.py +++ b/tools/backend_gap_report/build_site_bundle.py @@ -14,6 +14,10 @@ REPO_ROOT = Path(__file__).resolve().parents[2] BACKEND_GAP_DOCS_ROOT = REPO_ROOT / "docs" / "backend-gap" SITE_SOURCE_ROOT = REPO_ROOT / "tools" / "backend_gap_report" / "site" +# Standalone DeepSeek-V4 performance-projection site, published as a subpath of +# the same Pages bundle (served at //deepseek-v4-projection/). +PROJECTION_SITE_ROOT = REPO_ROOT / "deepseek-v4" / "projection" / "site" +PROJECTION_SITE_SUBDIR = "deepseek-v4-projection" SOURCE_DASHBOARD_DATA_DIR = BACKEND_GAP_DOCS_ROOT / "dashboard-data" METADATA_REPORTS_DIR = SOURCE_DASHBOARD_DATA_DIR / "reports" PDF_TEMPLATE = REPO_ROOT / "tools" / "backend_gap_report" / "templates" / "pdf-report.css" @@ -271,6 +275,9 @@ def build_site(output_dir: Path) -> None: copy_tree(SOURCE_DASHBOARD_DATA_DIR, output_dir / "dashboard-data") build_combined_reports_index(output_dir) build_pdf_artifacts(output_dir) + if PROJECTION_SITE_ROOT.exists(): + print("[projection] Copy DeepSeek-V4 projection site", flush=True) + copy_tree(PROJECTION_SITE_ROOT, output_dir / PROJECTION_SITE_SUBDIR) print("[backend-gap] Validate standalone dashboard bundle", flush=True) validate_bundle(output_dir) From 225fb8dc83d62937b6f7e437a16fccc0af51b07c Mon Sep 17 00:00:00 2001 From: wenxie-amd Date: Fri, 17 Jul 2026 17:27:56 +0800 Subject: [PATCH 041/127] Support Crusoe (Spur) cluster for DeepSeek-V4 multi-node runs (#885) ### Summary Enables the DeepSeek-V4 flash example to launch multi-node training on the Crusoe cluster, which runs the **Spur** scheduler (a SLURM-compatible controller) with `ionic` (AINIC) RDMA NICs. Getting a 4-node run off the ground surfaced several launcher/config issues that also affect any non-stock-SLURM cluster; this PR fixes them. ### Changes **`runner/primus-cli-slurm-entry.sh`** - Add a pure-bash `SLURM_NODELIST` expander used as a fallback when `scontrol show hostnames` is unavailable. Spur's `scontrol` only supports `job/node/partition/reservation/federation/config`, so the previous `command -v scontrol` gate still took the broken path and failed to resolve the master host. The expander handles comma lists and bracket forms (`prefix[01-04,06]`, `prefix[030,058]`). - Match only the leading hostname label in the `MASTER_ADDR` consistency check, so a scheduler-provided FQDN (Spur exports `MASTER_ADDR=.crusoe.amd.com`) is accepted against the short name produced by nodelist expansion. A genuine mismatch is still caught. **`examples/deepseek-v4/run_deepseek_v4_flash.sh`** - Remove the literal double-quotes wrapping `PRIMUS_PP_LAYOUT` and `PRIMUS_COMPRESS_RATIOS`. The quotes were passed through verbatim as CLI overrides, and Megatron's pipeline-layout parser rejected them (`AssertionError: Invalid layer character: "`). This only triggered for multi-stage PP layouts (`NNODES=4/8`), so single-stage smoke runs never hit it. Values are still single-quoted in bash to protect `|`, `*`, `(`, `)`, and spaces. **`examples/deepseek-v4/run_deepseek_v4.sh`** - Make `SLURM_PARTITION` overridable via env (was hard-coded to `Compute-DCPT`). - Add `PRIMUS_OUTPUT_ROOT` (default `output`) for the host-side launcher log directory, so runs work when the canonical `output/` tree is owned by root from earlier privileged runs. ### Notes for reviewers - No behavior change for stock-SLURM clusters: `scontrol show hostnames` is still preferred when it works, and all new env knobs default to the previous values. - For multi-node runs on Spur nodes, the socket bootstrap interface must be a routable NIC; pass `NCCL_SOCKET_IFNAME`/`GLOO_SOCKET_IFNAME=ens3` (the script otherwise falls back to `lo`). ### Testing - Verified nodelist expansion for comma / bracket-list / bracket-range / mixed forms. - Launched a 4-node DeepSeek-V4 flash run on `amd-spur`: passes scheduler launch, master resolution, layout parsing, NCCL init, model build, and dataset setup. (A separate image-side stall during optimizer/DeepEP setup with `tasimage/primus:pr-862` is still under investigation and is out of scope for this launcher/config PR.) --- Want me to create the PR into `main` with this description? Co-authored-by: Cursor --- examples/deepseek-v4/run_deepseek_v4.sh | 10 ++- examples/deepseek-v4/run_deepseek_v4_flash.sh | 10 +-- runner/primus-cli-slurm-entry.sh | 69 +++++++++++++++---- 3 files changed, 69 insertions(+), 20 deletions(-) diff --git a/examples/deepseek-v4/run_deepseek_v4.sh b/examples/deepseek-v4/run_deepseek_v4.sh index eefd55e80..87796618f 100755 --- a/examples/deepseek-v4/run_deepseek_v4.sh +++ b/examples/deepseek-v4/run_deepseek_v4.sh @@ -25,7 +25,7 @@ export NNODES=${NNODES:-1} export TRAIN_ITERS=${TRAIN_ITERS:-20} export DOCKER_IMAGE=${DOCKER_IMAGE:-"tasimage/primus:pr-715-ainic"} -export SLURM_PARTITION=Compute-DCPT +export SLURM_PARTITION=${SLURM_PARTITION:-Compute-DCPT} export SLURM_NODELIST="${SLURM_NODELIST:-smci355-ccs-aus-n01-21,smci355-ccs-aus-n01-33,smci355-ccs-aus-n02-21,smci355-ccs-aus-n02-25,smci355-ccs-aus-n02-29,smci355-ccs-aus-n02-33,smci355-ccs-aus-n03-33,smci355-ccs-aus-n04-21,smci355-ccs-aus-n04-25,smci355-ccs-aus-n04-29,smci355-ccs-aus-n04-33,smci355-ccs-aus-n05-21,smci355-ccs-aus-n05-29,smci355-ccs-aus-n05-33,smci355-ccs-aus-n06-25,smci355-ccs-aus-n06-33,smci355-ccs-aus-n10-29}" export MASTER_PORT=${MASTER_PORT:-29500} @@ -230,6 +230,10 @@ export BACKEND_PATH=${BACKEND_PATH:-"$(pwd)/third_party/Megatron-LM"} export PRIMUS_TEAM=${PRIMUS_TEAM:-amd} export PRIMUS_USER=${PRIMUS_USER:-tas-mi355x-$(date +%Y%m%d)} export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_smoke_${PRECISION_TYPE}_MBS${MBS}_GBS${GBS}_PP${PRIMUS_PP}_EP${PRIMUS_EP}} +# Host-side directory for the launcher's aggregated log. Defaults to the +# canonical "output" tree; override when that tree is not writable by the +# invoking user (e.g. it was created by an earlier root/sudo run). +export PRIMUS_OUTPUT_ROOT=${PRIMUS_OUTPUT_ROOT:-output} if [ ! -d "$BACKEND_PATH" ] || [ -z "$(ls -A "$BACKEND_PATH" 2>/dev/null)" ]; then echo "[ERROR] BACKEND_PATH does not exist or is empty: $BACKEND_PATH" @@ -237,7 +241,7 @@ if [ ! -d "$BACKEND_PATH" ] || [ -z "$(ls -A "$BACKEND_PATH" 2>/dev/null)" ]; th exit 1 fi -mkdir -p "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" +mkdir -p "$PRIMUS_OUTPUT_ROOT/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" export PRIMUS_EXIT_FAST=1 @@ -311,4 +315,4 @@ fi --profile_step_end 7 \ --profile_step_start 6 \ --bias_swiglu_fusion "$PRIMUS_BIAS_SWIGLU_FUSION" \ - 2>&1 | tee "output/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/log_node_${NODE_RANK:-0}.txt" + 2>&1 | tee "$PRIMUS_OUTPUT_ROOT/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/log_node_${NODE_RANK:-0}.txt" diff --git a/examples/deepseek-v4/run_deepseek_v4_flash.sh b/examples/deepseek-v4/run_deepseek_v4_flash.sh index 5c5a0daa9..ed77ae0e2 100644 --- a/examples/deepseek-v4/run_deepseek_v4_flash.sh +++ b/examples/deepseek-v4/run_deepseek_v4_flash.sh @@ -7,7 +7,7 @@ export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-256} export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6} export PRIMUS_MOE_FFN_HIDDEN_SIZE=${PRIMUS_MOE_FFN_HIDDEN_SIZE:-2048} export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-512} -export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-'"[0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]"'} +export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-'[0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]'} export MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-1} export NNODES=${NNODES:-8} @@ -18,9 +18,9 @@ if [ "$NNODES" -eq 8 ]; then export PRIMUS_EP=${PRIMUS_EP:-8} export PRIMUS_RECOMPUTE_LAYERS=0 if [ "$MTP_NUM_LAYERS" -eq 1 ]; then - export PRIMUS_PP_LAYOUT='"Et*4|t*5|(t*6|)*5,t*4mL"' + export PRIMUS_PP_LAYOUT='Et*4|t*5|(t*6|)*5,t*4mL' else - export PRIMUS_PP_LAYOUT='"Et*4|t*5|(t*6|)*5,t*4L"' + export PRIMUS_PP_LAYOUT='Et*4|t*5|(t*6|)*5,t*4L' fi elif [ "$NNODES" -eq 4 ]; then export PRIMUS_TP=${PRIMUS_TP:-1} @@ -28,9 +28,9 @@ elif [ "$NNODES" -eq 4 ]; then export PRIMUS_EP=${PRIMUS_EP:-8} export PRIMUS_RECOMPUTE_LAYERS=3 if [ "$MTP_NUM_LAYERS" -eq 1 ]; then - export PRIMUS_PP_LAYOUT='"Et*10|t*11|t*11|t*11mL"' + export PRIMUS_PP_LAYOUT='Et*10|t*11|t*11|t*11mL' else - export PRIMUS_PP_LAYOUT='"Et*10|t*11|t*11|t*11L"' + export PRIMUS_PP_LAYOUT='Et*10|t*11|t*11|t*11L' fi fi diff --git a/runner/primus-cli-slurm-entry.sh b/runner/primus-cli-slurm-entry.sh index dd8662fab..028f79658 100755 --- a/runner/primus-cli-slurm-entry.sh +++ b/runner/primus-cli-slurm-entry.sh @@ -118,18 +118,59 @@ if [[ -z "${SLURM_NODELIST:-}" ]]; then exit 2 fi -# Get all node hostnames (sorted, as needed). Prefer scontrol, which correctly -# expands compressed nodelists (e.g. "node[01-04]"). When scontrol is -# unavailable -- CI containers / dev VMs without the Slurm client tools -- fall -# back to parsing SLURM_NODELIST directly. This mirrors the scontrol-optional -# handling in primus-cli-direct.sh so every launcher behaves consistently -# off-cluster (range expansion is skipped in the fallback, which is fine for the -# single-host / comma-list forms used in CI and single-node runs). +# Pure-bash expander for a SLURM/Spur nodelist. Handles comma-separated lists +# and bracket forms like "prefix[01-04,06]" / "prefix[030,058]" (optional +# suffix). Used as the portable fallback below. +_expand_nodelist() { + local nl="$1" + local seg="" depth=0 i ch + local -a segs=() + for (( i=0; i<${#nl}; i++ )); do + ch="${nl:i:1}" + if [[ "$ch" == "[" ]]; then depth=$((depth+1)); seg+="$ch" + elif [[ "$ch" == "]" ]]; then depth=$((depth-1)); seg+="$ch" + elif [[ "$ch" == "," && $depth -eq 0 ]]; then segs+=("$seg"); seg="" + else seg+="$ch"; fi + done + [[ -n "$seg" ]] && segs+=("$seg") + + local s pre body suf tok a b width n + for s in "${segs[@]}"; do + [[ -n "$s" ]] || continue + if [[ "$s" == *"["* ]]; then + pre="${s%%[*}" + body="${s#*[}"; body="${body%%]*}" + suf="${s#*]}" + local OLD_IFS="$IFS"; IFS=',' + for tok in $body; do + if [[ "$tok" == *-* ]]; then + a="${tok%%-*}"; b="${tok##*-}"; width=${#a} + for (( n=10#$a; n<=10#$b; n++ )); do + printf "%s%0*d%s\n" "$pre" "$width" "$n" "$suf" + done + else + printf "%s%s%s\n" "$pre" "$tok" "$suf" + fi + done + IFS="$OLD_IFS" + else + printf "%s\n" "$s" + fi + done +} + +# Get all node hostnames. Prefer `scontrol show hostnames`, which correctly +# expands compressed nodelists on stock Slurm. Some Slurm-compatible schedulers +# (e.g. Spur) lack that subcommand, and CI containers / dev VMs may lack the +# client entirely -- in both cases fall back to the pure-bash expander above so +# every launcher behaves consistently, including bracket-range expansion. +NODE_ARRAY=() if command -v scontrol >/dev/null 2>&1; then - readarray -t NODE_ARRAY < <(scontrol show hostnames "$SLURM_NODELIST") -else - LOG_WARN "[slurm-entry] scontrol not found; parsing SLURM_NODELIST without range expansion" - readarray -t NODE_ARRAY < <(tr ',' '\n' <<< "$SLURM_NODELIST") + readarray -t NODE_ARRAY < <(scontrol show hostnames "$SLURM_NODELIST" 2>/dev/null || true) +fi +if [[ ${#NODE_ARRAY[@]} -eq 0 ]]; then + LOG_WARN "[slurm-entry] 'scontrol show hostnames' unavailable; expanding SLURM_NODELIST locally" + readarray -t NODE_ARRAY < <(_expand_nodelist "$SLURM_NODELIST") fi SLURM_MASTER_ADDR="${NODE_ARRAY[0]:-}" if [[ -z "$SLURM_MASTER_ADDR" ]]; then @@ -139,7 +180,11 @@ fi if [[ -z "${MASTER_ADDR:-}" ]]; then MASTER_ADDR="$SLURM_MASTER_ADDR" -elif [[ "$MASTER_ADDR" != "$SLURM_MASTER_ADDR" ]]; then +elif [[ "${MASTER_ADDR%%.*}" != "${SLURM_MASTER_ADDR%%.*}" ]]; then + # Compare only the leading hostname label so a scheduler-provided FQDN + # (e.g. Spur exports MASTER_ADDR=node.crusoe.amd.com) still matches the + # short name produced by nodelist expansion. A genuine mismatch (a + # different node entirely) is still caught. LOG_ERROR "[slurm-entry] MASTER_ADDR must match the first host in SLURM_NODELIST." LOG_ERROR "[slurm-entry] MASTER_ADDR=$MASTER_ADDR, expected=$SLURM_MASTER_ADDR" exit 2 From f96b1fbbf5fe83f2035422cb1565562501ba37d7 Mon Sep 17 00:00:00 2001 From: WangLingxun Date: Fri, 17 Jul 2026 17:54:48 +0800 Subject: [PATCH 042/127] feat(torchtitan): upgrade to v0.2.2 for torch 2.12 + GPT-OSS support (#871) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Upgrade the TorchTitan backend to upstream **v0.2.2** (`73a0e697`), which pins **torch 2.12** (`torch-2.12.0.dev+`, torchao 0.17), for the Primus 26.4 release. Adapt the Primus integration layer to the v0.2.2 `Attention` interface and add first-class **GPT-OSS** support (model + BF16/FP8 example configs + Primus-Turbo sink attention). ### Commits - `bump submodule to v0.2.2 for torch 2.12` — `third_party/torchtitan` `5fb7cc2e → 73a0e697`. - `adapt model mirrors to v0.2.2 Attention interface` — mirror `Attention.forward` gains `positions` (llama3/llama4/qwen3/deepseek_v3); deepseek classic `attn_type="flex"` + node-limited routing moved into `MoEArgs`. - `add GPT-OSS model support (v0.2.2) + example configs` — `gpt_oss_{20b,120b}{,-fp8}.yaml` model configs + 12 example pretrain configs (`{20B,120B}×{BF16,FP8}×{MI300X,MI325X,MI355X}`, turbo on by default, parallelism/batch sized per model & device). - `Primus-Turbo sink attention for GPT-OSS` — GPT-OSS mirror `Attention`/`TransformerBlock` routing through primus_turbo functional `flash_attn_func(sink=, window_size=)` (learned per-head sinks + even-layer sliding window); enabled via `gptoss_sink_attention` setup patch, gated on `enable_primus_turbo + use_turbo_attention + model.name==gpt_oss` (default FlexAttention path unaffected). - `add gpt_oss config, sink-attention unit tests + e2e`. ## Notes - Known upstream/torch-2.12 issue (not introduced here): qwen3 tied-embedding + FSDP2 raises a shared-parameter error; llama3/gpt_oss (untied) unaffected. --------- Co-authored-by: Xiaoming-AMD Co-authored-by: wenxie-amd --- .github/workflows/ci.yaml | 28 ++++ .../MI300X/gpt_oss_120B-BF16-pretrain.yaml | 57 +++++++ .../MI300X/gpt_oss_120B-FP8-pretrain.yaml | 58 +++++++ .../MI300X/gpt_oss_20B-BF16-pretrain.yaml | 58 +++++++ .../MI300X/gpt_oss_20B-FP8-pretrain.yaml | 58 +++++++ .../MI325X/gpt_oss_120B-BF16-pretrain.yaml | 57 +++++++ .../MI325X/gpt_oss_120B-FP8-pretrain.yaml | 58 +++++++ .../MI325X/gpt_oss_20B-BF16-pretrain.yaml | 58 +++++++ .../MI325X/gpt_oss_20B-FP8-pretrain.yaml | 58 +++++++ .../MI355X/gpt_oss_120B-BF16-pretrain.yaml | 57 +++++++ .../MI355X/gpt_oss_120B-FP8-pretrain.yaml | 58 +++++++ .../MI355X/gpt_oss_20B-BF16-pretrain.yaml | 58 +++++++ .../MI355X/gpt_oss_20B-FP8-pretrain.yaml | 58 +++++++ examples/torchtitan/prepare.py | 12 +- primus/_thirdparty.lock | 2 +- .../torchtitan/models/deepseek_v3/__init__.py | 14 +- .../models/deepseek_v3/model/model.py | 13 +- .../torchtitan/models/gpt_oss/__init__.py | 7 + .../models/gpt_oss/model/__init__.py | 5 + .../torchtitan/models/gpt_oss/model/model.py | 105 ++++++++++++ .../torchtitan/models/llama3/model/model.py | 10 +- .../torchtitan/models/llama4/model/model.py | 10 +- .../torchtitan/models/qwen3/model/model.py | 13 +- .../torchtitan/patches/turbo/__init__.py | 1 + .../turbo/gptoss_sink_attention_patches.py | 70 ++++++++ .../patches/turbo/moe_grouped_mm_patches.py | 27 +++- .../models/torchtitan/gpt_oss_120b-fp8.yaml | 14 ++ .../models/torchtitan/gpt_oss_120b.yaml | 12 ++ .../models/torchtitan/gpt_oss_20b-fp8.yaml | 14 ++ .../models/torchtitan/gpt_oss_20b.yaml | 12 ++ .../train/pretrain/torchtitan/prepare.py | 12 +- tests/conftest.py | 9 ++ tests/trainer/test_torchtitan_trainer.py | 31 ++++ .../torchtitan/test_gpt_oss_configs.py | 108 +++++++++++++ .../torchtitan/test_gpt_oss_sink_attention.py | 150 ++++++++++++++++++ .../torchtitan/test_moe_grouped_mm_patch.py | 103 ++++++++++++ tests/unit_tests/conftest.py | 9 ++ third_party/torchtitan | 2 +- 38 files changed, 1444 insertions(+), 42 deletions(-) create mode 100644 examples/torchtitan/configs/MI300X/gpt_oss_120B-BF16-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI300X/gpt_oss_120B-FP8-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI300X/gpt_oss_20B-BF16-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI300X/gpt_oss_20B-FP8-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI325X/gpt_oss_120B-BF16-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI325X/gpt_oss_120B-FP8-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI325X/gpt_oss_20B-BF16-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI355X/gpt_oss_120B-BF16-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI355X/gpt_oss_120B-FP8-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml create mode 100644 examples/torchtitan/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml create mode 100644 primus/backends/torchtitan/models/gpt_oss/__init__.py create mode 100644 primus/backends/torchtitan/models/gpt_oss/model/__init__.py create mode 100644 primus/backends/torchtitan/models/gpt_oss/model/model.py create mode 100644 primus/backends/torchtitan/patches/turbo/gptoss_sink_attention_patches.py create mode 100644 primus/configs/models/torchtitan/gpt_oss_120b-fp8.yaml create mode 100644 primus/configs/models/torchtitan/gpt_oss_120b.yaml create mode 100644 primus/configs/models/torchtitan/gpt_oss_20b-fp8.yaml create mode 100644 primus/configs/models/torchtitan/gpt_oss_20b.yaml create mode 100644 tests/unit_tests/backends/torchtitan/test_gpt_oss_configs.py create mode 100644 tests/unit_tests/backends/torchtitan/test_gpt_oss_sink_attention.py create mode 100644 tests/unit_tests/backends/torchtitan/test_moe_grouped_mm_patch.py diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index f65542997..002706592 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -285,6 +285,34 @@ jobs: - name: Install Primus run: | pip install -r requirements.txt + # v26.3-only: drop this whole step on v26.4. Unit tests import torchtitan + # from the submodule source via conftest sys.path injection (v0.2.2 has + # PEP-420 namespace packages that need the source root on sys.path). This + # just removes any stale torchtitan on the persistent runner that would + # shadow it (old install or leftover editable finder / .pth / dist-info). + - name: Purge stale TorchTitan install (v26.3-only, drop on v26.4) + run: | + pip uninstall -y torchtitan || true + python - <<'PY' + import glob, os, shutil, site + dirs = set(site.getsitepackages() + [site.getusersitepackages()]) + patterns = ( + "__editable__*torchtitan*", + "__editable__.torchtitan*.pth", + "torchtitan*.pth", + "torchtitan*.egg-link", + "torchtitan*.dist-info", + "torchtitan*.egg-info", + ) + for d in dirs: + for pat in patterns: + for f in glob.glob(os.path.join(d, pat)): + try: + shutil.rmtree(f) if os.path.isdir(f) else os.remove(f) + print("removed", f) + except OSError as e: + print("skip", f, e) + PY - name: Install fixed origami run: | # rocm/primus:v26.3 bundles origami 0.1.0, whose rank_configs() raises diff --git a/examples/torchtitan/configs/MI300X/gpt_oss_120B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI300X/gpt_oss_120B-BF16-pretrain.yaml new file mode 100644 index 000000000..3c923f249 --- /dev/null +++ b/examples/torchtitan/configs/MI300X/gpt_oss_120B-BF16-pretrain.yaml @@ -0,0 +1,57 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-BF16-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_120b.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 1.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 20 + + training: + local_batch_size: 1 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "always" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "full" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch. + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: false + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI300X/gpt_oss_120B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI300X/gpt_oss_120B-FP8-pretrain.yaml new file mode 100644 index 000000000..ae3fe1c10 --- /dev/null +++ b/examples/torchtitan/configs/MI300X/gpt_oss_120B-FP8-pretrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-FP8-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_120b-fp8.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 1.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 20 + + training: + local_batch_size: 1 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "always" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "full" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch. + # FP8 dense linears via the quantize.linear.float8 converter. + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: true + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI300X/gpt_oss_20B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI300X/gpt_oss_20B-BF16-pretrain.yaml new file mode 100644 index 000000000..2e87bd229 --- /dev/null +++ b/examples/torchtitan/configs/MI300X/gpt_oss_20B-BF16-pretrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-BF16-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_20b.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 2.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 10 + + training: + local_batch_size: 4 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "default" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "selective" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention is installed by the + # gptoss_sink_attention setup patch. BF16 (no fp8/mx). + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: false + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI300X/gpt_oss_20B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI300X/gpt_oss_20B-FP8-pretrain.yaml new file mode 100644 index 000000000..65593906e --- /dev/null +++ b/examples/torchtitan/configs/MI300X/gpt_oss_20B-FP8-pretrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-FP8-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_20b-fp8.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 2.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 10 + + training: + local_batch_size: 6 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "default" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "selective" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch. + # FP8 dense linears via the quantize.linear.float8 converter. + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: true + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI325X/gpt_oss_120B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI325X/gpt_oss_120B-BF16-pretrain.yaml new file mode 100644 index 000000000..24810a1f4 --- /dev/null +++ b/examples/torchtitan/configs/MI325X/gpt_oss_120B-BF16-pretrain.yaml @@ -0,0 +1,57 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-BF16-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_120b.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 1.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 20 + + training: + local_batch_size: 2 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "always" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "full" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch. + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: false + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI325X/gpt_oss_120B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI325X/gpt_oss_120B-FP8-pretrain.yaml new file mode 100644 index 000000000..ea03d4a13 --- /dev/null +++ b/examples/torchtitan/configs/MI325X/gpt_oss_120B-FP8-pretrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-FP8-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_120b-fp8.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 1.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 20 + + training: + local_batch_size: 2 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "always" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "full" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch. + # FP8 dense linears via the quantize.linear.float8 converter. + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: true + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI325X/gpt_oss_20B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI325X/gpt_oss_20B-BF16-pretrain.yaml new file mode 100644 index 000000000..c071aa0c3 --- /dev/null +++ b/examples/torchtitan/configs/MI325X/gpt_oss_20B-BF16-pretrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-BF16-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_20b.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 2.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 10 + + training: + local_batch_size: 6 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "default" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "selective" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention is installed by the + # gptoss_sink_attention setup patch. BF16 (no fp8/mx). + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: false + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml new file mode 100644 index 000000000..ade0bd8e3 --- /dev/null +++ b/examples/torchtitan/configs/MI325X/gpt_oss_20B-FP8-pretrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-FP8-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_20b-fp8.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 2.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 10 + + training: + local_batch_size: 8 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "default" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "selective" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch. + # FP8 dense linears via the quantize.linear.float8 converter. + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: true + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI355X/gpt_oss_120B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI355X/gpt_oss_120B-BF16-pretrain.yaml new file mode 100644 index 000000000..24810a1f4 --- /dev/null +++ b/examples/torchtitan/configs/MI355X/gpt_oss_120B-BF16-pretrain.yaml @@ -0,0 +1,57 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-BF16-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_120b.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 1.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 20 + + training: + local_batch_size: 2 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "always" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "full" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch. + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: false + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI355X/gpt_oss_120B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI355X/gpt_oss_120B-FP8-pretrain.yaml new file mode 100644 index 000000000..ea03d4a13 --- /dev/null +++ b/examples/torchtitan/configs/MI355X/gpt_oss_120B-FP8-pretrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_120B-FP8-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_120b-fp8.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 1.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 20 + + training: + local_batch_size: 2 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "always" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "full" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch. + # FP8 dense linears via the quantize.linear.float8 converter. + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: true + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml b/examples/torchtitan/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml new file mode 100644 index 000000000..e0db78d38 --- /dev/null +++ b/examples/torchtitan/configs/MI355X/gpt_oss_20B-BF16-pretrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-BF16-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_20b.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 2.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 10 + + training: + local_batch_size: 8 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "default" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "none" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention is installed by the + # gptoss_sink_attention setup patch. BF16 (no fp8/mx). + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: false + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml b/examples/torchtitan/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml new file mode 100644 index 000000000..c271776db --- /dev/null +++ b/examples/torchtitan/configs/MI355X/gpt_oss_20B-FP8-pretrain.yaml @@ -0,0 +1,58 @@ +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:gpt_oss_20B-FP8-pretrain} +workspace: ./output + +modules: + pre_trainer: + framework: torchtitan + config: pre_trainer.yaml + + model: gpt_oss_20b-fp8.yaml + overrides: + metrics: + log_freq: 1 + enable_wandb: false + + optimizer: + name: "AdamW" + lr: 2.2e-4 + eps: 1.0e-8 + + lr_scheduler: + warmup_steps: 10 + + training: + local_batch_size: 10 + seq_len: 4096 + max_norm: 1.0 # grad norm clipping + steps: 50 + + parallelism: + data_parallel_replicate_degree: 1 + data_parallel_shard_degree: -1 + fsdp_reshard_after_forward: "default" # default / never / always + tensor_parallel_degree: 1 + pipeline_parallel_degree: 1 + expert_parallel_degree: 8 + expert_tensor_parallel_degree: 1 + + activation_checkpoint: + mode: "selective" # ["none", "selective", "full"] + selective_ac_option: "op" + + compile: + enable: true + components: ["model", "loss"] + + # Primus-Turbo on; GPT-OSS sink attention via gptoss_sink_attention patch. + # FP8 dense linears via the quantize.linear.float8 converter. + primus_turbo: + enable_primus_turbo: true + use_turbo_attention: true + enable_attention_float8: false + use_turbo_float8_linear: true + use_turbo_mx_linear: false + use_moe_fp8: false + use_turbo_grouped_mm: false + use_classic_attention: false diff --git a/examples/torchtitan/prepare.py b/examples/torchtitan/prepare.py index aeff9043a..1777ab5ff 100644 --- a/examples/torchtitan/prepare.py +++ b/examples/torchtitan/prepare.py @@ -43,8 +43,16 @@ def parse_args(): def pip_install_editable(path: Path, name: str): - log_info(f"Installing {name} in editable mode via pip (path: {path})") - ret = subprocess.run(["pip", "install", "-e", ".", "-q"], cwd=path) + # TorchTitan v0.2.2 has PEP-420 namespace subpackages (e.g. torchtitan/tools, + # no __init__.py) that the default editable finder cannot resolve. Use compat + # mode (writes a .pth with the source root on sys.path) after uninstalling any + # stale torchtitan, so the training subprocess can import them. + log_info(f"Installing {name} in editable (compat) mode via pip (path: {path})") + subprocess.run(["pip", "uninstall", "-y", name.lower()], cwd=path) + ret = subprocess.run( + ["pip", "install", "-e", ".", "--config-settings", "editable_mode=compat", "-q"], + cwd=path, + ) if ret.returncode != 0: log_error_and_exit(f"Failed to install {name} via pip.") diff --git a/primus/_thirdparty.lock b/primus/_thirdparty.lock index 250b6dfb7..72f6be2f5 100644 --- a/primus/_thirdparty.lock +++ b/primus/_thirdparty.lock @@ -4,7 +4,7 @@ "name": "torchtitan", "path": "third_party/torchtitan", "url": "https://github.com/pytorch/torchtitan.git", - "commit": "5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021" + "commit": "73a0e6979dd10b6b1904098eb3c8f62c18ab87ce" }, { "name": "Megatron-LM", diff --git a/primus/backends/torchtitan/models/deepseek_v3/__init__.py b/primus/backends/torchtitan/models/deepseek_v3/__init__.py index 6ec1c8600..6ddf45f9e 100644 --- a/primus/backends/torchtitan/models/deepseek_v3/__init__.py +++ b/primus/backends/torchtitan/models/deepseek_v3/__init__.py @@ -33,7 +33,7 @@ qk_rope_head_dim=64, v_head_dim=128, mscale=0.70, - use_flex_attn=True, + attn_type="flex", attn_mask_type="block_causal", q_head=16, n_kv_heads=16, @@ -51,19 +51,19 @@ num_experts=160, num_shared_experts=2, top_k=6, + num_expert_groups=8, + num_limited_groups=3, score_func="softmax", route_norm=True, route_scale=16.0, score_before_experts=False, ), - n_expert_groups=8, - n_limited_groups=3, q_lora_rank=1536, kv_lora_rank=512, qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128, - use_flex_attn=True, + attn_type="flex", attn_mask_type="block_causal", q_head=40, n_kv_heads=8, @@ -81,19 +81,19 @@ num_experts=256, num_shared_experts=1, top_k=8, + num_expert_groups=8, + num_limited_groups=4, score_func="sigmoid", route_norm=True, route_scale=2.5, score_before_experts=False, ), - n_expert_groups=8, - n_limited_groups=4, q_lora_rank=1536, kv_lora_rank=512, qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128, - use_flex_attn=True, + attn_type="flex", attn_mask_type="block_causal", q_head=56, n_kv_heads=8, diff --git a/primus/backends/torchtitan/models/deepseek_v3/model/model.py b/primus/backends/torchtitan/models/deepseek_v3/model/model.py index ec587fe67..bb8797b30 100644 --- a/primus/backends/torchtitan/models/deepseek_v3/model/model.py +++ b/primus/backends/torchtitan/models/deepseek_v3/model/model.py @@ -26,6 +26,7 @@ def forward( x: torch.Tensor, freqs_cis: torch.Tensor, attention_masks: AttentionMasksType | None, + positions: torch.Tensor | None = None, ): """ Forward pass for the Multi-Head Latent Attention (MLA) Layer. @@ -33,6 +34,7 @@ def forward( Args: x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim). freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings. + positions (torch.Tensor | None): Position indices used to access/shuffle the RoPE cache. Returns: torch.Tensor: Output tensor with the same shape as the input. @@ -50,14 +52,14 @@ def forward( # the above linear ops. q = q.view(bsz, seqlen, -1, self.qk_head_dim) q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - q_pe = apply_rotary_emb(q_pe, freqs_cis) + q_pe = apply_rotary_emb(q_pe, freqs_cis, positions) q = torch.cat([q_nope, q_pe], dim=-1) # (bsz, seqlen, n_heads, qk_head_dim) # Key-value projection kv = self.wkv_a(x) # (bsz, seqlen, kv_lora_rank + qk_rope_head_dim) kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis) # (bsz, seqlen, 1, qk_rope_head_dim) + k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis, positions) # (bsz, seqlen, 1, qk_rope_head_dim) kv = self.wkv_b(self.kv_norm(kv)) # (bsz, seqlen, n_heads * (qk_nope_head_dim + v_head_dim)) kv = kv.view(bsz, seqlen, -1, self.qk_nope_head_dim + self.v_head_dim) @@ -99,7 +101,9 @@ def __init__(self, deepseek_args: DeepSeekV3ClassicModelArgs): self.dim = deepseek_args.dim self.head_dim = deepseek_args.head_dim - self.use_flex_attn = deepseek_args.use_flex_attn + # Upstream v0.2.2 Attention.__init__ dispatches on ``attn_type`` + # (the legacy ``use_flex_attn`` flag was removed). + self.attn_type = deepseek_args.attn_type # Initialize the parent class with the mock args super().__init__( @@ -112,10 +116,11 @@ def forward( x: torch.Tensor, freqs_cis: torch.Tensor, attention_masks: AttentionMasksType | None, + positions: torch.Tensor | None = None, ): # Always use llama4-style freqs_cis for this attention, regardless of input seqlen = x.shape[1] freqs_llama4 = llama4_precompute_freqs_cis(self.head_dim, seqlen, self.rope_theta) # Ensure freqs are on the same device as activations freqs_llama4 = freqs_llama4.to(x.device, dtype=x.dtype) - return super().forward(x, freqs_llama4, None) + return super().forward(x, freqs_llama4, None, positions) diff --git a/primus/backends/torchtitan/models/gpt_oss/__init__.py b/primus/backends/torchtitan/models/gpt_oss/__init__.py new file mode 100644 index 000000000..095aa107e --- /dev/null +++ b/primus/backends/torchtitan/models/gpt_oss/__init__.py @@ -0,0 +1,7 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Primus TorchTitan GPT-OSS backend extensions (Primus-Turbo sink attention).""" diff --git a/primus/backends/torchtitan/models/gpt_oss/model/__init__.py b/primus/backends/torchtitan/models/gpt_oss/model/__init__.py new file mode 100644 index 000000000..771e1b42f --- /dev/null +++ b/primus/backends/torchtitan/models/gpt_oss/model/__init__.py @@ -0,0 +1,5 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### diff --git a/primus/backends/torchtitan/models/gpt_oss/model/model.py b/primus/backends/torchtitan/models/gpt_oss/model/model.py new file mode 100644 index 000000000..b032e82d5 --- /dev/null +++ b/primus/backends/torchtitan/models/gpt_oss/model/model.py @@ -0,0 +1,105 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Primus-Turbo sink attention mirror for TorchTitan GPT-OSS. + +Upstream GPT-OSS attention (``torchtitan.models.gpt_oss.model.model.Attention``) +uses ``FlexAttentionWrapper`` with: + * learnable per-head attention sinks (``self.sinks``), applied as a post-hoc + ``sigmoid(lse - sinks)`` rescaling of the FlexAttention output, and + * a sliding-window mask applied on even layers (``layer_id % 2 == 0``). + +This mirror keeps the exact same parameters (``wq/wk/wv/wo`` incl. biases and the +learnable ``self.sinks``) but routes attention through the Primus-Turbo +*functional* kernel ``primus_turbo.pytorch.ops.flash_attn_func``, which natively +supports both: + * ``sink=`` (per-head learnable sinks; automatically dispatches to the Triton + backend), matching the Megatron ``PrimusTurboAttention`` semantics, and + * ``window_size=(left, 0)`` for the sliding window. + +Because ``self.sinks`` is inherited from the upstream module (not re-created), the +learned sink weights are preserved on checkpoint load/save. The mirror is only +installed when ``primus_turbo.enable_primus_turbo`` and +``primus_turbo.use_turbo_attention`` are both set for a ``gpt_oss`` run (see +``patches/turbo/gptoss_sink_attention_patches.py``); the default GPT-OSS path +keeps upstream FlexAttention untouched. +""" + +import torch +from torchtitan.models.gpt_oss.model.model import Attention as TTGptOssAttention +from torchtitan.models.gpt_oss.model.model import TransformerBlock as TTGptOssBlock +from torchtitan.models.gpt_oss.model.model import apply_rotary_emb + +# Sentinel meaning "no sliding window" for primus_turbo's flash_attn_func. +_FULL_WINDOW = (-1, -1) + + +class Attention(TTGptOssAttention): + """GPT-OSS attention backed by Primus-Turbo ``flash_attn_func`` (sink-aware).""" + + def __init__(self, model_args): + super().__init__(model_args) + # Per-layer sliding window is injected by the mirror TransformerBlock + # before each forward; default to full (causal, no window). + self.sliding_window_size = model_args.sliding_window_size + self._turbo_window = _FULL_WINDOW + + def forward( + self, + x: torch.Tensor, + rope_cache: torch.Tensor, + attention_masks=None, # noqa: ARG002 - flex masks unused on the turbo path + ): + bsz, seqlen, _ = x.size() + hidden_shape = (bsz, seqlen, -1, self.head_dim) + + q = self.wq(x).view(hidden_shape) + k = self.wk(x).view(hidden_shape) + v = self.wv(x).view(hidden_shape) + + # RoPE is applied on the (b, s, h, d) layout, exactly like upstream. + q, k = apply_rotary_emb(q, k, rope_cache) + + # primus_turbo.flash_attn_func consumes the (b, s, h, d) layout directly + # and handles GQA (n_heads > n_kv_heads) internally, so we do NOT + # transpose to (b, h, s, d) and do NOT repeat_kv. The learnable sinks are + # applied inside the kernel (Triton backend), so the upstream post-hoc + # sigmoid(lse - sinks) rescaling is not needed here. + import primus_turbo.pytorch as turbo + + output = turbo.ops.flash_attn_func( + q, + k, + v, + softmax_scale=self.softmax_scale, + causal=True, + window_size=self._turbo_window, + sink=self.sinks.to(q.dtype), + ) + + output = output.reshape(bsz, seqlen, -1).contiguous() + return self.wo(output) + + +class TransformerBlock(TTGptOssBlock): + """GPT-OSS block that selects the per-layer sliding window for turbo sink attn.""" + + def forward( + self, + x: torch.Tensor, + rope_cache: torch.Tensor, + attention_masks=None, # noqa: ARG002 - flex masks unused on the turbo path + ): + # gpt-oss pattern: even layers use sliding-window attention. + if self.use_sliding_attention: + self.attention._turbo_window = (self.attention.sliding_window_size, 0) + else: + self.attention._turbo_window = _FULL_WINDOW + + x = x + self.attention(self.attention_norm(x), rope_cache, None) + x = x + self.moe(self.ffn_norm(x)) + return x diff --git a/primus/backends/torchtitan/models/llama3/model/model.py b/primus/backends/torchtitan/models/llama3/model/model.py index daebde2fd..954b1bf51 100644 --- a/primus/backends/torchtitan/models/llama3/model/model.py +++ b/primus/backends/torchtitan/models/llama3/model/model.py @@ -18,6 +18,7 @@ def forward( x: torch.Tensor, freqs_cis: torch.Tensor, attention_masks: AttentionMasksType | None, + positions: torch.Tensor | None = None, ): bs, seqlen, _ = x.shape xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) @@ -29,12 +30,11 @@ def forward( xk = xk.view(bs, seqlen, -1, self.head_dim) xv = xv.view(bs, seqlen, -1, self.head_dim) - xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis) - - # repeat k/v heads if n_kv_heads < n_heads - # xk = repeat_kv(xk, self.n_rep) # (bs, seqlen, n_local_heads, head_dim) - # xv = repeat_kv(xv, self.n_rep) # (bs, seqlen, n_local_heads, head_dim) + xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis, positions=positions) + # Primus-Turbo path: inner_attention is replaced by TurboAttention, which + # consumes the (bs, seqlen, n_heads, head_dim) layout and handles GQA / + # causal masking internally, so we skip repeat_kv and the transpose. output = self.inner_attention(xq, xk, xv) output = output.contiguous().view(bs, seqlen, -1) diff --git a/primus/backends/torchtitan/models/llama4/model/model.py b/primus/backends/torchtitan/models/llama4/model/model.py index 7807fcbd6..5e80aeca6 100644 --- a/primus/backends/torchtitan/models/llama4/model/model.py +++ b/primus/backends/torchtitan/models/llama4/model/model.py @@ -18,6 +18,7 @@ def forward( x: torch.Tensor, freqs_cis: torch.Tensor, attention_masks: AttentionMasksType | None, + positions: torch.Tensor | None = None, ): bs, seqlen, _ = x.shape xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) @@ -30,12 +31,11 @@ def forward( xv = xv.view(bs, seqlen, -1, self.head_dim) if self.use_rope: - xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis) - - # repeat k/v heads if n_kv_heads < n_heads - # xk = repeat_kv(xk, self.n_rep) # (bs, seqlen, n_local_heads, head_dim) - # xv = repeat_kv(xv, self.n_rep) # (bs, seqlen, n_local_heads, head_dim) + xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis, positions=positions) + # Primus-Turbo path: inner_attention is replaced by TurboAttention, which + # consumes the (bs, seqlen, n_heads, head_dim) layout and handles GQA / + # causal masking internally, so we skip repeat_kv and the transpose. output = self.inner_attention(xq, xk, xv) output = output.contiguous().view(bs, seqlen, -1) diff --git a/primus/backends/torchtitan/models/qwen3/model/model.py b/primus/backends/torchtitan/models/qwen3/model/model.py index 632a33095..35f88bfda 100644 --- a/primus/backends/torchtitan/models/qwen3/model/model.py +++ b/primus/backends/torchtitan/models/qwen3/model/model.py @@ -18,6 +18,7 @@ def forward( x: torch.Tensor, rope_cache: torch.Tensor, attention_masks: AttentionMasksType | None, + positions: torch.Tensor | None = None, ): bs, seqlen, _ = x.shape xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) @@ -34,14 +35,12 @@ def forward( if self.k_norm is not None: xk = self.k_norm(xk) - xq, xk = apply_rotary_emb(xq, xk, rope_cache) + xq, xk = apply_rotary_emb(xq, xk, rope_cache, positions) - if self.use_flex_attn: - assert isinstance(attention_masks, BlockMask), attention_masks - output = self.inner_attention(xq, xk, xv, block_mask=attention_masks) - else: - assert attention_masks is None - output = self.inner_attention(xq, xk, xv) + # Primus-Turbo path: inner_attention is replaced by TurboAttention, which + # consumes the (bs, seqlen, n_heads, head_dim) layout and handles GQA / + # causal masking internally, so we skip the transpose and block_mask. + output = self.inner_attention(xq, xk, xv) output = output.contiguous().view(bs, seqlen, -1) return self.wo(output) diff --git a/primus/backends/torchtitan/patches/turbo/__init__.py b/primus/backends/torchtitan/patches/turbo/__init__.py index 5677368ef..dfc2b79bc 100644 --- a/primus/backends/torchtitan/patches/turbo/__init__.py +++ b/primus/backends/torchtitan/patches/turbo/__init__.py @@ -17,6 +17,7 @@ attention_patches, deepseek_v3_classic_attention_patches, fp8_linear_patches, + gptoss_sink_attention_patches, moe_grouped_mm_patches, mx_linear_patches, ) diff --git a/primus/backends/torchtitan/patches/turbo/gptoss_sink_attention_patches.py b/primus/backends/torchtitan/patches/turbo/gptoss_sink_attention_patches.py new file mode 100644 index 000000000..d90cee727 --- /dev/null +++ b/primus/backends/torchtitan/patches/turbo/gptoss_sink_attention_patches.py @@ -0,0 +1,70 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +""" +TorchTitan GPT-OSS Primus-Turbo Sink Attention Patch +==================================================== + +GPT-OSS uses FlexAttention with learnable per-head attention sinks and a +sliding-window mask on even layers. The module-form ``TurboAttention`` used by +the other models' turbo path does NOT expose ``sink`` / ``window_size``, so it +cannot model GPT-OSS attention. However, the Primus-Turbo *functional* kernel +``primus_turbo.pytorch.ops.flash_attn_func`` does support both (this is the same +API the Megatron ``PrimusTurboAttention`` uses for GPT-OSS sink attention). + +This setup patch swaps the upstream GPT-OSS ``Attention`` and ``TransformerBlock`` +classes for Primus mirrors that route attention through ``flash_attn_func`` with +``sink=self.sinks`` and the per-layer ``window_size``. It is gated on +``primus_turbo.enable_primus_turbo`` + ``primus_turbo.use_turbo_attention`` and +only applies to ``gpt_oss`` runs, so existing GPT-OSS FlexAttention configs +(turbo off) are unaffected. +""" + +from primus.core.patches import PatchContext, get_param, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +def _gptoss_turbo_enabled(ctx: PatchContext) -> bool: + return ( + get_param(ctx, "model.name", None) == "gpt_oss" + and get_param(ctx, "primus_turbo.enable_primus_turbo", False) + and get_param(ctx, "primus_turbo.use_turbo_attention", False) + ) + + +@register_patch( + "torchtitan.primus_turbo.gptoss_sink_attention", + backend="torchtitan", + phase="setup", + description="Use Primus-Turbo functional flash_attn_func (sink + sliding window) for GPT-OSS", + condition=_gptoss_turbo_enabled, +) +def patch_gptoss_sink_attention(ctx: PatchContext) -> None: + """Install the Primus-Turbo sink-attention mirror for GPT-OSS.""" + log_rank_0( + "[Patch:torchtitan.primus_turbo.gptoss_sink_attention] " + "Enabling Primus-Turbo sink attention (flash_attn_func) for GPT-OSS...", + ) + + import torchtitan.models.gpt_oss.model.model as gptoss_model_mod + + from primus.backends.torchtitan.models.gpt_oss.model.model import ( + Attention as GptOssTurboAttention, + ) + from primus.backends.torchtitan.models.gpt_oss.model.model import ( + TransformerBlock as GptOssTurboBlock, + ) + + # Replace the module-level classes so GptOssModel builds the mirror instances + # (the mirror TransformerBlock also constructs the mirror Attention, and both + # resolve these names from this module's globals at build time). + gptoss_model_mod.Attention = GptOssTurboAttention + gptoss_model_mod.TransformerBlock = GptOssTurboBlock + + log_rank_0( + "[Patch:torchtitan.primus_turbo.gptoss_sink_attention] " + "GPT-OSS Primus-Turbo sink attention successfully installed.", + ) diff --git a/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py b/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py index fcd2569bf..17cf70b8d 100644 --- a/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py +++ b/primus/backends/torchtitan/patches/turbo/moe_grouped_mm_patches.py @@ -19,11 +19,13 @@ delegate to Primus' grouped_mm implementation with ``use_fp8`` bound. """ -import functools - from primus.core.patches import PatchContext, get_param, register_patch from primus.core.utils.module_utils import log_rank_0 +# Upstream apply_compile skips torch.compile when the grouped_mm __qualname__ +# already contains this string. We reuse it to opt out of compile. +_ALREADY_PATCHED_SENTINEL = "_run_experts_grouped_mm_dynamic" + @register_patch( "torchtitan.primus_turbo.moe_grouped_mm", @@ -43,17 +45,26 @@ def patch_torchtitan_moe(ctx: PatchContext) -> None: from primus.backends.torchtitan.models.moe.moe import _run_experts_grouped_mm - # Get MoE FP8 configuration and create a partial function + # Get MoE FP8 configuration and bind it onto the replacement function. use_moe_fp8 = get_param(ctx, "primus_turbo.use_moe_fp8", False) log_rank_0( "[Patch:torchtitan.primus_turbo.moe_grouped_mm] " f"Set MoE FP8 mode: {use_moe_fp8}", ) - # Patch the grouped_mm function with use_fp8 parameter pre-set - torchtitan.models.moe.moe._run_experts_grouped_mm = functools.partial( - _run_experts_grouped_mm, - use_fp8=use_moe_fp8, - ) + # The turbo grouped_mm kernels are torch.compiler.disable'd, so letting + # upstream compile our wrapper under fullgraph traces into them and hard-errors + # (dynamo gb0098). Decorating the wrapper with torch.compiler.disable does not + # help (torch.compile unwraps it via innermost_fn). Instead we name the wrapper + # with the sentinel so apply_compile treats it as already patched and skips + # compile, keeping the turbo kernels in eager as designed. + def _run_experts_grouped_mm_dynamic(*args, **kwargs): + kwargs.setdefault("use_fp8", use_moe_fp8) + return _run_experts_grouped_mm(*args, **kwargs) + + _run_experts_grouped_mm_dynamic.__qualname__ = _ALREADY_PATCHED_SENTINEL + _run_experts_grouped_mm_dynamic.__name__ = _ALREADY_PATCHED_SENTINEL + + torchtitan.models.moe.moe._run_experts_grouped_mm = _run_experts_grouped_mm_dynamic log_rank_0( "[Patch:torchtitan.primus_turbo.moe_grouped_mm] " diff --git a/primus/configs/models/torchtitan/gpt_oss_120b-fp8.yaml b/primus/configs/models/torchtitan/gpt_oss_120b-fp8.yaml new file mode 100644 index 000000000..d35a4795d --- /dev/null +++ b/primus/configs/models/torchtitan/gpt_oss_120b-fp8.yaml @@ -0,0 +1,14 @@ +job: + dump_folder: "./outputs" + description: "GPT-OSS 120B FP8 training" + +model: + name: "gpt_oss" + flavor: "120b" + hf_assets_path: "openai/gpt-oss-120b" + # Attention (sinks + sliding window) is routed to Primus-Turbo via the + # gptoss_sink_attention setup patch. FP8 dense linears are enabled through the + # quantize.linear.float8 converter, which the turbo_float8_linear patch swaps + # for the Primus-Turbo FP8 converter when use_turbo_float8_linear is set. + converters: + - quantize.linear.float8 diff --git a/primus/configs/models/torchtitan/gpt_oss_120b.yaml b/primus/configs/models/torchtitan/gpt_oss_120b.yaml new file mode 100644 index 000000000..4cccb592b --- /dev/null +++ b/primus/configs/models/torchtitan/gpt_oss_120b.yaml @@ -0,0 +1,12 @@ +job: + dump_folder: "./outputs" + description: "GPT-OSS 120B training" + +model: + name: "gpt_oss" + flavor: "120b" + hf_assets_path: "openai/gpt-oss-120b" + # No model converter: GPT-OSS attention (learnable sinks + sliding window) is + # routed to Primus-Turbo via the torchtitan.primus_turbo.gptoss_sink_attention + # setup patch (functional flash_attn_func), not via the module-form converter. + converters: [] diff --git a/primus/configs/models/torchtitan/gpt_oss_20b-fp8.yaml b/primus/configs/models/torchtitan/gpt_oss_20b-fp8.yaml new file mode 100644 index 000000000..752ea87df --- /dev/null +++ b/primus/configs/models/torchtitan/gpt_oss_20b-fp8.yaml @@ -0,0 +1,14 @@ +job: + dump_folder: "./outputs" + description: "GPT-OSS 20B FP8 training" + +model: + name: "gpt_oss" + flavor: "20b" + hf_assets_path: "openai/gpt-oss-20b" + # Attention (sinks + sliding window) is routed to Primus-Turbo via the + # gptoss_sink_attention setup patch. FP8 dense linears are enabled through the + # quantize.linear.float8 converter, which the turbo_float8_linear patch swaps + # for the Primus-Turbo FP8 converter when use_turbo_float8_linear is set. + converters: + - quantize.linear.float8 diff --git a/primus/configs/models/torchtitan/gpt_oss_20b.yaml b/primus/configs/models/torchtitan/gpt_oss_20b.yaml new file mode 100644 index 000000000..f1ee457a9 --- /dev/null +++ b/primus/configs/models/torchtitan/gpt_oss_20b.yaml @@ -0,0 +1,12 @@ +job: + dump_folder: "./outputs" + description: "GPT-OSS 20B training" + +model: + name: "gpt_oss" + flavor: "20b" + hf_assets_path: "openai/gpt-oss-20b" + # No model converter: GPT-OSS attention (learnable sinks + sliding window) is + # routed to Primus-Turbo via the torchtitan.primus_turbo.gptoss_sink_attention + # setup patch (functional flash_attn_func), not via the module-form converter. + converters: [] diff --git a/runner/helpers/hooks/train/pretrain/torchtitan/prepare.py b/runner/helpers/hooks/train/pretrain/torchtitan/prepare.py index 799d2e820..459300b08 100644 --- a/runner/helpers/hooks/train/pretrain/torchtitan/prepare.py +++ b/runner/helpers/hooks/train/pretrain/torchtitan/prepare.py @@ -50,8 +50,16 @@ def parse_args(): def pip_install_editable(path: Path, name: str): - log_info(f"Installing {name} in editable mode via pip (path: {path})") - ret = subprocess.run(["pip", "install", "-e", ".", "-q"], cwd=path) + # TorchTitan v0.2.2 has PEP-420 namespace subpackages (e.g. torchtitan/tools, + # no __init__.py) that the default editable finder cannot resolve. Use compat + # mode (writes a .pth with the source root on sys.path) after uninstalling any + # stale torchtitan, so the training subprocess can import them. + log_info(f"Installing {name} in editable (compat) mode via pip (path: {path})") + subprocess.run(["pip", "uninstall", "-y", name.lower()], cwd=path) + ret = subprocess.run( + ["pip", "install", "-e", ".", "--config-settings", "editable_mode=compat", "-q"], + cwd=path, + ) if ret.returncode != 0: log_error_and_exit(f"Failed to install {name} via pip.") diff --git a/tests/conftest.py b/tests/conftest.py index b4048ee90..f3d9cef97 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -115,6 +115,15 @@ def pytest_configure(config): if str(megatron_path) not in sys.path: sys.path.append(str(megatron_path)) + # TorchTitan v0.2.2 has PEP-420 namespace subpackages (e.g. torchtitan/tools, + # no __init__.py) that only import with the source root on sys.path. Insert the + # submodule at the FRONT so it wins over any stale torchtitan in site-packages. + torchtitan_path = os.environ.get("TORCHTITAN_PATH") + if torchtitan_path is None or not os.path.exists(torchtitan_path): + torchtitan_path = project_root / "third_party" / "torchtitan" + if str(torchtitan_path) not in sys.path: + sys.path.insert(0, str(torchtitan_path)) + # Only needed for the megatron suite, so skip for unrelated selections # (fail-open: detection errors still run it). if _selection_includes_megatron(config): diff --git a/tests/trainer/test_torchtitan_trainer.py b/tests/trainer/test_torchtitan_trainer.py index 587c6edd6..9e116e621 100644 --- a/tests/trainer/test_torchtitan_trainer.py +++ b/tests/trainer/test_torchtitan_trainer.py @@ -230,3 +230,34 @@ def test_deepseek_v3_671b(self): "True", ], ) + + def test_gpt_oss_20B(self): + # Default Primus-Turbo path: GPT-OSS sink attention (flash_attn_func). + run_script( + self.__class__.__name__, + "gpt_oss_20B", + "examples/torchtitan/configs/MI300X/gpt_oss_20B-BF16-pretrain.yaml", + extra_args=[ + "--model.n_layers", + "4", + "--training.steps", + "3", + "--training.mock_data", + "True", + ], + ) + + def test_gpt_oss_20B_fp8(self): + run_script( + self.__class__.__name__, + "gpt_oss_20B_fp8", + "examples/torchtitan/configs/MI300X/gpt_oss_20B-FP8-pretrain.yaml", + extra_args=[ + "--model.n_layers", + "4", + "--training.steps", + "3", + "--training.mock_data", + "True", + ], + ) diff --git a/tests/unit_tests/backends/torchtitan/test_gpt_oss_configs.py b/tests/unit_tests/backends/torchtitan/test_gpt_oss_configs.py new file mode 100644 index 000000000..878384493 --- /dev/null +++ b/tests/unit_tests/backends/torchtitan/test_gpt_oss_configs.py @@ -0,0 +1,108 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for the GPT-OSS TorchTitan configs added for the v0.2.2 upgrade. + +CPU-only: these only parse YAML, so they run in CI without GPU or torchtitan. + +Covers: + - gpt_oss model configs (BF16 + FP8) resolve to name=gpt_oss with the right + flavor and converters (BF16 -> [], FP8 -> [quantize.linear.float8]). + - gpt_oss example configs (20B/120B x BF16/FP8 x MI300X/MI325X/MI355X) + default to Primus-Turbo on (sink attention) and wire the fp8 switch only + for the FP8 variant. +""" + +from pathlib import Path + +import pytest +import yaml + +REPO_ROOT = Path(__file__).resolve().parents[4] +MODEL_CFG_DIR = REPO_ROOT / "primus" / "configs" / "models" / "torchtitan" +EXAMPLE_DIR = REPO_ROOT / "examples" / "torchtitan" / "configs" +MACHINES = ["MI300X", "MI325X", "MI355X"] + + +def _load(path: Path) -> dict: + assert path.exists(), f"missing config: {path}" + with open(path) as f: + return yaml.safe_load(f) + + +# ----------------------------------------------------------------------------- +# Model configs +# ----------------------------------------------------------------------------- + + +@pytest.mark.parametrize("flavor", ["20b", "120b"]) +def test_gpt_oss_bf16_model_config(flavor): + cfg = _load(MODEL_CFG_DIR / f"gpt_oss_{flavor}.yaml") + model = cfg["model"] + assert model["name"] == "gpt_oss" + assert model["flavor"] == flavor + # BF16 uses no model converter (sink attention is a setup patch, and the + # module-form primus_turbo converter is incompatible with GPT-OSS FlexAttn). + assert model["converters"] == [] + + +@pytest.mark.parametrize("flavor", ["20b", "120b"]) +def test_gpt_oss_fp8_model_config(flavor): + cfg = _load(MODEL_CFG_DIR / f"gpt_oss_{flavor}-fp8.yaml") + model = cfg["model"] + assert model["name"] == "gpt_oss" + assert model["flavor"] == flavor + # FP8 enables fp8 dense linears through the quantize.linear.float8 converter. + assert model["converters"] == ["quantize.linear.float8"] + + +# ----------------------------------------------------------------------------- +# Example pretrain configs +# ----------------------------------------------------------------------------- + + +def _example(machine: str, name: str) -> dict: + return _load(EXAMPLE_DIR / machine / f"{name}-pretrain.yaml") + + +@pytest.mark.parametrize("machine", MACHINES) +@pytest.mark.parametrize("size", ["20B", "120B"]) +@pytest.mark.parametrize("precision", ["BF16", "FP8"]) +def test_gpt_oss_example_defaults_turbo_on(machine, size, precision): + cfg = _example(machine, f"gpt_oss_{size}-{precision}") + pre = cfg["modules"]["pre_trainer"] + assert pre["framework"] == "torchtitan" + + expected_model = ( + f"gpt_oss_{size.lower()}.yaml" if precision == "BF16" else f"gpt_oss_{size.lower()}-fp8.yaml" + ) + assert pre["model"] == expected_model + + turbo = pre["overrides"]["primus_turbo"] + # Turbo + sink attention are on by default for every gpt_oss example. + assert turbo["enable_primus_turbo"] is True + assert turbo["use_turbo_attention"] is True + + if precision == "FP8": + assert turbo["use_turbo_float8_linear"] is True + else: + assert turbo["use_turbo_float8_linear"] is False + + # MoE experts use GPT-OSS's own grouped-mm path; the shared moe fp8 switch + # must stay off so it is not silently assumed to apply. + assert turbo["use_moe_fp8"] is False + + +@pytest.mark.parametrize("machine", MACHINES) +@pytest.mark.parametrize("size", ["20B", "120B"]) +@pytest.mark.parametrize("precision", ["BF16", "FP8"]) +def test_gpt_oss_example_uses_expert_parallel(machine, size, precision): + cfg = _example(machine, f"gpt_oss_{size}-{precision}") + par = cfg["modules"]["pre_trainer"]["overrides"]["parallelism"] + # GPT-OSS is MoE: experts are sharded with Expert Parallel. + assert par["expert_parallel_degree"] == 8 + assert par["data_parallel_shard_degree"] == -1 diff --git a/tests/unit_tests/backends/torchtitan/test_gpt_oss_sink_attention.py b/tests/unit_tests/backends/torchtitan/test_gpt_oss_sink_attention.py new file mode 100644 index 000000000..49e2e5eae --- /dev/null +++ b/tests/unit_tests/backends/torchtitan/test_gpt_oss_sink_attention.py @@ -0,0 +1,150 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for the GPT-OSS Primus-Turbo sink-attention integration. + +Two groups: + * Registration / condition (CPU-only, no torchtitan needed): the setup patch + is registered and only applies for gpt_oss with turbo attention enabled. + * Mirror interface (needs torchtitan; auto-skipped otherwise): the mirror + Attention keeps the learnable per-head sinks and accepts ``positions``, + and the mirror TransformerBlock selects the sliding window on even layers. +""" + +import inspect +from types import SimpleNamespace + +import pytest + +# Importing the patch module registers the patch (top-level import pulls in only +# primus.core.*, torchtitan is imported lazily inside the handler). +import primus.backends.torchtitan.patches.turbo.gptoss_sink_attention_patches as gptoss_patch +from primus.core.patches import PatchContext +from primus.core.patches.patch_registry import PatchRegistry + +PATCH_ID = "torchtitan.primus_turbo.gptoss_sink_attention" + + +def _ctx(model_name, enable_turbo, use_turbo_attention, backend="torchtitan", phase="setup"): + params = SimpleNamespace( + model=SimpleNamespace(name=model_name), + primus_turbo=SimpleNamespace( + enable_primus_turbo=enable_turbo, + use_turbo_attention=use_turbo_attention, + ), + ) + module_config = SimpleNamespace(params=params) + return PatchContext(backend=backend, phase=phase, extra={"module_config": module_config}) + + +class TestSinkAttentionPatchRegistration: + def test_patch_registered(self): + assert PATCH_ID in PatchRegistry.list_ids() + patch = PatchRegistry.get(PATCH_ID) + assert patch is not None + assert patch.backend == "torchtitan" + assert patch.phase == "setup" + + def test_condition_enabled_for_gpt_oss(self): + assert gptoss_patch._gptoss_turbo_enabled(_ctx("gpt_oss", True, True)) is True + + def test_condition_disabled_for_other_models(self): + for name in ("llama3", "llama4", "qwen3", "deepseek_v3"): + assert gptoss_patch._gptoss_turbo_enabled(_ctx(name, True, True)) is False + + def test_condition_disabled_when_turbo_off(self): + assert gptoss_patch._gptoss_turbo_enabled(_ctx("gpt_oss", False, True)) is False + assert gptoss_patch._gptoss_turbo_enabled(_ctx("gpt_oss", True, False)) is False + + def test_applies_to_enabled_for_gpt_oss(self): + # End-to-end wiring (backend + phase + condition) resolves to True for a + # gpt_oss turbo run. Negative cases are covered by the condition tests + # above (applies_to logs on skip, which needs the Primus logger set up). + patch = PatchRegistry.get(PATCH_ID) + assert patch.applies_to(_ctx("gpt_oss", True, True)) is True + + +@pytest.fixture +def gpt_oss_args(): + pytest.importorskip("torchtitan") + from torchtitan.models.gpt_oss import GptOssModelArgs + from torchtitan.models.moe import MoEArgs + + return GptOssModelArgs( + dim=256, + n_heads=8, + n_kv_heads=2, + head_dim=32, + sliding_window_size=64, + vocab_size=512, + moe_inter_dim=256, + n_layers=2, + moe_args=MoEArgs( + num_experts=4, + num_shared_experts=0, + top_k=2, + use_grouped_mm=True, + score_func="softmax", + route_norm=True, + gate_bias=True, + ), + ) + + +class TestSinkAttentionMirror: + def test_attention_keeps_sinks_and_accepts_positions(self, gpt_oss_args): + from primus.backends.torchtitan.models.gpt_oss.model.model import Attention + + attn = Attention(gpt_oss_args) + # Learnable per-head sinks preserved from the upstream module. + assert tuple(attn.sinks.shape) == (gpt_oss_args.n_heads,) + # Full (no) window by default; per-layer window set by the block. + assert attn._turbo_window == (-1, -1) + assert attn.sliding_window_size == gpt_oss_args.sliding_window_size + # GPT-OSS attention forward matches the upstream signature (rope_cache + + # attention_masks; GPT-OSS does not take the llama/qwen positions arg). + params = inspect.signature(attn.forward).parameters + assert "rope_cache" in params + assert "attention_masks" in params + + def test_block_selects_sliding_window_on_even_layers(self, gpt_oss_args): + import torch + import torch.nn as nn + + from primus.backends.torchtitan.models.gpt_oss.model.model import ( + TransformerBlock, + ) + + captured = {} + + class AttnStub(nn.Module): + def __init__(self, sliding_window_size): + super().__init__() + self.sliding_window_size = sliding_window_size + self._turbo_window = (-1, -1) + + def forward(self, x, rope_cache, attention_masks=None): + captured["window"] = self._turbo_window + return torch.zeros_like(x) + + class MoeStub(nn.Module): + def forward(self, x): + return torch.zeros_like(x) + + for layer_id, expected in [ + (0, (gpt_oss_args.sliding_window_size, 0)), # even -> sliding + (1, (-1, -1)), # odd -> full + ]: + block = TransformerBlock(layer_id, gpt_oss_args) + assert block.use_sliding_attention == (layer_id % 2 == 0) + # Replace the heavy attention/moe with CPU stubs to test the block's + # window-selection logic without invoking the Triton flash kernel. + block.attention = AttnStub(gpt_oss_args.sliding_window_size) + block.moe = MoeStub() + x = torch.zeros(1, 4, gpt_oss_args.dim) + block(x, None, None) + assert captured["window"] == expected diff --git a/tests/unit_tests/backends/torchtitan/test_moe_grouped_mm_patch.py b/tests/unit_tests/backends/torchtitan/test_moe_grouped_mm_patch.py new file mode 100644 index 000000000..10ef794ba --- /dev/null +++ b/tests/unit_tests/backends/torchtitan/test_moe_grouped_mm_patch.py @@ -0,0 +1,103 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for the Primus-Turbo MoE grouped_mm patch. + +Regression focus: the replacement on +``torchtitan.models.moe.moe._run_experts_grouped_mm`` must expose a +``__qualname__`` containing the sentinel ``_run_experts_grouped_mm_dynamic`` so +upstream ``apply_compile`` treats it as already patched and skips +``torch.compile`` over the ``torch.compiler.disable``'d turbo kernels (dynamo +gb0098). A ``functools.partial`` (no ``__qualname__``) would crash that read. +""" + +import functools +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +import primus.backends.torchtitan.patches.turbo.moe_grouped_mm_patches as moe_patch +from primus.core.patches import PatchContext +from primus.core.patches.patch_registry import PatchRegistry + +PATCH_ID = "torchtitan.primus_turbo.moe_grouped_mm" + + +def _ctx(enable_turbo=True, use_turbo_grouped_mm=True, use_moe_fp8=False): + params = SimpleNamespace( + primus_turbo=SimpleNamespace( + enable_primus_turbo=enable_turbo, + use_turbo_grouped_mm=use_turbo_grouped_mm, + use_moe_fp8=use_moe_fp8, + ), + ) + module_config = SimpleNamespace(params=params) + return PatchContext(backend="torchtitan", phase="setup", extra={"module_config": module_config}) + + +class TestMoeGroupedMmPatch: + def test_patch_registered(self): + assert PATCH_ID in PatchRegistry.list_ids() + p = PatchRegistry.get(PATCH_ID) + assert p is not None and p.backend == "torchtitan" and p.phase == "setup" + + def test_condition_gates_on_turbo_flags(self): + p = PatchRegistry.get(PATCH_ID) + assert p.condition(_ctx(True, True)) is True + assert p.condition(_ctx(False, True)) is False + assert p.condition(_ctx(True, False)) is False + + def _run_patch(self, use_moe_fp8, stub_impl): + """Apply the real patch against the real torchtitan moe module. + + Saves/restores torchtitan.models.moe.moe._run_experts_grouped_mm and + stubs the Primus grouped_mm impl so no GPU kernel runs. + """ + import torchtitan.models.moe.moe as tt_moe + + import primus.backends.torchtitan.models.moe.moe as primus_moe + + orig_tt = tt_moe._run_experts_grouped_mm + orig_primus = primus_moe._run_experts_grouped_mm + try: + primus_moe._run_experts_grouped_mm = stub_impl + with patch.object(moe_patch, "log_rank_0"): + moe_patch.patch_torchtitan_moe(_ctx(use_moe_fp8=use_moe_fp8)) + return tt_moe._run_experts_grouped_mm + finally: + tt_moe._run_experts_grouped_mm = orig_tt + primus_moe._run_experts_grouped_mm = orig_primus + + def test_replacement_has_qualname_not_partial(self): + pytest.importorskip("torchtitan") + + def _stub(w1, w2, w3, x, num_tokens_per_expert, use_fp8=True): + return use_fp8 + + _stub.__qualname__ = "_run_experts_grouped_mm" + + replacement = self._run_patch(use_moe_fp8=True, stub_impl=_stub) + assert not isinstance(replacement, functools.partial) + assert hasattr(replacement, "__qualname__") + # Must match upstream's already_patched guard to skip torch.compile. + assert "_run_experts_grouped_mm_dynamic" in replacement.__qualname__ + + def test_replacement_binds_use_fp8(self): + pytest.importorskip("torchtitan") + captured = {} + + def _stub(w1, w2, w3, x, num_tokens_per_expert, use_fp8=True): + captured["use_fp8"] = use_fp8 + return "ok" + + _stub.__qualname__ = "_run_experts_grouped_mm" + + replacement = self._run_patch(use_moe_fp8=True, stub_impl=_stub) + # use_fp8 must be injected by the wrapper (upstream calls positionally). + assert replacement(1, 2, 3, 4, 5) == "ok" + assert captured["use_fp8"] is True diff --git a/tests/unit_tests/conftest.py b/tests/unit_tests/conftest.py index 24d07cf77..f3366df71 100644 --- a/tests/unit_tests/conftest.py +++ b/tests/unit_tests/conftest.py @@ -38,6 +38,15 @@ def pytest_configure(config): if str(megatron_path) not in sys.path: sys.path.append(str(megatron_path)) + # TorchTitan v0.2.2 has PEP-420 namespace subpackages (e.g. torchtitan/tools, + # no __init__.py) that only import with the source root on sys.path. Insert the + # submodule at the FRONT so it wins over any stale torchtitan in site-packages. + torchtitan_path = os.environ.get("TORCHTITAN_PATH") + if torchtitan_path is None or not os.path.exists(torchtitan_path): + torchtitan_path = project_root / "third_party" / "torchtitan" + if str(torchtitan_path) not in sys.path: + sys.path.insert(0, str(torchtitan_path)) + # Register custom markers used by the primus test suite. config.addinivalue_line( "markers", diff --git a/third_party/torchtitan b/third_party/torchtitan index 5fb7cc2e3..73a0e6979 160000 --- a/third_party/torchtitan +++ b/third_party/torchtitan @@ -1 +1 @@ -Subproject commit 5fb7cc2e3bbb9b9dc0ab7af34ed5cc58b5f32021 +Subproject commit 73a0e6979dd10b6b1904098eb3c8f62c18ab87ce From 8b5e8091899e723dd8118ca3946bf61674bb02f9 Mon Sep 17 00:00:00 2001 From: botaohu001 Date: Fri, 17 Jul 2026 17:55:48 +0800 Subject: [PATCH 043/127] =?UTF-8?q?fix(fsdp2):=20skip=20explicit=20forward?= =?UTF-8?q?=20prefetch=20when=20activation=20recompute=20=E2=80=A6=20(#884?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes a ~50% FSDP2 throughput regression (llama3.1-70B, MI355X) introduced in #808. The rewritten PrimusTorchFullyShardedDataParallel set explicit forward prefetch unconditionally; with activation recompute on, the recomputed forward triggers wrong-direction all-gathers that expose communication and inflate memory. Fix: skip explicit forward prefetch when recompute is enabled, matching Megatron upstream. Verified on 8×MI355X: ~1026 → ~2026 toks/s/gpu, iter-1 reserved 249GB → 201GB. --- .../distributed/torch_fully_sharded_data_parallel.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py b/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py index 292dd1e68..be24cae00 100644 --- a/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py +++ b/primus/backends/megatron/core/distributed/torch_fully_sharded_data_parallel.py @@ -323,10 +323,16 @@ def restore_custom_attrs(module, custom_attrs): log_rank_0(f"FSDP2: wrapped {len(wrapped_list)} inner modules + root") prefetch_depth = getattr(self.config, "fsdp_prefetch_depth", 1) + recompute_on = getattr(self.config, "recompute_granularity", None) is not None for i, mod in enumerate(wrapped_list): - fwd_targets = wrapped_list[i + 1 : i + 1 + prefetch_depth] - if fwd_targets: - mod.set_modules_to_forward_prefetch(fwd_targets) + # With activation recompute enabled, the recomputed forward pass triggers + # wrong-direction forward-prefetch all-gathers that expose communication and + # inflate peak memory (matches Megatron upstream, which only sets backward + # prefetch when recompute is on). Skip explicit forward prefetch in that case. + if not recompute_on: + fwd_targets = wrapped_list[i + 1 : i + 1 + prefetch_depth] + if fwd_targets: + mod.set_modules_to_forward_prefetch(fwd_targets) bwd_start = max(0, i - prefetch_depth) bwd_targets = list(reversed(wrapped_list[bwd_start:i])) From e2c0b9e41576722e52796e75d7f512e6f155d399 Mon Sep 17 00:00:00 2001 From: Jani Sainio Date: Sat, 18 Jul 2026 03:44:07 +0300 Subject: [PATCH 044/127] feat(flux): MLPerf logging/warmup/lr-schedule patches (#820) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets `feat/flux/trainers` and also merges `feat/flux/data` — review after both. Top of the stack. ## What this changes The MLPerf-alignment layer: logging patches, warmup-state and lr-schedule patches, and the training-log / wall-clock-timer patches. ## Why it lists the data parent One MLPerf validation test imports the synthetic dataset + image task encoder from `feat/flux/data`; this is redundant through `feat/flux/trainers` but kept explicit so the data dependency is self-documented. ## Dependencies Sequenced after the CI-pins PR (`feat/flux/ci-env`); its tests exercise no mxfp4/compile path and pass on the current CI pin (no turbo-bump dependency). Builds on `feat/flux/trainers` + `feat/flux/data`. ## Test plan `pytest tests/unit_tests/backends/megatron/diffusion -k "mlperf or warmup"` plus the runtime-hooks / training-log tests. Validated locally on an AMD GPU container: 51 passed. ## Files 14 (MLPerf logging/warmup/lr-schedule + training-log patches + tests). --------- Co-authored-by: Flux Split Trial Co-authored-by: Luiza Sayfullina Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- .../megatron/patches/lr_schedule_patches.py | 68 ++ .../patches/mlperf_logging_patches.py | 485 +++++++++++ .../megatron/patches/mlperf_warmup_patches.py | 471 +++++++++++ .../megatron/patches/training_log/__init__.py | 16 +- .../training_log/print_rank_last_patches.py | 367 +++++++- .../training_log/wall_clock_timer_patch.py | 85 ++ .../diffusion/test_mlperf_warmup_fp8_state.py | 281 +++++++ .../backends/megatron/test_mlperf_patches.py | 795 ++++++++++++++++++ .../megatron/test_mlperf_validation.py | 164 ++++ .../test_mlperf_warmup_state_equivalence.py | 86 ++ .../megatron/test_runtime_hooks_patches.py | 6 +- .../megatron/test_training_log_patches.py | 31 +- .../megatron/test_warmup_convergence.py | 363 ++++++++ .../megatron/test_warmup_prefetch_cache.py | 179 ++++ 14 files changed, 3333 insertions(+), 64 deletions(-) create mode 100644 primus/backends/megatron/patches/lr_schedule_patches.py create mode 100644 primus/backends/megatron/patches/mlperf_logging_patches.py create mode 100644 primus/backends/megatron/patches/mlperf_warmup_patches.py create mode 100644 primus/backends/megatron/patches/training_log/wall_clock_timer_patch.py create mode 100644 tests/unit_tests/backends/megatron/diffusion/test_mlperf_warmup_fp8_state.py create mode 100644 tests/unit_tests/backends/megatron/test_mlperf_patches.py create mode 100644 tests/unit_tests/backends/megatron/test_mlperf_validation.py create mode 100644 tests/unit_tests/backends/megatron/test_mlperf_warmup_state_equivalence.py create mode 100644 tests/unit_tests/backends/megatron/test_warmup_convergence.py create mode 100644 tests/unit_tests/backends/megatron/test_warmup_prefetch_cache.py diff --git a/primus/backends/megatron/patches/lr_schedule_patches.py b/primus/backends/megatron/patches/lr_schedule_patches.py new file mode 100644 index 000000000..616a274f1 --- /dev/null +++ b/primus/backends/megatron/patches/lr_schedule_patches.py @@ -0,0 +1,68 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +NeMo-aligned LR warmup patch. + +NeMo's WarmupHoldPolicy._get_warmup_lr (nemo/core/optim/lr_scheduler.py) +computes warmup as: + + lr = base_lr * (step + 1) / (warmup_steps + 1) + +Megatron's OptimizerParamScheduler.get_lr uses: + + lr = init_lr + (max_lr - init_lr) * num_steps / lr_warmup_steps + +In Megatron's sample-space (num_steps increments by GBS per iteration, +lr_warmup_steps = warmup_iters * GBS), the NeMo-equivalent formula is: + + lr = init_lr + (max_lr - init_lr) * (num_steps + GBS) / (lr_warmup_steps + GBS) + +This patch replaces the warmup branch of get_lr with the NeMo formula. +Enabled by setting nemo_aligned_lr_warmup: true in the YAML config. +""" + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +def _nemo_lr_enabled(ctx: PatchContext) -> bool: + args = get_args(ctx) + return args is not None and getattr(args, "nemo_aligned_lr_warmup", False) + + +@register_patch( + "megatron.lr_schedule.nemo_aligned", + backend="megatron", + phase="before_train", + description="Align LR warmup with NeMo's (step+1)/(warmup_steps+1) formula", + condition=_nemo_lr_enabled, + priority=50, +) +def patch_nemo_aligned_lr_warmup(ctx: PatchContext): + from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler + + _original_get_lr = OptimizerParamScheduler.get_lr + _gbs_cache = [None] + + def _nemo_get_lr(self, param_group): + if self.lr_warmup_steps > 0 and self.num_steps <= self.lr_warmup_steps: + if _gbs_cache[0] is None: + from megatron.training import get_args as megatron_get_args + + _gbs_cache[0] = megatron_get_args().global_batch_size + gbs = _gbs_cache[0] + max_lr = param_group.get("max_lr", self.max_lr) + return self.init_lr + ( + (max_lr - self.init_lr) * float(self.num_steps + gbs) / float(self.lr_warmup_steps + gbs) + ) + return _original_get_lr(self, param_group) + + OptimizerParamScheduler.get_lr = _nemo_get_lr + log_rank_0( + "[Patch:nemo_aligned_lr] Patched get_lr warmup: " + "(num_steps+GBS)/(lr_warmup_steps+GBS) = NeMo's (step+1)/(warmup+1)" + ) diff --git a/primus/backends/megatron/patches/mlperf_logging_patches.py b/primus/backends/megatron/patches/mlperf_logging_patches.py new file mode 100644 index 000000000..b53173d8a --- /dev/null +++ b/primus/backends/megatron/patches/mlperf_logging_patches.py @@ -0,0 +1,485 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +MLPerf Logging Patches for Flux Training. + +Installs MLPerf-compliant logging into Megatron's training loop by wrapping: + - training_log: emit INIT_STOP, RUN_START, tracked_stats, train_loss + - evaluate_and_print_results: emit EVAL events, convergence check + - print_rank_last / get_tensorboard_writer / get_wandb_writer: suppress + +Uses mlperf_logging.mllog library for structured event output. +""" + +import logging +import os +import sys +import time + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + +logger = logging.getLogger(__name__) + + +def _mlperf_logging_enabled(ctx: PatchContext) -> bool: + args = get_args(ctx) + return args is not None and getattr(args, "mlperf_mode", False) + + +class ThroughputTimer: + """Wall-clock throughput tracker with eval pause/resume.""" + + def __init__(self, gbs: int): + self.gbs = gbs + self.training_start_time: float | None = None + self.eval_cumulative_secs: float = 0.0 + self._eval_enter_time: float | None = None + self.consumed_samples: int = 0 + + def mark_training_start(self): + if self.training_start_time is None: + self.training_start_time = time.time() + + def update_samples(self, iteration: int): + self.consumed_samples = iteration * self.gbs + + def pause_for_eval(self): + self._eval_enter_time = time.time() + + def resume_after_eval(self): + if self._eval_enter_time is not None: + self.eval_cumulative_secs += time.time() - self._eval_enter_time + self._eval_enter_time = None + + def compute_throughput(self): + if self.training_start_time is None: + return 0.0 + wall = time.time() - self.training_start_time + training_secs = wall - self.eval_cumulative_secs + if training_secs <= 0: + return 0.0 + return self.consumed_samples / training_secs + + def compute_combined_throughput(self): + if self.training_start_time is None: + return 0.0 + wall = time.time() - self.training_start_time + if wall <= 0: + return 0.0 + return self.consumed_samples / wall + + +class FluxMLPerfLogger: + """MLPerf logger using mlperf_logging.mllog directly.""" + + def __init__( + self, + global_batch_size: int, + micro_batch_size: int, + target_val_loss: float = 0.586, + log_every_n_steps: int = 10, + ): + from mlperf_logging import mllog + + self._mllogger = mllog.get_mllogger() + self._constants = mllog.constants + self.gbs = global_batch_size + self.mbs = micro_batch_size + self.target_val_loss = target_val_loss + self.log_every_n_steps = log_every_n_steps + self.timer = ThroughputTimer(global_batch_size) + self._converged = False + + self.profiler = os.getenv("PROFILER", "") + self.profiler_warmup_steps = int(os.getenv("PROF_WARMUP_STEPS", "0")) + self.profiler_active_steps = int(os.getenv("PROF_ACTIVE_STEPS", "0")) + self.rpd = None + self.rpd_running = False + + if self.profiler == "rpd": + try: + from rpdTracerControl import rpdTracerControl + + rpdTracerControl.setFilename("trace.rpd", append=True) + self.rpd = rpdTracerControl() + logger.info("RPD profiler initialized") + except ImportError: + logger.warning("rpdTracerControl not available") + + def _event(self, key, value=None, metadata=None): + self._mllogger.event(key=key, value=value, metadata=metadata) + + def _start(self, key, value=None, metadata=None): + self._mllogger.start(key=key, value=value, metadata=metadata) + + def _end(self, key, value=None, metadata=None): + self._mllogger.end(key=key, value=value, metadata=metadata) + + def log_init(self, seed: int): + if int(os.environ.get("RANK", "0")) == 0: + self._start(key=self._constants.INIT_START) + self._event(key=self._constants.SUBMISSION_BENCHMARK, value="flux1") + self._event( + key=self._constants.SUBMISSION_ORG, + value=os.environ.get("MLLOG_SUBMISSION_ORG", "AMD"), + ) + self._event( + key=self._constants.SUBMISSION_DIVISION, + value=os.environ.get("MLLOG_SUBMISSION_DIVISION", "closed"), + ) + self._event( + key=self._constants.SUBMISSION_PLATFORM, + value=os.environ.get("MLLOG_SUBMISSION_PLATFORM", "MI355X"), + ) + self._event(key=self._constants.SUBMISSION_STATUS, value="onprem") + self._event(key="target_accuracy", value=self.target_val_loss) + self._event(key=self._constants.SEED, value=seed) + + def log_hyperparams(self, args): + if int(os.environ.get("RANK", "0")) != 0: + return + self._event(key=self._constants.GLOBAL_BATCH_SIZE, value=self.gbs) + self._event( + key=self._constants.TRAIN_SAMPLES, + value=getattr(args, "train_samples", 1099776), + ) + self._event( + key=self._constants.EVAL_SAMPLES, + value=getattr(args, "eval_samples", 29696), + ) + gas = max(self.gbs // self.mbs, 1) + self._event(key=self._constants.GRADIENT_ACCUMULATION_STEPS, value=gas) + self._event(key=self._constants.OPT_NAME, value="adamw") + self._event( + key=self._constants.OPT_BASE_LR, + value=getattr(args, "lr", 2e-4), + ) + self._event( + key="opt_adamw_beta_1", + value=getattr(args, "adam_beta1", 0.9), + ) + self._event( + key="opt_adamw_beta_2", + value=getattr(args, "adam_beta2", 0.95), + ) + self._event( + key="opt_adamw_epsilon", + value=getattr(args, "adam_eps", 1e-8), + ) + self._event( + key="opt_adamw_weight_decay", + value=getattr(args, "weight_decay", 0.1), + ) + + def log_init_stop_run_start(self): + if int(os.environ.get("RANK", "0")) == 0: + self._end(key=self._constants.INIT_STOP) + self._start(key=self._constants.RUN_START) + self._start(key=self._constants.EPOCH_START, metadata={"epoch_num": 0}) + self._start(key=self._constants.BLOCK_START, metadata={"first_epoch_num": 0}) + + def on_train_batch_end(self, global_step: int, loss: float, lr: float): + self.timer.mark_training_start() + self.timer.update_samples(global_step) + + self._handle_profiler(global_step) + + if int(os.environ.get("RANK", "0")) != 0: + return + if global_step % self.log_every_n_steps == 0: + self._event( + key="tracked_stats", + value={"train_loss": loss}, + metadata={ + "samples_count": global_step * self.gbs, + "lr": lr, + "step": global_step, + }, + ) + + def on_validation_start(self, global_step: int): + self.timer.update_samples(global_step) + self.timer.pause_for_eval() + + if int(os.environ.get("RANK", "0")) == 0: + if global_step > 0: + throughput = self.timer.compute_throughput() + self._event( + key="throughput", + value=throughput, + metadata={ + "samples_count": global_step * self.gbs, + "step": global_step, + }, + ) + self._end( + key=self._constants.BLOCK_STOP, + metadata={"first_epoch_num": 0}, + ) + self._start(key=self._constants.EVAL_START, metadata={"epoch_num": 0}) + + def on_validation_end(self, global_step: int, val_loss: float): + self.timer.resume_after_eval() + + if int(os.environ.get("RANK", "0")) == 0: + self._event( + key=self._constants.EVAL_ACCURACY, + value=val_loss, + metadata={ + "samples_count": global_step * self.gbs, + "step": global_step, + }, + ) + self._end(key=self._constants.EVAL_STOP, metadata={"epoch_num": 0}) + combined_throughput = self.timer.compute_combined_throughput() + self._event( + key="combined_throughput", + value=combined_throughput, + metadata={ + "samples_count": global_step * self.gbs, + "step": global_step, + }, + ) + + def _handle_profiler(self, global_step: int): + if self.profiler != "rpd": + return + if self.rpd and not self.rpd_running and global_step >= self.profiler_warmup_steps: + logger.info("Starting RPD profiler") + self.rpd.start() + self.rpd.rangePush("python", "Training", "") + self.rpd_running = True + if self.rpd_running and global_step > self.profiler_warmup_steps + self.profiler_active_steps: + logger.info("Stopping RPD profiler") + self.rpd.rangePop() + self.rpd.stop() + self.rpd = None + self.rpd_running = False + + @property + def converged(self): + return self._converged + + def log_run_stop(self, success: bool, global_step: int): + if success: + self._converged = True + if int(os.environ.get("RANK", "0")) == 0: + status = "success" if success else "aborted" + self._end( + key=self._constants.RUN_STOP, + value=status, + metadata={ + "samples_count": global_step * self.gbs, + "step": global_step, + "status": status, + }, + ) + + def teardown(self): + if self.rpd_running and self.rpd: + self.rpd.rangePop() + self.rpd.stop() + self.rpd = None + self.rpd_running = False + + +def _extract_val_loss(loss_dict): + """Extract scalar validation loss from captured total_loss_dict.""" + if not loss_dict or not isinstance(loss_dict, dict): + return None + for key in ("loss", "lm loss"): + if key in loss_dict: + val = loss_dict[key] + return val.item() if hasattr(val, "item") else float(val) + if loss_dict: + val = next(iter(loss_dict.values())) + return val.item() if hasattr(val, "item") else float(val) + return None + + +@register_patch( + "megatron.training.mlperf_logging", + backend="megatron", + phase="before_train", + description="Install MLPerf logging wrappers for Flux training", + condition=_mlperf_logging_enabled, + priority=15, +) +def patch_mlperf_logging(ctx: PatchContext): + """Install MLPerf logging: suppress Megatron output, wrap training_log and eval.""" + import megatron.training.training as megatron_training + + if getattr(megatron_training, "_primus_mlperf_logging_installed", False): + return + + args = get_args(ctx) + seed = getattr(args, "seed", 42) + gbs = getattr(args, "global_batch_size", 512) + mbs = getattr(args, "micro_batch_size", 64) + target_val_loss = getattr(args, "target_val_loss", 0.586) + log_interval = getattr(args, "log_interval", 10) + + mlperf_logger = FluxMLPerfLogger( + global_batch_size=gbs, + micro_batch_size=mbs, + target_val_loss=target_val_loss, + log_every_n_steps=log_interval, + ) + + mlperf_logger.log_init(seed=seed) + mlperf_logger.log_hyperparams(args) + + # --- Suppress Megatron's built-in logging --- + megatron_training.print_rank_last = lambda *a, **k: None + + for writer_fn in ("get_tensorboard_writer", "get_wandb_writer"): + if hasattr(megatron_training, writer_fn): + setattr(megatron_training, writer_fn, lambda: None) + + # --- Wrap training_log --- + _orig_training_log = megatron_training.training_log + _first_training_log_call = [True] + + def _mlperf_training_log(*args_tl, **kwargs_tl): + if _first_training_log_call[0]: + _first_training_log_call[0] = False + mlperf_logger.log_init_stop_run_start() + + result = _orig_training_log(*args_tl, **kwargs_tl) + + try: + loss_dict = args_tl[0] if len(args_tl) > 0 else kwargs_tl.get("loss_dict", {}) + learning_rate = args_tl[2] if len(args_tl) > 2 else kwargs_tl.get("learning_rate", 0.0) + # Upstream Megatron: training_log(loss_dict, total_loss_dict, learning_rate, iteration, ...) + iteration = args_tl[3] if len(args_tl) > 3 else kwargs_tl.get("iteration", 0) + + if loss_dict: + loss_val = next(iter(loss_dict.values())) + if hasattr(loss_val, "item"): + loss_val = loss_val.item() + mlperf_logger.on_train_batch_end(iteration, loss_val, learning_rate) + + if iteration % log_interval == 0: + lr_str = f"{learning_rate:.2e}" if learning_rate else "N/A" + sys.stdout.write( + f"step {iteration} | loss: {loss_val:.4f} | lr: {lr_str}" + f" | samples: {iteration * gbs}\n" + ) + sys.stdout.flush() + except Exception as e: + logger.debug("MLPerf training_log hook: %s", e) + + return result + + _mlperf_training_log._primus_mlperf_logging_wrapper = True + megatron_training.training_log = _mlperf_training_log + + # --- Wrap evaluate_and_print_results --- + _orig_eval = megatron_training.evaluate_and_print_results + + def _mlperf_evaluate_and_print_results(*eval_args, **eval_kwargs): + # evaluate_and_print_results(prefix, fwd, data, model, iteration[4], ...) + iteration = eval_kwargs.get("iteration", eval_args[4] if len(eval_args) > 4 else 0) + + mlperf_logger.on_validation_start(iteration) + + # Temporarily wrap whatever `evaluate` is at call time (e.g. + # primus_evaluate installed by the evaluate patch) so we can + # capture the total_loss_dict it returns. This avoids relying + # on a persistent hook that later patches can overwrite. + _loss_capture = {} + _current_eval = megatron_training.evaluate + + def _capture_wrapper(*a, **kw): + res = _current_eval(*a, **kw) + td = res[0] if isinstance(res, tuple) else res + if isinstance(td, dict): + _loss_capture.update(td) + return res + + megatron_training.evaluate = _capture_wrapper + try: + import gc + + result = _orig_eval(*eval_args, **eval_kwargs) + gc.collect() + finally: + megatron_training.evaluate = _current_eval + + try: + import torch + + torch.cuda.empty_cache() + except Exception: + pass + + val_loss = _extract_val_loss(_loss_capture) + if val_loss is not None: + # Megatron's `evaluate()` (training.py:3178-3180) divides the + # per-rank accumulated loss locally and does NOT all-reduce + # across the data-parallel group — the result is intended for + # `print_rank_last` / TensorBoard which only read on a single + # rank. We must reduce here so every rank evaluates the same + # global validation loss against `target_val_loss`; otherwise + # ranks can disagree on the early-stop branch and the + # divergent `args.train_iters` mutation below desyncs + # collective ordering at the next training step, producing a + # NCCL watchdog deadlock (observed on FLUX 12B MLPerf at + # step 2560 when val_loss landed near target). + import torch + import torch.distributed as dist + + if dist.is_initialized(): + try: + from megatron.core import parallel_state as mpu + + dp_group = mpu.get_data_parallel_group() + except Exception: + dp_group = None + _vl = torch.tensor(val_loss, dtype=torch.float64, device="cuda") + dist.all_reduce(_vl, op=dist.ReduceOp.AVG, group=dp_group) + val_loss = _vl.item() + + mlperf_logger.on_validation_end(iteration, val_loss) + log_rank_0( + f"[MLPerf] Validation loss at step {iteration}: {val_loss:.6f} " + f"(target: {target_val_loss:.6f})" + ) + + if val_loss <= target_val_loss: + log_rank_0( + f"[MLPerf] Convergence reached! val_loss={val_loss:.6f} " + f"<= target={target_val_loss:.6f}" + ) + mlperf_logger.log_run_stop(success=True, global_step=iteration) + try: + from megatron.training import get_args as megatron_get_args + + megatron_get_args().train_iters = iteration + except Exception: + logger.warning("Could not set args.train_iters for early stop") + else: + if int(os.environ.get("RANK", "0")) == 0: + mlperf_logger._start( + key=mlperf_logger._constants.BLOCK_START, + metadata={"first_epoch_num": 0}, + ) + else: + logger.warning("Could not extract validation loss from evaluate result") + + return result + + _mlperf_evaluate_and_print_results._primus_mlperf_eval_wrapper = True + megatron_training.evaluate_and_print_results = _mlperf_evaluate_and_print_results + + megatron_training._primus_mlperf_logging_installed = True + + log_rank_0( + f"[Patch:mlperf_logging] Installed MLPerf logging (gbs={gbs}, " + f"target_val_loss={target_val_loss}, log_interval={log_interval})" + ) diff --git a/primus/backends/megatron/patches/mlperf_warmup_patches.py b/primus/backends/megatron/patches/mlperf_warmup_patches.py new file mode 100644 index 000000000..5ac532978 --- /dev/null +++ b/primus/backends/megatron/patches/mlperf_warmup_patches.py @@ -0,0 +1,471 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +MLPerf Warmup Patches for Flux Training. + +Wraps megatron.training.training.train_step with a one-shot hook that runs +all warmup steps in a tight loop on the first invocation, then executes the +first real training step and self-removes. This avoids relying on Megatron's +local ``iteration`` variable (which cannot be controlled from a train_step +wrapper) and mirrors NeMo's approach of running warmup before real data is +touched. + +Priority 95 ensures this hook is the outermost wrapper around the full +train_step chain (FP8 cache, delayed scaling, wall-clock timer, etc.). +Self-removal restores the inner chain intact. +""" + +import logging + +import torch +import torch.distributed + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + +logger = logging.getLogger(__name__) + + +def _log(msg): + log_rank_0(f"[MLPerf_WARMUP] {msg}") + + +def _warmup_enabled(ctx: PatchContext) -> bool: + args = get_args(ctx) + return args is not None and getattr(args, "warmup_train_steps", 0) > 0 + + +def _reset_fp8_te_spec(models): + """Reset FP8 state for TransformerEngine spec modules. + + Recipe-agnostic: walks ``fp8_meta`` and only touches buffers that exist on + the recipe-specific state object. Skips TE's ``reset_fp8_meta_tensors`` + helper because TE 2.8.0.dev0 unconditionally derefs ``.scale`` / + ``.amax_history``, which crashes on ``Float8CurrentScalingRecipeState`` + (current/tensorwise scaling has no persistent state — see TE's + ``fp8.py``: *"Per-tensor current quantization does not require state"*). + """ + count = 0 + for m in models: + for module in m.modules(): + if not hasattr(module, "fp8_initialized"): + continue + module.fp8_initialized = False + count += 1 + if not hasattr(module, "fp8_meta"): + continue + meta = module.fp8_meta + for key in ("scaling_fwd", "scaling_bwd"): + if key not in meta: + continue + tm = meta[key] + if hasattr(tm, "amax_history"): + tm.amax_history.fill_(0.0) + if hasattr(tm, "scale"): + tm.scale.fill_(1.0) + if hasattr(tm, "scale_inv"): + tm.scale_inv.fill_(1.0) + return count + + +def _seed_fp8_amax(models, seed_value=1.0): + """Seed FP8 amax_history with a safe non-zero value to prevent scale=inf.""" + count = 0 + for m in models: + for module in m.modules(): + if not hasattr(module, "fp8_meta"): + continue + meta = module.fp8_meta + for key in ("scaling_fwd", "scaling_bwd"): + if key not in meta: + continue + tm = meta[key] + if hasattr(tm, "amax_history"): + tm.amax_history.fill_(seed_value) + count += 1 + return count + + +def _reset_fp8_local_spec(models): + """Reset FP8 state for local spec (tensorwise delayed-scaling) modules. + + Re-initialises per-module delayed-scaling buffers via the canonical + ``_init_delayed_scaling_state`` helper. The new buffers are separate + objects from the ``_DelayedScalingRegistry``'s global tensors, so the + pointer-check in ``_fast_update_scales`` / ``_fast_update_scales_with_history`` + will detect the mismatch and trigger ``registry.__init__(modules)`` on the + next training step, which re-creates global tensors with + ``_first_step = True`` and bootstraps weight amaxes from the restored weights. + + Buffers are moved to the module's device because ``_init_delayed_scaling_state`` + creates plain CPU tensors (bare ``torch.zeros`` / ``torch.tensor``). + """ + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _init_delayed_scaling_state, + ) + + _DELAYED_BUF_NAMES = ( + "scale_input", + "scale_weight", + "scale_grad", + "amax_history_input", + "amax_history_weight", + "amax_history_grad", + "staged_input_amax", + "staged_grad_amax", + "staged_weight_amax", + ) + + count = 0 + for m in models: + for module in m.modules(): + if not getattr(module, "_use_delayed_scaling", False): + continue + device = module.weight.device + _init_delayed_scaling_state(module) + for buf_name in _DELAYED_BUF_NAMES: + buf = module._buffers.get(buf_name) + if buf is not None and buf.device != device: + module._buffers[buf_name] = buf.to(device) + count += 1 + return count + + +def _neuter_optimizer(optimizer): + """Set optimizer to no-op mode: betas=[1,1], weight_decay=0.""" + saved = [] + inner = getattr(optimizer, "optimizer", optimizer) + _log( + f"Neutering optimizer: type={type(optimizer).__name__}, " + f"inner={type(inner).__name__}, " + f"param_groups={len(inner.param_groups)}" + ) + for group in inner.param_groups: + state = {} + for key in ("betas", "weight_decay", "bias_correction", "pre_mult_wd"): + if key in group: + state[key] = group[key] + saved.append(state) + + if "betas" in group: + group["betas"] = [1.0, 1.0] + if "weight_decay" in group: + group["weight_decay"] = 0.0 + if "bias_correction" in group: + group["bias_correction"] = False + if "pre_mult_wd" in group: + group["pre_mult_wd"] = 0.0 + return saved + + +def _restore_optimizer(optimizer, saved): + """Restore optimizer hyperparams (betas, weight_decay, etc.).""" + inner = getattr(optimizer, "optimizer", optimizer) + for group, state in zip(inner.param_groups, saved): + for key, val in state.items(): + group[key] = val + _log("Restored optimizer parameters") + + +def _reset_optimizer_state(optimizer): + """Zero per-parameter step counters so Adam acts as if no steps occurred. + + Handles both flat ``MegatronOptimizer`` wrappers and ``ChainedOptimizer`` + which wraps several sub-optimizers. Step counts live either in + ``param_groups[i]["step"]`` (Apex / TE FusedAdam) or in + ``optimizer.state[p]["step"]`` (stock PyTorch Adam). + """ + + def _reset_single(opt): + inner = getattr(opt, "optimizer", opt) + for group in inner.param_groups: + if "step" in group: + group["step"] = 0 + for state in inner.state.values(): + if isinstance(state, dict) and "step" in state: + if isinstance(state["step"], torch.Tensor): + state["step"].zero_() + else: + state["step"] = 0 + + if hasattr(optimizer, "chained_optimizers"): + for sub_opt in optimizer.chained_optimizers: + _reset_single(sub_opt) + else: + _reset_single(optimizer) + _log("Reset optimizer step counters") + + +@register_patch( + "megatron.training.mlperf_warmup", + backend="megatron", + phase="before_train", + description="MLPerf warmup: synthetic data steps before measured training", + condition=_warmup_enabled, + priority=95, +) +def patch_mlperf_warmup(ctx: PatchContext): + """Install warmup hook on train_step at priority 95 (outermost wrapper).""" + import megatron.training.training as mt + + if hasattr(mt.train_step, "_primus_warmup_hook"): + return + + primus_args = get_args(ctx) + warmup_steps = getattr(primus_args, "warmup_train_steps", 2) + + _lazy_state = { + "initialized": False, + "synthetic_iter": None, + "use_fsdp2_fp8": False, + "transformer_impl": "local", + } + + _wrapped_chain = mt.train_step + _warmup_done = [False] + + def _lazy_init(): + """One-time initialization on first train_step call, when Megatron args exist.""" + if _lazy_state["initialized"]: + return + + from megatron.training import get_args as megatron_get_args + + megatron_args = megatron_get_args() + + from torch.utils.data import DataLoader + + from primus.backends.megatron.data.dataloader import MegatronDataloaderWrapper + from primus.backends.megatron.data.synthetic.mock_datasets import ( + PreGeneratedMockFluxSchnellDataset, + ) + + image_size = getattr(primus_args, "image_size", 256) + vae_latent_mode = getattr(primus_args, "vae_latent_mode", "resample") + mbs = getattr(primus_args, "micro_batch_size", 64) + + mock_dataset = PreGeneratedMockFluxSchnellDataset( + num_samples=max(mbs * 4, 256), + image_size=image_size, + vae_latent_mode=vae_latent_mode, + ) + mock_loader = DataLoader(mock_dataset, batch_size=mbs, shuffle=False, drop_last=True) + _lazy_state["synthetic_iter"] = MegatronDataloaderWrapper(mock_loader) + + _lazy_state["use_fsdp2_fp8"] = getattr(megatron_args, "use_fsdp2_fp8_all_gather", False) + _lazy_state["transformer_impl"] = getattr(megatron_args, "transformer_impl", "local") + _lazy_state["initialized"] = True + _log(f"Lazy init complete (warmup_steps={warmup_steps})") + + def _hooked_train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=None, + ): + if _warmup_done[0]: + return _wrapped_chain( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, + ) + + _lazy_init() + + from megatron.training import get_args as megatron_get_args + + megatron_args = megatron_get_args() + models = model if isinstance(model, (list, tuple)) else [model] + synthetic_iter = _lazy_state["synthetic_iter"] + + # ---- 1. Snapshot model parameters to CPU ---- + _log("Saving model parameters to CPU before warmup") + saved_params = {} + for m in models: + for name, p in m.named_parameters(): + saved_params[name] = p.data.to("cpu", non_blocking=True) + torch.cuda.synchronize() + _log(f"Saved {len(saved_params)} parameter tensors") + + # ---- 2. Neuter optimizer ---- + saved_opt = _neuter_optimizer(optimizer) + + # ---- 3. Suppress training_log and eval during warmup ---- + saved_training_log = mt.training_log + saved_eval = mt.evaluate_and_print_results + mt.training_log = lambda *a, **k: None + mt.evaluate_and_print_results = lambda *a, **k: None + + # ---- 3b. Save LR scheduler state (NeMo never steps the scheduler during warmup) ---- + saved_lr_num_steps = opt_param_scheduler.num_steps + + # ---- 4. Run warmup steps with synthetic data ---- + for step_idx in range(warmup_steps): + _log(f"Warmup step {step_idx + 1}/{warmup_steps}") + _wrapped_chain( + forward_step_func, + synthetic_iter, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, + ) + _log(f"Completed {warmup_steps} warmup steps") + + # ---- 5. Restore optimizer ---- + _restore_optimizer(optimizer, saved_opt) + _reset_optimizer_state(optimizer) + + # ---- 6. Restore model parameters from CPU ---- + restored = 0 + for m in models: + for name, p in m.named_parameters(): + if name in saved_params: + p.data.copy_(saved_params[name]) + restored += 1 + del saved_params + _log(f"Restored {restored} parameter tensors from CPU snapshot") + + # ---- 7. FP8 reset (spec-aware) ---- + if _lazy_state["transformer_impl"] == "transformer_engine": + te_count = _reset_fp8_te_spec(models) + amax_count = _seed_fp8_amax(models) + _log(f"FP8 TE reset: {te_count} modules, " f"seeded {amax_count} amax tensors") + else: + local_count = _reset_fp8_local_spec(models) + _log(f"FP8 local spec reset: {local_count} modules") + + # ---- 8. FSDP2 FP8 all-gather recompute ---- + if _lazy_state["use_fsdp2_fp8"]: + try: + from primus.backends.megatron.core.distributed.fsdp2_fp8_all_gather import ( + precompute_fp8_scales_for_fsdp, + ) + + cache_data = getattr(megatron_args, "fp8_precompute_data_cache", True) + use_cpp = getattr(megatron_args, "use_cpp_fp8_quantize", False) + sr = getattr(megatron_args, "fp8_all_gather_stochastic_rounding", False) + precompute_fp8_scales_for_fsdp( + models[0], + cache_data=cache_data, + use_cpp_quantize=use_cpp, + stochastic_rounding=sr, + ) + _log("Recomputed FSDP2 FP8 all-gather scales") + except Exception as e: + _log(f"FSDP2 FP8 recompute failed (non-fatal): {e}") + + # ---- 9. Reload model params in optimizer (FSDP2 BF16 master weight) ---- + if hasattr(optimizer, "reload_model_params"): + optimizer.reload_model_params() + _log("Called optimizer.reload_model_params()") + + # ---- 10. Post-restore NaN check ---- + nan_params = 0 + for m in models: + for name, p in m.named_parameters(): + if p.data.is_floating_point() and torch.isnan(p.data).any(): + nan_params += 1 + _log(f"Post-restore parameter check: nan_params={nan_params}") + + # ---- 11. Zero gradients ---- + try: + optimizer.zero_grad(set_to_none=True) + except TypeError: + optimizer.zero_grad() + + # ---- 12. Reset counters ---- + megatron_args.consumed_train_samples = 0 + megatron_args.skipped_train_samples = 0 + opt_param_scheduler.num_steps = saved_lr_num_steps + _log( + f"Reset consumed_train_samples=0, skipped_train_samples=0, " + f"lr_scheduler.num_steps={saved_lr_num_steps}" + ) + + # ---- 13. Restore training_log and eval ---- + mt.training_log = saved_training_log + mt.evaluate_and_print_results = saved_eval + + # ---- 13b. Invalidate the CudaPrefetchIterator that was built around + # the SYNTHETIC iterator during warmup step 1. + # + # ``patch_grad_zero_and_data_prefetch`` builds a ``CudaPrefetchIterator`` + # the first time its ``_patched_train_step`` runs and caches it in a + # closure-local ``_prefetch_state["iter"]``. Because warmup step 1 + # is the first call into that train_step, the prefetch iterator gets + # bound to ``synthetic_iter``. ``MegatronDataloaderWrapper`` is + # cyclic (never raises ``StopIteration``), so subsequent real + # training steps would silently keep reading from the cycling + # synthetic dataset instead of the actual training dataset -- model + # overfits the mock samples and val_loss on real data stays stuck + # at ~1.38 forever. + # + # Dropping the cached entry forces the next train_step to rebuild + # the prefetch wrapper around its incoming ``data_iterator`` arg + # (the real iterator). + try: + from primus.backends.megatron.patches.delayed_fp8_scaling_patches import ( + reset_prefetch_state, + ) + + evicted = reset_prefetch_state() + if evicted is None: + _log(" Prefetch reset: no cached iterator to evict") + else: + _log( + f" Prefetch reset: evicted cached {type(evicted).__name__} " + f"(wrapped synthetic warmup iterator) -- next train_step " + f"will rebuild it around the real data_iterator" + ) + except Exception as _e: + _log(f" Prefetch reset failed (non-fatal): {_e}") + + # ---- 14. Synchronize ---- + torch.cuda.synchronize() + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + # ---- 15. Execute first real step ---- + _log("Executing first real train_step with training data") + result = _wrapped_chain( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, + ) + + # ---- 16. Self-remove ---- + _warmup_done[0] = True + mt.train_step = _wrapped_chain + _log("Self-removed warmup hook, train_step = inner wrapped chain") + + return result + + _hooked_train_step._primus_warmup_hook = True + mt.train_step = _hooked_train_step + + _log( + f"Installed MLPerf warmup hook (warmup_steps={warmup_steps}, " + f"deferred init until first train_step)" + ) + log_rank_0(f"[Patch:mlperf_warmup] Installed warmup hook " f"(warmup_steps={warmup_steps}, priority=95)") diff --git a/primus/backends/megatron/patches/training_log/__init__.py b/primus/backends/megatron/patches/training_log/__init__.py index 8935e1158..14d08c3a0 100644 --- a/primus/backends/megatron/patches/training_log/__init__.py +++ b/primus/backends/megatron/patches/training_log/__init__.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. +# Copyright (c) 2026, Advanced Micro Devices, Inc. # # See LICENSE for license information. ############################################################################### @@ -11,16 +11,18 @@ - print_rank_last_patches: scoped print_rank_last hook for training_log that injects ROCm memory and throughput statistics. + - wall_clock_timer_patch: wraps train_step with a wall-clock timer for + NeMo-comparable throughput measurement. -The actual patch registration is handled by ``print_rank_last_patches`` via -``@register_patch``; this ``__init__`` exists mainly to make -``training_log`` a proper package so that the auto-import logic in -``primus.backends.megatron.patches.__init__`` can discover and import -``print_rank_last_patches`` automatically. +The actual patch registration is handled via ``@register_patch``; this +``__init__`` exists mainly to make ``training_log`` a proper package so +that the auto-import logic in ``primus.backends.megatron.patches.__init__`` +can discover and import the patch modules automatically. """ from primus.backends.megatron.patches.training_log import ( # noqa: F401 print_rank_last_patches, + wall_clock_timer_patch, ) -__all__ = ["print_rank_last_patches"] +__all__ = ["print_rank_last_patches", "wall_clock_timer_patch"] diff --git a/primus/backends/megatron/patches/training_log/print_rank_last_patches.py b/primus/backends/megatron/patches/training_log/print_rank_last_patches.py index db0cedf04..3a1553b3d 100644 --- a/primus/backends/megatron/patches/training_log/print_rank_last_patches.py +++ b/primus/backends/megatron/patches/training_log/print_rank_last_patches.py @@ -1,5 +1,5 @@ ############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. # # See LICENSE for license information. ############################################################################### @@ -13,7 +13,12 @@ - ROCm/HIP memory stats. - Running average elapsed time per iteration (ms). - Running average throughput per GPU (TFLOP/s/GPU). - - Running average token throughput per GPU (tokens/s/GPU). + - Running average token throughput per GPU (tokens/s/GPU) (language models only). + - Diffusion-specific metrics (diffusion models only): + * Images per GPU (images/s/GPU): instant/average + * Latency per image (ms): instant + * Image resolution: height x width + * Average timestep Design: - We first parse Megatron's original ``log_string`` into a structured @@ -35,6 +40,40 @@ from primus.core.utils.rocm_mem_info import get_rocm_smi_mem_info +def _is_diffusion_model(args: Any, module_config: Any = None) -> bool: + """ + Detect if this is a diffusion model training run. + + Checks: + 1. model_type == 'diffusion_model' (from args/params) + 2. trainer_class contains 'Flux' or 'Diffusion' (from module_config) + + Args: + args: Megatron args (module_config.params) + module_config: Full module config (for trainer_class access) + + Returns: + True if diffusion model, False otherwise + """ + # Check model_type from args + model_type = getattr(args, "model_type", None) + if model_type == "diffusion_model": + return True + + # Check trainer_class from module_config (preferred source) + if module_config: + trainer_class = getattr(module_config, "trainer_class", None) + if trainer_class and ("Flux" in str(trainer_class) or "Diffusion" in str(trainer_class)): + return True + + # Fallback: check trainer_class in args (shouldn't be there if in reserved_keys) + trainer_class = getattr(args, "trainer_class", None) + if trainer_class and ("Flux" in str(trainer_class) or "Diffusion" in str(trainer_class)): + return True + + return False + + @dataclass class TrainingLogInfo: """Structured view of Megatron's training_log output line.""" @@ -250,15 +289,20 @@ def inject( # When pipeline parallelism (PP) is enabled, memory usage can vary across ranks. # Therefore, we report the maximum ROCm memory usage across all ranks. - r_used_tensor = torch.tensor([r_used], device="cuda", dtype=torch.int64) - world_size = torch.distributed.get_world_size() - gathered_r_used = [torch.zeros_like(r_used_tensor) for _ in range(world_size)] - torch.distributed.all_gather(gathered_r_used, r_used_tensor) - - total_r_used = [t.item() for t in gathered_r_used] - log_rank_0(f"total_r_used: {[round(r_used / 1024 ** 3, 2) for r_used in total_r_used]}") - max_r_used = max(total_r_used) - max_rank = total_r_used.index(max_r_used) + # Use constant-size all_reduce(MAX) instead of O(world_size) all_gather: + # one reduce for the max value, a second to recover the owning rank + # (the rank tensor is masked to -1 on non-max ranks). On ties the + # highest such rank wins (vs. the lowest under the previous all_gather). + max_used_tensor = torch.tensor([r_used], device="cuda", dtype=torch.int64) + torch.distributed.all_reduce(max_used_tensor, op=torch.distributed.ReduceOp.MAX) + max_r_used = max_used_tensor.item() + + my_rank = torch.distributed.get_rank() + rank_tensor = torch.tensor( + [my_rank if r_used == max_r_used else -1], device="cuda", dtype=torch.int64 + ) + torch.distributed.all_reduce(rank_tensor, op=torch.distributed.ReduceOp.MAX) + max_rank = rank_tensor.item() rocm_mem_str = ( f" | rocm mem usage/free/total/usage_ratio: " @@ -388,6 +432,38 @@ def __init__(self, args: Any): f"log_avg_reset_interval: {self._log_avg_reset_interval}" ) + def _inject_tflops(self, parsed: TrainingLogInfo) -> None: + """ + Shared TFLOP throughput logic (extracted for reuse by diffusion extension). + + Updates the throughput segment with running-average TFLOP throughput. + + Args: + parsed: Parsed training log information + """ + if parsed.throughput_tflops is not None: + tflops_value = parsed.throughput_tflops + iteration = parsed.iteration + + # Handle warmup & sliding window logic for TFLOPs. + if iteration is not None and ( + iteration == self._log_avg_skip_iterations + 1 + or len(self._recent_tflop_throughputs) >= self._log_avg_reset_interval + ): + self._recent_tflop_throughputs.clear() + + # Only accumulate after skip window. + if iteration is None or iteration > self._log_avg_skip_iterations: + self._recent_tflop_throughputs.append(tflops_value) + + if self._recent_tflop_throughputs: + avg_tflops = sum(self._recent_tflop_throughputs) / len(self._recent_tflop_throughputs) + idx = parsed.throughput_index + if idx is not None and 0 <= idx < len(parsed.segments): + parsed.segments[idx] = ( + f"throughput per GPU (TFLOP/s/GPU): {tflops_value:.1f}/{avg_tflops:.1f}" + ) + def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> str: """ Update ``parsed`` with running-average TFLOP and token throughput. @@ -406,27 +482,7 @@ def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> s iteration = parsed.iteration # ---------------- TFLOPs ---------------- - if parsed.throughput_tflops is not None: - tflops_value = parsed.throughput_tflops - - # Handle warmup & sliding window logic for TFLOPs. - if iteration is not None and ( - iteration == self._log_avg_skip_iterations + 1 - or len(self._recent_tflop_throughputs) >= self._log_avg_reset_interval - ): - self._recent_tflop_throughputs.clear() - - # Only accumulate after skip window. - if iteration is None or iteration > self._log_avg_skip_iterations: - self._recent_tflop_throughputs.append(tflops_value) - - if self._recent_tflop_throughputs: - avg_tflops = sum(self._recent_tflop_throughputs) / len(self._recent_tflop_throughputs) - idx = parsed.throughput_index - if idx is not None and 0 <= idx < len(parsed.segments): - parsed.segments[idx] = ( - f"compute per GPU (TFLOP/s/GPU): {tflops_value:.1f} (avg {avg_tflops:.1f})" - ) + self._inject_tflops(parsed) # ---------------- Tokens/s ---------------- if parsed.elapsed_ms is None: @@ -494,6 +550,219 @@ def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> s return log_string +class DiffusionThroughputAverageExtension(ThroughputAverageExtension): + """ + Helper extension for diffusion models: TFLOP throughput only, no tokens. + + Inherits from ThroughputAverageExtension to reuse TFLOP logic. + Skips token throughput calculation entirely (diffusion models use images, not tokens). + + Semantics mirror ThroughputAverageExtension: + - Ignore the first `log_avg_skip_iterations` iterations for averaging. + - Maintain a sliding window up to `log_avg_reset_interval` entries. + """ + + def __init__(self, args: Any): + super().__init__(args) + log_rank_0( + f"[Patch:megatron.training_log] DiffusionThroughputAverageExtension initialized " + f"(TFLOP only, no tokens)" + ) + + def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> str: + """ + Update ``parsed`` with running-average TFLOP throughput only. + + For diffusion models: TFLOPs only, no token throughput. + """ + try: + # If no parsed info is provided, keep the original string unchanged. + if parsed is None: + return log_string + + # Only inject TFLOP throughput (no tokens for diffusion models) + self._inject_tflops(parsed) + + # String result is ignored by the main patch when parsed is provided. + return log_string + except Exception: + # Any parsing / numeric issues should not break logging. + return log_string + + +class DiffusionMetricsExtension: + """ + Helper extension to compute and inject diffusion-specific metrics + (images per second per GPU and latency per image) into Megatron training logs. + + This extension only activates for diffusion models (model_type == 'diffusion_model'). + For language models, it early-returns without modifying logs. + + Semantics mirror ThroughputAverageExtension: + - Ignore the first `log_avg_skip_iterations` iterations for averaging. + - Maintain a sliding window up to `log_avg_reset_interval` entries. + """ + + def __init__(self, args: Any, module_config: Any = None, runtime_state: Any = None): + self._args = args + # Store module_config reference for diffusion detection + # (We only access trainer_class from it, which is a top-level attribute, not the nested params) + self._module_config = module_config + self._runtime_state = runtime_state + # Cache world_size once at construction time + self._world_size = getattr(args, "world_size", None) + # Track image throughput statistics across calls + self._recent_image_throughputs: list[float] = [] + # We follow the same warmup/reset semantics as ThroughputAverageExtension: + # - Ignore the first `log_avg_skip_iterations` iterations for averaging + # - Maintain a sliding window of size `log_avg_reset_interval` + self._log_avg_skip_iterations: int = int(getattr(args, "log_avg_skip_iterations", 0)) + self._log_avg_reset_interval: int = int(getattr(args, "log_avg_reset_interval", 1000)) + + def _calculate_image_metrics(self, parsed: TrainingLogInfo) -> Optional[tuple[float, float]]: + """ + Calculate image throughput metrics from parsed log info. + + Args: + parsed: Parsed training log information + + Returns: + Tuple of (images_per_second, latency_per_image_ms) or None if calculation fails + """ + if parsed.elapsed_ms is None or parsed.global_batch_size is None or self._world_size is None: + return None + + batch_size = int(parsed.global_batch_size) + elapsed_ms = float(parsed.elapsed_ms) + elapsed_s = elapsed_ms / 1000.0 + + if elapsed_s <= 0: + return None + + # Calculate images per second per GPU + images_per_second = batch_size / elapsed_s / self._world_size + + # Calculate latency per image (in milliseconds) + latency_per_image_ms = elapsed_ms / batch_size + + return (images_per_second, latency_per_image_ms) + + def _format_diffusion_metrics( + self, images_per_second: float, latency_per_image_ms: float, avg_images: float + ) -> list[str]: + """ + Format diffusion metrics as log segments. + + Args: + images_per_second: Current images per second per GPU + latency_per_image_ms: Current latency per image in milliseconds + avg_images: Average images per second per GPU + + Returns: + List of metric strings to append to log segments + """ + metrics = [] + + # Add images per GPU metrics (no trailing |, render function adds it) + images_metric = f"images per GPU (images/s/GPU): {images_per_second:.2f}/" f"{avg_images:.2f}" + metrics.append(images_metric) + + # Add latency per image metric (no trailing |, render function adds it) + latency_metric = f"latency per image (ms): {latency_per_image_ms:.1f}" + metrics.append(latency_metric) + + # Get metrics from runtime_state (required, no fallback) + last_metrics = None + if self._runtime_state: + last_metrics = self._runtime_state.last_metrics + else: + # Defensive: log warning but don't break logging + log_rank_0( + "[DiffusionMetricsExtension] WARNING: runtime_state not available, skipping diffusion metrics" + ) + return metrics # Return existing metrics without adding diffusion-specific ones + + if last_metrics: + if "image_height" in last_metrics and "image_width" in last_metrics: + resolution_metric = ( + f"image resolution: " + f"{int(last_metrics['image_height'])}x{int(last_metrics['image_width'])}" + ) + metrics.append(resolution_metric) + if "avg_timestep" in last_metrics: + timestep_metric = f"avg timestep: {last_metrics['avg_timestep']:.1f}" + metrics.append(timestep_metric) + + # Wall-clock step timer (from wall_clock_timer_patch) + if "wall_clock_step_ms" in last_metrics and self._world_size: + wc_ms = float(last_metrics["wall_clock_step_ms"]) + metrics.append(f"wall clock (ms): {wc_ms:.1f}") + gbs = getattr(self._args, "global_batch_size", None) + if gbs and wc_ms > 0: + wc_img_per_s = int(gbs) / (wc_ms / 1000.0) / self._world_size + metrics.append(f"wall clock img/s/GPU: {wc_img_per_s:.2f}") + + return metrics + + def inject(self, log_string: str, parsed: Optional[TrainingLogInfo] = None) -> str: + """ + Update ``parsed`` with images per second and latency per image metrics. + + Only activates for diffusion models (model_type == 'diffusion_model'). + For other models, early-returns without modification. + + Metrics: + - images per GPU (images/s/GPU): instant/average + - latency per image (ms): instant + """ + try: + # If no parsed info is provided, keep the original string unchanged + if parsed is None: + return log_string + + # Early return for non-diffusion models + if not _is_diffusion_model(self._args, self._module_config): + return log_string + + # Calculate image metrics + metrics_result = self._calculate_image_metrics(parsed) + if metrics_result is None: + return log_string + + images_per_second, latency_per_image_ms = metrics_result + iteration = parsed.iteration + + # Handle warmup & sliding window logic for images + if iteration is not None and ( + iteration == self._log_avg_skip_iterations + 1 + or len(self._recent_image_throughputs) >= self._log_avg_reset_interval + ): + self._recent_image_throughputs.clear() + + # Only accumulate after skip window + if iteration is None or iteration > self._log_avg_skip_iterations: + self._recent_image_throughputs.append(images_per_second) + + if self._recent_image_throughputs: + avg_images = sum(self._recent_image_throughputs) / len(self._recent_image_throughputs) + + # Format and append metrics + metrics = self._format_diffusion_metrics(images_per_second, latency_per_image_ms, avg_images) + parsed.segments.extend(metrics) + + # String result is ignored by the main patch when parsed is provided. + return log_string + except Exception as e: + # Log the exception to help debug, but don't break training + iteration = parsed.iteration if parsed else None + log_rank_0( + f"[Patch:megatron.training_log] DiffusionMetricsExtension ERROR " + f"(iteration={iteration}): {type(e).__name__}: {e}" + ) + # Any parsing / numeric issues should not break logging. + return log_string + + @register_patch( "megatron.training_log.unified_patch", backend="megatron", @@ -519,6 +788,9 @@ def patch_training_log_unified(ctx: PatchContext): # Get unified Megatron args (module_config.params) from context. config = get_args(ctx) + # Get runtime_state from context + runtime_state = ctx.extra.get("runtime_state") + # Check whether we should enable ROCm stats / throughput logging. use_rocm_mem = bool(getattr(config, "use_rocm_mem_info", False)) rocm_iters = getattr(config, "use_rocm_mem_info_iters", []) @@ -541,14 +813,25 @@ def patch_training_log_unified(ctx: PatchContext): if getattr(original_training_log, "_primus_training_log_print_rank_wrapper", False): return - # Create helper extensions only when stat injection is enabled so they - # keep state (ROCm cache, avg windows) across all training_log - # invocations. In forwarding-only mode they are not needed. - mem_ext = elapsed_ext = throughput_ext = None - if enable_rocm_stats: - mem_ext = MemoryStatsExtension(config) - elapsed_ext = ElapsedAverageExtension(config) + # Create helper extensions once so they keep state (ROCm cache, avg windows) + # across all training_log invocations. + # Get module_config from context to pass to diffusion extensions + # (needed to access trainer_class which is in reserved_keys, not in params) + module_config = ctx.extra.get("module_config") + + # Detect if this is a diffusion model to instantiate the correct throughput extension + is_diffusion = _is_diffusion_model(config, module_config) + + mem_ext = MemoryStatsExtension(config) + elapsed_ext = ElapsedAverageExtension(config) + # Use diffusion-specific throughput extension if this is a diffusion model + if is_diffusion: + throughput_ext = DiffusionThroughputAverageExtension(config) + else: throughput_ext = ThroughputAverageExtension(config) + diffusion_ext = DiffusionMetricsExtension( + config, module_config=module_config, runtime_state=runtime_state + ) call_count = 0 # Capture the original ``print_rank_last`` so we can delegate actual # printing back to Megatron after mutating the log string. @@ -579,12 +862,12 @@ def primus_print_rank_last(log_string: str) -> None: # Parse the original log string once and share across extensions. parsed = parse_training_log_line(log_string) - # Inject memory statistics, elapsed avg, throughput, and token - # throughput by mutating the parsed structure. These calls ignore - # their string return value when `parsed` is provided. + # Inject memory statistics, elapsed avg, throughput, and diffusion + # metrics by mutating the parsed structure. mem_ext.inject(log_string, call_count, parsed) elapsed_ext.inject(log_string, parsed) throughput_ext.inject(log_string, parsed) + diffusion_ext.inject(log_string, parsed) # Render the final line from the parsed structure. updated = render_training_log_line(parsed) diff --git a/primus/backends/megatron/patches/training_log/wall_clock_timer_patch.py b/primus/backends/megatron/patches/training_log/wall_clock_timer_patch.py new file mode 100644 index 000000000..5febc1cd3 --- /dev/null +++ b/primus/backends/megatron/patches/training_log/wall_clock_timer_patch.py @@ -0,0 +1,85 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Wall-clock timer patch for Megatron train_step. + +Wraps ``megatron.training.training.train_step`` with ``time.perf_counter()`` +to measure the actual forward + backward + optimizer wall-clock duration, +independent of Megatron's ``interval-time`` timer which includes collective +barriers, CUDA synchronizes, and logging overhead. + +The measured duration is stored on ``runtime_state.last_metrics`` so that +downstream log extensions (e.g. ``DiffusionMetricsExtension``) can surface +it alongside existing throughput metrics. +""" + +import time + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + + +def _wall_clock_timer_enabled(ctx: PatchContext) -> bool: + args = get_args(ctx) + return args is not None and getattr(args, "wall_clock_step_timer", False) + + +@register_patch( + "megatron.training.wall_clock_step_timer", + backend="megatron", + phase="before_train", + description="Wrap train_step with wall-clock timer for NeMo-comparable throughput measurement", + condition=_wall_clock_timer_enabled, + priority=90, +) +def patch_wall_clock_timer(ctx: PatchContext): + import megatron.training.training as megatron_training + + from primus.backends.megatron.patches._patch_guard import is_patched, mark_patched + + _PATCH_KEY = "megatron.training.train_step_wall_clock_timer" + if is_patched(megatron_training, _PATCH_KEY): + log_rank_0("[Patch:wall_clock_step_timer] Already applied; skipping re-wrap.") + return + + runtime_state = ctx.extra.get("runtime_state") + + _original_train_step = megatron_training.train_step + + def _patched_train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=None, + ): + t0 = time.perf_counter() + result = _original_train_step( + forward_step_func, + data_iterator, + model, + optimizer, + opt_param_scheduler, + config, + forward_backward_func, + iteration=iteration, + ) + dt_ms = (time.perf_counter() - t0) * 1000.0 + if runtime_state is not None: + runtime_state.update_metrics({"wall_clock_step_ms": dt_ms}) + return result + + megatron_training.train_step = _patched_train_step + mark_patched(megatron_training, _PATCH_KEY) + log_rank_0( + "[Patch:wall_clock_step_timer] " + "Patched train_step with wall-clock timer " + f"(runtime_state={'available' if runtime_state else 'unavailable'})" + ) diff --git a/tests/unit_tests/backends/megatron/diffusion/test_mlperf_warmup_fp8_state.py b/tests/unit_tests/backends/megatron/diffusion/test_mlperf_warmup_fp8_state.py new file mode 100644 index 000000000..99bc03c96 --- /dev/null +++ b/tests/unit_tests/backends/megatron/diffusion/test_mlperf_warmup_fp8_state.py @@ -0,0 +1,281 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +GPU unit tests for FP8 delayed-scaling state reset during MLPerf warmup. + +Validates that ``_reset_fp8_local_spec`` correctly breaks buffer pointers, +triggering ``_DelayedScalingRegistry`` reinitialisation with +``_first_step = True`` so that the first real step bootstraps weight amaxes +from the restored (pre-warmup) weights. + +Requires a CUDA device; skipped automatically when unavailable. + +Run: + python -m pytest tests/unit_tests/backends/megatron/diffusion/test_mlperf_warmup_fp8_state.py -v +""" + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + +DEVICE = "cuda:0" +FP8_FWD_MAX = torch.finfo(torch.float8_e4m3fn).max +FP8_BWD_MAX = torch.finfo(torch.float8_e5m2).max + + +class _FakeDelayedModule(torch.nn.Module): + """Minimal stand-in for a Float8*ParallelLinear with delayed scaling.""" + + def __init__(self, in_features=64, out_features=64, history_len=1): + super().__init__() + self.weight = torch.nn.Parameter( + torch.randn(out_features, in_features, dtype=torch.bfloat16, device=DEVICE) + ) + self._use_delayed_scaling = True + self._fp8_fwd_dtype = torch.float8_e4m3fn + self._fp8_bwd_dtype = torch.float8_e5m2 + self._fp8_fwd_max = FP8_FWD_MAX + self._fp8_bwd_max = FP8_BWD_MAX + self._amax_compute_algo = "most_recent" + self._first_delayed_step = True + self._history_idx = 0 + self.config = type( + "Config", + (), + { + "fp8_amax_history_len": history_len, + "fp8_amax_compute_algo": "most_recent", + }, + )() + + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _init_delayed_scaling_state, + ) + + _init_delayed_scaling_state(self) + for name in list(self._buffers): + buf = self._buffers[name] + if buf is not None and buf.device.type == "cpu": + self._buffers[name] = buf.to(DEVICE) + + +class TestResetFp8LocalSpecOnDevice: + """Verify _reset_fp8_local_spec creates new buffers on the correct device.""" + + def test_buffers_on_device_after_reset(self): + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_fp8_local_spec, + ) + + module = _FakeDelayedModule() + model = torch.nn.Sequential(module) + + count = _reset_fp8_local_spec([model]) + assert count == 1 + + for buf_name in ( + "scale_input", + "scale_weight", + "scale_grad", + "amax_history_input", + "amax_history_weight", + "amax_history_grad", + "staged_input_amax", + "staged_grad_amax", + "staged_weight_amax", + ): + buf = module._buffers.get(buf_name) + assert buf is not None, f"Buffer {buf_name} missing after reset" + assert ( + buf.device == module.weight.device + ), f"Buffer {buf_name} on {buf.device}, expected {module.weight.device}" + + def test_first_delayed_step_reset(self): + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_fp8_local_spec, + ) + + module = _FakeDelayedModule() + module._first_delayed_step = False + + _reset_fp8_local_spec([torch.nn.Sequential(module)]) + + assert module._first_delayed_step is True + + +class TestRegistryPointerBreak: + """Verify that reset breaks pointer identity with registry global tensors.""" + + def test_pointer_mismatch_triggers_reinit(self): + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _DelayedScalingRegistry, + ) + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_fp8_local_spec, + ) + + modules = [_FakeDelayedModule() for _ in range(4)] + model = torch.nn.Sequential(*modules) + registry = _DelayedScalingRegistry(modules) + + # The detection-via-aliasing check in _fast_update_scales_with_history + # is keyed on amax_history (the only registry tensor that the modules + # hold views into); per-module scale_* buffers have always been + # independent scalar tensors. Verify the alias holds pre-reset and is + # broken post-reset. + old_registry_amax_ptr = registry.amax_history.untyped_storage().data_ptr() + old_mod0_amax_ptr = modules[0].amax_history_input.untyped_storage().data_ptr() + assert ( + old_registry_amax_ptr == old_mod0_amax_ptr + ), "Sanity: registry.amax_history must alias each module's amax_history_input" + + _reset_fp8_local_spec([model]) + + new_mod0_amax_ptr = modules[0].amax_history_input.untyped_storage().data_ptr() + assert new_mod0_amax_ptr != old_registry_amax_ptr, ( + "After reset, module amax_history_input should have different storage " + "than the original registry.amax_history (so the next " + "_fast_update_scales* call triggers registry.__init__ via the " + "data_ptr mismatch check)" + ) + + +class TestRegistryReinitialisationOnFirstStep: + """End-to-end: reset → _fast_update_scales detects mismatch → re-creates registry.""" + + def test_fast_path_reinit(self): + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _DelayedScalingRegistry, + _fast_update_scales, + ) + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_fp8_local_spec, + ) + + modules = [_FakeDelayedModule(history_len=1) for _ in range(4)] + model = torch.nn.Sequential(*modules) + registry = _DelayedScalingRegistry(modules) + + _fast_update_scales(registry) + assert registry._first_step is False + + new_weights = [torch.randn_like(m.weight.data) for m in modules] + for m, w in zip(modules, new_weights): + m.weight.data.copy_(w) + + _reset_fp8_local_spec([model]) + _fast_update_scales(registry) + + assert registry._first_step is False, "Should have consumed _first_step" + for i, m in enumerate(modules): + expected_amax = new_weights[i].abs().amax().float().item() + # _fast_update_scales writes per-module weight amaxes to + # m.staged_weight_amax in the _first_step bootstrap path; the + # registry-batched staged_amaxes_3n is only populated by + # _fast_update_scales_with_history, so read from the module here. + actual_amax = m.staged_weight_amax.item() + assert ( + abs(actual_amax - expected_amax) < 1e-3 + ), f"Module {i}: weight amax {actual_amax} != expected {expected_amax}" + + def test_history_path_reinit(self): + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _DelayedScalingRegistry, + _fast_update_scales_with_history, + ) + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_fp8_local_spec, + ) + + H = 16 + modules = [_FakeDelayedModule(history_len=H) for _ in range(4)] + model = torch.nn.Sequential(*modules) + registry = _DelayedScalingRegistry(modules) + + _fast_update_scales_with_history(registry) + assert registry._first_step is False + + new_weights = [torch.randn_like(m.weight.data) for m in modules] + for m, w in zip(modules, new_weights): + m.weight.data.copy_(w) + + _reset_fp8_local_spec([model]) + _fast_update_scales_with_history(registry) + + assert registry._first_step is False + for i, m in enumerate(modules): + expected_amax = new_weights[i].abs().amax().float().item() + # _fast_update_scales_with_history mirrors the staged amaxes into + # registry.staged_amaxes_3n (rows 0/1/2 = input/weight/grad), so + # both registry.staged_amaxes_3n[1, i] and m.staged_weight_amax + # carry the value. Read from the module for consistency with the + # fast-path test above. + actual_amax = m.staged_weight_amax.item() + assert ( + abs(actual_amax - expected_amax) < 1e-3 + ), f"Module {i}: weight amax {actual_amax} != expected {expected_amax}" + + +class TestScaleLeakagePrevention: + """Verify that warmup-phase scales do not leak into post-reset state.""" + + def test_scales_recomputed_from_clean_state(self): + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _DelayedScalingRegistry, + _fast_update_scales, + ) + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_fp8_local_spec, + ) + + modules = [_FakeDelayedModule(history_len=1) for _ in range(4)] + model = torch.nn.Sequential(*modules) + registry = _DelayedScalingRegistry(modules) + + # Staging is now per-module: registry batches them into + # registry.staged_amaxes_3n inside _fast_update_scales_with_history, + # but _fast_update_scales reads directly from the per-module buffers. + # _fast_update_scales also scatters new scales back to per-module + # m.scale_* buffers (registry.scales_3n is no longer the + # source-of-truth for the fast path), so we read scales from there. + for _ in range(5): + for i, m in enumerate(modules): + m.staged_input_amax.fill_(100.0 + i) + m.staged_grad_amax.fill_(50.0 + i) + _fast_update_scales(registry) + + warmup_scales = torch.stack( + [ + torch.stack([m.scale_input.clone() for m in modules]), + torch.stack([m.scale_weight.clone() for m in modules]), + torch.stack([m.scale_grad.clone() for m in modules]), + ] + ) + assert (warmup_scales != 1.0).any(), "Scales should have changed during warmup" + + new_weights = [torch.randn_like(m.weight.data) * 0.01 for m in modules] + for m, w in zip(modules, new_weights): + m.weight.data.copy_(w) + + _reset_fp8_local_spec([model]) + _fast_update_scales(registry) + + post_reset_scales = torch.stack( + [ + torch.stack([m.scale_input.clone() for m in modules]), + torch.stack([m.scale_weight.clone() for m in modules]), + torch.stack([m.scale_grad.clone() for m in modules]), + ] + ) + assert not torch.equal( + warmup_scales, post_reset_scales + ), "Scales should differ after reset with new weights" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/test_mlperf_patches.py b/tests/unit_tests/backends/megatron/test_mlperf_patches.py new file mode 100644 index 000000000..077db6d25 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_mlperf_patches.py @@ -0,0 +1,795 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for MLPerf logging and warmup patches. + +Covers the production patch behavior (CPU-only): logging-patch monkey-patching +and idempotency, the INIT_STOP/RUN_START emission on the first post-warmup +training_log call, convergence detection, and the FP8/optimizer warmup helper +functions (_reset_fp8_te_spec, seed FP8 amax, optimizer neuter/restore/reset). +""" + +import types +from types import SimpleNamespace +from unittest.mock import MagicMock + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _install_fake_megatron(monkeypatch): + """Install a fake megatron.training.training module into sys.modules.""" + import sys + + megatron_mod = types.ModuleType("megatron") + training_pkg = types.ModuleType("megatron.training") + training_mod = types.ModuleType("megatron.training.training") + global_vars_mod = types.ModuleType("megatron.training.global_vars") + + def fake_train_step(fwd, data_iter, model, optimizer, sched, config, fwdbwd, iteration=None): + return {}, 0, False, False, 0, 0.0, 0, None + + def fake_training_log(*args, **kwargs): + return None + + def fake_eval(*args, **kwargs): + return None + + def fake_evaluate(*args, **kwargs): + return ({},) + + def fake_print_rank_last(msg): + pass + + def fake_get_tb_writer(): + return None + + def fake_get_wandb_writer(): + return None + + training_mod.train_step = fake_train_step + training_mod.training_log = fake_training_log + training_mod.evaluate_and_print_results = fake_eval + training_mod.evaluate = fake_evaluate + training_mod.print_rank_last = fake_print_rank_last + training_mod.get_tensorboard_writer = fake_get_tb_writer + training_mod.get_wandb_writer = fake_get_wandb_writer + + training_pkg.training = training_mod + megatron_mod.training = training_pkg + + # Store megatron args for get_args + _megatron_args = SimpleNamespace( + iteration=0, + curr_iteration=0, + consumed_train_samples=0, + skipped_train_samples=0, + train_iters=100, + eval_interval=10, + do_valid=True, + global_batch_size=512, + micro_batch_size=64, + ) + global_vars_mod.get_args = lambda: _megatron_args + + training_pkg.get_args = global_vars_mod.get_args + + monkeypatch.setitem(sys.modules, "megatron", megatron_mod) + monkeypatch.setitem(sys.modules, "megatron.training", training_pkg) + monkeypatch.setitem(sys.modules, "megatron.training.training", training_mod) + monkeypatch.setitem(sys.modules, "megatron.training.global_vars", global_vars_mod) + + return training_mod, _megatron_args + + +def _make_ctx( + mlperf_mode=False, + warmup_train_steps=0, + target_val_loss=0.586, + **kwargs, +): + """Build a minimal PatchContext-like object.""" + params = SimpleNamespace( + mlperf_mode=mlperf_mode, + warmup_train_steps=warmup_train_steps, + target_val_loss=target_val_loss, + global_batch_size=kwargs.get("global_batch_size", 512), + micro_batch_size=kwargs.get("micro_batch_size", 64), + seed=kwargs.get("seed", 42), + log_interval=kwargs.get("log_interval", 10), + lr=kwargs.get("lr", 2e-4), + adam_beta1=kwargs.get("adam_beta1", 0.9), + adam_beta2=kwargs.get("adam_beta2", 0.95), + adam_eps=kwargs.get("adam_eps", 1e-8), + weight_decay=kwargs.get("weight_decay", 0.1), + image_size=kwargs.get("image_size", 256), + vae_latent_mode=kwargs.get("vae_latent_mode", "resample"), + transformer_impl=kwargs.get("transformer_impl", "local"), + use_fsdp2_fp8_all_gather=kwargs.get("use_fsdp2_fp8_all_gather", False), + wall_clock_step_timer=False, + **{ + k: v + for k, v in kwargs.items() + if k + not in ( + "global_batch_size", + "micro_batch_size", + "seed", + "log_interval", + "lr", + "adam_beta1", + "adam_beta2", + "adam_eps", + "weight_decay", + "image_size", + "vae_latent_mode", + "transformer_impl", + "use_fsdp2_fp8_all_gather", + ) + }, + ) + module_config = SimpleNamespace(params=params) + return SimpleNamespace( + extra={"module_config": module_config}, + backend="megatron", + phase="before_train", + ) + + +# ============================================================================ +# Level 1: Patch registration and conditions +# ============================================================================ + + +class TestLoggingPatchMonkeyPatching: + """Verify that the logging patch replaces the expected functions.""" + + def test_installs_wrappers(self, monkeypatch): + mt, _ = _install_fake_megatron(monkeypatch) + + monkeypatch.setattr( + "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0", + lambda *a, **k: None, + ) + + mock_mllog = MagicMock() + mock_mllog.get_mllogger.return_value = MagicMock() + mock_mllog.constants = SimpleNamespace( + INIT_START="init_start", + INIT_STOP="init_stop", + RUN_START="run_start", + RUN_STOP="run_stop", + SUBMISSION_BENCHMARK="submission_benchmark", + SUBMISSION_ORG="submission_org", + SUBMISSION_DIVISION="submission_division", + SUBMISSION_PLATFORM="submission_platform", + SUBMISSION_STATUS="submission_status", + SEED="seed", + GLOBAL_BATCH_SIZE="global_batch_size", + TRAIN_SAMPLES="train_samples", + EVAL_SAMPLES="eval_samples", + GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps", + OPT_NAME="opt_name", + OPT_BASE_LR="opt_base_lr", + EVAL_ACCURACY="eval_accuracy", + EVAL_START="eval_start", + EVAL_STOP="eval_stop", + EPOCH_START="epoch_start", + BLOCK_START="block_start", + BLOCK_STOP="block_stop", + ) + mock_mlperf_pkg = MagicMock() + mock_mlperf_pkg.mllog = mock_mllog + monkeypatch.setitem(__import__("sys").modules, "mlperf_logging", mock_mlperf_pkg) + monkeypatch.setitem(__import__("sys").modules, "mlperf_logging.mllog", mock_mllog) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + original_tl = mt.training_log + original_eval = mt.evaluate_and_print_results + original_prl = mt.print_rank_last + + ctx = _make_ctx(mlperf_mode=True) + patch_mlperf_logging(ctx) + + assert mt.training_log is not original_tl + assert mt.evaluate_and_print_results is not original_eval + assert mt.print_rank_last is not original_prl + assert getattr(mt, "_primus_mlperf_logging_installed", False) is True + + def test_idempotent(self, monkeypatch): + mt, _ = _install_fake_megatron(monkeypatch) + + monkeypatch.setattr( + "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0", + lambda *a, **k: None, + ) + + mock_mllog = MagicMock() + mock_mllog.get_mllogger.return_value = MagicMock() + mock_mllog.constants = SimpleNamespace( + INIT_START="init_start", + INIT_STOP="init_stop", + RUN_START="run_start", + RUN_STOP="run_stop", + SUBMISSION_BENCHMARK="submission_benchmark", + SUBMISSION_ORG="submission_org", + SUBMISSION_DIVISION="submission_division", + SUBMISSION_PLATFORM="submission_platform", + SUBMISSION_STATUS="submission_status", + SEED="seed", + GLOBAL_BATCH_SIZE="global_batch_size", + TRAIN_SAMPLES="train_samples", + EVAL_SAMPLES="eval_samples", + GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps", + OPT_NAME="opt_name", + OPT_BASE_LR="opt_base_lr", + EVAL_ACCURACY="eval_accuracy", + EVAL_START="eval_start", + EVAL_STOP="eval_stop", + EPOCH_START="epoch_start", + BLOCK_START="block_start", + BLOCK_STOP="block_stop", + ) + mock_mlperf_pkg = MagicMock() + mock_mlperf_pkg.mllog = mock_mllog + monkeypatch.setitem(__import__("sys").modules, "mlperf_logging", mock_mlperf_pkg) + monkeypatch.setitem(__import__("sys").modules, "mlperf_logging.mllog", mock_mllog) + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + ctx = _make_ctx(mlperf_mode=True) + mt._primus_mlperf_logging_installed = False + + patch_mlperf_logging(ctx) + first_tl = mt.training_log + first_eval = mt.evaluate_and_print_results + + patch_mlperf_logging(ctx) + assert mt.training_log is first_tl + assert mt.evaluate_and_print_results is first_eval + + +class TestTrainingLogFirstCall: + """Verify INIT_STOP + RUN_START fire on first post-warmup training_log.""" + + def test_first_call_emits_init_stop_run_start(self, monkeypatch): + mt, args = _install_fake_megatron(monkeypatch) + + monkeypatch.setattr( + "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0", + lambda *a, **k: None, + ) + + emitted_events = [] + + mock_mllogger = MagicMock() + mock_constants = SimpleNamespace( + INIT_START="init_start", + INIT_STOP="init_stop", + RUN_START="run_start", + RUN_STOP="run_stop", + SUBMISSION_BENCHMARK="submission_benchmark", + SUBMISSION_ORG="submission_org", + SUBMISSION_DIVISION="submission_division", + SUBMISSION_PLATFORM="submission_platform", + SUBMISSION_STATUS="submission_status", + SEED="seed", + GLOBAL_BATCH_SIZE="global_batch_size", + TRAIN_SAMPLES="train_samples", + EVAL_SAMPLES="eval_samples", + GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps", + OPT_NAME="opt_name", + OPT_BASE_LR="opt_base_lr", + EVAL_ACCURACY="eval_accuracy", + EVAL_START="eval_start", + EVAL_STOP="eval_stop", + EPOCH_START="epoch_start", + BLOCK_START="block_start", + BLOCK_STOP="block_stop", + ) + + def track_start(key, value=None, metadata=None): + emitted_events.append(("start", key)) + + def track_end(key, value=None, metadata=None): + emitted_events.append(("end", key)) + + def track_event(key, value=None, metadata=None): + emitted_events.append(("event", key)) + + mock_mllogger.start = track_start + mock_mllogger.end = track_end + mock_mllogger.event = track_event + + mock_mllog_module = MagicMock() + mock_mllog_module.get_mllogger.return_value = mock_mllogger + mock_mllog_module.constants = mock_constants + + # Wire the top-level mlperf_logging mock so that + # `from mlperf_logging import mllog` resolves correctly + mock_mlperf_pkg = MagicMock() + mock_mlperf_pkg.mllog = mock_mllog_module + + import sys as _sys + + monkeypatch.setitem(_sys.modules, "mlperf_logging", mock_mlperf_pkg) + monkeypatch.setitem(_sys.modules, "mlperf_logging.mllog", mock_mllog_module) + monkeypatch.setenv("RANK", "0") + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + ctx = _make_ctx(mlperf_mode=True, log_interval=1) + patch_mlperf_logging(ctx) + + emitted_events.clear() + + # First call should emit INIT_STOP + RUN_START + mt.training_log({"loss": 0.5}, {}, 1e-4, 1, 1.0, False, False, 0.0, None, 0, None) + + event_keys = [e[1] for e in emitted_events] + assert "init_stop" in event_keys, f"INIT_STOP not emitted. Events: {emitted_events}" + assert "run_start" in event_keys, f"RUN_START not emitted. Events: {emitted_events}" + + init_stop_idx = event_keys.index("init_stop") + run_start_idx = event_keys.index("run_start") + assert init_stop_idx < run_start_idx + + # Second call should NOT emit INIT_STOP/RUN_START again + emitted_events.clear() + mt.training_log({"loss": 0.4}, {}, 1e-4, 2, 1.0, False, False, 0.0, None, 0, None) + + event_keys_2 = [e[1] for e in emitted_events] + assert "init_stop" not in event_keys_2 + assert "run_start" not in event_keys_2 + + +# ============================================================================ +# Level 3: Component tests +# ============================================================================ + + +class TestConvergenceDetection: + """Verify convergence detection in evaluate_and_print_results wrapper.""" + + def test_convergence_sets_train_iters(self, monkeypatch): + mt, megatron_args = _install_fake_megatron(monkeypatch) + + monkeypatch.setattr( + "primus.backends.megatron.patches.mlperf_logging_patches.log_rank_0", + lambda *a, **k: None, + ) + + # Mock mlperf_logging — wire .mllog attribute on the top-level package + mock_mllogger = MagicMock() + mock_constants = SimpleNamespace( + INIT_START="init_start", + INIT_STOP="init_stop", + RUN_START="run_start", + RUN_STOP="run_stop", + SUBMISSION_BENCHMARK="submission_benchmark", + SUBMISSION_ORG="submission_org", + SUBMISSION_DIVISION="submission_division", + SUBMISSION_PLATFORM="submission_platform", + SUBMISSION_STATUS="submission_status", + SEED="seed", + GLOBAL_BATCH_SIZE="global_batch_size", + TRAIN_SAMPLES="train_samples", + EVAL_SAMPLES="eval_samples", + GRADIENT_ACCUMULATION_STEPS="gradient_accumulation_steps", + OPT_NAME="opt_name", + OPT_BASE_LR="opt_base_lr", + EVAL_ACCURACY="eval_accuracy", + EVAL_START="eval_start", + EVAL_STOP="eval_stop", + EPOCH_START="epoch_start", + BLOCK_START="block_start", + BLOCK_STOP="block_stop", + ) + mock_mllogger.start = MagicMock() + mock_mllogger.end = MagicMock() + mock_mllogger.event = MagicMock() + + mock_mllog_module = MagicMock() + mock_mllog_module.get_mllogger.return_value = mock_mllogger + mock_mllog_module.constants = mock_constants + + mock_mlperf_pkg = MagicMock() + mock_mlperf_pkg.mllog = mock_mllog_module + + import sys as _sys + + monkeypatch.setitem(_sys.modules, "mlperf_logging", mock_mlperf_pkg) + monkeypatch.setitem(_sys.modules, "mlperf_logging.mllog", mock_mllog_module) + monkeypatch.setenv("RANK", "0") + + # Set evaluate to return a loss below target BEFORE patching. + # Also make evaluate_and_print_results call evaluate() internally, + # mirroring real Megatron behavior so _captured_loss gets populated. + mt.evaluate = lambda *a, **k: ({"loss": 0.500},) + + def fake_eval_and_print(*a, **k): + mt.evaluate(*a, **k) + + mt.evaluate_and_print_results = fake_eval_and_print + + from primus.backends.megatron.patches.mlperf_logging_patches import ( + patch_mlperf_logging, + ) + + target_val_loss = 0.586 + megatron_args.train_iters = 5000 + + ctx = _make_ctx(mlperf_mode=True, target_val_loss=target_val_loss) + patch_mlperf_logging(ctx) + + # Call eval at iteration 512 + mt.evaluate_and_print_results( + "iteration 512", + lambda: None, + None, + [MagicMock()], + 512, + None, + MagicMock(), + ) + + assert megatron_args.train_iters == 512 + + +# ============================================================================ +# Level 4: Helper function unit tests (CPU-only, no GPU required) +# ============================================================================ + + +class TestResetFp8TeSpec: + """Verify _reset_fp8_te_spec clears fp8_initialized and meta tensors.""" + + def test_resets_fp8_initialized_flag(self): + import torch + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_fp8_te_spec, + ) + + class FakeTeModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.fp8_initialized = True + self.fp8_meta = { + "scaling_fwd": SimpleNamespace( + amax_history=torch.ones(16), + scale=torch.full((4,), 3.14), + scale_inv=torch.full((4,), 0.318), + ), + "scaling_bwd": SimpleNamespace( + amax_history=torch.ones(16), + scale=torch.full((4,), 2.71), + scale_inv=torch.full((4,), 0.369), + ), + } + + model = torch.nn.Sequential(FakeTeModule(), FakeTeModule()) + count = _reset_fp8_te_spec([model]) + + assert count == 2 + for module in model.modules(): + if hasattr(module, "fp8_initialized"): + assert module.fp8_initialized is False + for key in ("scaling_fwd", "scaling_bwd"): + tm = module.fp8_meta[key] + assert (tm.amax_history == 0.0).all() + assert (tm.scale == 1.0).all() + assert (tm.scale_inv == 1.0).all() + + def test_skips_reset_fp8_meta_tensors_shortcut(self): + """ + TE 2.8.0.dev0's `reset_fp8_meta_tensors` unconditionally derefs `.scale` + on the recipe state, which crashes on `Float8CurrentScalingRecipeState` + (current/tensorwise scaling has no persistent state). The reset must + therefore go through the recipe-agnostic manual path even when a TE + module advertises the helper. + """ + from types import SimpleNamespace + + import torch + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_fp8_te_spec, + ) + + reset_called = [False] + + class FakeTeCurrentScalingModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.fp8_initialized = True + # Mimic Float8CurrentScalingRecipeState: no .scale, no .amax_history. + self.fp8_meta = { + "scaling_fwd": SimpleNamespace(), + "scaling_bwd": SimpleNamespace(), + } + + def reset_fp8_meta_tensors(self): + reset_called[0] = True + + model = torch.nn.Sequential(FakeTeCurrentScalingModule()) + count = _reset_fp8_te_spec([model]) + + assert not reset_called[ + 0 + ], "reset_fp8_meta_tensors must NOT be called (would crash on current scaling)" + assert count == 1 + assert model[0].fp8_initialized is False + + def test_falls_through_for_delayed_scaling_buffers(self): + """ + Companion to ``test_skips_reset_fp8_meta_tensors_shortcut``: confirms + the manual fallback still resets delayed-scaling buffers (.scale, + .amax_history, .scale_inv) when they exist on the recipe state. + """ + from types import SimpleNamespace + + import torch + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_fp8_te_spec, + ) + + reset_called = [False] + + class FakeTeDelayedScalingModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.fp8_initialized = True + self.fp8_meta = { + "scaling_fwd": SimpleNamespace( + amax_history=torch.full((16,), 7.0), + scale=torch.full((4,), 2.71), + scale_inv=torch.full((4,), 0.369), + ), + "scaling_bwd": SimpleNamespace( + amax_history=torch.full((16,), 7.0), + scale=torch.full((4,), 2.71), + scale_inv=torch.full((4,), 0.369), + ), + } + + def reset_fp8_meta_tensors(self): + reset_called[0] = True + + model = torch.nn.Sequential(FakeTeDelayedScalingModule()) + count = _reset_fp8_te_spec([model]) + + assert not reset_called[0] + assert count == 1 + for key in ("scaling_fwd", "scaling_bwd"): + tm = model[0].fp8_meta[key] + assert (tm.amax_history == 0.0).all() + assert (tm.scale == 1.0).all() + assert (tm.scale_inv == 1.0).all() + + +class TestSeedFp8Amax: + """Verify _seed_fp8_amax fills amax_history with the requested seed value.""" + + def _make_te_model(self): + import torch + + class FakeTeModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.fp8_meta = { + "scaling_fwd": SimpleNamespace( + amax_history=torch.zeros(16), + ), + "scaling_bwd": SimpleNamespace( + amax_history=torch.zeros(16), + ), + } + + return torch.nn.Sequential(FakeTeModule(), FakeTeModule()) + + def test_seeds_default_value(self): + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _seed_fp8_amax, + ) + + model = self._make_te_model() + count = _seed_fp8_amax([model]) + + assert count == 4 + for module in model.modules(): + if hasattr(module, "fp8_meta"): + for key in ("scaling_fwd", "scaling_bwd"): + assert (module.fp8_meta[key].amax_history == 1.0).all() + + def test_seeds_custom_value(self): + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _seed_fp8_amax, + ) + + model = self._make_te_model() + count = _seed_fp8_amax([model], seed_value=42.0) + + assert count == 4 + for module in model.modules(): + if hasattr(module, "fp8_meta"): + for key in ("scaling_fwd", "scaling_bwd"): + assert (module.fp8_meta[key].amax_history == 42.0).all() + + +class TestNeuterRestoreOptimizer: + """Verify _neuter_optimizer / _restore_optimizer roundtrip.""" + + def test_roundtrip_preserves_hyperparams(self): + import torch + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _neuter_optimizer, + _restore_optimizer, + ) + + model = torch.nn.Linear(4, 4) + opt = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.95), weight_decay=0.1) + + orig_betas = list(opt.param_groups[0]["betas"]) + orig_wd = opt.param_groups[0]["weight_decay"] + + wrapper = SimpleNamespace(optimizer=opt) + saved = _neuter_optimizer(wrapper) + + assert opt.param_groups[0]["betas"] == [1.0, 1.0] + assert opt.param_groups[0]["weight_decay"] == 0.0 + + _restore_optimizer(wrapper, saved) + + assert list(opt.param_groups[0]["betas"]) == orig_betas + assert opt.param_groups[0]["weight_decay"] == orig_wd + + def test_roundtrip_all_keys(self): + """All 4 production keys roundtrip: betas, weight_decay, bias_correction, pre_mult_wd.""" + import torch + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _neuter_optimizer, + _restore_optimizer, + ) + + model = torch.nn.Linear(4, 4) + opt = torch.optim.SGD(model.parameters(), lr=0.01) + opt.param_groups[0]["betas"] = [0.9, 0.999] + opt.param_groups[0]["weight_decay"] = 0.01 + opt.param_groups[0]["bias_correction"] = True + opt.param_groups[0]["pre_mult_wd"] = 0.05 + + wrapper = SimpleNamespace(optimizer=opt) + saved = _neuter_optimizer(wrapper) + + assert opt.param_groups[0]["betas"] == [1.0, 1.0] + assert opt.param_groups[0]["weight_decay"] == 0.0 + assert opt.param_groups[0]["bias_correction"] is False + assert opt.param_groups[0]["pre_mult_wd"] == 0.0 + + _restore_optimizer(wrapper, saved) + + assert opt.param_groups[0]["betas"] == [0.9, 0.999] + assert opt.param_groups[0]["weight_decay"] == 0.01 + assert opt.param_groups[0]["bias_correction"] is True + assert opt.param_groups[0]["pre_mult_wd"] == 0.05 + + def test_multi_param_group_roundtrip(self): + """Neuter/restore with 2 param groups preserves per-group values.""" + import torch + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _neuter_optimizer, + _restore_optimizer, + ) + + model = torch.nn.Linear(4, 4) + w_params = [model.weight] + b_params = [model.bias] + opt = torch.optim.Adam( + [ + {"params": w_params, "weight_decay": 0.1}, + {"params": b_params, "weight_decay": 0.0}, + ], + lr=1e-3, + betas=(0.9, 0.999), + ) + + wrapper = SimpleNamespace(optimizer=opt) + saved = _neuter_optimizer(wrapper) + + for g in opt.param_groups: + assert g["betas"] == [1.0, 1.0] + assert g["weight_decay"] == 0.0 + + _restore_optimizer(wrapper, saved) + + assert opt.param_groups[0]["weight_decay"] == 0.1 + assert opt.param_groups[1]["weight_decay"] == 0.0 + for g in opt.param_groups: + assert list(g["betas"]) == [0.9, 0.999] + + +class TestResetOptimizerState: + """Verify _reset_optimizer_state clears step counters.""" + + def test_resets_param_group_step(self): + import torch + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_optimizer_state, + ) + + model = torch.nn.Linear(4, 4) + opt = torch.optim.Adam(model.parameters(), lr=1e-3) + loss = model(torch.randn(2, 4)).sum() + loss.backward() + opt.step() + + opt.param_groups[0]["step"] = 42 + wrapper = SimpleNamespace(optimizer=opt) + + _reset_optimizer_state(wrapper) + + assert opt.param_groups[0]["step"] == 0 + + def test_resets_per_param_state_step_tensor(self): + import torch + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_optimizer_state, + ) + + model = torch.nn.Linear(4, 4) + opt = torch.optim.Adam(model.parameters(), lr=1e-3) + loss = model(torch.randn(2, 4)).sum() + loss.backward() + opt.step() + + for state in opt.state.values(): + assert state["step"].item() > 0 + + wrapper = SimpleNamespace(optimizer=opt) + _reset_optimizer_state(wrapper) + + for state in opt.state.values(): + assert state["step"].item() == 0 + + def test_handles_chained_optimizer(self): + import torch + + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _reset_optimizer_state, + ) + + m1 = torch.nn.Linear(4, 4) + m2 = torch.nn.Linear(4, 4) + opt1 = torch.optim.SGD(m1.parameters(), lr=0.01) + opt2 = torch.optim.SGD(m2.parameters(), lr=0.01) + + opt1.param_groups[0]["step"] = 10 + opt2.param_groups[0]["step"] = 20 + + w1 = SimpleNamespace(optimizer=opt1) + w2 = SimpleNamespace(optimizer=opt2) + chained = SimpleNamespace(chained_optimizers=[w1, w2]) + + _reset_optimizer_state(chained) + + assert opt1.param_groups[0]["step"] == 0 + assert opt2.param_groups[0]["step"] == 0 diff --git a/tests/unit_tests/backends/megatron/test_mlperf_validation.py b/tests/unit_tests/backends/megatron/test_mlperf_validation.py new file mode 100644 index 000000000..3330a4bc3 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_mlperf_validation.py @@ -0,0 +1,164 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for MLPerf validation support. + +Tests cover: + 1. Validation loss function shape and target computation + 2. Training loss function shape (unchanged) + 3. Mock dataset timestep generation (production MockDiffusionDataset) +""" + +import pytest +import torch +import torch.nn.functional as F + +# --------------------------------------------------------------------------- +# 1. Validation loss function shape +# --------------------------------------------------------------------------- + + +class TestValLossFunc: + """Test the validation loss function produces correct shapes and values.""" + + def _make_val_loss_func(self, noise, clean_latents): + """Build val_loss_func mirroring DiffusionPretrainTrainer.forward_step.""" + + class Holder: + _last_noise = noise + _last_clean_latents = clean_latents + + def val_loss_func(output_tensor, non_loss_data=False): + if non_loss_data: + return output_tensor + target = Holder._last_noise - Holder._last_clean_latents + loss = F.mse_loss(output_tensor.float(), target.float(), reduction="none") + loss_per_sample = loss.mean(dim=tuple(range(1, loss.ndim))) + loss_sum = loss_per_sample.sum() + sample_count = torch.tensor(loss_per_sample.numel(), dtype=loss_sum.dtype, device=loss_sum.device) + return loss_sum, {"loss": (loss_sum.detach(), sample_count.detach())} + + return val_loss_func + + def test_returns_tuple_with_loss_key(self): + B, C, H, W = 4, 16, 32, 32 + noise = torch.randn(B, C, H, W) + clean = torch.randn(B, C, H, W) + pred = torch.randn(B, C, H, W) + fn = self._make_val_loss_func(noise, clean) + + loss_scalar, metrics = fn(pred) + + assert isinstance(loss_scalar, torch.Tensor) + assert loss_scalar.ndim == 0 + assert "loss" in metrics + loss_sum, sample_count = metrics["loss"] + assert loss_sum.ndim == 0 + assert sample_count.item() == B + + def test_correct_target(self): + B, C, H, W = 2, 4, 8, 8 + noise = torch.randn(B, C, H, W) + clean = torch.randn(B, C, H, W) + target = noise - clean + fn = self._make_val_loss_func(noise, clean) + + loss_scalar, metrics = fn(target) + + assert loss_scalar.item() == pytest.approx(0.0, abs=1e-6) + + def test_non_loss_data_passthrough(self): + fn = self._make_val_loss_func(torch.zeros(1), torch.zeros(1)) + sentinel = torch.tensor(42.0) + result = fn(sentinel, non_loss_data=True) + assert result is sentinel + + +# --------------------------------------------------------------------------- +# 2. Training loss function shape +# --------------------------------------------------------------------------- + + +class TestTrainLossFunc: + """Verify the training loss path returns scalar with reduced_train_loss.""" + + def _make_train_loss_func(self, loss_fn, noise, clean, mask): + class Holder: + _last_noise = noise + _last_clean_latents = clean + _last_loss_mask = mask + + def diffusion_loss_func(output_tensor, non_loss_data=False): + if non_loss_data: + return output_tensor + loss = loss_fn( + output_tensor, Holder._last_clean_latents, Holder._last_noise, Holder._last_loss_mask + ) + return loss, {"reduced_train_loss": loss.detach().clone()} + + return diffusion_loss_func + + def test_returns_scalar_with_key(self): + from primus.backends.megatron.training.diffusion.loss_computation import ( + compute_flow_matching_loss, + ) + + B, C, H, W = 4, 16, 32, 32 + noise = torch.randn(B, C, H, W) + clean = torch.randn(B, C, H, W) + pred = torch.randn(B, C, H, W) + + fn = self._make_train_loss_func(compute_flow_matching_loss, noise, clean, None) + loss, metrics = fn(pred) + + assert loss.ndim == 0 + assert "reduced_train_loss" in metrics + assert metrics["reduced_train_loss"].ndim == 0 + + +# --------------------------------------------------------------------------- +# 3. Mock dataset timestep generation +# --------------------------------------------------------------------------- + + +class TestMockDatasetTimestep: + """Verify validation mock datasets produce cycling timestep 0-7.""" + + def test_validation_dataset_has_timestep(self): + from primus.backends.megatron.data.synthetic.mock_datasets import ( + MockDiffusionDataset, + ) + + ds = MockDiffusionDataset( + num_samples=16, + image_size=256, + model_preset="flux_schnell", + dtype=torch.float32, + device="cpu", + is_validation=True, + ) + for i in range(16): + sample = ds[i] + assert "timestep" in sample, f"Sample {i} missing 'timestep'" + assert sample["timestep"].item() == i % 8 + + def test_training_dataset_no_timestep(self): + from primus.backends.megatron.data.synthetic.mock_datasets import ( + MockDiffusionDataset, + ) + + ds = MockDiffusionDataset( + num_samples=4, + image_size=256, + model_preset="flux_schnell", + dtype=torch.float32, + device="cpu", + is_validation=False, + ) + for i in range(4): + sample = ds[i] + assert "timestep" not in sample diff --git a/tests/unit_tests/backends/megatron/test_mlperf_warmup_state_equivalence.py b/tests/unit_tests/backends/megatron/test_mlperf_warmup_state_equivalence.py new file mode 100644 index 000000000..7742e010f --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_mlperf_warmup_state_equivalence.py @@ -0,0 +1,86 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Integration-style tests verifying that optimizer state is clean after warmup +completion. + +Uses mocked Megatron training harnesses - no real GPU training loop, but +exercises the warmup helper pipeline (neuter -> restore -> reset) end-to-end. + +Run: + python -m pytest tests/unit_tests/backends/megatron/test_mlperf_warmup_state_equivalence.py -v +""" + +from types import SimpleNamespace + +import torch + + +class TestOptimizerStateCleanAfterWarmup: + """After neuter → steps → restore → reset, optimizer should be clean.""" + + def test_adam_state_clean_after_warmup_simulation(self): + """Simulate warmup: step with real betas, then restore + reset. + + Note: production uses Apex/TE FusedAdam which supports + bias_correction=False. Stock PyTorch Adam doesn't, so betas=[1,1] + causes division-by-zero. We use SGD for the neutered phase, then + switch to Adam to verify the reset path. + """ + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _neuter_optimizer, + _reset_optimizer_state, + _restore_optimizer, + ) + + model = torch.nn.Linear(16, 8, bias=True) + opt = torch.optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=0.01) + wrapper = SimpleNamespace(optimizer=opt) + + for _ in range(3): + loss = model(torch.randn(4, 16)).sum() + loss.backward() + opt.step() + opt.zero_grad() + + for state in opt.state.values(): + assert state["step"].item() > 0 + + saved_hyp = _neuter_optimizer(wrapper) + assert opt.param_groups[0]["betas"] == [1.0, 1.0] + assert opt.param_groups[0]["weight_decay"] == 0.0 + + _restore_optimizer(wrapper, saved_hyp) + _reset_optimizer_state(wrapper) + + assert list(opt.param_groups[0]["betas"]) == [0.9, 0.999] + assert opt.param_groups[0]["weight_decay"] == 0.01 + + for state in opt.state.values(): + step_val = state["step"] + if isinstance(step_val, torch.Tensor): + assert step_val.item() == 0 + else: + assert step_val == 0 + + def test_param_groups_step_cleared(self): + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _neuter_optimizer, + _reset_optimizer_state, + _restore_optimizer, + ) + + model = torch.nn.Linear(8, 4) + opt = torch.optim.SGD(model.parameters(), lr=0.01) + opt.param_groups[0]["step"] = 42 + wrapper = SimpleNamespace(optimizer=opt) + + saved_hyp = _neuter_optimizer(wrapper) + _restore_optimizer(wrapper, saved_hyp) + _reset_optimizer_state(wrapper) + + assert opt.param_groups[0].get("step", 0) == 0 diff --git a/tests/unit_tests/backends/megatron/test_runtime_hooks_patches.py b/tests/unit_tests/backends/megatron/test_runtime_hooks_patches.py index b29fe36e6..406a61051 100644 --- a/tests/unit_tests/backends/megatron/test_runtime_hooks_patches.py +++ b/tests/unit_tests/backends/megatron/test_runtime_hooks_patches.py @@ -26,9 +26,9 @@ def _orig_compile(): training_pkg.initialize = init_mod megatron_mod.training = training_pkg - sys.modules["megatron"] = megatron_mod - sys.modules["megatron.training"] = training_pkg - sys.modules["megatron.training.initialize"] = init_mod + monkeypatch.setitem(sys.modules, "megatron", megatron_mod) + monkeypatch.setitem(sys.modules, "megatron.training", training_pkg) + monkeypatch.setitem(sys.modules, "megatron.training.initialize", init_mod) return init_mod, _orig_compile diff --git a/tests/unit_tests/backends/megatron/test_training_log_patches.py b/tests/unit_tests/backends/megatron/test_training_log_patches.py index 6776c23ff..bcd781de9 100644 --- a/tests/unit_tests/backends/megatron/test_training_log_patches.py +++ b/tests/unit_tests/backends/megatron/test_training_log_patches.py @@ -45,9 +45,9 @@ def fake_get_model(*args, **kwargs): training_pkg.training = training_mod megatron_mod.training = training_pkg - sys.modules["megatron"] = megatron_mod - sys.modules["megatron.training"] = training_pkg - sys.modules["megatron.training.training"] = training_mod + monkeypatch.setitem(sys.modules, "megatron", megatron_mod) + monkeypatch.setitem(sys.modules, "megatron.training", training_pkg) + monkeypatch.setitem(sys.modules, "megatron.training.training", training_mod) return training_mod, fake_training_log @@ -140,9 +140,10 @@ def test_patch_training_log_wraps_and_stacks_extensions(monkeypatch: pytest.Monk def test_rocm_monitor_hooked_print_rank_last_injects_stats(monkeypatch: pytest.MonkeyPatch): - # Prepare fake torch and ROCm SMI helpers. MemoryStatsExtension also calls - # torch.tensor / torch.distributed.all_gather for cross-rank max ROCm mem; - # stub those so inject() does not fall into the exception path in unit tests. + # Prepare fake torch and ROCm SMI helpers. MemoryStatsExtension computes the + # cross-rank max ROCm mem via torch.tensor + torch.distributed.all_reduce(MAX) + # and torch.distributed.get_rank(); stub those so inject() does not fall into + # the exception path in unit tests. class _FakeTensor: __slots__ = ("_value",) @@ -163,13 +164,17 @@ def _zeros_like(_t): def _get_world_size(): return 8 - def _all_gather(gathered_list, tensor): - v = tensor.item() - for i in range(len(gathered_list)): - gathered_list[i] = _FakeTensor(v) + def _get_rank(): + return 0 + + def _all_reduce(tensor, op=None): + # Single-rank reduction: MAX over one rank is a no-op, so leave the + # in-place tensor value unchanged (matching real all_reduce semantics). + return None # `inject` passes `dtype=torch.int64`; SimpleNamespace must expose `int64` - # or attribute lookup fails before our fake `tensor()` runs. + # or attribute lookup fails before our fake `tensor()` runs. The refactored + # code also reads `torch.distributed.ReduceOp.MAX`, so provide that too. fake_torch = SimpleNamespace( cuda=SimpleNamespace( mem_get_info=lambda: (2 * 1024**3, 4 * 1024**3), @@ -180,7 +185,9 @@ def _all_gather(gathered_list, tensor): int64=int, distributed=SimpleNamespace( get_world_size=_get_world_size, - all_gather=_all_gather, + get_rank=_get_rank, + all_reduce=_all_reduce, + ReduceOp=SimpleNamespace(MAX=object()), ), ) monkeypatch.setattr( diff --git a/tests/unit_tests/backends/megatron/test_warmup_convergence.py b/tests/unit_tests/backends/megatron/test_warmup_convergence.py new file mode 100644 index 000000000..07ab49219 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_warmup_convergence.py @@ -0,0 +1,363 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Single-GPU end-to-end test proving that warmup is transparent. + +A model that undergoes warmup steps (with param snapshot, optimizer neutering, +FP8 reset) followed by a full state restore produces the **exact same +training trajectory** as a model that trains from scratch. + +Parametrised over both FP8 spec paths (local delayed-scaling and +TransformerEngine). The trainable modules are ``torch.nn.Linear`` subclasses +decorated with FP8 metadata; the forward pass runs standard bf16 matmuls. +Loss-matching validates param snapshot/restore + optimizer reset. FP8 state +health (no NaN scales, registry reinit, amax seeding) is checked separately. + +Run: + python -m pytest tests/unit_tests/backends/megatron/test_warmup_convergence.py -v +""" + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") + +DEVICE = "cuda:0" +FP8_FWD_MAX = torch.finfo(torch.float8_e4m3fn).max +FP8_BWD_MAX = torch.finfo(torch.float8_e5m2).max + +WARMUP_STEPS = 3 +TRAIN_STEPS = 20 +IN_FEATURES = 32 +HIDDEN = 64 +OUT_FEATURES = 16 +BATCH = 8 + + +# --------------------------------------------------------------------------- +# Trainable FP8 module helpers +# --------------------------------------------------------------------------- + + +class _DelayedScalingLinear(torch.nn.Linear): + """``torch.nn.Linear`` with local-spec delayed-scaling FP8 metadata.""" + + def __init__(self, in_features, out_features, history_len=1): + super().__init__(in_features, out_features) + self._use_delayed_scaling = True + self._fp8_fwd_dtype = torch.float8_e4m3fn + self._fp8_bwd_dtype = torch.float8_e5m2 + self._fp8_fwd_max = FP8_FWD_MAX + self._fp8_bwd_max = FP8_BWD_MAX + self._amax_compute_algo = "most_recent" + self._first_delayed_step = True + self._history_idx = 0 + self.config = SimpleNamespace( + fp8_amax_history_len=history_len, + fp8_amax_compute_algo="most_recent", + ) + + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _init_delayed_scaling_state, + ) + + _init_delayed_scaling_state(self) + + def _move_fp8_buffers(self, device): + for name in list(self._buffers): + buf = self._buffers[name] + if buf is not None and buf.device != device: + self._buffers[name] = buf.to(device) + + +class _TEScalingLinear(torch.nn.Linear): + """``torch.nn.Linear`` with TransformerEngine-style FP8 metadata.""" + + def __init__(self, in_features, out_features, device=DEVICE): + super().__init__(in_features, out_features) + self.fp8_initialized = True + self.fp8_meta = { + "scaling_fwd": SimpleNamespace( + amax_history=torch.zeros(16, 1, device=device), + scale=torch.ones(1, device=device), + scale_inv=torch.ones(1, device=device), + ), + "scaling_bwd": SimpleNamespace( + amax_history=torch.zeros(16, 1, device=device), + scale=torch.ones(1, device=device), + scale_inv=torch.ones(1, device=device), + ), + } + + +def _make_model(fp8_spec, seed, device=DEVICE): + """Create a 3-layer MLP with the requested FP8 metadata, on ``device``.""" + torch.manual_seed(seed) + if fp8_spec == "local": + model = torch.nn.Sequential( + _DelayedScalingLinear(IN_FEATURES, HIDDEN), + torch.nn.ReLU(), + _DelayedScalingLinear(HIDDEN, OUT_FEATURES), + ).to(device) + for m in model.modules(): + if isinstance(m, _DelayedScalingLinear): + m._move_fp8_buffers(torch.device(device)) + elif fp8_spec == "te": + model = torch.nn.Sequential( + _TEScalingLinear(IN_FEATURES, HIDDEN, device=device), + torch.nn.ReLU(), + _TEScalingLinear(HIDDEN, OUT_FEATURES, device=device), + ).to(device) + else: + raise ValueError(f"Unknown fp8_spec: {fp8_spec}") + return model + + +def _get_delayed_modules(model): + return [m for m in model.modules() if getattr(m, "_use_delayed_scaling", False)] + + +def _make_registry(model): + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _DelayedScalingRegistry, + ) + + modules = _get_delayed_modules(model) + assert modules, "No delayed-scaling modules found" + return _DelayedScalingRegistry(modules) + + +def _generate_fixed_data(num_batches, seed=99, device=DEVICE): + torch.manual_seed(seed) + return [ + ( + torch.randn(BATCH, IN_FEATURES, device=device), + torch.randn(BATCH, OUT_FEATURES, device=device), + ) + for _ in range(num_batches) + ] + + +def _run_warmup(model, adam_opt, fp8_spec, registry=None): + """Execute the full warmup pipeline: snapshot, neuter, SGD steps, restore. + + Returns the registry (possibly reinitialised for local spec). + """ + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _fast_update_scales, + ) + from primus.backends.megatron.patches.mlperf_warmup_patches import ( + _neuter_optimizer, + _reset_fp8_local_spec, + _reset_fp8_te_spec, + _reset_optimizer_state, + _restore_optimizer, + _seed_fp8_amax, + ) + + wrapper = SimpleNamespace(optimizer=adam_opt) + device = next(model.parameters()).device + + # 1. Snapshot params to CPU + saved_params = {name: p.data.to("cpu", non_blocking=True) for name, p in model.named_parameters()} + torch.cuda.synchronize() + + # 2. Neuter Adam + saved_hyp = _neuter_optimizer(wrapper) + + # 3-4. Throwaway SGD warmup steps + warmup_sgd = torch.optim.SGD(model.parameters(), lr=0.01) + for _ in range(WARMUP_STEPS): + if fp8_spec == "local" and registry is not None: + _fast_update_scales(registry) + x = torch.randn(BATCH, IN_FEATURES, device=device) + y = torch.randn(BATCH, OUT_FEATURES, device=device) + loss = (model(x) - y).pow(2).mean() + loss.backward() + warmup_sgd.step() + warmup_sgd.zero_grad() + + # 5. Restore Adam hyperparams + _restore_optimizer(wrapper, saved_hyp) + + # 6. Reset Adam step counters + _reset_optimizer_state(wrapper) + + # 7. Restore params from CPU snapshot + for name, p in model.named_parameters(): + if name in saved_params: + p.data.copy_(saved_params[name]) + del saved_params + + # 8. Reset FP8 state + if fp8_spec == "local": + _reset_fp8_local_spec([model]) + elif fp8_spec == "te": + _reset_fp8_te_spec([model]) + _seed_fp8_amax([model]) + + # 9. Synchronize + torch.cuda.synchronize() + + # 10. Zero gradients + adam_opt.zero_grad(set_to_none=True) + + return registry + + +def _train_loop(model, optimizer, data, fp8_spec, registry=None): + """Run training and return per-step losses.""" + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _fast_update_scales, + ) + + losses = [] + for x, y in data: + if fp8_spec == "local" and registry is not None: + _fast_update_scales(registry) + loss = (model(x) - y).pow(2).mean() + loss.backward() + optimizer.step() + optimizer.zero_grad() + losses.append(loss.detach()) + return losses + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("fp8_spec", ["local", "te"]) +class TestWarmupConvergence: + + def test_loss_matches_baseline(self, fp8_spec): + """After warmup + full reset, loss curve is identical to no-warmup baseline.""" + seed = 42 + + baseline_model = _make_model(fp8_spec, seed) + warmup_model = _make_model(fp8_spec, seed) + + # Verify identical starting params + for (n1, p1), (n2, p2) in zip(baseline_model.named_parameters(), warmup_model.named_parameters()): + assert torch.equal(p1.data, p2.data), f"Init mismatch on {n1}" + + baseline_opt = torch.optim.Adam( + baseline_model.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=0.01 + ) + warmup_opt = torch.optim.Adam( + warmup_model.parameters(), lr=1e-3, betas=(0.9, 0.999), weight_decay=0.01 + ) + + baseline_reg = _make_registry(baseline_model) if fp8_spec == "local" else None + warmup_reg = _make_registry(warmup_model) if fp8_spec == "local" else None + + # Run warmup on warmup_model + _run_warmup(warmup_model, warmup_opt, fp8_spec, registry=warmup_reg) + + # Generate fixed training data + data = _generate_fixed_data(TRAIN_STEPS) + + # Train both + baseline_losses = _train_loop(baseline_model, baseline_opt, data, fp8_spec, registry=baseline_reg) + warmup_losses = _train_loop(warmup_model, warmup_opt, data, fp8_spec, registry=warmup_reg) + + for i, (lb, lw) in enumerate(zip(baseline_losses, warmup_losses)): + if i < 5: + assert torch.equal( + lb, lw + ), f"Step {i}: bitwise mismatch baseline={lb.item()}, warmup={lw.item()}" + else: + assert abs(lb.item() - lw.item()) < 1e-5, ( + f"Step {i}: baseline={lb.item()}, warmup={lw.item()}, " + f"diff={abs(lb.item() - lw.item())}" + ) + + def test_fp8_state_healthy_after_reset(self, fp8_spec): + """FP8 metadata has no NaN/Inf after warmup + reset.""" + from primus.backends.megatron.core.extensions.primus_turbo_float8_local import ( + _fast_update_scales, + ) + + seed = 42 + model = _make_model(fp8_spec, seed) + opt = torch.optim.Adam(model.parameters(), lr=1e-3) + registry = _make_registry(model) if fp8_spec == "local" else None + + _run_warmup(model, opt, fp8_spec, registry=registry) + + if fp8_spec == "local": + _fast_update_scales(registry) + + assert ( + registry._first_step is False + ), "_first_step should be consumed after post-reset _fast_update_scales" + assert torch.isfinite( + registry.scales_3n + ).all(), f"Non-finite scales after reset: {registry.scales_3n}" + + for i, m in enumerate(_get_delayed_modules(model)): + expected_amax = m.weight.data.abs().amax().float().item() + actual_amax = m.staged_weight_amax.item() + assert ( + abs(actual_amax - expected_amax) < 1e-3 + ), f"Module {i}: weight amax {actual_amax} != expected {expected_amax}" + + elif fp8_spec == "te": + for m in model.modules(): + if hasattr(m, "fp8_initialized"): + assert m.fp8_initialized is False, "fp8_initialized should be False after TE reset" + if hasattr(m, "fp8_meta"): + meta = m.fp8_meta + for key in ("scaling_fwd", "scaling_bwd"): + if key not in meta: + continue + tm = meta[key] + if hasattr(tm, "amax_history"): + assert (tm.amax_history == 1.0).all(), ( + f"{key}.amax_history should be seeded to 1.0, " f"got {tm.amax_history}" + ) + if hasattr(tm, "scale"): + assert torch.isfinite(tm.scale).all(), f"{key}.scale has non-finite values" + + def test_optimizer_step_parity(self, fp8_spec): + """Optimizer step counters match baseline after same number of real steps.""" + seed = 42 + train_steps = 10 + + baseline_model = _make_model(fp8_spec, seed) + warmup_model = _make_model(fp8_spec, seed) + + baseline_opt = torch.optim.Adam(baseline_model.parameters(), lr=1e-3, betas=(0.9, 0.999)) + warmup_opt = torch.optim.Adam(warmup_model.parameters(), lr=1e-3, betas=(0.9, 0.999)) + + baseline_reg = _make_registry(baseline_model) if fp8_spec == "local" else None + warmup_reg = _make_registry(warmup_model) if fp8_spec == "local" else None + + _run_warmup(warmup_model, warmup_opt, fp8_spec, registry=warmup_reg) + + data = _generate_fixed_data(train_steps, seed=77) + + _train_loop(baseline_model, baseline_opt, data, fp8_spec, registry=baseline_reg) + _train_loop(warmup_model, warmup_opt, data, fp8_spec, registry=warmup_reg) + + for (p_b, state_b), (p_w, state_w) in zip(baseline_opt.state.items(), warmup_opt.state.items()): + step_b = state_b["step"] + step_w = state_w["step"] + if isinstance(step_b, torch.Tensor): + step_b = step_b.item() + if isinstance(step_w, torch.Tensor): + step_w = step_w.item() + assert step_b == step_w == train_steps, ( + f"Step count mismatch: baseline={step_b}, warmup={step_w}, " f"expected={train_steps}" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/unit_tests/backends/megatron/test_warmup_prefetch_cache.py b/tests/unit_tests/backends/megatron/test_warmup_prefetch_cache.py new file mode 100644 index 000000000..8905fb758 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_warmup_prefetch_cache.py @@ -0,0 +1,179 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +""" +Unit tests for the CudaPrefetchIterator-cache invalidation that the MLPerf +warmup hook performs in its epilogue. + +Background +---------- +``patch_grad_zero_and_data_prefetch`` (priority 41) lazily wraps the current +``data_iterator`` in a ``CudaPrefetchIterator`` on its first invocation and +caches the result in a closure-local ``_prefetch_state["iter"]``. The +MLPerf one-shot warmup hook (priority 95) drives the first invocation of +that inner ``_patched_train_step`` with ``data_iterator = synthetic_iter`` +for the synthetic warmup steps, which means the cached prefetch wrapper +ends up bound to the synthetic iterator -- and ``MegatronDataloaderWrapper`` +is cyclic (never raises ``StopIteration``). + +Without invalidation, every subsequent real training step's +``_synced_prefetch_fwd_bwd`` substitutes the cached +``CudaPrefetchIterator(synthetic_iter)`` for the incoming real +``data_iterator`` argument, so the model trains forever on the synthetic +mock dataset and val_loss on real data stays stuck near 1.38. + +The fix exposes a ``_PREFETCH_HANDLE`` plus a ``reset_prefetch_state`` +helper on ``delayed_fp8_scaling_patches`` and calls it from the warmup +epilogue. These tests cover the helper contract, the closure-handle +binding done at patch-installation time, and the integration semantics +that "after warmup epilogue, the next train_step rebuilds the prefetcher". + +Level 1: Pure Python, no GPU. + +Run: + python -m pytest \ + tests/unit_tests/backends/megatron/test_warmup_prefetch_cache.py -v +""" + +import pytest + +from primus.backends.megatron.patches import delayed_fp8_scaling_patches as dfp + + +@pytest.fixture(autouse=True) +def _isolate_prefetch_handle(): + """Snapshot+restore the module-level ``_PREFETCH_HANDLE`` so tests can + mutate it without bleeding into each other or into a real patch install. + """ + saved = dict(dfp._PREFETCH_HANDLE) + yield + dfp._PREFETCH_HANDLE.clear() + dfp._PREFETCH_HANDLE.update(saved) + + +# --------------------------------------------------------------------------- +# Helper contract +# --------------------------------------------------------------------------- + + +class _FakePrefetchIter: + """Stand-in for ``CudaPrefetchIterator`` -- the helper only inspects the + type name, never the behaviour.""" + + def __next__(self): + raise StopIteration + + +def test_get_prefetch_state_none_when_unbound(): + """Before any patch installation, the handle points at None.""" + dfp._PREFETCH_HANDLE["state"] = None + assert dfp.get_prefetch_state() is None + + +def test_get_prefetch_state_returns_closure_dict_when_bound(): + """``patch_grad_zero_and_data_prefetch`` binds its closure-local + ``_prefetch_state`` here at install time; the getter must return it.""" + bound = {"iter": _FakePrefetchIter()} + dfp._PREFETCH_HANDLE["state"] = bound + assert dfp.get_prefetch_state() is bound + + +def test_reset_prefetch_state_noop_when_unbound(): + """No patch installed -> reset is a no-op returning None.""" + dfp._PREFETCH_HANDLE["state"] = None + assert dfp.reset_prefetch_state() is None + + +def test_reset_prefetch_state_noop_when_bound_but_empty(): + """Patch installed but no iterator yet cached -> nothing to evict.""" + bound = {} + dfp._PREFETCH_HANDLE["state"] = bound + assert dfp.reset_prefetch_state() is None + assert bound == {} # state dict untouched + + +def test_reset_prefetch_state_evicts_cached_iter(): + """The cached iterator is returned AND popped from the state dict, so + the closure's ``if "iter" not in _prefetch_state`` guard will rebuild + on the next train_step.""" + fake_iter = _FakePrefetchIter() + bound = {"iter": fake_iter} + dfp._PREFETCH_HANDLE["state"] = bound + + evicted = dfp.reset_prefetch_state() + + assert evicted is fake_iter + assert "iter" not in bound, "iter key must be removed after reset" + + +def test_reset_prefetch_state_idempotent(): + """Calling reset twice doesn't crash and the second call returns None.""" + bound = {"iter": _FakePrefetchIter()} + dfp._PREFETCH_HANDLE["state"] = bound + + first = dfp.reset_prefetch_state() + second = dfp.reset_prefetch_state() + + assert first is not None + assert second is None + + +def test_reset_prefetch_state_preserves_other_keys(): + """Reset must only pop the ``iter`` key; any other state (e.g. metadata + a future patch revision might add) should survive.""" + fake_iter = _FakePrefetchIter() + bound = {"iter": fake_iter, "other_metadata": 42} + dfp._PREFETCH_HANDLE["state"] = bound + + evicted = dfp.reset_prefetch_state() + + assert evicted is fake_iter + assert bound == {"other_metadata": 42} + + +# --------------------------------------------------------------------------- +# Integration semantics: warmup epilogue + next-train_step rebuild +# --------------------------------------------------------------------------- + + +def test_warmup_epilogue_makes_next_step_rebuild_prefetcher(): + """End-to-end behavioural test of the contract we rely on. + + Simulates: + 1. Patch installation publishes a closure dict via ``_PREFETCH_HANDLE``. + 2. Warmup step 1 lazily caches a prefetch wrapper around the SYNTHETIC + iterator. + 3. Warmup epilogue calls ``reset_prefetch_state()``. + 4. The next ``_patched_train_step`` call sees ``"iter" not in + _prefetch_state`` and rebuilds the wrapper around the REAL + iterator. + + What we assert: step (4) sees an empty / iter-free state dict, which is + the precondition for the closure's lazy-rebuild guard to fire. + """ + + # Step 1: patch install publishes the closure dict. + closure_state: dict = {} + dfp._PREFETCH_HANDLE["state"] = closure_state + + # Step 2: warmup step 1 caches the prefetcher around synthetic_iter. + synthetic_iter_marker = object() + closure_state["iter"] = synthetic_iter_marker + assert dfp.get_prefetch_state() is closure_state + assert dfp.get_prefetch_state()["iter"] is synthetic_iter_marker + + # Step 3: warmup epilogue invalidates. + evicted = dfp.reset_prefetch_state() + assert evicted is synthetic_iter_marker + + # Step 4: next train_step's lazy-build guard now fires -- i.e. "iter" + # is not in the state dict, so the closure will construct a fresh + # CudaPrefetchIterator around the real data_iterator argument. + assert "iter" not in closure_state, ( + "After warmup epilogue invalidates the cache, the closure's " + "_prefetch_state['iter'] must be absent so the next train_step " + "rebuilds the prefetch wrapper around the real data_iterator." + ) From 6eb904a7e3c3fbd887e6ffee74b428aaddec4ba7 Mon Sep 17 00:00:00 2001 From: botaohu001 Date: Mon, 20 Jul 2026 10:52:21 +0800 Subject: [PATCH 045/127] feat: odc adapt (#864) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR adapts the On-Demand Communication (ODC) framework to the AMD ROCm stack (MI300X) inside Primus. ODC reduces gradient-communication synchronization from per-microbatch to per-minibatch and overlaps single-sided push/pull communication with the backward pass, enabling LB-Mini load balancing (variable microbatch counts per rank) on top of a rocSHMEM-based point-to-point backend. The rocSHMEM point-to-point ops are now provided by Primus-Turbo (primus_turbo.pytorch._C.odc_rocshmem_host / odc_rocshmem_gda) and consumed by _rocshmem_backend.py, instead of being built in-tree as ctypes .sos — ODC is therefore a pure-Python in-tree module now (the in-tree rocSHMEM sources and build_rocshmem_backend.sh are removed). This depends on Primus-Turbo PR #409. Validated to match the pure-Primus baseline on both single-node (host / XGMI IPC) and multi-node (GDA) runs. --------- Co-authored-by: botahu Co-authored-by: botahu Co-authored-by: Cursor Co-authored-by: Xiaoming-AMD --- LICENSE | 4 + .../MI355X/deepseek1.5B-odc-lbmini.yaml | 116 ++ .../configs/MI355X/qwen14B-odc-dn.yaml | 125 ++ examples/megatron/prepare.py | 9 + examples/run_pretrain.sh | 15 +- .../patches/distributed_init_patches.py | 13 +- .../patches/fused_linear_ce_patches.py | 186 +++ .../megatron/patches/odc_lb_mini_patches.py | 268 ++++ .../patches/odc_torch_fsdp2_patches.py | 457 +++++++ .../backends/megatron/sft/lb_mini_dataset.py | 333 +++++ .../backends/megatron/sft/lb_mini_packing.py | 398 ++++++ primus/backends/megatron/sft/packing.py | 13 +- .../deepseek_r1_distill_qwen_1.5B.yaml | 19 + .../configs/modules/megatron/sft_trainer.yaml | 26 + .../modules/megatron/trainer_base.yaml | 33 + primus/core/odc/.gitignore | 229 ++++ primus/core/odc/README.md | 163 +++ primus/core/odc/__init__.py | 57 + primus/core/odc/fsdp/__init__.py | 3 + primus/core/odc/fsdp/fsdp1.py | 281 +++++ primus/core/odc/fsdp/fsdp2.py | 1092 +++++++++++++++++ primus/core/odc/odc_early/sitecustomize.py | 40 + primus/core/odc/primitives/__init__.py | 43 + .../core/odc/primitives/_rocshmem_backend.py | 588 +++++++++ primus/core/odc/primitives/gather.py | 244 ++++ .../core/odc/primitives/scatter_accumulate.py | 431 +++++++ primus/core/odc/primitives/shmem_triton.py | 465 +++++++ primus/core/odc/primitives/utils.py | 488 ++++++++ primus/core/odc/rocshmem_runtime/README.md | 58 + .../rocshmem_runtime/scripts/cleanup_rs.sh | 37 + .../odc/rocshmem_runtime/scripts/run_odc.sh | 97 ++ primus/core/odc/runtime_config.py | 86 ++ pyproject.toml | 9 + 33 files changed, 6420 insertions(+), 6 deletions(-) create mode 100644 examples/megatron/configs/MI355X/deepseek1.5B-odc-lbmini.yaml create mode 100644 examples/megatron/configs/MI355X/qwen14B-odc-dn.yaml create mode 100644 primus/backends/megatron/patches/fused_linear_ce_patches.py create mode 100644 primus/backends/megatron/patches/odc_lb_mini_patches.py create mode 100644 primus/backends/megatron/patches/odc_torch_fsdp2_patches.py create mode 100644 primus/backends/megatron/sft/lb_mini_dataset.py create mode 100644 primus/backends/megatron/sft/lb_mini_packing.py create mode 100644 primus/configs/models/megatron/deepseek_r1_distill_qwen_1.5B.yaml create mode 100644 primus/core/odc/.gitignore create mode 100644 primus/core/odc/README.md create mode 100644 primus/core/odc/__init__.py create mode 100644 primus/core/odc/fsdp/__init__.py create mode 100644 primus/core/odc/fsdp/fsdp1.py create mode 100644 primus/core/odc/fsdp/fsdp2.py create mode 100644 primus/core/odc/odc_early/sitecustomize.py create mode 100644 primus/core/odc/primitives/__init__.py create mode 100644 primus/core/odc/primitives/_rocshmem_backend.py create mode 100644 primus/core/odc/primitives/gather.py create mode 100644 primus/core/odc/primitives/scatter_accumulate.py create mode 100644 primus/core/odc/primitives/shmem_triton.py create mode 100644 primus/core/odc/primitives/utils.py create mode 100644 primus/core/odc/rocshmem_runtime/README.md create mode 100755 primus/core/odc/rocshmem_runtime/scripts/cleanup_rs.sh create mode 100755 primus/core/odc/rocshmem_runtime/scripts/run_odc.sh create mode 100644 primus/core/odc/runtime_config.py diff --git a/LICENSE b/LICENSE index 297acbff2..96a4fb715 100644 --- a/LICENSE +++ b/LICENSE @@ -41,6 +41,10 @@ Primus uses or references the following third-party projects: - License: MIT License (Sea AI Lab, NVIDIA) - Repository: https://github.com/sail-sg/zero-bubble-pipeline-parallelism +5. odc (On-Demand Communication) + - License: MIT License (Sea AI Lab) + - Repository: https://github.com/sail-sg/odc + User must comply with the respective licenses of these third-party projects when using or distributing Primus. -------------------------------------------------------------------------------- diff --git a/examples/megatron/configs/MI355X/deepseek1.5B-odc-lbmini.yaml b/examples/megatron/configs/MI355X/deepseek1.5B-odc-lbmini.yaml new file mode 100644 index 000000000..a0ab10937 --- /dev/null +++ b/examples/megatron/configs/MI355X/deepseek1.5B-odc-lbmini.yaml @@ -0,0 +1,116 @@ +### DeepSeek-R1-Distill-Qwen-1.5B ODC LB-Mini SFT example (single-node). +### Sequence-length load balancing (LB-Mini) with the "fit" cost model on the +### torch-FSDP2 + ODC path; packed (thd) attention enabled. + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:deepseek1.5B-odc-lbmini} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + post_trainer: + framework: megatron + config: sft_trainer.yaml + + model: deepseek_r1_distill_qwen_1.5B.yaml + + overrides: + stage: sft + + hf_path: deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B + tokenizer_model: deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B + + sft_dataset_name: "zai-org/LongAlign-10k" + sft_conversation_format: "messages" + + # ODC on-demand communication (rocSHMEM/XGMI P2P gradient reduction) on the + # torch-FSDP2 path. Master switch (was the ODC_ENABLE env var). + enable_odc: true + odc_phase: 2 + + # thd segmented (packed) attention. + enable_packed_sequences: true + use_packed_attention: true + # LB-Mini sequence-length load balancing (was the ODC_LB_MINI env var). + enable_odc_lb_mini: true + # LB-Mini cost model: fit (attention-aware a*s^2 + b*s, 1.5B coefficients). + lb_mini_cost_model: fit + lb_mini_max_token_len: 32768 + + # ODC P2P backend: rocshmem (validated single-node host/XGMI-IPC). trainer_base default is mori. + odc_p2p_backend: rocshmem + + wandb_project: "Primus_ODC_DeepSeek1.5B" + stderr_sink_level: DEBUG + log_avg_skip_iterations: 2 + log_avg_reset_interval: 100 + + train_iters: 50 + micro_batch_size: 1 + global_batch_size: 16 + + seq_length: 32768 + max_position_embeddings: 32768 + + lr: 1.0e-5 + min_lr: 0.0 + lr_warmup_iters: 2 + lr_decay_iters: 100 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + + seed: 1234 + eod_mask_loss: false + init_method_std: 0.008 + norm_epsilon: 1.0e-6 + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: false + + enable_fused_linear_ce: false + + use_torch_fsdp2: true + use_megatron_fsdp: false + use_distributed_optimizer: false + overlap_grad_reduce: false + overlap_param_gather: false + gradient_accumulation_fusion: false + + finetune: true + load: null + save: null + save_interval: 1000 + eval_interval: 1000 + no_save_optim: null + no_save_rng: null + disable_last_saving: true + ckpt_format: torch_dist + + bf16: true + + enable_primus_turbo: false + use_turbo_attention: false + use_turbo_grouped_mlp: false + use_turbo_rms_norm: false + + eval_iters: 0 + + profile: false + use_pytorch_profiler: true + use_nsys_profiler: false + profile_step_start: 10 + profile_step_end: 11 + profile_ranks: [0] + pytorch_profiler_collect_shapes: true + pytorch_profiler_collect_callstack: false + pytorch_profiler_collect_chakra: false + record_shapes: false + record_memory_history: false + nvtx_ranges: false + + lora: + enabled: false diff --git a/examples/megatron/configs/MI355X/qwen14B-odc-dn.yaml b/examples/megatron/configs/MI355X/qwen14B-odc-dn.yaml new file mode 100644 index 000000000..88d3242cd --- /dev/null +++ b/examples/megatron/configs/MI355X/qwen14B-odc-dn.yaml @@ -0,0 +1,125 @@ +### DeepSeek-R1-Distill-Qwen-14B ODC dual-node SFT example. +### Sequence-length load balancing (LB-Mini) with the "fit" cost model on the +### torch-FSDP2 + ODC path; packed (thd) attention enabled. + +work_group: ${PRIMUS_TEAM:amd} +user_name: ${PRIMUS_USER:root} +exp_name: ${PRIMUS_EXP_NAME:qwen14B-odc-dn} +workspace: ${PRIMUS_WORKSPACE:./output} + +modules: + post_trainer: + framework: megatron + config: sft_trainer.yaml + + model: qwen2.5_14B.yaml + + overrides: + stage: sft + + hf_path: deepseek-ai/DeepSeek-R1-Distill-Qwen-14B + tokenizer_model: deepseek-ai/DeepSeek-R1-Distill-Qwen-14B + + sft_dataset_name: "zai-org/LongAlign-10k" + sft_conversation_format: "messages" + + # ODC on-demand communication (rocSHMEM/XGMI P2P gradient reduction) on the + # torch-FSDP2 path. Master switch (was the ODC_ENABLE env var). + enable_odc: true + odc_phase: 2 + + # thd segmented (packed) attention. + enable_packed_sequences: true + use_packed_attention: true + # LB-Mini sequence-length load balancing (was the ODC_LB_MINI env var). + enable_odc_lb_mini: true + # LB-Mini cost model: fit (attention-aware a*s^2 + b*s). + lb_mini_cost_model: fit + + # ODC P2P backend + dual-node GDA (validated). trainer_base defaults are mori / false. + odc_p2p_backend: rocshmem + odc_rocshmem_gda: true + + wandb_project: "Primus_ODC_Qwen14B" + stderr_sink_level: DEBUG + log_avg_skip_iterations: 2 + log_avg_reset_interval: 100 + + train_iters: 50 + micro_batch_size: 1 + global_batch_size: 16 + + seq_length: 65536 + max_position_embeddings: 65536 + + lr: 2.0e-6 + min_lr: 0.0 + lr_warmup_iters: 2 + lr_decay_iters: 100 + lr_decay_style: cosine + weight_decay: 0.1 + adam_beta1: 0.9 + adam_beta2: 0.95 + + seed: 1234 + eod_mask_loss: false + init_method_std: 0.008 + norm_epsilon: 1.0e-5 + add_qkv_bias: true + apply_rope_fusion: false + + tensor_model_parallel_size: 1 + pipeline_model_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: false + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 + + enable_fused_linear_ce: false + + use_torch_fsdp2: true + use_megatron_fsdp: false + use_distributed_optimizer: false + overlap_grad_reduce: false + overlap_param_gather: false + gradient_accumulation_fusion: false + + finetune: true + # Megatron checkpoint to warm-start from. Set PRETRAINED_CKPT to your own + # converted DeepSeek-R1-Distill-Qwen-14B checkpoint dir; null trains from + # the HF weights only. + pretrained_checkpoint: ${PRETRAINED_CKPT:null} + load: null + save: null + save_interval: 1000 + eval_interval: 1000 + no_save_optim: null + no_save_rng: null + disable_last_saving: true + ckpt_format: torch_dist + + bf16: true + + enable_primus_turbo: false + use_turbo_attention: false + use_turbo_grouped_mlp: false + use_turbo_rms_norm: false + + eval_iters: 0 + + profile: false + use_pytorch_profiler: true + use_nsys_profiler: false + profile_step_start: 8 + profile_step_end: 9 + profile_ranks: [0] + pytorch_profiler_collect_shapes: true + pytorch_profiler_collect_callstack: false + pytorch_profiler_collect_chakra: false + record_shapes: false + record_memory_history: false + nvtx_ranges: false + + lora: + enabled: false diff --git a/examples/megatron/prepare.py b/examples/megatron/prepare.py index bdad391c3..86725f311 100644 --- a/examples/megatron/prepare.py +++ b/examples/megatron/prepare.py @@ -282,6 +282,15 @@ def build_megatron_helper(primus_path: Path, patch_args: Path, backend_path: str f"if you actually need it). SFT does not require this dependency." ) return + # This prepare step may run once PER RANK (e.g. 8-16 concurrent processes on a + # node under torchrun), and concurrent `pip install -e` writes race on the shared + # venv .pth file (OSError). If the package is already importable (one-time + # pre-install), take a fast, write-free path so all ranks agree without racing. + import importlib.util + + if importlib.util.find_spec("emerging_optimizers") is not None: + log_info("Emerging-Optimizers already installed; skipping editable reinstall.") + return log_info(f"Building Emerging Optimizers in {emerging_optimizers_path}") ret = subprocess.run( ["pip", "install", "--no-build-isolation", "-e", str(emerging_optimizers_path)], check=True diff --git a/examples/run_pretrain.sh b/examples/run_pretrain.sh index d16d6baaa..a14d93edf 100755 --- a/examples/run_pretrain.sh +++ b/examples/run_pretrain.sh @@ -92,11 +92,18 @@ export HF_HOME=${HF_HOME:-"${DATA_PATH}/huggingface"} # shellcheck source=/dev/null source "${PRIMUS_PATH}/runner/helpers/envs/path_utils.sh" -LOG_INFO_RANK0 "Pip installing required packages ..." -if [ "${BACKEND:-}" != "MaxText" ]; then - pip install -r "$PRIMUS_PATH/requirements.txt" --quiet +# PRIMUS_SKIP_PIP=1 skips the per-run pip install (deps already ship in the base +# image). Use it when many ranks/nodes share one venv and concurrent pip installs +# would race, or simply to speed up a warm container. Not keyed on the launcher. +if [ "${PRIMUS_SKIP_PIP:-0}" == "1" ]; then + LOG_INFO_RANK0 "PRIMUS_SKIP_PIP=1: skipping pip install (deps from image)" else - pip install -r "$PRIMUS_PATH/requirements-jax.txt" --quiet + LOG_INFO_RANK0 "Pip installing required packages ..." + if [ "${BACKEND:-}" != "MaxText" ]; then + pip install -r "$PRIMUS_PATH/requirements.txt" --quiet + else + pip install -r "$PRIMUS_PATH/requirements-jax.txt" --quiet + fi fi diff --git a/primus/backends/megatron/patches/distributed_init_patches.py b/primus/backends/megatron/patches/distributed_init_patches.py index b86a8b3b5..df5c483a4 100644 --- a/primus/backends/megatron/patches/distributed_init_patches.py +++ b/primus/backends/megatron/patches/distributed_init_patches.py @@ -35,7 +35,18 @@ "prevent RCCL device mapping deadlocks on MI355X (FSDP2 only)." ), priority=10, - condition=lambda ctx: getattr(get_args(ctx), "use_torch_fsdp2", False), + # ODC (enable_odc=true) drives gradient exchange over rocSHMEM P2P, not RCCL, and + # relies on those P2P copy streams overlapping with compute in the backward pass. + # The device_id injection here eagerly creates the world + ~26 sub-group RCCL + # communicators, whose resident streams/DMA queues serialize ODC's XGMI copy + # streams onto the critical path (profiled: cross-stream overlap 120ms -> 2.4ms, + # ~+128ms/step on single-node 1.5B). nccl_pad is unaffected (it uses these RCCL + # comms as its native reduce-scatter). So skip the eager-RCCL device_id patch + # under ODC. Safe on MI300X (the MI355X deadlock this guards does not trigger + # here; commits before this patch existed ran ODC correctly). + condition=lambda ctx: ( + getattr(get_args(ctx), "use_torch_fsdp2", False) and not getattr(get_args(ctx), "enable_odc", False) + ), ) def patch_init_process_group_device_id(ctx: PatchContext): """ diff --git a/primus/backends/megatron/patches/fused_linear_ce_patches.py b/primus/backends/megatron/patches/fused_linear_ce_patches.py new file mode 100644 index 000000000..7333a2222 --- /dev/null +++ b/primus/backends/megatron/patches/fused_linear_ce_patches.py @@ -0,0 +1,186 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +############################################################################### +# Chunked linear cross-entropy for Megatron GPTModel (avoid full-logits OOM). +# +# PROBLEM: Megatron's stock GPTModel._postprocess materializes the FULL logits +# tensor [seq, batch, vocab] before cross-entropy. For a large vocab (Qwen +# 151936) and long sequence (e.g. 64k), that single tensor is tens of GB and +# OOMs -- even though a 1.5B model's weights/activations are tiny. This is the +# exact wall we hit validating ODC/LB-Mini on DeepSeek-R1-Distill-Qwen-1.5B. +# +# FIX (verl/Liger spirit): split along the sequence dim and compute +# logits+CE per chunk under activation checkpointing, so the full [seq, vocab] +# logits is NEVER resident -- peak logits memory is one chunk only. This lets +# TP=1 / DP=8 (the layout ODC requires) run 64k-token sequences without OOM. +# +# ZERO IMPACT on stock Megatron: +# * Pure monkey-patch in the Primus layer; third-party Megatron source is NOT +# touched. +# * Gated by enable_fused_linear_ce (yaml) or FUSED_LINEAR_CE=1 (env). +# DEFAULT OFF -> the original full-logits _postprocess runs byte-for-byte. +# * Even when ON, only the training-with-labels path is intercepted; +# inference / no-labels (generation) / MTP / non-post-process stages all +# fall through to the original implementation. +# * Numerically equivalent: per-token CE is independent across positions, so +# chunk-then-concatenate yields the identical [b, s] per-token loss. +############################################################################### + +import os + +import torch + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0 + +_PATCHED = False + + +def _fused_ce_enabled(args) -> bool: + """Switch can come from the YAML arg or, for experiments, the env var.""" + return bool(getattr(args, "enable_fused_linear_ce", False)) or ( + os.environ.get("FUSED_LINEAR_CE", "0") == "1" + ) + + +def _chunk_size() -> int: + """Per-chunk sequence length. Smaller = less peak logits mem, more recompute.""" + return int(os.environ.get("FUSED_CE_CHUNK", "0") or 0) or 4096 + + +class _ChunkedLinearCE(torch.autograd.Function): + """verl-style fused linear cross-entropy. + + forward: compute the LM loss in sequence chunks WITHOUT building an autograd + graph and WITHOUT saving per-chunk logits -> the full [s, vocab] logits + is never materialized (peak = one chunk). + backward: recompute each chunk's logits and obtain grads via a SINGLE + torch.autograd.grad call per chunk, inside ONE backward invocation. + + Why not torch.utils.checkpoint: its recompute is driven by an unpack hook + that runs INTERLEAVED with FSDP2/ODC's backward communication. Under ODC the + ranks run DIFFERENT micro-batch counts, so that interleaving desyncs and + DEADLOCKS (the iter-3 hang). Doing the recompute here -- a plain local + autograd.grad with NO collective -- keeps each rank independent in backward. + """ + + @staticmethod + def forward(ctx, hidden, weight, labels, model, chunk): + ctx.model = model + ctx.chunk = chunk + ctx.save_for_backward(hidden, weight, labels) + seq = hidden.size(0) + parts = [] + with torch.no_grad(): + for i in range(0, seq, chunk): + logits = model._scale_logits(torch.matmul(hidden[i : i + chunk], weight.t())) + parts.append(model.compute_language_model_loss(labels[:, i : i + chunk].contiguous(), logits)) + return torch.cat(parts, dim=1) # [b, s] + + @staticmethod + def backward(ctx, grad_out): # grad_out [b, s] + hidden, weight, labels = ctx.saved_tensors + model = ctx.model + chunk = ctx.chunk + seq = hidden.size(0) + grad_hidden = torch.empty_like(hidden) + grad_weight = torch.zeros_like(weight) + for i in range(0, seq, chunk): + h_c = hidden[i : i + chunk].detach().requires_grad_(True) + w = weight.detach().requires_grad_(True) + with torch.enable_grad(): + logits = model._scale_logits(torch.matmul(h_c, w.t())) + loss_c = model.compute_language_model_loss(labels[:, i : i + chunk].contiguous(), logits) + g_h, g_w = torch.autograd.grad(loss_c, (h_c, w), grad_out[:, i : i + chunk].contiguous()) + grad_hidden[i : i + chunk] = g_h + grad_weight = grad_weight + g_w + return grad_hidden, grad_weight, None, None, None + + +def _chunked_lm_loss(model, hidden_states, labels): + """Chunked linear+CE that never materializes the full [s, b, vocab] logits. + + hidden_states: [s, b, h] labels: [b, s] -> per-token loss [b, s] + """ + # PREFERRED PATH (ODC + DiffMicro): use the full output weight that ODC's + # train-loop hook already all-gathered ONCE at the minibatch boundary + # (pre_minibatch_start). Calling full_tensor() here -- inside the per-micro + # batch forward -- would issue a DTensor collective whose call-count differs + # across ranks under DiffMicro, and deadlock. The cached tensor is a leaf + # with requires_grad=True; its .grad accumulates across all this rank's + # micro-batches and is reduce-scattered back to the sharded param at + # pre_optimizer_step (see odc_torch_fsdp2_patches._odc_reduce_output_grad). + cached = getattr(model, "_odc_cached_output_weight", None) + if cached is not None: + return _ChunkedLinearCE.apply(hidden_states, cached, labels, model, _chunk_size()) + + # FALLBACK (no ODC, or SameMicro): same-count collectives are safe. + if model.share_embeddings_and_output_weights: + ow = model.shared_embedding_or_output_weight() + else: + ow = model.output_layer.weight + if hasattr(ow, "full_tensor"): + ow_full = ow.full_tensor() + else: + ow_full = ow.clone() + return _ChunkedLinearCE.apply(hidden_states, ow_full, labels, model, _chunk_size()) + + +def _install_fused_ce_patch(): + global _PATCHED + if _PATCHED: + return + import megatron.core.models.gpt.gpt_model as gpt_mod + + GPTModel = gpt_mod.GPTModel + if getattr(GPTModel._postprocess, "_fused_ce_hooked", False): + _PATCHED = True + return + orig_postprocess = GPTModel._postprocess + + def postprocess_with_fused_ce(self, *args, **kwargs): + # forward() calls _postprocess() with all-kwargs, but support positional + # too: signature is (hidden_states, input_ids, position_ids, labels, ...). + hidden_states = kwargs.get("hidden_states", args[0] if len(args) > 0 else None) + labels = kwargs.get("labels", args[3] if len(args) > 3 else None) + inference_context = kwargs.get("inference_context") + in_inference = inference_context is not None and not self.training + + if ( + labels is not None + and hidden_states is not None + and not in_inference + and getattr(self, "post_process", True) + and not getattr(self.config, "mtp_num_layers", 0) + ): + return _chunked_lm_loss(self, hidden_states, labels) + # Everything else -> stock path, unchanged. + return orig_postprocess(self, *args, **kwargs) + + postprocess_with_fused_ce._fused_ce_hooked = True + GPTModel._postprocess = postprocess_with_fused_ce + _PATCHED = True + log_rank_0( + f"[FusedCE] patched GPTModel._postprocess: chunked linear+CE " + f"(chunk={_chunk_size()}); full [seq, vocab] logits no longer materialized." + ) + + +@register_patch( + "megatron.fused_linear_ce", + backend="megatron", + phase="before_train", + description="Chunked linear cross-entropy for GPTModel to avoid full-logits OOM (large vocab + long seq).", + condition=lambda ctx: _fused_ce_enabled(get_args(ctx)), +) +def patch_fused_linear_ce(ctx: PatchContext): + log_rank_0( + "[FusedCE] enable_fused_linear_ce=true -> installing chunked linear+CE " + "(avoids materializing full [seq, vocab] logits; numerically equivalent)." + ) + _install_fused_ce_patch() + + +__all__ = ["patch_fused_linear_ce"] diff --git a/primus/backends/megatron/patches/odc_lb_mini_patches.py b/primus/backends/megatron/patches/odc_lb_mini_patches.py new file mode 100644 index 000000000..e240274e6 --- /dev/null +++ b/primus/backends/megatron/patches/odc_lb_mini_patches.py @@ -0,0 +1,268 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +############################################################################### +# LB-Mini (sequence-length load balancing) for Megatron's FSDP2 path. +# +# ``enable_odc_lb_mini`` is an ORTHOGONAL, standalone capability: it serves the +# variable-length, Karmarkar-Karp-balanced LB-Mini DATA on the torch-FSDP2 path. +# It is INDEPENDENT of the ODC communication switch (``enable_odc``); what the +# ODC switch selects is only HOW the per-rank micro-batch counts are aligned: +# +# * enable_odc_lb_mini=false (DEFAULT) -> this patch is a complete no-op; +# Megatron runs its stock fixed-num_microbatches, all-ranks-in-lockstep +# schedule on stock (padded) data. Byte-for-byte unchanged. +# +# * enable_odc_lb_mini=true + enable_odc=true -> DECOUPLED mode. Data is served +# variable length and KK-balanced across DP ranks; each rank runs its OWN +# (possibly different) number of micro-batches (same_micro_num=False). Only +# ODC's point-to-point comm can drive ranks out of lockstep without a +# collective deadlock, hence this mode requires ODC comm. +# +# * enable_odc_lb_mini=true + enable_odc=false -> ALIGNED mode. The SAME +# variable-length KK-balanced DATA is served, but the micro-batch count is +# all-reduce(MAX)-aligned so every rank runs the SAME number of steps +# (same_micro_num=True). Uniform per-rank counts keep standard FSDP2 + RCCL +# collectives in lockstep (no deadlock), so this is a fair "same data" nccl +# baseline WITHOUT ODC. This is the config-driven replacement for the removed +# LB_MINI_FORCE_DATA A/B env (which likewise served LB-Mini data under NCCL). +# +# In BOTH enabled modes we install the dataloader patch (variable-length data) +# AND the schedule patch (rank-local num_microbatches). The schedule patch is +# comm-agnostic -- it only overrides num_microbatches with the iterator's planned +# per-rank count. Under ALIGNED mode those counts are identical across ranks, so +# it is NCCL/RCCL-safe; it is NOT an ODC-specific reduction. (Running the +# dataloader patch WITHOUT the schedule patch would be incoherent: the stock +# schedule would pull a fixed num_microbatches while the iterator plans a +# possibly-different per-rank count, drifting the two out of sync.) +# +# All wiring is monkey-patch in the Primus layer; the third-party Megatron-LM +# source is NOT modified. +# +# Stage-1 scope (this file): make "different micro-batch count per rank" run end +# to end without deadlocking. Numerical normalization (loss_scale / consumed +# samples by real tokens) is Stage-2. +############################################################################### + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 + +# Global handle so the schedule patch can reach the LB-Mini iterator that the +# dataloader patch created. Stage-1 only drives the TRAIN iterator (eval_iters=0 +# in the aligned config), so a single handle is sufficient. +_LB_MINI_TRAIN_ITER = None +_FB_PATCHED = False + + +def _lb_mini_enabled(args) -> bool: + """LB-Mini DATA serving is driven by ``enable_odc_lb_mini`` ALONE. + + LB-Mini is ORTHOGONAL to the ODC comm switch: it only requires the explicit + ``enable_odc_lb_mini`` config item and the torch-FSDP2 path it patches. The + ``enable_odc`` switch does NOT gate whether LB-Mini data is served; it only + selects the micro-batch alignment mode (see ``_lb_mini_aligned``): + + * enable_odc=true -> DECOUPLED (ranks may run different micro-batch counts; + needs ODC point-to-point comm). + * enable_odc=false -> ALIGNED (all ranks run the same micro-batch count + via all_reduce(MAX); NCCL/RCCL-safe "same data" + baseline, no ODC). + """ + return bool(getattr(args, "enable_odc_lb_mini", False)) and bool(getattr(args, "use_torch_fsdp2", False)) + + +def _lb_mini_aligned(args) -> bool: + """Micro-batch alignment mode for LB-Mini. + + ALIGNED (True) when ODC comm is OFF: all ranks are forced to the same + micro-batch count (``same_micro_num=True``, all_reduce MAX) so standard + FSDP2 + RCCL collectives stay in lockstep. DECOUPLED (False) when ODC comm is + ON: ranks may run different counts, which only ODC's point-to-point comm can + drive without a collective deadlock. + """ + return not bool(getattr(args, "enable_odc", False)) + + +def _build_lb_mini_train_iterator(args): + """Build the variable-length, KK-balanced LB-Mini train iterator.""" + from megatron.core import mpu + from megatron.training import get_tokenizer + + from primus.backends.megatron.sft.lb_mini_dataset import ( + LBMiniDataIterator, + build_varlen_samples, + ) + from primus.backends.megatron.sft.packing import _resolve_pad_token_id + + tokenizer = get_tokenizer() + samples = build_varlen_samples( + dataset_name=getattr(args, "sft_dataset_name", "tatsu-lab/alpaca"), + tokenizer=tokenizer, + max_seq_length=args.seq_length, + # Some datasets (e.g. SWE-bench/SWE-smith-trajectories) have no "train" + # split; sft_dataset_split lets the config pick it (default "train"). + split=str(getattr(args, "sft_dataset_split", "train")), + formatter=getattr(args, "sft_conversation_format", "alpaca"), + seed=args.seed, + bridge_compat_inline_bos=bool(getattr(args, "sft_bridge_compat_inline_bos", False)), + ) + # Per-micro-batch token cap. Larger than a single sample lets short samples + # pack together and long samples stand alone -> creates per-rank micro-batch + # count differences (where DiffMicro saves comm rounds). Priority: + # yaml lb_mini_max_token_len > seq_length. + max_token_len = int(getattr(args, "lb_mini_max_token_len", 0) or 0) or int(args.seq_length) + # ALIGNED (enable_odc=false) -> same_micro_num=True: all ranks run the SAME + # micro-batch count (all_reduce MAX), keeping standard RCCL collectives in + # lockstep -> fair "same data" NCCL baseline. DECOUPLED (enable_odc=true) -> + # same_micro_num=False: ranks may differ; only ODC p2p comm can drive that. + aligned = _lb_mini_aligned(args) + it = LBMiniDataIterator( + samples=samples, + global_batch_size=args.global_batch_size, + max_token_len=max_token_len, + dp_rank=mpu.get_data_parallel_rank(), + dp_size=mpu.get_data_parallel_world_size(), + pad_id=_resolve_pad_token_id(tokenizer), + cost_model=str(getattr(args, "lb_mini_cost_model", "linear")), + seed=args.seed, + shuffle=True, + same_micro_num=aligned, + packing_method="kk", + ) + log_rank_0( + f"[ODC.lb_mini] built LB-Mini train iterator: {len(samples)} varlen samples, " + f"global_batch_size={args.global_batch_size}, max_token_len={max_token_len}, " + f"dp_size={mpu.get_data_parallel_world_size()}, " + f"cost_model={getattr(args, 'lb_mini_cost_model', 'linear')}, " + f"same_micro_num={aligned} " + f"({'ALIGNED baseline (nccl, no ODC)' if aligned else 'LB-Mini decoupled (ODC)'})" + ) + return it + + +def _install_dataloader_patch(): + """Patch build_pretraining_data_loader so the TRAIN loader is LB-Mini. + + Valid/test loaders (if any) fall through to the stock builder unchanged. + We tag the first (train) request via a module flag because the stock + signature does not carry the split explicitly. + """ + import megatron.training.training as mt_training + from megatron.training.datasets import data_samplers + + if getattr(data_samplers.build_pretraining_data_loader, "_lb_mini_hooked", False): + return + + orig_builder = data_samplers.build_pretraining_data_loader + + def lb_mini_builder(dataset, consumed_samples): + global _LB_MINI_TRAIN_ITER + # Runtime call: use Megatron's get_args() (no ctx). Primus' get_args(ctx) + # is only valid inside register_patch conditions / patch bodies. + from megatron.training import get_args as _mt_get_args + + args = _mt_get_args() + # Only replace the TRAIN loader, and only once (the first non-zero-len + # build). Identify train by: not yet built + dataset present. + if _lb_mini_enabled(args) and _LB_MINI_TRAIN_ITER is None and dataset is not None: + try: + _LB_MINI_TRAIN_ITER = _build_lb_mini_train_iterator(args) + log_rank_0("[ODC.lb_mini] TRAIN dataloader replaced by LB-Mini iterator") + return _LB_MINI_TRAIN_ITER + except Exception as e: # noqa: BLE001 + warning_rank_0( + f"[ODC.lb_mini] failed to build LB-Mini iterator, " + f"falling back to stock loader: {type(e).__name__}: {e}" + ) + return orig_builder(dataset, consumed_samples) + + lb_mini_builder._lb_mini_hooked = True + data_samplers.build_pretraining_data_loader = lb_mini_builder + # training.py imported the symbol into its own namespace; rebind there too. + if hasattr(mt_training, "build_pretraining_data_loader"): + mt_training.build_pretraining_data_loader = lb_mini_builder + log_rank_0("[ODC.lb_mini] hooked build_pretraining_data_loader") + + +def _install_schedule_patch(): + """Patch forward_backward_no_pipelining to use THIS rank's micro-batch count. + + At the top of every train_step the LB-Mini iterator plans one global + minibatch (KK balance across ranks); we read this rank's micro-batch count + and override the (globally-identical) ``num_microbatches`` argument so the + schedule's forward/backward loop runs the right rank-local number of steps. + """ + global _FB_PATCHED + import megatron.core.pipeline_parallel.schedules as sched + + if _FB_PATCHED or getattr(sched.forward_backward_no_pipelining, "_lb_mini_hooked", False): + return + + orig_fb = sched.forward_backward_no_pipelining + + def lb_mini_fb(*args, **kwargs): + it = _LB_MINI_TRAIN_ITER + forward_only = kwargs.get("forward_only", False) + # Only re-plan for the train path (an LB-Mini iterator exists) and when + # actually training (forward_only=False is the train_step path). + if it is not None and not forward_only: + try: + local_nmb = it.begin_minibatch() + if local_nmb > 0: + kwargs["num_microbatches"] = local_nmb + except Exception as e: # noqa: BLE001 + warning_rank_0( + f"[ODC.lb_mini] begin_minibatch failed, using stock " + f"num_microbatches: {type(e).__name__}: {e}" + ) + return orig_fb(*args, **kwargs) + + lb_mini_fb._lb_mini_hooked = True + sched.forward_backward_no_pipelining = lb_mini_fb + + # get_forward_backward_func returns the module-global by name; rebinding the + # module attribute is enough as long as it is fetched AFTER this patch. Also + # patch the function it returns defensively if it caches a reference. + _FB_PATCHED = True + log_rank_0("[ODC.lb_mini] hooked forward_backward_no_pipelining (rank-local num_microbatches)") + + +@register_patch( + "megatron.fsdp.odc_lb_mini", + backend="megatron", + phase="before_train", + description="LB-Mini sequence-length load balancing for Megatron FSDP2 (data decoupled from ODC comm).", + condition=lambda ctx: _lb_mini_enabled(get_args(ctx)), +) +def patch_odc_lb_mini(ctx: PatchContext): + aligned = _lb_mini_aligned(get_args(ctx)) + mode = "ALIGNED (nccl, no ODC)" if aligned else "DECOUPLED (ODC comm)" + log_rank_0( + f"[ODC.lb_mini] enable_odc_lb_mini=true, mode={mode} -> installing LB-Mini " + "(variable-length KK-balanced data + rank-local num_microbatches schedule)." + ) + # Loss normalization under torch FSDP2: we KEEP calculate_per_token_loss at + # its default (False). Each micro-batch loss is mean-reduced (/=num_tokens) + # then /=num_microbatches (this rank's count) -- the same per-minibatch mean + # ODC's own example uses, and it keeps gradients at the right magnitude. + # + # We deliberately do NOT force calculate_per_token_loss=True. Under + # use_torch_fsdp2 the per-token grad rescale lives in Megatron's + # finalize_model_grads (scale by the GLOBAL all-reduced token count), but + # FSDP2 does its OWN reduce-scatter and BYPASSES that path, so per-token + # leaves the summed (un-normalized) loss and gradients explode ~1000x + # (measured: grad norm ~45000 vs ~55). KK balancing keeps per-rank + # micro-batch counts nearly equal (and in ALIGNED mode they are EXACTLY + # equal), so the residual per-minibatch-mean weighting difference (e.g. a rank + # with 3 vs 4 micro-batches) is negligible. + # + # Both patches are installed in either mode. The schedule patch is + # comm-agnostic (it only sets this rank's num_microbatches); in ALIGNED mode + # (enable_odc=false) the counts are all_reduce(MAX)-uniform across ranks, so + # the standard FSDP2 + RCCL collectives stay in lockstep -- no ODC required. + _install_dataloader_patch() + _install_schedule_patch() + + +__all__ = ["patch_odc_lb_mini"] diff --git a/primus/backends/megatron/patches/odc_torch_fsdp2_patches.py b/primus/backends/megatron/patches/odc_torch_fsdp2_patches.py new file mode 100644 index 000000000..4e9e31951 --- /dev/null +++ b/primus/backends/megatron/patches/odc_torch_fsdp2_patches.py @@ -0,0 +1,457 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +############################################################################### +# ODC (On-Demand Communication) integration into Megatron's PyTorch-FSDP2 path. +# +# Gated by the enable_odc config item (so it is a no-op unless explicitly turned +# on) AND requires use_torch_fsdp2=true. +# +# Strategy (see primus/core/odc/ROCM_ADAPTATION_REPORT.md for context): +# Megatron's TorchFullyShardedDataParallel wraps the model with the STANDARD +# PyTorch `fully_shard` API. ODC's odc/fsdp/fsdp2.py monkey-patches the same +# torch.distributed.fsdp._fully_shard internals (foreach_all_gather, +# FSDPParamGroup.post_backward, ...), so ODC's comm replacement applies +# transparently to Megatron's FSDP2 modules. +# +# Integration points: +# patch_fsdp2() -> BEFORE the first fully_shard call +# patch_lazy_init(m) -> AFTER fully_shard, for every wrapped module +# pre_minibatch_start -> start of each train_step (PHASE 2) +# pre_optimizer_step -> before optimizer.step() (PHASE 2) +# +# The odc_phase config item controls how far we wire in: +# 1: only __init__ hook (patch_fsdp2 + patch_lazy_init). +# Verifies model construction + first forward/backward do not crash (i.e. +# ODC's symm-buffer replacement is compatible with Megatron's FSDP2 params). +# Gradients are NOT yet routed. +# 2 (production default): also hook train_step / optimizer.step for full grad +# routing. +############################################################################### + +import os + +from primus.core.patches import PatchContext, get_args, register_patch +from primus.core.utils.module_utils import log_rank_0, warning_rank_0 + +_ODC_READY = False +_FSDP2_PATCHED = False + + +def _ensure_odc_ready(): + """Initialize MORI-SHMEM and apply ODC's module-level FSDP2 patches. + + Must run BEFORE the first fully_shard() call. Idempotent. + """ + global _ODC_READY, _FSDP2_PATCHED + import odc + from odc.fsdp import fsdp2 as odc_fsdp2 + + if not _ODC_READY: + os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "8G") + if "MORI_SOCKET_IFNAME" not in os.environ and "NCCL_SOCKET_IFNAME" in os.environ: + os.environ["MORI_SOCKET_IFNAME"] = os.environ["NCCL_SOCKET_IFNAME"].lstrip("=") + odc.init_shmem() + _ODC_READY = True + log_rank_0("[ODC.torch_fsdp2] init_shmem (MORI) done") + + if not _FSDP2_PATCHED: + odc_fsdp2.patch_fsdp2(enable_hpz=False) + _FSDP2_PATCHED = True + log_rank_0("[ODC.torch_fsdp2] patch_fsdp2() applied (foreach_all_gather / post_backward replaced)") + + +def _apply_patch_lazy_init(root_module): + """Call ODC patch_lazy_init on every fully_shard-ed submodule that owns a + parameter group. + + A module is fully_shard-ed iff FSDP2 attached the ``_get_fsdp_state`` + accessor. We additionally require ``state._fsdp_param_group is not None``: + container wrappers (e.g. the root) have no param group of their own, and ODC + itself skips them everywhere (replace_sharded_param_with_symm_buffer / + pre_minibatch_start / pre_optimizer_step all filter ``_fsdp_param_group is + None``). Patching such a module would just raise AttributeError, so we skip + it explicitly instead of catching the error. + """ + from odc.fsdp import fsdp2 as odc_fsdp2 + + patched, skipped = 0, 0 + for m in root_module.modules(): + if not hasattr(m, "_get_fsdp_state"): + continue + state = m._get_fsdp_state() + if state is None or state._fsdp_param_group is None: + skipped += 1 + continue + try: + odc_fsdp2.patch_lazy_init(m) + patched += 1 + except Exception as e: # noqa: BLE001 + warning_rank_0( + f"[ODC.torch_fsdp2] patch_lazy_init failed on a param-group module: " + f"{type(e).__name__}: {e}" + ) + log_rank_0( + f"[ODC.torch_fsdp2] patch_lazy_init applied to {patched} param-group modules " + f"({skipped} container module(s) without param group skipped)" + ) + + +def _populate_odc_runtime_config(ctx: PatchContext): + """Bridge the odc_* trainer-config items into the ODC library runtime config. + + The ODC primitives (odc.primitives.*) are decoupled library code that reads its + tuning knobs from odc.runtime_config (populated here) rather than os.environ. + Only values explicitly set in the config are forwarded; unset items (None) keep + the library defaults, so behaviour is unchanged unless a knob is configured. + """ + import odc + + args = get_args(ctx) + + def _g(name): + return getattr(args, name, None) + + _defer = _g("odc_gda_defer_reduce") + odc.set_runtime_config( + p2p_backend=_g("odc_p2p_backend"), + mori_init=_g("odc_mori_init"), + max_buffer_size=_g("odc_max_buffer_size"), + rocshmem_gda=_g("odc_rocshmem_gda"), + rocshmem_lib=_g("odc_rocshmem_lib"), + gda_rs_blocks=_g("odc_gda_rs_blocks"), + gda_pipe=_g("odc_gda_pipe"), + gda_defer_reduce=(str(_defer) if _defer is not None else None), + gda_warmup_mode=_g("odc_gda_warmup_mode"), + gda_stride_bytes=_g("odc_gda_stride_bytes"), + ) + log_rank_0( + "[ODC.torch_fsdp2] runtime config populated from trainer config " + f"(p2p_backend={odc.get_runtime_config().p2p_backend}, " + f"rocshmem_gda={odc.get_runtime_config().rocshmem_gda}, " + f"warmup_mode={odc.get_runtime_config().gda_warmup_mode})" + ) + + +@register_patch( + "megatron.fsdp.odc_torch_fsdp2", + backend="megatron", + phase="before_train", + description="Integrate ODC on-demand communication into Megatron's PyTorch FSDP2 path (ROCm/MORI).", + condition=lambda ctx: getattr(get_args(ctx), "enable_odc", False) + and getattr(get_args(ctx), "use_torch_fsdp2", False), +) +def patch_odc_torch_fsdp2(ctx: PatchContext): + # Populate the ODC runtime config from the trainer config BEFORE any ODC + # primitive is imported. `import odc` is cheap (odc/__init__ is lazy and does + # not pull in odc.primitives), so this runs before _ensure_odc_ready() imports + # odc.fsdp.fsdp2 -> odc.primitives, whose import-time backend selection + # (odc_p2p_backend) must read the populated config, not a stale default. + _populate_odc_runtime_config(ctx) + + # PR #808 (feat(flux): FSDP2 optimizers + fp8 all-gather) replaced Megatron's FSDP2 + # wrapper with PrimusTorchFullyShardedDataParallel (installed as + # megatron.training.training.torch_FSDP), so the stock TorchFullyShardedDataParallel + # is never instantiated. We must hook the class the trainer actually uses, otherwise + # _ensure_odc_ready()/reduction_service never initializes and pre_minibatch_start + # crashes with 'NoneType' object has no attribute 'clear_accumulations'. + try: + from primus.backends.megatron.core.distributed.torch_fully_sharded_data_parallel import ( + PrimusTorchFullyShardedDataParallel as TorchFSDP, + ) + except ImportError: + import megatron.core.distributed.torch_fully_sharded_data_parallel as tfsdp_mod + + TorchFSDP = tfsdp_mod.TorchFullyShardedDataParallel + if getattr(TorchFSDP.__init__, "_odc_hooked", False): + log_rank_0("[ODC.torch_fsdp2] __init__ already hooked, skip") + return + + orig_init = TorchFSDP.__init__ + phase = str(getattr(get_args(ctx), "odc_phase", 2)) + + def odc_init(self, *args, **kwargs): + # 1) MORI init + ODC module-level FSDP2 patch, BEFORE any fully_shard. + _ensure_odc_ready() + # 2) Megatron's original __init__ runs all fully_shard() calls. + orig_init(self, *args, **kwargs) + # 3) AFTER fully_shard: install ODC's lazy_init hook (symm-buffer replace) + # on every wrapped module. + _apply_patch_lazy_init(self.module) + # 4) [ODC x #808] ODC uses a serial single-stream rocSHMEM/GDA transport, so + # #808's FSDP2 forward all-gather prefetch cannot overlap and is pure + # overhead (+3.4~3.9s/step gather_kernel, +26.6GB max reserved on dual-node + # 14B). ODC is active here (this patch only runs when enable_odc=true), so we + # opt every wrapped module out of the forward prefetch. Backward prefetch is + # kept -- it is fine/needed under ODC. + _disable_forward_prefetch(self.module) + log_rank_0(f"[ODC.torch_fsdp2] TorchFSDP wrapped with ODC (odc_phase={phase})") + + odc_init._odc_hooked = True + TorchFSDP.__init__ = odc_init + log_rank_0(f"[ODC.torch_fsdp2] hooked {TorchFSDP.__name__}.__init__") + + if phase == "2": + _install_train_loop_hooks() + + +def _disable_forward_prefetch(root_module): + """[ODC x #808] Opt ODC-managed modules out of #808's forward all-gather prefetch. + + PR #808's PrimusTorchFullyShardedDataParallel calls set_modules_to_forward_prefetch() + on every fully_shard-ed inner module, overlapping the *next* layer's all-gather with + the current layer's forward. That is a win for FSDP2-native NCCL all-gather, but ODC's + transport is a serial single-stream rocSHMEM/GDA path (overlap_factor 1.00x): the + prefetched all-gather cannot overlap and only adds gather_kernel time and peak memory + (profiled: +3.4~3.9s/step, +26.6GB max reserved on dual-node 14B). Since this runs only + when ODC is enabled, we clear forward prefetch on every param-group module. FORWARD + only -- backward prefetch is left intact (fine/needed under ODC). + """ + cleared, missing = 0, 0 + for m in root_module.modules(): + if not hasattr(m, "_get_fsdp_state"): + continue + state = m._get_fsdp_state() + if state is None or state._fsdp_param_group is None: + continue + if not hasattr(m, "set_modules_to_forward_prefetch"): + missing += 1 + continue + try: + m.set_modules_to_forward_prefetch([]) + cleared += 1 + except Exception as e: # noqa: BLE001 + warning_rank_0( + f"[ODC.torch_fsdp2] set_modules_to_forward_prefetch([]) failed: " f"{type(e).__name__}: {e}" + ) + log_rank_0( + f"[ODC.torch_fsdp2] skipped #808 forward all-gather prefetch on {cleared} module(s) " + f"(kept backward prefetch); ODC's serial transport cannot overlap it" + + (f"; {missing} module(s) lacked the API" if missing else "") + ) + + +def _find_gpt_model(root): + """Locate the GPTModel (owns output_layer + compute_language_model_loss).""" + m = getattr(root, "module", root) # Float16Module -> GPTModel + if hasattr(m, "output_layer") and hasattr(m, "compute_language_model_loss"): + return m + for sub in root.modules(): + if hasattr(sub, "output_layer") and hasattr(sub, "compute_language_model_loss"): + return sub + return None + + +def _fused_ce_hooked_on(gpt): + pp = getattr(type(gpt), "_postprocess", None) + return bool(getattr(pp, "_fused_ce_hooked", False)) + + +def _odc_gather_output_weight(root): + """At the minibatch boundary (a SYNC point: all ranks enter train_step + together), all-gather the sharded output weight ONCE and cache the full + tensor as a grad-tracking leaf. fused CE then reuses it for every micro + batch WITHOUT per-micro collectives (which would deadlock under DiffMicro). + Only active when fused CE has patched GPTModel._postprocess. + """ + gpt = _find_gpt_model(root) + if gpt is None or not _fused_ce_hooked_on(gpt): + return + try: + if getattr(gpt, "share_embeddings_and_output_weights", False): + ow = gpt.shared_embedding_or_output_weight() + else: + ow = gpt.output_layer.weight + if hasattr(ow, "full_tensor"): + gpt._odc_cached_output_weight = ow.full_tensor().detach().requires_grad_(True) + except Exception as e: # noqa: BLE001 + warning_rank_0(f"[ODC.fusedce] gather output weight failed: {type(e).__name__}: {e}") + + +def _odc_reduce_output_grad(root): + """At pre_optimizer_step (minibatch end, also a SYNC point): all-reduce the + cached full-weight grad across DP (AVG), scatter it back to the sharded + param's .grad (DTensor), and drop the cache. One collective per minibatch + per rank -> DiffMicro-safe. + """ + import torch.distributed as dist + + gpt = _find_gpt_model(root) + if gpt is None: + return + cached = getattr(gpt, "_odc_cached_output_weight", None) + if cached is None: + return + try: + if cached.grad is not None: + full_grad = cached.grad + if getattr(gpt, "share_embeddings_and_output_weights", False): + ow = gpt.shared_embedding_or_output_weight() + else: + ow = gpt.output_layer.weight + dist.all_reduce(full_grad, op=dist.ReduceOp.AVG) + from torch.distributed.tensor import DTensor, distribute_tensor + + if isinstance(ow, DTensor): + scattered = distribute_tensor(full_grad, ow.device_mesh, ow.placements) + ow.grad = scattered if ow.grad is None else (ow.grad + scattered) + elif ow.grad is None: + ow.grad = full_grad + else: + ow.grad = ow.grad + full_grad + except Exception as e: # noqa: BLE001 + warning_rank_0(f"[ODC.fusedce] reduce output grad failed: {type(e).__name__}: {e}") + finally: + gpt._odc_cached_output_weight = None + + +def _odc_grad_spike_guard(root): + """Skip the optimizer step when the global grad norm spikes abnormally. + + WHY: ODC's asynchronous P2P/MORI gradient reduction is numerically + NON-deterministic -- the same forward (identical loss) can occasionally + produce a huge grad spike (observed: same iter, loss 11.17512 vs 11.17508, + but grad norm 9.7 vs 442415). Most spikes are small and clip_grad absorbs + them, but an extreme one pushes params into a persistent divergent state + (clip only bounds magnitude, not the already-corrupted direction). This is + most visible under token-imbalanced SAME_MICRO (arm2); LB-Mini (token + balanced) rarely triggers it. + + FIX: at the minibatch sync point (optimizer.step), compute the global grad + norm; if it exceeds the odc_grad_spike_threshold config (default 1000, <=0 + disables), zero ALL grads so the ensuing step is a no-op -- i.e. SKIP this bad + iter, exactly like Megatron skips nan/inf iterations. Normal grad norm here is + ~3-42, recoverable spikes ~100-200, pathological ones ~1e5+, so 1000 cleanly + separates them. + + The all_reduce below is safe under DiffMicro: optimizer.step runs once per + minibatch at a sync point where all ranks are aligned. + """ + import torch + import torch.distributed as dist + from megatron.training import get_args as _mt_get_args + + thr = float(getattr(_mt_get_args(), "odc_grad_spike_threshold", 1000.0)) + if thr <= 0: + return + local_sq = None + for p in root.parameters(): + g = getattr(p, "grad", None) + if g is None: + continue + if hasattr(g, "to_local"): # DTensor -> this rank's shard + g = g.to_local() + s = g.detach().float().pow(2).sum() + local_sq = s if local_sq is None else (local_sq + s) + # [ODC-FIX] Every rank MUST enter the all_reduce below in lockstep. Under odc_nopad + # (variable micro-batch counts) a rank can legitimately finish a minibatch with no + # local grad (local_sq is None). The old early-return made that rank skip the + # collective -> ncclDevKernel_Generic_2 spins forever on the other ranks -> GPU + # deadlock at optimizer.step (confirmed via rocgdb + HIP trace). Contribute a 0 so the + # barrier stays aligned; only bail out when there is truly no process group. + if dist.is_initialized(): + if local_sq is None: + local_sq = torch.zeros((), device=next(root.parameters()).device, dtype=torch.float32) + dist.all_reduce(local_sq, op=dist.ReduceOp.SUM) + if local_sq is None: + return + gnorm = local_sq.sqrt().item() + if gnorm > thr: + for p in root.parameters(): + if getattr(p, "grad", None) is not None: + p.grad.zero_() + warning_rank_0( + f"[ODC.spike_guard] grad norm {gnorm:.1f} > {thr} -> SKIP step " + f"(grads zeroed; non-deterministic ODC async-reduce spike, params unchanged)" + ) + + +def _install_train_loop_hooks(): + """PHASE 2: wire pre_minibatch_start / pre_optimizer_step into the loop.""" + import megatron.training.training as mt_training + + from odc.fsdp import fsdp2 as odc_fsdp2 + + if getattr(mt_training.train_step, "_odc_hooked", False): + return + + orig_train_step = mt_training.train_step + + def _find_fsdp_root(model): + # model is a list of model_chunks (TorchFSDP instances). The fully_shard + # root is chunk.module. + chunk = model[0] if isinstance(model, (list, tuple)) else model + root = getattr(chunk, "module", chunk) + return root + + def odc_train_step(forward_step_func, data_iterator, model, optimizer, *a, **kw): + root = _find_fsdp_root(model) + # pre_minibatch_start: clear ODC accumulations at the start of the step. + try: + odc_fsdp2.pre_minibatch_start(root) + except Exception as e: # noqa: BLE001 + warning_rank_0(f"[ODC.torch_fsdp2] pre_minibatch_start failed: {type(e).__name__}: {e}") + # fused CE: all-gather full output weight ONCE per minibatch (sync point), + # so per-micro-batch fused CE never issues a DiffMicro-unsafe collective. + _odc_gather_output_weight(root) + + # Hook optimizer.step once to inject pre_optimizer_step before it. + if not getattr(optimizer.step, "_odc_hooked", False): + orig_step = optimizer.step + + def odc_step(*sa, **sk): + try: + odc_fsdp2.pre_optimizer_step(root) + except Exception as e: # noqa: BLE001 + warning_rank_0(f"[ODC.torch_fsdp2] pre_optimizer_step failed: {type(e).__name__}: {e}") + # fused CE: reduce-scatter cached full output-weight grad back to + # the sharded param (sync point, one collective per minibatch). + _odc_reduce_output_grad(root) + # guard against ODC async-reduce's occasional non-deterministic + # grad spike (can push params into a divergent state, esp. under + # token-imbalanced SAME_MICRO): skip the step if grad norm spikes. + _odc_grad_spike_guard(root) + return orig_step(*sa, **sk) + + odc_step._odc_hooked = True + optimizer.step = odc_step + log_rank_0("[ODC.torch_fsdp2] hooked optimizer.step (pre_optimizer_step injected)") + + return orig_train_step(forward_step_func, data_iterator, model, optimizer, *a, **kw) + + odc_train_step._odc_hooked = True + mt_training.train_step = odc_train_step + log_rank_0("[ODC.torch_fsdp2] hooked train_step (pre_minibatch_start injected)") + + +@register_patch( + "megatron.fsdp.odc_torch_fsdp2_teardown", + backend="megatron", + phase="after_train", + description="Tear down the ODC reduction service after training (ROCm/MORI).", + condition=lambda ctx: getattr(get_args(ctx), "enable_odc", False) + and getattr(get_args(ctx), "use_torch_fsdp2", False), +) +def patch_odc_torch_fsdp2_teardown(ctx: PatchContext): + """Tear down the ODC reduction service once training is done. + + The device-side (single-node XGMI pull-sum) and GPU-direct (GDA) reduce + paths run no host-side subprocess, so ``ReductionService.stop()`` is a + no-op kept for API compatibility. This hook is retained as the single + teardown call site (harmless today) so Primus keeps a place to release ODC + resources; we deliberately skip finalize_distributed() / + SymmBufferRegistry.finalize() so we do not tear down the process group + before Primus' own cleanup runs. + """ + from odc.fsdp import fsdp2 as odc_fsdp2 + + rs = odc_fsdp2.get_reduction_service() + if rs is None: + log_rank_0("[ODC.torch_fsdp2] teardown: no reduction_service, skip") + return + try: + rs.stop() + log_rank_0("[ODC.torch_fsdp2] reduction service torn down at after_train") + except Exception as e: # noqa: BLE001 + warning_rank_0(f"[ODC.torch_fsdp2] teardown stop() failed (non-fatal): " f"{type(e).__name__}: {e}") diff --git a/primus/backends/megatron/sft/lb_mini_dataset.py b/primus/backends/megatron/sft/lb_mini_dataset.py new file mode 100644 index 000000000..20124d3f5 --- /dev/null +++ b/primus/backends/megatron/sft/lb_mini_dataset.py @@ -0,0 +1,333 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""Variable-length SFT sample store + LB-Mini data iterator (ODC path). + +This is the data layer for LB-Mini. Unlike ``PackedSFTDataset`` (which bin-packs +everything into fixed ``max_seq_length`` blocks and hands each DP rank the SAME +number of blocks), LB-Mini keeps samples VARIABLE length and, for every global +minibatch, uses Karmarkar-Karp (see ``lb_mini_packing.plan_minibatch``) to: + + * balance the TOTAL token workload across DP ranks, and + * let each rank own a DIFFERENT number of micro-batches. + +The "different micro-batch count per rank" is exactly what a collective-comm +backend (NCCL all-gather/reduce-scatter) cannot tolerate, which is why this path +is gated behind ``enable_odc=true`` -- ODC's point-to-point comm lets ranks run +out of lockstep without deadlocking. + +Two pieces live here: + 1. ``build_varlen_samples`` -- tokenize the dataset ONCE into variable-length + ``{input_ids, labels, loss_mask, length}`` dicts (cached to disk, reusing + packing.py's tokenizer-identity cache-key machinery), and + 2. ``LBMiniDataIterator`` -- the iterator Megatron's schedule pulls from. It + refills one global minibatch at a time, plans this rank's micro-batches via + KK, and exposes ``current_num_microbatches`` so the schedule patch can run + the right (rank-local) number of forward/backward steps. +""" + +import os +from typing import Dict, List + +import numpy as np +import torch + +from primus.backends.megatron.sft.lb_mini_packing import ( + plan_minibatch, + resolve_cost_func, +) +from primus.core.utils.module_utils import log_rank_0 + + +def _shift_labels_loss_mask(input_ids: np.ndarray, loss_mask: np.ndarray): + """Next-token shift, mirroring packing._build_packed_sequence semantics. + + Megatron's ``compute_language_model_loss`` does NOT shift internally, so the + dataset MUST emit shifted labels: labels[i]=input_ids[i+1], and loss_mask + shifted the same way; the final position has no target and is masked. + """ + n = len(input_ids) + labels = input_ids.copy() + if n >= 2: + labels[:-1] = input_ids[1:] + shifted_mask = np.zeros_like(loss_mask) + if n >= 2: + shifted_mask[:-1] = loss_mask[1:] + shifted_mask[-1] = 0 + return labels, shifted_mask + + +def build_varlen_samples( + dataset_name: str, + tokenizer, + max_seq_length: int, + split: str, + formatter: str, + seed: int, + bridge_compat_inline_bos: bool = False, + **kwargs, +) -> List[Dict[str, np.ndarray]]: + """Tokenize the dataset into variable-length (un-padded) samples, cached. + + Reuses ``PackedSFTDataset``'s tokenize step (``_tokenize_no_pad``) but does + NOT bin-pack: every raw sample becomes one variable-length record. Samples + whose tokenized length is 0 are dropped; samples longer than max_seq_length + are truncated (same as the non-bridge packing path). + """ + from primus.backends.megatron.sft import packing as pack_mod + from primus.backends.megatron.sft.dataset import SFTDataset + + cache_disabled = os.environ.get("PRIMUS_DISABLE_PACK_CACHE", "0") not in ("0", "", "false", "False") + pad_id = pack_mod._resolve_pad_token_id(tokenizer) + cache_file = lock_file = None + if not cache_disabled: + # Distinct cache namespace from packed (different layout): prefix "varlen". + key = pack_mod._build_pack_cache_key( + dataset_name=dataset_name, + split=split, + formatter=formatter, + max_seq_length=max_seq_length, + pad_id=pad_id, + tokenizer_id=pack_mod._tokenizer_identity(tokenizer), + bridge_compat_inline_bos=bridge_compat_inline_bos, + ) + cache_dir = pack_mod._resolve_pack_cache_dir() + cache_file = cache_dir / f"sft_varlen_{key}.pt" + # Node-local filelock (PRIMUS_PACK_LOCK_DIR): avoid NFS "Stale file handle". + _lock_dir = os.environ.get("PRIMUS_PACK_LOCK_DIR") + if _lock_dir: + os.makedirs(_lock_dir, exist_ok=True) + lock_file = os.path.join(_lock_dir, f"sft_varlen_{key}.lock") + else: + lock_file = cache_dir / f"sft_varlen_{key}.lock" + + def _build() -> List[Dict[str, np.ndarray]]: + base = SFTDataset( + dataset_name=dataset_name, + tokenizer=tokenizer, + max_seq_length=max_seq_length, + split=split, + formatter=formatter, + seed=seed, + **kwargs, + ) + raw = base.dataset + log_rank_0(f"[LB-Mini] Tokenizing {len(raw)} samples (variable length, no packing)...") + out: List[Dict[str, np.ndarray]] = [] + for i in range(len(raw)): + sample = pack_mod.normalize_sft_sample(raw[i]) + formatted = base.formatter.format_sample(sample) + tok = pack_mod._tokenize_no_pad( + formatted, tokenizer, max_seq_length, bridge_compat_inline_bos=bridge_compat_inline_bos + ) + if tok["length"] <= 0: + continue + labels, shifted_mask = _shift_labels_loss_mask(tok["input_ids"], tok["loss_mask"]) + out.append( + { + "input_ids": tok["input_ids"].astype(np.int64), + "labels": labels.astype(np.int64), + "loss_mask": shifted_mask.astype(np.int64), + "length": int(tok["length"]), + } + ) + sup = sum(int(s["loss_mask"].sum()) for s in out) + log_rank_0( + f"[LB-Mini] Built {len(out)} variable-length samples; " + f"supervised tokens={sup} (sanity: must be > 0)." + ) + return out + + if cache_disabled: + return _build() + + try: + from filelock import FileLock + except ImportError: + return _build() + + with FileLock(str(lock_file)): + if cache_file.exists(): + log_rank_0(f"[LB-Mini] varlen CACHE HIT ({cache_file.name}); loading.") + samples = torch.load(cache_file, weights_only=False) + log_rank_0(f"[LB-Mini] Loaded {len(samples)} variable-length samples from cache.") + return samples + samples = _build() + tmp = cache_file.with_suffix(".pt.tmp") + torch.save(samples, tmp) + os.replace(tmp, cache_file) + log_rank_0(f"[LB-Mini] Cached variable-length samples to {cache_file}") + return samples + + +class LBMiniDataIterator: + """Iterator that yields THIS rank's variable-length micro-batches. + + Every ``global_batch_size`` raw samples form one *global minibatch*. For each + global minibatch we: + 1. read the per-sample effective lengths (identical view on all ranks), + 2. ``plan_minibatch`` -> KK-balance across DP ranks + split this rank into + micro-batches each <= ``max_token_len``, + 3. push this rank's micro-batches onto a queue and record + ``current_num_microbatches`` (may differ across ranks). + + The schedule patch calls ``begin_minibatch()`` once at the top of each + train_step to materialize the plan and read ``current_num_microbatches``; + then Megatron's ``forward_step`` pulls each micro-batch via ``__next__``. + + A micro-batch here is ONE packed 1-D sequence (the KK-chosen samples for that + micro-batch concatenated). With ``micro_batch_size`` typically 1 and short + samples concatenated up to ``max_token_len``, attention runs causal over the + concatenation (same implicit-multi-turn regime as Primus packed SFT with + ``use_packed_attention=false``); ``loss_mask`` keeps loss on response tokens. + """ + + def __init__( + self, + samples: List[Dict[str, np.ndarray]], + global_batch_size: int, + max_token_len: int, + dp_rank: int, + dp_size: int, + pad_id: int, + cost_model: str = "linear", + seed: int = 1234, + shuffle: bool = True, + same_micro_num: bool = False, + packing_method: str = "kk", + ): + self.samples = samples + self.global_batch_size = int(global_batch_size) + self.max_token_len = int(max_token_len) + self.dp_rank = int(dp_rank) + self.dp_size = int(dp_size) + self.pad_id = int(pad_id) + self.cost_func = resolve_cost_func(cost_model) + self.seed = int(seed) + self.shuffle = bool(shuffle) + # same_micro_num=True forces all ranks to the SAME micro-batch count + # (all_reduce MAX): this is the no-LB-Mini baseline (ranks aligned, the + # short-workload ranks pad/idle). False = LB-Mini (ranks may differ). + self.same_micro_num = bool(same_micro_num) + # packing_method: "kk" = Karmarkar-Karp load balance (LB-Mini); + # "round_robin" = ODC-example-style "None" packing (rank takes + # idx[rank::dp], NO balancing -> uneven workload, for A/B of ODC comm). + self.packing_method = str(packing_method) + + self._order: List[int] = [] + self._cursor = 0 + self._epoch = 0 + self._queue: List[Dict[str, torch.Tensor]] = [] + self.current_num_microbatches = 0 + self._dbg_count = 0 + self._reshuffle() + + def _reshuffle(self): + n = len(self.samples) + if self.shuffle: + g = np.random.default_rng(self.seed + self._epoch) + self._order = g.permutation(n).tolist() + else: + self._order = list(range(n)) + self._cursor = 0 + self._epoch += 1 + + def _next_global_indices(self) -> List[int]: + """Pull the next ``global_batch_size`` sample indices (wraps epochs).""" + if self._cursor + self.global_batch_size > len(self._order): + self._reshuffle() + idx = self._order[self._cursor : self._cursor + self.global_batch_size] + self._cursor += self.global_batch_size + return idx + + def _pack_microbatch(self, sample_indices: List[int]) -> Dict[str, torch.Tensor]: + """Concatenate chosen samples into one variable-length micro-batch. + + Emits thd-compatible fields (cu_seqlens padded to [1, MAX_SEGMENTS+1], + per-segment position_ids, max_sub_seqlen) so that with + ``use_packed_attention=true`` the forward path runs SEGMENTED (thd) + attention -- O(sum seqlen^2) instead of O(total^2). This is what makes a + large max_token_len (packed multi-sample micro-batches) fit in memory. + """ + from primus.backends.megatron.sft.packing import MAX_SEGMENTS_PER_PACK + + ids = np.concatenate([self.samples[i]["input_ids"] for i in sample_indices]) + lbl = np.concatenate([self.samples[i]["labels"] for i in sample_indices]) + msk = np.concatenate([self.samples[i]["loss_mask"] for i in sample_indices]) + seglens = [int(self.samples[i]["length"]) for i in sample_indices] + total = int(sum(seglens)) + # Per-segment position ids (each sub-sample restarts at 0). + pos = np.concatenate([np.arange(s, dtype=np.int64) for s in seglens]) + # cu_seqlens padded to a fixed width so default_collate can stack. + cu = [0] + off = 0 + for s in seglens: + off += s + cu.append(off) + n_seg = len(seglens) + while len(cu) < MAX_SEGMENTS_PER_PACK + 1: + cu.append(total) + cu_np = np.asarray(cu, dtype=np.int32) + max_sub = max(seglens) if seglens else 0 + return { + "input_ids": torch.from_numpy(ids).long().unsqueeze(0), # [1, T] + "labels": torch.from_numpy(lbl).long().unsqueeze(0), + "loss_mask": torch.from_numpy(msk).long().unsqueeze(0), + "position_ids": torch.from_numpy(pos).long().unsqueeze(0), + "cu_seqlens": torch.from_numpy(cu_np).unsqueeze(0), # [1, MAX_SEG+1] + "num_segments": torch.tensor([n_seg], dtype=torch.int32), # [1] + "max_sub_seqlen": torch.tensor([max_sub], dtype=torch.int32), + } + + def begin_minibatch(self) -> int: + """Plan ONE global minibatch; fill this rank's queue. Returns rank count.""" + global_idx = self._next_global_indices() + lengths = [int(self.samples[i]["length"]) for i in global_idx] + if self.packing_method == "round_robin": + # No load balancing (ODC example's default "None"): rank takes + # global_idx[rank::dp], one sample per micro-batch. Per-rank workload + # is UNEVEN (variable lengths) -> this is where ODC's on-demand comm + # is supposed to win by overlapping ranks instead of bulk-syncing. + local_positions = list(range(self.dp_rank, len(global_idx), self.dp_size)) + local_micro = [[p] for p in local_positions] + else: + # KK-balance across DP ranks + split this rank into micro-batches. + local_micro = plan_minibatch( + lengths, + rank=self.dp_rank, + world_size=self.dp_size, + max_token_len=self.max_token_len, + same_micro_num=self.same_micro_num, # False=LB-Mini, True=aligned baseline + get_seq_costs_func=self.cost_func, + ) + # local_micro entries index into ``lengths`` (== position in global_idx); + # map back to absolute sample indices. + self._queue = [self._pack_microbatch([global_idx[p] for p in micro]) for micro in local_micro] + self.current_num_microbatches = len(self._queue) + # Observability: show per-rank micro-batch count for the first few global + # minibatches -- this is where "different count per rank" becomes visible. + if self._dbg_count < 4: + tot = sum(int(mb["input_ids"].numel()) for mb in self._queue) + print( + f"[LB-Mini] rank{self.dp_rank}: num_microbatches={self.current_num_microbatches} " + f"total_tokens={tot}", + flush=True, + ) + self._dbg_count += 1 + return self.current_num_microbatches + + def __iter__(self): + return self + + def __next__(self) -> Dict[str, torch.Tensor]: + if not self._queue: + # Schedule patch should call begin_minibatch() first; as a fallback + # (e.g. eval loops) we refill transparently. + self.begin_minibatch() + return self._queue.pop(0) + + +__all__ = ["build_varlen_samples", "LBMiniDataIterator"] diff --git a/primus/backends/megatron/sft/lb_mini_packing.py b/primus/backends/megatron/sft/lb_mini_packing.py new file mode 100644 index 000000000..365ed80f1 --- /dev/null +++ b/primus/backends/megatron/sft/lb_mini_packing.py @@ -0,0 +1,398 @@ +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""LB-Mini sequence-length load balancing for Megatron-native SFT (ODC path). + +This is the Megatron-side port of the load-balancing core that ODC's official +example (``primus/core/odc/examples/llm_training/packing.py``) uses. The algorithm +itself is unchanged (it is already validated on the ODC example); we only strip +the example's ``from args import get_args`` dependency so it can be driven by +plain function arguments from a Primus monkey-patch. + +What it does +------------ +Given the *effective* sequence lengths of one global minibatch (``minibatch_size +* dp`` samples), it produces a balanced assignment so that: + + 1. ``get_seqlen_balanced_partitions`` (Karmarkar-Karp differencing) splits the + samples across the ``dp`` ranks so each rank's TOTAL token workload is as + equal as possible -- this is the "LB" (load balance) part. + + 2. Within a rank, ``rearrange_micro_batches`` re-packs that rank's samples into + micro-batches each capped at ``max_token_len``. With ``same_num_in_dp=False`` + (the LB-Mini / ODC mode) each rank ends up with a DIFFERENT number of + micro-batches -- the short-workload ranks get fewer, the long ones get more, + and crucially NO rank pads up to the global max. This is the "Mini" + (variable micro-batch count) part that only ODC's point-to-point comm can + drive without a collective deadlock. + +The baseline (``same_num_in_dp=True``) instead all-reduce(MAX)es the micro-batch +count so every rank runs the same number of steps -- exactly Megatron's current +behaviour, kept here so the two paths share one code base. + +IMPORTANT: this module is pure CPU index math; it never touches CUDA or the +process group except for the optional ``all_reduce`` used by the baseline +``same_num_in_dp=True`` path. +""" + +import heapq +import os +from typing import Callable, List, Optional + +import torch +from torch import distributed as dist + + +# --------------------------------------------------------------------------- +# Cost models: map a sequence length to a scalar "workload". +# --------------------------------------------------------------------------- +def get_seq_costs_linear(seq_len: List[int]) -> List[float]: + """Linear cost == token count. Good default for MLP-bound SFT.""" + return list(seq_len) + + +# Default fit coefficients (a, b) + normalizer = ODC example's DeepSeek-1.5B fit. +# A 7B/8B-class model has a different attention/MLP FLOP ratio, so the relative +# weighting of long vs short sequences differs. To keep the cost model +# model-appropriate WITHOUT forking this shared function (existing 1.5B runs stay +# byte-for-byte identical), the coefficients are overridable per-run via env: +# LB_MINI_FIT_A / LB_MINI_FIT_B / LB_MINI_FIT_NORM. +# Unset -> the 1.5B defaults below. For a 7B-class model use a=4.348357, b=8.768377. +_FIT_A_DEFAULT = 1.982122 +_FIT_B_DEFAULT = 2.611821 +_FIT_NORM_DEFAULT = 32000.0 + + +def get_seq_costs_fit(seq_len: List[int]) -> List[float]: + """Quadratic-ish fit (attention-aware): a*s^2 + b*s with s normalized. + + Coefficients default to the ODC example's DeepSeek-1.5B fit and may be + overridden per-run via env (LB_MINI_FIT_A/B/NORM) so a 7B/8B-class model can + use its own (a=4.348357, b=8.768377). Only meaningful when + use_packed_attention=true (segmented thd attn); under full O(total^2) + attention the linear model is the correct target. + """ + a = float(os.environ.get("LB_MINI_FIT_A", _FIT_A_DEFAULT)) + b = float(os.environ.get("LB_MINI_FIT_B", _FIT_B_DEFAULT)) + norm = float(os.environ.get("LB_MINI_FIT_NORM", _FIT_NORM_DEFAULT)) + normalized = [s / norm for s in seq_len] + return [a * s * s + b * s for s in normalized] + + +def resolve_cost_func(name: str) -> Callable[[List[int]], List[float]]: + if name == "linear": + return get_seq_costs_linear + if name == "fit": + return get_seq_costs_fit + raise ValueError(f"Unknown LB-Mini cost model: {name!r}. Supported: linear, fit") + + +# --------------------------------------------------------------------------- +# Karmarkar-Karp largest-differencing partitioning. +# https://en.wikipedia.org/wiki/Largest_differencing_method +# --------------------------------------------------------------------------- +def karmarkar_karp(seq_cost_list: List[float], k_partitions: int, equal_size: bool): + """Partition indices of ``seq_cost_list`` into ``k_partitions`` balanced sets.""" + + class Set: + def __init__(self) -> None: + self.sum = 0 + self.items = [] + + def add(self, idx: int, val: float): + self.items.append((idx, val)) + self.sum += val + + def merge(self, other): + for idx, val in other.items: + self.items.append((idx, val)) + self.sum += val + + def __lt__(self, other): + if self.sum != other.sum: + return self.sum < other.sum + if len(self.items) != len(other.items): + return len(self.items) < len(other.items) + return self.items < other.items + + class State: + def __init__(self, items, k: int) -> None: + self.k = k + self.sets = [Set() for _ in range(k)] + assert len(items) in [1, k], f"{len(items)} not in [1, {k}]" + for i, (idx, seqlen) in enumerate(items): + self.sets[i].add(idx=idx, val=seqlen) + self.sets = sorted(self.sets, reverse=True) + + def get_partitions(self): + partitions = [] + for i in range(len(self.sets)): + cur = [idx for idx, _ in self.sets[i].items] + partitions.append(cur) + return partitions + + def merge(self, other): + for i in range(self.k): + self.sets[i].merge(other.sets[self.k - 1 - i]) + self.sets = sorted(self.sets, reverse=True) + + @property + def spread(self) -> float: + return self.sets[0].sum - self.sets[-1].sum + + def __lt__(self, other): + if self.spread != other.spread: + return self.spread > other.spread + return self.sets[0] > other.sets[0] + + sorted_seq_cost_list = sorted([(seqlen, i) for i, seqlen in enumerate(seq_cost_list)]) + states_pq = [] + if equal_size: + assert len(seq_cost_list) % k_partitions == 0 + for offset in range(0, len(sorted_seq_cost_list), k_partitions): + items = [] + for i in range(k_partitions): + seqlen, idx = sorted_seq_cost_list[offset + i] + items.append((idx, seqlen)) + heapq.heappush(states_pq, State(items=items, k=k_partitions)) + else: + for seqlen, idx in sorted_seq_cost_list: + heapq.heappush(states_pq, State(items=[(idx, seqlen)], k=k_partitions)) + + while len(states_pq) > 1: + state0 = heapq.heappop(states_pq) + state1 = heapq.heappop(states_pq) + state0.merge(state1) + heapq.heappush(states_pq, state0) + + partitions = states_pq[0].get_partitions() + if equal_size: + for partition in partitions: + assert len(partition) * k_partitions == len(seq_cost_list) + return partitions + + +# --------------------------------------------------------------------------- +# Local-search refinement to tighten the KK partitions. +# --------------------------------------------------------------------------- +_EPS = 1e-6 + + +def _swap_max_partition(seq_cost_list, partitions): + max_cost = sum(seq_cost_list[_] for _ in partitions[0]) + min_cost = sum(seq_cost_list[_] for _ in partitions[-1]) + for i, item_idx in enumerate(partitions[0]): + item_cost = seq_cost_list[item_idx] + for j in range(1, len(partitions)): + cost_j = sum(seq_cost_list[_] for _ in partitions[j]) + for k, swap_idx in enumerate(partitions[j]): + swap_cost = seq_cost_list[swap_idx] + if item_cost - swap_cost <= _EPS: + continue + if cost_j - swap_cost + item_cost + _EPS < max_cost and ( + max_cost - item_cost + swap_cost > min_cost + _EPS + ): + partitions[j][k] = item_idx + partitions[0][i] = swap_idx + return True + return False + + +def _swap_min_partition(seq_cost_list, partitions): + max_cost = sum(seq_cost_list[_] for _ in partitions[0]) + min_cost = sum(seq_cost_list[_] for _ in partitions[-1]) + for i, item_idx in enumerate(partitions[-1]): + item_cost = seq_cost_list[item_idx] + for j in range(0, len(partitions) - 1): + cost_j = sum(seq_cost_list[_] for _ in partitions[j]) + for k, swap_idx in enumerate(partitions[j]): + swap_cost = seq_cost_list[swap_idx] + if swap_cost - item_cost <= _EPS: + continue + if cost_j - swap_cost + item_cost > min_cost + _EPS and ( + min_cost - item_cost + swap_cost + _EPS < max_cost + ): + partitions[j][k] = item_idx + partitions[-1][i] = swap_idx + return True + return False + + +def _balance_partition(seq_cost_list, partitions): + while True: + partitions = sorted( + partitions, + key=lambda x: (sum(seq_cost_list[i] for i in x), min(x) if x else 0), + reverse=True, + ) + if _swap_max_partition(seq_cost_list, partitions): + continue + if _swap_min_partition(seq_cost_list, partitions): + continue + break + return partitions + + +def get_seqlen_balanced_partitions( + seqlen_list: List[int], + k_partitions: int, + equal_size: bool, + get_seq_costs_func: Optional[Callable] = None, +) -> List[List[int]]: + """Balance ``seqlen_list`` indices into ``k_partitions`` by total workload. + + equal_size=True -> every partition has the same item count (baseline). + equal_size=False -> partitions may have different item counts (LB-Mini). + """ + if get_seq_costs_func is None: + get_seq_costs_func = get_seq_costs_linear + seq_cost_list = get_seq_costs_func(seqlen_list) + assert len(seq_cost_list) >= k_partitions, f"{len(seq_cost_list)} < {k_partitions}" + + def _check_and_sort(partitions): + assert len(partitions) == k_partitions + seen = set() + out = [None] * k_partitions + for i, partition in enumerate(partitions): + assert len(partition) > 0, f"partition {i} empty" + seen.update(partition) + out[i] = sorted(partition) + assert seen == set(range(len(seq_cost_list))) + return out + + def _diff(partitions): + loads = [sum(seq_cost_list[i] for i in p) for p in partitions] + return max(loads) - min(loads) + + partitions = karmarkar_karp(seq_cost_list, k_partitions, equal_size) + partitions = _balance_partition(seq_cost_list, partitions) + if not equal_size and len(seq_cost_list) % k_partitions == 0: + eq = karmarkar_karp(seq_cost_list, k_partitions, equal_size=True) + eq = _balance_partition(seq_cost_list, eq) + if _diff(partitions) > _diff(eq): + partitions = eq + return _check_and_sort(partitions) + + +def _ceildiv(a, b): + return -(a // -b) + + +def rearrange_micro_batches( + seq_len_effective: List[int], + max_token_len: int, + dp_group=None, + same_num_in_dp: bool = True, + sort_partition_workload: bool = True, + get_seq_costs_func: Optional[Callable] = None, +) -> List[List[int]]: + """Split one rank's samples into micro-batches each <= ``max_token_len``. + + same_num_in_dp=True -> all DP ranks agree on the micro-batch count via + all_reduce(MAX). Megatron's current behaviour. + same_num_in_dp=False -> each rank uses its own count (LB-Mini; ODC only). + """ + if get_seq_costs_func is None: + get_seq_costs_func = get_seq_costs_linear + total_seqlen = sum(seq_len_effective) + num_micro_batches = min(len(seq_len_effective), _ceildiv(total_seqlen, max_token_len)) + if dist.is_initialized() and same_num_in_dp: + t = torch.tensor([num_micro_batches]).cuda() + dist.all_reduce(t, op=dist.ReduceOp.MAX, group=dp_group) + num_micro_batches = t.cpu().item() + + while True: + assert num_micro_batches <= len(seq_len_effective) + micro_bsz_idx = get_seqlen_balanced_partitions( + seq_len_effective, + num_micro_batches, + equal_size=False, + get_seq_costs_func=get_seq_costs_func, + ) + check_failed = False + for partition in micro_bsz_idx: + if sum(seq_len_effective[i] for i in partition) > max_token_len: + check_failed = True + break + actual_size = num_micro_batches + 1 if check_failed else len(micro_bsz_idx) + if dist.is_initialized() and same_num_in_dp: + t = torch.tensor([actual_size]).cuda() + dist.all_reduce(t, op=dist.ReduceOp.MAX, group=dp_group) + actual_size = t.cpu().item() + if actual_size == num_micro_batches: + break + num_micro_batches += 1 + + if sort_partition_workload: + micro_bsz_idx.sort( + key=lambda partition: ( + sum(get_seq_costs_func([seq_len_effective[idx] for idx in partition])), + min(partition) if partition else 0, + ), + reverse=True, + ) + return micro_bsz_idx + + +def plan_minibatch( + lengths: List[int], + rank: int, + world_size: int, + max_token_len: int, + same_micro_num: bool = False, + get_seq_costs_func: Optional[Callable] = None, +) -> List[List[int]]: + """Top-level LB-Mini planner for ONE global minibatch. + + Args: + lengths: effective seq lengths of ALL samples in this global minibatch + (length == minibatch_size * world_size), identical on every rank. + rank/world_size: this rank within the DP group. + max_token_len: per-micro-batch token cap (memory constraint). + same_micro_num: False -> LB-Mini (ODC, ranks differ); True -> baseline. + + Returns: + local_idx: list of micro-batches for THIS rank; each micro-batch is a list + of indices into ``lengths``. ``len(local_idx)`` is this rank's + micro-batch count (may differ across ranks when same_micro_num + is False). + """ + if get_seq_costs_func is None: + get_seq_costs_func = get_seq_costs_linear + assert max(lengths) <= max_token_len, f"{max(lengths)} > max_token_len={max_token_len}" + + # Step 1: KK-balance the whole minibatch across DP ranks. + mini_partitions = get_seqlen_balanced_partitions( + lengths, + world_size, + equal_size=same_micro_num, + get_seq_costs_func=get_seq_costs_func, + ) + local_index = mini_partitions[rank] + local_lengths = [lengths[i] for i in local_index] + + # Step 2: split this rank's slice into micro-batches. + micro_indexes = rearrange_micro_batches( + local_lengths, + max_token_len, + same_num_in_dp=same_micro_num, + sort_partition_workload=True, + get_seq_costs_func=get_seq_costs_func, + ) + # Map rank-local positions back to global minibatch indices. + local_idx = [[local_index[p] for p in micro] for micro in micro_indexes] + return local_idx + + +__all__ = [ + "get_seq_costs_linear", + "get_seq_costs_fit", + "resolve_cost_func", + "karmarkar_karp", + "get_seqlen_balanced_partitions", + "rearrange_micro_batches", + "plan_minibatch", +] diff --git a/primus/backends/megatron/sft/packing.py b/primus/backends/megatron/sft/packing.py index d1dae01f4..4afb7ad64 100644 --- a/primus/backends/megatron/sft/packing.py +++ b/primus/backends/megatron/sft/packing.py @@ -526,7 +526,18 @@ def __init__( ) cache_dir = _resolve_pack_cache_dir() cache_file = cache_dir / f"sft_pack_{cache_key}.pt" - lock_file = cache_dir / f"sft_pack_{cache_key}.lock" + # Keep the lock on node-LOCAL storage when PRIMUS_PACK_LOCK_DIR is set. On a + # shared NFS cache dir, cross-node FileLock contention (e.g. 16 ranks over 2 + # nodes in dual-node ODC) triggers ESTALE ([Errno 116] Stale file handle). + # Routing the lock to node-local /tmp avoids the NFS flock race; the cache + # .pt itself stays shared on NFS (deterministic build, atomic rename). + _lock_base = os.environ.get("PRIMUS_PACK_LOCK_DIR") + if _lock_base: + _lock_dir = Path(_lock_base) + _lock_dir.mkdir(parents=True, exist_ok=True) + else: + _lock_dir = cache_dir + lock_file = _lock_dir / f"sft_pack_{cache_key}.lock" # Lazy import keeps filelock optional for non-cache code paths (and out # of unit tests that don't exercise PackedSFTDataset). diff --git a/primus/configs/models/megatron/deepseek_r1_distill_qwen_1.5B.yaml b/primus/configs/models/megatron/deepseek_r1_distill_qwen_1.5B.yaml new file mode 100644 index 000000000..3bb90f351 --- /dev/null +++ b/primus/configs/models/megatron/deepseek_r1_distill_qwen_1.5B.yaml @@ -0,0 +1,19 @@ +extends: + - qwen2.5_base.yaml + +tokenizer_type: HuggingFaceTokenizer +tokenizer_model: deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B + +# DeepSeek-R1-Distill-Qwen-1.5B (Qwen2 architecture) specific parameters. +# This is the exact model the ODC reference repo (sail-sg/odc) uses in its +# examples/llm_training default config, so we use it to validate ODC/LB-Mini +# on Megatron at a size that fits a 64k-token micro-batch without OOM. +hidden_size: 1536 +ffn_hidden_size: 8960 +num_layers: 28 +num_attention_heads: 12 +num_query_groups: 2 + +# IMPORTANT: DeepSeek-R1-Distill-Qwen-1.5B uses rope_theta=10000, NOT the +# Qwen2.5 default of 1000000 inherited from qwen2.5_base.yaml. Override it. +rotary_base: 10000 diff --git a/primus/configs/modules/megatron/sft_trainer.yaml b/primus/configs/modules/megatron/sft_trainer.yaml index 32da1b266..a0ba26557 100644 --- a/primus/configs/modules/megatron/sft_trainer.yaml +++ b/primus/configs/modules/megatron/sft_trainer.yaml @@ -63,6 +63,9 @@ use_distributed_optimizer: true # Note: These are typically overridden in the experiment config sft_dataset_name: "tatsu-lab/alpaca" sft_conversation_format: "alpaca" +# HuggingFace dataset split to load. Some datasets have no "train" split (e.g. +# SWE-bench/SWE-smith-trajectories use tool/xml/ticks); override here. +sft_dataset_split: "train" # Sequence packing: concatenate multiple short samples into a single # max_seq_length sequence. Big throughput win on short-sample datasets such @@ -84,6 +87,29 @@ enable_packed_sequences: false # short-sample datasets. use_packed_attention: false +# === LB-Mini (ODC sequence-length load balancing) === +# When true (and enable_odc=true + use_torch_fsdp2=true), the TRAIN data is +# served variable-length and Karmarkar-Karp balanced across DP ranks, and each +# rank runs its OWN (possibly different) number of micro-batches -- removing the +# workload-imbalance bubble on variable-length datasets (e.g. LongAlign). +# Requires ODC point-to-point comm; it would deadlock under NCCL collectives. +# DEFAULT false => stock Megatron fixed-num_microbatches, all-ranks-in-lockstep +# schedule, byte-for-byte unchanged. +enable_odc_lb_mini: false +# Per-micro-batch token cap for LB-Mini (0 => use seq_length). +lb_mini_max_token_len: 0 +# LB-Mini cost model: "linear" (token count) or "fit" (attention-aware, ~s^2). +lb_mini_cost_model: "linear" + +# === Chunked linear cross-entropy (avoid full-logits OOM) === +# When true (or env FUSED_LINEAR_CE=1), GPTModel computes the LM loss in +# sequence chunks under activation checkpointing, so the full [seq, vocab] +# logits is never materialized. Essential for large-vocab models (e.g. Qwen +# 151936) at long sequence lengths where the logits tensor alone would OOM. +# DEFAULT false => stock full-logits path, byte-for-byte unchanged. Numerically +# equivalent when on. Tune chunk via env FUSED_CE_CHUNK (default 4096). +enable_fused_linear_ce: false + # Opt-in: reproduce NeMo Megatron-Bridge's packed dataset byte-for-byte. # # When False (default) Native does a clean single-shot tokenize of the full diff --git a/primus/configs/modules/megatron/trainer_base.yaml b/primus/configs/modules/megatron/trainer_base.yaml index 8e412f038..6830dd8b0 100755 --- a/primus/configs/modules/megatron/trainer_base.yaml +++ b/primus/configs/modules/megatron/trainer_base.yaml @@ -193,6 +193,39 @@ keep_fp8_transpose_cache_when_using_custom_fsdp: false num_distributed_optimizer_instances: 1 # int data_parallel_replicate_degree: 1 use_torch_fsdp2: false + +# === ODC (On-Demand Communication) === +# Master switch for ODC's rocSHMEM/XGMI point-to-point gradient reduction on the +# torch-FSDP2 path (ROCm/MI300X). DEFAULT false => ODC is a complete no-op and +# every ODC integration patch is skipped (byte-for-byte stock behaviour). When +# true (requires use_torch_fsdp2=true), Primus: +# * routes FSDP2 gradient reduce-scatter through ODC (scatter_accumulate), +# * skips #856's eager-RCCL device_id injection (distributed_init_patches.py), +# * skips #808's FSDP2 forward all-gather prefetch (it cannot overlap ODC's +# serial transport); backward prefetch is kept. +# Formerly gated by the ODC_ENABLE env var; now config-driven. +enable_odc: false +# ODC wiring depth (integer): 1 = install the __init__/lazy_init hooks only +# (construction + first fwd/bwd sanity, gradients NOT routed); 2 = also hook +# train_step / optimizer.step for full gradient routing. Production = 2. +odc_phase: 2 +# Guard against ODC async-reduce's occasional non-deterministic grad spike: at +# the optimizer step, if the global grad norm exceeds this threshold, the step is +# skipped (grads zeroed) exactly like Megatron skips nan/inf iters. <= 0 disables. +odc_grad_spike_threshold: 1000.0 +# ODC symmetric-memory / comm-kernel tuning (consumed by the ODC library via its +# runtime config; these replace the former ODC_* env vars). Defaults reproduce the +# previous env defaults exactly. +odc_p2p_backend: mori # symmetric-memory backend: mori | rocshmem +odc_mori_init: pg # MORI init method: pg | uid +odc_max_buffer_size: 67108864 # BufferSplitter max global buffer (bytes, 64 MiB) +odc_rocshmem_gda: false # rocSHMEM GPU-direct (GDA) device path (multi-node) +odc_rocshmem_lib: null # optional librs_host_gda.so ctypes override path +odc_gda_rs_blocks: 64 # GDA reduce-scatter kernel grid blocks +odc_gda_pipe: 1 # GDA peer-pipeline depth (also bridged to turbo env) +odc_gda_defer_reduce: auto # defer cross-node reduce to per-minibatch: auto|1|0 +odc_gda_warmup_mode: strided # write-visibility settle: strided|full|hdp|fence|hdpfence +odc_gda_stride_bytes: 65536 # strided warm-up page stride (bytes) nccl_communicator_config_path: null use_tp_pp_dp_mapping: false replication: false diff --git a/primus/core/odc/.gitignore b/primus/core/odc/.gitignore new file mode 100644 index 000000000..4aab3e88c --- /dev/null +++ b/primus/core/odc/.gitignore @@ -0,0 +1,229 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py.cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +# Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +# poetry.lock +# poetry.toml + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. +# https://pdm-project.org/en/latest/usage/project/#working-with-version-control +# pdm.lock +# pdm.toml +.pdm-python +.pdm-build/ + +# pixi +# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. +# pixi.lock +# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one +# in the .venv directory. It is recommended not to include this directory in version control. +.pixi + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# Redis +*.rdb +*.aof +*.pid + +# RabbitMQ +mnesia/ +rabbitmq/ +rabbitmq-data/ + +# ActiveMQ +activemq-data/ + +# SageMath parsed files +*.sage.py + +# Environments +.env +.envrc +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +# .idea/ + +# Abstra +# Abstra is an AI-powered process automation framework. +# Ignore directories containing user credentials, local state, and settings. +# Learn more at https://abstra.io/docs +.abstra/ + +# Visual Studio Code +# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore +# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore +# and can be added to the global gitignore or merged into this file. However, if you prefer, +# you could uncomment the following to ignore the entire vscode folder +# .vscode/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Marimo +marimo/_static/ +marimo/_lsp/ +__marimo__/ + +# Streamlit +.streamlit/secrets.toml + +# Custom +.DS_Store + +# rocSHMEM build artifacts. The ODC rocSHMEM ops now live in Primus-Turbo; the +# rocSHMEM static library is an external dependency (ROCSHMEM_HOME), built +# out-of-tree. Never commit any rocSHMEM install tree / static lib / build log. +*.a +rocshmem_runtime/rocshmem_src/ +rocshmem_runtime/rocshmem_single/ +rocshmem_runtime/rocshmem_ro/ +rocshmem_runtime/rocshmem_gda/ +rocshmem_runtime/build.log diff --git a/primus/core/odc/README.md b/primus/core/odc/README.md new file mode 100644 index 000000000..0e8c81d16 --- /dev/null +++ b/primus/core/odc/README.md @@ -0,0 +1,163 @@ +# On-demand Communication (ODC) + +ODC is a patch to FSDP that adapts Parameter Server (PS) into Fully Sharded Data Parallel (FSDP) by replacing collective all-gather and +reduce-scatter with on-demand point-to-point communication. + +![Original-FSDP](./docs/readme/FSDP-ODC.jpg) + +With ODC, the synchronization frequency is reduced from per-iteration to per-minibatch, which fundamentally reduces the workload-imbalance bubbles in FSDP. + +ODC is accepted in ICLR 2026! Check out the [paper](https://openreview.net/pdf?id=iIEEgI6WsF) for more details. + +## Attribution / Provenance + +This ODC code is ported from the upstream open-source project +[sail-sg/odc](https://github.com/sail-sg/odc) (Sea AI Lab, ICLR 2026). Per its +package metadata the upstream project is released under the MIT License +(copyright held by the original authors, Sea AI Lab); the upstream repository +does not ship a standalone `LICENSE` file or per-file license headers. + +This version has been adapted for the AMD ROCm / MI300X platform, including +migrating the communication backend from nvshmem/CUDA to rocSHMEM/MORI + HIP, +device-side reduce, and torchrun-based launching. + +## ODC Primitives + +The key idea is to replace the collective all-gather and reduce-scatter with a on-demand point-to-point communication. + +![ODC Primitives](./docs/readme/ag_rs.png) + +To support transparent on-demand communications, we implement RDMA based primitives on ROCm using HIP GPU IPC (intra-node) and MORI-SHMEM / rocSHMEM (inter-node). Details are in [odc/primitives](./odc/primitives). + +## Support for FSDP +- FSDP1 +- FSDP2 + - HSDP + - `reshard_after_forward=int` + +## Usage + +### Prerequisites + +- PyTorch (ROCm build) +- ROCm 7.x +- Python >= 3.8 +- `amd_mori` (symmetric-memory / RDMA backend) + +We highly recommend using a ROCm PyTorch base image (e.g. the +`tasimage/primus-odc` images, ROCm 7.2.0 based). + +### Enable ODC +ODC ships as an **in-tree module of Primus** (no separate `pip install`; it has no +compiled extension). Put its parent directory on `PYTHONPATH` so `import odc` +resolves, plus the `odc_early` shim dir so the MORI/TE load-order fix runs at +interpreter startup: +``` +export PYTHONPATH=/primus/core:/primus/core/odc/odc_early:$PYTHONPATH +``` +(The launcher `rocshmem_runtime/scripts/run_odc.sh` sets this up automatically.) +ODC is pure Python. The rocSHMEM backend (single-node XGMI IPC host +API and multi-node GPU-direct GDA) is provided by Primus-Turbo as the +`primus_turbo.pytorch._C.odc_rocshmem_host` / `odc_rocshmem_gda` pybind +submodules and consumed by `odc/primitives/_rocshmem_backend.py`; select it with +the config items `odc_p2p_backend: rocshmem` (and `odc_rocshmem_gda: true` for +the multi-node GDA path). Ensure a Primus-Turbo build with the ODC rocSHMEM ops +is importable (installed, or on `PYTHONPATH`). + +> TODO(primus-turbo): pin the exact Primus-Turbo merge commit (`PRIMUS_TURBO_COMMIT`) +> once the PR adding the ODC rocSHMEM ops is merged. + +## Quick Start + +A complete example is provided in `examples/llm_training/`: +```shell +pip install -r examples/llm_training/requirements.txt + +bash examples/llm_training/run.sh +``` + +## Memory +> User may need to tune `MORI_SHMEM_HEAP_SIZE` for better memory usage. +When using ODC in FSDP, symmetric buffers are allocated for sharded parameters, sharded gradient accumulation buffer, miscellaneous buffers for `gather` and `scatter-accumulate`. +To achieve smallest memory footprint, +`MORI_SHMEM_HEAP_SIZE` should be set to be slightly higher than the size of sharded parameters and gradient: params.element_size() * params.numel() + grad_reduce_buf.element_size() * grad_reduce_buf.numel(). + +### Basic Usage with FSDP1 + +```python +import torch +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + +import odc +from odc.fsdp import fsdp1 + + +fsdp1.patch_fsdp1() + +torch.distributed.init_process_group(backend="nccl", device_id=device) +odc.init_shmem() + + +fsdp_model = FSDP( + model, + # ... +) + +for epoch in range(10): + for minibatch in dataset: + fsdp1.pre_minibatch_start(fsdp_model) + loss = loss_fn(model) + loss.backward() + fsdp1.pre_optimizer_step(model) + optimizer.step() + optimizer.zero_grad() + +fsdp1.stop() +``` + +### Basic Usage with FSDP2 + +```python +import torch +import odc +from odc.fsdp import fsdp2 + + +torch.distributed.init_process_group(backend="nccl", device_id=device) +odc.init_shmem() + +fsdp2.patch_fsdp2() + +for layer in model.layers: + fully_shard(layer, **fsdp_kwargs) +fsdp_model = fully_shard(model, **fsdp_kwargs) + +# Call patch_lazy_init just as how we call fully_shard above. +for layer in fsdp_model.layers: + fsdp2.patch_lazy_init(layer) +fsdp2.patch_lazy_init(fsdp_model) + +for epoch in range(10): + for minibatch in dataset: + fsdp2.pre_minibatch_start(fsdp_model) + loss = loss_fn(model) + loss.backward() + fsdp2.pre_optimizer_step(model) + optimizer.step() + optimizer.zero_grad() + +fsdp2.stop() +``` + + +## Development + +### Running Linter +``` +make lint +``` + +### Running Tests +```bash +make test +``` diff --git a/primus/core/odc/__init__.py b/primus/core/odc/__init__.py new file mode 100644 index 000000000..a649c84bc --- /dev/null +++ b/primus/core/odc/__init__.py @@ -0,0 +1,57 @@ +# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under +# the MIT License per its package metadata (pyproject.toml / setup.py +# classifiers). The upstream repository ships no LICENSE file or per-file +# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab). +# +# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# See LICENSE for license information. + +import logging + +# Eagerly expose the runtime-config API only. This module is a stdlib-only leaf, +# so importing `odc` stays cheap and does NOT pull in odc.primitives -- which is +# essential: the ODC integration patch calls odc.set_runtime_config(...) BEFORE +# the primitives are first imported, so their import-time backend selection +# (odc_p2p_backend) reads the populated config rather than a stale default. +from odc.runtime_config import OdcRuntimeConfig +from odc.runtime_config import get_config as get_runtime_config # noqa: F401 +from odc.runtime_config import set_config as set_runtime_config + +logger = logging.getLogger(__name__) +logger.addHandler(logging.NullHandler()) +logger.setLevel(logging.INFO) + + +# The heavy primitives (which import torch/triton/mori and read the runtime +# config at import time) are exposed lazily via PEP 562 module __getattr__, so +# they are imported only on first access -- after set_runtime_config has run. +_LAZY_EXPORTS = { + "init_shmem": ("odc.primitives.utils", "init_shmem"), + "finalize_distributed": ("odc.primitives.utils", "finalize_distributed"), + "SymmBufferRegistry": ("odc.primitives.utils", "SymmBufferRegistry"), + "ReductionService": ("odc.primitives.scatter_accumulate", "ReductionService"), + "GatherService": ("odc.primitives.gather", "GatherService"), +} + + +def __getattr__(name): + target = _LAZY_EXPORTS.get(name) + if target is None: + raise AttributeError(f"module 'odc' has no attribute '{name}'") + import importlib + + mod = importlib.import_module(target[0]) + return getattr(mod, target[1]) + + +__all__ = [ + "init_shmem", + "SymmBufferRegistry", + "ReductionService", + "GatherService", + "finalize_distributed", + "OdcRuntimeConfig", + "get_runtime_config", + "set_runtime_config", +] diff --git a/primus/core/odc/fsdp/__init__.py b/primus/core/odc/fsdp/__init__.py new file mode 100644 index 000000000..999672bdf --- /dev/null +++ b/primus/core/odc/fsdp/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. diff --git a/primus/core/odc/fsdp/fsdp1.py b/primus/core/odc/fsdp/fsdp1.py new file mode 100644 index 000000000..9b94101cd --- /dev/null +++ b/primus/core/odc/fsdp/fsdp1.py @@ -0,0 +1,281 @@ +# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under +# the MIT License per its package metadata (pyproject.toml / setup.py +# classifiers). The upstream repository ships no LICENSE file or per-file +# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab). +# +# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# See LICENSE for license information. + +import logging + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.distributed.fsdp._flat_param import FlatParamHandle, HandleShardingStrategy +from torch.distributed.fsdp._runtime_utils import _div_if_needed, _FSDPState + +import odc + +logger = logging.getLogger(__name__) + +reduction_service = None +gather_service = None + + +def get_reduction_service(): + return reduction_service + + +def get_gather_service(): + return gather_service + + +def custom_get_reduce_scatter_tensors( + state: _FSDPState, unsharded_grad: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Returns the input and output tensors to reduce-scatter, respectively. + """ + chunks = list(unsharded_grad.chunk(state.world_size)) + numel_to_pad = state.world_size * chunks[0].numel() - unsharded_grad.numel() + padded_unsharded_grad = F.pad(unsharded_grad, [0, numel_to_pad]) if numel_to_pad > 0 else unsharded_grad + return padded_unsharded_grad + + +def _reduce_grad(state, handle) -> None: + """ + For sharded strategies, this runs gradient reduction, sharded gradient + accumulation if needed, and the post-reduction callback. + """ + + flat_param = handle.flat_param + uses_hybrid_sharded_strategy = handle._sharding_strategy in ( + HandleShardingStrategy.HYBRID_SHARD, + HandleShardingStrategy._HYBRID_SHARD_ZERO2, + ) + # We clear `.grad` to permit multiple backwards. This avoids a race where + # the second backward pass computation precedes ahead of the first backward + # pass reduction, which is possible since the reduction is issued in a + # separate stream and is async and would result in reducing the wrong + # gradient. + unsharded_grad = flat_param.grad.data + flat_param.grad = None + padded_unsharded_grad = custom_get_reduce_scatter_tensors(state, unsharded_grad) + assert state._comm_hook is None, "ODC does not support comm hook" + if state._comm_hook is None: # default path + _div_if_needed(padded_unsharded_grad, state._gradient_predivide_factor) + pg = handle._fake_process_group if handle._use_fake_reduce else state.process_group + + rs_func = get_reduction_service().scatter_accumulate + rs_func(id(handle.flat_param), padded_unsharded_grad, pg) + handle.flat_param._saved_grad_shard = get_reduction_service().get_accumulation(id(handle.flat_param)) + + assert not uses_hybrid_sharded_strategy, "ODC does not support hybrid sharded strategy" + else: + pass + # NOTE: HSDP variants do not support communication hook. + + assert not handle._offload_params, "ODC does not support offloading" + assert not handle._use_orig_params, "ODC does not support using original parameters" + + +def prepare_gradient_for_optim(self): + """Prepare the gradient for optimizer computation by moving the sharded gradient to the ``.grad`` attribute.""" + from torch.distributed.utils import _p_assert + + def cast_grad_to_param_dtype_if_needed(flat_param): + # TODO (rohan-varma): test for full precision with keep_low_precision_grads + if not self._force_full_precision and self._keep_low_precision_grads: + _p_assert(flat_param.grad is not None, "Unexpected None grad!") + if flat_param.grad.dtype != self._fwd_bwd_param_dtype: + flat_param.grad.data = flat_param.grad.to(self._fwd_bwd_param_dtype) + if self._use_orig_params: + self._use_sharded_grad_views() + + flat_param = self.flat_param + # TODO (awgu): We should replace these conditional checks to encode + # the logical intention more directly. + if hasattr(flat_param, "_cpu_grad"): + # NOTE: This branch includes `NO_SHARD`. + self._check_sharded(flat_param) + self._check_on_cpu(flat_param) + flat_param.grad = flat_param._cpu_grad # type: ignore[attr-defined] + cast_grad_to_param_dtype_if_needed(flat_param) + elif hasattr(flat_param, "_saved_grad_shard"): + self._check_sharded(flat_param) + self._check_on_compute_device(flat_param) + if flat_param._saved_grad_shard is not None: + self._check_on_compute_device(flat_param._saved_grad_shard) # type: ignore[attr-defined] + # If no sharded gradient was computed this iteration, then there is + # no need to forward `_saved_grad_shard` to `grad` + if flat_param._post_backward_called: # type: ignore[attr-defined] + flat_param.grad = None # type: ignore[attr-defined] + if flat_param.grad is not None: + cast_grad_to_param_dtype_if_needed(flat_param) + else: + _p_assert( + not self.uses_sharded_strategy or not flat_param._post_backward_called, # type: ignore[attr-defined] + "All sharded parameters that received a gradient in the " + "post-backward should use `_saved_grad_shard`", + ) + # Delete `_saved_grad_shard` since its existence indicates a previous + # gradient to accumulate with in the post-backward hook + if hasattr(flat_param, "_saved_grad_shard"): + delattr(flat_param, "_saved_grad_shard") + + +def all_gather_flat_param(self, padded_unsharded_flat_param): + """ + All-gather the handle's flat parameter to the destination ``padded_unsharded_flat_param``. + + Then switch to use the all-gathered tensor. + """ + from torch.distributed.fsdp._common_utils import _no_dispatch_record_stream + from torch.distributed.utils import _p_assert + + _p_assert( + hasattr(self, "process_group") and hasattr(self, "world_size"), + "Expects a process group and world size to have been set via `shard()`", + ) + sharded_flat_param = self.flat_param.data + expected_numel = sharded_flat_param.numel() * self.world_size + _p_assert( + padded_unsharded_flat_param.numel() == expected_numel, + f"Expects {expected_numel} numel but got {padded_unsharded_flat_param.numel()}", + ) + + pg = self._fake_process_group if self._use_fake_all_gather else self.process_group + + # HACK this should be handled by C10D + if sharded_flat_param.is_cpu: # type: ignore[attr-defined] + tensor_list = list( + torch.chunk( + padded_unsharded_flat_param, + dist.get_world_size(pg), # type: ignore[arg-type] + ) + ) + dist.all_gather(tensor_list, sharded_flat_param, group=pg) + else: + # padded_unsharded_flat_param output torch.bfloat16 sharded_flat_param input torch.float32 + assert padded_unsharded_flat_param.dtype == self._fwd_bwd_param_dtype + assert sharded_flat_param.dtype == self._orig_param_dtype + # Handle dtype conversion for mixed precision + # We need to convert to output to orig_dtype for all-gather + # as the sharded_flat_param is in orig_dtype for all ranks, then convert back + # This is because when using ODC, the peer may have not reached here and be able to + # convert input to _fwd_bwd_param_dtype. + from torch.distributed.utils import _free_storage + + needs_dtype_conversion = padded_unsharded_flat_param.dtype != sharded_flat_param.dtype + if needs_dtype_conversion: + assert self._uses_param_mixed_precision + # Convert output to orig_dtype for all-gather (creates new tensor) + padded_unsharded_flat_param_orig_dtype = padded_unsharded_flat_param.to(self._orig_param_dtype) + else: + padded_unsharded_flat_param_orig_dtype = padded_unsharded_flat_param + + ag_func = gather_service.gather_into_tensor + ag_func( + padded_unsharded_flat_param_orig_dtype, # Could be fp32 + # padded_unsharded_flat_param, + sharded_flat_param, + pg, + ) + + # Convert back to fwd_bwd_param_dtype and free the temporary buffer + if needs_dtype_conversion: + # Convert the dtype conversion buffer back to fwd_bwd_param_dtype in-place + # copy_ will do the implicit conversion. + padded_unsharded_flat_param.copy_(padded_unsharded_flat_param_orig_dtype) + # Free the temporary buffer + _free_storage(padded_unsharded_flat_param_orig_dtype) + + if self._offload_params: + # In case of offloading, `flat_param.data` (i.e. sharded param) is + # created on the pre-unshard stream. We need to hand it over to the + # unshard stream for all-gather + _no_dispatch_record_stream( + sharded_flat_param, + self._device_handle.current_stream(), # unshard_stream + ) + return padded_unsharded_flat_param + + +old_get_shard = FlatParamHandle._get_shard + + +def custom_get_shard(tensor, rank, world_size): + from torch.distributed.fsdp._flat_param import FlatParameter + + sharded, padded = old_get_shard(tensor, rank, world_size) + assert isinstance(tensor, FlatParameter) + sharded_in_shmem = odc.SymmBufferRegistry.get_instance().update_symm_buffer(id(tensor), sharded) + return sharded_in_shmem, padded + + +def _use_low_precision_shard(self): + """ + Allocate on the compute device and switch to using the low precision sharded flat parameter. + call path: + _runtime_utils.py:_unshard() + -> pre_unshard() + -> self._use_low_precision_shard() + """ + self._check_low_precision_shard() + + +def _free_low_precision_sharded_param(self): + """ + Frees the low precision sharded flat parameter. + + call path: + _runtime_utils.py:_unshard() + -> handle.post_unshard() + -> self._free_low_precision_sharded_param() + """ + self._check_low_precision_shard() + + +def patch_fsdp1(reduce_dtype=None): + global reduction_service, gather_service + reduction_service = odc.ReductionService(accumulation_dtype=reduce_dtype) + gather_service = odc.GatherService() + from torch.distributed.fsdp import _runtime_utils + + _runtime_utils._reduce_grad = _reduce_grad + + FlatParamHandle.prepare_gradient_for_optim = prepare_gradient_for_optim + FlatParamHandle._get_shard = custom_get_shard + FlatParamHandle._all_gather_flat_param = all_gather_flat_param + FlatParamHandle._use_low_precision_shard = _use_low_precision_shard + FlatParamHandle._free_low_precision_sharded_param = _free_low_precision_sharded_param + + +def pre_optimizer_step(fsdp_module): + + assert isinstance(fsdp_module, _FSDPState) + with torch.cuda.nvtx.range("scatter_accumulate_sync"): + get_reduction_service().sync(fsdp_module.process_group) + + for acc in get_reduction_service().accumulations: + if hasattr(fsdp_module, "_inter_node_pg"): + dist.all_reduce(acc, group=fsdp_module._inter_node_pg) + _div_if_needed(acc, fsdp_module._gradient_postdivide_factor) + for handle in fsdp_module._all_handles: + handle.flat_param.grad = ( + get_reduction_service().get_accumulation(id(handle.flat_param)).to(handle.flat_param.dtype) + ) + + +def pre_minibatch_start(_fsdp_module): + get_reduction_service().clear_accumulations() + + # Make sure optimizer updates are visible to all ranks + dist.barrier() + + +def stop(): + get_reduction_service().stop() + odc.SymmBufferRegistry.get_instance().finalize() + odc.finalize_distributed() diff --git a/primus/core/odc/fsdp/fsdp2.py b/primus/core/odc/fsdp/fsdp2.py new file mode 100644 index 000000000..5d5e67733 --- /dev/null +++ b/primus/core/odc/fsdp/fsdp2.py @@ -0,0 +1,1092 @@ +# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under +# the MIT License per its package metadata (pyproject.toml / setup.py +# classifiers). The upstream repository ships no LICENSE file or per-file +# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab). +# +# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# See LICENSE for license information. + +import dataclasses +import logging +import operator +import types +from functools import reduce +from itertools import chain +from typing import Any, Callable, Optional, cast + +import torch +import torch.distributed as dist +from torch import nn +from torch.distributed.device_mesh import _get_device_handle +from torch.distributed.fsdp import fully_shard +from torch.distributed.fsdp._fully_shard import _fsdp_collectives +from torch.distributed.fsdp._fully_shard._fsdp_api import AllGather, ReduceScatter +from torch.distributed.fsdp._fully_shard._fsdp_collectives import ( + AllGatherResult, + DefaultAllocMixin, + _div_if_needed, + _get_all_gather_input_metadatas, + _get_gradient_divide_factors, + _get_param_all_gather_inputs, + foreach_reduce_scatter_copy_in, +) +from torch.distributed.fsdp._fully_shard._fsdp_common import ( + FSDPMeshInfo, + HSDPMeshInfo, + TrainingState, + _get_dim0_padded_size, + _raise_assert_with_print, + _to_dtype_if_needed, + compiled_autograd_enabled, +) +from torch.distributed.fsdp._fully_shard._fsdp_param import ( + FSDPParam, + ShardedState, + set_requires_grad_if_needed, +) +from torch.distributed.fsdp._fully_shard._fsdp_param_group import ( + AllReduceState, + FSDPParamGroup, + ReduceScatterState, +) +from torch.distributed.tensor import DTensor +from torch.profiler import record_function + +from odc.primitives.gather import GatherService +from odc.primitives.scatter_accumulate import ReductionService +from odc.primitives.utils import SymmBufferRegistry, finalize_distributed + +logger = logging.getLogger(__name__) + + +class nvtx_record_function(record_function): + def __enter__(self): + torch.cuda.nvtx.range_push(self.name) + return super().__enter__() + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + torch.cuda.nvtx.range_pop() + + +reduction_service = None +gather_service = None + + +def get_reduction_service(): + return reduction_service + + +def get_gather_service(): + return gather_service + + +class ODCAllGather(DefaultAllocMixin, AllGather): + def __call__( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + group: dist.ProcessGroup, + async_op: bool = False, + ) -> Optional[dist.Work]: + gather = get_gather_service() + gather.gather_into_tensor(output_tensor, input_tensor, group) + if async_op: + event = torch.cuda.Event() + event.record() + return event + return None + + +def get_fsdp_params_key(fsdp_params: list[FSDPParam]) -> str: + ids = "_".join([str(id(param)) for param in fsdp_params]) + return f"fsdp_params_gather_{ids}" + + +def get_hpz_params_key(fsdp_param_key: str) -> str: + return f"hpz_{fsdp_param_key}" + + +@torch.no_grad() +def patch_lazy_init(fsdp_model): + state = fsdp_model._get_fsdp_state() + group = state._fsdp_param_group + prev_lazy_init = group.lazy_init + + def patched_lazy_init(_self): + prev_lazy_init() + replace_sharded_param_with_symm_buffer(fsdp_model) + + group.lazy_init = types.MethodType(patched_lazy_init, group) + + +@torch.no_grad() +def replace_sharded_param_with_symm_buffer( + fsdp_model, +) -> torch.Tensor: + state = fsdp_model._get_fsdp_state() + fsdp_param_group = state._fsdp_param_group + if fsdp_param_group is None: + return + + is_hpz = fsdp_param_group._use_post_forward_mesh + if _enable_hpz is None: + raise ValueError("Need to run patch_fsdp2() first") + assert _enable_hpz == is_hpz, "If HPZ is enabled, reshard_after_forward must be set to int" + + dtype = fsdp_param_group.fsdp_params[0]._sharded_param_data.dtype + hpz_dtype = fsdp_param_group.fsdp_params[0].param_dtype + for fsdp_param in fsdp_param_group.fsdp_params: + if fsdp_param._sharded_param_data.dtype != dtype: + raise ValueError( + f"All FSDP parameters must have the same dtype: {fsdp_param._sharded_param_data.dtype=} {dtype=}" + ) + if fsdp_param.param_dtype != hpz_dtype: + raise ValueError( + f"All FSDP parameters must have the same param dtype: {fsdp_param.param_dtype=} {hpz_dtype=}" + ) + if is_hpz and hpz_dtype is None: + logger.warning( + f"mixed precision param_dtype is not set but HPZ is enabled, using the same as the original dtype {dtype}" + ) + hpz_dtype = dtype + + total_size = sum(fsdp_param._sharded_param_data.numel() for fsdp_param in fsdp_param_group.fsdp_params) + + hpz_sharded_param_total_size = 0 + num_nodes = 1 + hpz_symm_buffer = None + post_forward_mesh_info = fsdp_param_group.post_forward_mesh_info + # This key needs to be the same as the one used in `foreach_all_gather` + key = get_fsdp_params_key(fsdp_param_group.fsdp_params) + if not is_hpz: + symm_buffer = SymmBufferRegistry.get_instance().get_or_create_symm_buffer(key, (total_size,), dtype) + else: + # When HPZ is enabled, we keep the original parameters(mostly fp32) + # in the torch tensor just like the original FSDP2. + # Pre-allocate the symmetric buffer for HPZ (Hierarchical Partitioning for ZeRO in ZeRO++). + assert isinstance(post_forward_mesh_info, HSDPMeshInfo), f"{post_forward_mesh_info=}" + shard_world_size = post_forward_mesh_info.shard_mesh_size + world_size = fsdp_param_group._all_gather_process_group.size() + assert world_size % shard_world_size == 0, f"{world_size=} {shard_world_size=}" + num_nodes = world_size // shard_world_size + hpz_sharded_param_total_size = total_size * num_nodes + hpz_key = get_hpz_params_key(key) + hpz_symm_buffer = SymmBufferRegistry.get_instance().get_or_create_symm_buffer( + hpz_key, (hpz_sharded_param_total_size,), hpz_dtype + ) + logger.info( + f"Replacing HPZ params with symmetric buffer of shape {hpz_symm_buffer.shape} and dtype: {hpz_dtype}" + ) + + offset = 0 + hpz_offset = 0 + # Refer to the codes in `FSDPParam._init_sharded_param` + for fsdp_param in fsdp_param_group.fsdp_params: + param_size = fsdp_param._sharded_param_data.numel() + if not is_hpz: + numel = reduce(operator.mul, fsdp_param.padded_sharded_param_size, 1) + assert numel == param_size, f"{fsdp_param.padded_sharded_param_size=} {param_size=}" + symm_buffer[offset : offset + param_size].copy_(fsdp_param._sharded_param_data) + padded_sharded_param = symm_buffer[offset : offset + param_size] + padded_sharded_param = padded_sharded_param.view(fsdp_param.padded_sharded_param_size) + offset += param_size + + fsdp_param._sharded_param_data = padded_sharded_param.view(-1) + shard_dim = fsdp_param.fsdp_placement.dim + old_sharded_param = fsdp_param.sharded_param._local_tensor + length = old_sharded_param.size(shard_dim) if old_sharded_param.numel() > 0 else 0 + sharded_param = padded_sharded_param.narrow(dim=shard_dim, start=0, length=length) + fsdp_param.sharded_param._local_tensor = sharded_param + else: + hpz_param_size = param_size * num_nodes + fsdp_param._sharded_post_forward_param_data = hpz_symm_buffer[ + hpz_offset : hpz_offset + hpz_param_size + ] + hpz_offset += hpz_param_size + + +__odc_gather = None + + +def get_odc_gather_comm(): + global __odc_gather + if __odc_gather is None: + __odc_gather = ODCAllGather() + return __odc_gather + + +original_foreach_all_gather = _fsdp_collectives.foreach_all_gather + + +@torch.no_grad() +def foreach_all_gather( + fsdp_params: list[FSDPParam], + group: dist.ProcessGroup, + async_op: bool, + all_gather_copy_in_stream: torch.Stream, + all_gather_stream: torch.Stream, + device: torch.device, + all_gather_comm: AllGather, +) -> Optional[AllGatherResult]: + is_hpz_list = [ + param.post_forward_mesh_info is not None and param.mesh_info != param.post_forward_mesh_info + for param in fsdp_params + ] + assert len(set(is_hpz_list)) == 1, f"{is_hpz_list=}" + is_hpz = is_hpz_list[0] + + if is_hpz and fsdp_params[0].sharded_state == ShardedState.SHARDED: + assert ( + original_foreach_all_gather is not foreach_all_gather + ), "original_foreach_all_gather and foreach_all_gather are the same" + return original_foreach_all_gather( + fsdp_params, + group, + async_op, + all_gather_copy_in_stream, + all_gather_stream, + device, + all_gather_comm, + ) + + world_size, _rank = group.size(), group.rank() + device_handle = _get_device_handle(device.type) + # Override the all-gather comm with ODCAllGather + all_gather_comm = get_odc_gather_comm() + with device_handle.stream(all_gather_copy_in_stream): + key = get_fsdp_params_key(fsdp_params) + if fsdp_params[0].sharded_state == ShardedState.SHARDED_POST_FORWARD: + for fsdp_param in fsdp_params: + assert ( + fsdp_param.sharded_state == ShardedState.SHARDED_POST_FORWARD + ), f"{fsdp_param.sharded_state=}" + key = get_hpz_params_key(key) + assert SymmBufferRegistry.get_instance().has_key( + key + ), f"{key=} not found. The fsdp param group has not been replaced with symm buffer yet." + all_gather_input = SymmBufferRegistry.get_instance().get_symm_buffer(key) + all_gather_output = all_gather_comm.allocate( + (all_gather_input.numel() * world_size,), dtype=all_gather_input.dtype, device=device + ) + # _get_param_all_gather_inputs allocates some tensors + # But we don't use them here. Just get the metadata. + param_all_gather_inputs = _get_param_all_gather_inputs(fsdp_params) + ( + param_all_gather_input_dtypes, + param_all_gather_input_numels, + dtype, + ) = _get_all_gather_input_metadatas(param_all_gather_inputs) + if dtype == torch.uint8: + all_gather_inputs = [t.view(torch.uint8) for ts in param_all_gather_inputs for t in ts] + else: + all_gather_inputs = [*chain.from_iterable(param_all_gather_inputs)] + inp_split_sizes = [t.numel() for t in all_gather_inputs] + input_size_sum = sum(inp_split_sizes) + assert input_size_sum == all_gather_input.numel(), f"{input_size_sum=} != {all_gather_input.numel()}" + all_gather_stream.wait_stream(all_gather_copy_in_stream) + with device_handle.stream(all_gather_stream): + all_gather_work = all_gather_comm( + output_tensor=all_gather_output, + input_tensor=all_gather_input, + group=group, + async_op=async_op, + ) + all_gather_event = all_gather_stream.record_event() + return AllGatherResult( + all_gather_output, + all_gather_event, + all_gather_work, + param_all_gather_input_dtypes, + param_all_gather_input_numels, + inp_split_sizes, + ) + + +@torch.no_grad() +def foreach_all_gather_copy_out( + all_gather_result: AllGatherResult, + fsdp_params: list[FSDPParam], + group: dist.ProcessGroup, +) -> None: + ( + all_gather_output, + all_gather_event, + all_gather_work, + param_all_gather_input_dtypes, + param_all_gather_input_numels, + all_gather_input_split_sizes, + ) = all_gather_result + _dtype, device = all_gather_output.dtype, all_gather_output.device + device_handle = _get_device_handle(device.type) + + if all_gather_event is not None: # sync op + device_handle.current_stream().wait_event(all_gather_event) + if isinstance(all_gather_work, dist.distributed_c10d.Work): # async op + all_gather_work.wait() + world_size, device = group.size(), all_gather_output.device + + # In ODC, we use the original dtype for all-gather input and output. + # Then when the output needs to be used here, we cast it to the param_dtype. + param_dtype = fsdp_params[0].param_dtype + if param_dtype is not None: + all_gather_output = _to_dtype_if_needed(all_gather_output, param_dtype) + input_dtype = param_all_gather_input_dtypes[0][0] + for dtypes in param_all_gather_input_dtypes: + for dtype in dtypes: + if dtype != input_dtype: + raise ValueError(f"Input dtypes are not the same: {dtype} != {input_dtype}") + param_all_gather_input_dtypes = [ + [param_dtype] * len(dtypes) for dtypes in param_all_gather_input_dtypes + ] + + split_with_sizes_out: list[torch.Tensor] = [] + shard_i_copy_infos: list[tuple[FSDPParam, list[torch.Tensor]]] = [] + for all_gather_input_numels, all_gather_input_dtypes, fsdp_param in zip( + param_all_gather_input_numels, param_all_gather_input_dtypes, fsdp_params + ): + # NOTE: Under compile, make sure we always recreate all_gather_outputs + # per AllGather. See [Note: Invariants for torch.compile Traceable FSDP2]. + force_recreate = compiled_autograd_enabled() + fsdp_param.init_all_gather_outputs( + all_gather_input_numels, + all_gather_input_dtypes, + world_size, + device, + force_recreate=force_recreate, + ) + if not force_recreate: + fsdp_param.alloc_all_gather_outputs() + param_all_gather_outputs = fsdp_param.all_gather_outputs + if fsdp_param.fsdp_placement.dim != 0: + # Copy to a temporary and then chunk-cat into the final all-gather + # output tensors + param_all_gather_outputs = [torch.empty_like(t) for t in param_all_gather_outputs] + shard_i_copy_infos.append((fsdp_param, param_all_gather_outputs)) + split_with_sizes_out.extend(param_all_gather_outputs) + + all_gather_output = all_gather_output.view(world_size, -1) + if all_gather_output.dtype == torch.uint8: + out = [t.view(world_size, -1).view(torch.uint8) for t in split_with_sizes_out] + else: + out = [t.view(world_size, -1) for t in split_with_sizes_out] + + # only avoid VC bump if we are not in inference mode + if torch._dynamo.is_compiling(): + # For torch.compile, we turn off inference_mode for fake tensor + # propagation, and therefore graph break on is_inference. For `compile`, + # we don't care about VCs, so just skip the optimization. + non_inference_outs = [] + else: + non_inference_outs = [o for o in out if not o.is_inference()] + + if len(non_inference_outs) > 0: + with torch.autograd._unsafe_preserve_version_counter(tuple(non_inference_outs)): + torch.ops.fsdp.split_with_sizes_copy( + all_gather_output, all_gather_input_split_sizes, dim=1, out=out + ) + else: + torch.ops.fsdp.split_with_sizes_copy(all_gather_output, all_gather_input_split_sizes, dim=1, out=out) + + for fsdp_param, param_all_gather_outputs in shard_i_copy_infos: + # Chunk-cat from the temporary to the final all-gather output tensors + shard_dim = fsdp_param.fsdp_placement.dim + + with torch.autograd._unsafe_preserve_version_counter(tuple(fsdp_param.all_gather_outputs)): + for param_all_gather_output, target_all_gather_output in zip( + param_all_gather_outputs, fsdp_param.all_gather_outputs + ): + padded_sharded_size = ( + fsdp_param.padded_sharded_param_size + if fsdp_param.sharded_state == ShardedState.SHARDED + else cast(torch.Tensor, fsdp_param._sharded_post_forward_param_data).size() + ) + pre_param_size = list(padded_sharded_size) + pre_param_size[0] *= world_size + chunks = torch.chunk(param_all_gather_output.view(pre_param_size), world_size, dim=0) + post_param_size = list(padded_sharded_size) + post_param_size[shard_dim] *= world_size + cat_out = target_all_gather_output.view(post_param_size) + torch.cat(chunks, dim=shard_dim, out=cat_out) + + +def pre_minibatch_start(fsdp_module): + get_reduction_service().clear_accumulations() + + # Make sure optimizer updates are visible to all ranks + dist.barrier() + + ensure_resharded_within_node(fsdp_module) + + +def is_bw() -> bool: + return torch._C._current_graph_task_id() != -1 + + +# Old version of pytorch does not skip resharding after the recomputation in backward, +# resulting in duplicated all-gather in backward. +# Patch the new pytorch version here. +def post_forward(self, _module: nn.Module, _input: Any, output: Any): + if not compiled_autograd_enabled(): + logger.debug("%s", self._with_fqn("FSDP::post_forward")) + with record_function(self._with_fqn("FSDP::post_forward")): + if not compiled_autograd_enabled(): + # for AC(fully_shard(model)), AC runs fsdp's _pre_forward + # it shouldn't change post_forward_order + if not is_bw(): + self.reshard() + self._record_post_forward() + else: + self.reshard() + self._record_post_forward() + self._training_state = TrainingState.IDLE + return output + + +def post_backward(self, *_unused: Any): + # This method should be idempotent and safe to call even when this + # FSDP parameter group was not used in backward (should be a no-op) + if not compiled_autograd_enabled(): + logger.debug("%s", self._with_fqn("FSDP::post_backward")) + self._training_state = TrainingState.POST_BACKWARD + with record_function(self._with_fqn("FSDP::post_backward_accumulate")): + for fsdp_param in self.fsdp_params: + fsdp_param.accumulate_unsharded_grad_if_needed() + with record_function(self._with_fqn("FSDP::post_backward_reshard")): + if not self.reduce_grads: + if self.reshard_after_backward: + self.reshard() + for fsdp_param in self.fsdp_params: + fsdp_param.to_accumulated_grad_if_needed() + return + # Save the autograd-computed gradients before resharding to only + # access the unsharded parameters when their data is present + fsdp_params_with_grad: list[FSDPParam] = [] + unsharded_grads: list[torch.Tensor] = [] + for fsdp_param in self.fsdp_params: + if not hasattr(fsdp_param, "_unsharded_param"): + continue + # May have an accumulated gradient of the reduce dtype if the + # previous backward did not reduce-scatter + if fsdp_param.unsharded_accumulated_grad is not None: + fsdp_params_with_grad.append(fsdp_param) + unsharded_grads.append(fsdp_param.unsharded_accumulated_grad_data) + fsdp_param.unsharded_accumulated_grad = None + elif fsdp_param.unsharded_param.grad is not None: + fsdp_params_with_grad.append(fsdp_param) + unsharded_grads.append(fsdp_param.unsharded_grad_data) + fsdp_param.unsharded_param.grad = None + if self.reshard_after_backward: + self.reshard() + if len(fsdp_params_with_grad) == 0: + return + with record_function(self._with_fqn("FSDP::post_backward_reduce")): + if ( + self.comm_ctx.reduce_scatter_state is not None + and self.comm_ctx.reduce_scatter_state.event is not None + ): + self.device_handle.current_stream().wait_event(self.comm_ctx.reduce_scatter_state.event) + self.comm_ctx.reduce_scatter_state = None + all_reduce_pg = self._all_reduce_process_group if self._is_hsdp else None + all_reduce_stream: torch.cuda.Stream + if all_reduce_pg is None and self._all_reduce_hook_stream is not None: + # this means the native HSDP is not enabled, + # but user may want to have a custom HSDP setup + assert ( + self._all_reduce_hook is not None + ), "all reduce hook stream is specified but hook itself is missing." + all_reduce_stream = self._all_reduce_hook_stream + else: + all_reduce_stream = self.comm_ctx.all_reduce_stream + + self._wait_for_post_backward() + ( + reduce_scatter_input, + reduce_scatter_event, + self._post_reduce_event, + all_reduce_input, + all_reduce_event, + self._partial_reduce_output, + ) = foreach_reduce( + self, + fsdp_params_with_grad, + unsharded_grads, + self._reduce_scatter_process_group, + self.comm_ctx.reduce_scatter_stream, + self._reduce_scatter_comm, + self._orig_dtype, + self._reduce_dtype, + self.device, + self.gradient_divide_factor, + self._all_reduce_process_group if self._is_hsdp else None, + all_reduce_stream, + self.all_reduce_grads, + self._partial_reduce_output, + self._all_reduce_hook, + self.force_sum_reduction_for_comms, + ) + self.comm_ctx.reduce_scatter_state = ReduceScatterState(reduce_scatter_input, reduce_scatter_event) + if all_reduce_input is not None: + if self.device.type != "cpu": + assert all_reduce_event is not None + self._all_reduce_state = AllReduceState(all_reduce_input, all_reduce_event) + + +def reshard(self, refresh_post_forward_data: bool = False): + """ + This will only be patched in HPZ mode. + Keep params sharded on the post-forward mesh (within-node) outside forward + when reshard_after_forward is int (HPZ mode). + """ + # Supports gather parameters from local GPUs at the same node + # even between forward for different microbatches. + # So even in backward, we still does not shard it back to fully-sharded. + # We just shard it within each node. + # After all the backward is done, we will shard it back to fully-sharded. + # Only reshard to post-forward if we currently have unsharded params. + if self._training_state in (TrainingState.FORWARD, TrainingState.POST_BACKWARD): + if not self._reshard_after_forward: + return + if self._use_post_forward_mesh: + self._to_sharded_post_forward(refresh_post_forward_data=refresh_post_forward_data) + self._reshard_after_forward_event = self.device_handle.Event() + if self._reshard_after_forward_event is not None: + self._reshard_after_forward_event.record() + return + self._to_sharded() + + +def _to_sharded_post_forward(self, refresh_post_forward_data: bool = False): + """This patch is used to supports refresh_post_forward_data argument""" + if not self.is_sharded_post_forward: + for fsdp_param in self.fsdp_params: + fsdp_param.to_sharded_post_forward(refresh_post_forward_data=refresh_post_forward_data) + self._sharded_state = ShardedState.SHARDED_POST_FORWARD + + +@dataclasses.dataclass +class ReduceScatterContext: + # arguments + fsdp_params: list[FSDPParam] + unsharded_grads: list[torch.Tensor] + reduce_scatter_group: dist.ProcessGroup + reduce_scatter_stream: torch.Stream + orig_dtype: Optional[torch.dtype] + reduce_dtype: Optional[torch.dtype] + device: torch.device + gradient_divide_factor: Optional[float] + all_reduce_group: Optional[dist.ProcessGroup] # not `None` iff HSDP + all_reduce_stream: torch.Stream + all_reduce_grads: bool + partial_reduce_output: Optional[torch.Tensor] # only used for HSDP + all_reduce_hook: Optional[Callable[[torch.Tensor], None]] + force_sum_reduction_for_comms: bool + # others + padded_unsharded_sizes: tuple + grad_dtype: torch.dtype + + +@torch.no_grad() +def foreach_reduce( + fsdp_param_group: FSDPParamGroup, + fsdp_params: list[FSDPParam], + unsharded_grads: list[torch.Tensor], + reduce_scatter_group: dist.ProcessGroup, + reduce_scatter_stream: torch.Stream, + reduce_scatter_comm: ReduceScatter, + orig_dtype: Optional[torch.dtype], + reduce_dtype: Optional[torch.dtype], + device: torch.device, + gradient_divide_factor: Optional[float], + all_reduce_group: Optional[dist.ProcessGroup], # not `None` iff HSDP + all_reduce_stream: torch.Stream, + all_reduce_grads: bool, + partial_reduce_output: Optional[torch.Tensor], # only used for HSDP + all_reduce_hook: Optional[Callable[[torch.Tensor], None]], + force_sum_reduction_for_comms: bool = False, +) -> tuple[ + torch.Tensor, + torch.Event, + torch.Event, + Optional[torch.Tensor], + Optional[torch.Event], + Optional[torch.Tensor], +]: + """ + ``unsharded_grads`` owns the references to the gradients computed by + autograd, so clearing the list frees the gradients. + """ + + grad_dtypes = {grad.dtype for grad in unsharded_grads} + if len(grad_dtypes) != 1: + # Check this at runtime since it could be a real runtime error if e.g. + # fp8 weights do not produce the correct higher precision gradients + _raise_assert_with_print(f"FSDP reduce-scatter expects uniform gradient dtype but got {grad_dtypes}") + grad_dtype = unsharded_grads[0].dtype + reduce_dtype = reduce_dtype or grad_dtype + (predivide_factor, _postdivide_factor, _reduce_scatter_op, _all_reduce_op) = _get_gradient_divide_factors( + reduce_scatter_group, + all_reduce_group, + reduce_dtype, + device.type, + gradient_divide_factor, + force_sum_reduction_for_comms, + ) + world_size = reduce_scatter_group.size() + device_handle = _get_device_handle(device.type) + current_stream = device_handle.current_stream() + + if world_size > 1: + for i, (fsdp_param, unsharded_grad) in enumerate(zip(fsdp_params, unsharded_grads)): + if (shard_dim := fsdp_param.fsdp_placement.dim) == 0: + continue + assert ( + unsharded_grad.size(shard_dim) % world_size == 0 + ), f"Shard({shard_dim}) requires even sharding: {unsharded_grad.size()=} {world_size=}" + chunks = torch.chunk(unsharded_grad, world_size, dim=shard_dim) + unsharded_grads[i] = torch.cat(chunks, dim=0) + + padded_unsharded_sizes = tuple(_get_dim0_padded_size(grad.size(), world_size) for grad in unsharded_grads) + reduce_scatter_input_numel = sum(s.numel() for s in padded_unsharded_sizes) + reduce_scatter_input = reduce_scatter_comm.allocate( + (reduce_scatter_input_numel,), + dtype=reduce_dtype, + device=device, + ) + + foreach_reduce_scatter_copy_in(unsharded_grads, reduce_scatter_input, world_size) + + # Only after the copy-in finishes can we free the gradients + unsharded_grads.clear() + reduce_scatter_stream.wait_stream(current_stream) + all_reduce_input = None + all_reduce_event = None + + with device_handle.stream(reduce_scatter_stream): + _div_if_needed(reduce_scatter_input, predivide_factor) + key = id(fsdp_param_group) + scatter = get_reduction_service() + scatter.scatter_accumulate(key, reduce_scatter_input, reduce_scatter_group) + + post_reduce_stream = reduce_scatter_stream + reduce_scatter_event = reduce_scatter_stream.record_event() + + # Save the context for the next update_gradients call + fsdp_param_group.__odc_reduce_scatter_context = ReduceScatterContext( + fsdp_params=fsdp_params, + unsharded_grads=unsharded_grads, + reduce_scatter_group=reduce_scatter_group, + reduce_scatter_stream=reduce_scatter_stream, + orig_dtype=orig_dtype, + reduce_dtype=reduce_dtype, + device=device, + gradient_divide_factor=gradient_divide_factor, + all_reduce_group=all_reduce_group, + all_reduce_stream=all_reduce_stream, + all_reduce_grads=all_reduce_grads, + partial_reduce_output=partial_reduce_output, + all_reduce_hook=all_reduce_hook, + force_sum_reduction_for_comms=force_sum_reduction_for_comms, + padded_unsharded_sizes=padded_unsharded_sizes, + grad_dtype=grad_dtype, + ) + + return ( + reduce_scatter_input, + reduce_scatter_event, + post_reduce_stream.record_event(), + all_reduce_input, + all_reduce_event, + partial_reduce_output, + ) + + +def ensure_resharded_within_node(fsdp_module): + """ + Ensure all parameters are sharded within each node at the beginning of epoch. + This is needed for reshard_after_forward=int mode with ODC to ensure all GPUs + have finished sharding before backward gather operations start. + """ + root_state = fully_shard.state(fsdp_module) + root_state._lazy_init() + all_fsdp_states = root_state._state_ctx.all_states + all_fsdp_param_groups = [ + state._fsdp_param_group for state in all_fsdp_states if state._fsdp_param_group is not None + ] + + hpz = any(fsdp_param_group._use_post_forward_mesh for fsdp_param_group in all_fsdp_param_groups) + if not hpz: + return + + per_node_pg = None + for fsdp_param_group in all_fsdp_param_groups: + # Only do this if reshard_after_forward is an int (HPZ mode) + if not fsdp_param_group._use_post_forward_mesh: + continue + if per_node_pg is None: + assert isinstance( + fsdp_param_group.post_forward_mesh_info, HSDPMeshInfo + ), f"{fsdp_param_group.post_forward_mesh_info=}" + per_node_pg = fsdp_param_group.post_forward_mesh_info.shard_process_group + + # Set training state to FORWARD so reshard() will do post-forward resharding + old_state = fsdp_param_group._training_state + fsdp_param_group._training_state = TrainingState.FORWARD + + # Unshard (all-gather on all GPUs) - what pre_forward does + with torch.cuda.nvtx.range("unshard"): + fsdp_param_group.unshard(async_op=False) + with torch.cuda.nvtx.range("wait_for_unshard"): + fsdp_param_group.wait_for_unshard() + + # Reshard (shard within each node) - what post_forward does + with torch.cuda.nvtx.range("reshard"): + fsdp_param_group.reshard(refresh_post_forward_data=True) + + # Restore training state + fsdp_param_group._training_state = old_state + + assert per_node_pg is not None, "HPZ enabled but per-node process group not found" + torch.distributed.barrier(group=per_node_pg) + + +@torch.no_grad() +def update_gradients(fsdp_param_group: FSDPParamGroup): + if not hasattr(fsdp_param_group, "__odc_reduce_scatter_context"): + # This is to support that in some iteration, there is no microbatch, + # so no reduce-scatter is needed. + # __odc_reduce_scatter_context does not exists here in this case. + return + reduce_scatter_context = fsdp_param_group.__odc_reduce_scatter_context + del fsdp_param_group.__odc_reduce_scatter_context + fsdp_params = reduce_scatter_context.fsdp_params + reduce_scatter_group = reduce_scatter_context.reduce_scatter_group + reduce_scatter_stream = reduce_scatter_context.reduce_scatter_stream + orig_dtype = reduce_scatter_context.orig_dtype + reduce_dtype = reduce_scatter_context.reduce_dtype + device = reduce_scatter_context.device + gradient_divide_factor = reduce_scatter_context.gradient_divide_factor + all_reduce_group = reduce_scatter_context.all_reduce_group + all_reduce_stream = reduce_scatter_context.all_reduce_stream + all_reduce_hook = reduce_scatter_context.all_reduce_hook + force_sum_reduction_for_comms = reduce_scatter_context.force_sum_reduction_for_comms + padded_unsharded_sizes = reduce_scatter_context.padded_unsharded_sizes + grad_dtype = reduce_scatter_context.grad_dtype + + device_handle = _get_device_handle(device.type) + + reduce_dtype = reduce_dtype or grad_dtype + (_predivide_factor, postdivide_factor, reduce_scatter_op, all_reduce_op) = _get_gradient_divide_factors( + reduce_scatter_group, + all_reduce_group, + reduce_dtype, + device.type, + gradient_divide_factor, + force_sum_reduction_for_comms, + ) + world_size = reduce_scatter_group.size() + device_handle = _get_device_handle(device.type) + current_stream = device_handle.current_stream() + + with device_handle.stream(reduce_scatter_stream): + post_reduce_stream = all_reduce_stream + + scatter = get_reduction_service() + key = id(fsdp_param_group) + reduce_output = scatter.get_accumulation(key) + assert reduce_scatter_op in [ + torch.distributed.ReduceOp.SUM, + torch.distributed.ReduceOp.AVG, + ], f"reduce_scatter_op {reduce_scatter_op} is not supported" + if reduce_scatter_op == torch.distributed.ReduceOp.AVG: + reduce_output /= world_size + + if all_reduce_group is not None: # HSDP + # ODC defers the inter-node all-reduce during HSDP gradient + # accumulation (all_reduce_grads=False) inside scatter-accumulate, so + # the original implementation's partial-reduce bookkeeping is not + # needed here. + post_reduce_stream = all_reduce_stream + if world_size >= 1: + all_reduce_stream.wait_stream(reduce_scatter_stream) + else: + all_reduce_stream.wait_stream(current_stream) + with device_handle.stream(all_reduce_stream): + dist.all_reduce( + reduce_output, + group=all_reduce_group, + op=all_reduce_op, + ) + # -- END: ops in reduce_scatter stream + + if all_reduce_hook is not None: + # Execute user-specified all reduce hook. + # If native HSDP is used, this is executed after the HSDP all reduce. + # If 1-d FSDP is used, this is executed post reduce-scatter. + post_reduce_stream = all_reduce_stream + all_reduce_stream.wait_stream(reduce_scatter_stream) + with device_handle.stream(all_reduce_stream): + all_reduce_hook(reduce_output) + # -- END: ops post reduce_scatter + + with device_handle.stream(post_reduce_stream): + _div_if_needed(reduce_output, postdivide_factor) + reduce_output = _to_dtype_if_needed(reduce_output, orig_dtype) + # View out and accumulate sharded gradients + flat_grad_offset = 0 + for padded_unsharded_size, fsdp_param in zip(padded_unsharded_sizes, fsdp_params): + # Assume even sharding for Shard(i), i > 0; otherwise would require + # copy-out for contiguous strides + new_sharded_grad = torch.as_strided( + reduce_output, + size=fsdp_param.sharded_size, + stride=fsdp_param.contiguous_sharded_stride, + storage_offset=flat_grad_offset, + ) + to_accumulate_grad = fsdp_param.sharded_param.grad is not None + if fsdp_param.offload_to_cpu: + # Only overlap the D2H copy (copying to pinned memory) if not + # accumulating gradients since the CPU add kernel depends on + # the copy result and we cannot run the add as a callback + non_blocking = fsdp_param.pin_memory and not to_accumulate_grad + # Since the GPU sharded gradient is allocated in the RS stream, + # we can free it here by not keeping a ref without waiting for + # the D2H copy since future RS-stream ops run after the copy + new_sharded_grad = new_sharded_grad.to(torch.device("cpu"), non_blocking=non_blocking) + if non_blocking: + # Record an event on which to block the CPU thread to + # ensure that the D2H copy finishes before the optimizer + fsdp_param.grad_offload_event = post_reduce_stream.record_event() + if to_accumulate_grad: + assert isinstance(fsdp_param.sharded_param.grad, DTensor) + fsdp_param.sharded_param.grad._local_tensor += new_sharded_grad + else: + new_sharded_dtensor_grad = fsdp_param.to_sharded_dtensor(new_sharded_grad) + fsdp_param.sharded_param.grad = new_sharded_dtensor_grad + if not compiled_autograd_enabled(): + for hook in ( + getattr(fsdp_param.sharded_param, "_post_accumulate_grad_hooks", {}) or {} + ).values(): + hook(fsdp_param.sharded_param) + padded_sharded_numel = padded_unsharded_size.numel() // world_size + flat_grad_offset += padded_sharded_numel + post_reduce_event = post_reduce_stream.record_event() + # The RS output is allocated in the RS stream and used in the default + # stream (for optimizer). To ensure its memory is not reused for later + # RSs, we do not need extra synchronization since the sharded parameters + # hold refs through the end of backward. + + # Synchronize the default stream with post_reduce_stream to ensure + # gradient writes are visible before optimizer step + current_stream = device_handle.current_stream() + current_stream.wait_event(post_reduce_event) + + +# FSDPParam +def to_sharded_post_forward(self, refresh_post_forward_data: bool = False) -> None: + if self.is_dtensor: + raise NotImplementedError("Resharding to smaller mesh with TP is not supported yet") + self._assert_in_states(ShardedState.UNSHARDED) + assert self.post_forward_mesh_info is not None # mypy + assert len(self.all_gather_outputs) == 1 + shard_world_size = self.post_forward_mesh_info.shard_mesh_size + if (numel := self.all_gather_outputs[0].numel()) % shard_world_size != 0: + _raise_assert_with_print( + f"All-gather output size ({numel}) must be divisible by the shard " + f"world size ({shard_world_size})" + ) + shard_rank = self.post_forward_mesh_info.shard_mesh_rank + # pyrefly: ignore # unbound-name + sharded_numel = numel // shard_world_size + # Don't replace the symmetric buffer _sharded_post_forward_param_data here. + + # If hpz is enabled, self._sharded_post_forward_param_data + # only needs to be updated (copy) + # on the first unshard in ensure_resharded_within_node. + # Later unshard in forward and backward doing gather + # from the _sharded_post_forward_param_data won't change + # _sharded_post_forward_param_data itself. + if refresh_post_forward_data: + self._sharded_post_forward_param_data.copy_( + self.all_gather_outputs[0].narrow(0, sharded_numel * shard_rank, sharded_numel) + ) + sharded_post_forward_tensor = torch.as_strided( + self._sharded_post_forward_param_data, + size=self.sharded_post_forward_size, + stride=self.contiguous_sharded_post_forward_stride, + storage_offset=0, + ) + self._sharded_post_forward_param = nn.Parameter( + self.to_sharded_post_forward_dtensor(sharded_post_forward_tensor) + ) + self._setattr_on_modules(self._sharded_post_forward_param) + self.free_unsharded_param() + self.sharded_state = ShardedState.SHARDED_POST_FORWARD + + +# FSDPParam +def to_unsharded(self) -> None: + # Assume that the data has been allocated and all-gathered + set_requires_grad_if_needed(self.sharded_param, self._unsharded_param) + self._setattr_on_modules(self._unsharded_param) + if self.sharded_state == ShardedState.SHARDED_POST_FORWARD: + # The data is allocated in the default stream via the post-forward + # reshard and must be kept alive for the next all-gather copy-in. + # Since we call this method after the copy-out, the data's lifetime + # is ensured without further synchronization. + self._sharded_post_forward_param = None + # Do not free the symmetric buffer for HPZ. + self.sharded_state = ShardedState.UNSHARDED + + +def pre_optimizer_step(fsdp_module): + scatter = get_reduction_service() + + root_state = fully_shard.state(fsdp_module) + all_fsdp_states = root_state._state_ctx.all_states + all_fsdp_param_groups = [ + state._fsdp_param_group for state in all_fsdp_states if state._fsdp_param_group is not None + ] + for i, fsdp_param_group in enumerate(all_fsdp_param_groups): + if i == 0: + # Scatter-accumulate uses the global shard group even in HPZ since + # gradients are sharded across the full world size. + mesh_info = fsdp_param_group.mesh_info + assert isinstance(mesh_info, FSDPMeshInfo) + reduce_scatter_group = mesh_info.shard_process_group + with torch.cuda.nvtx.range("scatter_accumulate_sync"): + scatter.sync(reduce_scatter_group) + + with torch.cuda.nvtx.range(f"update_gradients:{fsdp_param_group._module_fqn}"): + update_gradients(fsdp_param_group) + if fsdp_param_group._use_post_forward_mesh and fsdp_param_group.is_sharded_post_forward: + # After gradient sync, return to fully-sharded params so the next + # minibatch performs cross-node all-gather again. + fsdp_param_group._to_sharded() + + +def _get_post_forward_mesh_info_no_convert(reshard_after_forward, mesh_info): + """Variant of FSDP's helper that preserves int semantics even on 1 node. + This is mainly for development purpose. + Running HPZ in 1 node increase memory without any benefit. + Actually we don't need to use HPZ in 1 node. + """ + from torch._logging import warning_once + from torch.distributed.tensor import DeviceMesh + + shard_mesh_size = mesh_info.shard_mesh_size + if not isinstance(reshard_after_forward, (bool, int)): + raise ValueError( + "reshard_after_forward should be a bool or an int representing the " + f"group size to reshard to, not {reshard_after_forward}" + ) + # NOTE: `isinstance(False, int)` returns `True`. + if not isinstance(reshard_after_forward, bool) and isinstance(reshard_after_forward, int): + if ( + reshard_after_forward < 1 + or reshard_after_forward > shard_mesh_size + or shard_mesh_size % reshard_after_forward != 0 + ): + raise ValueError( + "If passing reshard_after_forward as an int, it should be a " + f"factor of {shard_mesh_size}, not {reshard_after_forward}" + ) + if reshard_after_forward == 1: + msg = ( + "reshard_after_forward=1 (int) means resharding parameters to world size 1, " + "instead of reshard_after_forward=True (bool)" + ) + warning_once(logger, msg, stacklevel=2) + reshard_after_forward = False + # In the original pytorch implementation, + # if reshard_after_forward == shard_mesh_size, + # it is actually equivalent to True but use more memory. + # So it sets it to True. + # For us, for easier development with just 1 node, + # we disable this behavior. + post_forward_mesh_info = None + if reshard_after_forward is True: + post_forward_mesh_info = mesh_info + elif reshard_after_forward is not False: # int case + post_forward_mesh_tensor = mesh_info.mesh.mesh.view(-1, reshard_after_forward) + post_forward_mesh = DeviceMesh(mesh_info.mesh.device_type, post_forward_mesh_tensor) + post_forward_mesh_info = HSDPMeshInfo(post_forward_mesh, shard_mesh_dim=1, replicate_mesh_dim=0) + return post_forward_mesh_info + + +_enable_hpz = None + + +def patch_fsdp2(enable_hpz: bool = False) -> None: + from torch.distributed._composable import replicate_with_fsdp + from torch.distributed.fsdp._fully_shard import ( + _fsdp_init, + _fsdp_param_group, + _fully_shard, + ) + + global _enable_hpz + if _enable_hpz is None: + _enable_hpz = enable_hpz + else: + assert ( + _enable_hpz == enable_hpz + ), f"HPZ mode is already set to {_enable_hpz}, cannot change to {enable_hpz}" + + _fsdp_collectives.foreach_all_gather = foreach_all_gather + _fsdp_param_group.foreach_all_gather = foreach_all_gather + _fsdp_collectives.foreach_all_gather_copy_out = foreach_all_gather_copy_out + _fsdp_param_group.foreach_all_gather_copy_out = foreach_all_gather_copy_out + FSDPParamGroup.post_backward = post_backward + FSDPParamGroup.post_forward = post_forward + if enable_hpz: + _fsdp_init._get_post_forward_mesh_info = _get_post_forward_mesh_info_no_convert + _fully_shard._get_post_forward_mesh_info = _get_post_forward_mesh_info_no_convert + replicate_with_fsdp._get_post_forward_mesh_info = _get_post_forward_mesh_info_no_convert + FSDPParamGroup.reshard = reshard + FSDPParamGroup._to_sharded_post_forward = _to_sharded_post_forward + FSDPParam.to_sharded_post_forward = to_sharded_post_forward + FSDPParam.to_unsharded = to_unsharded + _fsdp_param_group.record_function = nvtx_record_function + torch.profiler.record_function = nvtx_record_function + torch.autograd.profiler.record_function = nvtx_record_function + + global reduction_service + global gather_service + reduction_service = ReductionService() + gather_service = GatherService() + + +def stop(): + get_reduction_service().stop() + SymmBufferRegistry.get_instance().finalize() + finalize_distributed() + + +def get_symm_buffer_memory_breakdown() -> dict[str, float]: + """Return ODC symmetric-buffer memory breakdown in GB.""" + breakdown = { + "fsdp_params": 0, + "accumulation": 0, + "shared_buffer": 0, + "ag_buffer": 0, + "hpz_buffer": 0, + "other": 0, + } + registry = SymmBufferRegistry.get_instance() + for key, tensor in registry.local_tensor.items(): + nbytes = tensor.nbytes + if key.startswith("fsdp_params_gather_"): + breakdown["fsdp_params"] += nbytes + elif key.startswith("rs_accumulation_"): + breakdown["accumulation"] += nbytes + elif key.startswith("shared_buffer_"): + breakdown["shared_buffer"] += nbytes + elif key.startswith("ag_buffer_"): + breakdown["ag_buffer"] += nbytes + elif key.startswith("hpz_"): + breakdown["hpz_buffer"] += nbytes + else: + breakdown["other"] += nbytes + return {k: round(v / (1024**3), 2) for k, v in breakdown.items()} diff --git a/primus/core/odc/odc_early/sitecustomize.py b/primus/core/odc/odc_early/sitecustomize.py new file mode 100644 index 000000000..dd74b1448 --- /dev/null +++ b/primus/core/odc/odc_early/sitecustomize.py @@ -0,0 +1,40 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""ODC early-init shim (auto-executed by Python's site machinery). + +Root cause (verified on MI300X + ROCm 7.2 + PyTorch 2.10a): + Importing `transformer_engine` BEFORE MORI's C++ runtime is loaded leaves + the process in a state where MORI init (or even just allocating its + symmetric heap) aborts with `free(): invalid pointer`. It is a C++ + dynamic-library load-order / global-ctor conflict, NOT an init-order or + init-method issue. + + The fix is simply to load MORI's C++ runtime (via `import mori`) BEFORE + Megatron imports transformer_engine. This file runs at interpreter + startup (site.py imports `sitecustomize` if it is on sys.path), which is + earlier than any Primus/Megatron/TE import, guaranteeing the correct + order. + +Enable by putting this directory on PYTHONPATH (the ODC launcher does this): + export PYTHONPATH=/primus/core/odc/odc_early:$PYTHONPATH + +This shim is a pure bootstrap (a launch-time load-order workaround, NOT feature +logic): it only runs when its directory is on PYTHONPATH, which the ODC launcher +adds exactly for ODC runs, so it is a complete no-op for normal runs. Whether ODC +is actually active for training is decided by the enable_odc config item, not by +this file. +""" + +import sys + +try: + import mori # noqa: F401 -- loads libmori / MORI C++ runtime + import mori.shmem # noqa: F401 + + sys.stderr.write("[ODC sitecustomize] pre-imported mori before TE (load-order fix)\n") + sys.stderr.flush() +except Exception as _e: # noqa: BLE001 + sys.stderr.write(f"[ODC sitecustomize] WARNING: pre-import mori failed: {_e}\n") + sys.stderr.flush() diff --git a/primus/core/odc/primitives/__init__.py b/primus/core/odc/primitives/__init__.py new file mode 100644 index 000000000..ef3f8fb6f --- /dev/null +++ b/primus/core/odc/primitives/__init__.py @@ -0,0 +1,43 @@ +# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under +# the MIT License per its package metadata (pyproject.toml / setup.py +# classifiers). The upstream repository ships no LICENSE file or per-file +# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab). +# +# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# See LICENSE for license information. + +from .shmem_triton import ( + LIB_SHMEM_PATH, + SHMEM_EXTERN_LIBS, + __syncthreads, + getmem_nbi_block, + int_atomic_compare_swap, + int_atomic_swap, + int_g, + int_p, + int_p_remote, + int_wait_until_equals, + int_wait_until_equals_remote, + putmem_nbi_block, + quiet, + tid, +) + +__all__ = [ + # shmem_triton + "int_atomic_compare_swap", + "int_atomic_swap", + "putmem_nbi_block", + "getmem_nbi_block", + "quiet", + "int_p", + "int_p_remote", + "int_g", + "int_wait_until_equals", + "int_wait_until_equals_remote", + "tid", + "__syncthreads", + "LIB_SHMEM_PATH", + "SHMEM_EXTERN_LIBS", +] diff --git a/primus/core/odc/primitives/_rocshmem_backend.py b/primus/core/odc/primitives/_rocshmem_backend.py new file mode 100644 index 000000000..682ba2d28 --- /dev/null +++ b/primus/core/odc/primitives/_rocshmem_backend.py @@ -0,0 +1,588 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""ODC rocSHMEM P2P backend, consuming the rocSHMEM ops from Primus-Turbo. + +Activated by ``ODC_P2P_BACKEND=rocshmem`` (the default is ``mori``). The +rocSHMEM host/GDA surface is no longer loaded from an in-tree ``librs_host*.so`` +via ctypes; it is now provided by Primus-Turbo as two pybind submodules of +``primus_turbo.pytorch._C``: + + * ``odc_rocshmem_host`` — single-node host-API surface (XGMI IPC path): + ``rs_get_uid`` / ``rs_init_uid`` / ``rs_malloc`` / ``rs_ptr`` / + ``rs_barrier`` / ``rs_finalize``. Selected when odc_rocshmem_gda is false. + * ``odc_rocshmem_gda`` — multi-node GPU-direct (GDA) surface: the same + host-compatible ``rs_*`` bootstrap plus the device ``gda_gather`` / + ``gda_reduce_scatter_acc`` launchers. Selected when odc_rocshmem_gda is true. + +Single-node (host) path: this backend uses ONLY the host API to manage the +symmetric heap and resolve peer pointers. It deliberately does NOT link any +rocSHMEM *device* bitcode into Triton — every on-device P2P op is plain Triton +``tl.load`` / ``tl.store`` / ``tl.atomic_*`` to XGMI-mapped peer addresses (see +the ``rocshmem`` branch in ``shmem_triton.py``). The heavy gather / scatter data +movement is host-side torch ``.copy_()`` on peer-view tensors (XGMI peer +load/store); the only device primitives exercised are ``int_p`` / +``int_wait_until_equals`` for the scatter-accumulate hand-shake. Those translate +a local symmetric address to a peer address with a per-PE affine delta — +``rocshmem_ptr`` was empirically verified to be affine (``rs_ptr(x, pe) - x`` is +constant across symmetric addresses ``x``), so a single ``delta[pe]`` resolves +every symmetric pointer. + +pybind surface consumed (see Primus-Turbo csrc/pytorch/dist/odc_rocshmem_*.cpp):: + + int rs_uid_bytes(); + bytes rs_get_uid(); # returns the uid as `bytes` + void rs_init_uid(int rank, int nranks, bytes uid); + int rs_my_pe(); + int rs_n_pes(); + long long rs_malloc(size_t n); # symmetric-heap device ptr + long long rs_ptr(long long p, int pe); # peer mapping of symmetric ptr p + void rs_barrier(); + void rs_finalize(); + # GDA submodule only, device launchers (peers passed as a Python list): + int gda_gather(target, src, nbytes, list peers, stride_bytes); + int gda_reduce_scatter_acc(...); # etc. +""" + +import ctypes +import logging +import os +from functools import reduce + +import torch +import torch.distributed as dist + +from odc.runtime_config import get_config + +logger = logging.getLogger(__name__) + +# c10::ScalarType enum codes (stable across torch versions). +_DTYPE_CODE = { + torch.uint8: 0, + torch.int8: 1, + torch.int16: 2, + torch.int32: 3, + torch.int64: 4, + torch.float16: 5, + torch.float32: 6, + torch.float64: 7, + torch.bool: 11, + torch.bfloat16: 15, +} + +# Max same-node PEs supported by the device delta table (MI300X node == 8 GPUs). +MAX_LOCAL_PES = 8 + +_FROM_BLOB_CPP = r""" +#include +// Wrap an externally-owned device pointer as a torch.Tensor (no deleter: the +// rocSHMEM symmetric heap owns the memory and frees it at rs_finalize()). +torch::Tensor odc_rs_from_blob(int64_t ptr, std::vector sizes, + int64_t dtype_code, int64_t device_index) { + auto dtype = static_cast(dtype_code); + auto options = torch::TensorOptions().dtype(dtype).device(torch::kCUDA, device_index); + return torch::from_blob(reinterpret_cast(ptr), sizes, options); +} +""" + +# --------------------------------------------------------------------------- +# Module state +# --------------------------------------------------------------------------- +_lib = None +_from_blob = None +_my_pe = -1 +_n_pes = -1 +_ref_base = None # reference symmetric address used for affine peer deltas +_allocations = [] # keep python tensor objects alive for the run +_initialized = False + +_gda_enabled = False # GPU-direct (GDA) device path: config odc_rocshmem_gda=true +_ctypes_gda_lib = False # True when odc_rocshmem_lib overrides embedded turbo GDA + + +def gda_enabled(): + """True when the GPU-direct (GDA) device-kernel cross-node path is active.""" + return _gda_enabled + + +def _is_gda(): + # config item odc_rocshmem_gda: GPU-direct (GDA) device path (multi-node). + return bool(get_config().rocshmem_gda) + + +def _load_ctypes_gda(so): + """Load proven in-tree librs_host_gda.so when odc_rocshmem_lib is set. + + The embedded turbo ``odc_rocshmem_gda`` path linked against rocshmem_combined + can fault on first device gather; this override uses the monolithic RDC .so + that matches the in-tree dual baseline. + """ + lib = ctypes.CDLL(so) + lib.rs_uid_bytes.restype = ctypes.c_int + lib.rs_get_uid.argtypes = [ctypes.c_char_p] + lib.rs_init_uid.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_char_p] + lib.rs_my_pe.restype = ctypes.c_int + lib.rs_n_pes.restype = ctypes.c_int + lib.rs_malloc.restype = ctypes.c_longlong + lib.rs_malloc.argtypes = [ctypes.c_size_t] + lib.rs_ptr.restype = ctypes.c_longlong + lib.rs_ptr.argtypes = [ctypes.c_longlong, ctypes.c_int] + lib.rs_barrier.restype = None + lib.rs_finalize.restype = None + if hasattr(lib, "gda_gather"): + lib.gda_gather.restype = ctypes.c_int + lib.gda_gather.argtypes = [ + ctypes.c_longlong, + ctypes.c_longlong, + ctypes.c_size_t, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_size_t, + ] + lib.gda_reduce_scatter_acc.restype = ctypes.c_int + lib.gda_reduce_scatter_acc.argtypes = [ + ctypes.c_longlong, + ctypes.c_longlong, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_int, + ctypes.c_longlong, + ctypes.c_size_t, + ctypes.c_int, + ctypes.c_int, + ] + if hasattr(lib, "gda_stage_fence"): + lib.gda_stage_fence.restype = ctypes.c_int + lib.gda_stage_fence.argtypes = [ctypes.c_longlong, ctypes.c_longlong, ctypes.c_size_t] + if hasattr(lib, "gda_hdp_flush"): + lib.gda_hdp_init.restype = ctypes.c_int + lib.gda_hdp_init.argtypes = [] + lib.gda_hdp_flush.restype = ctypes.c_int + lib.gda_hdp_flush.argtypes = [] + if hasattr(lib, "gda_strided_touch"): + lib.gda_strided_touch.restype = ctypes.c_int + lib.gda_strided_touch.argtypes = [ + ctypes.c_longlong, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_int, + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_longlong, + ctypes.c_size_t, + ctypes.c_int, + ] + if hasattr(lib, "gda_reduce_scatter_acc_async"): + lib.gda_reduce_scatter_acc_async.restype = ctypes.c_int + lib.gda_reduce_scatter_acc_async.argtypes = lib.gda_reduce_scatter_acc.argtypes + lib.gda_rs_overlap_sync.restype = ctypes.c_int + lib.gda_rs_overlap_sync.argtypes = [] + if hasattr(lib, "gda_gather_async"): + lib.gda_gather_async.restype = ctypes.c_int + lib.gda_gather_async.argtypes = [ + ctypes.c_longlong, + ctypes.c_longlong, + ctypes.c_size_t, + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_size_t, + ctypes.c_longlong, + ] + return lib + + +def _load_backend(): + """Resolve the ODC rocSHMEM backend submodule from Primus-Turbo. + + When the odc_rocshmem_lib config points at a GDA ``librs_host_gda.so``, load + that monolithic binding via ctypes instead of the embedded turbo GDA kernels. + """ + global _ctypes_gda_lib + # config item odc_rocshmem_lib: optional monolithic librs_host_gda.so override. + so = get_config().rocshmem_lib + if so and os.path.isfile(so) and _is_gda(): + _ctypes_gda_lib = True + logger.info("rocSHMEM GDA backend: ctypes override %s", so) + return _load_ctypes_gda(so) + _ctypes_gda_lib = False + try: + import primus_turbo # noqa: F401 -- package init (also loads the _C ext) + import primus_turbo.pytorch._C as _C + except ImportError as e: + raise ImportError( + "ODC rocSHMEM backend now consumes the rocSHMEM ops from Primus-Turbo " + "(primus_turbo.pytorch._C.odc_rocshmem_host / odc_rocshmem_gda). Install " + "Primus-Turbo built with the ODC rocSHMEM ops, or put its build tree on " + "PYTHONPATH so `import primus_turbo` resolves it." + ) from e + name = "odc_rocshmem_gda" if _is_gda() else "odc_rocshmem_host" + mod = getattr(_C, name, None) + if mod is None: + available = [n for n in dir(_C) if "odc" in n.lower()] + raise RuntimeError( + f"primus_turbo.pytorch._C has no submodule '{name}'. This primus_turbo " + f"build was compiled without the ODC rocSHMEM ops (DISABLE_ROCSHMEM). " + f"ODC submodules present: {available}" + ) + return mod + + +def _build_from_blob(): + from torch.utils.cpp_extension import load_inline + + mod = load_inline( + name="odc_rs_from_blob", + cpp_sources=[_FROM_BLOB_CPP], + functions=["odc_rs_from_blob"], + verbose=False, + ) + return mod.odc_rs_from_blob + + +def _wrap(ptr, shape, dtype, dev_index): + if isinstance(shape, int): + shape = (shape,) + return _from_blob(int(ptr), list(shape), _DTYPE_CODE[dtype], int(dev_index)) + + +def _nbytes(shape, dtype): + if isinstance(shape, int): + shape = (shape,) + numel = reduce(lambda a, b: a * b, shape, 1) + return numel * torch._utils._element_size(dtype) + + +def init(): + """Bootstrap rocSHMEM from PyTorch's WORLD process group via a unique-id + broadcast. + + BOTH the single-node host (``odc_rocshmem_host``) and the multi-node + GPU-direct (``odc_rocshmem_gda``) submodules use the SAME unique-id + bootstrap: rank 0 generates the uid (``rs_get_uid``, returned as ``bytes``), + it is broadcast over the torch WORLD process group, and every rank calls + ``rs_init_uid`` (which maps to ``rocshmem_init_attr(ROCSHMEM_INIT_WITH_UNIQUEID)``). + The GDA transport exchanges the uid over a TCP socket bootstrap + (``ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME``), so NO MPI job / mpirun is required — + the run can be launched with plain torchrun. + """ + global _lib, _from_blob, _my_pe, _n_pes, _initialized, _gda_enabled + if _initialized: + return + assert dist.is_initialized(), "torch.distributed must be initialized first" + _lib = _load_backend() + _from_blob = _build_from_blob() + + _gda_enabled = _is_gda() and hasattr(_lib, "gda_gather") + if _gda_enabled: + # GDA multi-node bootstrap carries rank-0's address inside the uid and + # exchanges it over a TCP socket. Default the bootstrap NIC to eth0 unless + # the deployment overrode it; this replaces the old MPI/mpirun bootstrap. + os.environ.setdefault("ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME", "eth0") + + rank = dist.get_rank() + world = dist.get_world_size() + if _ctypes_gda_lib: + n = _lib.rs_uid_bytes() + buf = (ctypes.c_char * n)() + uidb = None + if rank == 0: + _lib.rs_get_uid(buf) + uidb = bytes(buf) + obj = [uidb] + dist.broadcast_object_list(obj, src=0) + uidb = obj[0] + ctypes.memmove(buf, uidb, n) + _lib.rs_init_uid(rank, world, buf) + else: + # rs_get_uid() returns the uid as `bytes` (pybind); broadcast it as a python + # object and hand the raw bytes back to rs_init_uid on every rank. + uidb = _lib.rs_get_uid() if rank == 0 else None + obj = [uidb] + dist.broadcast_object_list(obj, src=0) + uidb = obj[0] + _lib.rs_init_uid(rank, world, uidb) + + _my_pe = _lib.rs_my_pe() + _n_pes = _lib.rs_n_pes() + assert ( + _my_pe == rank and _n_pes == world + ), f"rocSHMEM PE mismatch: my_pe={_my_pe} rank={rank} n_pes={_n_pes} world={world}" + logger.info( + "init_shmem (rocSHMEM %s, uid bootstrap): my_pe=%d n_pes=%d", + "GDA" if _gda_enabled else "host-API", + _my_pe, + _n_pes, + ) + _initialized = True + + +def _ensure_peer_deltas(base, local_world_size, rank): + """On the first allocation, compute the affine peer deltas and push them to + the Triton device layer so int_p / int_g / int_wait_until_equals can + translate a local symmetric address to a peer address on-device.""" + global _ref_base + if _ref_base is not None: + return + assert local_world_size <= MAX_LOCAL_PES, ( + f"rocshmem backend supports up to {MAX_LOCAL_PES} same-node PEs, got " + f"local_world_size={local_world_size}" + ) + _ref_base = base + node_start = rank - rank % local_world_size + # Index deltas by LOCAL position (pe - node_start), NOT global pe: on node>0 + # the same-node PEs are [node_start, node_start+lws), which would overflow a + # global-indexed 8-slot table. The device kernel translates a runtime global + # pe to a local position via the baked node_start (see _rs_peer_delta). + deltas = [0] * MAX_LOCAL_PES # indexed by LOCAL position + for i in range(local_world_size): + pe = node_start + i + if pe == _my_pe: + deltas[i] = 0 + else: + peer = _lib.rs_ptr(base, pe) + if peer == 0: + raise RuntimeError(f"rs_ptr(base, {pe}) == 0: no XGMI P2P route to same-node peer") + deltas[i] = peer - base + from .shmem_triton import set_rocshmem_peer_deltas + + set_rocshmem_peer_deltas(deltas, node_start) + logger.info( + "rocSHMEM peer deltas (node_start=%d, local_pos -> delta bytes): %s", + node_start, + deltas, + ) + + +def alloc_peer_tensors(shape, dtype, local_world_size, rank): + """Allocate one symmetric buffer and return ``(local_tensor, peer_tensors)`` + where ``peer_tensors[i]`` is the same allocation viewed from the address + space of same-node PE ``node_start + i`` (length ``local_world_size``). + + Equivalent to MORI's ``mori_shmem_create_tensor_list_intra_node``. + """ + base = _lib.rs_malloc(_nbytes(shape, dtype)) + if base == 0: + raise RuntimeError( + f"rs_malloc({_nbytes(shape, dtype)} bytes) returned NULL — symmetric " + f"heap exhausted? Raise the rocSHMEM heap size." + ) + dev = torch.cuda.current_device() + if _gda_enabled: + # GDA (USE_IPC=OFF): no intra-node IPC peer views; all cross-node AND + # same-node transfers go through device get/put. Return the local view + # plus dummy peer entries (never used on the GDA path). + local_tensor = _wrap(base, shape, dtype, dev) + _allocations.append(local_tensor) + return local_tensor, [local_tensor] * local_world_size + _ensure_peer_deltas(base, local_world_size, rank) + + node_start = rank - rank % local_world_size + peer_tensors = [] + for i in range(local_world_size): + pe = node_start + i + p = base if pe == _my_pe else _lib.rs_ptr(base, pe) + if p == 0: + raise RuntimeError(f"rs_ptr(base, {pe}) == 0 (no XGMI route to peer {pe})") + peer_tensors.append(_wrap(p, shape, dtype, dev)) + local_tensor = peer_tensors[rank % local_world_size] + _allocations.append(local_tensor) + return local_tensor, peer_tensors + + +def create_tensor(shape, dtype): + """Bare local symmetric tensor (no peer list).""" + base = _lib.rs_malloc(_nbytes(shape, dtype)) + if base == 0: + raise RuntimeError("rs_malloc returned NULL (symmetric heap exhausted?)") + t = _wrap(base, shape, dtype, torch.cuda.current_device()) + _allocations.append(t) + return t + + +def free_tensor(tensor): + # rocSHMEM host binding exposes no per-allocation free; the whole symmetric + # heap is released by rs_finalize(). No-op here (buffers live for the run). + pass + + +def barrier(): + _lib.rs_barrier() + + +# --------------------------------------------------------------------------- +# GPU-direct (GDA) device-kernel launchers (only valid when gda_enabled()). +# Pointers are raw device addresses (int) into the GDA symmetric heap. +# --------------------------------------------------------------------------- +def gda_gather(target_ptr, src_ptr, nbytes, peers, stride_bytes): + """Device gather: for each cross-node peer in ``peers`` (global PE/rank), pull + its shard (at symmetric ``src_ptr``) into ``target_ptr + peer*stride_bytes``.""" + if len(peers) == 0: + return + if _ctypes_gda_lib: + n = len(peers) + arr = (ctypes.c_int * n)(*[int(p) for p in peers]) + rc = _lib.gda_gather(int(target_ptr), int(src_ptr), int(nbytes), arr, n, int(stride_bytes)) + else: + # pybind binds gda_gather(target, src, nbytes, std::vector peers, stride) + rc = _lib.gda_gather( + int(target_ptr), int(src_ptr), int(nbytes), [int(p) for p in peers], int(stride_bytes) + ) + if rc != 0: + raise RuntimeError(f"gda_gather hipError={rc}") + + +def gda_reduce_scatter_acc( + acc_ptr, + input_ptr, + seg_off_bytes, + shard_elems, + n_pes, + scratch_ptr, + scratch_stride_bytes, + dtype_code, + nblocks, +): + """Device pull-based reduce-scatter accumulate: acc_fp32[i] += sum over all + PEs of input[seg_off + i] (pulled from each PE). Race-free (on-chip sum).""" + rc = _lib.gda_reduce_scatter_acc( + int(acc_ptr), + int(input_ptr), + int(seg_off_bytes), + int(shard_elems), + int(n_pes), + int(scratch_ptr), + int(scratch_stride_bytes), + int(dtype_code), + int(nblocks), + ) + if rc != 0: + raise RuntimeError(f"gda_reduce_scatter_acc hipError={rc}") + + +def gda_stage_fence(dst_ptr, src_ptr, nbytes): + """Copy src->dst (device) and SYSTEM-fence so the staged symmetric buffer is + visible to the NIC before a remote PE's device getmem (HDP-flush substitute).""" + rc = _lib.gda_stage_fence(int(dst_ptr), int(src_ptr), int(nbytes)) + if rc != 0: + raise RuntimeError(f"gda_stage_fence hipError={rc}") + + +def gda_hdp_init(): + """Resolve this rank's GPU HDP flush register (call after set_device). + Returns 0 on success; nonzero means the register could not be resolved.""" + if not hasattr(_lib, "gda_hdp_init"): + return -99 + return int(_lib.gda_hdp_init()) + + +def gda_hdp_flush(): + """Flush this GPU's HDP cache so prior symmetric writes are NIC-visible (the + proper GPUDirect-RDMA write-visibility primitive). Returns 1 if flushed.""" + return int(_lib.gda_hdp_flush()) + + +def gda_strided_touch( + input_ptr, + seg_off_bytes, + seg_bytes, + n_pes, + stride_bytes, + touch_bytes, + scratch_ptr, + scratch_stride_bytes, + nblocks, +): + """Strided page-touch warm-up: a tiny RDMA read at every ``stride_bytes`` page + of my shard segment on every PE, priming all pages/NICs (deterministic + read-triggered settle) at minimal volume vs the full-shard throwaway RS.""" + rc = _lib.gda_strided_touch( + int(input_ptr), + int(seg_off_bytes), + int(seg_bytes), + int(n_pes), + int(stride_bytes), + int(touch_bytes), + int(scratch_ptr), + int(scratch_stride_bytes), + int(nblocks), + ) + if rc != 0: + raise RuntimeError(f"gda_strided_touch hipError={rc}") + + +def gda_reduce_scatter_acc_async( + acc_ptr, + input_ptr, + seg_off_bytes, + shard_elems, + n_pes, + scratch_ptr, + scratch_stride_bytes, + dtype_code, + nblocks, +): + """Comm/compute OVERLAP: launch reduce-scatter on a side stream WITHOUT syncing. + Caller must gda_rs_overlap_sync() before re-staging input or consuming acc.""" + rc = _lib.gda_reduce_scatter_acc_async( + int(acc_ptr), + int(input_ptr), + int(seg_off_bytes), + int(shard_elems), + int(n_pes), + int(scratch_ptr), + int(scratch_stride_bytes), + int(dtype_code), + int(nblocks), + ) + if rc != 0: + raise RuntimeError(f"gda_reduce_scatter_acc_async launch err={rc}") + + +def gda_rs_overlap_sync(): + """Wait for all pending overlapped reduce-scatter kernels on the side stream.""" + rc = _lib.gda_rs_overlap_sync() + if rc != 0: + raise RuntimeError(f"gda_rs_overlap_sync hipError={rc}") + + +def gda_gather_async(target_ptr, src_ptr, nbytes, peers, stride_bytes, stream): + """Approach 1: launch all-gather kernel on the given HIP `stream` WITHOUT syncing, + so FSDP2 prefetch overlaps it with compute. Reassembly + consumer order via the + same stream (gather reads stable params -> no settle/barrier needed).""" + if len(peers) == 0: + return + if _ctypes_gda_lib: + n = len(peers) + arr = (ctypes.c_int * n)(*[int(p) for p in peers]) + rc = _lib.gda_gather_async( + int(target_ptr), + int(src_ptr), + int(nbytes), + arr, + n, + int(stride_bytes), + int(stream), + ) + else: + rc = _lib.gda_gather_async( + int(target_ptr), + int(src_ptr), + int(nbytes), + [int(p) for p in peers], + int(stride_bytes), + int(stream), + ) + if rc != 0: + raise RuntimeError(f"gda_gather_async launch err={rc}") + + +def dtype_code(dtype): + return _DTYPE_CODE[dtype] + + +def finalize(): + global _initialized + if _lib is not None and _initialized: + _lib.rs_finalize() + _initialized = False diff --git a/primus/core/odc/primitives/gather.py b/primus/core/odc/primitives/gather.py new file mode 100644 index 000000000..752da1659 --- /dev/null +++ b/primus/core/odc/primitives/gather.py @@ -0,0 +1,244 @@ +# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under +# the MIT License per its package metadata (pyproject.toml / setup.py +# classifiers). The upstream repository ships no LICENSE file or per-file +# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab). +# +# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# See LICENSE for license information. + +import logging +import math + +import torch +import torch.distributed as dist +import triton +import triton.language as tl +from torch import Tensor + +from odc.primitives import ( + SHMEM_EXTERN_LIBS, + __syncthreads, + getmem_nbi_block, + quiet, + tid, +) +from odc.primitives.utils import ( + BufferSplitter, + SymmBufferRegistry, + get_comm_stream, + get_local_world_size, + sync_cta, +) + +logger = logging.getLogger(__name__) + +from odc.primitives.utils import _USE_ROCSHMEM # noqa: E402 + +if _USE_ROCSHMEM: + from odc.primitives import _rocshmem_backend as _rs +else: + _rs = None + + +def _gda_active(): + return _rs is not None and _rs.gda_enabled() + + +@triton.jit +def shmem_device_producer_gather_2d_get_block_kernel_chunked_synced( + remote_tensor_ptr, + target_tensor_ptr, + elem_per_rank, + size_per_elem, + rank, + num_ranks_per_node, + world_size, + chunk_size, + signal_ptr, +): + pid = tl.program_id(axis=0) + # np = tl.num_programs(axis=0) + assert num_ranks_per_node == tl.num_programs(axis=0) + np = num_ranks_per_node + num_nodes = world_size // np + + tidx = tid(axis=0) + expected = 0 + for i in range(1, num_nodes): + peer_node = (i + rank // np) % num_nodes + peer = (pid + peer_node * np) % world_size + # chunk_size = elem_per_rank // num_chunks + num_chunks = tl.cdiv(elem_per_rank, chunk_size) + for chunk in range(num_chunks): + this_chunk_size = chunk_size + if chunk == num_chunks - 1: + this_chunk_size = elem_per_rank - chunk * chunk_size + getmem_nbi_block( + target_tensor_ptr + peer * elem_per_rank + (chunk * chunk_size), + remote_tensor_ptr + (chunk * chunk_size), + this_chunk_size * size_per_elem, + peer, + ) + expected += np + sync_cta(signal_ptr, expected) + if tidx == 0 and pid == 0: + quiet() + __syncthreads() + + expected += np + sync_cta(signal_ptr, expected) + + +class GatherService: + def __init__(self): + self.shaped_buffer = {} + self.buffer_splitter = BufferSplitter() + self.chunk_size_bytes = 2**20 + + def get_chunk_size(self, buffer_dtype): + return self.chunk_size_bytes // buffer_dtype.itemsize + + def gather_into_tensor(self, output_tensor: Tensor, input_tensor: Tensor, pg: dist.ProcessGroup): + buf_size = self.buffer_splitter.get_global_buffer_size(output_tensor.shape) + buffer_shape = (buf_size,) + output_size = output_tensor.numel() + assert output_size >= buf_size, f"output_size: {output_size} < buf_size: {buf_size}" + + rank = torch.distributed.get_rank() + if (buffer_shape, output_tensor.dtype) not in self.shaped_buffer: + logger.info( + f"Rank {rank} create buffer: output_size: {output_size} num_sub_buffers: {math.ceil(output_size / buf_size)} buf_size: {buf_size}" + ) + self.shaped_buffer[ + (buffer_shape, output_tensor.dtype) + ] = SymmBufferRegistry.get_instance().allocate_symm_buffer( + f"ag_buffer_{buffer_shape}_{output_tensor.dtype}", + buffer_shape, + output_tensor.dtype, + ) + target_tensor = self.shaped_buffer[(buffer_shape, output_tensor.dtype)] + + assert (input_tensor.numel() * input_tensor.element_size()) % ( + 2**6 + ) == 0 or input_tensor.numel() < 2**6, "better align to 64 for efficiency" + chunk_size = self.get_chunk_size(input_tensor.dtype) + # assert input_tensor.numel() % chunk_size == 0 + + registry = SymmBufferRegistry.get_instance() + peer_tensors = registry.get_peer_tensors(input_tensor) + + group_world_size = torch.distributed.get_world_size(pg) + local_world_size = get_local_world_size() + assert group_world_size in ( + torch.distributed.get_world_size(), + local_world_size, + ), f"{group_world_size=} {torch.distributed.get_world_size()=} {local_world_size=}" + + get_comm_stream().wait_stream(torch.cuda.current_stream()) + # GPU-direct (GDA): IPC is off, so there are no XGMI peer views; ALL ranks + # (incl. same-node and self) are pulled via device rocshmem_getmem. Handle + # the whole cross-group gather here and return. + gda = _gda_active() and local_world_size != group_world_size + if gda: + self._gda_gather_into_tensor( + output_tensor, + input_tensor, + target_tensor, + buf_size, + group_world_size, + rank, + ) + return + + with torch.cuda.stream(get_comm_stream()): + output_tensor_split = output_tensor.view(group_world_size, -1) + assert local_world_size == len(peer_tensors) + local_rank = rank % local_world_size + rank_same_node_start = rank - local_rank + rank_same_node_end = rank_same_node_start + local_world_size + for r_offset in range(local_world_size): + src_local_rank = (local_rank + r_offset) % local_world_size + if group_world_size == local_world_size: + output_tensor_split[src_local_rank].copy_(peer_tensors[src_local_rank]) + else: + src_rank = rank_same_node_start + src_local_rank + output_tensor_split[src_rank].copy_(peer_tensors[src_local_rank]) + + assert buf_size % group_world_size == 0 + local_buf_size = buf_size // group_world_size + signal_ptr = torch.empty(1, dtype=torch.int32, device="cuda") + for start in range(0, input_tensor.numel(), local_buf_size): + if local_world_size == group_world_size: + continue + size = min(local_buf_size, input_tensor.numel() - start) + sub_input_tensor = input_tensor.view(-1)[start : start + size] + assert (sub_input_tensor.numel() * sub_input_tensor.element_size()) % ( + 2**6 + ) == 0 or sub_input_tensor.numel() < 2**6, "better align to 64 for efficiency" + target_buf_size = size * group_world_size + assert target_buf_size <= buf_size + target_tensor_split = target_tensor[:target_buf_size].view(group_world_size, size) + + signal_ptr.fill_(0) + assert group_world_size % 8 == 0 or group_world_size < 8 + # grid_size = 8 if world_size == 32 else world_size + grid_size = local_world_size + shmem_device_producer_gather_2d_get_block_kernel_chunked_synced[(grid_size,)]( + remote_tensor_ptr=sub_input_tensor, + target_tensor_ptr=target_tensor_split.view(-1), + elem_per_rank=sub_input_tensor.numel(), + size_per_elem=sub_input_tensor.element_size(), + rank=rank, + num_ranks_per_node=local_world_size, + world_size=group_world_size, + chunk_size=chunk_size, + signal_ptr=signal_ptr, + num_warps=32, + extern_libs=SHMEM_EXTERN_LIBS, + ) + if buf_size == output_size: + local_world_data_size = size * local_world_size + local_world_idx = rank // local_world_size + data_start_idx = local_world_data_size * local_world_idx + data_end_idx = data_start_idx + local_world_data_size + output_tensor[:data_start_idx].copy_(target_tensor[:data_start_idx]) + output_tensor[data_end_idx:].copy_(target_tensor[data_end_idx:]) + # output_tensor.copy_(target_tensor) + else: + for r in range(group_world_size): + if rank_same_node_start <= r < rank_same_node_end: + continue + output_tensor_split[r, start : start + size].copy_(target_tensor_split[r, :]) + torch.cuda.current_stream().wait_stream(get_comm_stream()) + + def _gda_gather_into_tensor( + self, output_tensor, input_tensor, target_tensor, buf_size, group_world_size, rank + ): + """GPU-direct all-gather: every rank's shard (incl same-node and self) is + pulled via a single device rocshmem_getmem kernel into target slot[r], + then copied into output_tensor_split[r]. No XGMI peer views / host loop.""" + es = input_tensor.element_size() + assert buf_size % group_world_size == 0 + local_buf_size = buf_size // group_world_size + output_split = output_tensor.view(group_world_size, -1) + peers = list(range(group_world_size)) + comm = get_comm_stream() + comm.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(comm): + torch.cuda.synchronize() + # NO per-gather rocshmem_barrier_all here. Gather reads PARAMS (symmetric + # shard written by the optimizer a full step earlier and read-only through + # fwd/bwd) -> peers' data is already stable/visible, so the cross-PE + # rendezvous is unnecessary (unlike scatter's just-written staging). + for start in range(0, input_tensor.numel(), local_buf_size): + size = min(local_buf_size, input_tensor.numel() - start) + sub_input = input_tensor.view(-1)[start : start + size] + tb = size * group_world_size + tsplit = target_tensor[:tb].view(group_world_size, size) + _rs.gda_gather(target_tensor.data_ptr(), sub_input.data_ptr(), size * es, peers, size * es) + torch.cuda.synchronize() + # reassembly on the SAME comm stream -> ordered AFTER the gather kernel + for r in range(group_world_size): + output_split[r, start : start + size].copy_(tsplit[r]) + torch.cuda.current_stream().wait_stream(comm) diff --git a/primus/core/odc/primitives/scatter_accumulate.py b/primus/core/odc/primitives/scatter_accumulate.py new file mode 100644 index 000000000..1dc3ca7cf --- /dev/null +++ b/primus/core/odc/primitives/scatter_accumulate.py @@ -0,0 +1,431 @@ +# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under +# the MIT License per its package metadata (pyproject.toml / setup.py +# classifiers). The upstream repository ships no LICENSE file or per-file +# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab). +# +# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# See LICENSE for license information. + +import logging + +import torch +import torch.distributed as dist + +from odc.primitives.utils import ( + SymmBufferRegistry, + get_comm_stream, + get_local_world_size, +) +from odc.runtime_config import get_config + +logger = logging.getLogger(__name__) + +from odc.primitives.utils import _USE_ROCSHMEM # noqa: E402 + +if _USE_ROCSHMEM: + from odc.primitives import _rocshmem_backend as _rs +else: + _rs = None + + +def _gda_active(): + return _rs is not None and _rs.gda_enabled() + + +class ReductionService: + """Reduce-scatter-accumulate service (device-side, no host-polling subprocess). + + Two paths, both device-side and free of any host-polling subprocess: + * single node -> owner-side XGMI PULL + on-chip fp32 sum + (``_single_device_scatter_accumulate``). + * multi node -> rocSHMEM GPU-direct (GDA) reduce-scatter + (``_gda_scatter_accumulate``). + + Both DEFER the collective reduce to once per param-group at + ``get_accumulation`` (matched barrier count across ranks -> deadlock-free + under variable-length / nopad micro-batching), pre-accumulating each + micro-batch's grad locally in fp32. + """ + + def __init__(self, accumulation_dtype=None): + self.accumulations = [] + self.accumulation_indices = {} + self.input_buffer = {} + self.dispatched_tasks = 0 + # Accepted for API compatibility with the FSDP1 caller; the device and + # GDA reduce paths always accumulate in fp32. + self.accumulation_dtype = accumulation_dtype + + def clear_accumulations(self): + for acc in self.accumulations: + acc.fill_(0) + if hasattr(self, "_gda_deferred"): # reset per-minibatch deferred grads + self._gda_deferred = {} + self._gda_deferred_pg = {} + if hasattr(self, "_sdr_deferred"): # reset per-minibatch deferred grads + self._sdr_deferred = {} + self._sdr_deferred_pg = {} + + def _gda_scatter_accumulate(self, key, input_tensor, pg: dist.ProcessGroup): + """GPU-direct pull-based reduce-scatter accumulate (race-free, no host-polling subprocess). + + Stages this rank's full input into a symmetric buffer, barriers, then a + device kernel pulls every PE's contribution to MY output shard and sums + it on-chip into the fp32 accumulation buffer (acc += reduce_scatter(input)). + """ + gws = torch.distributed.get_world_size(pg) + assert ( + gws == _rs._n_pes + ), f"GDA reduce-scatter requires a full-world group: gws={gws} n_pes={_rs._n_pes}" + assert input_tensor.numel() % gws == 0, f"{input_tensor.numel()=} % {gws=}" + shard_elems = input_tensor.numel() // gws + dt = input_tensor.dtype + es = input_tensor.element_size() + reg = SymmBufferRegistry.get_instance() + + if key not in self.accumulation_indices: + acc = reg.get_or_create_symm_buffer(f"gda_acc_{key}", (shard_elems,), torch.float32) + acc.fill_(0) + self.accumulation_indices[key] = len(self.accumulations) + self.accumulations.append(acc) + acc = self.accumulations[self.accumulation_indices[key]] + + in_key = ("gda_in", dt, input_tensor.numel()) + if in_key not in self.input_buffer: + self.input_buffer[in_key] = reg.get_or_create_symm_buffer( + f"gda_in_{dt}_{input_tensor.numel()}", (input_tensor.numel(),), dt + ) + input_sym = self.input_buffer[in_key] + + cfg = get_config() + # GRID geometry (config odc_gda_rs_blocks): reduce-scatter kernel grid (one + # block per disjoint shard chunk -> # concurrent cross-node getmem_wg = QP/NIC + # parallelism lever). Scratch auto-resizes (sc_key includes chunk*nblk), so + # this stays correct for any nblk. Default 64 (current behavior). + nblk = int(cfg.gda_rs_blocks) + if nblk < 1: + nblk = 1 + # PIPE (config odc_gda_pipe): peer-pipeline batch depth. The pipelined rs_acc + # needs `pipe` scratch slots PER BLOCK (issues `pipe` peers' nbi getmem + # concurrently), so scratch grows pipe x and the main call passes + # scratch_stride = pipe*chunk. This is also a Primus-Turbo device-kernel knob + # (odc_rocshmem_gda.cu env_pipe); set_runtime_config bridges odc_gda_pipe back + # to the PRIMUS_TURBO_ODC_GDA_PIPE env var so the C++ side agrees. + pipe = int(cfg.gda_pipe) + if pipe < 1: + pipe = 1 + chunk = (shard_elems + nblk - 1) // nblk + sc_slots = nblk * pipe + sc_key = ("gda_scr", dt, chunk * sc_slots) + if sc_key not in self.input_buffer: + self.input_buffer[sc_key] = reg.get_or_create_symm_buffer( + f"gda_scr_{dt}_{chunk * sc_slots}", (chunk * sc_slots,), dt + ) + scratch = self.input_buffer[sc_key] + + rank = torch.distributed.get_rank(pg) + + # Cross-node write-visibility strategy for the just-staged grad + # (config odc_gda_warmup_mode). Modes: + # "strided" (default) - page-strided tiny throwaway READ that keeps + # full-warmup's deterministic "read-triggered settle" (validated 0 grad + # spikes, loss == single-node) at minimal volume; stride via + # odc_gda_stride_bytes (default 64KB). + # "full" - full-shard throwaway reduce-scatter: most robust, slowest. + # "hdp" - O(1) HDP_MEM_FLUSH_CNTL register write; auto-falls back to + # "full" if no HDP register. "fence"/"hdpfence" - device system-scope + # fence (+ optional HDP flush) instead of the read settle. + _warm_mode = cfg.gda_warmup_mode + if _warm_mode == "hdp" and getattr(self, "_hdp_fallback", False): + _warm_mode = "full" + if not getattr(self, "_logged_warm_mode", False): + logger.warning( + "[GDA] reduce-scatter warm-up mode=%s stride_bytes=%s", + _warm_mode, + cfg.gda_stride_bytes, + ) + self._logged_warm_mode = True + if _warm_mode in ("fence", "hdpfence"): + # Ensure the grad (input_tensor) is fully produced before copy+fence, + # then copy into the symmetric staging buffer with a trailing system-scope + # fence (gda_stage_fence self-syncs the device). "hdpfence" additionally + # flushes this GPU's HDP so the remote NIC's RDMA read sees fresh data. + torch.cuda.current_stream().synchronize() + _flat = input_tensor.view(-1) + _rs.gda_stage_fence(input_sym.data_ptr(), _flat.data_ptr(), _flat.numel() * es) + if _warm_mode == "hdpfence": + if not getattr(self, "_hdp_inited", False) and not getattr(self, "_hdp_fallback", False): + rc = _rs.gda_hdp_init() + if rc == 0: + self._hdp_inited = True + else: + logger.warning("gda_hdp_init failed rc=%d -> hdpfence falls back to fence-only", rc) + self._hdp_fallback = True + if getattr(self, "_hdp_inited", False): + _rs.gda_hdp_flush() + else: + get_comm_stream().wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(get_comm_stream()): + input_sym.copy_(input_tensor.view(-1)) + torch.cuda.synchronize() + if _warm_mode == "hdp": + if not getattr(self, "_hdp_inited", False): + rc = _rs.gda_hdp_init() + if rc != 0: + logger.warning( + "gda_hdp_init failed rc=%d -> falling back to full-shard warm-up RS", rc + ) + self._hdp_fallback = True + _warm_mode = "full" + else: + self._hdp_inited = True + if _warm_mode == "hdp": + _rs.gda_hdp_flush() + _rs.barrier() + + # Cross-node write-visibility settle (FIX for the intermittent stale-read + # grad spikes): a throwaway "warm-up" reduce-scatter + barrier BEFORE the + # real one. A single barrier after staging is insufficient on this + # mlx5/GDA path (peer's just-staged GPU write isn't yet NIC-visible to the + # first device getmem -> stale read -> huge wrong gradient ~half the time). + # The warm-up getmem + barrier forces the staged data visible; the real + # reduce-scatter then reads fresh data. (Verified: eliminates spikes, + # grad norms normal, loss matches single-node. Reduce-scatter is a small + # fraction of the step, so the extra pass is cheap vs the gather.) + # Opt1: the warm-up only needs to settle cross-node write-visibility (drain + # the HDP so peers' staged writes are NIC-visible), NOT move the full + # shard. A tiny getmem touch per peer + barrier achieves that at a + # fraction of the cost (full-shard warm-up was ~9s/iter of pure overhead). + # warm-up settles cross-node write-visibility (HDP) before the real RS. + # Single-NIC: a 1024-elem touch suffices. Multi-NIC: each peer routes via a + # different NIC and HDP is per-NIC/per-page, so a tiny touch leaves most + # pages stale -> spikes; use the FULL shard (default) so every NIC/page is + # settled. The full warm-up is itself parallelized across NICs (cheap). + if _warm_mode in ("fence", "hdp", "hdpfence"): + # The staged write was already made NIC-visible (HDP flush / fence), + # so the throwaway warm-up reduce-scatter (the ~59%-of-RS overhead) is + # skipped entirely. The barrier after staging still rendezvouses PEs. + pass + elif _warm_mode == "strided": + # Page-strided tiny throwaway READ: keeps full-warmup's deterministic + # "read-triggered settle" (covers every page of my segment on every PE + # -> all 8 NICs) but at minimal volume (one touch per odc_gda_stride_bytes + # page, not the whole shard). Staging above used plain copy_ (no flush). + stride_b = int(cfg.gda_stride_bytes) + touch_b = int(es) + seg_bytes = int(shard_elems * es) + npages = (seg_bytes + stride_b - 1) // stride_b + total_touch = _rs._n_pes * max(npages, 1) + sstride = 256 # bytes/throwaway scratch slot (>= touch_b; avoids collide) + scratch_cap = scratch.numel() * scratch.element_size() + nblk_t = min(int(total_touch), 4096, max(1, scratch_cap // sstride)) + _rs.gda_strided_touch( + input_sym.data_ptr(), + rank * shard_elems * es, + seg_bytes, + _rs._n_pes, + stride_b, + touch_b, + scratch.data_ptr(), + sstride, + int(nblk_t), + ) + _rs.barrier() + else: # "full": full-shard throwaway reduce-scatter settle + n_warm, w_nblk, w_stride = int(shard_elems), nblk, chunk * es + wkey = ("gda_warmup", n_warm) + if wkey not in self.input_buffer: + self.input_buffer[wkey] = torch.zeros(n_warm, dtype=torch.float32, device="cuda") + warmup = self.input_buffer[wkey] + warmup.zero_() + _rs.gda_reduce_scatter_acc( + warmup.data_ptr(), + input_sym.data_ptr(), + rank * shard_elems * es, + n_warm, + _rs._n_pes, + scratch.data_ptr(), + w_stride, + _rs.dtype_code(dt), + w_nblk, + ) + _rs.barrier() + _rs.gda_reduce_scatter_acc( + acc.data_ptr(), + input_sym.data_ptr(), + rank * shard_elems * es, + shard_elems, + _rs._n_pes, + scratch.data_ptr(), + pipe * chunk * es, + _rs.dtype_code(dt), + nblk, + ) + torch.cuda.synchronize() + self.dispatched_tasks += 1 + + def _single_device_scatter_accumulate(self, key, input_tensor, pg: dist.ProcessGroup): + """Single-node device-side reduce-scatter accumulate (no host-polling subprocess). + + Mechanism (owner-side PULL + on-chip fp32 sum over same-node XGMI peer + views): each PE stages its (locally pre-accumulated) full grad into a + symmetric fp32 buffer whose same-node peer views are already resolved on + allocation (the same XGMI peer-view machinery gather.py uses); after a + rendezvous barrier, PE r reads every same-node peer's segment destined + for its shard (a plain XGMI ``.copy_`` peer read) and sums it on-chip. + There are NO cross-rank writes -> no device atomics -> no MI300X + write-visibility hazard; only XGMI reads gated by a barrier. No second + process -> no IPC handle and no host-side reduction subprocess. + + ``input_tensor`` is this PE's full (locally pre-accumulated) grad, laid + out as ``[shard_0 | shard_1 | ... | shard_{gws-1}]``; PE r owns output + shard r. + + Steps (all same-node, no subprocess, no IPC handle, no atomics): + 1. Stage my full grad into a symmetric fp32 buffer (peer views resolved + on allocation). + 2. cuda sync + collective barrier -> every peer's stage is retired and + visible before any peer reads it over XGMI. + 3. For each same-node peer p, XGMI-``.copy_`` p's segment destined for + MY shard into a local temp, then acc += temp (fp32 on-chip sum). + 4. cuda sync + collective barrier -> all reads done before the shared + staging buffer can be reused by the next key/minibatch. + """ + gws = torch.distributed.get_world_size(pg) + lws = get_local_world_size() + assert gws == lws, f"single-device reduce is same-node only: gws={gws} lws={lws}" + assert input_tensor.numel() % gws == 0, f"{input_tensor.numel()=} % {gws=}" + shard_elems = input_tensor.numel() // gws + reg = SymmBufferRegistry.get_instance() + rank = torch.distributed.get_rank(pg) # single node -> group rank == local pos + + # fp32 accumulator = MY output shard (matches GDA's fp32 acc semantics). + if key not in self.accumulation_indices: + acc = reg.get_or_create_symm_buffer(f"sdr_acc_{key}", (shard_elems,), torch.float32) + acc.fill_(0) + self.accumulation_indices[key] = len(self.accumulations) + self.accumulations.append(acc) + acc = self.accumulations[self.accumulation_indices[key]] + + # Symmetric fp32 staging for my full grad, WITH same-node peer views. One + # per grad-numel; reused across keys/minibatches (guarded by the trailing + # barrier so no peer reuses it mid-read). + in_key = ("sdr_in", input_tensor.numel()) + if in_key not in self.input_buffer: + self.input_buffer[in_key] = reg.get_or_create_symm_buffer( + f"sdr_in_{input_tensor.numel()}", (input_tensor.numel(),), torch.float32 + ) + input_sym = self.input_buffer[in_key] + peer_inputs = reg.get_peer_tensors(input_sym) + assert len(peer_inputs) == lws + + tmp_key = ("sdr_tmp", shard_elems) + if tmp_key not in self.input_buffer: + self.input_buffer[tmp_key] = torch.empty(shard_elems, dtype=torch.float32, device="cuda") + tmp = self.input_buffer[tmp_key] + + input_sym.copy_(input_tensor.view(-1).to(torch.float32)) + torch.cuda.synchronize() + torch.distributed.barrier(group=pg) + + lo = rank * shard_elems + hi = lo + shard_elems + # Rotate the peer order by local rank so all 8 PEs don't hammer the same + # peer's HBM first (matches the gather/scatter round-robin peer ordering). + for off in range(lws): + p = (rank + off) % lws + tmp.copy_(peer_inputs[p][lo:hi]) + acc.add_(tmp) + torch.cuda.synchronize() + torch.distributed.barrier(group=pg) + self.dispatched_tasks += 1 + + def scatter_accumulate(self, key, input_tensor, pg: dist.ProcessGroup): + if _gda_active(): + # DEFER the cross-node reduce to once-per-minibatch. The per-call + # _gda_scatter_accumulate does rocshmem_barrier_all (collective); + # calling it per-microbatch deadlocks under nopad (ranks have + # different micro-batch counts -> mismatched barrier counts). Fix: + # accumulate the unsharded grad LOCALLY here (no comm/barrier -> + # backward is lockstep-free), then do ONE barriered reduce-scatter + # per group at get_accumulation (count == #groups, matched across + # ranks). Config odc_gda_defer_reduce: "auto" (default) => defer iff + # multi-node (n_pes > local_world_size); "1"/"0" force on/off. + _defer = str(get_config().gda_defer_reduce) + if _defer == "auto": + _defer_on = _rs._n_pes > get_local_world_size() + else: + _defer_on = _defer in ("1", "true", "True") + if _defer_on: + if not hasattr(self, "_gda_deferred"): + self._gda_deferred = {} + self._gda_deferred_pg = {} + cur = self._gda_deferred.get(key) + if cur is None: + self._gda_deferred[key] = input_tensor.detach().clone() + else: + cur.add_(input_tensor) + self._gda_deferred_pg[key] = pg + return + return self._gda_scatter_accumulate(key, input_tensor, pg) + # Single-node device-side reduce (replaces the removed host-side + # reduction subprocess path). Cross-node reduce always goes through the + # rocSHMEM GDA path above. + assert torch.distributed.get_world_size(pg) == get_local_world_size(), ( + f"non-GDA reduce-scatter is single-node only " + f"(gws={torch.distributed.get_world_size(pg)} lws={get_local_world_size()}); " + f"multi-node requires the rocSHMEM GDA backend" + ) + # DEFER exactly like the multi-node GDA path: pre-accumulate this + # micro-batch's grad LOCALLY (no collective -> nopad-safe), then run ONE + # barriered owner-side pull-sum per group at get_accumulation. + # Pre-accumulate in fp32 so cross-micro-batch accumulation matches the + # fp32 device accumulator. + if not hasattr(self, "_sdr_deferred"): + self._sdr_deferred = {} + self._sdr_deferred_pg = {} + cur = self._sdr_deferred.get(key) + if cur is None: + self._sdr_deferred[key] = input_tensor.detach().to(torch.float32) + else: + cur.add_(input_tensor) + self._sdr_deferred_pg[key] = pg + return + + def get_accumulation(self, key): + # GDA DEFER: run the single per-minibatch cross-node reduce-scatter now + # (once per group, matched barrier count across ranks -> deadlock-free). + if _gda_active() and getattr(self, "_gda_deferred", None) is not None: + pending = self._gda_deferred.get(key) + if pending is not None: + self._gda_deferred[key] = None + self._gda_scatter_accumulate(key, pending, self._gda_deferred_pg[key]) + # Single-node device DEFER: run the single per-minibatch owner-side + # pull-sum now (once per group, matched barrier count across ranks). + if not _gda_active() and getattr(self, "_sdr_deferred", None) is not None: + pending = self._sdr_deferred.get(key) + if pending is not None: + self._sdr_deferred[key] = None + self._single_device_scatter_accumulate(key, pending, self._sdr_deferred_pg[key]) + acc = self.accumulations[self.accumulation_indices[key]] + return acc + + def sync(self, pg: dist.ProcessGroup): + if _gda_active(): + # GDA path is fully synchronous (device kernels + barriers). + torch.cuda.synchronize() + torch.distributed.barrier(group=pg) + self.dispatched_tasks = 0 + return + # Single-node device path: the owner-side pull-sum already ran + # (synchronously, with its own barriers) inside get_accumulation. Just a + # final cuda sync + barrier before the optimizer reads acc. + torch.cuda.synchronize() + torch.distributed.barrier(group=pg) + self.dispatched_tasks = 0 + + def stop(self): + # The device and GDA reduce paths run no host-polling subprocess, so there is + # nothing to tear down. Kept for API compatibility with the callers. + pass diff --git a/primus/core/odc/primitives/shmem_triton.py b/primus/core/odc/primitives/shmem_triton.py new file mode 100644 index 000000000..2747b2158 --- /dev/null +++ b/primus/core/odc/primitives/shmem_triton.py @@ -0,0 +1,465 @@ +# Adapted from ODC (https://github.com/sail-sg/odc), whose upstream Triton +# device API (odc/primitives/nvshmem_triton.py) this ROCm reimplementation is +# derived from. ODC is distributed under the MIT License per its package +# metadata (pyproject.toml / setup.py classifiers); the upstream repository +# ships no LICENSE file or per-file copyright headers, and upstream copyright is +# held by the ODC authors (Sea AI Lab). +# +# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# See LICENSE for license information. + +""" +ODC Triton device API — ROCm implementation. + +ODC's Triton kernels reach symmetric peer memory through a small set of +device primitives. On ROCm there are two backends, selected by the +``ODC_P2P_BACKEND`` environment variable: + +* ``mori`` (default): re-exports MORI-SHMEM's device API. MORI provides the + device bitcode (``libmori_shmem_device.bc``, JIT-compiled by + ``mori.ir.find_bitcode()``) that is linked into the Triton kernels. +* ``rocshmem``: a host-API rocSHMEM backend (single-node / IPC-only) that + links NO device bitcode — every P2P op is plain Triton to an XGMI-mapped + peer address. + +Module-level public names (stable across backends): + putmem_nbi_block, getmem_nbi_block, quiet, int_p, int_g, + int_atomic_compare_swap, int_atomic_swap, + tid, __syncthreads, + LIB_SHMEM_PATH, SHMEM_EXTERN_LIBS + +Design notes +------------ +1. ``putmem_nbi_block``, ``getmem_nbi_block``, ``int32_p`` and other MORI + device functions take an *extra* ``qp`` (queue-pair index) argument. The + ODC wrappers below thread a fixed ``qp=0`` inside ``@triton.jit``-decorated + thin shims so callers in ODC need not change. +2. MORI has no single-int blocking get, so ``int_g`` is implemented via a + 4-byte ``getmem_thread`` into a scratch register. +3. ``tid(axis=0)`` was used by ``utils.sync_cta`` to elect a "first thread" + in a wave that performs an atomic_add. ``llvm.nvvm.read.ptx.sreg.tid.x`` + does not exist on AMDGPU. We provide a Triton-native alternative: an + ``arange(0, 1)`` lane vector with mask ``offsets == 0`` restricts execution + to lane 0, so ``tid`` is implemented as a constant that the caller compares + against 0. +4. ``__syncthreads`` is implemented via Triton's IR barrier + (``create_barrier``), which lowers to ``s_barrier`` on AMDGPU. +5. ``SHMEM_EXTERN_LIBS`` is ``mori.ir.triton.get_extern_libs()`` which locates + the JIT-compiled MORI device bitcode for the current GPU + NIC. The Triton + compile hook ``install_hook()`` is called at import time. +""" + +import triton +import triton.language as tl +from packaging import version +from triton.language import core + +# --------------------------------------------------------------------------- +# ODC P2P backend selector (config item odc_p2p_backend). Default "mori" keeps +# the existing, verified behaviour byte-for-byte; "rocshmem" selects the host-API +# backend (single-node / IPC-only) defined in the dedicated branch below. The ODC +# integration patch populates the runtime config before this module is imported. +# --------------------------------------------------------------------------- +from odc.runtime_config import get_config # noqa: E402 + +_P2P_BACKEND = get_config().p2p_backend.lower() + + +def is_triton_version_supported(): + return version.parse(triton.__version__) >= version.parse("3.4.0") + + +# =========================================================================== +# Branch A: rocSHMEM host-API backend (single-node / IPC-only). +# +# No device bitcode is linked into Triton (SHMEM_EXTERN_LIBS == {}); every +# P2P op is plain Triton to an XGMI-mapped peer address. A LOCAL symmetric +# address is translated to a peer address by adding a per-PE affine byte delta +# (rocshmem_ptr is affine — empirically verified). The delta table is populated +# once at init by _rocshmem_backend.alloc_peer_tensors (before any kernel +# launch) via set_rocshmem_peer_deltas() below. +# +# REGRESSION GUARD (mirrors the mori branch): no RDMA / device-library symbol +# ever enters these kernels — they are pure Triton. Use a per-rank/per-run +# TRITON_CACHE_DIR so a rank never loads another rank's cached kernel (the +# baked deltas differ per process). +# =========================================================================== +if _P2P_BACKEND == "rocshmem": + # Empty extern-libs: Triton compiles these kernels with NO device bitcode. + SHMEM_EXTERN_LIBS = {} + LIB_NAME = "librocshmem_host" + LIB_SHMEM_PATH = "" + + # Per-PE affine pointer deltas (bytes), indexed by LOCAL position + # (global pe - node_start), baked into the int_p / int_g / + # int_wait_until_equals kernels at compile time. Triton only lets @jit + # functions read globals that are tl.constexpr, so the deltas are stored as + # constexpr (their values are captured AND versioned in the compile cache + # key per the used_global_vals mechanism -> a rank never reuses another + # rank's kernel even on a shared TRITON_CACHE_DIR). ``_RS_NODE_START`` is the + # first global rank on this node so a runtime global pe maps to a local slot. + _RS_D0 = tl.constexpr(0) + _RS_D1 = tl.constexpr(0) + _RS_D2 = tl.constexpr(0) + _RS_D3 = tl.constexpr(0) + _RS_D4 = tl.constexpr(0) + _RS_D5 = tl.constexpr(0) + _RS_D6 = tl.constexpr(0) + _RS_D7 = tl.constexpr(0) + _RS_NODE_START = tl.constexpr(0) + + def set_rocshmem_peer_deltas(deltas, node_start=0): + """Populate the device peer-delta table. Called once at init, BEFORE any + kernel launch. ``deltas[i]`` == peer_addr(x, node_start+i) - x for + symmetric x (indexed by LOCAL position i). ``node_start`` is this node's + first global rank, used to map a runtime global pe to a local slot.""" + global _RS_D0, _RS_D1, _RS_D2, _RS_D3, _RS_D4, _RS_D5, _RS_D6, _RS_D7 + global _RS_NODE_START + d = list(deltas) + [0] * (8 - len(deltas)) + _RS_D0 = tl.constexpr(d[0]) + _RS_D1 = tl.constexpr(d[1]) + _RS_D2 = tl.constexpr(d[2]) + _RS_D3 = tl.constexpr(d[3]) + _RS_D4 = tl.constexpr(d[4]) + _RS_D5 = tl.constexpr(d[5]) + _RS_D6 = tl.constexpr(d[6]) + _RS_D7 = tl.constexpr(d[7]) + _RS_NODE_START = tl.constexpr(node_start) + + @triton.jit + def _rs_peer_delta(pe): + # Select the baked delta for runtime global ``pe`` as int64 via a masked + # sum over LOCAL positions (exactly one term is non-zero for a same-node + # pe; all terms are int64). + p = pe.to(tl.int64) - _RS_NODE_START + d = (p == 0).to(tl.int64) * _RS_D0 + d += (p == 1).to(tl.int64) * _RS_D1 + d += (p == 2).to(tl.int64) * _RS_D2 + d += (p == 3).to(tl.int64) * _RS_D3 + d += (p == 4).to(tl.int64) * _RS_D4 + d += (p == 5).to(tl.int64) * _RS_D5 + d += (p == 6).to(tl.int64) * _RS_D6 + d += (p == 7).to(tl.int64) * _RS_D7 + return d + + @triton.jit + def _rs_peer_addr(local_ptr, pe): + # Translate a LOCAL symmetric pointer to the peer's XGMI-mapped address. + delta = _rs_peer_delta(pe).to(tl.uint64, bitcast=True) + return local_ptr.to(tl.uint64, bitcast=True) + delta + + @triton.jit + def int_p(dest, value, pe): + """Same-node single-int put: SYSTEM-scope release atomic store to the + peer's XGMI-mapped fine-grained slot (visible to a peer CPU-side read). + ``pe == my_pe`` -> delta 0 -> local store.""" + peer = _rs_peer_addr(dest, pe).to(tl.pointer_type(tl.int32), bitcast=True) + tl.atomic_xchg(peer, value, sem="release", scope="sys") + + @triton.jit + def int_g(src, pe): + """Single int read from a peer slot. A volatile load is fine for a + one-shot read; spin loops must use int_wait_until_equals.""" + peer = _rs_peer_addr(src, pe).to(tl.pointer_type(tl.int32), bitcast=True) + return tl.load(peer, volatile=True) + + @triton.jit + def int_wait_until_equals(ptr, expected, pe): + """Spin until the peer slot == ``expected`` using a SYSTEM-scope acquire + atomic load (atomic_add of 0). This bypasses the MI300X L2 staleness that + makes a naive volatile-load spin never observe a peer's ack + (verified: ~sub-ms wait instead of a hang). The pure-Triton equivalent + of mori's int32_wait_until_equals.""" + peer = _rs_peer_addr(ptr, pe).to(tl.pointer_type(tl.int32), bitcast=True) + got = expected - 1 + while got != expected: + got = tl.atomic_add(peer, 0, sem="acquire", scope="sys") + + @triton.jit + def quiet(): + # Single-node: bulk transfer is host-side copy_ and the same-node + # signalling kernels never call quiet(). No-op shim so the module + # imports; never on the single-node hot path. + pass + + @triton.jit + def putmem_nbi_block(dest, source, nbytes, pe): + # Single-node bulk moves are host-side copy_; this shim exists only so + # the (never-launched-on-single-node) cross-node kernels compile. + pass + + @triton.jit + def getmem_nbi_block(dest, source, nbytes, pe): + pass + + @triton.jit + def int_p_remote(dest, value, pe): + # Cross-node is unsupported in the single-node host-API backend. + pass + + @triton.jit + def int_wait_until_equals_remote(scratch_ptr, src_ptr, expected, pe): + pass + + @triton.jit + def int_atomic_compare_swap(dest, cond, value, pe): + peer = _rs_peer_addr(dest, pe).to(tl.pointer_type(tl.int32), bitcast=True) + return tl.atomic_cas(peer, cond, value, sem="acq_rel", scope="sys") + + @triton.jit + def int_atomic_swap(dest, value, pe): + peer = _rs_peer_addr(dest, pe).to(tl.pointer_type(tl.int32), bitcast=True) + return tl.atomic_xchg(peer, value, sem="acq_rel", scope="sys") + + @core.extern + def tid(axis: core.constexpr, _semantic=None): + # Platform-portable: return constant 0 (see mori branch rationale). + del axis + return core.full((), 0, dtype=tl.int32, _semantic=_semantic) + + @core.extern + def __syncthreads(_semantic=None): + return tl.tensor(_semantic.builder.create_barrier(), tl.void) + + +# =========================================================================== +# Branch B: MORI backend (default) — re-export MORI device API. +# =========================================================================== +else: + import mori.shmem as _ms # noqa: F401 (used by users importing this module) + from mori.ir import triton as _mt + from mori.ir.triton import get_extern_libs as _mori_get_extern_libs + from mori.ir.triton import install_hook as _mori_install_hook + + # Install MORI's Triton compile hook. After install_hook(), Triton + # compilation of any kernel that uses MORI device functions will be linked + # against libmori_shmem_device.bc. + _mori_install_hook() + + # SHMEM_EXTERN_LIBS is a dict {LIB_NAME: LIB_SHMEM_PATH}; MORI returns the + # same shape from ``get_extern_libs()``. + SHMEM_EXTERN_LIBS = _mori_get_extern_libs() + LIB_NAME = "libmori_shmem" + # LIB_SHMEM_PATH kept for backward compatibility with code that probes the + # dict; on ROCm this is the path to the JIT-compiled MORI bitcode. + if isinstance(SHMEM_EXTERN_LIBS, dict) and SHMEM_EXTERN_LIBS: + LIB_SHMEM_PATH = next(iter(SHMEM_EXTERN_LIBS.values())) + else: + LIB_SHMEM_PATH = "" + + # ----------------------------------------------------------------------- + # Device API: thin @triton.jit shims that fold the extra ``qp`` argument. + # ----------------------------------------------------------------------- + # NOTE: MORI device functions take an extra qp index. ODC's signatures pass + # (dest, src, nbytes, pe). We hard-code qp=0 for now (single-QP is the + # default); if/when ODC wants multi-QP, change here. + + @triton.jit + def putmem_nbi_block(dest, source, nbytes, pe): + return _mt.putmem_nbi_block(dest, source, nbytes, pe, 0) + + @triton.jit + def getmem_nbi_block(dest, source, nbytes, pe): + return _mt.getmem_nbi_block(dest, source, nbytes, pe, 0) + + @triton.jit + def quiet(): + return _mt.quiet_thread() + + @triton.jit + def int_p(dest, value, pe): + """Write a 32-bit int to ``dest`` on rank ``pe`` from this rank. + + SAME-NODE ONLY. + ------------------------------------------------------------------ + Resolve the peer pointer via ``ptr_p2p`` (XGMI direct mapping) + and write through it with ``tl.store``. This is the recommended + same-node path in MORI's own examples. + + IMPORTANT (regression guard): this primitive is compiled into the + SAME-NODE request kernel that runs on every single-node iteration. + It must NOT reference any RDMA device function (``int32_p`` etc.). + Linking an RDMA primitive into this kernel activates MORI's IBGDA + path, which on a host without a usable RoCE/IB NIC silently breaks + the same-node store (empirically: ODC grad reduce produced zero + gradients on ~half the iters). Cross-node puts live in the separate + ``int_p_remote`` below, which is only ever compiled into the + cross-node-only kernels (never launched on a single node). + + Known limitation (investigated empirically on MI300X + ROCm 7.2) + --------------------------------------------------------------- + ROCm's XGMI cache coherence between the writer GPU and the + reader GPU's CPU-side D2H copy is not strict: a P2P-written + int32 can take O(1 second) before the destination GPU's + ``cudaMemcpy(D2H)`` returns the new value. + """ + my_pe = _mt.my_pe() + if pe == my_pe: + tl.store(dest.to(tl.pointer_type(tl.int32)), value) + else: + peer_raw = _mt.ptr_p2p(dest.to(tl.uint64, bitcast=True), my_pe, pe) + peer_ptr = peer_raw.to(tl.pointer_type(tl.int32), bitcast=True) + tl.store(peer_ptr, value) + + @triton.jit + def int_p_remote(dest, value, pe): + """Cross-node single-int put: write ``value`` to the symmetric address + ``dest`` on the remote PE ``pe`` over RDMA. + + Uses MORI's ``int32_p`` which the device API dispatches onto the RDMA + (IBGDA) transport for a cross-node PE (and P2P for a same-node PE, so + it is also safe if called with a same-node peer). ``dest`` is the + LOCAL symmetric address; MORI resolves the remote PE's matching + symmetric offset internally (qp=0). + + This primitive is intentionally SEPARATE from ``int_p`` so that the + RDMA device symbol is only ever linked into the cross-node-only + kernels (``shmem_*_remote_node_kernel`` / ``shmem_cross_node_*``), + which are never compiled/launched on a single node. This keeps the + single-node code path unchanged. + """ + _mt.int32_p(dest.to(tl.uint64, bitcast=True), value, pe, 0) + + @triton.jit + def int_g(src, pe): + """Read a 32-bit int from ``src`` on rank ``pe`` to this rank (single read). + + SAME-NODE ONLY (same regression-guard rationale as ``int_p``). + ------------------------------------------------------------------ + Resolve the peer pointer via ``ptr_p2p`` (XGMI direct mapping) and + read with ``tl.load``. NOTE: a *single* read is fine, but do NOT + use this inside a hot in-kernel spin loop — on ROCm the GPU L2 + caches the peer address and a repeated volatile load never + observes a cross-process update. Use ``int_wait_until_equals`` + (same node) or ``int_wait_until_equals_remote`` (cross node) for + spin-waiting instead. + """ + my_pe = _mt.my_pe() + if pe == my_pe: + return tl.load(src.to(tl.pointer_type(tl.int32)), volatile=True) + peer_raw = _mt.ptr_p2p(src.to(tl.uint64, bitcast=True), my_pe, pe) + peer_ptr = peer_raw.to(tl.pointer_type(tl.int32), bitcast=True) + return tl.load(peer_ptr, volatile=True) + + @triton.jit + def int_wait_until_equals(ptr, expected, pe): + """Block until the int32 at ``ptr`` on rank ``pe`` equals ``expected``. + + This is the correct replacement for a naive + ``while r != expected: quiet(); r = int_g(...)`` + spin loop on ROCm. + + Root cause it fixes + ------------------- + A naive spin ``while ...: tl.load(volatile=True)`` does not observe + cross-process updates on MI300X: the GPU L2 caches the peer address + and a long-lived in-kernel spin never sees a cross-process peer's + ack write, causing the scatter_accumulate "60s hang". + + MORI's ``int32_wait_until_equals`` is implemented as + ``while (AtomicLoadRelaxedSystem(addr) != val) {}`` + i.e. a SYSTEM-SCOPE atomic load that bypasses/refreshes the L2, so + the cross-process / cross-GPU write becomes visible. Verified on + MI300X: replacing the volatile spin with this turns a hang into a + ~0.19 ms wait. We pass the PEER's P2P-resolved address so the + same-node read path is exercised. + """ + my_pe = _mt.my_pe() + if pe == my_pe: + target = ptr.to(tl.uint64, bitcast=True) + else: + target = _mt.ptr_p2p(ptr.to(tl.uint64, bitcast=True), my_pe, pe) + _mt.int32_wait_until_equals(target, expected) + + @triton.jit + def int_wait_until_equals_remote(scratch_ptr, src_ptr, expected, pe): + """Cross-node spin-wait: block until the int32 at ``src_ptr`` on rank + ``pe`` equals ``expected``. + + Why a separate primitive (vs ``int_wait_until_equals``) + ------------------------------------------------------ + MORI's ``int32_wait_until_equals(addr, val)`` polls a *local* + address (``while AtomicLoadRelaxedSystem(addr) != val``). The + same-node ``int_wait_until_equals`` makes that work by passing the + peer's XGMI-mapped address (the peer's memory is visible in this + GPU's address space). A cross-node peer has **no** such mapping — + ``ptr_p2p`` returns 0 and there is no local address that aliases + the remote slot — so we must actively pull the value over RDMA. + + Implementation + -------------- + Repeatedly RDMA-``get`` the 4-byte remote slot into a local + symmetric scratch and compare. ``getmem`` carries no atomic-alignment + constraint (unlike a remote atomic), so any int32 slot works. + + Args: + scratch_ptr: a LOCAL symmetric/registered int32 scratch slot + (RDMA needs the get destination to be a registered MR; + MORI symmetric-heap tensors qualify). + src_ptr: the LOCAL symmetric address of the slot to watch; the + value is fetched from the same symmetric offset on ``pe``. + expected: the int32 value to wait for. + pe: the remote PE that owns the authoritative slot. + """ + got = expected - 1 + while got != expected: + _mt.getmem_nbi_block( + scratch_ptr.to(tl.uint64, bitcast=True), + src_ptr.to(tl.uint64, bitcast=True), + 4, + pe, + 0, + ) + _mt.quiet_thread() + got = tl.load(scratch_ptr.to(tl.pointer_type(tl.int32)), volatile=True) + + @triton.jit + def int_atomic_compare_swap(dest, cond, value, pe): + # MORI: atomic_uint32_fetch_thread(dest, val, cmp, op, pe, qp) + # CAS op code = 1 (per mori/src/shmem/atomic_ops.hpp; if different, + # adjust). Only used by ODC tests, not by the production gather/ + # scatter path, so a slightly-incorrect op code would not affect + # the FSDP integration. + # We expose the same call shape; the op code is hard-coded as the + # MORI CAS-fetch primitive. + # TODO: cross-check op code with mori headers when wiring tests. + return _mt.atomic_uint32_fetch_thread(dest, value, cond, 1, pe, 0) + + @triton.jit + def int_atomic_swap(dest, value, pe): + # MORI atomic SWAP op code = 0 (placeholder; not used in production). + return _mt.atomic_uint32_fetch_thread(dest, value, 0, 0, pe, 0) + + # ----------------------------------------------------------------------- + # tid / __syncthreads — platform-portable shims + # ----------------------------------------------------------------------- + # ODC uses tid(axis=0) only to elect "first thread does the atomic_add" + # in sync_cta. The Triton-native way to do this is a mask pattern: + # offsets = tl.arange(0, 1) + # tl.atomic_add(signal_ptr + offsets, 1, mask=offsets == 0) + # rather than ``if tidx == 0: tl.atomic_add(signal_ptr, 1)``. + # We patch sync_cta in utils.py to use the mask form; here, ``tid`` + # returns a constant 0 so that ``if tidx == 0`` branches taken in + # ODC's existing kernels (request_accumulation_same_node_kernel and + # wait_accumulation_same_node_kernel) still execute as before. In a + # ``@triton.jit`` kernel, all lanes follow the same control-flow path + # at the IR level, so ``if tid == 0`` becomes a conditional that + # produces the same result on every lane — semantically equivalent to + # "everybody does it" plus relying on the operation being idempotent + # under repetition. For ``int_p`` (memory store) this is fine because + # we issue the same write from every lane and MORI's underlying + # ``mori_shmem_int32_p`` is a wave-collective store (single transaction + # per wave). + @core.extern + def tid(axis: core.constexpr, _semantic=None): + # axis is a Triton constexpr (0/1/2). Return constant 0 — all lanes + # are treated identically. See note above for correctness rationale. + # We return a scalar int32 value. + del axis # unused on ROCm path + return core.full((), 0, dtype=tl.int32, _semantic=_semantic) + + @core.extern + def __syncthreads(_semantic=None): + # Platform-portable: Triton IR barrier. Lowers to s_barrier on AMDGPU. + return tl.tensor(_semantic.builder.create_barrier(), tl.void) diff --git a/primus/core/odc/primitives/utils.py b/primus/core/odc/primitives/utils.py new file mode 100644 index 000000000..fed1ebee2 --- /dev/null +++ b/primus/core/odc/primitives/utils.py @@ -0,0 +1,488 @@ +# Adapted from ODC (https://github.com/sail-sg/odc), which is distributed under +# the MIT License per its package metadata (pyproject.toml / setup.py +# classifiers). The upstream repository ships no LICENSE file or per-file +# copyright headers; upstream copyright is held by the ODC authors (Sea AI Lab). +# +# Modifications Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# See LICENSE for license information. + +""" +ODC primitives utilities — ROCm symmetric-memory management. + +ODC allocates symmetric memory and enumerates same-node peer views. On ROCm +there are two backends, selected by ``ODC_P2P_BACKEND``: + +* ``mori`` (default): + - ``init_shmem`` initializes MORI-SHMEM from the PyTorch process + group via ``mori.shmem.shmem_torch_process_group_init("default")``. + - ``shmem_create_tensor`` calls ``mori.shmem.mori_shmem_create_tensor``. + - Same-node peer views are obtained in a single MORI call: + ``mori_shmem_create_tensor_list_intra_node`` which returns the + list of symmetric views across all same-node ranks. Because that + API allocates *and* returns the peer-view list together, we + structure ``SymmBufferRegistry.allocate_symm_buffer`` to use it, + keeping the same external contract (a single tensor handed back + to the caller, plus a separately stored peer list). + - ``shmem_free_tensor_sync`` and ``finalize_distributed`` map to + ``mori.shmem.mori_shmem_free_tensor`` / ``shmem_finalize``. +* ``rocshmem``: a host-API rocSHMEM backend (single-node / IPC-only) + implemented in ``_rocshmem_backend.py``. + +NOTE on environment variables (mori backend): + Set ``MORI_SHMEM_HEAP_SIZE`` (e.g. ``"4G"``) *before* the first MORI + call. Also set ``MORI_SOCKET_IFNAME`` (defaults to the value of + ``NCCL_SOCKET_IFNAME`` if available) for the MORI bootstrap. +""" + +import logging +import os +from functools import reduce +from typing import List + +import torch +import triton +import triton.language as tl + +from odc.primitives import __syncthreads +from odc.runtime_config import get_config + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# ODC P2P backend selector (config item odc_p2p_backend). Default "mori" +# preserves the existing, verified behaviour byte-for-byte; "rocshmem" selects +# the host-API backend (single-node / IPC-only) implemented in +# _rocshmem_backend.py. Read HERE at import time -- the ODC integration patch +# populates the runtime config before odc.primitives is first imported. +# --------------------------------------------------------------------------- +_P2P_BACKEND = get_config().p2p_backend.lower() +_USE_ROCSHMEM = _P2P_BACKEND == "rocshmem" + + +# --------------------------------------------------------------------------- +# Backend-specific imports +# --------------------------------------------------------------------------- +if _USE_ROCSHMEM: + from . import _rocshmem_backend as _rs +else: + import mori.shmem as _mori_shmem + + +# --------------------------------------------------------------------------- +# init_shmem — initialize the symmetric-memory runtime. +# --------------------------------------------------------------------------- +def init_shmem(): + """Initialize the symmetric-memory runtime on the global process group. + + mori backend: + Bootstrap MORI-SHMEM from PyTorch's WORLD process group by name. + Requires that PyTorch distributed is already initialized. + ``MORI_SHMEM_HEAP_SIZE`` and ``MORI_SOCKET_IFNAME`` should be + exported in the environment before this call. + + rocshmem backend: + Bootstrap the rocSHMEM host-API runtime via a unique-id broadcast. + """ + assert torch.distributed.is_initialized() + + if _USE_ROCSHMEM: + # rocSHMEM host-API backend: bootstrap via unique-id broadcast. + _rs.init() + return + + # MORI requires MORI_SHMEM_HEAP_SIZE to be set. Inherit a sensible + # default if the user did not export one — this is purely a guard; + # production paths should always set it explicitly. + os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "4G") + + # MORI uses its own TCP bootstrap socket. If the user set + # NCCL_SOCKET_IFNAME but not MORI_SOCKET_IFNAME, mirror it so they + # stay aligned. + if "MORI_SOCKET_IFNAME" not in os.environ and "NCCL_SOCKET_IFNAME" in os.environ: + os.environ["MORI_SOCKET_IFNAME"] = os.environ["NCCL_SOCKET_IFNAME"].lstrip("=") + + # Two init methods (select via the odc_mori_init config item): + # "pg" (default): shmem_torch_process_group_init("default") + # "uid": shmem_init_attr(WITH_UNIQUEID, rank, world, uid) — does NOT + # touch PyTorch process-group registration. Use this when the + # host framework (e.g. Megatron) has already created many + # named process groups and the PG-based init crashes + # (observed: `free(): invalid pointer` during + # shmem_torch_process_group_init under Megatron). + init_method = get_config().mori_init + logger.info( + "init_shmem (MORI): heap=%s, sock_ifname=%s, method=%s", + os.environ.get("MORI_SHMEM_HEAP_SIZE"), + os.environ.get("MORI_SOCKET_IFNAME", ""), + init_method, + ) + + if init_method == "uid": + rank_id = torch.distributed.get_rank() + num_ranks = torch.distributed.get_world_size() + # rank 0 generates the 128-byte unique id; broadcast to all ranks. + uid_holder = [_mori_shmem.shmem_get_unique_id() if rank_id == 0 else None] + torch.distributed.broadcast_object_list(uid_holder, src=0) + torch.distributed.barrier() + unique_id = uid_holder[0] + _mori_shmem.shmem_init_attr( + _mori_shmem.MORI_SHMEM_INIT_WITH_UNIQUEID, + rank_id, + num_ranks, + unique_id, + ) + return + + # method == "pg": register WORLD as "default", then PG-based init. + world_group = torch.distributed.group.WORLD + try: + torch._C._distributed_c10d._register_process_group("default", world_group) + except RuntimeError: + # Already registered — ignore. + pass + _mori_shmem.shmem_torch_process_group_init("default") + + +# --------------------------------------------------------------------------- +# Symmetric tensor creation / peer-view enumeration / free +# --------------------------------------------------------------------------- +def shmem_create_tensor(shape, dtype) -> torch.Tensor: + """Allocate a symmetric tensor on the symmetric heap. + + The full allocate + same-node peer enumeration is done in one call inside + ``SymmBufferRegistry.allocate_symm_buffer`` (see below). This top-level + helper is kept for callers that want a bare local tensor without the + peer list (e.g. when only the local view is needed). On the mori backend + it calls ``mori_shmem_create_tensor``, which allocates from MORI's + symmetric heap. + """ + torch.cuda.synchronize() + if _USE_ROCSHMEM: + tensor = _rs.create_tensor(shape, dtype) + else: + tensor = _mori_shmem.mori_shmem_create_tensor(shape, dtype) + torch.cuda.synchronize() + return tensor + + +def get_same_node_tensors(tensor, rank, local_world_size) -> List[torch.Tensor]: + """Return the list of peer-view tensors for the given symmetric ``tensor`` + across the same-node ranks (``rank // local_world_size``'s peers). + + This API is only safe when ``tensor`` was created by the paired + ``mori_shmem_create_tensor_list_intra_node`` (see + ``SymmBufferRegistry.allocate_symm_buffer``); the list is retrieved + from the registry. As a stand-alone fall-back, we synthesize it via + ``shmem_ptr_p2p`` and wrap raw device pointers with + ``torch.frombuffer``-style construction. + """ + # When the tensor was created via the registry path, the registry + # already stores the peer list — callers should prefer + # ``SymmBufferRegistry.get_peer_tensors``. + # As a stand-alone fallback for ad-hoc symmetric tensors, build + # peer views via shmem_ptr_p2p + torch.Tensor reconstruction. + peer_tensors: List[torch.Tensor] = [] + local_rank = rank % local_world_size + rank_on_same_node_start = rank - local_rank + for peer in range(rank_on_same_node_start, rank_on_same_node_start + local_world_size): + if peer == rank: + peer_tensors.append(tensor) + continue + peer_ptr = _mori_shmem.shmem_ptr_p2p(tensor.data_ptr(), rank, peer) + if peer_ptr == 0: + raise RuntimeError( + f"shmem_ptr_p2p returned 0 for peer {peer}: no XGMI P2P route " + "from PE {rank}. This usually means peer is on a different node " + "(requires RDMA, not P2P)." + ) + # Build a torch.Tensor view at peer_ptr. + peer_tensors.append(_tensor_from_raw_ptr(peer_ptr, tensor.shape, tensor.dtype, tensor.device)) + return peer_tensors + + +def _tensor_from_raw_ptr(raw_ptr: int, shape, dtype, device) -> torch.Tensor: + """Wrap a raw device pointer as a torch.Tensor without taking ownership. + + Used on ROCm to materialize peer-view tensors from + ``mori.shmem.shmem_ptr_p2p``. The returned tensor shares storage with + the original symmetric allocation (which is owned by the registry); + no deleter is attached, so the caller must guarantee that the symm + allocation outlives all peer views. + """ + import ctypes + + # Compute total bytes + if isinstance(shape, int): + shape = (shape,) + nbytes = reduce(lambda a, b: a * b, shape) * torch._utils._element_size(dtype) + + # Build a ctypes array view at raw_ptr (zero-copy) — only used to + # pass through PyTorch's from_blob. PyTorch will create a tensor that + # references this storage. The CObject is intentionally kept alive + # via attribute attachment so GC does not free the wrapper. + array_t = (ctypes.c_uint8 * nbytes).from_address(raw_ptr) + t = torch.frombuffer(array_t, dtype=torch.uint8, count=nbytes).view(dtype).view(*shape).to(device) + # NOTE: ``torch.frombuffer`` materializes a CPU tensor; ``.to(device)`` would + # COPY. We can't use that path for a real peer-view (would defeat the purpose). + # Instead use the explicit cuda path: + raise NotImplementedError( + "Direct raw-ptr → torch.Tensor wrapping is not implemented in this " + "iteration. Use SymmBufferRegistry.allocate_symm_buffer (which uses " + "mori_shmem_create_tensor_list_intra_node) — that path is the one " + "exercised by ODC's gather/scatter primitives. The bare " + "get_same_node_tensors fall-back is provided only as a placeholder." + ) + + +def shmem_free_tensor_sync(tensor): + torch.cuda.synchronize() + if _USE_ROCSHMEM: + _rs.free_tensor(tensor) + else: + _mori_shmem.mori_shmem_free_tensor(tensor) + torch.cuda.synchronize() + + +def finalize_distributed(): + if _USE_ROCSHMEM: + _rs.finalize() + else: + _mori_shmem.shmem_finalize() + + +# --------------------------------------------------------------------------- +# Symmetric buffer registry (the heart of ODC memory management) +# --------------------------------------------------------------------------- +class SymmBufferRegistry: + def __init__(self): + self.local_tensor = {} + self.local_tensor_to_keys = {} + self.updated = set() + self.peer_tensors = {} + self.allocations = [] + + @classmethod + def get_instance(cls): + if not hasattr(cls, "_instance"): + cls._instance = SymmBufferRegistry() + return cls._instance + + # we'll mark all symm buffer as dirty, and next update_symm_buffer will copy the data to the symm buffer + def flush(self): + self.updated.clear() + + def update_symm_buffer(self, buffer_key, values): + values = values.contiguous() + if buffer_key not in self.local_tensor: + self.allocate_symm_buffer(buffer_key, values.shape, values.dtype) + + if buffer_key not in self.updated: + self.updated.add(buffer_key) + self.local_tensor[buffer_key].copy_(values) + # Make sure updated buffer is visible to all ranks + torch.distributed.barrier() + return self.local_tensor[buffer_key] + + @classmethod + def set_shmem_flag(cls, tensor): + tensor._odc_is_shmem = True + + @classmethod + def is_shmem_tensor(cls, tensor): + return hasattr(tensor, "_odc_is_shmem") and tensor._odc_is_shmem + + def allocate_symm_buffer(self, key, shape, dtype): + """Allocate a symmetric tensor and record same-node peer views. + + mori backend: one-shot via ``mori_shmem_create_tensor_list_intra_node``, + which returns ``list[torch.Tensor]`` of length ``local_world_size``, + each entry being the same symmetric allocation viewed from the + corresponding peer's address space. ``tensors[my_local_rank]`` is + the local view; the entire list is recorded as the peer-view list. + + rocshmem backend: one-shot ``_rs.alloc_peer_tensors`` (equivalent). + """ + assert key not in self.local_tensor + local_world_size = get_local_world_size() + rank = torch.distributed.get_rank() + local_rank = rank % local_world_size + + if _USE_ROCSHMEM: + # rocSHMEM host-API: one-shot symmetric alloc + same-node peer-view + # list (equivalent to MORI's mori_shmem_create_tensor_list_intra_node). + tensor, peer_tensors = _rs.alloc_peer_tensors(shape, dtype, local_world_size, rank) + self.allocations.append(tensor) + else: + # node_start == 0 covers BOTH the single-node case and node 0 of a + # multi-node run. In that case the same-node PEs are exactly + # [0, local_world_size), so we call MORI's one-shot helper VERBATIM + # — byte-for-byte the original single-node behaviour (single-node + # protected, zero change). + # + # node_start > 0 (only reachable on node>0 of a multi-node run): + # MORI's helper hardcodes GLOBAL PEs range(lws) == [0, lws), which + # are cross-node from this node (ptr_p2p -> 0 -> bogus tensor at + # address 0, also corrupting the "local" view). We instead allocate + # the symmetric tensor once and enumerate THIS node's real PE range + # [node_start, node_start + lws). + node_start = rank - local_rank + if node_start == 0: + peer_tensors = _mori_shmem.mori_shmem_create_tensor_list_intra_node( + shape, dtype, local_world_size + ) + assert len(peer_tensors) == local_world_size, ( + f"mori_shmem_create_tensor_list_intra_node returned " + f"{len(peer_tensors)} tensors, expected {local_world_size}" + ) + tensor = peer_tensors[local_rank] + else: + base_tensor = _mori_shmem.mori_shmem_create_tensor(shape, dtype) + peer_tensors = [ + _mori_shmem.symm_mori_shmem_tensor(base_tensor, node_start + i) + for i in range(local_world_size) + ] + assert len(peer_tensors) == local_world_size, ( + f"expected {local_world_size} same-node peer views, " f"got {len(peer_tensors)}" + ) + # peer_tensors[local_rank] resolves to peer == my_pe, i.e. the + # true local allocation (base_tensor) — right object to free. + tensor = peer_tensors[local_rank] + # Keep the same allocation list to track for finalize. + self.allocations.append(tensor) + + self.set_shmem_flag(tensor) + assert len(peer_tensors) == local_world_size + + # ranks inside the same node must be contiguous. + self.local_tensor[key] = tensor + self.peer_tensors[key] = peer_tensors + self.local_tensor_to_keys[self.local_tensor[key].data_ptr()] = key + logger.info( + f"Rank {rank} create tensor {key} with shape {shape} and dtype {dtype} " + f"and ptr {self.local_tensor[key].data_ptr():#x}" + ) + return tensor + + def has_key(self, key): + return key in self.local_tensor + + def get_or_create_symm_buffer(self, key, shape, dtype): + if self.has_key(key): + return self.local_tensor[key] + return self.allocate_symm_buffer(key, shape, dtype) + + def get_symm_buffer(self, key): + if self.has_key(key): + return self.local_tensor[key] + raise ValueError(f"Symm buffer {key} not found") + + def get_peer_tensors(self, local_tensor): + # Returns tensors in the same node. + buffer_key = self.local_tensor_to_keys[local_tensor.data_ptr()] + return self.peer_tensors[buffer_key] + + def finalize(self): + for t in self.allocations: + shmem_free_tensor_sync(t) + self.local_tensor.clear() + self.local_tensor_to_keys.clear() + self.updated.clear() + self.peer_tensors.clear() + + def memory_allocated(self): + return sum(t.nbytes for t in self.allocations) + + +# --------------------------------------------------------------------------- +# Process group helpers — unchanged from upstream +# --------------------------------------------------------------------------- +local_world_pg = None + + +def get_local_world_pg(pg: torch.distributed.ProcessGroup): + local_world_size = get_local_world_size() + assert torch.distributed.get_world_size() == torch.distributed.get_world_size( + group=pg + ), "Cached AG only supports pure data parallelism" + rank = torch.distributed.get_rank() + global local_world_pg + if local_world_pg is None: + for i in range(0, torch.distributed.get_world_size(), local_world_size): + ranks = list(range(i, i + local_world_size)) + new_gp = torch.distributed.new_group(ranks=ranks, backend="nccl") + if rank in ranks: + local_world_pg = new_gp + assert local_world_pg is not None + return local_world_pg + + +def get_local_world_size(): + if "RAY_LOCAL_WORLD_SIZE" in os.environ: + return int(os.environ["RAY_LOCAL_WORLD_SIZE"]) + else: + return int(os.environ["LOCAL_WORLD_SIZE"]) + + +stream = None + + +def get_comm_stream(): + global stream + if stream is None: + stream = torch.cuda.Stream() + return stream + + +class BufferSplitter: + def __init__(self): + self.round_data_size = 2**6 + + def get_max_global_buffer_size(self): + # config item odc_max_buffer_size (default 64 MiB). + max_buffer_size = int(get_config().max_buffer_size) + return max_buffer_size + + def get_global_buffer_size(self, original_buffer_shape): + original_size = reduce(lambda x, y: x * y, original_buffer_shape) + max_buffer_size = self.get_max_global_buffer_size() + if max_buffer_size <= 0: + return original_size + buf_size = min(max_buffer_size, original_size) + return buf_size + + def round_to_data_size(self, size): + return (size + self.round_data_size - 1) // self.round_data_size * self.round_data_size + + def get_local_buffer_size(self, original_buffer_shape, world_size): + original_size = reduce(lambda x, y: x * y, original_buffer_shape) + max_buffer_size = self.get_max_global_buffer_size() + if max_buffer_size <= 0: + return self.round_to_data_size(original_size) + assert ( + max_buffer_size % world_size == 0 + ), f"ODC_MAX_BUFFER_SIZE: {max_buffer_size} % world_size: {world_size} != 0" + local_max_buffer_size = max_buffer_size // world_size + buf_size = min(local_max_buffer_size, original_size) + return self.round_to_data_size(buf_size) + + +# --------------------------------------------------------------------------- +# sync_cta — CTA-internal synchronization helper +# --------------------------------------------------------------------------- +# Elect a single thread (lane 0) to do the atomic_add using the Triton vector +# mask pattern (no NVVM tid intrinsic exists on AMDGPU). +@triton.jit +def sync_cta(signal_ptr, expected): + # Use a length-1 lane vector with mask to ensure only lane 0 of + # the wave performs the atomic_add. This is the Triton-native way + # to elect a "first thread". + offsets = tl.arange(0, 1) + mask = offsets == 0 + tl.atomic_add(signal_ptr + offsets, 1, mask=mask) + __syncthreads() + r = 0 + while r < expected: + signals = tl.load(signal_ptr + offsets, mask=mask, volatile=True) + r = tl.max(signals) diff --git a/primus/core/odc/rocshmem_runtime/README.md b/primus/core/odc/rocshmem_runtime/README.md new file mode 100644 index 000000000..98e65f248 --- /dev/null +++ b/primus/core/odc/rocshmem_runtime/README.md @@ -0,0 +1,58 @@ +# rocshmem_runtime — ODC runtime helpers (rocSHMEM ops now in Primus-Turbo) + +The ODC rocSHMEM host/GDA ops were migrated to **Primus-Turbo** +(`primus_turbo.pytorch._C.odc_rocshmem_host` / `odc_rocshmem_gda`) and are +consumed by `odc/primitives/_rocshmem_backend.py`. The former in-tree ctypes +bindings (`host_bindings/rs_host.cpp`, `gda_backend/rs_host_gda.cpp`) and the +`build_rocshmem_backend.sh` build script have been **removed**. + +The rocSHMEM **static library** (`librocshmem.a` + headers) that Primus-Turbo +links against is an **external dependency**: point the Turbo build's +`ROCSHMEM_HOME` at a rocSHMEM install. Any local rocSHMEM build tree +(`rocshmem_src/`, `rocshmem_{single,gda}/`, `*.a`, `*.so`) is `.gitignore`d and +must never be committed. + +> TODO(primus-turbo): pin the exact Primus-Turbo merge commit +> (`PRIMUS_TURBO_COMMIT`) once the ODC-ops PR merges. + +This directory now only holds runtime helpers: + +- `scripts/run_odc.sh` — single-node ODC launcher (plain torchrun). Default P2P + backend is `mori`; pass `rocshmem` to use the Turbo-provided ops. Set + `PRIMUS_TURBO_PATH` to a Primus-Turbo build tree if it is not already importable. +- `scripts/cleanup_rs.sh` — leftover-process cleanup. + +## Selecting the rocSHMEM backend + +These are Primus config items (read by the ODC library via its runtime config): + +- `odc_p2p_backend: rocshmem` — use the rocSHMEM ops (default is `mori`). +- `odc_rocshmem_gda: true` — multi-node GPU-direct (GDA) path; otherwise + single-node XGMI IPC host path. + +The ops resolve from `primus_turbo.pytorch._C`. As an escape hatch, +`odc_rocshmem_lib: /librs_host_gda.so` loads an external monolithic ctypes +binding instead of the Turbo submodule (see `_rocshmem_backend.py`). + +## Multi-node (GDA) correctness & deployment env + +**Correctness (automatic, no env needed).** Multi-node (`n_pes > GPUs-per-node`) +defaults to `ODC_GDA_DEFER_REDUCE=1`: each micro-batch's unsharded grad is +accumulated locally, then ONE barriered reduce-scatter runs per minibatch. This +keeps the collective-barrier count equal across ranks, so variable-length +(`nopad`) packing does not deadlock. Single-node defaults to `0`. + +**Cluster deployment env** (set per your cluster; not hardcoded in the repo): + +| env | example | notes | +|-----|---------|-------| +| `ROCSHMEM_HCA_LIST` | `mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9` | this cluster's compute NICs (else all traffic funnels through `mlx5_0`) | +| `ROCSHMEM_GDA_PROVIDER` | `mlx5` | RDMA provider | +| `ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME` | `eth0` | uid-over-socket bootstrap NIC | +| `ROCSHMEM_HEAP_SIZE` | `8589934592` | symmetric heap RAW bytes (decimal only, no K/M/G) | +| `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME` | `eth0` | torch.distributed control plane | +| `NCCL_IB_GID_INDEX` | `3` | RoCE v2 GID (fabric-dependent) | +| `ODC_GDA_WARMUP_MODE` / `ODC_GDA_STRIDE_BYTES` | `strided` / `65536` | GDA connection warmup | + +The GDA backend bootstraps rocSHMEM with a unique-id over a TCP socket, so the +job launches with **plain torchrun** — no MPI / mpirun. diff --git a/primus/core/odc/rocshmem_runtime/scripts/cleanup_rs.sh b/primus/core/odc/rocshmem_runtime/scripts/cleanup_rs.sh new file mode 100755 index 000000000..8206cc2b7 --- /dev/null +++ b/primus/core/odc/rocshmem_runtime/scripts/cleanup_rs.sh @@ -0,0 +1,37 @@ +#!/bin/bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +# shellcheck disable=SC2009 # need full args column; pgrep can't match python -m patterns +# Terminate all residual training/compile processes from the abandoned rocSHMEM +# smoke in THIS container. Patterns live in this file (not in the invoking +# shell's argv), so `pkill -f` cannot match its own parent shell. +echo "=== BEFORE: residual training/compile procs ===" +ps -eo pid,etimes,comm,args 2>/dev/null \ + | grep -iE 'torchrun|primus/cli|torch.distributed|hipcc|clang|cmake|ninja|llvm|multiprocessing' \ + | grep -vE 'grep|cleanup_rs.sh' || echo "(none before)" + +for pat in 'torchrun' 'torch.distributed.run' 'primus/cli' 'run_pretrain.sh' \ + 'run_odc.sh' 'multiprocessing.spawn' 'multiprocessing.resource_tracker' \ + 'hipcc' 'cmake' 'ninja'; do + pkill -9 -f "$pat" 2>/dev/null +done +pkill -9 -x python3 2>/dev/null +pkill -9 -x pt_main_thread 2>/dev/null +pkill -9 clang 2>/dev/null +pkill -9 'clang++' 2>/dev/null +pkill -9 llvm-as 2>/dev/null +pkill -9 lld 2>/dev/null +sleep 4 + +echo "=== AFTER: residual training/compile procs (want NONE) ===" +LEFT=$(ps -eo pid,etimes,comm,args 2>/dev/null \ + | grep -iE 'torchrun|primus/cli|torch.distributed|hipcc|clang|cmake|ninja|llvm|multiprocessing|python3' \ + | grep -vE 'grep|cleanup_rs.sh') +if [ -z "$LEFT" ]; then echo "(none)"; else echo "$LEFT"; fi + +echo "=== rocm-smi GPU use% ===" +rocm-smi --showuse 2>/dev/null | grep -E 'GPU use \(%\)' +echo "=== rocm-smi compute pids ===" +rocm-smi --showpids 2>/dev/null | sed -n '1,12p' diff --git a/primus/core/odc/rocshmem_runtime/scripts/run_odc.sh b/primus/core/odc/rocshmem_runtime/scripts/run_odc.sh new file mode 100755 index 000000000..5d0bc76a9 --- /dev/null +++ b/primus/core/odc/rocshmem_runtime/scripts/run_odc.sh @@ -0,0 +1,97 @@ +#!/bin/bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +# ============================================================================= +# run_odc.sh — portable ODC LB-mini launcher (single-node) +# +# Runs an ODC LB-mini training job. Default P2P backend is `mori`; passing +# `rocshmem` switches to the rocSHMEM backend. The rocSHMEM ops are consumed +# from Primus-Turbo (primus_turbo.pytorch._C.odc_rocshmem_host / _gda); set +# PRIMUS_TURBO_PATH to a Primus-Turbo build tree if it is not already importable +# in the environment (it is prepended to PYTHONPATH). +# +# usage: run_odc.sh [KEY=VAL ...] +# pad|nopad is retained for backwards-compatible invocation but is now a no-op +# (the aligned-vs-decoupled A/B study knob was removed; LB-Mini always runs +# decoupled when enabled). ODC feature switches live in the EXP yaml config. +# extra KEY=VAL args are exported verbatim. +# +# Overridable env (all have portable defaults): +# PRIMUS_ROOT project root (auto-derived from this script's path) +# PRIMUS_TURBO_PATH Primus-Turbo build tree to prepend to PYTHONPATH so +# `import primus_turbo` (with the ODC rocSHMEM ops) +# resolves; leave unset if already installed. +# HF_HOME HF cache dir (default: /workspace/hf_cache) +# PRIMUS_PACK_CACHE_DIR packed-sequence cache (default: $HOME/primus_packed) +# TRITON_CACHE_DIR triton cache (rocshmem: fresh per-run unless pinned) +# TRAIN_LOG_DIR where to write runlog_*.log (default: $HOME/odc_logs) +# MASTER_PORT default 29600 +# ROCSHMEM_HEAP_SIZE symmetric heap RAW BYTES (default 8 GiB) +# ============================================================================= +set -u +BACKEND=$1; PAD=$2; EXP_REL=$3; EXPNAME=$4; shift 4 + +# --- derive project root from this script's location (portable) ------------- +# scripts/ -> rocshmem_runtime/ -> odc/ -> core/ -> primus/ -> +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +RUNTIME_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" # rocshmem_runtime/ +ODC_ROOT="$(cd "$RUNTIME_DIR/.." && pwd)" # primus/core/odc/ +PRIMUS_ROOT="${PRIMUS_ROOT:-$(cd "$ODC_ROOT/../../.." && pwd)}" +cd "$PRIMUS_ROOT" || exit 1 + +export EXP=$EXP_REL +# ODC arm env. Prepend a Primus-Turbo build tree if provided so the rocSHMEM +# backend (primus_turbo.pytorch._C.odc_rocshmem_*) is importable. +# The `odc` package now lives directly at $ODC_ROOT (primus/core/odc/), so put +# its PARENT (primus/core/) on PYTHONPATH for `import odc`; odc_early holds the +# sitecustomize load-order shim. +export PYTHONPATH="${PRIMUS_TURBO_PATH:+$PRIMUS_TURBO_PATH:}$ODC_ROOT/odc_early:${ODC_ROOT%/*}" +# ODC feature switches (enable_odc, odc_phase, enable_odc_lb_mini, ...) are now +# CONFIG items set in the EXP yaml, NOT env vars. Only genuine infra env is set +# here. MORI_SHMEM_HEAP_SIZE is MORI runtime infra (symmetric heap size). +export MORI_SHMEM_HEAP_SIZE=8G +# NOTE: the $PAD positional arg is retained for backwards-compatible invocation +# but no longer toggles an aligned-vs-decoupled A/B baseline (that was a study +# knob and has been removed); LB-Mini always runs decoupled when enabled. +# public env +export HF_HOME=${HF_HOME:-/workspace/hf_cache} DATA_PATH=${DATA_PATH:-/workspace} +export GLOO_SOCKET_IFNAME=lo NCCL_SOCKET_IFNAME=lo NCCL_IB_DISABLE=1 +export FUSED_LINEAR_CE=1 +export PRIMUS_PACK_CACHE_DIR=${PRIMUS_PACK_CACHE_DIR:-$HOME/primus_packed} +export PRIMUS_EXP_NAME=$EXPNAME +export MASTER_PORT=${MASTER_PORT:-29600} +mkdir -p "$PRIMUS_PACK_CACHE_DIR" +# backend selection +if [ "$BACKEND" = "rocshmem" ]; then + export ODC_P2P_BACKEND=rocshmem + # The rocSHMEM ops are consumed from Primus-Turbo (see PRIMUS_TURBO_PATH above); + # no in-tree librs_host*.so is loaded anymore. + export ROCSHMEM_BOOTSTRAP_SOCKET_IFNAME=lo + # rocSHMEM symmetric heap size, RAW BYTES (the env parser is decimal-only and + # does NOT accept K/M/G suffixes). 8 GiB matches MORI_SHMEM_HEAP_SIZE. + export ROCSHMEM_HEAP_SIZE=${ROCSHMEM_HEAP_SIZE:-8589934592} + # IMPORTANT: the rocshmem device kernels bake per-PE peer deltas as Triton + # constexpr. Reusing a Triton cache built by a *different* toolchain (or a + # different launch) was observed to silently load mismatched kernels -> + # garbage int_p/wait signalling -> NaN grads from iter 1. Always start from a + # fresh per-run cache for rocshmem unless the caller pins TRITON_CACHE_DIR. + export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-/tmp/tcache_rocshmem_$(date +%Y%m%d_%H%M%S)_$$} +else + export ODC_P2P_BACKEND=mori + export TRITON_CACHE_DIR=${TRITON_CACHE_DIR:-/tmp/tcache_mori} +fi +# extra KEY=VAL env +# shellcheck disable=SC2163 # kv is a literal KEY=VAL token, so export works +for kv in "$@"; do export "$kv"; done + +# Unique per-run timestamp so reruns never clobber earlier logs. Override +# TRAIN_LOG_TS to pin a specific stamp (e.g. shared across multinode ranks). +TRAIN_LOG_TS=${TRAIN_LOG_TS:-$(date +%Y%m%d_%H%M%S)} +TRAIN_LOG_DIR=${TRAIN_LOG_DIR:-$HOME/odc_logs}; mkdir -p "$TRAIN_LOG_DIR" +export TRAIN_LOG="$TRAIN_LOG_DIR/runlog_${EXPNAME}_${TRAIN_LOG_TS}.log" +echo "[run_odc] ROOT=$PRIMUS_ROOT BACKEND=$BACKEND PAD=$PAD P2P=$ODC_P2P_BACKEND EXP=$EXP NAME=$EXPNAME TS=$TRAIN_LOG_TS" +echo "[run_odc] TURBO_PATH=${PRIMUS_TURBO_PATH:-} TRITON_CACHE_DIR=$TRITON_CACHE_DIR LOG=$TRAIN_LOG" +bash examples/run_pretrain.sh +echo "[run_odc] DONE exit=$? log=$TRAIN_LOG" diff --git a/primus/core/odc/runtime_config.py b/primus/core/odc/runtime_config.py new file mode 100644 index 000000000..adf9970e4 --- /dev/null +++ b/primus/core/odc/runtime_config.py @@ -0,0 +1,86 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. + +"""ODC runtime configuration. + +The ODC primitives (``odc.primitives.*``) are decoupled library code with no +knowledge of Primus' config system. Historically they read their tuning knobs +straight from ``os.environ`` (ODC_P2P_BACKEND, ODC_GDA_*, ...), which violates +the Primus rule that formal feature logic must be config-driven, not env-driven. + +This leaf module (imports only stdlib, so it can be populated BEFORE the heavy +``odc.primitives`` modules are first imported) holds a single ``OdcRuntimeConfig`` +instance. The Primus ODC integration patch reads the ``odc_*`` config items from +the trainer config and calls :func:`set_config` at ``before_train`` -- i.e. after +the Primus config is available but before ``odc.primitives`` (and their +import-time backend selection) run. The primitives then read values from here +instead of the environment. + +Defaults below are byte-for-byte the previous env defaults, so an unset config +reproduces the prior behaviour exactly. +""" + +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class OdcRuntimeConfig: + # --- symmetric-memory backend selection (was ODC_P2P_BACKEND) --- + # "mori" (default) | "rocshmem" + p2p_backend: str = "mori" + # MORI init method (was ODC_MORI_INIT): "pg" (default) | "uid" + mori_init: str = "pg" + # BufferSplitter max global buffer size in bytes (was ODC_MAX_BUFFER_SIZE). + max_buffer_size: int = 64 * 1024 * 1024 + + # --- rocSHMEM backend (was ODC_ROCSHMEM_GDA / ODC_ROCSHMEM_LIB) --- + # GPU-direct (GDA) device-kernel cross-node path; required for multi-node. + rocshmem_gda: bool = False + # Optional path to a monolithic librs_host_gda.so ctypes override; None -> + # consume the rocSHMEM ops from Primus-Turbo. + rocshmem_lib: Optional[str] = None + + # --- GDA reduce-scatter tuning knobs --- + # reduce-scatter kernel grid blocks (was ODC_GDA_RS_BLOCKS). + gda_rs_blocks: int = 64 + # peer-pipeline batch depth (was PRIMUS_TURBO_ODC_GDA_PIPE). ALSO consumed by + # the Primus-Turbo device kernel via getenv, so set_config bridges it back to + # the PRIMUS_TURBO_ODC_GDA_PIPE env var for the C++ side. + gda_pipe: int = 1 + # defer the cross-node reduce to once-per-minibatch (was ODC_GDA_DEFER_REDUCE): + # "auto" (default: on iff multi-node, i.e. n_pes > local_world_size) | "1" | "0". + gda_defer_reduce: str = "auto" + # cross-node write-visibility warm-up mode (was ODC_GDA_WARMUP_MODE): + # "strided" (default) | "full" | "hdp" | "fence" | "hdpfence". + gda_warmup_mode: str = "strided" + # strided warm-up page stride in bytes (was ODC_GDA_STRIDE_BYTES). + gda_stride_bytes: int = 65536 + + +_CONFIG = OdcRuntimeConfig() + + +def get_config() -> OdcRuntimeConfig: + """Return the process-wide ODC runtime config (populated by set_config).""" + return _CONFIG + + +def set_config(**kwargs) -> OdcRuntimeConfig: + """Update the ODC runtime config in place from keyword overrides. + + Only known fields are applied; None values are ignored (keep the default). + Returns the config instance. Also bridges gda_pipe back to the + PRIMUS_TURBO_ODC_GDA_PIPE env var, which the Primus-Turbo device kernel reads + via getenv (the one knob that must stay visible to the C++ side). + """ + import os + + valid = set(OdcRuntimeConfig.__dataclass_fields__) + for k, v in kwargs.items(): + if k in valid and v is not None: + setattr(_CONFIG, k, v) + # Bridge the pipe depth to the env var the turbo C++ kernel reads. + os.environ["PRIMUS_TURBO_ODC_GDA_PIPE"] = str(int(_CONFIG.gda_pipe)) + return _CONFIG diff --git a/pyproject.toml b/pyproject.toml index 19880d9a4..4851a3253 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,15 @@ Repository = "https://github.com/AMD-AGI/Primus" line-length = 110 target-version = ["py38"] +[tool.isort] +# Mirror the pre-commit hook (isort --profile black) so local runs and the +# pure-checkout CI agree. `odc` lives under primus/core/odc/odc/, so a clean +# checkout has no top-level `odc` package and isort would otherwise treat +# `import odc` as third-party. Declaring it first-party keeps import grouping +# deterministic and identical between local and CI. +profile = "black" +known_first_party = ["odc"] + ############################################################################### # Hatch build configuration ############################################################################### From d27016541bbf4a537b91a647f7128fd9a64347b9 Mon Sep 17 00:00:00 2001 From: HuangWei-95 Date: Mon, 20 Jul 2026 16:55:24 +0800 Subject: [PATCH 046/127] fix(config): coerce true/false env interpolation to bool in yaml loader (#886) Env substitution previously only parsed int/float, so values like ${PRIMUS_PROFILE:false} became the string "false" and were truthy in Megatron if-checks. Parse whole-string true/false case-insensitively while keeping 0/1 as integers for backward compatibility. Co-authored-by: HuangWei-95 Co-authored-by: Cursor --- primus/core/config/yaml_loader.py | 21 ++++++++----- .../core/config/test_yaml_loader.py | 31 +++++++++++++++++++ 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/primus/core/config/yaml_loader.py b/primus/core/config/yaml_loader.py index 4ec3a5722..b8b715c5c 100644 --- a/primus/core/config/yaml_loader.py +++ b/primus/core/config/yaml_loader.py @@ -12,6 +12,7 @@ ENV_PATTERN = re.compile(r"\${([^:{}]+)(?::([^}]*))?}") # Matches floats like: 1.2, .5, 1., 1e5, 1.2e-5, -1.2 FLOAT_PATTERN = re.compile(r"^-?(?:(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?|\d+[eE][+-]?\d+)$") +_BOOL_MAP = {"true": True, "false": False} def parse_yaml(path: str) -> dict: @@ -61,9 +62,10 @@ def _resolve_env_in_string(s: str): Returns ------- - str or int or float + str, bool, int, or float The resolved value. - If environment variable substitution occurs: + - If the resulting string is ``true``/``false`` (case-insensitive), it is converted to bool. - If the resulting string represents a number, it is converted to int or float. - Otherwise, returns the substituted string. - If no substitution occurs: returns the original string unchanged (even if it looks numeric). @@ -98,26 +100,31 @@ def replace_match(m): return os.environ.get(var, default) replaced = ENV_PATTERN.sub(replace_match, s) - return _try_numeric(replaced) if replaced != s else replaced + return _try_scalar(replaced) if replaced != s else replaced -def _try_numeric(v: str): +def _try_scalar(v: str): """ - Attempt to convert a string value to int or float. + Attempt to convert a string value to bool, int, or float. Returns the original string if conversion fails. """ - # 1. Integer check + # 1. Boolean check (whole-string, case-insensitive) + lowered = v.lower() + if lowered in _BOOL_MAP: + return _BOOL_MAP[lowered] + + # 2. Integer check if re.fullmatch(r"-?\d+", v): return int(v) - # 2. Float check using regex to avoid exceptions for non-numeric strings + # 3. Float check using regex to avoid exceptions for non-numeric strings if FLOAT_PATTERN.fullmatch(v): try: return float(v) except ValueError: return v - # 3. Return original string + # 4. Return original string return v diff --git a/tests/unit_tests/core/config/test_yaml_loader.py b/tests/unit_tests/core/config/test_yaml_loader.py index 038cef09a..9558de7f6 100644 --- a/tests/unit_tests/core/config/test_yaml_loader.py +++ b/tests/unit_tests/core/config/test_yaml_loader.py @@ -110,3 +110,34 @@ def test_resolve_env_in_string_helpers(self, monkeypatch): assert _resolve_env_in_string("${NUM}") == 42 assert _resolve_env_in_string("${FLOAT}") == 3.14 assert _resolve_env_in_string("prefix_${NUM}") == "prefix_42" + + def test_env_resolution_bool_strings(self, monkeypatch, tmp_path): + monkeypatch.setenv("FLAG_TRUE", "true") + monkeypatch.setenv("FLAG_FALSE", "false") + monkeypatch.setenv("FLAG_UPPER", "False") + monkeypatch.delenv("UNSET_BOOL", raising=False) + + assert _resolve_env_in_string("${FLAG_TRUE}") is True + assert _resolve_env_in_string("${FLAG_FALSE}") is False + assert _resolve_env_in_string("${FLAG_UPPER}") is False + assert _resolve_env_in_string("${UNSET_BOOL:false}") is False + assert _resolve_env_in_string("${UNSET_BOOL:true}") is True + + content = """ + native_false: false + env_false_default: ${UNSET_BOOL:false} + env_false_set: ${FLAG_FALSE} + env_true_set: ${FLAG_TRUE} + env_zero: ${UNSET_ZERO:0} + """ + monkeypatch.delenv("UNSET_ZERO", raising=False) + cfg_file = tmp_path / "bool_env.yaml" + cfg_file.write_text(content) + + result = parse_yaml(str(cfg_file)) + assert result["native_false"] is False + assert result["env_false_default"] is False + assert result["env_false_set"] is False + assert result["env_true_set"] is True + assert result["env_zero"] == 0 + assert isinstance(result["env_zero"], int) From 0e4cc728cab462867d2469bbaf1e99bbce0c7cc6 Mon Sep 17 00:00:00 2001 From: wenxie-amd Date: Mon, 20 Jul 2026 18:08:26 +0800 Subject: [PATCH 047/127] Adapt Primus launch to the spur (amd-spur) cluster (#887) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Adapt Primus launch to the spur (amd-spur) cluster ### Summary Makes `primus-cli slurm` and the DeepSeek‑V4 launch scripts work on the **spur** ("AI‑native") scheduler used by the `amd-spur` partition, including multi‑node runs that need the burst QOS and ionic (AINIC) RDMA. All spur‑specific behavior is gated behind `command -v spur`, so standard‑Slurm clusters are unaffected. ### Background / problem On spur, the existing launch path failed for multi‑node training: - **QOS**: spur's `srun` has no `--qos`; only `sbatch`/`spur submit` accept `-q`. - **sbatch arg handling**: spur's `sbatch` takes a *single* positional (the script), rejects arguments passed to the script, and has no `--wrap`. So `primus-cli slurm sbatch … -- --image …` failed with `unexpected argument '--image'`. - **Per‑node execution**: spur runs a batch script on *every* allocated node (not node‑0 only like standard Slurm), so naive nesting of an inner `srun` exploded into nested job steps (SIGKILL). - **ionic RDMA**: the AINIC image's bundled `libionic` advertises ionic uverbs **ABI 1**, but the host kernel ionic driver is **ABI 4** → `ibv` rejected all `ionic_*` devices → NCCL fell back to a loopback socket → multi‑node NCCL crashed (`Connection refused`). - **Container name collisions**: the fixed container name `primus-training` caused `name already in use` failures from orphaned containers left on reused nodes. ### Changes - **`runner/primus-cli-slurm.sh`** — detect spur via `command -v spur`. For `sbatch` on spur, emit a self‑contained per‑node batch script (entry + all args baked in via `printf %q`) and submit it with `--qos`/`--partition`/`--account` as sbatch options. `srun` and standard‑Slurm `sbatch` paths are unchanged. - **`runner/primus-cli-container.sh`** — unique per‑job container name (`primus-training-$SLURM_JOB_ID`) to avoid name conflicts from orphaned containers; make `--clean` tolerant of per‑container removal failures so a stale/racing container no longer aborts the run. - **`examples/deepseek-v4/run_deepseek_v4{,_flash}.sh`** — spur‑gated launcher block (`sbatch` + `amd-burst-qos` + `amd-primus` + `--exclusive`), NCCL/GLOO socket interface via `ens3` (instead of loopback), default image `docker.io/tasimage/primus:pr-882-ainic`, and `SBATCH_OUTPUT`/`SBATCH_ERROR` into the experiment output dir. Socket‑iface default simplified to `lo` (overridable). - **`runner/helpers/patches/10_fix_libionic_abi4.sh`** — new `--patch` script (gated by `PRIMUS_LIBIONIC_SRC_ABI4_SO`) that swaps the ionic libibverbs provider for an ABI‑4 build so ionic RDMA enumerates. Skips cleanly when the env var is unset. - Moved `rccl_avg_workaround/` under `examples/deepseek-v4/` and updated its references. - Added `skills/spur-cluster-status` (cluster status tooling). ### Validation End‑to‑end 4‑node DeepSeek‑V4 Flash run on `amd-spur` (job 15965): allocated 4 exclusive nodes under `amd-burst-qos`; the libionic patch swapped the provider (ABI‑mismatch warning gone); NCCL formed the full **32‑rank / 4‑node** communicator over ionic **RDMA (RCCL‑ANP)** — no loopback fallback, no container name conflict. Standard‑Slurm `srun`/`sbatch` output verified byte‑for‑byte unchanged via `--dry-run`. ### Compatibility / notes - Spur‑specific logic is behind `command -v spur`; non‑spur clusters see no behavior change. - The libionic swap and `--exclusive` are opt‑in / env‑gated (`PRIMUS_LIBIONIC_SRC_ABI4_SO`, `SLURM_EXCLUSIVE`). - The libionic swap is a runtime stop‑gap for AINIC images whose bundled `libionic` is ABI‑1; the durable fix is to ship an ABI‑4 provider in the image. Co-authored-by: Cursor --- .../rccl_avg_workaround}/.gitignore | 0 .../rccl_avg_workaround}/sitecustomize.py | 0 examples/deepseek-v4/run_deepseek_v4.sh | 37 +- examples/deepseek-v4/run_deepseek_v4_flash.sh | 24 +- .../run_deepseek_v4_pro_muon_1gpu.sh | 2 +- .../deepseek-v4/run_dsv4_projection_1gpu.sh | 2 +- .../helpers/patches/10_fix_libionic_abi4.sh | 47 ++ runner/primus-cli-container.sh | 13 +- runner/primus-cli-slurm.sh | 64 ++- skills/spur-cluster-status/SKILL.md | 83 +++ .../scripts/spur_status.py | 477 ++++++++++++++++++ 11 files changed, 725 insertions(+), 24 deletions(-) rename {rccl_avg_workaround => examples/deepseek-v4/rccl_avg_workaround}/.gitignore (100%) rename {rccl_avg_workaround => examples/deepseek-v4/rccl_avg_workaround}/sitecustomize.py (100%) mode change 100644 => 100755 examples/deepseek-v4/run_deepseek_v4_flash.sh create mode 100755 runner/helpers/patches/10_fix_libionic_abi4.sh create mode 100644 skills/spur-cluster-status/SKILL.md create mode 100644 skills/spur-cluster-status/scripts/spur_status.py diff --git a/rccl_avg_workaround/.gitignore b/examples/deepseek-v4/rccl_avg_workaround/.gitignore similarity index 100% rename from rccl_avg_workaround/.gitignore rename to examples/deepseek-v4/rccl_avg_workaround/.gitignore diff --git a/rccl_avg_workaround/sitecustomize.py b/examples/deepseek-v4/rccl_avg_workaround/sitecustomize.py similarity index 100% rename from rccl_avg_workaround/sitecustomize.py rename to examples/deepseek-v4/rccl_avg_workaround/sitecustomize.py diff --git a/examples/deepseek-v4/run_deepseek_v4.sh b/examples/deepseek-v4/run_deepseek_v4.sh index 87796618f..bf9b736fd 100755 --- a/examples/deepseek-v4/run_deepseek_v4.sh +++ b/examples/deepseek-v4/run_deepseek_v4.sh @@ -24,25 +24,17 @@ export WANDB_API_KEY="${WANDB_API_KEY:-your_wandb_api_key}" export NNODES=${NNODES:-1} export TRAIN_ITERS=${TRAIN_ITERS:-20} -export DOCKER_IMAGE=${DOCKER_IMAGE:-"tasimage/primus:pr-715-ainic"} +export DOCKER_IMAGE=${DOCKER_IMAGE:-"docker.io/tasimage/primus:pr-882-ainic"} export SLURM_PARTITION=${SLURM_PARTITION:-Compute-DCPT} -export SLURM_NODELIST="${SLURM_NODELIST:-smci355-ccs-aus-n01-21,smci355-ccs-aus-n01-33,smci355-ccs-aus-n02-21,smci355-ccs-aus-n02-25,smci355-ccs-aus-n02-29,smci355-ccs-aus-n02-33,smci355-ccs-aus-n03-33,smci355-ccs-aus-n04-21,smci355-ccs-aus-n04-25,smci355-ccs-aus-n04-29,smci355-ccs-aus-n04-33,smci355-ccs-aus-n05-21,smci355-ccs-aus-n05-29,smci355-ccs-aus-n05-33,smci355-ccs-aus-n06-25,smci355-ccs-aus-n06-33,smci355-ccs-aus-n10-29}" +export SLURM_NODELIST="${SLURM_NODELIST-smci355-ccs-aus-n01-21,smci355-ccs-aus-n01-33,smci355-ccs-aus-n02-21,smci355-ccs-aus-n02-25,smci355-ccs-aus-n02-29,smci355-ccs-aus-n02-33,smci355-ccs-aus-n03-33,smci355-ccs-aus-n04-21,smci355-ccs-aus-n04-25,smci355-ccs-aus-n04-29,smci355-ccs-aus-n04-33,smci355-ccs-aus-n05-21,smci355-ccs-aus-n05-29,smci355-ccs-aus-n05-33,smci355-ccs-aus-n06-25,smci355-ccs-aus-n06-33,smci355-ccs-aus-n10-29}" export MASTER_PORT=${MASTER_PORT:-29500} export USING_AINIC=${USING_AINIC:-1} export NCCL_IB_HCA="ionic_0:1,ionic_1:1,ionic_2:1,ionic_3:1,ionic_4:1,ionic_5:1,ionic_6:1,ionic_7:1" -# "fenic" is the cluster RDMA NIC. Prefer it when present (cluster / multi-node), -# but fall back to a real local interface (lo) so single-node / direct in-container -# smoke runs work out of the box — otherwise gloo aborts with -# "Unable to find address for: fenic". Still override-guarded: pass -# GLOO_SOCKET_IFNAME / NCCL_SOCKET_IFNAME explicitly to force a specific NIC. -if [ -d /sys/class/net/fenic ]; then - _PRIMUS_DEFAULT_IFNAME=fenic -else - _PRIMUS_DEFAULT_IFNAME=lo -fi -export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-$_PRIMUS_DEFAULT_IFNAME} -export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-$_PRIMUS_DEFAULT_IFNAME} +# Default socket interface to loopback (single-node fallback). Override with +# GLOO_SOCKET_IFNAME / NCCL_SOCKET_IFNAME for multi-node (e.g. ens3 on amd-spur). +export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-lo} +export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-lo} export NCCL_IB_GID_INDEX=1 export HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM:-1} export NVTE_CK_USES_BWD_V3=${NVTE_CK_USES_BWD_V3:-1} @@ -243,6 +235,14 @@ fi mkdir -p "$PRIMUS_OUTPUT_ROOT/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME" +# On spur, direct the sbatch job's aggregated stdout+stderr (the actual per-node +# training log) into the experiment output dir. sbatch returns right after +# submission, so we can't `tee` it here; spur's sbatch reads SBATCH_OUTPUT/ERROR. +if command -v spur >/dev/null 2>&1; then + export SBATCH_OUTPUT="$PRIMUS_OUTPUT_ROOT/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/train_sbatch_output.log" + export SBATCH_ERROR="$PRIMUS_OUTPUT_ROOT/$PRIMUS_TEAM/$PRIMUS_USER/$PRIMUS_EXP_NAME/train_sbatch_error.log" +fi + export PRIMUS_EXIT_FAST=1 # Launcher: slurm (default, multi-node cluster) or direct (single-node, already @@ -252,10 +252,15 @@ export PRIMUS_LAUNCHER=${PRIMUS_LAUNCHER:-slurm} if [ "$PRIMUS_LAUNCHER" = "direct" ]; then LAUNCHER_ARGS=(direct) else - LAUNCHER_ARGS=(slurm -N "$NNODES") + LAUNCHER_ARGS=(slurm "${SLURM_LAUNCH_CMD:-srun}" -N "$NNODES") [ -n "${SLURM_PARTITION:-}" ] && LAUNCHER_ARGS+=(--partition="${SLURM_PARTITION}") [ -n "${SLURM_NODELIST:-}" ] && LAUNCHER_ARGS+=(--nodelist="${SLURM_NODELIST}") - LAUNCHER_ARGS+=(-- --image "${DOCKER_IMAGE}" --clean -- --numa) + [ -n "${SLURM_QOS:-}" ] && LAUNCHER_ARGS+=(--qos="${SLURM_QOS}") + [ -n "${SLURM_ACCOUNT:-}" ] && LAUNCHER_ARGS+=(--account="${SLURM_ACCOUNT}") + # --exclusive = whole-node allocation. On spur only sbatch accepts it (srun does + # not), so add it only in sbatch mode. Set SLURM_EXCLUSIVE=0 to disable. + [ "${SLURM_LAUNCH_CMD:-srun}" = "sbatch" ] && [ "${SLURM_EXCLUSIVE:-1}" != "0" ] && LAUNCHER_ARGS+=(--exclusive) + LAUNCHER_ARGS+=(-- --image "${DOCKER_IMAGE}" --clean -- --numa --patch runner/helpers/patches/10_fix_libionic_abi4.sh) fi ./primus-cli "${LAUNCHER_ARGS[@]}" \ diff --git a/examples/deepseek-v4/run_deepseek_v4_flash.sh b/examples/deepseek-v4/run_deepseek_v4_flash.sh old mode 100644 new mode 100755 index ed77ae0e2..054b7051c --- a/examples/deepseek-v4/run_deepseek_v4_flash.sh +++ b/examples/deepseek-v4/run_deepseek_v4_flash.sh @@ -2,6 +2,26 @@ set -euo pipefail +# On the spur cluster (amd-spur), set the launcher / QOS / networking defaults for +# the DeepSeek-V4 Flash multi-node runs. Detected via the `spur` command. +if command -v spur >/dev/null 2>&1; then + export PRIMUS_LAUNCHER=slurm + export SLURM_LAUNCH_CMD=sbatch + export SLURM_PARTITION=amd-spur + export SLURM_QOS=amd-burst-qos + export SLURM_ACCOUNT=amd-primus + # Empty = let the scheduler allocate nodes (skip the hardcoded smci355 default + # in run_deepseek_v4.sh, whose SLURM_NODELIST uses ${VAR-default} so empty wins). + export SLURM_NODELIST="" + # ABI-4 libionic provider .so to swap into the container at launch (fixes ionic + # RDMA on AINIC images whose bundled libionic only advertises uverbs ABI 1). + # The tools/patches/fix_libionic_abi4.sh patch reads this; set empty to disable. + export PRIMUS_LIBIONIC_SRC_ABI4_SO="bak/ainic/libionic-rdmav34.so.host-abi4/libionic.so.1.0.54.0-149.g3304be71" + export NCCL_DEBUG=INFO + export GLOO_SOCKET_IFNAME=ens3 + export NCCL_SOCKET_IFNAME=ens3 +fi + export PRIMUS_TOTAL_LAYERS=${PRIMUS_TOTAL_LAYERS:-43} export PRIMUS_NUM_EXPERTS=${PRIMUS_NUM_EXPERTS:-256} export PRIMUS_MOE_TOPK=${PRIMUS_MOE_TOPK:-6} @@ -10,7 +30,7 @@ export PRIMUS_INDEX_TOPK=${PRIMUS_INDEX_TOPK:-512} export PRIMUS_COMPRESS_RATIOS=${PRIMUS_COMPRESS_RATIOS:-'[0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 0]'} export MTP_NUM_LAYERS=${MTP_NUM_LAYERS:-1} -export NNODES=${NNODES:-8} +export NNODES=${NNODES:-4} if [ "$NNODES" -eq 8 ]; then export PRIMUS_TP=${PRIMUS_TP:-1} @@ -61,4 +81,4 @@ export PROFILE=${PROFILE:-False} export PRIMUS_EXP_NAME=${PRIMUS_EXP_NAME:-deepseek_v4_flash_proxy_pp${PRIMUS_PP}_ep${PRIMUS_EP}_seq${PRIMUS_SEQ_LENGTH}} SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -exec "${SCRIPT_DIR}/run_deepseek_v4.sh" +exec "${SCRIPT_DIR}/run_deepseek_v4.sh" 2>&1 | tee train_flash_ainic.log diff --git a/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh b/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh index 1ac99e498..655852e19 100755 --- a/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh +++ b/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh @@ -109,7 +109,7 @@ else fi # ---------- REQUIRED gfx1250 RCCL AVG->SUM workaround ----------------------- -export PYTHONPATH="$SCRIPT_DIR/rccl_avg_workaround:${PYTHONPATH:-}" +export PYTHONPATH="$SCRIPT_DIR/examples/deepseek-v4/rccl_avg_workaround:${PYTHONPATH:-}" # Real primus_turbo imports flydsl at import time; put FLYDSL_PKG_DIR on PYTHONPATH. export FLYDSL_PKG_DIR=${FLYDSL_PKG_DIR:-} if [ -n "$FLYDSL_PKG_DIR" ] && [ -d "$FLYDSL_PKG_DIR/flydsl" ]; then diff --git a/examples/deepseek-v4/run_dsv4_projection_1gpu.sh b/examples/deepseek-v4/run_dsv4_projection_1gpu.sh index d3e6b98d3..fe0ec6886 100755 --- a/examples/deepseek-v4/run_dsv4_projection_1gpu.sh +++ b/examples/deepseek-v4/run_dsv4_projection_1gpu.sh @@ -59,7 +59,7 @@ export NVTE_ROCM_ENABLE_MXFP8=${NVTE_ROCM_ENABLE_MXFP8:-1} # RCCL all_reduce(AVG) hangs even at world_size=1 -> sitecustomize rewrites AVG->SUM/ws. # Also put the vendored Emerging-Optimizers on PYTHONPATH so the muon optimizer # (emerging_optimizers.*) imports — required for OPTIMIZER=muon. -export PYTHONPATH_IN="$SCRIPT_DIR/rccl_avg_workaround:$SCRIPT_DIR/third_party/Emerging-Optimizers" +export PYTHONPATH_IN="$SCRIPT_DIR/examples/deepseek-v4/rccl_avg_workaround:$SCRIPT_DIR/third_party/Emerging-Optimizers" LOG=${LOG:-dsv4-projection-${MODE}.log} diff --git a/runner/helpers/patches/10_fix_libionic_abi4.sh b/runner/helpers/patches/10_fix_libionic_abi4.sh new file mode 100755 index 000000000..1ca367df7 --- /dev/null +++ b/runner/helpers/patches/10_fix_libionic_abi4.sh @@ -0,0 +1,47 @@ +#!/bin/bash +############################################################################### +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +# +# primus-cli --patch script: swap the ionic libibverbs provider (.so) for an +# ABI-4-capable build at container launch (before torchrun starts). +# +# Why: some AINIC container images ship a libionic provider built against stock +# rdma-core that only advertises ionic uverbs ABI 1, while the host kernel ionic +# driver exposes ABI 4 -- so libibverbs rejects every ionic_* device and RDMA +# falls back to TCP. Copying in an ABI-4 provider .so fixes device enumeration +# without rebuilding the image. +# +# Controlled by PRIMUS_LIBIONIC_SRC_ABI4_SO (set it in your launch script, e.g. +# run_flash.sh). If unset/empty the patch is skipped. The PRIMUS_ prefix makes it +# auto-forward into the container (see primus-cli-container.sh env passthrough). +############################################################################### +set -euo pipefail + +if [[ -z "${PRIMUS_LIBIONIC_SRC_ABI4_SO:-}" ]]; then + echo "[fix_libionic_abi4] PRIMUS_LIBIONIC_SRC_ABI4_SO not set -- skipping" + exit 2 # 2 = skip (not an error), per runner/helpers/execute_patches.sh +fi + +SRC="$PRIMUS_LIBIONIC_SRC_ABI4_SO" +if [[ ! -f "$SRC" ]]; then + echo "[fix_libionic_abi4] source .so not found: $SRC" >&2 + exit 1 +fi + +PROVIDER_LINK=/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so +if [[ ! -e "$PROVIDER_LINK" ]]; then + echo "[fix_libionic_abi4] ionic provider not present ($PROVIDER_LINK); nothing to patch" >&2 + exit 2 +fi +DST="$(readlink -f "$PROVIDER_LINK")" + +if cmp -s "$SRC" "$DST"; then + echo "[fix_libionic_abi4] provider already matches source -- skipping" + exit 2 +fi + +cp --remove-destination "$SRC" "$DST" +echo "[fix_libionic_abi4] swapped ionic provider: $DST <- $SRC" diff --git a/runner/primus-cli-container.sh b/runner/primus-cli-container.sh index 9a06e01f5..97928f15e 100755 --- a/runner/primus-cli-container.sh +++ b/runner/primus-cli-container.sh @@ -408,6 +408,14 @@ for key in "${!container_config[@]}"; do [[ "$opt_name" == "image" ]] && continue [[ "$opt_value" == "[]" ]] && continue + # Make the container name unique per job so an orphaned container left by a + # previous/cancelled job on a reused node doesn't cause a "name already in + # use" conflict. SLURM_JOB_ID is per-job (same on all its nodes, unique across + # jobs); fall back to the PID when not under Slurm. + if [[ "$opt_name" == "name" && -n "$opt_value" && "$opt_value" != *$'\n'* ]]; then + opt_value="${opt_value}-${SLURM_JOB_ID:-$$}" + fi + # Check if this is a cumulative option is_cumulative=0 for cum_opt in "${CUMULATIVE_OPTIONS[@]}"; do @@ -452,7 +460,10 @@ if [[ "$CLEAN_DOCKER_CONTAINER" == "true" ]]; then LOG_INFO_RANK0 "[container] Cleaning up existing containers..." CONTAINERS="$($CONTAINER_RUNTIME ps -aq)" if [[ -n "$CONTAINERS" ]]; then - printf '%s\n' "$CONTAINERS" | xargs -r -n1 "$CONTAINER_RUNTIME" rm -f + # Tolerate per-container removal failures ("removal already in progress", + # "No such container") so a stale/concurrently-removing container on a + # shared node does not abort the whole run via xargs' exit code 123. + printf '%s\n' "$CONTAINERS" | xargs -r -n1 -I{} sh -c "$CONTAINER_RUNTIME rm -f {} 2>/dev/null || true" LOG_INFO_RANK0 "[container] Removed containers: $CONTAINERS" else LOG_INFO_RANK0 "[container] No containers to remove." diff --git a/runner/primus-cli-slurm.sh b/runner/primus-cli-slurm.sh index 4c8194ffb..fb83f553b 100755 --- a/runner/primus-cli-slurm.sh +++ b/runner/primus-cli-slurm.sh @@ -294,8 +294,54 @@ fi ENTRY="$RUNNER_DIR/primus-cli-slurm-entry.sh" require_file "$ENTRY" "[slurm] Entry script not found: $ENTRY" -# Build full command -CMD=("$LAUNCH_CMD" "${SLURM_FLAGS[@]}" "$ENTRY" "${ENTRY_ARGS[@]}" -- "$@") +# Build full command. +# +# Scheduler-flavor handling for the sbatch path: +# - Standard Slurm `sbatch` accepts `script [args...]` (and runs the script on +# node 0 only), so we can pass the entry + args directly, as before. +# - Spur's `sbatch` reimplementation accepts ONLY one positional (the script): +# it rejects args passed to the batch script, has no --wrap, and runs the +# script on EVERY allocated node. So we can't use `sbatch ENTRY ` there +# (it fails with "unexpected argument '--image'"). Instead we emit a +# self-contained script that bakes in the entry + all args and runs the +# per-node entry directly (no inner srun -- that would nest job steps); QOS / +# partition / account are passed as sbatch options (which spur does accept). +# Spur is detected by the presence of the `spur` command. srun is unchanged on +# both flavors (its COMMAND accepts args and fans out per task). +IS_SPUR=false +command -v spur >/dev/null 2>&1 && IS_SPUR=true + +# Emit a self-contained batch-script body: `bash ENTRY -- ` +# with every token safely quoted so it re-runs verbatim inside the batch job. +_emit_spur_batch_script() { + # "$@" here are the Primus args (everything after the first '--'). + printf '#!/bin/bash\n' + printf 'set -euo pipefail\n' + printf 'exec bash %q' "$ENTRY" + local _a + for _a in "${ENTRY_ARGS[@]}"; do printf ' %q' "$_a"; done + printf ' --' + for _a in "$@"; do printf ' %q' "$_a"; done + printf '\n' +} + +JOB_SCRIPT="" +if [[ "$LAUNCH_CMD" == "sbatch" && "$IS_SPUR" == "true" ]]; then + if [[ "$DRY_RUN_MODE" == "true" ]]; then + LOG_INFO "[slurm] Spur detected: sbatch would submit this self-contained per-node script:" + _emit_spur_batch_script "$@" | sed 's/^/[slurm] | /' >&2 + LOG_INFO "[slurm] [DRY RUN] Would execute: $LAUNCH_CMD ${SLURM_FLAGS[*]} " + LOG_INFO "[slurm] Dry-run mode: command not executed" + exit 0 + fi + JOB_SCRIPT="$(mktemp "${TMPDIR:-/tmp}/primus-sbatch.XXXXXX.sh")" + _emit_spur_batch_script "$@" > "$JOB_SCRIPT" + chmod +x "$JOB_SCRIPT" + LOG_INFO "[slurm] Spur sbatch: submitting self-contained per-node script: $JOB_SCRIPT" + CMD=("$LAUNCH_CMD" "${SLURM_FLAGS[@]}" "$JOB_SCRIPT") +else + CMD=("$LAUNCH_CMD" "${SLURM_FLAGS[@]}" "$ENTRY" "${ENTRY_ARGS[@]}" -- "$@") +fi # Display command if [[ "$DRY_RUN_MODE" == "true" ]]; then @@ -305,4 +351,16 @@ if [[ "$DRY_RUN_MODE" == "true" ]]; then fi LOG_INFO "[slurm] Executing: ${CMD[*]}" -exec "${CMD[@]}" +if [[ -n "$JOB_SCRIPT" ]]; then + # Don't exec here: run sbatch, then remove the temp script. Spur copies the + # script into its spool dir at submit time, so it is safe to delete once + # sbatch returns. + set +e + "${CMD[@]}" + _rc=$? + set -e + rm -f "$JOB_SCRIPT" + exit "$_rc" +else + exec "${CMD[@]}" +fi diff --git a/skills/spur-cluster-status/SKILL.md b/skills/spur-cluster-status/SKILL.md new file mode 100644 index 000000000..c6ee019df --- /dev/null +++ b/skills/spur-cluster-status/SKILL.md @@ -0,0 +1,83 @@ +--- +name: spur-cluster-status +description: Inspect the current Spur (AMD SLURM-compatible) cluster node-allocation state and produce a Markdown report covering the caller's account/QoS permissions, partitions and per-state node counts, per-QoS and per-account node usage, reservations, queue pressure, GPU capacity, and the caller's own jobs. Use when the user asks about Spur/SLURM cluster status, node allocation, which QoS/account/partition holds how many nodes, how many idle nodes are available, or wants a cluster snapshot report. +--- + +# Spur Cluster Status Report + +Generate a read-only snapshot of the Spur cluster's node-allocation state plus the +caller's account/QoS permissions, and emit a single Markdown report (with an +appendix of common commands). Spur is AMD's SLURM-compatible scheduler exposed via +`sinfo` / `squeue` / `scontrol` / `spur accounts`. + +## Workflow + +### Step 1: Generate the report + +Run the collector from the repo root (it is read-only and takes a few seconds): + +```bash +python3 .claude/skills/spur-cluster-status/scripts/spur_status.py +``` + +- Reports on the current user by default. To target another user: `--user `. +- The script prints a complete Markdown report to stdout and never mutates cluster state. + +### Step 2: Save and present + +1. Save the output under the repo root (create the dir if missing): + +```bash +mkdir -p output/skills +python3 .claude/skills/spur-cluster-status/scripts/spur_status.py \ + > "output/skills/spur-cluster-status-$(date +%Y%m%d.%H%M).md" +``` + +2. Show the report to the user. Present the tables inline; the idle-node list can be + long, so summarize it (count + a short sample) unless the user wants the full list. +3. Print the saved file path. + +### Step 3: Add insights (optional) + +After the tables, add a short analysis when relevant, e.g.: +- Whether enough idle nodes exist for the user's target job size. +- Whether the user's jobs are on shared (`mix`) nodes (GPU-contention risk) vs `--exclusive`. +- Whether a low-priority QoS (e.g. `amd-burst-qos`, Prio=1) is being used for a job that + should use the account's normal QoS. + +## Report contents + +The collector emits these sections (see `scripts/spur_status.py`). The report is +written in **English**: + +1. **Account / QoS (you)** — the caller's account(s), default account, default QoS, the + permission-model note, and the global QoS list with priority/preempt. +2. **Partitions & Node States** — per-partition total + per-state node counts, overall + state breakdown, utilization %, GPU capacity, idle nodelist, and drain/down nodes. +3. **Node Usage by QoS** — for running jobs, per QoS: #jobs, distinct nodes, node-job + slots, GPUs, plus pending jobs and pending node demand. +4. **Node Usage by Account** — same breakdown grouped by account. +5. **Reservations** — name, node count, users, end time. +6. **Queue Pressure** — running vs pending job counts and pending node demand. +7. **My Jobs** — the caller's running/pending jobs. +8. **Jobs by User (all users)** — one row per user: total jobs (running/pending), + distinct nodes held, number of accounts + which, which QoS, and which partitions. +9. **Appendix: Common Commands** — the command reference used to build the report. + +## Spur quirks (important) + +These are baked into the collector; keep them in mind if you extend it: + +- `-o` format strings render **space-separated** regardless of literal delimiters — parse by whitespace column, not by a custom separator. +- `spur accounts show user|account ` ignores the positional filter and lists everything — filter with `grep`/`awk`. +- `sacctmgr` policy queries are blocked ("Please ask your administrator"); use `spur accounts show qos` for the QoS list. +- There is **no** `scontrol show hostnames`; expand compressed nodelists locally (the script does this). +- The association table is empty and **QoS is not enforced per account** — only the **account** is enforced at submit time. QoS selects scheduling priority, not permission. + +## Extending + +To add a metric, add a parser + a section in `build_report()` in +`scripts/spur_status.py`, and add the underlying command to `COMMANDS_APPENDIX` +so the report's command list stays in sync. Candidate additions: top users by node +count, largest contiguous idle block for N-node jobs, `spur report cluster` historical +utilization, and per-node `CPUAlloc` from `scontrol show node`. diff --git a/skills/spur-cluster-status/scripts/spur_status.py b/skills/spur-cluster-status/scripts/spur_status.py new file mode 100644 index 000000000..a8160925f --- /dev/null +++ b/skills/spur-cluster-status/scripts/spur_status.py @@ -0,0 +1,477 @@ +#!/usr/bin/env python3 +"""Collect Spur cluster node-allocation status and print a Markdown report. + +Targets the AMD "Spur" scheduler (SLURM-compatible CLI: sinfo/squeue/scontrol/ +spur accounts). Read-only. Handles Spur quirks: `-o` delimiters are rendered as +spaces, positional filters are ignored, there is no `scontrol show hostnames`, +and QoS is not enforced per account (only the account is). + +Usage: + python3 spur_status.py [--user USER] + +The report is printed to stdout (redirect it to save a file). +""" + +import argparse +import getpass +import re +import subprocess +from collections import defaultdict +from datetime import datetime + +# Node-state buckets used for the state breakdown / utilization math. +BUSY_STATES = {"alloc", "mix"} +FREE_STATES = {"idle"} +DOWN_STATES = {"down", "drain", "drained", "fail", "failing", "unknown"} + + +def run(cmd): + """Run a command list; return stdout text, or '' on any failure.""" + try: + out = subprocess.run(cmd, capture_output=True, text=True, timeout=60, check=False) + return out.stdout or "" + except Exception: + return "" + + +# --------------------------------------------------------------------------- +# Hostlist expansion (no `scontrol show hostnames` on Spur, so do it locally). +# Handles: "p-030", "p-[036,089,110]", "p-[036-040]", and top-level commas. +# --------------------------------------------------------------------------- +def _split_top(s): + parts, depth, cur = [], 0, "" + for ch in s: + if ch == "[": + depth += 1 + cur += ch + elif ch == "]": + depth -= 1 + cur += ch + elif ch == "," and depth == 0: + parts.append(cur) + cur = "" + else: + cur += ch + if cur: + parts.append(cur) + return parts + + +def _expand_one(tok): + m = re.match(r"^(.*?)\[([^\]]*)\](.*)$", tok) + if not m: + return [tok] if tok else [] + pre, body, post = m.groups() + out = [] + for item in body.split(","): + item = item.strip() + if "-" in item: + a, b = item.split("-", 1) + width = len(a) + for i in range(int(a), int(b) + 1): + out.append(f"{pre}{str(i).zfill(width)}{post}") + elif item: + out.append(f"{pre}{item}{post}") + return out + + +def expand_hostlist(s): + if not s or s in ("(null)", "N/A", "-"): + return [] + res = [] + for tok in _split_top(s): + res.extend(_expand_one(tok)) + return res + + +# --------------------------------------------------------------------------- +# Parsers +# --------------------------------------------------------------------------- +def parse_user_accounts(user): + """Return (accounts, default_acct, def_qos) for the given user.""" + text = run(["spur", "accounts", "show", "user"]) + accounts, default_acct, def_qos = [], "", "" + for line in text.splitlines(): + toks = line.split() + if len(toks) < 2 or toks[0] in ("User", "----"): + continue + if toks[0] != user: + continue + accounts.append(toks[1]) + if len(toks) >= 4: + default_acct = toks[3] + if len(toks) >= 5: + def_qos = toks[4] + return sorted(set(accounts)), default_acct, def_qos + + +def parse_qos(): + """Return list of dicts: {name, prio, preempt}.""" + text = run(["spur", "accounts", "show", "qos"]) + rows = [] + for line in text.splitlines(): + toks = line.split() + if len(toks) < 2 or toks[0] in ("Name", "----"): + continue + rows.append( + { + "name": toks[0], + "prio": toks[1] if len(toks) > 1 else "", + "preempt": toks[2] if len(toks) > 2 else "", + } + ) + return rows + + +def parse_nodes(): + """Return dict node -> {partition, state, gpus} from `sinfo -N`.""" + text = run(["sinfo", "-N", "-h", "-o", "%N %P %t %G"]) + nodes = {} + for line in text.splitlines(): + toks = line.split() + if len(toks) < 3: + continue + name, part, state = toks[0], toks[1], toks[2].rstrip("*") + gres = toks[3] if len(toks) > 3 else "" + gpus = gres.count("gpu:") + nodes[name] = {"partition": part, "state": state, "gpus": gpus} + return nodes + + +def parse_jobs(nodes): + """Return list of job dicts from squeue. + + Spur renders `-o` fields space-separated and DROPS empty fields (e.g. a job + with no QoS), which shifts columns and breaks naive positional parsing. So we + only trust the always-present leading fields (jobid, partition, account, user, + state, nnodes) and disambiguate the trailing optional tokens (qos and/or + nodelist) against the known node set: the token that expands to a real node is + the nodelist; the other is the qos. `account` is always present because the + scheduler enforces it at submit time. + """ + text = run(["squeue", "-h", "-o", "%i %P %a %u %T %D %q %N"]) + jobs = [] + for line in text.splitlines(): + toks = line.split() + if len(toks) < 6: + continue + jobid, partition, account, user, state = toks[:5] + nnodes = int(toks[5]) if toks[5].isdigit() else 0 + qos, nodelist = "", "" + for t in toks[6:]: + expanded = expand_hostlist(t) + if expanded and expanded[0] in nodes: + nodelist = t + else: + qos = t + jobs.append( + { + "jobid": jobid, + "partition": partition, + "qos": qos, + "account": account, + "user": user, + "state": state, + "nnodes": nnodes, + "nodelist": nodelist, + } + ) + return jobs + + +def parse_controller(): + """Extract the controller address from `scontrol show config`.""" + for line in run(["scontrol", "show", "config"]).splitlines(): + if "Addr=" in line: + return line.split("Addr=", 1)[1].strip() + return "n/a" + + +def parse_reservations(): + """Return list of reservation dicts from scontrol show reservation.""" + text = run(["scontrol", "show", "reservation"]) + resvs, cur = [], {} + for line in text.splitlines(): + if line.startswith("ReservationName="): + if cur: + resvs.append(cur) + cur = {} + for m in re.finditer(r"(\w+)=([^\s]+)", line): + cur[m.group(1)] = m.group(2) + if cur: + resvs.append(cur) + return resvs + + +# --------------------------------------------------------------------------- +# Report +# --------------------------------------------------------------------------- +def group_usage(jobs, nodes, key, state): + """Aggregate running/pending jobs by a key (qos/account). + + Returns dict key -> {jobs, node_slots, distinct_nodes(set), gpus}. + distinct_nodes/gpus are only meaningful for RUNNING jobs. + """ + agg = defaultdict(lambda: {"jobs": 0, "node_slots": 0, "nodes": set(), "gpus": 0}) + for j in jobs: + if j["state"] != state: + continue + k = j[key] or "(none)" + agg[k]["jobs"] += 1 + agg[k]["node_slots"] += j["nnodes"] + for n in expand_hostlist(j["nodelist"]): + if n in nodes: + agg[k]["nodes"].add(n) + for k, v in agg.items(): + v["gpus"] = sum(nodes[n]["gpus"] for n in v["nodes"]) + return agg + + +def h(title): + return f"\n## {title}\n" + + +def group_by_user(jobs, nodes): + """Aggregate every job per user across all states.""" + agg = defaultdict( + lambda: { + "total": 0, + "running": 0, + "pending": 0, + "nodes": set(), + "accounts": set(), + "qos": set(), + "partitions": set(), + } + ) + for j in jobs: + a = agg[j["user"]] + a["total"] += 1 + if j["state"] == "RUNNING": + a["running"] += 1 + elif j["state"] == "PENDING": + a["pending"] += 1 + a["accounts"].add(j["account"]) + a["qos"].add(j["qos"]) + a["partitions"].add(j["partition"]) + for n in expand_hostlist(j["nodelist"]): + if n in nodes: + a["nodes"].add(n) + return agg + + +def build_report(user): + now = datetime.now().strftime("%Y-%m-%d %H:%M") + accounts, default_acct, def_qos = parse_user_accounts(user) + qos_rows = parse_qos() + nodes = parse_nodes() + jobs = parse_jobs(nodes) + resvs = parse_reservations() + + L = [] + L.append("# Spur Cluster Node-Allocation Report") + L.append("") + L.append(f"- **Generated**: {now}") + L.append(f"- **Cluster**: spur (controller: {parse_controller()})") + L.append(f"- **Report user**: {user}") + + # 1. Account / QoS + L.append(h("1. Account / QoS (you)")) + L.append(f"- **Accounts**: {', '.join(accounts) if accounts else '(none found)'}") + L.append(f"- **Default account**: {default_acct or '-'}") + L.append(f"- **Default QoS**: {def_qos or '(unset)'}") + L.append("") + L.append( + "> Permission model (verified): the **account is enforced** - you can only " + "submit under the accounts listed above (submitting under another account " + "fails with `user ... is not associated with account`). **QoS is NOT restricted " + "per account**; any QoS from the global list below is accepted, so QoS selects " + "scheduling priority, not permission." + ) + L.append("") + L.append("Global QoS (higher Prio = higher priority):") + L.append("") + L.append("| QoS | Prio | Preempt |") + L.append("|-----|------|---------|") + for q in sorted(qos_rows, key=lambda x: -int(x["prio"]) if x["prio"].isdigit() else 0): + L.append(f"| {q['name']} | {q['prio']} | {q['preempt']} |") + + # 2. Partitions & node states + L.append(h("2. Partitions & Node States")) + part_state = defaultdict(lambda: defaultdict(int)) + part_total = defaultdict(int) + state_total = defaultdict(int) + gpu_total = gpu_free = 0 + for n, info in nodes.items(): + part_state[info["partition"]][info["state"]] += 1 + part_total[info["partition"]] += 1 + state_total[info["state"]] += 1 + gpu_total += info["gpus"] + if info["state"] in FREE_STATES: + gpu_free += info["gpus"] + all_states = sorted(state_total) + L.append("| Partition | Total | " + " | ".join(all_states) + " |") + L.append("|" + "---|" * (len(all_states) + 2)) + for p in sorted(part_total): + cells = " | ".join(str(part_state[p].get(s, 0)) for s in all_states) + L.append(f"| {p} | {part_total[p]} | {cells} |") + total_nodes = sum(part_total.values()) + busy = sum(state_total.get(s, 0) for s in BUSY_STATES) + L.append("") + L.append(f"- **Total nodes**: {total_nodes}") + L.append("- **State breakdown**: " + ", ".join(f"{s}={state_total[s]}" for s in all_states)) + if total_nodes: + L.append(f"- **Utilization (alloc+mix)**: {busy}/{total_nodes} = {busy*100//total_nodes}%") + L.append(f"- **GPU capacity (whole-node)**: {gpu_total} total, ~{gpu_free} free (on idle nodes)") + + # idle / unhealthy nodelist + idle_nodes = sorted(n for n, i in nodes.items() if i["state"] in FREE_STATES) + bad_nodes = sorted(n for n, i in nodes.items() if i["state"] in DOWN_STATES) + L.append("") + L.append(f"- **Idle nodes ({len(idle_nodes)})**: `{','.join(idle_nodes) if idle_nodes else 'none'}`") + if bad_nodes: + L.append(f"- **Drain/down nodes ({len(bad_nodes)})**: `{','.join(bad_nodes)}`") + + # 3. Per-QoS usage + L.append(h("3. Node Usage by QoS (running jobs)")) + L.append( + "> Nodes can be shared (mix), so the sum of per-QoS distinct nodes may exceed the number of physically busy nodes." + ) + L.append("") + qos_run = group_usage(jobs, nodes, "qos", "RUNNING") + qos_pend = group_usage(jobs, nodes, "qos", "PENDING") + L.append( + "| QoS | Running jobs | Distinct nodes | Node-job slots | GPUs | Pending jobs | Pending nodes req |" + ) + L.append("|-----|------|------|------|------|------|------|") + for k in sorted( + set(qos_run) | set(qos_pend), key=lambda x: -len(qos_run.get(x, {"nodes": set()})["nodes"]) + ): + r = qos_run.get(k, {"jobs": 0, "node_slots": 0, "nodes": set(), "gpus": 0}) + p = qos_pend.get(k, {"jobs": 0, "node_slots": 0}) + L.append( + f"| {k} | {r['jobs']} | {len(r['nodes'])} | {r['node_slots']} | {r['gpus']} | {p['jobs']} | {p['node_slots']} |" + ) + + # 4. Per-account usage + L.append(h("4. Node Usage by Account (running jobs)")) + acc_run = group_usage(jobs, nodes, "account", "RUNNING") + acc_pend = group_usage(jobs, nodes, "account", "PENDING") + L.append( + "| Account | Running jobs | Distinct nodes | Node-job slots | GPUs | Pending jobs | Pending nodes req |" + ) + L.append("|-----|------|------|------|------|------|------|") + for k in sorted( + set(acc_run) | set(acc_pend), key=lambda x: -len(acc_run.get(x, {"nodes": set()})["nodes"]) + ): + r = acc_run.get(k, {"jobs": 0, "node_slots": 0, "nodes": set(), "gpus": 0}) + p = acc_pend.get(k, {"jobs": 0, "node_slots": 0}) + L.append( + f"| {k} | {r['jobs']} | {len(r['nodes'])} | {r['node_slots']} | {r['gpus']} | {p['jobs']} | {p['node_slots']} |" + ) + + # 5. Reservations + L.append(h("5. Reservations")) + if resvs: + L.append("| Name | Nodes | Users | End time |") + L.append("|-----|------|------|------|") + for r in resvs: + nlist = expand_hostlist(r.get("Nodes", "")) + L.append( + f"| {r.get('ReservationName','?')} | {len(nlist)} | {r.get('Users','-')} | {r.get('EndTime','-')} |" + ) + else: + L.append("No active reservations.") + + # 6. Queue pressure + L.append(h("6. Queue Pressure")) + pend = [j for j in jobs if j["state"] == "PENDING"] + run_j = [j for j in jobs if j["state"] == "RUNNING"] + L.append(f"- **Running jobs**: {len(run_j)}") + L.append(f"- **Pending jobs**: {len(pend)} (requesting {sum(j['nnodes'] for j in pend)} nodes total)") + + # 7. My jobs + L.append(h("7. My Jobs (user=%s)" % user)) + mine = [j for j in jobs if j["user"] == user] + if mine: + L.append("| JobID | QoS | Account | State | Nodes | Nodelist |") + L.append("|-----|------|------|------|------|------|") + for j in mine: + L.append( + f"| {j['jobid']} | {j['qos']} | {j['account']} | {j['state']} | {j['nnodes']} | {j['nodelist'] or '-'} |" + ) + else: + L.append("No running/pending jobs.") + + # 8. Jobs by user (all users) + L.append(h("8. Jobs by User (all users)")) + L.append( + "> One row per user. `Jobs` = total (R running / P pending); `Nodes` = distinct nodes held by running jobs." + ) + L.append("") + by_user = group_by_user(jobs, nodes) + L.append("| User | Jobs (R/P) | Nodes | #Accts | Accounts | QoS | Partitions |") + L.append("|-----|------|------|------|------|------|------|") + for u in sorted(by_user, key=lambda x: (-len(by_user[x]["nodes"]), -by_user[x]["total"], x)): + a = by_user[u] + acct_disp = ",".join(sorted(a["accounts"])) + qos_disp = ",".join(sorted((q or "(none)") for q in a["qos"])) + part_disp = ",".join(sorted(a["partitions"])) + L.append( + f"| {u} | {a['total']} ({a['running']}R/{a['pending']}P) | {len(a['nodes'])} | " + f"{len(a['accounts'])} | {acct_disp} | {qos_disp} | {part_disp} |" + ) + + # Appendix: common commands + L.append(h("Appendix: Common Commands")) + L.append(COMMANDS_APPENDIX) + + return "\n".join(L) + "\n" + + +COMMANDS_APPENDIX = """```bash +# --- Account / QoS --- +spur accounts show user # all user->account rows (grep yourself) +sacctmgr show user $(whoami) # equivalent alias +spur accounts show qos # global QoS list (Prio/Preempt) +spur accounts show account # all accounts +sshare -u $(whoami) # fair-share info + +# --- Cluster / nodes --- +sinfo # partitions + node counts per state +sinfo -N -o "%N %P %t %G" # per node: name/partition/state/GRES +sinfo -N -h -o "%t" | sort | uniq -c # node count per state +scontrol show node # single-node detail (CPUAlloc/GRES/State) +scontrol show reservation # reservations + +# --- Jobs / queue --- +squeue # queue +squeue -o "%i %P %q %a %u %T %D %N" # custom columns (NOTE: delimiters render as spaces) +squeue -u $(whoami) # only your jobs +squeue -t PENDING # only pending jobs +scontrol show job # single-job detail (QOS/Account/NodeList) + +# --- Submit (whole-node exclusive) --- +sbatch -A amd-primus -q amd-primus-qos -p amd-spur --exclusive -N -t